Issue 248: Add prediction and log-likelihood methods for the latent_gamma family (#273)

* Add quick go at these three functions for the latent_gamma family

* Lint package

* Rename predict tests as latent_lognormal tests

* Make generate_examples reproducible and add gamma data

* Add unit tests for the latent_gamma methods

* Use the fit_gamma rather than fit

* Upload gamma fit

* Add named documentation to latent_lognormal

* Reorder prep and i correctly

* Use [s] not [i]! Bug

* Rerun of generate test data required

* Use scale parameterisation

Former-commit-id: dd22e59
Former-commit-id: a2cd30a3f99b55572aeb286351078f133d2cf630
Former-commit-id: 3789a48cc40ec22f4128b208318686e7de538a27 [formerly 878955d]
Former-commit-id: 42ab0638f573e74ff809fc46db450a3cd86b9b29

Author: athowes

NAMESPACE

@@ -42,6 +42,7 @@ export(filter_obs_by_ptime)
 export(is_latent_individual)
 export(linelist_to_cases)
 export(linelist_to_counts)
+export(log_lik_latent_gamma)
 export(log_lik_latent_lognormal)
 export(make_relative_to_truth)
 export(observe_process)
@@ -53,7 +54,9 @@ export(plot_empirical_delay)
 export(plot_mean_posterior_pred)
 export(plot_recovery)
 export(plot_relative_recovery)
+export(posterior_epred_latent_gamma)
 export(posterior_epred_latent_lognormal)
+export(posterior_predict_latent_gamma)
 export(posterior_predict_latent_lognormal)
 export(predict_delay_parameters)
 export(predict_dpar)

R/latent_gamma.R

@@ -0,0 +1,85 @@
+#' Draws from the posterior predictive distribution of the `latent_gamma` family
+#'
+#' See [brms::posterior_predict()].
+#'
+#' @param i The index of the observation to predict
+#' @param prep The result of a call to [brms::posterior_predict()]
+#' @param ... Additional arguments
+#' @family postprocess
+#' @autoglobal
+#' @export
+posterior_predict_latent_gamma <- function(i, prep, ...) { # nolint: object_length_linter
+  mu <- brms::get_dpar(prep, "mu", i = i)
+  shape <- brms::get_dpar(prep, "shape", i = i)
+
+  obs_t <- prep$data$vreal1[i]
+  pwindow_width <- prep$data$vreal2[i]
+  swindow_width <- prep$data$vreal3[i]
+
+  .predict <- function(s) {
+    d_censored <- obs_t + 1
+    # while loop to impose the truncation
+    while (d_censored > obs_t) {
+      p_latent <- runif(1, 0, 1) * pwindow_width
+      d_latent <- rgamma(1, shape = shape[s], scale = mu[s] / shape[s])
+      s_latent <- p_latent + d_latent
+      p_censored <- floor_mult(p_latent, pwindow_width)
+      s_censored <- floor_mult(s_latent, swindow_width)
+      d_censored <- s_censored - p_censored
+    }
+    return(d_censored)
+  }
+
+  # Within brms this is a helper function called rblapply
+  do.call(rbind, lapply(seq_len(prep$ndraws), .predict))
+}
+
+#' Draws from the expected value of the posterior predictive distribution of the
+#' `latent_gamma` family
+#'
+#' See [brms::posterior_epred()].
+#'
+#' @param prep The result of a call to [`brms::prepare_predictions`]
+#' @family postprocess
+#' @autoglobal
+#' @export
+posterior_epred_latent_gamma <- function(prep) { # nolint: object_length_linter
+  mu <- brms::get_dpar(prep, "mu")
+  mu
+}
+
+#' Calculate the pointwise log likelihood of the `latent_gamma` family
+#'
+#' See [brms::log_lik()].
+#'
+#' @param i The index of the observation to calculate the log likelihood of
+#' @param prep The result of a call to [brms::prepare_predictions()]
+#' @family postprocess
+#' @autoglobal
+#' @export
+log_lik_latent_gamma <- function(i, prep) {
+  mu <- brms::get_dpar(prep, "mu", i = i)
+  shape <- brms::get_dpar(prep, "shape", i = i)
+  y <- prep$data$Y[i]
+  obs_t <- prep$data$vreal1[i]
+  pwindow_width <- prep$data$vreal2[i]
+  swindow_width <- prep$data$vreal3[i]
+
+  swindow_raw <- runif(prep$ndraws)
+  pwindow_raw <- runif(prep$ndraws)
+
+  swindow <- swindow_raw * swindow_width
+
+  # For no overlap calculate as usual, for overlap ensure pwindow < swindow
+  if (i %in% prep$data$noverlap) {
+    pwindow <- pwindow_raw * pwindow_width
+  } else {
+    pwindow <- pwindow_raw * swindow
+  }
+
+  d <- y - pwindow + swindow
+  obs_time <- obs_t - pwindow
+  lpdf <- dgamma(d, shape = shape, scale = mu / shape, log = TRUE)
+  lcdf <- pgamma(obs_time, shape = shape, scale = mu / shape, log.p = TRUE)
+  return(lpdf - lcdf)
+}

R/latent_lognormal.R

@@ -1,4 +1,5 @@
-#' Draws from the posterior predictive distribution
+#' Draws from the posterior predictive distribution of the `latent_lognormal`
+#' family
 #'
 #' See [brms::posterior_predict()].
 #'
@@ -34,7 +35,8 @@ posterior_predict_latent_lognormal <- function(i, prep, ...) { # nolint: object_
   do.call(rbind, lapply(seq_len(prep$ndraws), .predict))
 }
 
-#' Draws from the expected value of the posterior predictive distribution
+#' Draws from the expected value of the posterior predictive distribution of the
+#' `latent_gamma` family
 #'
 #' See [brms::posterior_epred()].
 #'
@@ -48,7 +50,7 @@ posterior_epred_latent_lognormal <- function(prep) { # nolint: object_length_lin
   exp(mu + sigma^2 / 2)
 }
 
-#' Calculate the pointwise log likelihood
+#' Calculate the pointwise log likelihood of the `latent_gamma` family
 #'
 #' See [brms::log_lik()].
 #'

inst/extdata/fit_gamma.rds.REMOVED.git-id

@@ -0,0 +1 @@
+ed630da445eec989cc02d70f845a0ec5272a3d82

inst/generate_examples.R

@@ -1,4 +1,11 @@
 source("tests/testthat/setup.R")
+set.seed(1)
 prep_obs <- as_latent_individual(sim_obs)
-fit <- epidist(prep_obs)
-saveRDS(fit, "inst/extdata/fit.rds")
+fit <- epidist(prep_obs, seed = 1)
+prep_obs_gamma <- as_latent_individual(sim_obs_gamma)
+fit_gamma <- epidist(
+  prep_obs_gamma,
+  family = stats::Gamma(link = "log"),
+  seed = 1
+)
+saveRDS(fit_gamma, "inst/extdata/fit_gamma.rds")