Issue 188: Add delay samples functionality

NAMESPACE

@@ -34,7 +34,6 @@ export(epidist_stancode)
 export(epidist_validate)
 export(epidist_version_stanvar)
 export(event_to_incidence)
-export(extract_lognormal_draws)
 export(filter_obs_by_obs_time)
 export(filter_obs_by_ptime)
 export(is_latent_individual)
@@ -50,6 +49,8 @@ export(plot_empirical_delay)
 export(plot_mean_posterior_pred)
 export(plot_recovery)
 export(plot_relative_recovery)
+export(predict_delay_parameters)
+export(predict_dpar)
 export(reverse_obs_at)
 export(simulate_double_censored_pmf)
 export(simulate_exponential_cases)
@@ -69,7 +70,6 @@ importFrom(checkmate,assert_names)
 importFrom(checkmate,assert_numeric)
 importFrom(cli,cli_abort)
 importFrom(cli,cli_inform)
-importFrom(posterior,as_draws_df)
 importFrom(rstan,lookup)
 importFrom(stats,ecdf)
 importFrom(stats,integrate)

R/globals.R

@@ -5,15 +5,6 @@ utils::globalVariables(c(
   "max_treedepth", # <epidist_diagnostics>
   "per_at_max_treedepth", # <epidist_diagnostics>
   "samples", # <epidist_diagnostics>
-  "mu", # <extract_lognormal_draws>
-  "Intercept", # <extract_lognormal_draws>
-  "sigma", # <extract_lognormal_draws>
-  "Intercept_sigma", # <extract_lognormal_draws>
-  "id", # <extract_lognormal_draws>
-  "true_value", # <make_relative_to_truth>
-  "value", # <make_relative_to_truth>
-  "rel_value", # <make_relative_to_truth>
-  "value", # <summarise_variable>
   "id", # <as_latent_individual.data.frame>
   "obs_t", # <as_latent_individual.data.frame>
   "obs_at", # <as_latent_individual.data.frame>
@@ -55,6 +46,12 @@ utils::globalVariables(c(
   "censored_obs_time", # <pad_zero>
   "delay_lwr", # <pad_zero>
   "delay_daily", # <pad_zero>
+  "delay_daily", # <calculate_cohort_mean>
+  "ptime_daily", # <calculate_cohort_mean>
+  "n", # <calculate_cohort_mean>
+  "obs_horizon", # <calculate_truncated_means>
+  "meanlog", # <calculate_truncated_means>
+  "sdlog", # <calculate_truncated_means>
   "value", # <plot_relative_recovery>
   "rel_value", # <plot_relative_recovery>
   "case_type", # <plot_cases_by_obs_window>
@@ -70,6 +67,10 @@ utils::globalVariables(c(
   "ptime_daily", # <plot_mean_posterior_pred>
   "n", # <plot_mean_posterior_pred>
   "obs_horizon", # <plot_mean_posterior_pred>
+  "true_value", # <make_relative_to_truth>
+  "value", # <make_relative_to_truth>
+  "rel_value", # <make_relative_to_truth>
+  "value", # <summarise_variable>
   "mu", # <add_mean_sd.lognormal_samples>
   "sigma", # <add_mean_sd.lognormal_samples>
   "sd", # <add_mean_sd.lognormal_samples>

R/plot-helpers.R

@@ -5,6 +5,7 @@
 #' @param by ...
 #' @param obs_at ...
 #' @family plot
+#' @autoglobal
 #' @export
 calculate_cohort_mean <- function(data, type = c("cohort", "cumulative"),
                                   by = c(), obs_at) {
@@ -34,6 +35,7 @@ calculate_cohort_mean <- function(data, type = c("cohort", "cumulative"),
 #' @param distribution ...
 #' @family plot
 #' @importFrom stats integrate
+#' @autoglobal
 #' @export
 calculate_truncated_means <- function(draws, obs_at, ptime,
                                       distribution = function(x, y, z) {

R/postprocess.R

@@ -1,3 +1,136 @@
+#' Extract samples of the delay distribution parameters
+#'
+#' @param fit A model fit with `epidist::epidist`
+#' @inheritParams brms::prepare_predictions
+#' @param ... Additional arguments passed to `brms::prepare_predictions`
+#' @family postprocess
+#' @autoglobal
+#' @export
+predict_delay_parameters <- function(fit, newdata = NULL, ...) {
+  if (!is.null(newdata)) {
+    newdata <- brms:::validate_newdata(newdata, fit)
+  }
+  pp <- brms::prepare_predictions(fit, newdata = newdata, ...)
+  # Every brms model has the parameter mu
+  lp_mu <- brms::get_dpar(pp, dpar = "mu", inv_link = TRUE)
+  df <- expand.grid(
+    "index" = seq_len(ncol(lp_mu)),
+    "draw" = seq_len(nrow(lp_mu))
+  )
+  df[["mu"]] <- as.vector(lp_mu)
+  for (dpar in setdiff(names(pp$dpars), "mu")) {
+    lp_dpar <- brms::get_dpar(pp, dpar = dpar, inv_link = TRUE)
+    df[[dpar]] <- as.vector(lp_dpar)
+  }
+  class(df) <- c(
+    class(df), paste0(sub(".*_", "", fit$family$name), "_samples")
+  )
+  dt <- as.data.table(df)
+  dt <- add_mean_sd(dt)
+  return(dt)
+}
+
+#' @rdname predict_delay_parameters
+#' @export
+predict_dpar <- predict_delay_parameters
+
+#' Convert posterior lognormal samples to long format
+#'
+#' @param draws ...
+#' @family postprocess
+#' @export
+draws_to_long <- function(draws) {
+  long_draws <- data.table::melt(
+    draws,
+    measure.vars = c("mu", "sigma", "mean", "sd"),
+    variable.name = "parameter"
+  )
+  return(long_draws[])
+}
+
+#' Make posterior lognormal samples relative to true values
+#'
+#' @param draws ...
+#' @param secondary_dist ...
+#' @param by ...
+#' @family postprocess
+#' @autoglobal
+#' @export
+make_relative_to_truth <- function(draws, secondary_dist, by = "parameter") {
+  draws <- merge(
+    draws,
+    secondary_dist[, true_value := value][, value := NULL],
+    by = by
+  )
+
+  draws[, rel_value := value / true_value]
+
+  return(draws[])
+}
+
+#' Summarise posterior draws
+#'
+#' @param draws A data.table of posterior draws
+#' @param sf The number of significant figures to use
+#' @param not_by ...
+#' @param by A vector of columns to group by
+#' @family postprocess
+#' @importFrom stats median quantile
+#' @export
+summarise_draws <- function(draws, sf, not_by = "value", by) {
+  if (missing(by)) {
+    by_cols <- setdiff(
+      colnames(draws), not_by
+    )
+  }else {
+    by_cols <- by
+  }
+
+  summarised_draws <- draws[,
+    list(
+      mean = mean(value, na.rm = TRUE),
+      median = median(value, na.rm = TRUE),
+      q2.5 = quantile(value, 0.025, na.rm = TRUE),
+      q5 = quantile(value, 0.05, na.rm = TRUE),
+      q20 = quantile(value, 0.2, na.rm = TRUE),
+      q35 = quantile(value, 0.35, na.rm = TRUE),
+      q65 = quantile(value, 0.65, na.rm = TRUE),
+      q80 = quantile(value, 0.8, na.rm = TRUE),
+      q95 = quantile(value, 0.95, na.rm = TRUE),
+      q97.5 = quantile(value, 0.975, na.rm = TRUE)
+    ),
+    by = by_cols
+  ]
+
+  if (!missing(sf)) {
+    cols <- setdiff(colnames(summarised_draws), by_cols)
+    summarised_draws <- summarised_draws[,
+      (cols) := lapply(.SD, signif, digits = sf),
+      .SDcols = cols
+    ]
+  }
+
+  return(summarised_draws[])
+}
+
+#' Summarise a variable
+#'
+#' @param variable The variable to summarise
+#' @inheritParams summarise_draws
+#'
+#' @family postprocess
+#' @autoglobal
+#' @export
+summarise_variable <- function(draws, variable, sf = 6, by = c()) {
+  if (missing(variable)) {
+    stop("variable must be specified")
+  }
+  summarised_draws <- data.table::copy(draws)
+  summarised_draws[, value := variable, env = list(variable = variable)]
+  summarised_draws <- summarise_draws(summarised_draws, sf = sf, by = by)
+  return(summarised_draws[])
+}
+
 #' Add natural scale mean and standard deviation parameters
 #'
 #' @param data A dataframe of distributional parameters