Enhance extinction functionality (#161)

* remove cap_cases argument from detect_extinct and extinct_prob and assign cap_cases as attribute in scenario_sim

* rename outbreak_df_week argument to scenario

* replace as.data.table with setDT in detect_extinct

* rename week_range argument to extinction_week

* remove assigning returned variable in detect_extinct from rebasing branch

* add extinct_prob check for extinction_week in scenario data

* update default of extinction_week of extinct_prob and detect_extinct to last 2 weeks of outbreak

* merge documentation of extinct_prob and detect_extinct, relates #118

* remove outdated cap in extinct_prob test

* enable single integer or bound vector extinction_week and add function documentation, fixes #132

* DRY-out extinct_prob and detect_extinct

* append extinct attribute to output in outbreak_model and scenario_sim, fixes #132

* use extinct attribute if present in detect_extinct, fixes #132

* temporarily remove attribute from scenario_sim output in extinction unit tests

* update extinction_week default to NULL and update function documentation

* add messages to detect_extinct for method of extinction calculation

* check cap_cases attribute in detect_extinct

* add extinct_prob examples using different extinction_week input

* update extinct_prob test to fix extinction_week error

* update RoxygenNote in DESCRIPTION

* fix extinct variable scope issue

Author: joshwlambert

DESCRIPTION

@@ -42,7 +42,7 @@ Encoding: UTF-8
 LazyData: true
 Roxygen: list(markdown = TRUE, roclets = c("collate", "namespace", "rd",
     "roxyglobals::global_roclet"))
-RoxygenNote: 7.3.2
+RoxygenNote: 7.3.3
 Depends:
     R (>= 4.4.0)
 Imports: 

NAMESPACE

@@ -20,11 +20,12 @@ importFrom(data.table,.I)
 importFrom(data.table,.N)
 importFrom(data.table,.NGRP)
 importFrom(data.table,.SD)
-importFrom(data.table,as.data.table)
 importFrom(data.table,data.table)
 importFrom(data.table,fcase)
 importFrom(data.table,fifelse)
 importFrom(data.table,rbindlist)
+importFrom(data.table,setDT)
+importFrom(data.table,setattr)
 importFrom(sn,rsn)
 importFrom(stats,rnbinom)
 importFrom(stats,runif)

R/aux_functions.R

@@ -111,67 +111,63 @@ presymptomatic_transmission_to_alpha <- function(presymptomatic_transmission) {
   res$minimum
 }
 
-#' Calculate proportion of runs that have controlled outbreak
+#' Outbreak extinction functions
 #'
-#' @inherit detect_extinct details
+#' @description
+#' `extinct_prob()`: Calculate proportion of runs that have controlled outbreak
 #'
-#' @inheritParams detect_extinct
+#' `detect_extinct()`: Calculate whether outbreaks went extinct or not
 #'
-#' @return a single `numeric` with the probability of extinction
-#' @export
+#' @details
+#' The data passed to `scenario` has to be produced by [scenario_sim()].
+#' It cannot be produced by [outbreak_model()] as it requires the `sim` column,
+#' which is only appended in [scenario_sim()].
 #'
-#' @examples
-#' res <- scenario_sim(
-#'   n = 10,
-#'   initial_cases = 1,
-#'   offspring = offspring_opts(
-#'     community = \(n) rnbinom(n = n, mu = 2.5, size = 0.16),
-#'     isolated = \(n) rnbinom(n = n, mu = 0.5, size = 1)
-#'   ),
-#'   delays = delay_opts(
-#'     incubation_period = \(n) rweibull(n = n, shape = 2.32, scale = 6.49),
-#'     onset_to_isolation = \(n) rweibull(n = n, shape = 1.65, scale = 4.28)
-#'   ),
-#'   event_probs = event_prob_opts(
-#'     asymptomatic = 0,
-#'     presymptomatic_transmission = 0.5,
-#'     symptomatic_ascertained = 0.2
-#'   ),
-#'   interventions = intervention_opts(quarantine = FALSE),
-#'   sim = sim_opts(cap_max_days = 350, cap_cases = 4500)
-#' )
-#' extinct_prob(res, cap_cases = 4500)
-extinct_prob <- function(outbreak_df_week, cap_cases, week_range = 12:16) {
-
-  checkmate::assert_data_frame(outbreak_df_week)
-  checkmate::assert_number(cap_cases, lower = 0)
-  checkmate::assert_numeric(week_range)
-
-  n <- max(outbreak_df_week$sim)
-
-  extinct_runs <- detect_extinct(outbreak_df_week, cap_cases, week_range)
-  sum(extinct_runs$extinct) / n
-}
-
-
-#' Calculate whether outbreaks went extinct or not
+#' ***Warning***: the output from [scenario_sim()] contains an `cap_cases`
+#' attribute which is used by [extinct_prob()] and [detect_extinct()],
+#' therefore if you modify the output of [scenario_sim()] before passing
+#' to [extinct_prob()] be careful not to drop the attribute (e.g.
+#' from subsetting the `data.table`).
 #'
-#' @details
-#' The `cap_cases` argument should be equal to the value supplied to
-#' [outbreak_model()] (possibly passed from [scenario_sim()]).
+#' @param scenario a `data.table`: weekly cases output by [scenario_sim()]
+#' @param extinction_week By default `NULL` but also accepts a positive
+#'   `integer` scalar or `integer` vector to test if the outbreak has gone
+#'   extinct (i.e. no new cases) by this week (zero indexed):
+#'   * By default (`NULL`) the extinction status is stored in the output of
+#'     `scenario_sim()` which is supplied to the `scenario` argument. If
+#'     `extinction_week` is not specified or specified as `NULL` then the
+#'     pre-computed extinction status from the outbreak simulation will be used.
+#'     This is defined as all infectious cases have had the opportunity to
+#'     transmit but no new cases are generated.
+#'   * A single `integer` to test if extinction has occurred by this week.
+#'     For example, `extinction_week = 5` tests whether the outbreak is
+#'     extinct by week 5 (inclusive) until the end of the outbreak.
+#'   * An `integer` vector of length two can be supplied to provide the lower
+#'     and upper bounds (inclusive) of the week range to test for whether
+#'     extinction occurred by this window. For example
+#'     `extinction_week = c(5, 10)` will test whether the outbreak went extinct
+#'     by week 5 and there we no new cases between weeks 5 and 10 (inclusive).
+#'   * An `integer` vector of length _n_ can be supplied to provide the weeks
+#'     to test for whether extinction occurred by this window. For example
+#'     `extinction_week = 12:16` will test that there are no new infections
+#'     between 12 and 16 weeks after the initial cases (inclusive). These
+#'     integer sequences will most likely be contiguous but the function
+#'     does allow non-contiguous integer sequences.
+#'
+#'   If extinction occurs before the `extinction_week` window then the outbreak
+#'   extinction is considered extinct, however, if the extinction occurs within
+#'   the `extinction_week` window it is not considered extinct. Therefore,
+#'   using a single `integer` for `extinction_week` and thinking of this as
+#'   "_has the outbreak gone extinct by week X_".
 #'
-#' @param outbreak_df_week a `data.table`: weekly cases produced by the
-#'   outbreak model
-#' @inheritParams sim_opts
-#' @param week_range a positive `integer` vector: giving the (zero indexed)
-#'   week range to test for whether an extinction occurred. Default is `12:16`.
-#' @importFrom data.table as.data.table fifelse
+#' @importFrom data.table setDT fifelse data.table
 #'
-#' @return A `data.table`, with two columns `sim` and `extinct`, for a binary
+#' @return
+#' `extinct_prob()`: a single `numeric` with the probability of extinction
+#'
+#' `detect_extinct()`: a `data.table`, with two columns `sim` and `extinct`, for a binary
 #' classification of whether the outbreak went extinct in each simulation
 #' replicate. `1` is an outbreak that went extinct, `0` if not.
-#' @autoglobal
-#' @export
 #'
 #' @examples
 #' res <- scenario_sim(
@@ -193,16 +189,89 @@ extinct_prob <- function(outbreak_df_week, cap_cases, week_range = 12:16) {
 #'   interventions = intervention_opts(quarantine = FALSE),
 #'   sim = sim_opts(cap_max_days = 350, cap_cases = 4500)
 #' )
-#' detect_extinct(outbreak_df_week = res, cap_cases = 4500)
-detect_extinct <- function(outbreak_df_week, cap_cases, week_range = 12:16) {
+#'
+#' # calculate probability of extinction
+#' extinct_prob(res)
+#'
+#' # determine if each outbreak simulation replicate has gone extinct
+#' detect_extinct(res)
+#'
+#' # calculate extinction in the last 2 weeks of the simulated outbreak
+#' # (i.e. the penultimate and last week of the outbreak)
+#' extinct_prob(res, extinction_week = max(res$week) - 1)
+#'
+#' # calculate extinction as no new cases between weeks 12 and 16 of the outbreak
+#' extinct_prob(res, extinction_week = 12:16)
+#' @name extinction
+NULL
 
-  checkmate::assert_data_frame(outbreak_df_week)
-  checkmate::assert_number(cap_cases, lower = 0)
-  checkmate::assert_integerish(week_range)
+#' @rdname extinction
+#' @export
+extinct_prob <- function(scenario,
+                         extinction_week = NULL) {
+
+  extinct_runs <- detect_extinct(
+    scenario = scenario,
+    extinction_week = extinction_week
+  )
+  sum(extinct_runs$extinct) / max(scenario$sim)
+}
+
+#' @rdname extinction
+#' @autoglobal
+#' @export
+detect_extinct <- function(scenario,
+                           extinction_week = NULL) {
+
+  extinct <- attr(scenario, which = "extinct", exact = TRUE)
+  if (is.null(extinction_week)) {
+    if (!is.null(extinct)) {
+      message(
+        "Calculating extinction using the extinction status from ",
+        "the simulation."
+      )
+      return(
+        data.table(
+          sim = 1:max(scenario$sim),
+          extinct = as.integer(extinct)
+        )
+      )
+    } else {
+      stop(
+        "`extinction_week` not specified and `scenario` is missing the ",
+        "`extinct` attribute.\n Use `scenario_sim()` to simulate `scenario`, ",
+        "or specify `extinction_week`.",
+        call. = FALSE
+      )
+    }
+  }
+
+  checkmate::assert_data_frame(scenario)
+  checkmate::assert_integerish(extinction_week, min.len = 1)
+
+  if (length(extinction_week) == 1) {
+    extinction_week <- extinction_week:max(scenario$week)
+  } else if (length(extinction_week) == 2) {
+    extinction_week <- min(extinction_week):max(extinction_week)
+  }
+  stopifnot(
+    "`extinction_week` not in simulated outbreak data" =
+      all(extinction_week %in% scenario$week)
+  )
+  message(
+    "Calculating extinction as no new cases within weeks: ",
+    min(extinction_week), " to ", max(extinction_week), " (inclusive)."
+  )
+
+  cap_cases <- attr(scenario, which = "cap_cases", exact = TRUE)
+  stopifnot(
+    "`scenario` is missing the `cap_cases` attribute.
+    Use `scenario_sim()` to simulate `scenario`" = !is.null(cap_cases)
+  )
 
-  outbreak_df_week <- as.data.table(outbreak_df_week)
-  outbreak_df_week <- outbreak_df_week[week %in% week_range]
-  outbreak_df_week[, list(
+  scenario <- setDT(scenario)
+  scenario <- scenario[week %in% extinction_week]
+  scenario[, list(
     extinct = fifelse(all(weekly_cases == 0 & cumulative < cap_cases), 1, 0)
   ), by = sim][]
 }

R/outbreak_model.R

@@ -15,7 +15,7 @@
 #' @autoglobal
 #' @export
 #'
-#' @importFrom data.table rbindlist
+#' @importFrom data.table rbindlist setattr
 #'
 #' @examples
 #' set.seed(1)
@@ -143,6 +143,9 @@ outbreak_model <- function(initial_cases,
   weekly_cases <- weekly_cases[, `:=`(effective_r0 = mean(effective_r0_vect,
                                                           na.rm = TRUE),
                                         cases_per_gen = list(cases_in_gen_vect))]
+
+  setattr(weekly_cases, name = "extinct", value = all(case_data$isolated))
+
   # return
   weekly_cases[]
 }

R/scenario_sim.R

@@ -5,7 +5,7 @@
 #' @inheritParams outbreak_step
 #' @inheritParams outbreak_model
 #'
-#' @importFrom data.table rbindlist
+#' @importFrom data.table rbindlist setattr
 #' @return A `data.table` object returning the results for multiple simulations
 #'   using the same set of parameters. The table has columns
 #' * week: The week in the simulation.
@@ -77,6 +77,15 @@ scenario_sim <- function(n,
     simplify = FALSE
   )
 
-  # bind output together and add simulation index and return
-  data.table::rbindlist(res, idcol = "sim")[]
+  extinct <- vapply(
+    res, attr, FUN.VALUE = logical(1), which = "extinct", exact = TRUE
+  )
+
+  # bind output together and add simulation index
+  res <- data.table::rbindlist(res, idcol = "sim")
+
+  setattr(res, name = "cap_cases", value = sim$cap_cases)
+  setattr(res, name = "extinct", value = extinct)
+
+  res[]
 }

README.Rmd

@@ -99,7 +99,7 @@ ggplot(
 ### Estimate extinction probability
 
 ```{r extinct_prob}
-extinct_prob(res, cap_cases = 4500)
+extinct_prob(res)
 ```
 
 ## Contributors