|
| 1 | +#' Generate forecast data file containing all forecast hub |
| 2 | +#' model submissions. |
| 3 | +#' |
| 4 | +#' This function fetches all forecast submissions from a |
| 5 | +#' forecast hub based on the reference date. The forecast |
| 6 | +#' data is then pivoted to create a wide format with |
| 7 | +#' quantile levels as columns. |
| 8 | +#' |
| 9 | +#' The resulting file contains the following columns: |
| 10 | +#' - `location_name`: full state name (including "US" for |
| 11 | +#' the US state) |
| 12 | +#' - `abbreviation`: state abbreviation |
| 13 | +#' - `horizon`: forecast horizon |
| 14 | +#' - `forecast_date`: date the forecast was generated |
| 15 | +#' - `target_end_date`: target date for the forecast |
| 16 | +#' - `model`: model name |
| 17 | +#' - `quantile_*`: forecast values for various quantiles |
| 18 | +#' (e.g., 0.025, 0.5, 0.975) |
| 19 | +#' - `forecast_teams`: name of the team that generated the |
| 20 | +#' model |
| 21 | +#' - `forecast_fullnames`: full model name |
| 22 | +#' |
| 23 | +#' @param reference_date character, the reference date for |
| 24 | +#' the forecast in YYYY-MM-DD format (ISO-8601). |
| 25 | +#' @param base_hub_path character, path to the forecast |
| 26 | +#' hub directory. |
| 27 | +#' @param hub_reports_path character, path to forecast hub |
| 28 | +#' reports directory. |
| 29 | +#' @param disease character, disease name ("covid" or |
| 30 | +#' "rsv"). Used to derive target name and file prefix. |
| 31 | +#' @param horizons_to_include integer vector, horizons to |
| 32 | +#' include in the output. Default: c(0, 1, 2). |
| 33 | +#' @param excluded_locations character vector of location |
| 34 | +#' codes to exclude from the output. Default: character(0). |
| 35 | +#' @param output_format character, output file format. One |
| 36 | +#' of "csv", "tsv", or "parquet". Default: "csv". |
| 37 | +#' @param targets character vector, target name(s) to filter |
| 38 | +#' forecasts. If NULL (default), does not filter by target. |
| 39 | +#' Can be a single target like "wk inc covid hosp" or |
| 40 | +#' multiple targets like c("wk inc covid hosp", "wk inc |
| 41 | +#' covid prop ed visits"). |
| 42 | +#' |
| 43 | +#' @export |
| 44 | +get_forecast_data <- function( |
| 45 | + reference_date, |
| 46 | + base_hub_path, |
| 47 | + hub_reports_path, |
| 48 | + disease, |
| 49 | + horizons_to_include = c(0, 1, 2), |
| 50 | + excluded_locations = character(0), |
| 51 | + output_format = "csv", |
| 52 | + targets = NULL |
| 53 | +) { |
| 54 | + checkmate::assert_choice(disease, choices = c("covid", "rsv")) |
| 55 | + checkmate::assert_subset(horizons_to_include, choices = c(-1, 0, 1, 2, 3)) |
| 56 | + checkmate::assert_character(excluded_locations) |
| 57 | + checkmate::assert_choice(output_format, choices = c("csv", "tsv", "parquet")) |
| 58 | + checkmate::assert_character(targets, null.ok = TRUE) |
| 59 | + |
| 60 | + reference_date <- lubridate::as_date(reference_date) |
| 61 | + |
| 62 | + model_metadata <- hubData::load_model_metadata( |
| 63 | + base_hub_path, |
| 64 | + model_ids = NULL |
| 65 | + ) |
| 66 | + |
| 67 | + hub_content <- hubData::connect_hub(base_hub_path) |
| 68 | + |
| 69 | + current_forecasts <- hub_content |> |
| 70 | + dplyr::filter( |
| 71 | + .data$reference_date == !!reference_date, |
| 72 | + !(.data$location %in% !!excluded_locations), |
| 73 | + .data$horizon %in% !!horizons_to_include |
| 74 | + ) |> |
| 75 | + hubData::collect_hub() |> |
| 76 | + dplyr::filter(forecasttools::nullable_comparison( |
| 77 | + .data$target, |
| 78 | + "%in%", |
| 79 | + !!targets |
| 80 | + )) |
| 81 | + |
| 82 | + all_forecasts_data <- forecasttools::pivot_hubverse_quantiles_wider( |
| 83 | + hubverse_table = current_forecasts, |
| 84 | + pivot_quantiles = c( |
| 85 | + "quantile_0.025" = 0.025, |
| 86 | + "quantile_0.25" = 0.25, |
| 87 | + "quantile_0.5" = 0.5, |
| 88 | + "quantile_0.75" = 0.75, |
| 89 | + "quantile_0.975" = 0.975 |
| 90 | + ) |
| 91 | + ) |> |
| 92 | + dplyr::mutate( |
| 93 | + location_name = forecasttools::us_location_recode( |
| 94 | + .data$location, |
| 95 | + "hub", |
| 96 | + "name" |
| 97 | + ), |
| 98 | + abbreviation = forecasttools::us_location_recode( |
| 99 | + .data$location, |
| 100 | + "hub", |
| 101 | + "abbr" |
| 102 | + ), |
| 103 | + dplyr::across( |
| 104 | + tidyselect::starts_with("quantile_"), |
| 105 | + round, |
| 106 | + .names = "{.col}_rounded" |
| 107 | + ), |
| 108 | + forecast_due_date = as.Date(!!reference_date) - 3, |
| 109 | + location_sort_order = ifelse(.data$location_name == "United States", 0, 1) |
| 110 | + ) |> |
| 111 | + dplyr::mutate( |
| 112 | + location_name = dplyr::case_match( |
| 113 | + .data$location_name, |
| 114 | + "United States" ~ "US", |
| 115 | + .default = .data$location_name |
| 116 | + ) |
| 117 | + ) |> |
| 118 | + dplyr::arrange(.data$location_sort_order, .data$location_name) |> |
| 119 | + dplyr::left_join( |
| 120 | + dplyr::distinct( |
| 121 | + model_metadata, |
| 122 | + .data$model_id, |
| 123 | + .keep_all = TRUE |
| 124 | + ), |
| 125 | + by = "model_id" |
| 126 | + ) |> |
| 127 | + dplyr::select( |
| 128 | + "location_name", |
| 129 | + "abbreviation", |
| 130 | + "horizon", |
| 131 | + forecast_date = "reference_date", |
| 132 | + "target_end_date", |
| 133 | + model = "model_id", |
| 134 | + "quantile_0.025", |
| 135 | + "quantile_0.25", |
| 136 | + "quantile_0.5", |
| 137 | + "quantile_0.75", |
| 138 | + "quantile_0.975", |
| 139 | + "quantile_0.025_rounded", |
| 140 | + "quantile_0.25_rounded", |
| 141 | + "quantile_0.5_rounded", |
| 142 | + "quantile_0.75_rounded", |
| 143 | + "quantile_0.975_rounded", |
| 144 | + forecast_team = "team_name", |
| 145 | + "forecast_due_date", |
| 146 | + model_full_name = "model_name" |
| 147 | + ) |
| 148 | + |
| 149 | + output_folder_path <- fs::path( |
| 150 | + hub_reports_path, |
| 151 | + "weekly-summaries", |
| 152 | + reference_date |
| 153 | + ) |
| 154 | + output_filename <- glue::glue("{reference_date}_{disease}_forecasts_data") |
| 155 | + output_filepath <- fs::path( |
| 156 | + output_folder_path, |
| 157 | + output_filename, |
| 158 | + ext = output_format |
| 159 | + ) |
| 160 | + |
| 161 | + fs::dir_create(output_folder_path) |
| 162 | + |
| 163 | + if (!fs::file_exists(output_filepath)) { |
| 164 | + forecasttools::write_tabular(all_forecasts_data, output_filepath) |
| 165 | + cli::cli_inform("File saved as: {output_filepath}") |
| 166 | + } else { |
| 167 | + cli::cli_abort("File already exists: {output_filepath}") |
| 168 | + } |
| 169 | +} |
0 commit comments