Final size with susceptibility groups (#42)

R implementation of final size of an epidemic with susceptibility groups.

* WIP - first draft finalsize_grps in R

* WIP - first draft tests for finalsize_grps in R

* WIP - missing comma for args

* Split off epi_spread and remove Newton solver

* Implement iterative solver for susc grps, fixes #44

* Use iterative solver in final_size_grps, fixes #45

* Test for epi_spread, WIP #46, #47

* Add tests for iterative solver, fixes #46

* Test final_size_grps, fixes #47

* Documentation for final_size_grps and related fns, fixes #48

* Add extra test for epi_spread, WIP #46, #47

* Minor refactor of p_susc in epi_spread_data, remove commented code

* Housekeeping: No explicit namespacing for testthat

* Kronecker con_mat replication, vectors instead of mats, fixes #49

* Vectorise con_mat filling in iterative solver, fixes #51

* Vectorise fn_f in iterative solver, and pi adjust, fixes #52

* Squash stopifnot calls into one, fixes #50

* Pass function args re solver steps, solves #53

* Update tests for epi_spread and final_size_grps

* Refactor final size from pi to epi_final_size, fixes #54

* Remove test for Newton solver

NAMESPACE

@@ -2,7 +2,7 @@
 
 export(final_size)
 export(final_size_cpp)
-exportPattern("^[[:alpha:]]+")
+export(final_size_grps)
 import(RcppEigen)
 importFrom(Rcpp,evalCpp)
 useDynLib(finalsize)

R/epi_spread.R

@@ -0,0 +1,47 @@
+#' @title Prepare contact and demography data for multiple risk groups
+#'
+#' @param contact_matrix Social contact matrix. Entry $mm_{ij}$ gives average
+#' number of contacts in group $i$ reported by participants in group j
+#' @param demography_vector Demography vector. Entry $pp_{i}$ gives proportion
+#' of total population in group $i$ (model will normalise if needed)
+#' @param p_susceptibility A matrix giving the probability that an individual
+#' in demography group $i$ is in risk (or susceptibility) group $j$.
+#' Each row represents the overall distribution of individuals in demographic
+#' group $i$ across risk groups, and each row *must sum to 1.0*.
+#' @param susceptibility A matrix giving the susceptibility of individuals in
+#' demographic group $i$ and risk group $j$.
+#'
+#' @return A list object with named elements: "contact_matrix",
+#' "demography_vector", "p_susceptibility_", and "susceptibility".
+#' The contact matrix is replicated row and column wise for each risk group
+#' and the demography vector is replicated for each risk group.
+epi_spread <- function(contact_matrix,
+                       demography_vector,
+                       p_susceptibility,
+                       susceptibility) {
+  # count risk groups
+  n_susc_groups <- ncol(p_susceptibility)
+  # make p_susceptibility matrix of ones
+  p_susceptibility_ <- rep(1.0, length(p_susceptibility))
+  # make lps, a 1 col matrix of all p_susc values
+  lps <- as.vector(p_susceptibility)
+  # replicate the demography vector and multiply by p_susceptibility
+  demography_vector_spread <- rep(demography_vector, n_susc_groups)
+  demography_vector_spread <- demography_vector_spread * lps
+
+  # replicate contact matrix
+  contact_matrix_spread <- kronecker(
+    X = matrix(1, nrow = n_susc_groups, ncol = n_susc_groups),
+    Y = contact_matrix
+  )
+
+  # unroll the susceptibility matrix
+  susceptibility_ <- as.vector(susceptibility)
+
+  list(
+    contact_matrix = contact_matrix_spread,
+    demography_vector = demography_vector_spread,
+    p_susceptibility = p_susceptibility_,
+    susceptibility = susceptibility_
+  )
+}

R/finalsize_risk_grps.R

@@ -0,0 +1,113 @@
+#' Final size of epidemic by susceptibility groups
+#'
+#' @param contact_matrix Social contact matrix. Entry $mm_{ij}$ gives average
+#' number of contacts in group $i$ reported by participants in group j
+#' @param demography_vector Demography vector. Entry $pp_{i}$ gives proportion
+#' of total population in group $i$ (model will normalise if needed)
+#' @param p_susceptibility A matrix giving the probability that an individual
+#' in demography group $i$ is in risk (or susceptibility) group $j$.
+#' Each row represents the overall distribution of individuals in demographic
+#' group $i$ across risk groups, and each row *must sum to 1.0*.
+#' @param susceptibility A matrix giving the susceptibility of individuals in
+#' demographic group $i$ and risk group $j$.
+#' @param iterations Number of solver iterations. Defaults to 1,000.
+#' @param tolerance Solver error tolerance.
+#' @param step_rate The solver step rate. Defaults to 1.9 as a value found to
+#' work well.
+#' @param adapt_step Boolean, whether the solver step rate should be changed
+#' based on the solver error. Defaults to TRUE.
+#'
+#' @keywords epidemic model
+#' @return A vector of final sizes by demography group.
+#' @export
+#' @examples
+#' library(socialmixr)
+#' data(polymod)
+#' r0 <- 2.0
+#' contact_data <- contact_matrix(
+#'   polymod,
+#'   countries = "United Kingdom",
+#'   age.limits = c(0, 20, 40)
+#' )
+#' c_matrix <- t(contact_data$matrix)
+#' d_vector <- contact_data$participants$proportion
+#' # Scale contact matrix to demography
+#' c_matrix <- apply(
+#'   c_matrix, 1, function(r) r / d_vector
+#' )
+#' n_demo_grps <- length(d_vector)
+#' n_risk_grps <- 3
+#' # prepare p_susceptibility and susceptibility
+#' psusc <- matrix(
+#'   data = 1, nrow = n_demo_grps, ncol = n_risk_grps
+#' )
+#' psusc <- t(apply(psusc, 1, \(x) x / sum(x)))
+#' susc <- matrix(
+#'   data = 1, nrow = n_demo_grps, ncol = n_risk_grps
+#' )
+#' final_size_grps(
+#'   contact_matrix = r0 * c_matrix,
+#'   demography_vector = d_vector,
+#'   p_susceptibility = psusc,
+#'   susceptibility = susc
+#' )
+#'
+final_size_grps <- function(contact_matrix,
+                            demography_vector,
+                            p_susceptibility,
+                            susceptibility,
+                            iterations = 1000,
+                            tolerance = 1e-6,
+                            step_rate = 1.9,
+                            adapt_step = TRUE) {
+  # check arguments input
+  stopifnot(
+    "Error: contact matrix must have as many rows as demography groups" =
+      (nrow(contact_matrix) == length(demography_vector)),
+    "Error: p_susceptibility must have as many rows as demography groups" =
+      (nrow(p_susceptibility) == length(demography_vector)),
+    "Error: susceptibility must have as many rows as demography groups" =
+      (nrow(susceptibility) == length(demography_vector)),
+    "Error: susceptibility must have same dimensions as p_susceptibility" =
+      (all(dim(p_susceptibility) == dim(susceptibility))),
+    "Error: p_susceptibility rows must sum to 1.0" =
+      (
+        all(abs(rowSums(p_susceptibility) - 1) < 1e-6)
+      )
+  )
+
+  # prepare epi spread object
+  epi_spread_data <- epi_spread(
+    contact_matrix = contact_matrix,
+    demography_vector = demography_vector,
+    p_susceptibility = p_susceptibility,
+    susceptibility = susceptibility
+  )
+
+  # solve for final size
+  epi_final_size <- solve_final_size_iterative(
+    contact_matrix = epi_spread_data[["contact_matrix"]],
+    demography_vector = epi_spread_data[["demography_vector"]],
+    p_susceptibility = epi_spread_data[["p_susceptibility"]],
+    susceptibility = epi_spread_data[["susceptibility"]],
+    iterations = iterations,
+    tolerance = tolerance,
+    step_rate = step_rate,
+    adapt_step = adapt_step
+  )
+
+  # unroll p_susceptibility data
+  lps <- as.vector(p_susceptibility)
+
+  # multiply demo-risk specific final sizes by corresponding pop proportions
+  epi_final_size <- epi_final_size * lps
+
+  # final sizes mapped to matrix with dims (n_demo_grp, n_risk_grps)
+  epi_final_size <- matrix(
+    data = epi_final_size,
+    nrow = length(demography_vector),
+    ncol = ncol(p_susceptibility)
+  )
+  # return row-wise sum, i.e., the demo-grp wise sum
+  rowSums(epi_final_size)
+}

R/iterative_solver.R

@@ -0,0 +1,107 @@
+#' Iterative solver implemented in R
+#'
+#' @param contact_matrix Social contact matrix. Entry $mm_{ij}$ gives average
+#' number of contacts in group $i$ reported by participants in group j
+#' @param demography_vector Demography vector. Entry $pp_{i}$ gives proportion
+#' of total population in group $i$ (model will normalise if needed)
+#' @param p_susceptibility A matrix giving the probability that an individual
+#' in demography group $i$ is in risk (or susceptibility) group $j$.
+#' Each row represents the overall distribution of individuals in demographic
+#' group $i$ across risk groups, and each row *must sum to 1.0*.
+#' @param susceptibility A matrix giving the susceptibility of individuals in
+#' demographic group $i$ and risk group $j$.
+#' @param iterations Number of solver iterations. Defaults to 1,000.
+#' @param tolerance Solver error tolerance.
+#' @param step_rate The solver step rate. Defaults to 1.9 as a value found to
+#' work well.
+#' @param adapt_step Boolean, whether the solver step rate should be changed
+#' based on the solver error. Defaults to TRUE.
+#'
+#' @return A vector of final sizes, of the size (N demography groups *
+#' N risk groups).
+solve_final_size_iterative <- function(contact_matrix,
+                                       demography_vector,
+                                       p_susceptibility,
+                                       susceptibility,
+                                       iterations = 1000,
+                                       tolerance = 1e-6,
+                                       step_rate = 1.9,
+                                       adapt_step = TRUE) {
+  # count demography groups
+  n_dim <- length(demography_vector)
+
+  # make vector of zeros
+  zeros <- rep(0.0, n_dim)
+  # make vector of initial final size guesses = 0.5
+  epi_final_size <- rep(0.5, n_dim)
+
+  # replicate contact matrix
+  contact_matrix_ <- 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 == 0 | susceptibility == 0 |
+    rowSums(contact_matrix) == 0
+  zeros[i_here] <- 1.0
+  epi_final_size[i_here] <- 0.0
+
+  # matrix filled by columns
+  contact_matrix_ <- matrix(
+    as.vector(contact_matrix_) *
+      (susceptibility %x% demography_vector), # note Kronecker product
+    nrow = nrow(contact_matrix_),
+    ncol = ncol(contact_matrix_)
+  )
+
+  contact_matrix_[i_here, zeros == 1] <- 0.0
+
+  # make a vector to hold final size, empty numeric of size n_dim
+  epi_final_size_return <- numeric(n_dim)
+
+  # define functions to minimise error in final size estimate
+  fn_f <- function(epi_final_size_, epi_final_size_return_) {
+    # contact_matrix_ captured from environment
+    s <- as.vector(contact_matrix_ %*% (-epi_final_size_))
+
+    epi_final_size_return_ <- ifelse(zeros == 1.0, 0.0, 1.0)
+    epi_final_size_return_ <- epi_final_size_return_ - (p_susceptibility *
+      exp(susceptibility * s))
+
+    epi_final_size_return_
+  }
+  # define initial current error
+  current_error <- step_rate * n_dim
+  # run over fn_f over epi_final_size (intial guess)
+  # and epi_final_size_return (current estimate?)
+  for (i in seq(iterations)) {
+    epi_final_size_return <- fn_f(epi_final_size, epi_final_size_return)
+    # get current error
+    dpi <- epi_final_size - epi_final_size_return
+    error <- sum(abs(dpi))
+    # break loop if error below tolerance
+    if (error < tolerance) {
+      epi_final_size <- epi_final_size - dpi
+      break
+    }
+    # adapt the step size based on the change in error
+    change <- current_error - error
+    if (change > 0.0) {
+      epi_final_size <- epi_final_size - step_rate * dpi
+      if (adapt_step) {
+        step_rate <- step_rate * 1.1
+      }
+    } else {
+      epi_final_size <- epi_final_size - dpi
+      if (adapt_step) {
+        step_rate <- max(0.9 * step_rate, 1.0)
+      }
+    }
+    current_error <- error
+  }
+
+  # adjust numerical errors
+  # relies on TRUE equivalent to 1
+  epi_final_size <- ifelse(zeros, 0.0, epi_final_size)
+
+  # return what
+  epi_final_size
+}