From 4c495c2320d8471105efd0de8e7853de65c7e6b2 Mon Sep 17 00:00:00 2001
From: jamesaazam <james.azam@lshtm.ac.uk>
Date: Mon, 5 Feb 2024 16:33:21 +0000
Subject: [PATCH 01/38] Add and apply function to initialise susceptible pop

---
 R/helpers.R                   | 18 ++++++++++++++
 R/simulate.r                  |  4 +--
 man/dot-init_susc_pop.Rd      | 30 +++++++++++++++++++++++
 tests/testthat/test-helpers.R | 46 +++++++++++++++++++++++++++++++++++
 4 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 man/dot-init_susc_pop.Rd

diff --git a/R/helpers.R b/R/helpers.R
index 4393cfcb..37abf752 100644
--- a/R/helpers.R
+++ b/R/helpers.R
@@ -35,6 +35,24 @@
 
 #' 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)
+}
+}
+#'
 #' @param next_gen numeric; vector of next generation offspring
 #' @param susc_pop numeric; susceptible population size
 #'
diff --git a/R/simulate.r b/R/simulate.r
index ee4e08ea..a21114f7 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -194,7 +194,7 @@ simulate_chains <- function(index_cases,
     )
   )
   # Initialise susceptible population
-  susc_pop <- max(round(pop * (1 - percent_immune)) - index_cases, 0)
+  susc_pop <- .init_susc_pop(pop, percent_immune, index_cases)
 
   # Add optional columns
   if (!missing(generation_time)) {
@@ -424,7 +424,7 @@ simulate_summary <- function(index_cases,
   sim <- seq_len(index_cases) ## track trees that are still being simulated
 
   # Initialise susceptible population
-  susc_pop <- max(round(pop * (1 - percent_immune)) - index_cases, 0)
+  susc_pop <- .init_susc_pop(pop, percent_immune, index_cases)
 
   ## next, simulate transmission chains from index cases
   while (length(sim) > 0 && susc_pop > 0) {
diff --git a/man/dot-init_susc_pop.Rd b/man/dot-init_susc_pop.Rd
new file mode 100644
index 00000000..3aac3c52
--- /dev/null
+++ b/man/dot-init_susc_pop.Rd
@@ -0,0 +1,30 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/helpers.R
+\name{.init_susc_pop}
+\alias{.init_susc_pop}
+\title{Initialize the susceptible population size}
+\usage{
+.init_susc_pop(pop, percent_immune, index_cases)
+}
+\arguments{
+\item{pop}{\verb{<Integer>}; Population size. Used alongside \code{percent_immune}. to
+define the susceptible population. Defaults to \code{Inf}.}
+
+\item{percent_immune}{\verb{<numeric>}; Percent of the population immune to
+infection at the start of the simulation. Used alongside \code{pop} to initialise
+the susceptible population. Accepted values lie between 0 and 1.
+Defaults to 0.}
+
+\item{index_cases}{Number of index cases to simulate transmission chains for.}
+}
+\value{
+numeric; initial 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.
+}
+\keyword{internal}
diff --git a/tests/testthat/test-helpers.R b/tests/testthat/test-helpers.R
index c7c05e86..61640b34 100644
--- a/tests/testthat/test-helpers.R
+++ b/tests/testthat/test-helpers.R
@@ -41,3 +41,49 @@ test_that("get_statistic_func works correctly", {
     "chain_statistic must be 'size' or 'length'"
   )
 })
+
+test_that(".init_susc_pop works correctly", {
+  expect_identical(
+    .init_susc_pop(
+      pop = 10,
+      percent_immune = 0.5,
+      index_cases = 1
+    ),
+    4
+  )
+  expect_identical(
+    .init_susc_pop(
+      pop = 0,
+      percent_immune = 0.5,
+      index_cases = 1
+    ),
+    0
+  )
+  expect_identical(
+    .init_susc_pop(
+      pop = 10,
+      percent_immune = 0,
+      index_cases = 0
+    ),
+    10
+  )
+  expect_length(
+    .init_susc_pop(
+      pop = 10,
+      percent_immune = 0,
+      index_cases = 1
+    ),
+    1
+  )
+  expect_type(
+    .init_susc_pop(
+      pop = 10,
+      percent_immune = 0,
+      index_cases = 1
+    ),
+    "double"
+  )
+})
+})
+  )
+})
\ No newline at end of file

From 0027e53f15d72252b18993b77183fbdc84f7fb69 Mon Sep 17 00:00:00 2001
From: jamesaazam <james.azam@lshtm.ac.uk>
Date: Mon, 5 Feb 2024 16:33:35 +0000
Subject: [PATCH 02/38] Update wordlist

---
 inst/WORDLIST | 1 +
 1 file changed, 1 insertion(+)

diff --git a/inst/WORDLIST b/inst/WORDLIST
index 7ebaebb1..12392f36 100644
--- a/inst/WORDLIST
+++ b/inst/WORDLIST
@@ -30,6 +30,7 @@ gborel
 geq
 infectee
 infectees
+infectible
 infector
 infectors
 json

From 293a08037884e6ae8e133b88be77f3428e7a921b Mon Sep 17 00:00:00 2001
From: jamesaazam <james.azam@lshtm.ac.uk>
Date: Mon, 5 Feb 2024 16:35:49 +0000
Subject: [PATCH 03/38] Add function to sample all possible next offspring

---
 R/helpers.R                          | 35 ++++++++++++++++++++++++++++
 R/simulate.r                         | 26 ++++++++++-----------
 man/dot-sample_possible_offspring.Rd | 32 +++++++++++++++++++++++++
 tests/testthat/test-helpers.R        | 19 +++++++++++++++
 4 files changed, 99 insertions(+), 13 deletions(-)
 create mode 100644 man/dot-sample_possible_offspring.Rd

diff --git a/R/helpers.R b/R/helpers.R
index 37abf752..35e8eab1 100644
--- a/R/helpers.R
+++ b/R/helpers.R
@@ -51,7 +51,42 @@
   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)
 }
+
 #'
 #' @param next_gen numeric; vector of next generation offspring
 #' @param susc_pop numeric; susceptible population size
diff --git a/R/simulate.r b/R/simulate.r
index a21114f7..0fd397a3 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -206,12 +206,12 @@ simulate_chains <- function(index_cases,
   }
   # next, simulate n trees
   while (length(sim) > 0 && susc_pop > 0) {
-    # sample next generation of offspring
-    next_gen <- do.call(
-      offspring_dist,
-      c(
-        list(n = sum(n_offspring[sim])),
-        pars
+    # simulate the next possible offspring
+    next_gen <- .sample_possible_offspring(
+      offspring_func = roffspring_name,
+      offspring_func_pars = pars,
+      n_offspring = n_offspring,
+      chains = sim
       )
     )
     # check that offspring distribution returns integers
@@ -428,14 +428,14 @@ simulate_summary <- function(index_cases,
 
   ## next, simulate transmission chains from index cases
   while (length(sim) > 0 && susc_pop > 0) {
-    ## simulate next generation
-    next_gen <- do.call(
-      offspring_dist,
-      c(
-        list(n = sum(n_offspring[sim])),
-        pars
-      )
+    # simulate the possible next generation of offspring
+    next_gen <- .sample_possible_offspring(
+      offspring_func = roffspring_name,
+      offspring_func_pars = pars,
+      n_offspring = n_offspring,
+      chains = sim
     )
+    # from all possible offspring, get those that are infectible
     # check that offspring distribution returns integers
     stopifnot(
       "Offspring distribution must return integers" =
diff --git a/man/dot-sample_possible_offspring.Rd b/man/dot-sample_possible_offspring.Rd
new file mode 100644
index 00000000..3aa60b67
--- /dev/null
+++ b/man/dot-sample_possible_offspring.Rd
@@ -0,0 +1,32 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/helpers.R
+\name{.sample_possible_offspring}
+\alias{.sample_possible_offspring}
+\title{Sample all possible offspring for the next generation}
+\usage{
+.sample_possible_offspring(
+  offspring_func,
+  offspring_func_pars,
+  n_offspring,
+  chains
+)
+}
+\arguments{
+\item{offspring_func}{A function to sample offspring}
+
+\item{offspring_func_pars}{A list of parameters for the offspring function}
+
+\item{n_offspring}{A vector of the number of offspring per chain}
+
+\item{chains}{Indices of chains/infectors being simulated}
+}
+\value{
+A vector of the number of offspring per chain
+}
+\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.
+}
+\keyword{internal}
diff --git a/tests/testthat/test-helpers.R b/tests/testthat/test-helpers.R
index 61640b34..16302d3c 100644
--- a/tests/testthat/test-helpers.R
+++ b/tests/testthat/test-helpers.R
@@ -84,6 +84,25 @@ test_that(".init_susc_pop works correctly", {
     "double"
   )
 })
+
+test_that(".init_susc_pop works correctly", {
+  expect_length(
+    .sample_possible_offspring(
+      offspring_func = "rpois",
+      offspring_func_pars = list(lambda = 1),
+      n_offspring = 10,
+      chains = 1
+    ), 10
+  )
+  expect_error(
+    .sample_possible_offspring(
+      offspring_func = "rnorm",
+      offspring_func_pars = list(mean = 0, sd = 1),
+      n_offspring = 10,
+      chains = 1
+    ),
+    "Offspring distribution must return integers"
+  )
 })
   )
 })
\ No newline at end of file

From 8009648b2dcd360f34f86dfa76382cf27542303e Mon Sep 17 00:00:00 2001
From: jamesaazam <james.azam@lshtm.ac.uk>
Date: Mon, 5 Feb 2024 16:36:39 +0000
Subject: [PATCH 04/38] Add function to sample potential infections that will
 be infected

---
 R/helpers.R                         | 32 ++++++++++++++++++++
 R/simulate.r                        | 46 +++++++----------------------
 man/dot-get_infectible_offspring.Rd | 26 ++++++++++++++++
 tests/testthat/test-helpers.R       | 21 ++++++++++++-
 4 files changed, 89 insertions(+), 36 deletions(-)
 create mode 100644 man/dot-get_infectible_offspring.Rd

diff --git a/R/helpers.R b/R/helpers.R
index 35e8eab1..4b843045 100644
--- a/R/helpers.R
+++ b/R/helpers.R
@@ -87,7 +87,39 @@
   return(possible_new_offspring)
 }
 
+#' Sample the number of infectible 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_infectible_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
+  infectible_offspring <- stats::rbinom(
+    n = length(new_offspring),
+    size = new_offspring,
+    prob = binom_prob
+  )
+  return(infectible_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
 #'
diff --git a/R/simulate.r b/R/simulate.r
index 0fd397a3..c060adf0 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -213,23 +213,13 @@ simulate_chains <- function(index_cases,
       n_offspring = n_offspring,
       chains = sim
       )
+    # from all possible offspring, get those that could be infected
+    next_gen <- .get_infectible_offspring(
+      new_offspring = next_gen,
+      susc_pop = susc_pop,
+      pop = pop
     )
-    # check that offspring distribution returns integers
-    stopifnot(
-      "Offspring distribution must return integers" =
-        !all(next_gen %% 1 > 0)
-    )
-    # 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 <- min(1, susc_pop / pop, na.rm = TRUE)
-    next_gen <- stats::rbinom(
-      n = length(next_gen),
-      size = next_gen,
-      prob = binom_prob
-    )
-    # Adjust next_gen if the number of offspring is greater than the
-    # susceptible population.
+    # Adjust the infectibles if they exceed the susceptible population
     if (sum(next_gen) > susc_pop) {
       next_gen <- .adjust_next_gen(
         next_gen = next_gen,
@@ -436,27 +426,13 @@ simulate_summary <- function(index_cases,
       chains = sim
     )
     # from all possible offspring, get those that are infectible
-    # check that offspring distribution returns integers
-    stopifnot(
-      "Offspring distribution must return integers" =
-        !all(next_gen %% 1 > 0)
-    )
-    # 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 <- min(
-      1,
-      susc_pop / pop,
-      na.rm = TRUE
-    )
-
-    next_gen <- stats::rbinom(
-      n = length(next_gen),
-      size = next_gen,
-      prob = binom_prob
+    next_gen <- .get_infectible_offspring(
+      new_offspring = next_gen,
+      susc_pop = susc_pop,
+      pop = pop
     )
     # Adjust next_gen if the number of offspring is greater than the
-    # susceptible population.
+    # susceptible population
     if (sum(next_gen) > susc_pop) {
       next_gen <- .adjust_next_gen(
         next_gen = next_gen,
diff --git a/man/dot-get_infectible_offspring.Rd b/man/dot-get_infectible_offspring.Rd
new file mode 100644
index 00000000..770d13bf
--- /dev/null
+++ b/man/dot-get_infectible_offspring.Rd
@@ -0,0 +1,26 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/helpers.R
+\name{.get_infectible_offspring}
+\alias{.get_infectible_offspring}
+\title{Sample the number of infectible offspring from all possible offspring}
+\usage{
+.get_infectible_offspring(new_offspring, susc_pop, pop)
+}
+\arguments{
+\item{new_offspring}{A vector of the possible new offspring per chain
+produced by \code{\link[=.sample_possible_offspring]{.sample_possible_offspring()}}}
+
+\item{pop}{\verb{<Integer>}; Population size. Used alongside \code{percent_immune}. to
+define the susceptible population. Defaults to \code{Inf}.}
+}
+\value{
+A vector of the number of offspring that can be infected given the
+current susceptible population size
+}
+\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.
+}
+\keyword{internal}
diff --git a/tests/testthat/test-helpers.R b/tests/testthat/test-helpers.R
index 16302d3c..07d01170 100644
--- a/tests/testthat/test-helpers.R
+++ b/tests/testthat/test-helpers.R
@@ -104,5 +104,24 @@ test_that(".init_susc_pop works correctly", {
     "Offspring distribution must return integers"
   )
 })
+
+test_that(".init_susc_pop works correctly", {
+  next_gen <- c(1, 2, 5)
+  expect_length(
+    .get_infectible_offspring(
+      new_offspring = next_gen,
+      susc_pop = 1,
+      pop = 20
+    ),
+    length(next_gen)
   )
-})
\ No newline at end of file
+  # If the susceptible population in infinite, next_gen should be returned
+  expect_identical(
+    .get_infectible_offspring(
+      new_offspring = next_gen,
+      susc_pop = Inf,
+      pop = Inf
+    ),
+    as.integer(next_gen)
+  )
+})

From e2fb3d7e46630b1006a218c219854fd671fc7b43 Mon Sep 17 00:00:00 2001
From: jamesaazam <james.azam@lshtm.ac.uk>
Date: Mon, 5 Feb 2024 16:37:40 +0000
Subject: [PATCH 05/38] Improve documentation of function that adjusts
 offspring to susc pop size

---
 man/dot-adjust_next_gen.Rd | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/man/dot-adjust_next_gen.Rd b/man/dot-adjust_next_gen.Rd
index f213cfc5..ae9af4fe 100644
--- a/man/dot-adjust_next_gen.Rd
+++ b/man/dot-adjust_next_gen.Rd
@@ -2,7 +2,7 @@
 % Please edit documentation in R/helpers.R
 \name{.adjust_next_gen}
 \alias{.adjust_next_gen}
-\title{Adjust next generation vector to match susceptible population size}
+\title{Adjust new offspring if it exceeds the susceptible population size}
 \usage{
 .adjust_next_gen(next_gen, susc_pop)
 }
@@ -15,6 +15,8 @@
 numeric; adjusted next generation offspring vector
 }
 \description{
-Adjust next generation vector to match susceptible population size
+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.
 }
 \keyword{internal}

From 89ec8790b450bdb8f2a9142a57cef0f06649a303 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 14:36:55 +0000
Subject: [PATCH 06/38] Add input checking function for simulate_* functions

---
 R/checks.R | 66 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/R/checks.R b/R/checks.R
index c4bc7060..8546b1eb 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -18,3 +18,69 @@
   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(
+    func_name = c("simulate_chains", "simulate_summary"),
+    index_cases,
+    statistic,
+    offspring_dist,
+    stat_max,
+    pop,
+    percent_immune,
+    generation_time = NULL,
+    t0 = NULL,
+    tf = NULL) {
+  func_name <- match.arg(func_name)
+  # Input checking
+  checkmate::assert_count(
+    index_cases,
+    positive = TRUE
+  )
+  checkmate::assert_choice(
+    statistic,
+    choices = c("size", "length")
+  )
+  checkmate::assert_string(offspring_dist)
+  # check that offspring function exists in the environment
+  roffspring_name <- paste0(
+    "r",
+    offspring_dist
+  )
+  .check_offspring_func_valid(roffspring_name)
+  checkmate::assert(
+    is.infinite(stat_max) ||
+      checkmate::assert_integerish(stat_max, lower = 0)
+  )
+  checkmate::assert(
+    is.infinite(pop) ||
+      checkmate::assert_integerish(pop, lower = 1)
+  )
+  checkmate::assert_number(
+    percent_immune,
+    lower = 0, upper = 1
+  )
+
+  if (func_name == "simulate_chains") {
+    if (!missing(generation_time)) {
+      .check_generation_time_valid(generation_time)
+    } else if (!missing(tf)) {
+      stop("If `tf` is specified, `generation_time` must be specified too.")
+    }
+    checkmate::assert_numeric(
+      t0,
+      lower = 0, finite = TRUE
+    )
+    checkmate::assert_number(
+      tf,
+      lower = 0
+    )
+  }
+  invisible(NULL)
+}

From 387e1091541e866b1167b753d7f705894cbf8ce2 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 14:37:11 +0000
Subject: [PATCH 07/38] Apply input checking function

---
 R/simulate.r | 78 ++++++++++++++++------------------------------------
 1 file changed, 23 insertions(+), 55 deletions(-)

diff --git a/R/simulate.r b/R/simulate.r
index c060adf0..e54561b6 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -139,40 +139,19 @@ simulate_chains <- function(index_cases,
                             generation_time = NULL,
                             t0 = 0,
                             tf = Inf) {
-  # Input checking
-  checkmate::assert_count(
-    index_cases,
-    positive = TRUE
-  )
-  checkmate::assert_choice(
-    statistic,
-    choices = c("size", "length")
-  )
-  .check_offspring_func_valid(offspring_dist)
-  checkmate::assert(
-    is.infinite(stat_max) ||
-      checkmate::assert_integerish(stat_max, lower = 0)
-  )
-  checkmate::assert(
-    is.infinite(pop) ||
-      checkmate::assert_integerish(pop, lower = 1)
-  )
-  checkmate::assert_number(
-    percent_immune,
-    lower = 0, upper = 1
-  )
-  if (!missing(generation_time)) {
-    .check_generation_time_valid(generation_time)
-  } else if (!missing(tf)) {
-    stop("If `tf` is specified, `generation_time` must be specified too.")
-  }
-  checkmate::assert_numeric(
-    t0,
-    lower = 0, finite = TRUE
-  )
-  checkmate::assert_number(
-    tf,
-    lower = 0
+  # Check inputs
+  func_name <- as.character(match.call()[[1]])
+  .check_sim_args(
+    func_name = func_name,
+    index_cases = index_cases,
+    statistic = statistic,
+    offspring_dist = offspring_dist,
+    stat_max = stat_max,
+    pop = pop,
+    percent_immune = percent_immune,
+    generation_time = generation_time,
+    t0 = t0,
+    tf = tf
   )
   # Gather offspring distribution parameters
   pars <- list(...)
@@ -383,27 +362,16 @@ simulate_summary <- function(index_cases,
                              stat_max = Inf,
                              pop = Inf,
                              percent_immune = 0) {
-  # Input checking
-  checkmate::assert_count(index_cases, positive = TRUE)
-  statistic <- match.arg(statistic)
-  checkmate::assert_choice(
-    statistic,
-    choices = c("size", "length")
-  )
-
-  # check that offspring is properly specified
-  .check_offspring_func_valid(offspring_dist)
-
-  checkmate::assert_number(
-    stat_max, lower = 0
-  )
-  checkmate::assert(
-    is.infinite(pop) ||
-      checkmate::assert_integerish(pop, lower = 1)
-  )
-  checkmate::assert_number(
-    percent_immune,
-    lower = 0, upper = 1
+  # Check inputs
+  func_name <- as.character(match.call()[[1]])
+  .check_sim_args(
+    func_name = func_name,
+    index_cases = index_cases,
+    statistic = statistic,
+    offspring_dist = offspring_dist,
+    stat_max = stat_max,
+    pop = pop,
+    percent_immune = percent_immune
   )
   # Gather offspring distribution parameters
   pars <- list(...)

From 2414ec2f0b2eee0591a6115737b2e6d09efcfdb9 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 14:57:39 +0000
Subject: [PATCH 08/38] Remove redundant input checking

---
 R/stat_likelihoods.R | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/R/stat_likelihoods.R b/R/stat_likelihoods.R
index 6581d79b..15368b34 100644
--- a/R/stat_likelihoods.R
+++ b/R/stat_likelihoods.R
@@ -169,13 +169,7 @@ offspring_ll <- function(x, offspring_dist, statistic,
   checkmate::assert_numeric(
     x, lower = 0, any.missing = FALSE
   )
-  # check that offspring is properly specified
-  .check_offspring_func_valid(offspring_dist)
-  checkmate::assert_character(statistic)
-  checkmate::assert_numeric(
-    nsim_offspring, lower = 1
-  )
-
+  # Remaining checks are done in simulate_summary()
   # Simulate the chains
   dist <- simulate_summary(
     index_cases = nsim_offspring,

From a6f2ca4fe6f026088f2c3ea722300f931f91b443 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 17:18:42 +0000
Subject: [PATCH 09/38] Fix call to offspring_dist

---
 R/simulate.r | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/R/simulate.r b/R/simulate.r
index e54561b6..31f1b4ff 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -187,7 +187,7 @@ simulate_chains <- function(index_cases,
   while (length(sim) > 0 && susc_pop > 0) {
     # simulate the next possible offspring
     next_gen <- .sample_possible_offspring(
-      offspring_func = roffspring_name,
+      offspring_func = offspring_dist,
       offspring_func_pars = pars,
       n_offspring = n_offspring,
       chains = sim
@@ -388,7 +388,7 @@ simulate_summary <- function(index_cases,
   while (length(sim) > 0 && susc_pop > 0) {
     # simulate the possible next generation of offspring
     next_gen <- .sample_possible_offspring(
-      offspring_func = roffspring_name,
+      offspring_func = offspring_dist,
       offspring_func_pars = pars,
       n_offspring = n_offspring,
       chains = sim

From 8e095efa05321d080df7196d92e194954d6aa104 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 17:19:20 +0000
Subject: [PATCH 10/38] Combine checks on statistic and stat_max

---
 R/checks.R     | 43 ++++++++++++++++++++++++++++++-------------
 R/likelihood.R | 11 +++++------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index 8546b1eb..51ae75cd 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -37,26 +37,19 @@
     generation_time = NULL,
     t0 = NULL,
     tf = NULL) {
+  # Get the function name
   func_name <- match.arg(func_name)
   # Input checking
   checkmate::assert_count(
     index_cases,
     positive = TRUE
   )
-  checkmate::assert_choice(
+  # 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,
-    choices = c("size", "length")
-  )
-  checkmate::assert_string(offspring_dist)
-  # check that offspring function exists in the environment
-  roffspring_name <- paste0(
-    "r",
-    offspring_dist
-  )
-  .check_offspring_func_valid(roffspring_name)
-  checkmate::assert(
-    is.infinite(stat_max) ||
-      checkmate::assert_integerish(stat_max, lower = 0)
+    stat_max
   )
   checkmate::assert(
     is.infinite(pop) ||
@@ -68,6 +61,8 @@
   )
 
   if (func_name == "simulate_chains") {
+    # Check generation time is properly specified and if tf
+    # is specified, generation_time is also specified
     if (!missing(generation_time)) {
       .check_generation_time_valid(generation_time)
     } else if (!missing(tf)) {
@@ -84,3 +79,25 @@
   }
   invisible(NULL)
 }
+
+#' Check that the statistic and stat_max arguments are valid
+#'
+#' @inheritParams simulate_chains
+#'
+#' @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::assert_integerish(
+        stat_max,
+        lower = 0,
+        null.ok = FALSE
+    )
+  )
+}
diff --git a/R/likelihood.R b/R/likelihood.R
index 879d11d6..2d1668a4 100644
--- a/R/likelihood.R
+++ b/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")

From 9736d4b027f12525d9e8bfb04b66fcd4f6de77a9 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 17:38:43 +0000
Subject: [PATCH 11/38] Add description to checker function

---
 R/checks.R | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/R/checks.R b/R/checks.R
index 51ae75cd..3d51407b 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -83,6 +83,11 @@
 #' 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

From 8357cec057eb2603b997afbced99934acc267f29 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 17:39:27 +0000
Subject: [PATCH 12/38] Add checks to see if generation_time and tf are
 specified in the main function

---
 R/checks.R   | 8 ++++----
 R/simulate.r | 8 ++++++++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index 3d51407b..bc217456 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -34,8 +34,10 @@
     stat_max,
     pop,
     percent_immune,
+    generation_time_specified,
     generation_time = NULL,
     t0 = NULL,
+    tf_specified,
     tf = NULL) {
   # Get the function name
   func_name <- match.arg(func_name)
@@ -61,11 +63,9 @@
   )
 
   if (func_name == "simulate_chains") {
-    # Check generation time is properly specified and if tf
-    # is specified, generation_time is also specified
-    if (!missing(generation_time)) {
+    if (!generation_time_specified) {
       .check_generation_time_valid(generation_time)
-    } else if (!missing(tf)) {
+    } else if (!tf_specified) {
       stop("If `tf` is specified, `generation_time` must be specified too.")
     }
     checkmate::assert_numeric(
diff --git a/R/simulate.r b/R/simulate.r
index 31f1b4ff..a57fb26d 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -141,6 +141,12 @@ simulate_chains <- function(index_cases,
                             tf = Inf) {
   # Check inputs
   func_name <- as.character(match.call()[[1]])
+  # We need to check if the generation time and tf are specified here before
+  # calling .check_sim_args() to not obfuscate the meaning of missing()
+  # in the checker function.
+  generation_time_specified <- missing(generation_time)
+  tf_specified <- missing(tf)
+  # Run checks
   .check_sim_args(
     func_name = func_name,
     index_cases = index_cases,
@@ -149,8 +155,10 @@ simulate_chains <- function(index_cases,
     stat_max = stat_max,
     pop = pop,
     percent_immune = percent_immune,
+    generation_time_specified = generation_time_specified,
     generation_time = generation_time,
     t0 = t0,
+    tf_specified = tf_specified,
     tf = tf
   )
   # Gather offspring distribution parameters

From 2adb92c78644b8530fbd14e1d043b9a7c26669d3 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 17:39:50 +0000
Subject: [PATCH 13/38] Generate docs

---
 man/dot-check_sim_args.Rd       | 75 +++++++++++++++++++++++++++++++++
 man/dot-check_statistic_args.Rd | 31 ++++++++++++++
 man/dot-init_susc_pop.Rd        |  2 +-
 3 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 man/dot-check_sim_args.Rd
 create mode 100644 man/dot-check_statistic_args.Rd

diff --git a/man/dot-check_sim_args.Rd b/man/dot-check_sim_args.Rd
new file mode 100644
index 00000000..84463ce3
--- /dev/null
+++ b/man/dot-check_sim_args.Rd
@@ -0,0 +1,75 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/checks.R
+\name{.check_sim_args}
+\alias{.check_sim_args}
+\title{Check inputs to \code{simulate_chains()} and \code{simulate_summary()}}
+\usage{
+.check_sim_args(
+  func_name = c("simulate_chains", "simulate_summary"),
+  index_cases,
+  statistic,
+  offspring_dist,
+  stat_max,
+  pop,
+  percent_immune,
+  generation_time_specified,
+  generation_time = NULL,
+  t0 = NULL,
+  tf_specified,
+  tf = NULL
+)
+}
+\arguments{
+\item{index_cases}{Number of index cases to simulate transmission chains for.}
+
+\item{statistic}{\verb{<String>}; Chain statistic to track as the stopping
+criteria for each chain being simulated when \code{stat_max} is not \code{Inf}.
+Can be one of:
+\itemize{
+\item "size": the total number of cases produced by a chain before it goes
+extinct.
+\item "length": the total number of ancestors produced by a chain before
+it goes extinct.
+}}
+
+\item{offspring_dist}{Offspring distribution: a function like the ones
+provided by R to generate random numbers from given distributions (e.g.,
+\code{\link{rpois}} for Poisson). More specifically, the function needs to
+accept at least one argument, \code{n}, which is the number of random
+numbers to generate. It can accept further arguments, which will be passed
+on to the random number generating functions. Examples that can be provided
+here are \code{rpois} for Poisson distributed offspring, \code{rnbinom} for negative
+binomial offspring, or custom functions.}
+
+\item{stat_max}{A cut off for the chain statistic (size/length) being
+computed. Results above \code{stat_max} are set to \code{stat_max}. Defaults to \code{Inf}.}
+
+\item{pop}{\verb{<Integer>}; Population size. Used alongside \code{percent_immune}. to
+define the susceptible population. Defaults to \code{Inf}.}
+
+\item{percent_immune}{\verb{<numeric>}; Percent of the population immune to
+infection at the start of the simulation. Used alongside \code{pop} to initialise
+the susceptible population. Accepted values lie between 0 and 1.
+Defaults to 0.}
+
+\item{generation_time}{The generation time function; the name
+of a user-defined named or anonymous function with only one argument \code{n},
+representing the number of generation times to sample.}
+
+\item{t0}{Start time (if generation time is given); either a single value
+or a vector of same length as \code{index_cases} (number of initial cases) with
+corresponding initial times. Defaults to 0, meaning all cases started at
+time 0.}
+
+\item{tf}{Cut-off for the infection times (if generation time is given).
+Defaults to \code{Inf}.}
+
+\item{sim_func}{\if{html}{\out{<character>}}; The simulation function to check}
+}
+\value{
+NULL; called for side effects
+}
+\description{
+Check inputs to \code{simulate_chains()} and \code{simulate_summary()}
+}
+\keyword{internal}
diff --git a/man/dot-check_statistic_args.Rd b/man/dot-check_statistic_args.Rd
new file mode 100644
index 00000000..0afd00a2
--- /dev/null
+++ b/man/dot-check_statistic_args.Rd
@@ -0,0 +1,31 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/checks.R
+\name{.check_statistic_args}
+\alias{.check_statistic_args}
+\title{Check that the statistic and stat_max arguments are valid}
+\usage{
+.check_statistic_args(statistic, stat_max)
+}
+\arguments{
+\item{statistic}{\verb{<String>}; Chain statistic to track as the stopping
+criteria for each chain being simulated when \code{stat_max} is not \code{Inf}.
+Can be one of:
+\itemize{
+\item "size": the total number of cases produced by a chain before it goes
+extinct.
+\item "length": the total number of ancestors produced by a chain before
+it goes extinct.
+}}
+
+\item{stat_max}{A cut off for the chain statistic (size/length) being
+computed. Results above \code{stat_max} are set to \code{stat_max}. Defaults to \code{Inf}.}
+}
+\value{
+NULL; called for side effects
+}
+\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.
+}
+\keyword{internal}
diff --git a/man/dot-init_susc_pop.Rd b/man/dot-init_susc_pop.Rd
index 3aac3c52..502b4aa9 100644
--- a/man/dot-init_susc_pop.Rd
+++ b/man/dot-init_susc_pop.Rd
@@ -2,7 +2,7 @@
 % Please edit documentation in R/helpers.R
 \name{.init_susc_pop}
 \alias{.init_susc_pop}
-\title{Initialize the susceptible population size}
+\title{Adjust next generation vector to match susceptible population size}
 \usage{
 .init_susc_pop(pop, percent_immune, index_cases)
 }

From 2c3ec6db99aff20f0566434dd4ed18c70fa12933 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:28:05 +0000
Subject: [PATCH 14/38] Use statistic explicitly

---
 vignettes/interventions.Rmd | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/vignettes/interventions.Rmd b/vignettes/interventions.Rmd
index a667913f..4bdc701d 100644
--- a/vignettes/interventions.Rmd
+++ b/vignettes/interventions.Rmd
@@ -46,7 +46,7 @@ We simulate 200 chains tracking up to 99 infections:
 ```{r simulate_chains}
 sims <- simulate_summary(
   index_cases = 200, offspring_dist = rnbinom, stat_max = 99, mu = 1.2,
-  size = 0.5
+  size = 0.5, statistic = "size"
 )
 ```
 
@@ -72,7 +72,7 @@ For example, to reduce R by 25% at the population level we scale the `mu` parame
 ```{r simulate_chains_pop_control}
 sims <- simulate_summary(
   index_cases = 200, offspring_dist = rnbinom, stat_max = 99, mu = 0.9,
-  size = 0.5
+  size = 0.5, statistic = "size"
 )
 sims[is.infinite(sims)] <- 100 # Replace infections > 99 with 100 for plotting.
 ggplot(data.frame(x = sims), aes(x = x)) +
@@ -110,7 +110,7 @@ Having defined this, we can generate simulations as before:
 ```{r simulate_chains_ind_control}
 sims <- simulate_summary(
   index_cases = 200, offspring_dist = rnbinom_ind, stat_max = 99, mu = 1.2,
-  size = 0.5, control = 0.25
+  size = 0.5, control = 0.25, statistic = "size"
 )
 sims[is.infinite(sims)] <- 100 # Replace infections > 99 with 100 for plotting.
 ggplot(data.frame(x = sims), aes(x = x)) +
@@ -140,7 +140,7 @@ This can be likened to a disease control strategy where gatherings are limited t
 ```{r simulate_chains_truncated}
 sims <- simulate_summary(
   index_cases = 200, offspring_dist = rnbinom_truncated, stat_max = 99,
-  mu = 1.2, size = 0.5, max =  10
+  mu = 1.2, size = 0.5, max =  10, statistic = "size"
 )
 sims[is.infinite(sims)] <- 100 # Replace infections > 99 with 100 for plotting.
 ggplot(data.frame(x = sims), aes(x = x)) +

From 78aad53ad80136d22b83af47be7d2c2d912d99e8 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:29:01 +0000
Subject: [PATCH 15/38] Make boolean result align with object name

---
 R/simulate.r | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/R/simulate.r b/R/simulate.r
index a57fb26d..232b38ca 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -144,8 +144,8 @@ simulate_chains <- function(index_cases,
   # We need to check if the generation time and tf are specified here before
   # calling .check_sim_args() to not obfuscate the meaning of missing()
   # in the checker function.
-  generation_time_specified <- missing(generation_time)
-  tf_specified <- missing(tf)
+  generation_time_specified <- !missing(generation_time)
+  tf_specified <- !missing(tf)
   # Run checks
   .check_sim_args(
     func_name = func_name,

From 5774f767c3f8a9bde82e70db750a78b9a196b9ec Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:29:21 +0000
Subject: [PATCH 16/38] Styling

---
 R/checks.R | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/checks.R b/R/checks.R
index bc217456..f598daca 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -92,7 +92,7 @@
 #' @return NULL; called for side effects
 #' @keywords internal
 .check_statistic_args <- function(statistic,
-                                  stat_max){
+                                  stat_max) {
   checkmate::assert_choice(
     statistic,
     choices = c("size", "length")

From 5fab467b90320e41bd692a99bdf90bcc108f96d1 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:29:30 +0000
Subject: [PATCH 17/38] Add tests

---
 tests/testthat/test-checks.R | 161 +++++++++++++++++++++++++++++++++++
 1 file changed, 161 insertions(+)

diff --git a/tests/testthat/test-checks.R b/tests/testthat/test-checks.R
index 9871a94d..f7f25365 100644
--- a/tests/testthat/test-checks.R
+++ b/tests/testthat/test-checks.R
@@ -15,4 +15,165 @@ test_that("Checks work", {
     .check_generation_time_valid(function(x) 3),
     "Must have length"
   )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_chains",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1,
+      generation_time_specified = FALSE,
+      t0 = 0,
+      tf_specified = TRUE,
+      tf = 10
+    ),
+    "If `tf` is specified, `generation_time` must be specified too."
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 0,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1
+    ),
+    "Must be >= 1."
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 10,
+      statistic = "a",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1
+    ),
+    "Must be element of set"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = r,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1
+    ),
+    "object 'r' not found"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = -1,
+      pop = 10,
+      percent_immune = 0.1
+    ),
+    "Element 1 is not >= 0"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 1,
+      pop = -1,
+      percent_immune = 0.1
+    ),
+    "Element 1 is not >= 1"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_summary",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 1,
+      pop = 10,
+      percent_immune = -0.1
+    ),
+    "Element 1 is not >= 0"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_chains",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1,
+      generation_time_specified = TRUE,
+      generation_time = NULL,
+      t0 = 0,
+      tf_specified = TRUE,
+      tf = 10
+    ),
+    "Must be a function, not 'NULL'"
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_chains",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1,
+      generation_time_specified = TRUE,
+      generation_time = function(x) rep(3, 10),
+      t0 = -1,
+      tf_specified = TRUE,
+      tf = Inf
+    ),
+    "Element 1 is not >= 0."
+  )
+  expect_error(
+    .check_sim_args(
+      func_name = "simulate_chains",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1,
+      generation_time_specified = TRUE,
+      generation_time = function(x) rep(3, 10),
+      t0 = 0,
+      tf_specified = TRUE,
+      tf = -1
+    ),
+    "Element 1 is not >= 0."
+  )
+  expect_no_error(
+    .check_sim_args(
+      func_name = "simulate_chains",
+      index_cases = 10,
+      statistic = "size",
+      offspring_dist = rpois,
+      stat_max = 10,
+      pop = 10,
+      percent_immune = 0.1,
+      generation_time_specified = TRUE,
+      generation_time = function(x) rep(3, 10),
+      t0 = 0,
+      tf_specified = TRUE,
+      tf = 10
+    )
+  )
+  expect_no_error(
+    .check_statistic_args(
+      statistic = "size",
+      stat_max = 10
+    )
+  )
 })

From 5018a5af872b76f5509a7f911d736866e9a6730e Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:30:00 +0000
Subject: [PATCH 18/38] Use monospace to avoid spell check failures

---
 R/checks.R                      | 2 +-
 man/dot-check_statistic_args.Rd | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index f598daca..e2e850cd 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -80,7 +80,7 @@
   invisible(NULL)
 }
 
-#' Check that the statistic and stat_max arguments are valid
+#' Check that the `statistic` and `stat_max` arguments are valid
 #'
 #' @inheritParams simulate_chains
 #' @description
diff --git a/man/dot-check_statistic_args.Rd b/man/dot-check_statistic_args.Rd
index 0afd00a2..52ba7ea2 100644
--- a/man/dot-check_statistic_args.Rd
+++ b/man/dot-check_statistic_args.Rd
@@ -2,7 +2,7 @@
 % Please edit documentation in R/checks.R
 \name{.check_statistic_args}
 \alias{.check_statistic_args}
-\title{Check that the statistic and stat_max arguments are valid}
+\title{Check that the \code{statistic} and \code{stat_max} arguments are valid}
 \usage{
 .check_statistic_args(statistic, stat_max)
 }

From ab2d7e5ae49c8fd36b9c469d73b4757c10bc9d5e Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Wed, 14 Feb 2024 18:30:16 +0000
Subject: [PATCH 19/38] Use the right boolean

---
 R/checks.R | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index e2e850cd..fc7baf65 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -63,9 +63,9 @@
   )
 
   if (func_name == "simulate_chains") {
-    if (!generation_time_specified) {
+    if (generation_time_specified) {
       .check_generation_time_valid(generation_time)
-    } else if (!tf_specified) {
+    } else if (tf_specified) {
       stop("If `tf` is specified, `generation_time` must be specified too.")
     }
     checkmate::assert_numeric(

From b6759b5087a4a5e6907c39069b8e971a5f848425 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:06:07 +0000
Subject: [PATCH 20/38] Fix bracket indentation

Co-authored-by: Sebastian Funk <sebastian.funk@lshtm.ac.uk>
---
 R/simulate.r | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/simulate.r b/R/simulate.r
index 232b38ca..4d5ed950 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -199,7 +199,7 @@ simulate_chains <- function(index_cases,
       offspring_func_pars = pars,
       n_offspring = n_offspring,
       chains = sim
-      )
+    )
     # from all possible offspring, get those that could be infected
     next_gen <- .get_infectible_offspring(
       new_offspring = next_gen,

From 42099af4c804758163702ed71abb6d004459e84e Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:08:26 +0000
Subject: [PATCH 21/38] Replace word "infectible" with "susceptible"

Co-authored-by: Sebastian Funk <sebastian.funk@lshtm.ac.uk>
---
 R/helpers.R  | 8 ++++----
 R/simulate.r | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/R/helpers.R b/R/helpers.R
index 4b843045..b613c564 100644
--- a/R/helpers.R
+++ b/R/helpers.R
@@ -87,7 +87,7 @@
   return(possible_new_offspring)
 }
 
-#' Sample the number of infectible offspring from all possible offspring
+#' Sample the number of susceptible offspring from all possible offspring
 #'
 #' @description
 #' Sample susceptible offspring to be infected from all possible offspring.
@@ -100,19 +100,19 @@
 #' @return A vector of the number of offspring that can be infected given the
 #' current susceptible population size
 #' @keywords internal
-.get_infectible_offspring <- function(new_offspring,
+.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
-  infectible_offspring <- stats::rbinom(
+  susceptible_offspring <- stats::rbinom(
     n = length(new_offspring),
     size = new_offspring,
     prob = binom_prob
   )
-  return(infectible_offspring)
+  return(susceptible_offspring)
 }
 
 #' Adjust new offspring if it exceeds the susceptible population size
diff --git a/R/simulate.r b/R/simulate.r
index 4d5ed950..f3060719 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -201,7 +201,7 @@ simulate_chains <- function(index_cases,
       chains = sim
     )
     # from all possible offspring, get those that could be infected
-    next_gen <- .get_infectible_offspring(
+    next_gen <- .get_susceptible_offspring(
       new_offspring = next_gen,
       susc_pop = susc_pop,
       pop = pop

From de468343a3e38ad5cdae488603d0480f7852aceb Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:05:13 +0000
Subject: [PATCH 22/38] Remove cyclocomp_linter guards

---
 R/simulate.r | 2 --
 1 file changed, 2 deletions(-)

diff --git a/R/simulate.r b/R/simulate.r
index f3060719..82840526 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -128,7 +128,6 @@
 #' "Branching Process Models for Surveillance of Infectious Diseases
 #' Controlled by Mass Vaccination.” Biostatistics (Oxford, England)
 #' 4 (2): 279–95. \doi{https://doi.org/10.1093/biostatistics/4.2.279}.
-# nolint start: cyclocomp_linter.
 simulate_chains <- function(index_cases,
                             statistic = c("size", "length"),
                             offspring_dist,
@@ -299,7 +298,6 @@ simulate_chains <- function(index_cases,
   )
   return(out)
 }
-# nolint end
 
 #' Simulate a vector of transmission chains sizes/lengths
 #'

From 5d059b3c3f3cc5eaf0e2f39e2873626d91072213 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:34:28 +0000
Subject: [PATCH 23/38] Remove "infectible" from WORDLIST

---
 inst/WORDLIST | 1 -
 1 file changed, 1 deletion(-)

diff --git a/inst/WORDLIST b/inst/WORDLIST
index 12392f36..7ebaebb1 100644
--- a/inst/WORDLIST
+++ b/inst/WORDLIST
@@ -30,7 +30,6 @@ gborel
 geq
 infectee
 infectees
-infectible
 infector
 infectors
 json

From d907e43f6a7ad7b8244093e8aa76ca950a5495de Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:45:23 +0000
Subject: [PATCH 24/38] Remove need to check if generation_time is specified in
 main function

---
 R/checks.R                   | 6 +++---
 R/simulate.r                 | 7 ++-----
 tests/testthat/test-checks.R | 8 ++------
 3 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index fc7baf65..10edc85f 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -34,8 +34,8 @@
     stat_max,
     pop,
     percent_immune,
-    generation_time_specified,
-    generation_time = NULL,
+    generation_time,
+    t0,
     t0 = NULL,
     tf_specified,
     tf = NULL) {
@@ -63,7 +63,7 @@
   )
 
   if (func_name == "simulate_chains") {
-    if (generation_time_specified) {
+    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.")
diff --git a/R/simulate.r b/R/simulate.r
index 82840526..82596606 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -140,10 +140,8 @@ simulate_chains <- function(index_cases,
                             tf = Inf) {
   # Check inputs
   func_name <- as.character(match.call()[[1]])
-  # We need to check if the generation time and tf are specified here before
-  # calling .check_sim_args() to not obfuscate the meaning of missing()
-  # in the checker function.
-  generation_time_specified <- !missing(generation_time)
+  # Determine if tf is specified. Use to check if tf is specified
+  # but generation_time is not, which is an error.
   tf_specified <- !missing(tf)
   # Run checks
   .check_sim_args(
@@ -154,7 +152,6 @@ simulate_chains <- function(index_cases,
     stat_max = stat_max,
     pop = pop,
     percent_immune = percent_immune,
-    generation_time_specified = generation_time_specified,
     generation_time = generation_time,
     t0 = t0,
     tf_specified = tf_specified,
diff --git a/tests/testthat/test-checks.R b/tests/testthat/test-checks.R
index f7f25365..247ec576 100644
--- a/tests/testthat/test-checks.R
+++ b/tests/testthat/test-checks.R
@@ -24,7 +24,7 @@ test_that("Checks work", {
       stat_max = 10,
       pop = 10,
       percent_immune = 0.1,
-      generation_time_specified = FALSE,
+      generation_time = NULL,
       t0 = 0,
       tf_specified = TRUE,
       tf = 10
@@ -112,13 +112,12 @@ test_that("Checks work", {
       stat_max = 10,
       pop = 10,
       percent_immune = 0.1,
-      generation_time_specified = TRUE,
       generation_time = NULL,
       t0 = 0,
       tf_specified = TRUE,
       tf = 10
     ),
-    "Must be a function, not 'NULL'"
+    "If `tf` is specified, `generation_time` must be specified too."
   )
   expect_error(
     .check_sim_args(
@@ -129,7 +128,6 @@ test_that("Checks work", {
       stat_max = 10,
       pop = 10,
       percent_immune = 0.1,
-      generation_time_specified = TRUE,
       generation_time = function(x) rep(3, 10),
       t0 = -1,
       tf_specified = TRUE,
@@ -146,7 +144,6 @@ test_that("Checks work", {
       stat_max = 10,
       pop = 10,
       percent_immune = 0.1,
-      generation_time_specified = TRUE,
       generation_time = function(x) rep(3, 10),
       t0 = 0,
       tf_specified = TRUE,
@@ -163,7 +160,6 @@ test_that("Checks work", {
       stat_max = 10,
       pop = 10,
       percent_immune = 0.1,
-      generation_time_specified = TRUE,
       generation_time = function(x) rep(3, 10),
       t0 = 0,
       tf_specified = TRUE,

From cd7a58a6006e245e192ffc3f8a80b814a4cd1769 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:47:00 +0000
Subject: [PATCH 25/38] Move up check on tf

---
 R/checks.R | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index 10edc85f..9ed122e2 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -68,14 +68,14 @@
     } else if (tf_specified) {
       stop("If `tf` is specified, `generation_time` must be specified too.")
     }
-    checkmate::assert_numeric(
-      t0,
-      lower = 0, finite = TRUE
-    )
     checkmate::assert_number(
       tf,
       lower = 0
     )
+    checkmate::assert_numeric(
+      t0,
+      lower = 0, finite = TRUE
+    )
   }
   invisible(NULL)
 }

From aabef0e94b6b76eadffd4868cf7e6091ffbc60c1 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:47:19 +0000
Subject: [PATCH 26/38] Improve comments

---
 R/simulate.r | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/R/simulate.r b/R/simulate.r
index 82596606..7f6e38ff 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -138,7 +138,7 @@ simulate_chains <- function(index_cases,
                             generation_time = NULL,
                             t0 = 0,
                             tf = Inf) {
-  # Check inputs
+  # Get function name (safeguard against function name changes)
   func_name <- as.character(match.call()[[1]])
   # Determine if tf is specified. Use to check if tf is specified
   # but generation_time is not, which is an error.
@@ -365,8 +365,9 @@ simulate_summary <- function(index_cases,
                              stat_max = Inf,
                              pop = Inf,
                              percent_immune = 0) {
-  # Check inputs
+  # Get function name (safeguard against function name changes)
   func_name <- as.character(match.call()[[1]])
+  # Run checks
   .check_sim_args(
     func_name = func_name,
     index_cases = index_cases,

From 2b3a707f45a8d6cd6f065aef49e1b0a702cd9eff Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:51:21 +0000
Subject: [PATCH 27/38] Apply new function name

---
 R/simulate.r | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/simulate.r b/R/simulate.r
index 7f6e38ff..05f96b7e 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -398,7 +398,7 @@ simulate_summary <- function(index_cases,
       chains = sim
     )
     # from all possible offspring, get those that are infectible
-    next_gen <- .get_infectible_offspring(
+    next_gen <- .get_susceptible_offspring(
       new_offspring = next_gen,
       susc_pop = susc_pop,
       pop = pop

From 5984b87cc812bcf51432f599c52fe756bc7de588 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:54:36 +0000
Subject: [PATCH 28/38] Apply new function name

---
 R/checks.R                    | 3 +--
 tests/testthat/test-helpers.R | 4 ++--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index 9ed122e2..b2bca32e 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -36,9 +36,8 @@
     percent_immune,
     generation_time,
     t0,
-    t0 = NULL,
     tf_specified,
-    tf = NULL) {
+    tf) {
   # Get the function name
   func_name <- match.arg(func_name)
   # Input checking
diff --git a/tests/testthat/test-helpers.R b/tests/testthat/test-helpers.R
index 07d01170..0370478c 100644
--- a/tests/testthat/test-helpers.R
+++ b/tests/testthat/test-helpers.R
@@ -108,7 +108,7 @@ test_that(".init_susc_pop works correctly", {
 test_that(".init_susc_pop works correctly", {
   next_gen <- c(1, 2, 5)
   expect_length(
-    .get_infectible_offspring(
+    .get_susceptible_offspring(
       new_offspring = next_gen,
       susc_pop = 1,
       pop = 20
@@ -117,7 +117,7 @@ test_that(".init_susc_pop works correctly", {
   )
   # If the susceptible population in infinite, next_gen should be returned
   expect_identical(
-    .get_infectible_offspring(
+    .get_susceptible_offspring(
       new_offspring = next_gen,
       susc_pop = Inf,
       pop = Inf

From 384a5e5e333d53e7f62e73d546792671ab64aa80 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 12:55:09 +0000
Subject: [PATCH 29/38] Make extra arguments to optional

---
 R/checks.R | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index b2bca32e..78f81dc7 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -34,10 +34,10 @@
     stat_max,
     pop,
     percent_immune,
-    generation_time,
-    t0,
     tf_specified,
-    tf) {
+    tf = NULL,
+    generation_time = NULL,
+    t0 = NULL) {
   # Get the function name
   func_name <- match.arg(func_name)
   # Input checking

From c040b47c056956f47fc3006fd7c6858334789ea8 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:15:13 +0000
Subject: [PATCH 30/38] Generate doc

---
 ...ible_offspring.Rd => dot-get_susceptible_offspring.Rd} | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)
 rename man/{dot-get_infectible_offspring.Rd => dot-get_susceptible_offspring.Rd} (79%)

diff --git a/man/dot-get_infectible_offspring.Rd b/man/dot-get_susceptible_offspring.Rd
similarity index 79%
rename from man/dot-get_infectible_offspring.Rd
rename to man/dot-get_susceptible_offspring.Rd
index 770d13bf..b67f14d3 100644
--- a/man/dot-get_infectible_offspring.Rd
+++ b/man/dot-get_susceptible_offspring.Rd
@@ -1,10 +1,10 @@
 % Generated by roxygen2: do not edit by hand
 % Please edit documentation in R/helpers.R
-\name{.get_infectible_offspring}
-\alias{.get_infectible_offspring}
-\title{Sample the number of infectible offspring from all possible offspring}
+\name{.get_susceptible_offspring}
+\alias{.get_susceptible_offspring}
+\title{Sample the number of susceptible offspring from all possible offspring}
 \usage{
-.get_infectible_offspring(new_offspring, susc_pop, pop)
+.get_susceptible_offspring(new_offspring, susc_pop, pop)
 }
 \arguments{
 \item{new_offspring}{A vector of the possible new offspring per chain

From d94b92db2a8543f52dc70c58b76993db566d4412 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:15:48 +0000
Subject: [PATCH 31/38] Restructure tests to use a default function

---
 tests/testthat/test-checks.R | 312 ++++++++++++++++-------------------
 1 file changed, 138 insertions(+), 174 deletions(-)

diff --git a/tests/testthat/test-checks.R b/tests/testthat/test-checks.R
index 247ec576..ae8699cd 100644
--- a/tests/testthat/test-checks.R
+++ b/tests/testthat/test-checks.R
@@ -1,175 +1,139 @@
-test_that("Checks work", {
-  expect_error(
-    .check_offspring_func_valid(rrpois),
-    "not found"
-  )
-  expect_error(
-    .check_generation_time_valid("a"),
-    "Must be a function"
-  )
-  expect_error(
-    .check_generation_time_valid(function(x) rep("a", 10)),
-    "numeric"
-  )
-  expect_error(
-    .check_generation_time_valid(function(x) 3),
-    "Must have length"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_chains",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1,
-      generation_time = NULL,
-      t0 = 0,
-      tf_specified = TRUE,
-      tf = 10
-    ),
-    "If `tf` is specified, `generation_time` must be specified too."
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 0,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1
-    ),
-    "Must be >= 1."
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 10,
-      statistic = "a",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1
-    ),
-    "Must be element of set"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = r,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1
-    ),
-    "object 'r' not found"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = -1,
-      pop = 10,
-      percent_immune = 0.1
-    ),
-    "Element 1 is not >= 0"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 1,
-      pop = -1,
-      percent_immune = 0.1
-    ),
-    "Element 1 is not >= 1"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_summary",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 1,
-      pop = 10,
-      percent_immune = -0.1
-    ),
-    "Element 1 is not >= 0"
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_chains",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1,
-      generation_time = NULL,
-      t0 = 0,
-      tf_specified = TRUE,
-      tf = 10
-    ),
-    "If `tf` is specified, `generation_time` must be specified too."
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_chains",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1,
-      generation_time = function(x) rep(3, 10),
-      t0 = -1,
-      tf_specified = TRUE,
-      tf = Inf
-    ),
-    "Element 1 is not >= 0."
-  )
-  expect_error(
-    .check_sim_args(
-      func_name = "simulate_chains",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1,
-      generation_time = function(x) rep(3, 10),
-      t0 = 0,
-      tf_specified = TRUE,
-      tf = -1
-    ),
-    "Element 1 is not >= 0."
-  )
-  expect_no_error(
-    .check_sim_args(
-      func_name = "simulate_chains",
-      index_cases = 10,
-      statistic = "size",
-      offspring_dist = rpois,
-      stat_max = 10,
-      pop = 10,
-      percent_immune = 0.1,
-      generation_time = function(x) rep(3, 10),
-      t0 = 0,
-      tf_specified = TRUE,
-      tf = 10
-    )
-  )
-  expect_no_error(
-    .check_statistic_args(
-      statistic = "size",
-      stat_max = 10
-    )
-  )
+# Default function to check wrongly specified arguments and return an error
+# message
+.check_sim_args_default <- function(...) {
+    default_args <- list(
+        func_name = "simulate_chains",
+        index_cases = 10,
+        statistic = "size",
+        offspring_dist = rpois,
+        stat_max = 10,
+        pop = 10,
+        percent_immune = 0.1,
+        generation_time = function(x) rep(3, 10),
+        t0 = 0,
+        tf_specified = TRUE,
+        tf = 10
+    )
+    # Modify the default arguments with the user's arguments
+    new_args <- modifyList(
+        default_args,
+        list(...)
+    )
+    # Run the check_sim_args function and capture the output
+    out <- tryCatch(
+        expr = {
+            do.call(
+                .check_sim_args,
+                new_args
+            )
+        },
+        error = function(e) {
+            stop(e)
+        }
+    )
+    # Return the output
+    return(out)
+}
+
+test_that("Smaller checker functions work", {
+    expect_error(
+        .check_offspring_func_valid(rrpois),
+        "not found"
+    )
+    expect_error(
+        .check_generation_time_valid("a"),
+        "Must be a function"
+    )
+    expect_error(
+        .check_generation_time_valid(function(x) rep("a", 10)),
+        "numeric"
+    )
+    expect_error(
+        .check_generation_time_valid(function(x) 3),
+        "Must have length"
+    )
+    expect_no_error(
+        .check_statistic_args(
+            statistic = "size",
+            stat_max = 10
+        )
+    )
+})
+
+test_that(".check_sim_args() returns errors", {
+    expect_no_error(
+        .check_sim_args_default()
+    )
+    # index_cases must be >= 1
+    expect_error(
+        .check_sim_args_default(
+            index_cases = 0
+        ),
+        "Must be >= 1."
+    )
+    # statistic can only be "size" or "length"
+    expect_error(
+        .check_sim_args_default(
+            statistic = "duration"
+        ),
+        "Must be element of set \\{'size','length'\\}"
+    )
+    # offspring_dist must be a function
+    expect_error(
+        .check_sim_args_default(
+            offspring_dist = "rpois"
+        ),
+        "Must be a function, not 'character'"
+    )
+    # offspring_dist must be a known function (in the environment)
+    expect_error(
+        .check_sim_args_default(
+            offspring_dist = r
+        ),
+        "object 'r' not found"
+    )
+    # stat_max must be >= 1
+    expect_error(
+        .check_sim_args_default(
+            stat_max = 0
+        ),
+        "Assertion failed."
+    )
+    # pop cannot be negative
+    expect_error(
+        .check_sim_args_default(
+            pop = -1
+        ),
+        "Element 1 is not >= 1."
+    )
+    # percent_immune must be in between 0 and 1
+    expect_error(
+        .check_sim_args_default(
+            percent_immune = 1.1
+        ),
+        "Element 1 is not <= 1."
+    )
+    # t0 cannot be negative
+    expect_error(
+        .check_sim_args_default(
+            t0 = -1
+        ),
+        "Element 1 is not >= 0."
+    )
+    # tf cannot be negative
+    expect_error(
+        .check_sim_args_default(
+            tf = -1
+        ),
+        "Element 1 is not >= 0."
+    )
+    # If tf is specified, generation_time must be specified too
+    expect_error(
+        .check_sim_args_default(
+            generation_time = NULL,
+            tf_specified = TRUE,
+            tf = 10
+        ),
+        "If `tf` is specified, `generation_time` must be specified too."
+    )
 })

From 730ed365187f6646be2517d2c8dffad855aab589 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:16:06 +0000
Subject: [PATCH 32/38] Make lower bound of stat_max 1

---
 R/checks.R | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/checks.R b/R/checks.R
index 78f81dc7..5d8877cc 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -100,7 +100,7 @@
     is.infinite(stat_max) ||
       checkmate::assert_integerish(
         stat_max,
-        lower = 0,
+        lower = 1,
         null.ok = FALSE
     )
   )

From 4444341ba884fce918b7595644dba02018647a51 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:16:44 +0000
Subject: [PATCH 33/38] Use combine = or instead of ||

---
 R/checks.R | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/R/checks.R b/R/checks.R
index 5d8877cc..28be7ed2 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -97,11 +97,12 @@
     choices = c("size", "length")
   )
   checkmate::assert(
-    is.infinite(stat_max) ||
+    is.infinite(stat_max),
       checkmate::assert_integerish(
         stat_max,
         lower = 1,
         null.ok = FALSE
-    )
+    ),
+    combine = "or"
   )
 }

From 68c7a8dcce9f1a360a0fbe248e3826a68ca52c99 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:17:15 +0000
Subject: [PATCH 34/38] Use check_ instead of assert_ inside
 checkmate::assert()

---
 R/checks.R | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/R/checks.R b/R/checks.R
index 28be7ed2..0534f36f 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -98,7 +98,7 @@
   )
   checkmate::assert(
     is.infinite(stat_max),
-      checkmate::assert_integerish(
+      checkmate::check_integerish(
         stat_max,
         lower = 1,
         null.ok = FALSE

From c6baad3623330fead221fd0449f2d84c4317f834 Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Thu, 15 Feb 2024 17:17:32 +0000
Subject: [PATCH 35/38] Generate doc

---
 man/dot-check_sim_args.Rd | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/man/dot-check_sim_args.Rd b/man/dot-check_sim_args.Rd
index 84463ce3..047a2b45 100644
--- a/man/dot-check_sim_args.Rd
+++ b/man/dot-check_sim_args.Rd
@@ -12,11 +12,10 @@
   stat_max,
   pop,
   percent_immune,
-  generation_time_specified,
-  generation_time = NULL,
-  t0 = NULL,
   tf_specified,
-  tf = NULL
+  tf = NULL,
+  generation_time = NULL,
+  t0 = NULL
 )
 }
 \arguments{
@@ -52,6 +51,9 @@ infection at the start of the simulation. Used alongside \code{pop} to initialis
 the susceptible population. Accepted values lie between 0 and 1.
 Defaults to 0.}
 
+\item{tf}{Cut-off for the infection times (if generation time is given).
+Defaults to \code{Inf}.}
+
 \item{generation_time}{The generation time function; the name
 of a user-defined named or anonymous function with only one argument \code{n},
 representing the number of generation times to sample.}
@@ -61,9 +63,6 @@ or a vector of same length as \code{index_cases} (number of initial cases) with
 corresponding initial times. Defaults to 0, meaning all cases started at
 time 0.}
 
-\item{tf}{Cut-off for the infection times (if generation time is given).
-Defaults to \code{Inf}.}
-
 \item{sim_func}{\if{html}{\out{<character>}}; The simulation function to check}
 }
 \value{

From 856eb1b73035a4904ec77c38629a11e7df200f7a Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Mon, 19 Feb 2024 13:31:17 +0000
Subject: [PATCH 36/38] Use latest version of Roxygen

---
 DESCRIPTION | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index d1caca5b..2273eb84 100644
--- a/DESCRIPTION
+++ b/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

From 8bd77e8aac80fba17786a0ecc11f83483aac57df Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Mon, 19 Feb 2024 13:33:51 +0000
Subject: [PATCH 37/38] Move time checks in .check_sim_args() into
 .check_time_args()

---
 R/checks.R                 | 62 ++++++++++++++++++++++----------------
 R/simulate.r               | 22 ++++++--------
 man/dot-check_sim_args.Rd  | 19 +-----------
 man/dot-check_time_args.Rd | 35 +++++++++++++++++++++
 4 files changed, 82 insertions(+), 56 deletions(-)
 create mode 100644 man/dot-check_time_args.Rd

diff --git a/R/checks.R b/R/checks.R
index 0534f36f..55097a2f 100644
--- a/R/checks.R
+++ b/R/checks.R
@@ -27,19 +27,12 @@
 #' @return NULL; called for side effects
 #' @keywords internal
 .check_sim_args <- function(
-    func_name = c("simulate_chains", "simulate_summary"),
     index_cases,
     statistic,
     offspring_dist,
     stat_max,
     pop,
-    percent_immune,
-    tf_specified,
-    tf = NULL,
-    generation_time = NULL,
-    t0 = NULL) {
-  # Get the function name
-  func_name <- match.arg(func_name)
+    percent_immune) {
   # Input checking
   checkmate::assert_count(
     index_cases,
@@ -60,22 +53,6 @@
     percent_immune,
     lower = 0, upper = 1
   )
-
-  if (func_name == "simulate_chains") {
-    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
-    )
-  }
   invisible(NULL)
 }
 
@@ -86,8 +63,6 @@
 #' 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,
@@ -106,3 +81,38 @@
     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
+  )
+}
diff --git a/R/simulate.r b/R/simulate.r
index 05f96b7e..d1317921 100644
--- a/R/simulate.r
+++ b/R/simulate.r
@@ -138,20 +138,21 @@ simulate_chains <- function(index_cases,
                             generation_time = NULL,
                             t0 = 0,
                             tf = Inf) {
-  # Get function name (safeguard against function name changes)
-  func_name <- as.character(match.call()[[1]])
-  # Determine if tf is specified. Use to check if tf is specified
-  # but generation_time is not, which is an error.
-  tf_specified <- !missing(tf)
-  # Run checks
+  # Check offspring and population-related arguments
   .check_sim_args(
-    func_name = func_name,
     index_cases = index_cases,
     statistic = statistic,
     offspring_dist = offspring_dist,
     stat_max = stat_max,
     pop = pop,
-    percent_immune = percent_immune,
+    percent_immune = percent_immune
+  )
+  # Check time-related arguments
+  # Since tf is passed to .check_time_args, we need to check if it is specified
+  # in this function environment. If tf is specified, we expect generation_time
+  # to be specified too.
+  tf_specified <- !missing(tf)
+  .check_time_args(
     generation_time = generation_time,
     t0 = t0,
     tf_specified = tf_specified,
@@ -365,11 +366,8 @@ simulate_summary <- function(index_cases,
                              stat_max = Inf,
                              pop = Inf,
                              percent_immune = 0) {
-  # Get function name (safeguard against function name changes)
-  func_name <- as.character(match.call()[[1]])
-  # Run checks
+  # Check offspring and population-related arguments
   .check_sim_args(
-    func_name = func_name,
     index_cases = index_cases,
     statistic = statistic,
     offspring_dist = offspring_dist,
diff --git a/man/dot-check_sim_args.Rd b/man/dot-check_sim_args.Rd
index 047a2b45..fd3172e2 100644
--- a/man/dot-check_sim_args.Rd
+++ b/man/dot-check_sim_args.Rd
@@ -5,17 +5,12 @@
 \title{Check inputs to \code{simulate_chains()} and \code{simulate_summary()}}
 \usage{
 .check_sim_args(
-  func_name = c("simulate_chains", "simulate_summary"),
   index_cases,
   statistic,
   offspring_dist,
   stat_max,
   pop,
-  percent_immune,
-  tf_specified,
-  tf = NULL,
-  generation_time = NULL,
-  t0 = NULL
+  percent_immune
 )
 }
 \arguments{
@@ -51,18 +46,6 @@ infection at the start of the simulation. Used alongside \code{pop} to initialis
 the susceptible population. Accepted values lie between 0 and 1.
 Defaults to 0.}
 
-\item{tf}{Cut-off for the infection times (if generation time is given).
-Defaults to \code{Inf}.}
-
-\item{generation_time}{The generation time function; the name
-of a user-defined named or anonymous function with only one argument \code{n},
-representing the number of generation times to sample.}
-
-\item{t0}{Start time (if generation time is given); either a single value
-or a vector of same length as \code{index_cases} (number of initial cases) with
-corresponding initial times. Defaults to 0, meaning all cases started at
-time 0.}
-
 \item{sim_func}{\if{html}{\out{<character>}}; The simulation function to check}
 }
 \value{
diff --git a/man/dot-check_time_args.Rd b/man/dot-check_time_args.Rd
new file mode 100644
index 00000000..6c6c31a7
--- /dev/null
+++ b/man/dot-check_time_args.Rd
@@ -0,0 +1,35 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/checks.R
+\name{.check_time_args}
+\alias{.check_time_args}
+\title{Check inputs that control time events}
+\usage{
+.check_time_args(tf_specified, tf, generation_time, t0)
+}
+\arguments{
+\item{tf_specified}{\if{html}{\out{<logical>}}; Whether the \code{tf} argument is specified. Only
+makes sense in the context where this function is called, i.e., in
+\code{\link[=simulate_chains]{simulate_chains()}}. If \code{tf} is specified, generation_time must be specified.}
+
+\item{tf}{Cut-off for the infection times (if generation time is given).
+Defaults to \code{Inf}.}
+
+\item{generation_time}{The generation time function; the name
+of a user-defined named or anonymous function with only one argument \code{n},
+representing the number of generation times to sample.}
+
+\item{t0}{Start time (if generation time is given); either a single value
+or a vector of same length as \code{index_cases} (number of initial cases) with
+corresponding initial times. Defaults to 0, meaning all cases started at
+time 0.}
+}
+\value{
+NULL; called for side effects
+}
+\description{
+This function checks the time-related inputs, i.e., start time of each chain,
+\code{t0}, the end time of the simulation, \code{tf}, and the generation time,
+generation_time. It also checks that the generation_time argument is
+specified if \code{tf} is specified as these go hand-in-hand.
+}
+\keyword{internal}

From 62df189d9f118defc130cc5fa052aa19046bb2bb Mon Sep 17 00:00:00 2001
From: James Azam <james.m.azam@gmail.com>
Date: Mon, 19 Feb 2024 13:34:09 +0000
Subject: [PATCH 38/38] Add tests for .check_time_args()

---
 tests/testthat/test-checks.R | 54 ++++++++++++++++++++++++++++--------
 1 file changed, 42 insertions(+), 12 deletions(-)

diff --git a/tests/testthat/test-checks.R b/tests/testthat/test-checks.R
index ae8699cd..9e234c94 100644
--- a/tests/testthat/test-checks.R
+++ b/tests/testthat/test-checks.R
@@ -1,18 +1,13 @@
-# Default function to check wrongly specified arguments and return an error
-# message
+# A function that passes all checks in .check_sim_args but returns an error
+# message if the supplied arguments are invalid
 .check_sim_args_default <- function(...) {
     default_args <- list(
-        func_name = "simulate_chains",
         index_cases = 10,
         statistic = "size",
         offspring_dist = rpois,
         stat_max = 10,
         pop = 10,
-        percent_immune = 0.1,
-        generation_time = function(x) rep(3, 10),
-        t0 = 0,
-        tf_specified = TRUE,
-        tf = 10
+        percent_immune = 0.1
     )
     # Modify the default arguments with the user's arguments
     new_args <- modifyList(
@@ -35,6 +30,34 @@
     return(out)
 }
 
+# A function that passes all checks in `.check_time_args` but returns an error
+# message if the supplied arguments are invalid
+.check_time_args_default <- function(...) {
+    default_args <- list(
+        t0 = 0,
+        tf_specified = FALSE, # tf is not specified but default tf = Inf below
+        tf = Inf,
+        generation_time = NULL
+    )
+    new_args <- modifyList(
+        default_args,
+        list(...)
+    )
+    out <- tryCatch(
+        expr = {
+            do.call(
+                .check_time_args,
+                new_args
+            )
+        },
+        error = function(e) {
+            stop(e)
+        }
+    )
+    return(out)
+}
+
+
 test_that("Smaller checker functions work", {
     expect_error(
         .check_offspring_func_valid(rrpois),
@@ -61,6 +84,7 @@ test_that("Smaller checker functions work", {
 })
 
 test_that(".check_sim_args() returns errors", {
+    # Checks with .check_sim_args
     expect_no_error(
         .check_sim_args_default()
     )
@@ -113,24 +137,30 @@ test_that(".check_sim_args() returns errors", {
         ),
         "Element 1 is not <= 1."
     )
+})
+
+test_that(".check_time_args() returns errors", {
+    # Checks with .check_time_args
+    expect_no_error(
+        .check_time_args_default()
+    )
     # t0 cannot be negative
     expect_error(
-        .check_sim_args_default(
+        .check_time_args_default(
             t0 = -1
         ),
         "Element 1 is not >= 0."
     )
     # tf cannot be negative
     expect_error(
-        .check_sim_args_default(
+        .check_time_args_default(
             tf = -1
         ),
         "Element 1 is not >= 0."
     )
     # If tf is specified, generation_time must be specified too
     expect_error(
-        .check_sim_args_default(
-            generation_time = NULL,
+        .check_time_args_default(
             tf_specified = TRUE,
             tf = 10
         ),