initial data for quantreg

Author: ‘topepo’

R/linear_reg_data.R

@@ -1,6 +1,7 @@
 set_new_model("linear_reg")
 
 set_model_mode("linear_reg", "regression")
+set_model_mode("linear_reg", "quantile regression")
 
 # ------------------------------------------------------------------------------
 
@@ -582,3 +583,74 @@ set_pred(
   )
 )
 
+# ------------------------------------------------------------------------------
+
+set_model_engine("linear_reg", "quantile regression", "quantreg")
+set_dependency("linear_reg", "quantreg", "quantreg")
+
+set_fit(
+  model = "linear_reg",
+  eng = "quantreg",
+  mode = "quantile regression",
+  value = list(
+    interface = "formula",
+    protect = c("formula", "data", "weights"),
+    func = c(pkg = "quantreg", fun = "rq"),
+    defaults = list()
+  )
+)
+
+set_encoding(
+  model = "linear_reg",
+  eng = "quantreg",
+  mode = "quantile regression",
+  options = list(
+    predictor_indicators = "traditional",
+    compute_intercept = TRUE,
+    remove_intercept = TRUE,
+    allow_sparse_x = FALSE
+  )
+)
+
+set_pred(
+  model = "linear_reg",
+  eng = "quantreg",
+  mode = "quantile regression",
+  type = "numeric",
+  value = list(
+    pre = NULL,
+    post = NULL,
+    func = c(fun = "predict"),
+    args =
+      list(
+        object = expr(object$fit),
+        newdata = expr(new_data),
+        type = "response",
+        rankdeficient = "simple"
+      )
+  )
+)
+
+set_pred(
+  model = "linear_reg",
+  eng = "quantreg",
+  mode = "quantile regression",
+  type = "conf_int",
+  value = list(
+    pre = NULL,
+    post = function(results, object) {
+      tibble::as_tibble(results) %>%
+        dplyr::select(-fit) %>%
+        setNames(c(".pred_lower", ".pred_upper"))
+    },
+    func = c(fun = "predict"),
+    args =
+      list(
+        object = expr(object$fit),
+        newdata = expr(new_data),
+        interval = "confidence",
+        level = expr(level)
+      )
+  )
+)
+