create a point prediction layer.

* currently bugs in the workflow hierarchy. For some reason the layers are deeply nested.

Author: dajmcdon

NAMESPACE

@@ -18,6 +18,8 @@ S3method(prep,step_epi_ahead)
 S3method(prep,step_epi_lag)
 S3method(print,step_epi_ahead)
 S3method(print,step_epi_lag)
+S3method(slather,layer_naomit)
+S3method(slather,layer_predict)
 export("%>%")
 export(add_epi_recipe)
 export(apply_frosting)

R/frosting.R

@@ -88,8 +88,8 @@ apply_frosting.epi_workflow <- function(workflow, components, the_fit, ...) {
     components$preds <- predict(the_fit, components$forged$predictors, ...)
     return(components)
   }
-  layers <- workflow$post$actions$frosting
-  for (l in seq_along(layers)) {
+  layers <- workflow$post$actions$frosting$frosting
+  for (l in seq_along(layers$layers)) {
     layer <- layers$layers[[l]]
     components <- slather(layer, components = components, the_fit)
   }

R/layer_naomit.R

@@ -12,7 +12,12 @@ layer_naomit_new <- function(terms, id) {
   layer("naomit", terms = terms, id = id)
 }
 
+#' @export
 slather.layer_naomit <- function(object, components, the_fit) {
+  exprs <- rlang::expr(c(!!!object$terms))
+  pos <- tidyselect::eval_select(exprs, components$predictions)
+  col_names <- names(pos)
   components$predictions <- components$predictions %>%
-    filter()
+    dplyr::filter(dplyr::if_any(dplyr::all_of(col_names), ~ !is.na(.x)))
+  components
 }

R/layer_predict.R

@@ -0,0 +1,21 @@
+layer_predict <-
+  function(frosting, id = rand_id("predict_default")) {
+    add_layer(
+      frosting,
+      layer_predict_new(
+        id = id
+      )
+    )
+  }
+
+
+layer_predict_new <- function(id) {
+  layer("predict", id = id)
+}
+
+#' @export
+slather.layer_predict <- function(object, components, the_fit) {
+  components$predictions <- predict(the_fit, components$forged$predictors)
+  components
+}
+