polished/updated parametric survival models.

Author: topepo

R/surv_reg.R

@@ -25,16 +25,39 @@
 #'  `strata` function cannot be used. To achieve the same effect,
 #'  the extra parameter roles can be used (as described above).
 #'
-#' The model can be created using the `fit()` function using the
-#'  following _engines_:
-#' \itemize{
-#' \item \pkg{R}:  `"flexsurv"`
-#' }
 #' @inheritParams boost_tree
 #' @param mode A single character string for the type of model.
 #'  The only possible value for this model is "regression".
 #' @param dist A character string for the outcome distribution. "weibull" is
 #'  the default.
+#' @details
+#' For `surv_reg`, the mode will always be "regression".
+#'
+#' The model can be created using the `fit()` function using the
+#'  following _engines_:
+#' \itemize{
+#' \item \pkg{R}:  `"flexsurv"`, `"survreg"`
+#' }
+#'
+#' @section Engine Details:
+#'
+#' Engines may have pre-set default arguments when executing the
+#'  model fit call. These can be changed by using the `...`
+#'  argument to pass in the preferred values. For this type of
+#'  model, the template of the fit calls are:
+#'
+#' \pkg{flexsurv}
+#'
+#' \Sexpr[results=rd]{parsnip:::show_fit(parsnip:::surv_reg(), "flexsurv")}
+#'
+#' \pkg{survreg}
+#'
+#' \Sexpr[results=rd]{parsnip:::show_fit(parsnip:::surv_reg(), "survreg")}
+#'
+#' Note that `model = TRUE` is needed to produce quantile
+#'  predictions when there is a stratification variable and can be
+#'  overridden in other cases.
+#'
 #' @seealso [varying()], [fit()], [survival::Surv()]
 #' @references Jackson, C. (2016). `flexsurv`: A Platform for Parametric Survival
 #'  Modeling in R. _Journal of Statistical Software_, 70(8), 1 - 33.
@@ -160,3 +183,51 @@ check_args.surv_reg <- function(object) {
 
   invisible(object)
 }
+
+# ------------------------------------------------------------------------------
+
+#' @importFrom stats setNames
+#' @importFrom dplyr mutate
+survreg_quant <- function(results, object) {
+  pctl <- object$spec$method$quantile$args$p
+  n <- nrow(results)
+  p <- ncol(results)
+  results <-
+    results %>%
+    as_tibble() %>%
+    setNames(names0(p)) %>%
+    mutate(.row = 1:n) %>%
+    gather(.label, .pred, -.row) %>%
+    arrange(.row, .label) %>%
+    mutate(.quantile = rep(pctl, n)) %>%
+    dplyr::select(-.label)
+  .row <- results[[".row"]]
+  results <-
+    results %>%
+    dplyr::select(-.row)
+  results <- split(results, .row)
+  names(results) <- NULL
+  tibble(.pred = results)
+}
+
+# ------------------------------------------------------------------------------
+
+#' @importFrom dplyr bind_rows
+flexsurv_mean <- function(results, object) {
+  results <- unclass(results)
+  results <- bind_rows(results)
+  results$est
+}
+
+#' @importFrom stats setNames
+flexsurv_quant <- function(results, object) {
+  results <- map(results, as_tibble)
+  names(results) <- NULL
+  results <- map(results, setNames, c(".quantile", ".pred", ".pred_lower", ".pred_upper"))
+}
+
+# ------------------------------------------------------------------------------
+
+#' @importFrom utils globalVariables
+utils::globalVariables(".label")
+

R/surv_reg_data.R

@@ -1,6 +1,8 @@
 
 surv_reg_arg_key <- data.frame(
-  flexsurv    =  c("dist"),
+  flexsurv =  c("dist"),
+  survreg  =  c("dist"),
+  stan     =  c("family"),
   stringsAsFactors = FALSE,
   row.names =  c("dist")
 )
@@ -9,6 +11,8 @@ surv_reg_modes <- "regression"
 
 surv_reg_engines <- data.frame(
   flexsurv = TRUE,
+  survreg  = TRUE,
+  stan     = TRUE,
   stringsAsFactors    = TRUE,
   row.names =  c("regression")
 )
@@ -23,5 +27,96 @@ surv_reg_flexsurv_data <-
       protect = c("formula", "data", "weights"),
       func = c(pkg = "flexsurv", fun = "flexsurvreg"),
       defaults = list()
+    ),
+    pred = list(
+      pre = NULL,
+      post = flexsurv_mean,
+      func = c(fun = "summary"),
+      args =
+        list(
+          object = expr(object$fit),
+          newdata = expr(new_data),
+          type = "mean"
+        )
+    ),
+    quantile = list(
+      pre = NULL,
+      post = flexsurv_quant,
+      func = c(fun = "summary"),
+      args =
+        list(
+          object = expr(object$fit),
+          newdata = expr(new_data),
+          type = "quantile",
+          quantiles = expr(quantile)
+        )
     )
   )
+
+# ------------------------------------------------------------------------------
+
+surv_reg_survreg_data <-
+  list(
+    libs = c("survival"),
+    fit = list(
+      interface = "formula",
+      protect = c("formula", "data", "weights"),
+      func = c(pkg = "survival", fun = "survreg"),
+      defaults = list(model = TRUE)
+    ),
+    pred = list(
+      pre = NULL,
+      post = NULL,
+      func = c(fun = "predict"),
+      args =
+        list(
+          object = expr(object$fit),
+          newdata = expr(new_data),
+          type = "response"
+        )
+    ),
+    quantile = list(
+      pre = NULL,
+      post = survreg_quant,
+      func = c(fun = "predict"),
+      args =
+        list(
+          object = expr(object$fit),
+          newdata = expr(new_data),
+          type = "quantile",
+          p = expr(quantile)
+        )
+    )
+  )
+
+# ------------------------------------------------------------------------------
+
+surv_reg_stan_data <-
+  list(
+    libs = c("brms"),
+    fit = list(
+      interface = "formula",
+      protect = c("formula", "data", "weights"),
+      func = c(pkg = "brms", fun = "brm"),
+      defaults = list(
+        family = expr(brms::weibull()),
+        seed = expr(sample.int(10^5, 1))
+      )
+    ),
+    pred = list(
+      pre = NULL,
+      post = function(results, object) {
+        tibble::as_tibble(results) %>%
+          dplyr::select(Estimate) %>%
+          setNames(".pred")
+      },
+      func = c(fun = "predict"),
+      args =
+        list(
+          object = expr(object$fit),
+          newdata = expr(new_data),
+          type = "response"
+        )
+    )
+  )
+