Make likelihood() work with <epichains_summary>, <epichains>, and <numeric> objects

R/checks.R

@@ -69,12 +69,12 @@
     statistic,
     choices = c("size", "length")
   )
+  # check that stat_threshold is an integer or Inf.
   checkmate::assert(
-    is.infinite(stat_threshold),
+    checkmate::anyInfinite(stat_threshold),
       checkmate::check_integerish(
         stat_threshold,
-        lower = 1,
-        null.ok = FALSE
+        lower = 1
     ),
     combine = "or"
   )

R/helpers.R

@@ -150,3 +150,11 @@
   next_gen[as.integer(names(next_gen_count))] <- unname(next_gen_count)
   return(next_gen)
 }
+
+#' Extract the `stat_threshold` attribute from an `<epichains>` object.
+#'
+#' @param obj An <epichains> or <epichains_summary> object.
+#' @keywords internal
+.get_stat_threshold <- function(obj) {
+  attr(obj, "stat_threshold")
+}

R/likelihood.R

@@ -2,7 +2,10 @@
 #'
 #' @inheritParams .offspring_ll
 #' @inheritParams simulate_chain_stats
-#' @param chains A numeric vector of chain summaries (sizes/lengths).
+#' @param chains Vector of chain summaries (sizes/lengths). Can be a
+#' `<numeric>` vector or an object of class `<epichains>` or
+#' `<epichains_summary>` (obtained from [simulate_chains()] or
+#' [simulate_chain_stats()]). See examples below.
 #' @param nsim_obs Number of simulations to be used to approximate the
 #' log-likelihood/likelihood if `obs_prob < 1` (imperfect observation). If
 #' `obs_prob == 1`, this argument is ignored.
@@ -11,6 +14,14 @@
 #' simulations will be used to approximate the (log)likelihood. This will
 #' also require specifying `nsim_obs`. In the simulation, the observation
 #' process is assumed to be constant.
+#' @param stat_threshold A stopping criterion for individual chain simulations;
+#' a positive number coercible to integer. When any chain's cumulative statistic
+#' reaches or surpasses `stat_threshold`, that chain ends. It also serves as a
+#' censoring limit so that results above the specified value, are set to `Inf`.
+#' Defaults to `Inf`. NOTE: For objects of class `<epichains>` or
+#' `<epichains_summary>`, the `stat_threshold` used in the simulation is
+#' extracted and used here so if this argument is specified, it is ignored and
+#' a warning is thrown.
 #' @param log Should the log-likelihoods be transformed to
 #' likelihoods? Logical. Defaults to `TRUE`.
 #' @param exclude A vector of indices of the sizes/lengths to exclude from the
@@ -34,31 +45,82 @@
 #' except that likelihoods, instead of log-likelihoods, are calculated in all
 #' cases. Moreover, the joint likelihoods are the product, instead of the sum,
 #' of the individual likelihoods.
-#' @author Sebastian Funk
+#' @author Sebastian Funk, James M. Azam
 #' @examples
 #' # example of observed chain sizes
 #' set.seed(121)
 #' # randomly generate 20 chains of size 1 to 10
 #' chain_sizes <- sample(1:10, 20, replace = TRUE)
 #' likelihood(
-#'   chains = chain_sizes, statistic = "size",
-#'   offspring_dist = rpois, nsim_obs = 100, lambda = 0.5
+#'   chains = chain_sizes,
+#'   statistic = "size",
+#'   offspring_dist = rpois,
+#'   nsim_obs = 100,
+#'   lambda = 0.5
 #' )
+#' # Example using an <epichains> object
+#' set.seed(32)
+#' epichains_obj_eg <- simulate_chains(
+#'  n_chains = 10,
+#'  pop = 100,
+#'  percent_immune = 0,
+#'  statistic = "size",
+#'  offspring_dist = rnbinom,
+#'  stat_threshold = 10,
+#'  generation_time = function(n) rep(3, n),
+#'  mu = 2,
+#'  size = 0.2
+#')
+#'
+#' epichains_obj_eg_lik <- likelihood(
+#'   chains = epichains_obj_eg,
+#'   statistic = "size",
+#'   offspring_dist = rnbinom,
+#'   mu = 2,
+#'   size = 0.2,
+#'   stat_threshold = 10
+#' )
+#' epichains_obj_eg_lik
+#'
+#' # Example using a <epichains_summary> object
+#' set.seed(32)
+#' epichains_summary_eg <- simulate_chain_stats(
+#'   n_chains = 10,
+#'   pop = 100,
+#'   percent_immune = 0,
+#'   statistic = "size",
+#'   offspring_dist = rnbinom,
+#'   stat_threshold = 10,
+#'   mu = 2,
+#'   size = 0.2
+#' )
+#' epichains_summary_eg_lik <- likelihood(
+#'   chains = epichains_summary_eg,
+#'   statistic = "size",
+#'   offspring_dist = rnbinom,
+#'   mu = 2,
+#'   size = 0.2,
+#'   stat_threshold = 10
+#' )
+#' epichains_summary_eg_lik
 #' @export
+#nolint start: cyclocomp_linter
 likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
-                       nsim_obs, obs_prob = 1, log = TRUE, stat_threshold = Inf,
+                       nsim_obs, obs_prob = 1, stat_threshold = Inf, log = TRUE,
                        exclude = NULL, individual = FALSE, ...) {
   statistic <- match.arg(statistic)
 
   ## Input checking
-  ## Check nsim_obs when specified
-  if (!missing(nsim_obs)) {
-    checkmate::assert_number(
-      nsim_obs, lower = 1, finite = TRUE, na.ok = FALSE
+  checkmate::assert(
+    checkmate::check_numeric(
+      chains, lower = 0, upper = Inf, any.missing = FALSE
+    ),
+    checkmate::check_class(
+      chains, "epichains"
+    ),
+    checkmate::check_class(
+      chains, "epichains_summary"
     )
-  }
-  checkmate::assert_numeric(
-    chains, lower = 0, upper = Inf, any.missing = FALSE
   )
   # check that arguments related to the statistic are valid
   .check_statistic_args(
@@ -78,9 +140,52 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
   checkmate::assert_numeric(
     exclude, null.ok = TRUE
   )
+  # likelihood cannot work with an <epichains> object so convert to
+  # <epichains_summary>
+  if (.is_epichains(chains)) {
+    chains <- summary(chains)
+  }
+
+  # Logic:
+  # * If the object is an <epichains>/<epichains_summary>, we'll give
+  # preference to the stat_threshold used in the simulation.
+  if (.is_epichains_summary(chains)) {
+    if (!missing(stat_threshold)) {
+    warning(
+      "`stat_threshold` specified but will be ignored. ",
+      "Using `stat_threshold` = ",
+      .get_stat_threshold(chains),
+      " as used in the simulation.",
+      call. = FALSE
+      )
+    }
+    stat_threshold <- .get_stat_threshold(chains)
+  }
+
+  if (is.finite(stat_threshold)) {
+    # censor the chains to be at most stat_threshold
+    chains <- pmin(chains, stat_threshold)
+  } else if (any(is.infinite(chains)) && !.is_epichains_summary(chains)) {
+      # Developer note: <epichains_summary> objects also pass the`is.numeric()`
+      # check, so we need to check if the object is not an <epichains_summary>
+      # object.
+      # Numeric chains can only contain Inf values based on human error/decision
+      # to censor it with infinite values, so we ask the user to fix that.
+      stop(
+        "`chains` must be censored with a finite value. ",
+        "Replace the `Inf` values with a finite value.",
+        call. = FALSE
+      )
+  }
+
+  # Apply the observation process
   if (obs_prob < 1) {
     if (missing(nsim_obs)) {
       stop("'nsim_obs' must be specified if 'obs_prob' is < 1")
+    } else {
+      checkmate::assert_integerish(
+        nsim_obs, lower = 1
+      )
     }
 
     statistic_func <- .get_statistic_func(statistic)
@@ -102,7 +207,7 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
     stat_rep_list <- list(chains)
   }
 
-  ## determine for which sizes to calculate the log-likelihood
+  ## Determine for which sizes to calculate the log-likelihood
   ## (for true chain size)
   if (any(stat_rep_vect == stat_threshold)) {
     calc_sizes <- seq_len(stat_threshold - 1)
@@ -160,7 +265,7 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
   ## assign likelihoods
   chains_likelihood <- lapply(stat_rep_list, function(sx) {
     likelihoods[sx[!(sx %in% exclude)]]
-  }
+    }
   )
 
   ## if individual == FALSE, return the joint log-likelihood
@@ -180,3 +285,4 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
 
   return(chains_likelihood)
 }
+#nolint end