feat: expose "method" arg of quantile_reg

cmu-delphi · Sep 4, 2024 · 985f6a0 · 985f6a0
1 parent 991d70b
commit 985f6a0
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/R/make_quantile_reg.R b/R/make_quantile_reg.R
@@ -22,7 +22,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'")
@@ -35,7 +35,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(
@@ -54,9 +54,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")
 
@@ -68,6 +65,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",
@@ -78,7 +83,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
       )
@@ -101,7 +105,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
@@ -117,7 +120,6 @@ make_quantile_reg <- function() {
     return(data.frame(.pred = out))
   }
 
-
   parsnip::set_pred(
     model = "quantile_reg",
     eng = "rq",