refactor step_adjust_ahead to be early step

Author: dsweber2

NAMESPACE

@@ -229,9 +229,9 @@ importFrom(dplyr,all_of)
 importFrom(dplyr,group_by)
 importFrom(dplyr,join_by)
 importFrom(dplyr,left_join)
+importFrom(dplyr,mutate)
 importFrom(dplyr,n)
 importFrom(dplyr,pull)
-importFrom(dplyr,rowwise)
 importFrom(dplyr,select)
 importFrom(dplyr,summarise)
 importFrom(dplyr,tibble)
@@ -242,6 +242,7 @@ importFrom(generics,fit)
 importFrom(generics,forecast)
 importFrom(ggplot2,autoplot)
 importFrom(glue,glue)
+importFrom(hardhat,extract_recipe)
 importFrom(hardhat,refresh_blueprint)
 importFrom(hardhat,run_mold)
 importFrom(magrittr,"%>%")
@@ -250,6 +251,7 @@ importFrom(recipes,bake)
 importFrom(recipes,detect_step)
 importFrom(recipes,prep)
 importFrom(rlang,"!!!")
+importFrom(recipes,recipes_eval_select)
 importFrom(rlang,"!!")
 importFrom(rlang,"%@%")
 importFrom(rlang,"%||%")
@@ -276,6 +278,7 @@ importFrom(stats,quantile)
 importFrom(stats,residuals)
 importFrom(tibble,tibble)
 importFrom(tidyr,drop_na)
+importFrom(tidyr,expand_grid)
 importFrom(tidyr,unnest)
 importFrom(vctrs,as_list_of)
 importFrom(vctrs,field)

R/arx_classifier.R

@@ -193,8 +193,6 @@ arx_class_epi_workflow <- function(
     }
   }
   o2 <- rlang::sym(paste0("ahead_", args_list$ahead, "_", o))
-  r <- r %>%
-    step_epi_ahead(!!o, ahead = args_list$ahead, role = "pre-outcome")
   method_adjust_latency <- args_list$adjust_latency
   if (!is.null(method_adjust_latency)) {
     # only extend_ahead is supported atm
@@ -203,6 +201,8 @@ arx_class_epi_workflow <- function(
       method = method_adjust_latency
     )
   }
+  r <- r %>%
+    step_epi_ahead(!!o, ahead = args_list$ahead, role = "pre-outcome")
   r <- r %>%
     step_mutate(
       outcome_class = cut(!!o2, breaks = args_list$breaks),
@@ -220,10 +220,6 @@ arx_class_epi_workflow <- function(
         drop_na = FALSE
       )
   }
-  if (!is.null) {
-    forecast_date <- args_list$forecast_date %||% max(epi_data$time_value)
-  }
-  target_date <- args_list$target_date %||% (forecast_date + args_list$ahead)
 
   # --- postprocessor
   f <- frosting() %>% layer_predict() # %>% layer_naomit()
@@ -308,9 +304,6 @@ arx_class_args_list <- function(
 
   arg_is_scalar(ahead, n_training, horizon, log_scale)
   arg_is_scalar(forecast_date, target_date, adjust_latency, allow_null = TRUE)
-  if (adjust_latency == "adjust_lags") {
-    cli::cli_abort("step_adjust_latency is not yet implemented for lagged differences and growth rates")
-  }
   arg_is_date(forecast_date, target_date, allow_null = TRUE)
   arg_is_nonneg_int(ahead, lags, horizon)
   arg_is_numeric(breaks)

R/arx_forecaster.R

@@ -142,12 +142,7 @@ arx_fcast_epi_workflow <- function(
 
   # --- preprocessor
   r <- epi_recipe(epi_data)
-  for (l in seq_along(lags)) {
-    p <- predictors[l]
-    r <- step_epi_lag(r, !!p, lag = lags[[l]])
-  }
-  r <- r %>%
-    step_epi_ahead(!!outcome, ahead = args_list$ahead)
+  # adjust latency if the user asks
   method_adjust_latency <- args_list$adjust_latency
   if (!is.null(method_adjust_latency)) {
     if (method_adjust_latency == "extend_ahead") {
@@ -162,6 +157,12 @@ arx_fcast_epi_workflow <- function(
       )
     }
   }
+  for (l in seq_along(lags)) {
+    p <- predictors[l]
+    r <- step_epi_lag(r, !!p, lag = lags[[l]])
+  }
+  r <- r %>%
+    step_epi_ahead(!!outcome, ahead = args_list$ahead)
   r <- r %>%
     step_epi_naomit() %>%
     step_training_window(n_recent = args_list$n_training)

R/canned-epipred.R

@@ -63,6 +63,7 @@ print.alist <- function(x, ...) {
 }
 
 #' @export
+#' @importFrom hardhat extract_recipe
 print.canned_epipred <- function(x, name, ...) {
   d <- cli::cli_div(theme = list(rule = list("line-type" = "double")))
   cli::cli_rule("A basic forecaster of type {name}")
@@ -110,19 +111,33 @@ print.canned_epipred <- function(x, name, ...) {
     "At forecast date{?s}: {.val {fds}},",
     "For target date{?s}: {.val {tds}},"
   ))
-  if (detect_step(x$epi_workflow$pre$actions$recipe$recipe, "adjust_latency")) {
-    latency_step <- keep(x$epi_workflow$pre$mold$blueprint$recipe$steps,
-                         \(x) inherits(x, "step_adjust_latency"))[[1]]
-    latency_per_base_col <- latency_step$shift_cols %>%
-      group_by(latency) %>%
-      reframe(variable = parent_name) %>%
-      distinct() %>%
-      mutate(latency = abs(latency)) %>%
-      relocate(variable, latency)
-    if (nrow(latency_per_base_col)>1) {
-      intro_text <- "Latency adjusted per column: "
+  fit_recipe <- extract_recipe(x$epi_workflow)
+  if (detect_step(fit_recipe, "adjust_latency")) {
+    is_adj_latency <- map_lgl(fit_recipe$steps, \(x) inherits(x, "step_adjust_latency"))
+    latency_step <- fit_recipe$steps[is_adj_latency][[1]]
+    # all steps after adjust_latency
+    later_steps <- fit_recipe$steps[-(1:which(is_adj_latency))]
+    if (latency_step$method == "extend_ahead") {
+      step_names <- "step_epi_ahead"
+      type_str <- "Aheads"
+    } else if (latency_step$method == "extend_lags") {
+      step_names <- "step_epi_lag"
+      type_str <- "Lags"
     } else {
-      intro_text <- "Latency adjusted for "
+      step_names <- ""
+      type_str <- "columns locf"
+    }
+    later_steps[[1]]$columns
+    valid_columns <- later_steps %>%
+      keep(\(x) inherits(x, step_names)) %>%
+      purrr::map("columns") %>%
+      reduce(c)
+    latency_per_base_col <- latency_step$latency_table %>%
+      filter(col_name %in% valid_columns) %>% mutate(latency = abs(latency))
+    if (latency_step$method != "locf" && nrow(latency_per_base_col) > 1) {
+      intro_text <- glue::glue("{type_str} adjusted per column: ")
+    } else if (latency_step$method != "locf") {
+      intro_text <- glue::glue("{type_str} adjusted for ")
     }
     latency_info <- paste0(intro_text, paste(apply(latency_per_base_col, 1, paste0, collapse = "="), collapse = ", "))
     cli::cli_ul(latency_info)

R/epi_shift.R

@@ -21,7 +21,13 @@ epi_shift_single <- function(x, col, shift_val, newname, key_cols) {
 #' the future back to today
 #' @keywords internal
 get_sign <- function(object) {
-  if (object$prefix == "lag_") {
+  if (!is.null(object$prefix)) {
+    if (object$prefix == "lag_") {
+      return(1)
+    } else {
+      return(-1)
+    }
+  } else if (object$method == "extend_lags") {
     return(1)
   } else {
     return(-1)
@@ -31,13 +37,28 @@ get_sign <- function(object) {
 #' backend for both `bake.step_epi_ahead` and `bake.step_epi_lag`, performs the
 #' checks missing in `epi_shift_single`
 #' @keywords internal
+#' @importFrom tidyr expand_grid
+#' @importFrom dplyr mutate left_join join_by
 add_shifted_columns <- function(new_data, object, amount) {
   sign_shift <- get_sign(object)
-  grid <- tidyr::expand_grid(col = object$columns, amount = amount) %>%
+  latency_table <- attributes(new_data)$metadata$latency_table
+  shift_sign_lat <- attributes(new_data)$metadata$shift_sign
+  if (!is.null(latency_table) &&
+    shift_sign_lat == sign_shift) {
+    #TODO this doesn't work on lags of transforms
+    rel_latency <- latency_table %>% filter(col_name %in% object$columns)
+  } else {
+    rel_latency <- tibble(col_name = object$columns, latency = 0L)
+  }
+  grid <- expand_grid(col = object$columns, amount = sign_shift *amount) %>%
+    left_join(rel_latency, by = join_by(col == col_name), ) %>%
+    tidyr::replace_na(list(latency = 0)) %>%
     dplyr::mutate(
-      newname = glue::glue("{object$prefix}{amount}_{col}"),
-      shift_val = sign_shift * amount,
-      amount = NULL
+      shift_val = amount + latency) %>%
+    mutate(
+      newname = glue::glue("{object$prefix}{abs(shift_val)}_{col}"), # name is always positive
+      amount = NULL,
+      latency = NULL
     )
 
   ## ensure no name clashes
@@ -59,8 +80,12 @@ add_shifted_columns <- function(new_data, object, amount) {
     dplyr::full_join,
     by = ok
   )
-  dplyr::full_join(new_data, shifted, by = ok) %>%
+  processed <- new_data %>%
+    dplyr::full_join(shifted, by = ok) %>%
     dplyr::group_by(dplyr::across(dplyr::all_of(ok[-1]))) %>%
     dplyr::arrange(time_value) %>%
-    dplyr::ungroup()
+    dplyr::ungroup() %>%
+    as_epi_df()
+  attributes(processed)$metadata <- attributes(new_data)$metadata
+  return(processed)
 }

R/layer_add_forecast_date.R

@@ -92,7 +92,7 @@ slather.layer_add_forecast_date <- function(object, components, workflow, new_da
   rlang::check_dots_empty()
   forecast_date <- object$forecast_date %||%
     get_forecast_date_in_layer(
-      extract_preprocessor(workflow),
+      extract_recipe(workflow),
       workflow$fit$meta$max_time_value,
       new_data
     )