Refactor likelihood functions

R/likelihood.R

@@ -1,93 +1,104 @@
-#' Estimate the (log) likelihood for observed branching processes
+#' Estimate the log-likelihood/likelihood for observed branching processes
 #'
 #' @inheritParams simulate_summary
-#' @param chains Vector of sizes/lengths of transmission chains.
-#' @param nsim_obs Number of simulations if the likelihood is to be
-#' approximated for imperfect observations.
-#' @param log Logical; Should the results be log-transformed? (Defaults
-#' to TRUE).
+#' @inheritParams offspring_ll
+#' @param nsim_obs Number of simulations if the log-likelihood/likelihood is to
+#' be approximated for imperfect observations.
+#' @param log Logical; Should the log-likelihoods be transformed to
+#' likelihoods? (Defaults to TRUE).
 #' @param obs_prob Observation probability (assumed constant)
 #' @param exclude A vector of indices of the sizes/lengths to exclude from the
-#' likelihood calculation.
-#' @param individual If TRUE, a vector of individual (log)likelihood
+#' log-likelihood calculation.
+#' @param individual If TRUE, a vector of individual log-likelihood/likelihood
 #' contributions will be returned rather than the sum.
-#' @param ... Parameters for the offspring distribution.
 #' @return
-#' * A log-likelihood, if \code{log = TRUE} (the default)
 #' * A vector of log-likelihoods, if \code{log = TRUE} (the default) and
 #' \code{obs_prob < 1}, or
 #' * A list of individual log-likelihood contributions, if
 #' \code{log = TRUE} (the default) and \code{individual = TRUE}.
-#' else raw likelihoods, or vector of likelihoods
-#' @seealso offspring_ll, pois_size_ll, nbinom_size_ll, gborel_size_ll,
-#' pois_length_ll, geom_length_ll.
+#' The interpretation follows for the other combinations of `log` and
+#' `individual`.
+#' @seealso offspring_ll(), pois_size_ll(), nbinom_size_ll(), gborel_size_ll(),
+#' pois_length_ll(), geom_length_ll()
 #' @author Sebastian Funk
 #' @examples
 #' # example of observed chain sizes
-#' chain_sizes <- c(1, 1, 4, 7)
+#' set.seed(121)
+#' chain_sizes <- sample(1:10, 20, replace = TRUE)
 #' likelihood(
 #'   chains = chain_sizes, statistic = "size",
 #'   offspring_dist = "pois", nsim_obs = 100, lambda = 0.5
 #' )
 #' @export
 likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
-                       nsim_obs, log = TRUE, obs_prob = 1, stat_max = Inf,
+                       nsim_obs = 100, log = TRUE, obs_prob = 1, stat_max = Inf,
                        exclude = NULL, individual = FALSE, ...) {
   statistic <- match.arg(statistic)
 
   ## checks
   check_offspring_valid(offspring_dist)
+  checkmate::assert_logical(
+    log,
+    any.missing = FALSE,
+    all.missing = FALSE,
+    len = 1
+  )
 
-  if (obs_prob <= 0 || obs_prob > 1) stop("'obs_prob' must be within (0,1]")
+  if (obs_prob <= 0 || obs_prob > 1) {
+    stop("'obs_prob' is a probability and must be between 0 and 1 inclusive")
+  }
   if (obs_prob < 1) {
     if (missing(nsim_obs)) {
       stop("'nsim_obs' must be specified if 'obs_prob' is < 1")
     }
 
-    sample_func <- get_statistic_func(statistic)
+    statistic_func <- get_statistic_func(statistic)
 
     sampled_x <- replicate(nsim_obs, pmin(
-      sample_func(
+      statistic_func(
         length(chains),
         chains, obs_prob
       ),
       stat_max
     ), simplify = FALSE)
     size_x <- unlist(sampled_x)
-    if (!is.finite(stat_max)) {
-      stat_max <- max(size_x) + 1
-    }
+    stat_max <- max(size_x) + 1
   } else {
     chains[chains >= stat_max] <- stat_max
     size_x <- chains
     sampled_x <- list(chains)
   }
 
-  ## determine for which sizes to calculate the likelihood (for true chain size)
+  ## determine for which sizes to calculate the log-likelihood
+  ## (for true chain size)
   if (any(size_x == stat_max)) {
     calc_sizes <- seq_len(stat_max - 1)
   } else {
     calc_sizes <- unique(c(size_x, exclude))
   }
 
-  ## get likelihood function as given by offspring_dist and statistic
+  ## get log-likelihood function as given by offspring_dist and statistic
   likelihoods <- vector(mode = "numeric")
   ll_func <- construct_offspring_ll_name(offspring_dist, statistic)
   pars <- as.list(unlist(list(...))) ## converts vectors to lists
 
-  ## calculate likelihoods
+  ## calculate log-likelihoods
   if (exists(ll_func, where = asNamespace("epichains"), mode = "function")) {
     func <- get(ll_func)
     likelihoods[calc_sizes] <- do.call(func, c(list(x = calc_sizes), pars))
   } else {
     likelihoods[calc_sizes] <-
       do.call(
         offspring_ll,
-        c(list(
-          chains = calc_sizes, offspring_dist = offspring_dist,
-          statistic = statistic, stat_max = stat_max,
-          log = log
-        ), pars)
+        c(
+          list(
+            chains = calc_sizes,
+            offspring_dist = offspring_dist,
+            statistic = statistic,
+            stat_max = stat_max
+          ),
+          pars
+        )
       )
   }
 
@@ -110,8 +121,20 @@ likelihood <- function(chains, statistic = c("size", "length"), offspring_dist,
     likelihoods[sx[!(sx %in% exclude)]]
   })
 
-  if (!individual) {
-    chains_likelihood <- vapply(chains_likelihood, sum, 0)
+  ## transform log-likelihoods into likelihoods if required
+  if (!isTRUE(log)) {
+    chains_likelihood <- lapply(chains_likelihood, exp)
+  }
+
+  ## if individual == FALSE, return the joint log-likelihood
+  ## (sum of the log-likelihoods), if log == TRUE, else
+  ## multiply the likelihoods
+  if (!isTRUE(individual)) {
+    if (isTRUE(log)) {
+      chains_likelihood <- vapply(chains_likelihood, sum, 0)
+    } else {
+      chains_likelihood <- vapply(chains_likelihood, prod, 0)
+    }
   }
 
   return(chains_likelihood)

R/stat_likelihoods.R

@@ -1,4 +1,4 @@
-#' Likelihood of the size of chains with Poisson offspring distribution
+#' Log-likelihood of the size of chains with Poisson offspring distribution
 #'
 #' @param x vector of sizes
 #' @param lambda rate of the Poisson distribution
@@ -9,7 +9,7 @@ pois_size_ll <- function(x, lambda) {
   (x - 1) * log(lambda) - lambda * x + (x - 2) * log(x) - lgamma(x)
 }
 
-#' Likelihood of the size of chains with Negative-Binomial offspring
+#' Log-likelihood of the size of chains with Negative-Binomial offspring
 #' distribution
 #'
 #' @param x vector of sizes
@@ -31,7 +31,7 @@ nbinom_size_ll <- function(x, size, prob, mu) {
     (size * x + (x - 1)) * log(1 + mu / size)
 }
 
-#' Likelihood of the size of chains with gamma-Borel offspring distribution
+#' Log-likelihood of the size of chains with gamma-Borel offspring distribution
 #'
 #' @param x vector of sizes
 #' @param size the dispersion parameter (often called \code{k} in ecological
@@ -52,9 +52,9 @@ gborel_size_ll <- function(x, size, prob, mu) {
     (x - 1) * log(x) - (size + x - 1) * log(x + size / mu)
 }
 
-#' Likelihood of the length of chains with Poisson offspring distribution
+#' Log-likelihood of the length of chains with Poisson offspring distribution
 #'
-#' @param x vector of sizes
+#' @param x vector of lengths
 #' @param lambda rate of the Poisson distribution
 #' @return log-likelihood values
 #' @author Sebastian Funk
@@ -70,9 +70,9 @@ pois_length_ll <- function(x, lambda) {
   log(Gk[x + 1] - Gk[x])
 }
 
-#' Likelihood of the length of chains with geometric offspring distribution
+#' Log-likelihood of the length of chains with geometric offspring distribution
 #'
-#' @param x vector of sizes
+#' @param x vector of lengths
 #' @param prob probability of the geometric distribution with mean
 #' \code{1/prob}
 #' @return log-likelihood values
@@ -86,43 +86,61 @@ geom_length_ll <- function(x, prob) {
   log(GkmGkm1)
 }
 
-#' Likelihood of the length of chains with generic offspring distribution
+#' Log-likelihood of the summary (size/length) of chains with generic offspring
+#' distribution
 #'
-#' The likelihoods are calculated with a crude approximation using simulated
-#' chains by linearly approximating any missing values in the empirical
+#' The log-likelihoods are calculated with a crude approximation using simulated
+#' chain summaries by linearly approximating any missing values in the empirical
 #' cumulative distribution function (ecdf).
 #' @inheritParams likelihood
-#' @inheritParams simulate_vec
-#' @param chains Vector of sizes/lengths
+#' @inheritParams simulate_summary
+#' @param chains Vector of chain summaries (sizes/lengths)
 #' @param nsim_offspring Number of simulations of the offspring distribution
-#' for approximating the statistic (size/length) distribution
-#' @param log Logical; Should the results be log-transformed? (Defaults
-#' to TRUE).
-#' @param ... any parameters to pass to \code{\link{simulate_tree}}
-#' @return If \code{log = TRUE} (the default), log-likelihood values,
-#' else raw likelihoods
+#' for approximating the distribution of the chain statistic summary
+#' (size/length)
+#' @param ... any parameters to pass to \code{\link{simulate_summary}}
+#' @return log-likelihood values
 #' @author Sebastian Funk
 #' @export
+#' @seealso [simulate_summary()] for simulating a summary of the transmission
+#' chains statistic (without the tree of infections)
+#' @examples
+#' set.seed(123)
+#' chain_size_ll <- offspring_ll(
+#'   chains = c(1, 5, 6, 8, 7, 8, 10),
+#'   offspring_dist = "pois",
+#'   statistic = "size",
+#'   lambda = 0.82
+#' )
 offspring_ll <- function(chains, offspring_dist, statistic,
-                         nsim_offspring = 100, log = TRUE, ...) {
+                         nsim_offspring = 100, ...) {
   # Simulate the chains
-  chains <- simulate_summary(
-    nsim_offspring, offspring_dist,
-    statistic, ...
+  dist <- simulate_summary(
+    nchains = nsim_offspring,
+    offspring_dist = offspring_dist,
+    statistic = statistic,
+    ...
   )
 
   # Compute the empirical Cumulative Distribution Function of the
   # simulated chains
-  chains_empirical_cdf <- stats::ecdf(chains)
+  chains_empirical_cdf <- stats::ecdf(dist)
 
   # Perform a lagged linear interpolation of the points
-  acdf <-
-    diff(c(0, stats::approx(
-      unique(chains), chains_empirical_cdf(unique(chains)),
-      seq_len(max(chains[is.finite(chains)]))
-    )$y))
+  acdf <- diff(
+    c(
+      0,
+      stats::approx(
+        unique(dist),
+        chains_empirical_cdf(unique(dist)),
+        seq_len(
+          max(dist[is.finite(dist)])
+        )
+      )$y
+    )
+  )
   lik <- acdf[chains]
   lik[is.na(lik)] <- 0
-  out <- ifelse(base::isTRUE(log), log(lik), lik)
+  out <- log(lik)
   return(out)
 }