Merge pull request #88 from epiforecasts/81-add-option-for-real-time-hosp-data

Issue 81: Add option to fit to real-time hospital admissions data

Author: kaitejohnson

DESCRIPTION

@@ -25,6 +25,7 @@ Depends:
     R (>= 4.0.0)
 Imports:
     cli,
+    data.table,
     forecast,
     dplyr,
     tidyr,
@@ -34,7 +35,8 @@ Imports:
     wwinference,
     lubridate,
     fs,
-    purrr,
+    arrow,
+    tidyselect,
     glue,
     scoringutils,
     zoo,

NAMESPACE

@@ -5,7 +5,10 @@ export(quantile_metrics)
 export(quiet)
 export(sample_metrics)
 export(save_csv)
+export(trajectories_to_quantiles)
 importFrom(cli,cli_warn)
+importFrom(data.table,as.data.table)
+importFrom(data.table,setattr)
 importFrom(dplyr,arrange)
 importFrom(dplyr,case_when)
 importFrom(dplyr,desc)
@@ -23,6 +26,7 @@ importFrom(forecast,auto.arima)
 importFrom(forecast,forecast)
 importFrom(fs,dir_create)
 importFrom(ggplot2,aes)
+importFrom(ggplot2,element_text)
 importFrom(ggplot2,facet_wrap)
 importFrom(ggplot2,geom_bar)
 importFrom(ggplot2,geom_line)
@@ -32,6 +36,7 @@ importFrom(ggplot2,geom_vline)
 importFrom(ggplot2,ggplot)
 importFrom(ggplot2,ggsave)
 importFrom(ggplot2,ggtitle)
+importFrom(ggplot2,theme)
 importFrom(ggplot2,theme_bw)
 importFrom(ggplot2,xlab)
 importFrom(ggplot2,ylab)

R/EDA_plots.R

@@ -152,7 +152,7 @@ get_plot_model_comparison <- function(
   }
   full_fp <- file.path(fig_fp, this_location)
   if (!file.exists(full_fp)) {
-    dir_create(full_fp, recursive = TRUE, showWarnings = FALSE)
+    dir_create(full_fp, recurse = TRUE)
   }
   ggsave(
     plot = p,
@@ -169,14 +169,15 @@ get_plot_model_comparison <- function(
 #'
 #' @param draws_w_data Data.frame of draws with data
 #' @param full_fp Directory to save
-#' @importFrom ggplot2 geom_vline
+#' @importFrom ggplot2 geom_vline theme element_text
 #' @returns ggplot object
 #' @autoglobal
 get_plot_draws_w_calib_data <- function(draws_w_data,
                                         full_fp) {
   loc <- unique(draws_w_data$location)
   include_ww <- unique(draws_w_data$include_ww)
   forecast_date <- unique(draws_w_data$forecast_date)
+  hosp_data_real_time <- unique(draws_w_data$hosp_data_real_time)
   n_draws <- max(draws_w_data$draw, na.rm = TRUE)
   draws <- draws_w_data |> dplyr::filter(
     draw %in% sample.int(n_draws, size = min(100, n_draws))
@@ -194,14 +195,15 @@ get_plot_draws_w_calib_data <- function(draws_w_data,
     ) +
     xlab("") +
     theme_bw() +
+    theme(plot.title = element_text(size = 10)) +
     geom_vline(aes(xintercept = forecast_date), linetype = "dashed") +
     ylab("7-day rolling sum of hospital admissions") +
-    ggtitle(glue("location: {loc}, include_ww: {include_ww}, forecast_date: {forecast_date}")) # nolint
+    ggtitle(glue("location: {loc}, include_ww: {include_ww}, forecast_date: {forecast_date}, hosp data real time: {hosp_data_real_time}")) # nolint
   ggsave(
     plot = p,
     filename = file.path(
       full_fp,
-      glue::glue("7d_hosp_draws_w_data_ww_{include_ww}.png")
+      glue::glue("7d_hosp_draws_w_data_ww_{include_ww}_rt_{hosp_data_real_time}.png") # nolint
     )
   )
   return(p)
@@ -217,8 +219,8 @@ get_plot_draws_w_calib_data <- function(draws_w_data,
 #' @autoglobal
 get_bar_chart_overall_scores <- function(scores) {
   scores_summarised <- scores |>
-    summarise_scores(by = c("model", "include_ww")) |>
-    mutate(model_ww = glue::glue("{model}-{include_ww}"))
+    summarise_scores(by = c("model", "include_ww", "hosp_data_real_time")) |>
+    mutate(model_ww = glue::glue("{model}-{include_ww}-{hosp_data_real_time}"))
 
   p <- ggplot(scores_summarised) +
     geom_bar(
@@ -232,5 +234,56 @@ get_bar_chart_overall_scores <- function(scores) {
     ) +
     theme_bw() +
     ggtitle("Scores across all locations and forecast dates")
+
+  scores_by_loc <- scores |>
+    summarise_scores(by = c(
+      "model", "include_ww",
+      "hosp_data_real_time", "forecast_date"
+    )) |>
+    mutate(model_ww = glue::glue("{model}-{include_ww}-{hosp_data_real_time}"))
+  p <- ggplot(scores_by_loc) +
+    geom_bar(
+      aes(
+        x = forecast_date,
+        y = wis,
+        fill = model_ww
+      ),
+      stat = "identity",
+      position = "dodge"
+    ) +
+    theme_bw() +
+    theme(legend.position = "bottom") +
+    ggtitle("Scores across all locations by forecast dates")
+  return(p)
+}
+
+#' Get bar chart of the scores by forecast date
+#'
+#' @param scores Data.frame of scores from across locations and forecast dates
+#'
+#' @importFrom ggplot2 geom_bar
+#' @importFrom scoringutils summarise_scores
+#' @returns ggplot object
+#' @autoglobal
+get_plot_scores_by_date <- function(scores) {
+  scores_by_loc <- scores |>
+    summarise_scores(by = c(
+      "model", "include_ww",
+      "hosp_data_real_time", "forecast_date"
+    )) |>
+    mutate(model_ww = glue::glue("{model}-{include_ww}-{hosp_data_real_time}"))
+  p <- ggplot(scores_by_loc) +
+    geom_bar(
+      aes(
+        x = forecast_date,
+        y = wis,
+        fill = model_ww
+      ),
+      stat = "identity",
+      position = "dodge"
+    ) +
+    theme_bw() +
+    theme(legend.position = "bottom") +
+    ggtitle("Scores across all locations by forecast dates")
   return(p)
 }

R/convert_rolling_sum_to_inc.R

@@ -38,10 +38,6 @@ convert_rolling_sum_to_inc <- function(rolling_sums,
     stop("rolling_sums cannot be empty", call. = FALSE)
   }
 
-  if (anyNA(rolling_sums)) {
-    warning("rolling_sums contains NA values. Function expects right-aligned rolling sums", call. = FALSE) # nolint
-  }
-
   # Handle initial values
   if (is.null(initial_values)) {
     # If no initial values provided, assume the first 6 days were zeros

R/convert_to_su_object.R

@@ -0,0 +1,21 @@
+#' Convert the scores to a scoringutils object
+#'
+#' @param scores_data Data.frame from Variant Nowcast Hub GitHub
+#' @importFrom data.table setattr as.data.table
+#' @importFrom rlang arg_match
+#' @importFrom dplyr rename select
+#' @returns scoringutils object
+convert_to_su_object <- function(scores_data) {
+  scores2 <- data.table::as.data.table(scores_data)
+  class(scores2) <- c("scores", class(scores2))
+  scores_su <- data.table::setattr(
+    scores2,
+    "metrics",
+    c(
+      "wis", "underprediction", "overprediction", "dispersion",
+      "bias", "interval_coverage_50", "interval_coverage_90",
+      "ae_median"
+    )
+  )
+  return(scores_su)
+}

R/fit_arima.R

@@ -31,7 +31,7 @@ fit_arima <- function(hosp_data_for_fit,
       date >= ymd(forecast_date),
       date <= ymd(forecast_date) + days(forecast_horizon)
     )
-
+  hosp_data_real_time <- unique(hosp_data_for_fit$hosp_data_real_time)
   auto_arima_model <- auto.arima(hosp_data_for_fit$updated_hosp_7d_count,
     seasonal = FALSE,
     stepwise = FALSE,
@@ -68,6 +68,9 @@ fit_arima <- function(hosp_data_for_fit,
     left_join(hosp_data_eval_forecast,
       by = "date"
     ) |>
-    mutate(forecast_date = ymd(forecast_date))
+    mutate(
+      forecast_date = ymd(forecast_date),
+      hosp_data_real_time = hosp_data_real_time
+    )
   return(forecast_df)
 }

R/fit_wwinference_wrapper.R

@@ -14,6 +14,8 @@
 #' @param quantiles_to_save Vector of numerics indicating the quantiles
 #' @param ind_filepath Character string of the file path to save the outputs
 #'   from each model run
+#' @param save_draws Boolean indicating whether or not to save the draws,
+#'   default is FALSE.
 #'
 #' @returns Data.frame of the quantiles alongside the evaluation data.
 #' @autoglobal
@@ -34,9 +36,18 @@ fit_wwinference_wrapper <- function(
     calibration_time = 90,
     forecast_horizon = 28,
     quantiles_to_save = c(0.025, 0.1, 0.25, 0.5, 0.75, 0.9, 0.975),
-    ind_filepath = file.path("output")) {
+    ind_filepath = file.path("output"),
+    save_draws = FALSE) {
   loc <- unique(count_data$state)
+  if ((nrow(ww_data) == 0 || is.null(ww_data)) &&
+    isTRUE(model_spec$include_ww)) {
+    model_spec$include_ww <- FALSE
+    flag_missing_ww <- TRUE
+  } else {
+    flag_missing_ww <- FALSE
+  }
   include_ww <- model_spec$include_ww
+  hosp_data_real_time <- unique(count_data$hosp_data_real_time)
   ww_fit_obj <- wwinference(
     ww_data = ww_data,
     count_data = count_data,
@@ -54,32 +65,37 @@ fit_wwinference_wrapper <- function(
   # Save plots
   full_fp <- file.path(ind_filepath, this_forecast_date, loc)
   if (!file.exists(file.path(full_fp))) {
-    dir_create(full_fp, recursive = TRUE, showWarnings = FALSE)
+    dir_create(full_fp, recurse = TRUE)
   }
   fig_fp <- file.path(full_fp, "figs")
   if (!file.exists(file.path(fig_fp))) {
-    dir_create(fig_fp, recursive = TRUE, showWarnings = FALSE)
+    dir_create(fig_fp, recurse = TRUE)
   }
 
   plot_hosp_draws <- get_plot_forecasted_counts(
     draws = hosp_draws,
     forecast_date = this_forecast_date
-  ) + ggtitle(glue("{loc}, wastewater: {include_ww}"))
+  ) + ggtitle(glue("{loc}, wastewater: {include_ww}, hosp data real-time: {hosp_data_real_time}")) # nolint
 
   ggsave(
     plot = plot_hosp_draws,
     filename = file.path(
       fig_fp,
-      glue("hosp_draws_ww_{include_ww}.png")
+      glue("hosp_draws_ww_{include_ww}_rt_{hosp_data_real_time}.png")
     )
   )
   ww_draws <- if (!is.null(ww_fit_obj$raw_input_data$input_ww_data)) {
     get_draws(ww_fit_obj, what = "predicted_ww")$predicted_ww
   } else {
     NULL
   }
+  data_fp <- file.path(full_fp, "data")
+  if (!file.exists(file.path(data_fp))) {
+    dir_create(data_fp, recurse = TRUE)
+  }
 
   if (!is.null(ww_draws)) {
+    # Plot
     plot_ww_draws <- get_plot_ww_conc(
       draws = ww_draws,
       forecast_date = this_forecast_date
@@ -91,6 +107,43 @@ fit_wwinference_wrapper <- function(
         "ww_draws.png"
       )
     )
+    ww_data_obs <- select(
+      ww_data,
+      date, site, lab,
+      log_genome_copies_per_ml, below_lod,
+      log_lod, flag_as_ww_outlier
+    )
+    ww_metadata <- ww_data |>
+      select(
+        site, lab, site_pop,
+        location_name, location_abbr,
+        forecast_date, lab_site_name
+      ) |>
+      distinct()
+
+    # Get and save quantiles
+    ww_quantiles <- ww_draws |>
+      trajectories_to_quantiles(
+        quantiles = quantiles_to_save,
+        timepoint_cols = "date",
+        value_col = "pred_value",
+        quantile_value_name = "predicted",
+        quantile_level_name = "quantile_level",
+        id_cols = c("site", "lab")
+      ) |>
+      left_join(ww_data_obs,
+        by = c("date", "site", "lab")
+      ) |>
+      left_join(ww_metadata,
+        by = c("site", "lab")
+      )
+    write_csv(
+      ww_quantiles,
+      file.path(
+        data_fp,
+        "ww_quantiles.csv"
+      )
+    )
   }
 
   draws_w_data <- get_model_draws_w_data(
@@ -100,18 +153,18 @@ fit_wwinference_wrapper <- function(
     model = "wwinference",
     forecast_date = this_forecast_date,
     location = loc,
+    hosp_data_real_time = hosp_data_real_time,
     eval_data = hosp_data_eval
   )
-  data_fp <- file.path(full_fp, "data")
-  if (!file.exists(file.path(data_fp))) {
-    dir_create(data_fp, recursive = TRUE, showWarnings = FALSE)
+
+  if (isTRUE(save_draws)) {
+    arrow::write_parquet(
+      draws_w_data,
+      file.path(data_fp, glue::glue(
+        "hosp_draws_ww_{include_ww}_rt_{hosp_data_real_time}.parquet"
+      ))
+    )
   }
-  write_csv(
-    draws_w_data,
-    file.path(data_fp, glue::glue(
-      "hosp_draws_ww_{include_ww}.csv"
-    ))
-  )
   # Make a plot here with calibration and evaluation data and save it.
   get_plot_draws_w_calib_data(
     draws_w_data,
@@ -124,7 +177,8 @@ fit_wwinference_wrapper <- function(
     offset = 1,
     quantiles = TRUE,
     probs = quantiles_to_save
-  )
+  ) |>
+    mutate(flag_missing_ww = flag_missing_ww)
 
   write_csv(
     hosp_quantiles,