predict method implemented for epi_workflow

* needs tests/validation
* needs improved documentation

Author: dajmcdon

NAMESPACE

@@ -25,6 +25,7 @@ export(epi_workflow)
 export(get_precision)
 export(grab_names)
 export(is_epi_recipe)
+export(is_epi_workflow)
 export(knn_iteraive_ar_args_list)
 export(knn_iteraive_ar_forecaster)
 export(knnarx_args_list)
@@ -37,6 +38,7 @@ import(recipes)
 importFrom(magrittr,"%>%")
 importFrom(rlang,"!!")
 importFrom(rlang,":=")
+importFrom(rlang,is_null)
 importFrom(stats,as.formula)
 importFrom(stats,lm)
 importFrom(stats,model.frame)

R/epi_keys.R

@@ -23,3 +23,4 @@ epi_keys.epi_df <- function(x) {
 epi_keys.recipe <- function(x) {
   x$var_info$variable[x$var_info$role %in% c("time_value", "geo_value", "key")]
 }
+

R/epi_recipe.R

@@ -212,7 +212,7 @@ epi_form2args <- function(formula, data, ...) {
 
 
 
-#' Test for `epi_df` format
+#' Test for `epi_recipe`
 #'
 #' @param x An object.
 #' @return `TRUE` if the object inherits from `epi_recipe`.
@@ -256,6 +256,7 @@ is_epi_recipe <- function(x) {
 #' @examples
 #' library(recipes)
 #' library(magrittr)
+#' library(workflows)
 #'
 #' recipe <- epi_recipe(mpg ~ cyl, mtcars) %>%
 #'   step_log(cyl)

R/epi_workflow.R

@@ -1,53 +1,93 @@
+#' Create an epi_workflow
+#'
+#' This is a container object that unifies preprocessing, fitting, prediction,
+#' and postprocessing for predictive modeling on epidemiological data. It extends
+#' the functionality of a [`workflows::workflow()`] to handle the typical panel
+#' data structures found in this field. This extension is handled completely
+#' internally, and should be invisible to the user. For all intents and purposes,
+#' this operates exactly like a [`workflows::workflow()`]. For more details
+#' and numerous examples, see there.
+#'
+#' @inheritParams workflows::workflow
+#'
+#' @return A new `epi_workflow` object.
+#' @seealso workflows::workflow
+#' @importFrom rlang is_null
 #' @export
 epi_workflow <- function(preprocessor = NULL, spec = NULL) {
-  out <- workflow(preprocessor, spec)
+  out <- workflows::workflow(spec = spec)
   class(out) <- c("epi_workflow", class(out))
-}
 
-predict.epi_workflow <-
-  function(object, new_data, type = NULL, opts = list(), forecast_date, ...) {
-    out <- predict(object, new_data, type = type, opts = opts, ...)
-    if (is_epi_df(new_data)) {
-      ek <- epi_keys(new_data)
+  if (is_epi_recipe(preprocessor)) {
+    return(add_epi_recipe(out, preprocessor))
+  }
 
-    }
+  if (!is_null(preprocessor)) {
+    return(workflows:::add_preprocessor(out, preprocessor))
   }
+  out
+}
 
+#' Test for an `epi_workflow`
+#'
+#' @param x An object.
+#' @return `TRUE` if the object inherits from `epi_workflow`.
+#'
+#' @export
 is_epi_workflow <- function(x) {
   inherits(x, "epi_workflow")
 }
 
-workflow <- function(preprocessor = NULL, spec = NULL) {
-  out <- new_workflow()
-
-  if (!is_null(preprocessor)) {
-    out <- add_preprocessor(out, preprocessor)
-  }
 
-  if (!is_null(spec)) {
-    out <- add_model(out, spec)
+predict.epi_workflow <-
+  function(object, new_data, type = NULL, opts = list(),
+           forecast_date = NULL, ...) {
+    if (!workflows::is_trained_workflow(object)) {
+      rlang::abort(
+        c("Can't predict on an untrained epi_workflow.",
+          i = "Do you need to call `fit()`?"))
+    }
+    the_fit <- workflows::extract_fit_parsnip(object)
+    mold <- workflows::extract_mold(object)
+    forged <- hardhat::forge(new_data, blueprint = mold$blueprint)
+    preds <- predict(the_fit, forged$predictors, type = type, opts = opts, ...)
+    keys <- grab_forged_keys(forged, mold, new_data)
+    out <- dplyr::bind_cols(keys, preds, forecast_date)
+    out
   }
 
-  out
-}
-
-add_preprocessor <- function(x, preprocessor, ..., call = caller_env()) {
-  check_dots_empty()
-
-  if (is_formula(preprocessor)) {
-    return(add_formula(x, preprocessor))
+grab_forged_keys <- function(forged, mold, new_data) {
+  keys <- c("time_value", "geo_value", "key")
+  forged_names <- names(forged$extras$roles)
+  molded_names <- names(mold$extras$roles)
+  extras <- dplyr::bind_cols(forged$extras$roles[forged_names %in% keys])
+  # 1. these are the keys in the test data after prep/bake
+  new_keys <- names(extras)
+  # 2. these are the keys in the training data
+  old_keys <- purrr::map_chr(mold$extras$roles[molded_names %in% keys], names)
+  # 3. these are the keys in the test data as input
+  new_df_keys <- epi_keys(new_data)
+  if (! (setequal(old_keys, new_df_keys) && setequal(new_keys, new_df_keys))) {
+    rlang::warn(c(
+      "Not all epi keys that were present in the training data are available",
+      "in `new_data`. Predictions will have only the available keys.")
+    )
   }
-
-  if (is_recipe(preprocessor)) {
-    return(add_recipe(x, preprocessor))
+  if (epiprocess::is_epi_df(new_data) || keys[1:2] %in% new_keys) {
+    l <- list()
+    if (length(new_keys) > 2) l <- list(other_keys = new_keys[-c(1:2)])
+    extras <- as_epi_df(extras, additional_metadata = l)
   }
+  extras
+}
 
-  if (is_workflow_variables(preprocessor)) {
-    return(add_variables(x, variables = preprocessor))
-  }
+new_epi_workflow <- function(
+    pre = workflows:::new_stage_pre(),
+    fit = workflows:::new_stage_fit(),
+    post = workflows:::new_stage_post(),
+    trained = FALSE) {
 
-  abort(
-    "`preprocessor` must be a formula, recipe, or a set of workflow variables.",
-    call = call
-  )
+  out <- workflows:::new_workflow(
+    pre = pre, fit = fit, post = post, trained = trained)
+  class(out) <- c("epi_workflow", class(out))
 }