Refactor diphtheria model to vectorised form, WIP #166, WIP #167

R/model_diphtheria.R

@@ -15,15 +15,15 @@
 #' group to model influxes or evacuations from camps.
 #'
 #' @inheritParams model_default
-#' @param reporting_rate A single number for the proportion of infectious cases
+#' @param reporting_rate A numeric for the proportion of infectious cases
 #' that is reported; this is a precursor to hospitalisation as only reported
 #' cases are hospitalised.
-#' @param prop_hosp A single number for the proportion of reported cases that is
+#' @param prop_hosp A numeric for the proportion of reported cases that is
 #' hospitalised.
-#' @param hosp_entry_rate A single number for the rate at which reported cases
+#' @param hosp_entry_rate A numeric for the rate at which reported cases
 #' of infectious individuals are hospitalised.
 #' This is calculated as 1 / time to hospitalisation, denoted \eqn{\tau_1}.
-#' @param hosp_exit_rate A single number for the rate at which individuals are
+#' @param hosp_exit_rate A numeric for the rate at which individuals are
 #' discharged from hospital to enter the 'recovered' compartment.
 #' This is calculated as 1 / time to discharge, denoted \eqn{\tau_2}.
 #' @param prop_vaccinated A numeric vector of the same length as the number of
@@ -151,16 +151,22 @@ model_diphtheria_cpp <- function(population,
                                  prop_hosp = 0.01,
                                  hosp_entry_rate = 0.2,
                                  hosp_exit_rate = 0.2,
-                                 prop_vaccinated = 0.0 * get_parameter(
-                                   population, "demography_vector"
-                                 ),
+                                 prop_vaccinated = 0.0 *
+                                   population[["demography_vector"]],
                                  intervention = NULL,
                                  time_dependence = NULL,
                                  population_change = NULL,
                                  time_end = 100,
                                  increment = 1) {
-  # check class on required inputs
+  # TODO: ensure population is properly vectorised
   checkmate::assert_class(population, "population")
+  # get compartment names
+  compartments <- c(
+    "susceptible", "exposed", "infectious", "hospitalised", "recovered"
+  )
+  assert_population(population, compartments)
+
+  # NOTE: model rates very likely bounded 0 - 1 but no upper limit set for now
   checkmate::assert_numeric(transmissibility, lower = 0, finite = TRUE)
   checkmate::assert_numeric(infectiousness_rate, lower = 0, finite = TRUE)
   checkmate::assert_numeric(recovery_rate, lower = 0, finite = TRUE)
@@ -188,107 +194,127 @@ model_diphtheria_cpp <- function(population,
     hosp_exit_rate = hosp_exit_rate,
     time_end = time_end
   )
-  stopifnot(
-    "All parameters must be of the same length, or must have length 1" =
-      test_recyclable(params)
-  )
+  # take parameter names here as names(DT) updates by reference!
+  param_names <- names(params)
 
   # check that prop_vaccinated is the same length as demography_vector
   # TODO: treat this as a composable element for vectorisation
   checkmate::assert_numeric(
     prop_vaccinated,
     lower = 0, upper = 1,
-    len = length(get_parameter(population, "demography_vector"))
+    len = length(population[["demography_vector"]])
   )
+  # convert to list for data.table
+  prop_vaccinated <- list(prop_vaccinated)
 
-  # allow only rate interventions in the intervention list
-  checkmate::assert_list(
+  # Check if parameters can be recycled;
+  # Check if `population` is a single population or a list of such
+  # and convert to list for a data.table list column;
+  # also check if `intervention` is a list of interventions or a list-of-lists
+  # and convert to a list for a data.table list column. NULL is allowed;
+  is_lofints <- checkmate::test_list(
+    intervention, "intervention",
+    all.missing = FALSE, null.ok = TRUE
+  )
+  # allow some NULLs (a valid no intervention scenario) but not all NULLs
+  is_lofls <- checkmate::test_list(
     intervention,
-    min.len = 1, names = "unique", any.missing = FALSE,
-    types = "rate_intervention", null.ok = TRUE
+    types = c("list", "null"), all.missing = FALSE
+  ) && all(
+    vapply(
+      unlist(intervention, recursive = FALSE),
+      FUN = function(x) {
+        is_intervention(x) || is.null(x)
+      }, TRUE
+    )
+  )
+
+  stopifnot(
+    "All parameters must be of the same length, or must have length 1" =
+      test_recyclable(params),
+    "`intervention` must be a list of <intervention>s or a list of such lists" =
+      is_lofints || is_lofls
   )
 
+  # make lists if not lists
+  population <- list(population)
+  if (is_lofints) {
+    intervention <- list(intervention)
+  }
+
   # check that time-dependence functions are passed as a list with at least the
-  # arguments `time` and `x`
-  # time must be before x, and they must be first two args
+  # arguments `time` and `x`, in order as the first two args
+  # NOTE: this functionality is not vectorised;
+  # convert to list for data.table list column
   checkmate::assert_list(
     time_dependence, "function",
     null.ok = TRUE,
-    names = "unique", any.missing = FALSE
+    any.missing = FALSE, names = "unique"
   )
+  # lapply on null returns an empty list
   invisible(
-    lapply(time_dependence, checkmate::check_function,
-      args = c("time", "x"),
-      ordered = TRUE, null.ok = TRUE
+    lapply(time_dependence, checkmate::assert_function,
+      args = c("time", "x"), ordered = TRUE
     )
   )
-
-  # check population change and create
-  checkmate::assert_list(
-    population_change,
-    null.ok = TRUE, len = 2L, types = c("numeric", "list")
-  )
-  if (!is.null(population_change)) {
-    checkmate::assert_names(
-      names(population_change),
-      identical.to = c("time", "values")
-    )
-  }
-
-  # collect population, infection
-  # combine parameters and composable elements into a list of lists of mod args
-  params <- .recycle_vectors(params)
-  params <- .transpose_base(params)
-  model_arguments <- lapply(
-    params, function(x) {
+  time_dependence <- list(
+    .cross_check_timedep(
+      time_dependence,
       c(
-        list(
-          population = population,
-          prop_vaccinated = prop_vaccinated,
-          intervention = intervention,
-          time_dependence = time_dependence,
-          pop_change_times = population_change[["time"]],
-          pop_change_values = population_change[["values"]],
-          increment = increment
-        ),
-        x
+        "transmissibility", "infectiousness_rate", "prop_hosp",
+        "reporting_rate", "hosp_entry_rate", "hosp_exit_rate", "recovery_rate"
       )
-    }
+    )
   )
 
-  # cross-check model arguments
-  # TODO: simplify to only check composable elements
-  model_arguments <- lapply(
-    model_arguments, function(l) {
-      .prepare_args_model_diphtheria(
-        .check_args_model_diphtheria(l)
-      )
-    }
-  )
+  # TODO: allow vectorised input for population_change
+  # convert to list for data.table
+  population_change <- list(population_change)
 
-  # get compartment names
-  compartments <- c(
-    "susceptible", "exposed", "infectious", "hospitalised", "recovered"
-  )
+  # collect parameters and add a parameter set identifier
+  params <- data.table::as.data.table(params)
+  params[, "param_set" := .I]
 
-  # run model over arguments, prepare output, and return list
-  # assign run number as list index - this is the specific combination of
-  # parameters
-  output <- Map(
-    model_arguments, seq_along(model_arguments),
-    f = function(l, i) {
-      output_ <- .output_to_df(
-        do.call(.model_diphtheria_cpp, l),
-        population = population,
-        compartments = compartments
-      )
-      output_$run <- i
-      output_
-    }
+  # this nested data.table will be returned
+  model_output <- data.table::CJ(
+    population = population,
+    prop_vaccinated = prop_vaccinated,
+    intervention = intervention,
+    time_dependence = time_dependence,
+    population_change = population_change,
+    increment = increment,
+    sorted = FALSE
   )
-  output <- data.table::rbindlist(output)
 
-  # convert to data.frame and return
-  # TODO: parameters need to be pulled along
-  data.table::setDF(output)[]
+  # process the population, interventions, and vaccinations, after
+  # cross-checking them agains the relevant population
+  model_output[, args := apply(model_output, 1, function(x) {
+    .check_prepare_args_diphtheria(c(x))
+  })]
+  model_output[, "scenario" := .I]
+
+  # combine infection parameters and scenarios
+  # NOTE: join X[Y] must have params as X as list cols not supported for X
+  model_output <- params[, as.list(model_output), by = names(params)]
+
+  # collect model arguments in column data, then overwrite
+  model_output[, args := apply(model_output, 1, function(x) {
+    c(x[["args"]], x[param_names]) # avoid including col "param_set"
+  })]
+  model_output[, data := Map(population, args, f = function(p, l) {
+    .output_to_df(
+      do.call(.model_diphtheria_cpp, l),
+      population = p, # taken from local scope/env
+      compartments = compartments
+    )
+  })]
+
+  # check for single row output, i.e., scalar arguments, and return data.frame
+  # do not return the parameters in this case
+  if (nrow(model_output) == 1L) {
+    model_output <- model_output[["data"]][[1L]] # hardcoded for special case
+  }
+
+  # return data.table
+  model_output[]
 }