enable quantile prediction

Author: ‘topepo’

R/quantiles.R renamed to R/aaa_quantiles.R

@@ -9,6 +9,7 @@ check_quantile_level <- function(x, object, call) {
       {.arg quantile_level} must be specified for quantile regression models.")
     }
   }
+  x <- sort(unique(x))
   # TODO we need better vectorization here, otherwise we get things like:
   # "Error during wrapup: i In index: 2." in the traceback.
   res <-
@@ -17,7 +18,24 @@ check_quantile_level <- function(x, object, call) {
                                       arg = "quantile_level", call = call,
                                       allow_infinite = FALSE)
   )
-  return(invisible(TRUE))
+  x
 }
 
+# Assumes the columns have the same order as quantile_level
+restructure_rq_pred <- function(x, object) {
+  n <- nrow(x)
+  p <- ncol(x)
+  # TODO check p = length(quantile_level)
+  # check p = 1 case
+  quantile_level <- object$spec$quantile_level
+  res <-
+    tibble::tibble(
+    .pred_quantile = as.vector(x),
+    .quantile_level = rep(quantile_level, each = n),
+    .row = rep(1:n, p))
+  res <- vctrs::vec_split(x = res[,1:2], by = res[, ".row"])
+  res <- vctrs::vec_cbind(res$key, tibble::new_tibble(list(.pred_quantile = res$val)))
+  res$.row <- NULL
+  res
+}
 

R/arguments.R

@@ -109,7 +109,8 @@ set_mode.model_spec <- function(object, mode, quantile_level = NULL, ...) {
 
   object$mode <- mode
   object$user_specified_mode <- TRUE
-  check_quantile_level(quantile_level, object, call = caller_env(0))
+  quantile_level <-
+    check_quantile_level(quantile_level, object, call = caller_env(0))
   object$quantile_level <- quantile_level
   object
 }

R/fit.R

@@ -169,6 +169,10 @@ fit.model_spec <-
     eval_env$formula <- formula
     eval_env$weights <- wts
 
+    if ( !is.null(object$quantile_level) ) {
+      eval_env$quantile_level <- object$quantile_level
+    }
+
     fit_interface <-
       check_interface(eval_env$formula, eval_env$data, cl, object)
 
@@ -282,6 +286,10 @@ fit_xy.model_spec <-
     eval_env$y_var <- y_var
     eval_env$weights <- weights_to_numeric(case_weights, object)
 
+    if ( !is.null(object$quantile_level) ) {
+      eval_env$quantile_level <- object$quantile_level
+    }
+
     # TODO case weights: pass in eval_env not individual elements
     fit_interface <- check_xy_interface(eval_env$x, eval_env$y, cl, object)
 

R/linear_reg_data.R

@@ -596,7 +596,7 @@ set_fit(
     interface = "formula",
     protect = c("formula", "data", "weights"),
     func = c(pkg = "quantreg", fun = "rq"),
-    defaults = list()
+    defaults = list(tau = expr(quantile_level))
   )
 )
 
@@ -635,22 +635,15 @@ set_pred(
   model = "linear_reg",
   eng = "quantreg",
   mode = "quantile regression",
-  type = "conf_int",
+  type = "quantile",
   value = list(
     pre = NULL,
-    post = function(results, object) {
-      tibble::as_tibble(results) %>%
-        dplyr::select(-fit) %>%
-        setNames(c(".pred_lower", ".pred_upper"))
-    },
+    post = restructure_rq_pred,
     func = c(fun = "predict"),
     args =
       list(
         object = expr(object$fit),
-        newdata = expr(new_data),
-        interval = "confidence",
-        level = expr(level)
+        newdata = expr(new_data)
       )
   )
 )
-

R/predict_quantile.R

@@ -6,40 +6,36 @@
 #' @method predict_quantile model_fit
 #' @export predict_quantile.model_fit
 #' @export
-predict_quantile.model_fit <- function(object,
-                                       new_data,
-                                       quantile = (1:9)/10,
-                                       interval = "none",
-                                       level = 0.95,
-                                       ...) {
+predict_quantile.model_fit <- function(object, new_data, ...) {
 
-    check_spec_pred_type(object, "quantile")
+  check_spec_pred_type(object, "quantile")
 
-    if (inherits(object$fit, "try-error")) {
-      rlang::warn("Model fit failed; cannot make predictions.")
-      return(NULL)
-    }
-
-    new_data <- prepare_data(object, new_data)
+  if (inherits(object$fit, "try-error")) {
+    cli::cli_warn("Model fit failed; cannot make predictions.")
+    return(NULL)
+  }
 
-    # preprocess data
-    if (!is.null(object$spec$method$pred$quantile$pre))
-      new_data <- object$spec$method$pred$quantile$pre(new_data, object)
+  new_data <- prepare_data(object, new_data)
 
-    # Pass some extra arguments to be used in post-processor
-    object$spec$method$pred$quantile$args$p <- quantile
-    pred_call <- make_pred_call(object$spec$method$pred$quantile)
+  # preprocess data
+  if (!is.null(object$spec$method$pred$quantile$pre)) {
+    new_data <- object$spec$method$pred$quantile$pre(new_data, object)
+  }
 
-    res <- eval_tidy(pred_call)
+  # Pass some extra arguments to be used in post-processor
+  object$spec$method$pred$quantile$args$quantile_level <- object$quantile_level
+  pred_call <- make_pred_call(object$spec$method$pred$quantile)
 
-    # post-process the predictions
-    if(!is.null(object$spec$method$pred$quantile$post)) {
-      res <- object$spec$method$pred$quantile$post(res, object)
-    }
+  res <- eval_tidy(pred_call)
 
-    res
+  # post-process the predictions
+  if(!is.null(object$spec$method$pred$quantile$post)) {
+    res <- object$spec$method$pred$quantile$post(res, object)
   }
 
+  res
+}
+
 # @export
 # @keywords internal
 # @rdname other_predict