Offspring distributions sample once per generation (#169)

* add sampled col to case_data and check in offspring sampling, relates #148

* remove subsetting by exposures before infector isolation time in outbreak_step, relates #148

* add sample_offspring function

* change isolated to sampled in case_data and make new_cases an integer

* update outbreak_step to call sample_offspring and use sampled instead of isolated

* update globals.R

* add alpha and latent_period arguments to sample_offspring and update incubation_to_generation_time calls

* update sample_offspring call in outbreak_step

* fix outbreak_continue to use sampled after merging main into iso-offspring-fix

* update globals.R

* remove old commented code from outbreak_step

* fix bug in outbreak_step early return and simplify output list

* fix ordering of generation times when constructing prob_samples in outbreak_step

* $isolated -> $sampled for outbreak_model extinct attribute

* make a copy of case_data in outbreak_step due to sample_offspring modifying by reference

* update outbreak_setup test expectations and snapshot

* update outbreak_step tests given new offspring sampling

* update scenario_sim snapshot

Author: joshwlambert

NAMESPACE

@@ -20,6 +20,7 @@ importFrom(data.table,.I)
 importFrom(data.table,.N)
 importFrom(data.table,.NGRP)
 importFrom(data.table,.SD)
+importFrom(data.table,copy)
 importFrom(data.table,data.table)
 importFrom(data.table,fcase)
 importFrom(data.table,fifelse)

R/aux_functions.R

@@ -299,7 +299,7 @@ outbreak_continue <- function(case_data, sim) {
 
   latest_onset <- max(case_data$onset)
   total_cases <- nrow(case_data)
-  extinct <- all(case_data$isolated)
+  extinct <- all(case_data$sampled)
 
   return(
     latest_onset < sim$cap_max_days &&

R/globals.R

@@ -16,14 +16,15 @@ utils::globalVariables(c(
   "isolated_time", # <outbreak_setup>
   "onset", # <outbreak_setup>
   "new_cases", # <outbreak_step>
-  "isolated", # <outbreak_step>
+  "sampled", # <outbreak_step>
+  "caseid", # <outbreak_step>
+  "isolated_time", # <outbreak_step>
   "asymptomatic", # <outbreak_step>
   "onset", # <outbreak_step>
   "exposure", # <outbreak_step>
-  "caseid", # <outbreak_step>
-  "isolated_time", # <outbreak_step>
-  "infector_isolation_time", # <outbreak_step>
   "infector_asymptomatic", # <outbreak_step>
   "missed", # <outbreak_step>
+  "infector_isolation_time", # <outbreak_step>
+  "sampled", # <sample_offspring>
   NULL
 ))

R/outbreak_model.R

@@ -144,7 +144,7 @@ outbreak_model <- function(initial_cases,
                                                           na.rm = TRUE),
                                         cases_per_gen = list(cases_in_gen_vect))]
 
-  setattr(weekly_cases, name = "extinct", value = all(case_data$isolated))
+  setattr(weekly_cases, name = "extinct", value = all(case_data$sampled))
 
   # return
   weekly_cases[]

R/outbreak_setup.R

@@ -11,9 +11,9 @@
 #' * `$infector`: `numeric`
 #' * `$missed`: `logical`
 #' * `$onset`: `numeric`
-#' * `$new_cases`: `logical`
+#' * `$new_cases`: `integer`
 #' * `$isolated_time`: `numeric`
-#' * `$isolated`: `logical`
+#' * `$sampled`: `logical`
 #' @autoglobal
 #' @export
 #' @importFrom data.table data.table
@@ -49,11 +49,11 @@ outbreak_setup <- function(initial_cases, delays, event_probs) {
     asymptomatic = runif(initial_cases) < event_probs$asymptomatic,
     caseid = seq_len(initial_cases), # set case id
     infector = 0,
-    isolated = FALSE,
     missed = TRUE,
     onset = delays$incubation_period(initial_cases),
-    new_cases = NA,
-    isolated_time = Inf
+    new_cases = NA_integer_,
+    isolated_time = Inf,
+    sampled = FALSE
   )
 
   # isolate each symptomatic case after an onset-to-isolation delay after

R/outbreak_step.R

@@ -18,7 +18,7 @@
 #'   the intervention settings for the \pkg{ringbp} model, returned by
 #'   [intervention_opts()]. Contains one element: `quarantine`
 #'
-#' @importFrom data.table data.table rbindlist fcase fifelse
+#' @importFrom data.table data.table rbindlist fcase fifelse copy
 #' @importFrom stats runif
 #' @importFrom stats rnbinom
 #'
@@ -76,44 +76,40 @@ outbreak_step <- function(case_data,
   checkmate::assert_class(event_probs, "ringbp_event_prob_opts")
   checkmate::assert_class(interventions, "ringbp_intervention_opts")
 
-  # For each case in case_data, draw new_cases from a negative binomial
-  # distribution with an R0 and dispersion dependent on if isolated = TRUE
-  case_data[, new_cases := fcase(
-    isolated, offspring$isolated(.N),
-    asymptomatic, offspring$asymptomatic(.N),
-    default = offspring$community(.N)
-  )]
+  # work with copy of case_data in outbreak_step due to
+  # sample_offspring modify by reference
+  case_data <- copy(case_data)
+
+  # case_data is modified by reference and generation times are returned
+  gt <- sample_offspring(
+    case_data = case_data,
+    offspring = offspring,
+    alpha = event_probs$alpha,
+    latent_period = delays$latent_period
+  )
 
   # Select cases that have generated any new cases
-  new_case_data <- case_data[new_cases > 0]
-  # The total new cases generated
-  total_new_cases <- case_data[, sum(new_cases), ]
+  new_case_data <- case_data[new_cases > 0 & !sampled]
 
   # If no new cases drawn, outbreak is over so return case_data
-  if (total_new_cases == 0) {
-    # If everyone is isolated it means that either control has worked or
+  if (nrow(new_case_data) == 0) {
+    # If everyone is sampled it means that either control has worked or
     # everyone has had a chance to infect but didn't
-    case_data$isolated <- TRUE
-
-    effective_r0 <- 0
-    cases_in_gen <- 0
-    out <- list(case_data, effective_r0, cases_in_gen)
-    names(out) <- c("cases", "effective_r0", "cases_in_gen")
-
+    case_data[, sampled := TRUE]
+    out <- list(
+      cases = case_data,
+      effective_r0 = 0,
+      cases_in_gen = 0
+    )
     return(out)
   }
 
   # Compile a data.table for all new cases, new_cases is the amount of people
   # that each infector has infected
   prob_samples <- new_case_data[, list(
     # time when new cases were exposed, a draw from generation time based on
-    # infector's onset
-    exposure = incubation_to_generation_time(
-      symptom_onset_time = rep(onset, new_cases),
-      exposure_time = rep(exposure, new_cases),
-      alpha = event_probs$alpha,
-      latent_period = delays$latent_period
-    ),
+    # infector's onset, ordered to match order of caseid
+    exposure = gt[order(as.numeric(names(gt)))],
     # records the infector of each new person
     infector = rep(caseid, new_cases),
     # records when infector was isolated
@@ -123,14 +119,11 @@ outbreak_step <- function(case_data,
     # cases whose parents are asymptomatic are automatically missed;
     # will draw this for infector_asymptomatic == FALSE
     missed = TRUE,
-    isolated = FALSE,
-    new_cases = NA
+    new_cases = NA_integer_,
+    sampled = FALSE
   )][,
     # draws a sample to see if this person is asymptomatic
     asymptomatic := runif(.N) < event_probs$asymptomatic
-  ][
-    # keep only news cases that are pre-isolation
-    exposure < infector_isolation_time
   ][,
     # onset of new case is exposure + incubation period sample
     onset := exposure + delays$incubation_period(.N)
@@ -167,11 +160,11 @@ outbreak_step <- function(case_data,
   cases_in_gen <- nrow(prob_samples)
 
   ## Estimate the effective r0
-  effective_r0 <- cases_in_gen / case_data[isolated == FALSE, .N]
+  effective_r0 <- cases_in_gen / case_data[sampled == FALSE, .N]
 
   # Everyone in case_data so far has had their chance to infect and are
   # therefore considered isolated
-  case_data$isolated <- TRUE
+  case_data[, sampled := TRUE]
 
   # bind original cases + new secondary cases
   case_data <- data.table::rbindlist(list(case_data, prob_samples),

R/sample_offspring.R

@@ -0,0 +1,101 @@
+#' Sample the offspring distributions for cases that can be in either the
+#' _community_, _isolated_ or _asymptomatic_ states and transition between
+#' states
+#'
+#' Samples from the offspring distributions (see [offspring_opts()]) and adds
+#'   the next generation of transmission events by reference to `case_data`.
+#'   The generation times for each infector-infectee pair is also sampled and
+#'   is returned from the function.
+#'
+#' @details The offspring distribution for a case in the community cannot simply
+#'   be sampled from the `community` offspring distribution as it might become
+#'   isolated before infecting some or all of those cases.
+#'   To account for cases that transition between states (for now only
+#'   _community_ -> _isolated_) we draw from the offspring from both
+#'   distributions, assign all new cases a generation time, and then discard
+#'   the ones that have generation time <= isolation time (for those from the
+#'   isolated offspring distribution) or generation time > isolation time
+#'   (for those from the community offspring distribution), respectively.
+#'
+#' @inheritParams outbreak_step
+#' @inheritParams incubation_to_generation_time
+#' @inheritParams delay_opts
+#'
+#' @autoglobal
+#'
+#' @return A `numeric` vector with the generation times for the new cases
+#'   exposure/infection times.
+#'
+#'   ***Note*** The `case_data` supplied to the function is modified by
+#'   references, see [data.table::set()] for more information.
+#' @keywords internal
+sample_offspring <- function(case_data, offspring, alpha, latent_period) {
+
+  # subset to cases in current generation
+  new_cases <- case_data[sampled == FALSE]
+
+  # logical vectors with the row index of asymptomatic and symptomatic cases
+  asymptomatic_idx <- new_cases$asymptomatic
+  symptomatic_idx <- !asymptomatic_idx
+
+  # cases cannot be isolated at the start of their generation so both community
+  # and isolated offspring distribution are sampled for community cases and
+  # are subset by whether the generation time is before or after the isolation
+  # time
+  community <- offspring$community(sum(symptomatic_idx))
+  isolated <- offspring$isolated(sum(symptomatic_idx))
+
+  # asymptomatic cases are known from the start of their generation so their
+  # offspring can be sampled directly from the asymptomatic offspring
+  # distribution
+  asymptomatic <- offspring$asymptomatic(sum(asymptomatic_idx))
+
+  # get generation times for community and isolated cases
+  community_exposure <- incubation_to_generation_time(
+    symptom_onset_time = rep(new_cases$onset[symptomatic_idx], community),
+    exposure_time = rep(new_cases$exposure[symptomatic_idx], community),
+    alpha = alpha,
+    latent_period = latent_period
+  )
+  names(community_exposure) <- rep(new_cases$caseid[symptomatic_idx], community)
+  isolated_exposure <- incubation_to_generation_time(
+    symptom_onset_time = rep(new_cases$onset[symptomatic_idx], isolated),
+    exposure_time = rep(new_cases$exposure[symptomatic_idx], isolated),
+    alpha = alpha,
+    latent_period = latent_period
+  )
+  names(isolated_exposure) <- rep(new_cases$caseid[symptomatic_idx], isolated)
+
+  # if there is any transmission from asymptomatic cases get generation time
+  if (length(asymptomatic) > 0) {
+    asymptomatic_exposure <- incubation_to_generation_time(
+      symptom_onset_time = rep(new_cases$onset[asymptomatic_idx], asymptomatic),
+      exposure_time = rep(new_cases$exposure[asymptomatic_idx], asymptomatic),
+      alpha = alpha,
+      latent_period = latent_period
+    )
+    names(asymptomatic_exposure) <- rep(new_cases$caseid[asymptomatic_idx], asymptomatic)
+  } else {
+    # create asymptomatic_exposure for all exposure vector below (NULL dropped)
+    asymptomatic_exposure <- NULL
+  }
+
+  # subset transmission events in community and isolation based on infector
+  # isolation time and infectee exposure time
+  infect_before_isolate <- community_exposure < rep(new_cases$isolated_time[symptomatic_idx], community)
+  community_exposure <- community_exposure[infect_before_isolate]
+  infect_after_isolate <- isolated_exposure > rep(new_cases$isolated_time[symptomatic_idx], isolated)
+  isolated_exposure <- isolated_exposure[infect_after_isolate]
+
+  # infectee exposure time for all transmission events
+  exposure <- c(community_exposure, isolated_exposure, asymptomatic_exposure)
+
+  next_gen <- table(names(exposure))
+
+  # assign next generation of cases by reference
+  case_data[sampled == FALSE, new_cases := 0L]
+  case_data[as.numeric(names(next_gen)), new_cases := next_gen]
+
+  # return generation times
+  exposure
+}

man/outbreak_setup.Rd

@@ -8,18 +8,18 @@
         symptomatic_ascertained = 0.8))
     Output
       Index: <asymptomatic>
-         exposure asymptomatic caseid infector isolated missed    onset new_cases
-            <num>       <lgcl>  <int>    <num>   <lgcl> <lgcl>    <num>    <lgcl>
-      1:        0        FALSE      1        0    FALSE   TRUE 3.696253        NA
-      2:        0        FALSE      2        0    FALSE   TRUE 6.837342        NA
-      3:        0        FALSE      3        0    FALSE   TRUE 4.369195        NA
-      4:        0        FALSE      4        0    FALSE   TRUE 6.009202        NA
-      5:        0        FALSE      5        0    FALSE   TRUE 4.157019        NA
-         isolated_time
-                 <num>
-      1:      5.984245
-      2:     12.350205
-      3:      9.761168
-      4:      7.443792
-      5:      7.509917
+         exposure asymptomatic caseid infector missed    onset new_cases
+            <num>       <lgcl>  <int>    <num> <lgcl>    <num>     <int>
+      1:        0        FALSE      1        0   TRUE 3.696253        NA
+      2:        0        FALSE      2        0   TRUE 6.837342        NA
+      3:        0        FALSE      3        0   TRUE 4.369195        NA
+      4:        0        FALSE      4        0   TRUE 6.009202        NA
+      5:        0        FALSE      5        0   TRUE 4.157019        NA
+         isolated_time sampled
+                 <num>  <lgcl>
+      1:      5.984245   FALSE
+      2:     12.350205   FALSE
+      3:      9.761168   FALSE
+      4:      7.443792   FALSE
+      5:      7.509917   FALSE