Quantile predictions output constructor

NAMESPACE

@@ -36,6 +36,7 @@ S3method(extract_spec_parsnip,model_fit)
 S3method(fit,model_spec)
 S3method(fit_xy,gen_additive_mod)
 S3method(fit_xy,model_spec)
+S3method(format,vctrs_quantiles)
 S3method(glance,model_fit)
 S3method(has_multi_predict,default)
 S3method(has_multi_predict,model_fit)
@@ -54,6 +55,7 @@ S3method(multi_predict_args,default)
 S3method(multi_predict_args,model_fit)
 S3method(multi_predict_args,workflow)
 S3method(nullmodel,default)
+S3method(obj_print_footer,vctrs_quantiles)
 S3method(predict,"_elnet")
 S3method(predict,"_glmnetfit")
 S3method(predict,"_lognet")
@@ -172,6 +174,8 @@ S3method(update,svm_rbf)
 S3method(varying_args,model_spec)
 S3method(varying_args,recipe)
 S3method(varying_args,step)
+S3method(vec_ptype_abbr,vctrs_quantiles)
+S3method(vec_ptype_full,vctrs_quantiles)
 export("%>%")
 export(.censoring_weights_graf)
 export(.check_glmnet_penalty_fit)
@@ -280,6 +284,7 @@ export(new_model_spec)
 export(null_model)
 export(null_value)
 export(nullmodel)
+export(obj_print_footer)
 export(parsnip_addin)
 export(pls)
 export(poisson_reg)
@@ -350,6 +355,7 @@ export(update_model_info_file)
 export(update_spec)
 export(varying)
 export(varying_args)
+export(vec_quantiles)
 export(xgb_predict)
 export(xgb_train)
 import(rlang)
@@ -396,6 +402,8 @@ importFrom(purrr,map)
 importFrom(purrr,map_chr)
 importFrom(purrr,map_dbl)
 importFrom(purrr,map_lgl)
+importFrom(rlang,"!!!")
+importFrom(rlang,is_double)
 importFrom(stats,.checkMFClasses)
 importFrom(stats,.getXlevels)
 importFrom(stats,as.formula)
@@ -426,5 +434,6 @@ importFrom(utils,globalVariables)
 importFrom(utils,head)
 importFrom(utils,methods)
 importFrom(utils,stack)
+importFrom(vctrs,obj_print_footer)
 importFrom(vctrs,vec_size)
 importFrom(vctrs,vec_unique)

R/aaa_quantiles.R

@@ -21,23 +21,100 @@ check_quantile_level <- function(x, object, call) {
   x
 }
 
-# Assumes the columns have the same order as quantile_level
-restructure_rq_pred <- function(x, object) {
-  num_quantiles <- NCOL(x)
-  if ( num_quantiles == 1L ){
-    x <- matrix(x, ncol = 1)
+
+# -------------------------------------------------------------------------
+# A column vector of quantiles with an attribute
+
+#' @export
+vec_ptype_abbr.vctrs_quantiles <- function(x, ...) "qntls"
+
+#' @export
+vec_ptype_full.vctrs_quantiles <- function(x, ...) "quantiles"
+
+#' @importFrom rlang is_double !!!
+new_vec_quantiles <- function(values = list(), quantile_levels = double()) {
+  quantile_levels <- vctrs::vec_cast(quantile_levels, double())
+  vctrs::new_vctr(
+    values, quantile_levels = quantile_levels, class = "vctrs_quantiles"
+  )
+}
+
+
+#' Create a vector containing sets of quantiles
+#'
+#' @param values A matrix of values. Each column should correspond to one of
+#'   the quantile levels.
+#' @param quantile_levels A vector of probabilities corresponding to `values`.
+#'
+#' @export
+#' @return A vector of values associated with the quantile levels.
+#'
+#' @examples
+#' v <- vec_quantiles(matrix(rnorm(20), 5), c(.2, .4, .6, .8))
+#'
+#' # Access the underlying information
+#' attr(v, "quantile_levels")
+#' vctrs::vec_data(v)
+vec_quantiles <- function(values, quantile_levels = double()) {
+  check_vec_quantiles_inputs(values, quantile_levels)
+  quantile_levels <- vctrs::vec_cast(quantile_levels, double())
+  num_lvls <- length(quantile_levels)
+
+  if (ncol(values) != num_lvls) {
+    cli::cli_abort(
+      "The number of columns in {.arg values} must be equal to the length of
+        {.arg quantile_levels}."
+    )
   }
-  n <- nrow(x)
+  values <- lapply(vctrs::vec_chop(values), drop)
+  new_vec_quantiles(values, quantile_levels)
+}
 
+check_vec_quantiles_inputs <- function(values, levels) {
+  if (!is.matrix(values)) {
+    cls <- class(values)[1]
+    cli::cli_abort("{.arg values} must be a {.cls matrix} not a {.cls {cls}}.")
+  }
+  purrr::walk(levels,
+    ~ check_number_decimal(.x, min = 0, max = 1, arg = "quantile_levels")
+  )
+  if (is.unsorted(levels)) {
+    cli::cli_abort("{.arg quantile_levels} must be sorted in increasing order.")
+  }
+  invisible(NULL)
+}
+
+#' @export
+format.vctrs_quantiles <- function(x, ...) {
+  quantile_levels <- attr(x, "levels")
+  if (length(quantile_levels) == 1L) {
+    x <- unlist(x)
+    out <- round(x, 3L)
+    out[is.na(x)] <- NA
+  } else {
+    rng <- sapply(x, range)
+    out <- paste0("[", round(rng[1, ], 3L), ", ", round(rng[2, ], 3L), "]")
+    out[is.na(rng[1, ]) | is.na(rng[2, ])] <- NA
+  }
+  out
+}
+
+#' @importFrom vctrs obj_print_footer
+#' @export
+vctrs::obj_print_footer
+
+#' @export
+obj_print_footer.vctrs_quantiles <- function(x, ...) {
+  lvls <- attr(x, "quantile_levels")
+  cat("# Quantile levels: ", format(lvls, digits = 3), "\n", sep = " ")
+}
+
+restructure_rq_pred <- function(x, object) {
+  if (!is.matrix(x)) x <- as.matrix(x)
+  rownames(x) <- NULL
+  n_pred_quantiles <- ncol(x)
+  # TODO check p = length(quantile_level)
   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, num_quantiles))
-  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
+  tibble::tibble(.pred_quantile = vec_quantiles(x, quantile_level))
 }
 

tests/testthat/test-linear_reg_quantreg.R

@@ -24,19 +24,21 @@ test_that('linear quantile regression via quantreg - single quantile', {
   expect_true(nrow(one_quant_pred) == nrow(sac_test))
   expect_named(one_quant_pred, ".pred_quantile")
   expect_true(is.list(one_quant_pred[[1]]))
-  expect_s3_class(one_quant_pred$.pred_quantile[[1]], c("tbl_df", "tbl", "data.frame"))
-  expect_named(one_quant_pred$.pred_quantile[[1]], c(".pred_quantile", ".quantile_level"))
-  expect_true(nrow(one_quant_pred$.pred_quantile[[1]]) == 1L)
+  expect_s3_class(one_quant_pred$.pred_quantile[1], c("vctrs_quantiles", "vctrs_vctr", "list"))
+  expect_identical(class(one_quant_pred$.pred_quantile[[1]]), "numeric")
+  expect_true(length(one_quant_pred$.pred_quantile[[1]]) == 1L)
+  expect_identical(attr(one_quant_pred$.pred_quantile, "quantile_levels"), .5)
 
   ###
 
   one_quant_one_row <- predict(one_quant, new_data = sac_test[1,])
   expect_true(nrow(one_quant_one_row) == 1L)
   expect_named(one_quant_one_row, ".pred_quantile")
   expect_true(is.list(one_quant_one_row[[1]]))
-  expect_s3_class(one_quant_one_row$.pred_quantile[[1]], c("tbl_df", "tbl", "data.frame"))
-  expect_named(one_quant_one_row$.pred_quantile[[1]], c(".pred_quantile", ".quantile_level"))
-  expect_true(nrow(one_quant_one_row$.pred_quantile[[1]]) == 1L)
+  expect_s3_class(one_quant_one_row$.pred_quantile[1], c("vctrs_quantiles", "vctrs_vctr", "list"))
+  expect_identical(class(one_quant_one_row$.pred_quantile[[1]]), "numeric")
+  expect_true(length(one_quant_one_row$.pred_quantile[[1]]) == 1L)
+  expect_identical(attr(one_quant_pred$.pred_quantile, "quantile_levels"), .5)
 })
 
 test_that('linear quantile regression via quantreg - multiple quantiles', {
@@ -65,19 +67,21 @@ test_that('linear quantile regression via quantreg - multiple quantiles', {
   expect_true(nrow(ten_quant_pred) == nrow(sac_test))
   expect_named(ten_quant_pred, ".pred_quantile")
   expect_true(is.list(ten_quant_pred[[1]]))
-  expect_s3_class(ten_quant_pred$.pred_quantile[[1]], c("tbl_df", "tbl", "data.frame"))
-  expect_named(ten_quant_pred$.pred_quantile[[1]], c(".pred_quantile", ".quantile_level"))
-  expect_true(nrow(ten_quant_pred$.pred_quantile[[1]]) == 10L)
+  expect_s3_class(ten_quant_pred$.pred_quantile[1], c("vctrs_quantiles", "vctrs_vctr", "list"))
+  expect_identical(class(ten_quant_pred$.pred_quantile[[1]]), "numeric")
+  expect_true(length(ten_quant_pred$.pred_quantile[[1]]) == 10L)
+  expect_identical(attr(ten_quant_pred$.pred_quantile, "quantile_levels"), (0:9)/9)
 
   ###
 
   ten_quant_one_row <- predict(ten_quant, new_data = sac_test[1,])
   expect_true(nrow(ten_quant_one_row) == 1L)
   expect_named(ten_quant_one_row, ".pred_quantile")
   expect_true(is.list(ten_quant_one_row[[1]]))
-  expect_s3_class(ten_quant_one_row$.pred_quantile[[1]], c("tbl_df", "tbl", "data.frame"))
-  expect_named(ten_quant_one_row$.pred_quantile[[1]], c(".pred_quantile", ".quantile_level"))
-  expect_true(nrow(ten_quant_one_row$.pred_quantile[[1]]) == 10L)
+  expect_s3_class(ten_quant_one_row$.pred_quantile[1], c("vctrs_quantiles", "vctrs_vctr", "list"))
+  expect_identical(class(ten_quant_one_row$.pred_quantile[[1]]), "numeric")
+  expect_true(length(ten_quant_one_row$.pred_quantile[[1]]) == 10L)
+  expect_identical(attr(ten_quant_one_row$.pred_quantile, "quantile_levels"), (0:9)/9)
 })
 
 

tests/testthat/test-vec_quantiles.R

@@ -0,0 +1,25 @@
+test_that("vec_quantiles error types", {
+  expect_error(vec_quantiles(1:10, 1:4 / 5), "matrix")
+  expect_error(
+    vec_quantiles(matrix(1:20, 5), -1:4 / 5),
+    "`quantile_levels` must be a number between 0 and 1"
+  )
+  expect_error(
+    vec_quantiles(matrix(1:20, 5), 1:5 / 6),
+    "The number of columns in `values` must be equal to"
+  )
+  expect_error(
+    vec_quantiles(matrix(1:20, 5), 4:1 / 5),
+    "must be sorted in increasing order"
+  )
+})
+
+test_that("vec_quantiles outputs", {
+  v <- vec_quantiles(matrix(1:20, 5), 1:4 / 5)
+  expect_s3_class(v, "vctrs_quantiles")
+  expect_identical(attr(v, "quantile_levels"), 1:4 / 5)
+  expect_identical(
+    vctrs::vec_data(v),
+    lapply(vctrs::vec_chop(matrix(1:20, 5)), drop)
+  )
+})