Misc suggestions for simulate_chains()

DESCRIPTION

@@ -5,6 +5,8 @@ Version: 0.0.0.9999
 Authors@R: c(
     person("James M.", "Azam", , "[email protected]", role = c("aut", "cre"),
            comment = c(ORCID = "https://orcid.org/0000-0001-5782-7330")),
+    person("Hugo", "Gruson", , "[email protected]", role = c("ctb"),
+           comment = c(ORCID = "https://orcid.org/0000-0002-4094-1476")),
     person("Zhian N.", "Kamvar", , "[email protected]", role = "ctb",
            comment = c(ORCID = "https://orcid.org/0000-0003-1458-7108")),
     person("Flavio", "Finger", , "[email protected]", role = "aut",

R/simulate.r

@@ -224,11 +224,7 @@ simulate_chains <- function(index_cases,
     # Sample susceptible offspring to be infected from all possible offspring
     # We first adjust for the case where susceptible can be Inf but prob is max
     # 1.
-    binom_prob <- ifelse(
-      is.infinite(susc_pop),
-      1,
-      susc_pop / pop
-    )
+    binom_prob <- min(1, susc_pop / pop, na.rm = TRUE)
     next_gen <- stats::rbinom(
       n = length(next_gen),
       size = next_gen,
@@ -258,19 +254,19 @@ simulate_chains <- function(index_cases,
     # Adorn the new offspring with their information: their ids, their
     # infector's ids, and the generation they were infected in.
     # Also update the susceptible population and generation.
-    if (sum(n_offspring[sim]) > 0) {
+    if (sum(n_offspring) > 0) {
       infectors <- rep(infector_ids, next_gen)
       current_max_id <- unname(tapply(infector_ids, parent_ids, max))
       parent_ids <- rep(sim, n_offspring[sim])
 
       # create new ids
       ids <- rep(current_max_id, n_offspring[sim]) +
-        unlist(lapply(n_offspring[sim], seq_len))
+        sequence(n_offspring[sim])
 
       # increment the generation
       generation <- generation + 1L
       # Update susceptible population
-      susc_pop <- susc_pop - sum(n_offspring[sim])
+      susc_pop <- susc_pop - sum(n_offspring)
 
       # store new simulation results
       tree_df[[generation]] <-
@@ -300,8 +296,6 @@ simulate_chains <- function(index_cases,
       if (!missing(generation_time)) {
         ## only continue to simulate trees that don't go beyond tf
         sim <- intersect(sim, unique(parent_ids)[current_min_time < tf])
-      }
-      if (!missing(generation_time)) {
         times <- times[parent_ids %in% sim]
       }
       infector_ids <- ids[parent_ids %in% sim]
@@ -311,6 +305,7 @@ simulate_chains <- function(index_cases,
   # Combine the results
   tree_df <- do.call(rbind, tree_df)
 
+  # time column only exists if tf was specified
   if (!missing(tf)) {
     tree_df <- tree_df[tree_df$time < tf, ]
   }

README.md

@@ -71,7 +71,7 @@ library("epichains")
   a generation time function is specified, the time of infection.
 
 - `simulate_summary()`: provides a performant version of
-  `simulate_chains()` that only tracks and returns a vector of realized
+  `simulate_chains()` that only tracks and return a vector of realized
   chain sizes or lengths/durations for each index case without details
   of the infection tree.