Merge pull request #374 from cmu-delphi/ds/quantreg-method

feat: expose "method" arg of quantile_reg

Author: dajmcdon

R/make_quantile_reg.R

@@ -13,6 +13,8 @@
 #'   "rq" and "grf" are supported.
 #' @param quantile_levels A scalar or vector of values in (0, 1) to determine which
 #'   quantiles to estimate (default is 0.5).
+#' @param method A fitting method used by [quantreg::rq()]. See the
+#'   documentation for a list of options.
 #'
 #' @export
 #'
@@ -25,7 +27,7 @@
 #' rq_spec <- quantile_reg(quantile_levels = c(.2, .8)) %>% set_engine("rq")
 #' ff <- rq_spec %>% fit(y ~ ., data = tib)
 #' predict(ff, new_data = tib)
-quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0.5) {
+quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0.5, method = "br") {
   # Check for correct mode
   if (mode != "regression") {
     cli_abort("`mode` must be 'regression'")
@@ -38,7 +40,7 @@ quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0
     cli::cli_warn("Sorting `quantile_levels` to increasing order.")
     quantile_levels <- sort(quantile_levels)
   }
-  args <- list(quantile_levels = rlang::enquo(quantile_levels))
+  args <- list(quantile_levels = rlang::enquo(quantile_levels), method = rlang::enquo(method))
 
   # Save some empty slots for future parts of the specification
   parsnip::new_model_spec(
@@ -57,9 +59,6 @@ make_quantile_reg <- function() {
     parsnip::set_new_model("quantile_reg")
   }
   parsnip::set_model_mode("quantile_reg", "regression")
-
-
-
   parsnip::set_model_engine("quantile_reg", "regression", eng = "rq")
   parsnip::set_dependency("quantile_reg", eng = "rq", pkg = "quantreg")
 
@@ -71,6 +70,14 @@ make_quantile_reg <- function() {
     func = list(pkg = "quantreg", fun = "rq"),
     has_submodel = FALSE
   )
+  parsnip::set_model_arg(
+    model = "quantile_reg",
+    eng = "rq",
+    parsnip = "method",
+    original = "method",
+    func = list(pkg = "quantreg", fun = "rq"),
+    has_submodel = FALSE
+  )
 
   parsnip::set_fit(
     model = "quantile_reg",
@@ -81,7 +88,6 @@ make_quantile_reg <- function() {
       protect = c("formula", "data", "weights"),
       func = c(pkg = "quantreg", fun = "rq"),
       defaults = list(
-        method = "br",
         na.action = rlang::expr(stats::na.omit),
         model = FALSE
       )
@@ -104,7 +110,6 @@ make_quantile_reg <- function() {
     object <- parsnip::extract_fit_engine(object)
     type <- class(object)[1]
 
-
     # can't make a method because object is second
     out <- switch(type,
       rq = dist_quantiles(unname(as.list(x)), object$quantile_levels), # one quantile
@@ -120,7 +125,6 @@ make_quantile_reg <- function() {
     return(dplyr::tibble(.pred = out))
   }
 
-
   parsnip::set_pred(
     model = "quantile_reg",
     eng = "rq",