feat: add forecast method #293

Author: dshemetov

NAMESPACE

@@ -152,6 +152,7 @@ export(flatline)
 export(flatline_args_list)
 export(flatline_forecaster)
 export(flusight_hub_formatter)
+export(forecast)
 export(frosting)
 export(get_test_data)
 export(grab_names)

R/epi_workflow.R

@@ -197,7 +197,11 @@ update_model.epi_workflow <- function(x, spec, ..., formula = NULL) {
 #'
 #' @export
 fit.epi_workflow <- function(object, data, ..., control = workflows::control_workflow()) {
-  object$fit$meta <- list(max_time_value = max(data$time_value), as_of = attributes(data)$metadata$as_of)
+  object$fit$meta <- list(
+    max_time_value = max(data$time_value),
+    as_of = attributes(data)$metadata$as_of,
+    train_data = data
+  )
 
   NextMethod()
 }

R/forecast.R

@@ -0,0 +1,22 @@
+#' Produce forecast from an epi workflow
+#' @param epi_workflow An epi workflow
+#'
+#' @return A forecast tibble.
+#'
+#' @export
+forecast <- function(epi_workflow, fill_locf = FALSE, n_recent = NULL, forecast_date = NULL) {
+  # Find data inside the epi_workflow
+  if (!epi_workflow$trained) {
+    cli_abort("The epi_workflow is not trained.")
+  }
+
+  test_data <- get_test_data(
+    hardhat::extract_preprocessor(epi_workflow),
+    epi_workflow$fit$meta$train_data,
+    fill_locf = fill_locf,
+    n_recent = n_recent %||% Inf,
+    forecast_date = forecast_date %||% max(epi_workflow$fit$meta$train_data$time_value)
+  )
+
+  predict(epi_workflow, new_data = test_data)
+}

man/forecast.Rd

@@ -0,0 +1,19 @@
+test_that("forecast method works", {
+  jhu <- case_death_rate_subset %>%
+    filter(time_value > "2021-11-01", geo_value %in% c("ak", "ca", "ny"))
+  r <- epi_recipe(jhu) %>%
+    step_epi_lag(death_rate, lag = c(0, 7, 14)) %>%
+    step_epi_ahead(death_rate, ahead = 7) %>%
+    step_epi_naomit()
+  wf <- epi_workflow(r, parsnip::linear_reg()) %>% fit(jhu)
+
+  latest <- get_test_data(
+    hardhat::extract_preprocessor(wf),
+    jhu
+  )
+
+  expect_equal(
+    forecast(wf),
+    predict(wf, new_data = latest)
+  )
+})