predict layer debugged

Author: dajmcdon

R/epi_workflow.R

@@ -130,8 +130,7 @@ predict.epi_workflow <- function(object, new_data, ...) {
   components$keys <- grab_forged_keys(components$forged,
                                       components$mold, new_data)
   components <- apply_frosting(object, components, the_fit, ...)
-  out <- dplyr::bind_cols(components$keys, components$preds)
-  out
+  components$predictions
 }
 
 grab_forged_keys <- function(forged, mold, new_data) {

R/frosting.R

@@ -79,23 +79,26 @@ apply_frosting.default <- function(workflow, components, ...) {
 #' @export
 apply_frosting.epi_workflow <- function(workflow, components, the_fit, ...) {
   if (!has_postprocessor(workflow)) {
-    components$preds <- predict(the_fit, components$forged$predictors, ...)
+    components$predictions <- predict(the_fit, components$forged$predictors, ...)
+    components$predictions <- dplyr::bind_cols(components$keys, components$predictions)
     return(components)
   }
   if (!has_postprocessor_frosting(workflow)) {
     rlang::warn(c("Only postprocessors of class frosting are allowed.",
                   "Returning unpostprocessed predictions."))
-    components$preds <- predict(the_fit, components$forged$predictors, ...)
+    components$predictions <- predict(the_fit, components$forged$predictors, ...)
+    components$predictions <- dplyr::bind_cols(components$keys, components$predictions)
     return(components)
   }
-  layers <- workflow$post$actions$frosting$frosting
-  for (l in seq_along(layers$layers)) {
-    layer <- layers$layers[[l]]
-    components <- slather(layer, components = components, the_fit)
+  layers <- workflow$post$actions$frosting$frosting$layers
+  for (l in seq_along(layers)) {
+    la <- layers[[l]]
+    components <- slather(la, components = components, the_fit)
   }
   # last for the moment, anticipating that layer[1] will do the prediction...
-  if (is_null(components$preds)) {
-    components$preds <- predict(the_fit, components$forged$predictors, ...)
+  if (is_null(components$predictions)) {
+    components$predictions <- predict(the_fit, components$forged$predictors, ...)
+    components$predictions <- dplyr::bind_cols(components$keys, components$predictions)
   }
   return(components)
 }

R/layer_predict.R

@@ -16,6 +16,7 @@ layer_predict_new <- function(id) {
 #' @export
 slather.layer_predict <- function(object, components, the_fit) {
   components$predictions <- predict(the_fit, components$forged$predictors)
+  components$predictions <- dplyr::bind_cols(components$keys, components$predictions)
   components
 }
 

tests/testthat/test-frosting.R

@@ -34,3 +34,5 @@ test_that("prediction works without any postprocessor", {
   expect_equal(tail(p$time_value, 1), as.Date("2021-12-31"))
   expect_equal(unique(p$geo_value), c("ak", "ca", "ny"))
 })
+
+

tests/testthat/test-layer_predict.R

@@ -0,0 +1,21 @@
+test_that("predict layer works alone", {
+  jhu <- case_death_rate_subset %>%
+    dplyr::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_naomit(all_predictors()) %>%
+    step_naomit(all_outcomes(), skip = TRUE)
+  wf <- epi_workflow(r, linear_reg()) %>% fit(jhu)
+  latest <- jhu %>% filter(time_value >= max(time_value) - 14)
+
+  f <- frosting() %>% layer_predict()
+  wf <- wf %>% add_frosting(f)
+
+  expect_silent(p <- predict(wf, latest))
+  expect_equal(ncol(p), 3L)
+  expect_s3_class(p, "epi_df")
+  expect_equal(nrow(p), 108L)
+  expect_named(p, c("geo_value", "time_value", ".pred"))
+
+})