using check_enough_train_data in practice

Author: dsweber2

R/arx_forecaster.R

@@ -171,7 +171,9 @@ arx_fcast_epi_workflow <- function(
     step_epi_ahead(!!outcome, ahead = args_list$ahead)
   r <- r %>%
     step_epi_naomit() %>%
-    step_training_window(n_recent = args_list$n_training)
+    step_training_window(n_recent = args_list$n_training) %>%
+    check_enough_train_data(all_predictors(), skip = FALSE)
+
   if (!is.null(args_list$check_enough_data_n)) {
     r <- r %>% check_enough_train_data(
       all_predictors(),

R/check_enough_train_data.R

@@ -47,7 +47,7 @@ check_enough_train_data <-
            role = NA,
            trained = FALSE,
            columns = NULL,
-           skip = TRUE,
+           skip = FALSE,
            id = rand_id("enough_train_data")) {
     recipes::add_check(
       recipe,
@@ -90,7 +90,7 @@ prep.check_enough_train_data <- function(x, training, info = NULL, ...) {
   }
 
   if (x$drop_na) {
-    training <- tidyr::drop_na(training)
+    training <- tidyr::drop_na(training, any_of(unname(col_names)))
   }
   cols_not_enough_data <- training %>%
     group_by(across(all_of(.env$x$epi_keys))) %>%
@@ -101,7 +101,8 @@ prep.check_enough_train_data <- function(x, training, info = NULL, ...) {
 
   if (length(cols_not_enough_data) > 0) {
     cli_abort(
-      "The following columns don't have enough data to predict: {cols_not_enough_data}."
+      "The following columns don't have enough data to predict: {cols_not_enough_data}.",
+      class = "epipredict__not_enough_train_data"
     )
   }
 
@@ -120,6 +121,23 @@ prep.check_enough_train_data <- function(x, training, info = NULL, ...) {
 
 #' @export
 bake.check_enough_train_data <- function(object, new_data, ...) {
+  col_names <- object$columns
+  if (object$drop_na) {
+    newish_data <- tidyr::drop_na(new_data, any_of(unname(col_names)))
+  }
+  cols_not_enough_data <- newish_data %>%
+    group_by(across(all_of(.env$object$epi_keys))) %>%
+    summarise(across(all_of(.env$col_names), ~ dplyr::n() < .env$object$n), .groups = "drop") %>%
+    summarise(across(all_of(.env$col_names), any), .groups = "drop") %>%
+    unlist() %>%
+    names(.)[.]
+
+  if (length(cols_not_enough_data) > 0) {
+    cli_abort(
+      "The following columns don't have enough data to predict: {cols_not_enough_data}.",
+      class = "epipredict__not_enough_train_data"
+    )
+  }
   new_data
 }
 

R/epi_workflow.R

@@ -270,6 +270,6 @@ forecast.epi_workflow <- function(object, ..., n_recent = NULL, forecast_date =
     hardhat::extract_preprocessor(object),
     object$original_data
   )
-
+  test_data
   predict(object, new_data = test_data)
 }

tests/testthat/test-arx_forecaster.R

@@ -24,3 +24,22 @@ test_that("arx_forecaster errors if forecast date, target date, and ahead are in
     class = "epipredict__arx_args__inconsistent_target_ahead_forecaste_date"
   )
 })
+
+test_that("warns if there's not enough data to predict", {
+  edf <- tibble(
+    geo_value = "ct",
+    time_value = seq(as.Date("2020-10-01"), as.Date("2023-05-31"), by = "day"),
+  ) %>%
+    mutate(value = seq_len(nrow(.)) + rnorm(nrow(.))) %>%
+    # Oct to May (flu season, ish) only:
+    filter(!between(as.POSIXlt(time_value)$mon + 1L, 6L, 9L)) %>%
+    # and actually, pretend we're around mid-October 2022:
+    filter(time_value <= as.Date("2022-10-12")) %>%
+    as_epi_df(as_of = as.Date("2022-10-12"))
+  edf %>% filter(time_value > "2022-08-01")
+
+  expect_error(
+    edf %>% arx_forecaster("value"),
+    class = "epipredict__not_enough_train_data"
+  )
+})