Convert duplicated code into helper functions

DESCRIPTION

@@ -50,5 +50,5 @@ Config/testthat/edition: 3
 Encoding: UTF-8
 LazyData: true
 Roxygen: list(markdown = TRUE)
-RoxygenNote: 7.3.0
+RoxygenNote: 7.3.1
 Language: en-GB

R/checks.R

@@ -18,3 +18,101 @@
   x <- generation_time(10)
   checkmate::assert_numeric(x, len = 10)
 }
+
+#' Check inputs to `simulate_chains()` and `simulate_summary()`
+#'
+#' @param sim_func <character>; The simulation function to check
+#'
+#' @inheritParams simulate_chains
+#' @return NULL; called for side effects
+#' @keywords internal
+.check_sim_args <- function(
+    index_cases,
+    statistic,
+    offspring_dist,
+    stat_max,
+    pop,
+    percent_immune) {
+  # Input checking
+  checkmate::assert_count(
+    index_cases,
+    positive = TRUE
+  )
+  # check that offspring is a function with argument "n"
+  .check_offspring_func_valid(offspring_dist)
+  # check that arguments related to the statistic are valid
+  .check_statistic_args(
+    statistic,
+    stat_max
+  )
+  checkmate::assert(
+    is.infinite(pop) ||
+      checkmate::assert_integerish(pop, lower = 1)
+  )
+  checkmate::assert_number(
+    percent_immune,
+    lower = 0, upper = 1
+  )
+  invisible(NULL)
+}
+
+#' Check that the `statistic` and `stat_max` arguments are valid
+#'
+#' @inheritParams simulate_chains
+#' @description
+#' The function treats these two arguments as related and checks
+#' them in one place to remove repeated checks in several places in the
+#' package.
+#' @return NULL; called for side effects
+#' @keywords internal
+.check_statistic_args <- function(statistic,
+                                  stat_max) {
+  checkmate::assert_choice(
+    statistic,
+    choices = c("size", "length")
+  )
+  checkmate::assert(
+    is.infinite(stat_max),
+      checkmate::check_integerish(
+        stat_max,
+        lower = 1,
+        null.ok = FALSE
+    ),
+    combine = "or"
+  )
+}
+
+#' Check inputs that control time events
+#'
+#' @description
+#' This function checks the time-related inputs, i.e., start time of each chain,
+#' `t0`, the end time of the simulation, `tf`, and the generation time,
+#' generation_time. It also checks that the generation_time argument is
+#' specified if `tf` is specified as these go hand-in-hand.
+#'
+#' @param tf_specified <logical>; Whether the `tf` argument is specified. Only
+#' makes sense in the context where this function is called, i.e., in
+#' [simulate_chains()]. If `tf` is specified, generation_time must be specified.
+#' @inheritParams simulate_chains
+#' @return NULL; called for side effects
+#' @keywords internal
+.check_time_args <- function(tf_specified,
+                             tf,
+                             generation_time,
+                             t0) {
+  # if tf is specified, generation_time must be specified too
+  if (!is.null(generation_time)) {
+    .check_generation_time_valid(generation_time)
+  } else if (tf_specified) {
+      stop("If `tf` is specified, `generation_time` must be specified too.")
+  }
+  checkmate::assert_number(
+    tf,
+    lower = 0
+  )
+  checkmate::assert_numeric(
+    t0,
+    lower = 0,
+    finite = TRUE
+  )
+}

R/helpers.R

@@ -35,6 +35,91 @@
 
 #' Adjust next generation vector to match susceptible population size
 #'
+#' @description Calculates the initial susceptible population size given
+#' the total population size, the percent immune, and the number of index
+#' cases. This function is used internally, and input checking is not
+#' performed here, only in the context where it is used. Using it directly
+#' is not recommended.
+#'
+#' @inheritParams simulate_chains
+#'
+#' @return numeric; initial susceptible population size
+#' @keywords internal
+.init_susc_pop <- function(pop,
+                           percent_immune,
+                           index_cases) {
+  ss <- max(round(pop * (1 - percent_immune)) - index_cases, 0)
+  return(ss)
+}
+
+#' Sample all possible offspring for the next generation
+#'
+#' @description
+#' Sample next generation of offspring using offspring distribution and
+#' associated parameters. This function is used internally, and input
+#' checking is not performed here, only in the context where it is used.
+#' Using it directly is not recommended.
+#' @param offspring_func A function to sample offspring
+#' @param offspring_func_pars A list of parameters for the offspring function
+#' @param n_offspring A vector of the number of offspring per chain
+#' @param chains Indices of chains/infectors being simulated
+#'
+#' @return A vector of the number of offspring per chain
+#' @keywords internal
+.sample_possible_offspring <- function(offspring_func,
+                                       offspring_func_pars,
+                                       n_offspring,
+                                       chains) {
+
+  possible_new_offspring <- do.call(
+    offspring_func,
+    c(
+      list(n = sum(n_offspring[chains])),
+      offspring_func_pars
+    )
+  )
+  # check that offspring distribution returns integers
+  stopifnot(
+    "Offspring distribution must return integers" =
+      !all(possible_new_offspring %% 1 > 0)
+  )
+
+  return(possible_new_offspring)
+}
+
+#' Sample the number of susceptible offspring from all possible offspring
+#'
+#' @description
+#' Sample susceptible offspring to be infected from all possible offspring.
+#' This function is used internally, and input checking is not
+#' performed here, only in the context where it is used. Using it directly
+#' is not recommended.
+#' @inheritParams simulate_chains
+#' @param new_offspring A vector of the possible new offspring per chain
+#' produced by [.sample_possible_offspring()]
+#' @return A vector of the number of offspring that can be infected given the
+#' current susceptible population size
+#' @keywords internal
+.get_susceptible_offspring <- function(new_offspring,
+                                      susc_pop,
+                                      pop) {
+  # We first adjust for the case where susceptible can be Inf but prob can only
+  # be maximum 1.
+  binom_prob <- min(1, susc_pop / pop, na.rm = TRUE)
+  # Sample the number of infectible offspring from all possible offspring
+  susceptible_offspring <- stats::rbinom(
+    n = length(new_offspring),
+    size = new_offspring,
+    prob = binom_prob
+  )
+  return(susceptible_offspring)
+}
+
+#' Adjust new offspring if it exceeds the susceptible population size
+#' @description
+#' This function is used internally, and input checking is not
+#' performed here, only in the context where it is used. Using it directly
+#' is not recommended.
 #' @param next_gen numeric; vector of next generation offspring
 #' @param susc_pop numeric; susceptible population size
 #'

R/likelihood.R

@@ -53,18 +53,18 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
       nsim_obs, lower = 1, finite = TRUE, na.ok = FALSE
     )
   }
-
   checkmate::assert_numeric(
     chains, lower = 0, upper = Inf, any.missing = FALSE
   )
-  checkmate::assert_character(statistic)
+  # check that arguments related to the statistic are valid
+  .check_statistic_args(
+    statistic,
+    stat_max
+  )
   .check_offspring_func_valid(offspring_dist)
   checkmate::assert_number(
     obs_prob, lower = 0, upper = 1, finite = TRUE, na.ok = FALSE
   )
-  checkmate::assert_number(
-    stat_max, lower = 0, na.ok = FALSE
-  )
   checkmate::assert_logical(
     log, any.missing = FALSE, all.missing = FALSE, len = 1
   )
@@ -74,7 +74,6 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
   checkmate::assert_numeric(
     exclude, null.ok = TRUE
   )
-
   if (obs_prob < 1) {
     if (missing(nsim_obs)) {
       stop("'nsim_obs' must be specified if 'obs_prob' is < 1")