Add some model components

NAMESPACE

@@ -12,6 +12,7 @@ export(generate_simulated_data)
 export(get_count_data_sizes)
 export(get_count_indices)
 export(get_count_values)
+export(get_draws_df)
 export(get_ind_m)
 export(get_mcmc_options)
 export(get_model_diagnostic_flags)
@@ -30,7 +31,6 @@ export(indicate_ww_exclusions)
 export(make_hospital_onset_delay_pmf)
 export(make_incubation_period_pmf)
 export(make_reporting_delay_pmf)
-export(postprocess)
 export(preprocess_hosp_data)
 export(preprocess_ww_data)
 export(simulate_double_censored_pmf)

R/delay_distribs.R

@@ -1,4 +1,4 @@
-#' Drop the first element of a simplex
+#' Drop the first element of a simplex and re-normalize the result to sum to 1.
 #'
 #' When this vector corresponds to the generation interval distribution, we
 #' want to drop this first bin. The renewal equation assumes that same-day
@@ -77,9 +77,9 @@ simulate_double_censored_pmf <- function(
 }
 
 #' @title Make incubation period pmf
-#' @description This makes a pmf corresponding to
-#' the incubation period for COVID after Omicron used in Park et al 2023
-#' These estimates are from early Omicron.
+#' @description When the default arguments are used, this returns a pmf
+#' corresponding to the incubation period for COVID after Omicron used in
+#' Park et al 2023. These estimates are from early Omicron.
 #' @param backward_scale numeric indicating the scale parameter for the Weibull
 #' used in producing the incubation period distribution. default is `3.60` for
 #' COVID
@@ -113,15 +113,15 @@ make_incubation_period_pmf <- function(backward_scale = 3.60,
   # Relies on fundamental assumption about epidemic growth rate.
 
 
-  corrected_sgtf <- tibble::tibble(
+  discr_gr_adj_weibull <- tibble::tibble(
     time = seq(0, 23, by = 1), # 23 seems to get most of the distribution mass
     density0 = dweibull(time,
       shape = backward_shape,
       scale = backward_scale
     ) * exp(r * time)
   )
 
-  inc_period_pmf <- corrected_sgtf$density0 / sum(corrected_sgtf$density0)
+  inc_period_pmf <- to_simplex(discr_gr_adj_weibull$density0)
   return(inc_period_pmf)
 }
 
@@ -180,6 +180,6 @@ make_reporting_delay_pmf <- function(incubation_period_pmf,
   )
 
   infection_to_hosp_delay_pmf <- add_pmfs(pmfs) |>
-    (\(x) x / sum(x))()
+    to_simplex()
   return(infection_to_hosp_delay_pmf)
 }

R/figures.R

@@ -1,22 +1,25 @@
 #' Get plot of fit and forecasted counts
 #'
 #' @param draws A dataframe containing the posterior draws with the data joined
-#' to it. This is `draws_df` output of the call to `wwinference()`
+#' to it. This is the `draws_df` output of a call to [wwinference()]
 #' @param count_data_eval A dataframe containing the count data we will
 #' evaluate the forecasts against.
+#' @param count_data_eval_col_name string indicating the name of the count
+#' data to evaluate against the forecasted count data
 #' @param forecast_date A string indicating the date we made the forecast, for
 #' plotting, in ISO8601 format YYYY-MM-DD
 #' @param count_type A string indicating what data the counts refer to,
 #' default is `hospital admissions`
 #'
 #' @return A ggplot object containing the posterior draw of the estimated,
-#' nowcasted, and forecasted hospital admissions alongside the data used to
-#' calibrate the hospital admissions model and the later observed hospital
-#' admissions used to evaluate the forecast performance
+#' nowcasted, and forecasted counts alongside the data used to
+#' calibrate the model and subsequently observed counts (if any) against which
+#' to evaluate the forecast performance.
 #' @export
 #'
 get_plot_forecasted_counts <- function(draws,
                                        count_data_eval,
+                                       count_data_eval_col_name,
                                        forecast_date,
                                        count_type = "hospital admissions") {
   p <- ggplot(draws |> dplyr::filter(
@@ -28,7 +31,7 @@ get_plot_forecasted_counts <- function(draws,
     ) +
     geom_point(
       data = count_data_eval,
-      aes(x = date, y = daily_hosp_admits_for_eval),
+      aes(x = date, y = .data[[count_data_eval_col_name]]),
       shape = 21, color = "black", fill = "white"
     ) +
     geom_point(aes(x = date, y = observed_value)) +
@@ -61,7 +64,7 @@ get_plot_forecasted_counts <- function(draws,
 #' Get plot of fit and forecasted wastewater concentrations
 #'
 #' @param draws A dataframe containing the posterior draws with the data joined
-#' to it. This is `draws_df` output of the call to `wwinference()`
+#' to it. This is the `draws_df` output of a call to [wwinference()]
 #' @param forecast_date A string indicating the date we made the forecast, for
 #' plotting, in ISO8601 format YYYY-MM-DD
 #'
@@ -120,7 +123,7 @@ get_plot_ww_conc <- function(draws,
 #' Get plot of fit, nowcasted, and forecasted "global" R(t)
 #'
 #' @param draws A dataframe containing the posterior draws with the data joined
-#' to it. This is `draws_df` output of the call to `wwinference()`
+#' to it. This is the `draws_df` output of a call to [wwinference()]
 #' @param forecast_date A string indicating the date we made the forecast, for
 #' plotting, in ISO8601 format YYYY-MM-DD
 #'
@@ -143,8 +146,8 @@ get_plot_global_rt <- function(draws,
     ) +
     geom_hline(aes(yintercept = 1), linetype = "dashed") +
     xlab("") +
-    ylab("R(t) of hypothetical state") +
-    ggtitle("R(t) estimate") +
+    ylab("Global R(t)") +
+    ggtitle("Global R(t) estimate") +
     scale_x_date(
       date_breaks = "2 weeks",
       labels = scales::date_format("%Y-%m-%d")

R/model_diagnostics.R

@@ -13,8 +13,10 @@
 #' (estimated bayesian fraction of missing information), default is `0.2`
 #' @param divergences_tolerance float indicating the tolerance for the
 #' proportion of sampling iterations that are divergent, default is `0.01`
-#' @param p_high_rhat_tolerance float indicating the tolerance for the
-#' proportion of parameters rhats>1.05, default is `0.05`
+#' @param frac_high_rhat_tolerance float indicating the tolerance for the
+#' proportion of parameters rhats>rhat_tolderance, default is `0.05`
+#' @param rhat_tolerance float indicating the tolerance for the rhat for
+#' individual parameters, default is `1.05`
 #' @param max_tree_depth_tol float indicating the tolerance for the proportion
 #' of iterations that exceed the maximum tree depth, default is `0.01`
 #'
@@ -25,7 +27,8 @@
 get_model_diagnostic_flags <- function(stan_fit_object,
                                        ebmfi_tolerance = 0.2,
                                        divergences_tolerance = 0.01,
-                                       p_high_rhat_tolerance = 0.05,
+                                       frac_high_rhat_tolerance = 0.05,
+                                       rhat_tolerance = 1.05,
                                        max_tree_depth_tol = 0.01) {
   n_chains <- stan_fit_object$num_chains()
   iter_sampling <- stan_fit_object$metadata()$iter_sampling
@@ -39,9 +42,11 @@ get_model_diagnostic_flags <- function(stan_fit_object,
   flag_too_many_divergences <- any(
     diagnostic_summary$num_divergent >= max_n_divergences
   )
-  p_high_rhat <- as.numeric(mean(summary[, "rhat"]$rhat > 1.05, na.rm = TRUE))
-  flag_high_rhat <- p_high_rhat >=
-    p_high_rhat_tolerance
+  frac_high_rhat <- as.numeric(mean(summary[, "rhat"]$rhat > rhat_tolerance,
+    na.rm = TRUE
+  ))
+  flag_high_rhat <- frac_high_rhat >=
+    frac_high_rhat_tolerance
   max_n_max_treedepth <- n_chains * iter_sampling * max_tree_depth_tol
   flag_high_max_treedepth <- any(
     diagnostic_summary$num_max_tree_depth >= max_n_max_treedepth

R/postprocess.R

@@ -29,12 +29,12 @@
 #' are observations, the data will be joined to each draw of the predicted
 #' observation to facilitate plotting.
 #' @export
-postprocess <- function(ww_data,
-                        count_data,
-                        fit_obj,
-                        date_time_spine,
-                        lab_site_spine,
-                        subpop_spine) {
+get_draws_df <- function(ww_data,
+                         count_data,
+                         fit_obj,
+                         date_time_spine,
+                         lab_site_spine,
+                         subpop_spine) {
   draws <- fit_obj$result$draws()
 
   count_draws <- draws |>

R/preprocessing.R

@@ -1,4 +1,5 @@
-#' Get input wastewater data
+#' Pre-process wastewater input data, adding needed indices and flagging
+#' potential outliers
 #' @param ww_data dataframe containing the following columns: site, lab,
 #' date, a column for concentration, and lod
 #' @param conc_col_name string indicating the name of the column containing
@@ -69,7 +70,8 @@ preprocess_ww_data <- function(ww_data,
   return(ww_preprocessed)
 }
 
-#' Get input hospital admissions data
+#' Pre-process hospital admissions data, converting column names to those
+#' that [get_stan_data()] expects.
 #' @param hosp_data dataframe containing the following columns: date,
 #' a count column, and a population size column
 #' @param count_col_name name of the column containing the epidemiological

R/utils.R

@@ -53,12 +53,12 @@ add_pmfs <- function(pmfs) {
 }
 
 #' @title Get index matrix
-#' @description Get the matrix needed to convert a vetor from weekly to daily
+#' @description Get a matrix to broadcast a vector from weekly to daily
 #' @param n_days number of days we will expand to
 #' @param n_weeks number of weeks those days correspond to
 #'
 #' @return a n_day x n_week matrix for multiplying by weekly estimated
-#' value to conver it to daily
+#' value to broadcast it to daily
 #' @export
 #'
 #' @examples
@@ -81,7 +81,7 @@ get_ind_m <- function(n_days, n_weeks) {
 #' @title Create a new directory if one doesn't exist
 #' @description
 #' Function to create a directory for the specified output file path if needed.
-#' dir_create won't throw a warning if its already made though!
+#' Does nothing if the target directory already exists.
 #'
 #'
 #' @param output_file_path file path that may or may not need to be created
@@ -94,30 +94,35 @@ create_dir <- function(output_file_path) {
   }
 }
 
-#' @title Convert to logmean in lognorm distribution
+#' @title Get the mean of a Normal distribution for a random variable Y
+#' needed to ensure that the distribution of X = exp(Y) (which is Log-Normal)
+#' has a specified mean and sd.
 #' @description
 #'  see arithmetic moments here
 #' https://en.wikipedia.org/wiki/Log-normal_distribution
 #'
-#' @param mean mean of the normal distribution
-#' @param sd sd of the normal distribution
+#' @param mean target mean for the Log-Normal distribution of X
+#' @param sd target sd for the Log-Normal distribution X
 #'
-#' @return corresponding mean of the lognormal distribution
+#' @return corresponding mean for the underlying Normal
+#' distribution of Y = log(X).
 #' @export
 convert_to_logmean <- function(mean, sd) {
   logmean <- log(mean^2 / sqrt(sd^2 + mean^2))
   return(logmean)
 }
 
 
-#' @title Convert to logsd in lognormal distribution
-#' @description@description see arithmetic moments here
+#' @title Get the sd of a Normal distribution for a random variable Y
+#' needed to ensure that the distribution of X = exp(Y) (which is Log-Normal)
+#' has a specified mean and sd.
+#' @description see arithmetic moments here
 #' https://en.wikipedia.org/wiki/Log-normal_distribution
 #'
-#' @param mean mean of the normal distribution
-#' @param sd sd of the normal distribution
+#' @param mean target mean for the Log-Normal distribution of X
+#' @param sd target sd for the Log-Normal distribution of X
 #'
-#' @return corresponding stdev of the lognormal distribution
+#' @return corresponding sd for the underlying Normal distribution of Y = log(X)
 #' @export
 convert_to_logsd <- function(mean, sd) {
   logsd <- sqrt(log(1 + (sd^2 / mean^2)))

R/wwinference.R

@@ -149,7 +149,7 @@ wwinference <- function(ww_data,
         ]
       ))
 
-    draws <- postprocess(
+    draws <- get_draws_df(
       ww_data = ww_data,
       count_data = count_data,
       fit_obj = fit,