Process contact matrix in main fn, WIP #68

R/final_size.R

@@ -189,6 +189,19 @@ final_size <- function(r0,
     Y = r0 * contact_matrix
   )
 
+  # set contact_matrix values to zero if there are no contacts among ...
+  # demography groups, or if demography groups are empty
+  i_here <- demography_vector_spread == 0 | as.vector(susceptibility) == 0 |
+    rowSums(contact_matrix_spread) == 0
+  contact_matrix_spread[i_here, i_here] <- 0.0
+
+  # process contact matrix in solver specific ways
+  contact_matrix_spread <- switch(solver,
+    iterative = t(t(contact_matrix_spread) * demography_vector_spread),
+    newton = t(t(contact_matrix_spread * as.vector(susceptibility)) *
+      demography_vector_spread)
+  )
+
   # get group wise final sizes
   epi_final_size <- .final_size(
     c(