Skip to content

Commit 985f6a0

Browse files
committed
feat: expose "method" arg of quantile_reg
1 parent 991d70b commit 985f6a0

File tree

1 file changed

+10
-8
lines changed

1 file changed

+10
-8
lines changed

R/make_quantile_reg.R

+10-8
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
#' rq_spec <- quantile_reg(quantile_levels = c(.2, .8)) %>% set_engine("rq")
2323
#' ff <- rq_spec %>% fit(y ~ ., data = tib)
2424
#' predict(ff, new_data = tib)
25-
quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0.5) {
25+
quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0.5, method = "br") {
2626
# Check for correct mode
2727
if (mode != "regression") {
2828
cli_abort("`mode` must be 'regression'")
@@ -35,7 +35,7 @@ quantile_reg <- function(mode = "regression", engine = "rq", quantile_levels = 0
3535
cli::cli_warn("Sorting `quantile_levels` to increasing order.")
3636
quantile_levels <- sort(quantile_levels)
3737
}
38-
args <- list(quantile_levels = rlang::enquo(quantile_levels))
38+
args <- list(quantile_levels = rlang::enquo(quantile_levels), method = rlang::enquo(method))
3939

4040
# Save some empty slots for future parts of the specification
4141
parsnip::new_model_spec(
@@ -54,9 +54,6 @@ make_quantile_reg <- function() {
5454
parsnip::set_new_model("quantile_reg")
5555
}
5656
parsnip::set_model_mode("quantile_reg", "regression")
57-
58-
59-
6057
parsnip::set_model_engine("quantile_reg", "regression", eng = "rq")
6158
parsnip::set_dependency("quantile_reg", eng = "rq", pkg = "quantreg")
6259

@@ -68,6 +65,14 @@ make_quantile_reg <- function() {
6865
func = list(pkg = "quantreg", fun = "rq"),
6966
has_submodel = FALSE
7067
)
68+
parsnip::set_model_arg(
69+
model = "quantile_reg",
70+
eng = "rq",
71+
parsnip = "method",
72+
original = "method",
73+
func = list(pkg = "quantreg", fun = "rq"),
74+
has_submodel = FALSE
75+
)
7176

7277
parsnip::set_fit(
7378
model = "quantile_reg",
@@ -78,7 +83,6 @@ make_quantile_reg <- function() {
7883
protect = c("formula", "data", "weights"),
7984
func = c(pkg = "quantreg", fun = "rq"),
8085
defaults = list(
81-
method = "br",
8286
na.action = rlang::expr(stats::na.omit),
8387
model = FALSE
8488
)
@@ -101,7 +105,6 @@ make_quantile_reg <- function() {
101105
object <- parsnip::extract_fit_engine(object)
102106
type <- class(object)[1]
103107

104-
105108
# can't make a method because object is second
106109
out <- switch(type,
107110
rq = dist_quantiles(unname(as.list(x)), object$quantile_levels), # one quantile
@@ -117,7 +120,6 @@ make_quantile_reg <- function() {
117120
return(data.frame(.pred = out))
118121
}
119122

120-
121123
parsnip::set_pred(
122124
model = "quantile_reg",
123125
eng = "rq",

0 commit comments

Comments
 (0)