better appraoch to creating formulas from recipes based on the model type

Author: topepo

R/fitter.R

@@ -226,18 +226,17 @@ formula_to_matrix <- function(object, formula, data, control, ...) {
 
 #' @importFrom recipes prep juice all_predictors all_outcomes
 
-recipe_data <- function(recipe, data, control, output = "matrix", combine = FALSE) {
+# add case weights as extra object returned (out$weights)
+recipe_data <- function(recipe, data, object, control, output = "matrix", combine = FALSE) {
   recipe <-
     prep(recipe, training = data, retain = TRUE, verbose = control$verbosity > 1)
 
   if (combine) {
-    out <- list(data = juice(recipe, all_predictors(), all_outcomes(), composition = output))
-    data_info <- summary(recipe)
-    y_names <- data_info$variable[data_info$role == "outcome"]
-    if (length(y_names) > 1)
-      out$form <- paste0("cbind(", paste0(y_names, collapse = ","), ")~.")
-    else
-      out$form <- paste0(y_names, "~.")
+    out <- list(
+      data = juice(recipe, composition = output),
+      form = formula(object, (recipe))
+    )
+
   } else {
     out <-
       list(
@@ -257,16 +256,16 @@ recipe_data <- function(recipe, data, control, output = "matrix", combine = FALS
 
 recipe_to_formula <-
   function(object, recipe, data, control, ...) {
-    info <- recipe_data(recipe, data, control, output = "tibble", combine = TRUE)
+    info <- recipe_data(recipe, data, object, control, output = "tibble", combine = TRUE)
     formula_to_formula(object, info$form, info$data, control, ...)
   }
 
 recipe_to_data.frame <- function(object, recipe, data, control, ...) {
-  info <- recipe_data(recipe, data, control, output = "tibble", combine = FALSE)
+  info <- recipe_data(recipe, data, object, control, output = "tibble", combine = FALSE)
   xy_to_xy(object, info$x, info$y, control, ...)
 }
 
 recipe_to_matrix <- function(object, recipe, data, control, ...) {
-  info <- recipe_data(recipe, data, control, output = "matrix", combine = FALSE)
+  info <- recipe_data(recipe, data, object, control, output = "matrix", combine = FALSE)
   xy_to_xy(object, info$x, info$y, control, ...)
 }

R/form_recipe.R

@@ -0,0 +1,67 @@
+# Class for creating an appropriate formula for a model based on the roles
+# contained in the recipe. 
+
+# maybe use `formula` generic with `formula.model_spec` acting as the generic?
+
+formula.model_spec <- function(x, recipe, ...) {
+  rec_vars <- summary(recipe)
+  y_names <- rec_vars$variable[rec_vars$role == "outcome"]
+  if (length(y_names) > 1)
+    form_text <- paste0("cbind(", paste0(y_names, collapse = ","), ")~.")
+  else
+    form_text <- paste0(y_names, "~.")
+  
+  form <- try(as.formula(form_text), silent = TRUE)
+  if(inherits(form, "try-error"))
+    stop("Could not parse the model formula: ", form_text, call. = FALSE)
+  form
+}
+
+
+formula.surv_reg <- function(x, recipe, ...) {
+  rec_vars <- summary(recipe)
+  y_names <- rec_vars$variable[rec_vars$role == "outcome"]
+  if (length(y_names) > 2 | length(y_names) < 1)
+    stop("There should be 1-2 variables in the `outcome` role.", call. = FALSE)
+  cens_names <- rec_vars$variable[rec_vars$role == "censoring var"]
+  if (length(cens_names) > 1)
+    stop("There should be 0-1 variables in the `censoring` role.", call. = FALSE)
+  x_names <- rec_vars$variable[rec_vars$role == "predictor"]
+  
+  # construct basic formula
+  form_text <- paste0("Surv(", paste0(y_names, collapse = ", "))
+  if (length(cens_names) == 1)
+    form_text <- paste0(form_text, ", ", cens_names, ") ~")
+  else
+    form_text <- paste0(form_text, ") ~")
+  
+  if (length(x_names) == 0)
+    form_text <- paste0(form_text, "1")
+  else
+    form_text <- paste0(form_text, paste0(x_names, collapse = "+"))
+  
+  # engine-speciifc options (e.g. spark needing censor var in text)
+  if (!is.null(x$engine) && x$engine == "flexsurv") {
+    extra_ind <- which(rec_vars$role %in% flexsurv_params)
+    if (length(extra_ind) > 0) {
+      extra_terms <- paste0(
+        rec_vars$role[extra_ind], "(",
+        rec_vars$variable[extra_ind], ")"
+        )
+      form_text <- paste0(form_text, "+", paste0(extra_terms, collapse = "+"))
+    }
+    if (any(rec_vars$role == "strata"))
+      warning(
+        "`flexsurv` does not use the `strata` function; instead use ",
+        "the parameter roles for differential values (e.g. `sigma`).",
+        call. = FALSE
+      )
+  }
+  
+  form <- try(as.formula(form_text), silent = TRUE)
+  if(inherits(form, "try-error"))
+    stop("Could not parse the model formula: ", form_text, call. = FALSE)
+  form
+}
+
+flexsurv_params <- c("sigma", "shape", "sdlog", "Q", "k", "P", "S1", "s2")

R/misc.R

@@ -101,20 +101,3 @@ resolve_args <- function(args, ...) {
   }
   args
 }
-
-
-# A function to convert a R formula to spark format
-surv_to_spark_formula <- function(f) {
-  if (!inherits(f, "formula"))
-    stop("A formula is required.")
-  if (length(f[[2]]) != 3)
-    stop("spark requires the `Surv` object to have a ",
-         "censoring indicator.")
-  if (!all.equal(f[[2]][[1]], as.name("Surv")))
-    stop("The formula sould contain a `Surv` object.")
-  f2 <- f
-  f2[[2]] <- f[[2]][[2]]
-  list(formula = f2, censor = deparse(f[[2]][[3]]))
-}
-
-