Adds support for flexsurvspline engine for survival_reg model spec (#831)

mattwarkentin · hfrick · topepo · web-flow · commit bdc28548fd46 · 2022-11-03T10:18:29.000-04:00
* Adds support for flexsurvspline engine for survival_reg model spec

* Add PR number to NEWS

* leave deprecated `surv_reg()` as is

* make `k` tunable

the arg name of `flexsurvspline()` is `k`, not `num_knots`

the method does not need conditional registration because it's new and was never registered in tune

* update `inst/models.tsv`

so that `uses_extension()` in the engine docs works

* update engine docs and knit

* document()

* update news, add contributor

* import also the generic

* remove engine arg from template

* notes about case weights

* render docs

Co-authored-by: Hannah Frick &lt;hannah@rstudio.com&gt;
Co-authored-by: topepo &lt;mxkuhn@gmail.com&gt;
diff --git a/NAMESPACE b/NAMESPACE
@@ -89,6 +89,7 @@ S3method(translate,survival_reg)
 S3method(translate,svm_linear)
 S3method(translate,svm_poly)
 S3method(translate,svm_rbf)
+S3method(tunable,survival_reg)
 S3method(type_sum,model_fit)
 S3method(type_sum,model_spec)
 S3method(update,C5_rules)
@@ -316,6 +317,7 @@ importFrom(generics,fit_xy)
 importFrom(generics,glance)
 importFrom(generics,required_pkgs)
 importFrom(generics,tidy)
+importFrom(generics,tunable)
 importFrom(generics,varying_args)
 importFrom(ggplot2,autoplot)
 importFrom(glue,glue_collapse)
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,7 @@
 # parsnip (development version)
 
+* Adds documentation and tuning infrastructure for the new `flexsurvspline` engine for the `survival_reg()` model specification from the `censored` package (@mattwarkentin, #831).
+
 * The matrix interface for fitting `fit_xy()` now works for the `"censored regression"` mode (#829).
 
 * The `num_leaves` argument of `boost_tree()`s `lightgbm` engine (via the bonsai package) is now tunable.
diff --git a/R/parsnip-package.R b/R/parsnip-package.R
@@ -10,7 +10,7 @@
 ## usethis namespace: start
 #' @importFrom dplyr arrange bind_cols bind_rows collect full_join group_by
 #' @importFrom dplyr mutate pull rename select starts_with summarise tally
-#' @importFrom generics varying_args
+#' @importFrom generics tunable varying_args
 #' @importFrom glue glue_collapse
 #' @importFrom pillar type_sum
 #' @importFrom purrr as_vector imap imap_lgl map map_chr map_dbl map_df map_dfr
diff --git a/R/survival_reg_flexsurvspline.R b/R/survival_reg_flexsurvspline.R
@@ -0,0 +1,11 @@
+#' Flexible parametric survival regression
+#'
+#' [flexsurv::flexsurvspline()] fits a flexible parametric survival model.
+#'
+#' @includeRmd man/rmd/survival_reg_flexsurvspline.md details
+#'
+#' @name details_survival_reg_flexsurvspline
+#' @keywords internal
+NULL
+
+# See inst/README-DOCS.md for a description of how these files are processed
diff --git a/R/tunable.R b/R/tunable.R
@@ -203,6 +203,17 @@ brulee_multinomial_engine_args <-
   brulee_mlp_engine_args %>%
   dplyr::filter(name %in% c("momentum", "batch_size", "stop_iter", "class_weights"))
 
+flexsurvspline_engine_args <-
+  tibble::tibble(
+    name = c("k"),
+    call_info = list(
+      list(pkg = "dials", fun = "num_knots")
+    ),
+    source = "model_spec",
+    component = "survival_reg",
+    component_id = "engine"
+  )
+
 # ------------------------------------------------------------------------------
 
 # Lazily registered in .onLoad()
@@ -324,5 +335,14 @@ tunable_mlp <- function(x, ...) {
   res
 }
 
+#' @export
+tunable.survival_reg <- function(x, ...) {
+  res <- NextMethod()
+  if (x$engine == "flexsurvspline") {
+    res <- add_engine_parameters(res, flexsurvspline_engine_args)
+  }
+  res
+}
+
 # nocov end
 
diff --git a/inst/models.tsv b/inst/models.tsv
@@ -123,6 +123,7 @@
 "surv_reg"	"regression"	"flexsurv"	NA
 "surv_reg"	"regression"	"survival"	NA
 "survival_reg"	"censored regression"	"flexsurv"	"censored"
+"survival_reg"	"censored regression"	"flexsurvspline"	"censored"
 "survival_reg"	"censored regression"	"survival"	"censored"
 "svm_linear"	"classification"	"kernlab"	NA
 "svm_linear"	"classification"	"LiblineaR"	NA
diff --git a/man/details_auto_ml_h2o.Rd b/man/details_auto_ml_h2o.Rd
diff --git a/man/details_survival_reg_flexsurv.Rd b/man/details_survival_reg_flexsurv.Rd
diff --git a/man/details_survival_reg_flexsurvspline.Rd b/man/details_survival_reg_flexsurvspline.Rd
diff --git a/man/rmd/survival_reg_flexsurv.Rmd b/man/rmd/survival_reg_flexsurv.Rmd
@@ -45,6 +45,11 @@ For this engine, stratification cannot be specified via [`strata()`], please see
 ```{r child = "template-survival-mean.Rmd"}
 ```
 
+## Case weights
+
+```{r child = "template-uses-case-weights.Rmd"}
+```
+
 ## Saving fitted model objects
 
 ```{r child = "template-butcher.Rmd"}
diff --git a/man/rmd/survival_reg_flexsurv.md b/man/rmd/survival_reg_flexsurv.md
@@ -48,6 +48,13 @@ For this engine, stratification cannot be specified via [`strata()`], please see
 
 Predictions of type `"time"` are predictions of the mean survival time.
 
+## Case weights
+
+
+This model can utilize case weights during model fitting. To use them, see the documentation in [case_weights] and the examples on `tidymodels.org`. 
+
+The `fit()` and `fit_xy()` arguments have arguments called `case_weights` that expect vectors of case weights. 
+
 ## Saving fitted model objects
 
 
diff --git a/man/rmd/survival_reg_flexsurvspline.Rmd b/man/rmd/survival_reg_flexsurvspline.Rmd
@@ -0,0 +1,48 @@
+```{r, child = "aaa.Rmd", include = FALSE}
+```
+
+`r descr_models("survival_reg", "flexsurvspline")`
+
+## Tuning Parameters
+
+This model has one engine-specific tuning parameter:
+
+ * `k`: Number of knots in the spline. The default is `k = 0`.
+
+## Translation from parsnip to the original package
+
+`r uses_extension("survival_reg", "flexsurvspline", "censored regression")`
+
+```{r flexsurvspline-creg}
+library(censored)
+
+survival_reg() %>% 
+  set_engine("flexsurvspline") %>% 
+  set_mode("censored regression") %>% 
+  translate()
+```
+
+## Other details
+
+The main interface for this model uses the formula method since the model specification typically involved the use of [survival::Surv()]. 
+
+For this engine, stratification cannot be specified via [`strata()`], please see [flexsurv::flexsurvspline()] for alternative specifications.
+
+```{r child = "template-survival-mean.Rmd"}
+```
+
+## Case weights
+
+```{r child = "template-uses-case-weights.Rmd"}
+```
+
+
+## Saving fitted model objects
+
+```{r child = "template-butcher.Rmd"}
+```
+
+
+## References
+
+- Jackson, C. 2016. `flexsurv`: A Platform for Parametric Survival  Modeling in R. _Journal of Statistical Software_, 70(8), 1 - 33.
diff --git a/man/rmd/survival_reg_flexsurvspline.md b/man/rmd/survival_reg_flexsurvspline.md
@@ -0,0 +1,62 @@
+
+
+
+For this engine, there is a single mode: censored regression
+
+## Tuning Parameters
+
+This model has one engine-specific tuning parameter:
+
+ * `k`: Number of knots in the spline. The default is `k = 0`.
+
+## Translation from parsnip to the original package
+
+The **censored** extension package is required to fit this model.
+
+
+```r
+library(censored)
+
+survival_reg() %>% 
+  set_engine("flexsurvspline") %>% 
+  set_mode("censored regression") %>% 
+  translate()
+```
+
+```
+## Parametric Survival Regression Model Specification (censored regression)
+## 
+## Computational engine: flexsurvspline 
+## 
+## Model fit template:
+## flexsurv::flexsurvspline(formula = missing_arg(), data = missing_arg(), 
+##     weights = missing_arg())
+```
+
+## Other details
+
+The main interface for this model uses the formula method since the model specification typically involved the use of [survival::Surv()]. 
+
+For this engine, stratification cannot be specified via [`strata()`], please see [flexsurv::flexsurvspline()] for alternative specifications.
+
+
+
+Predictions of type `"time"` are predictions of the mean survival time.
+
+## Case weights
+
+
+This model can utilize case weights during model fitting. To use them, see the documentation in [case_weights] and the examples on `tidymodels.org`. 
+
+The `fit()` and `fit_xy()` arguments have arguments called `case_weights` that expect vectors of case weights. 
+
+
+## Saving fitted model objects
+
+
+This model object contains data that are not required to make predictions. When saving the model for the purpose of prediction, the size of the saved object might be substantially reduced by using functions from the [butcher](https://butcher.tidymodels.org) package.
+
+
+## References
+
+- Jackson, C. 2016. `flexsurv`: A Platform for Parametric Survival  Modeling in R. _Journal of Statistical Software_, 70(8), 1 - 33.
