refactor: only pass quantiles

R/LearnerRegrRanger.R

@@ -126,27 +126,24 @@ LearnerRegrRanger = R6Class("LearnerRegrRanger",
       pv = self$param_set$get_values(tags = "predict")
       newdata = ordered_features(task, self)
 
-      type = switch(
-        self$predict_type,
-        response = "response",
-        se = "se",
-        quantile = "quantiles"
-      )
-
-      prediction = invoke(predict, self$model, data = newdata, type = type, quantiles = private$.quantile, .args = pv)
+      prediction = invoke(predict, self$model,
+        data = newdata,
+        type = if (self$predict_type == "quantile") "quantiles" else pv$type,
+        quantiles = private$.quantile,
+        .args = pv)
 
-      if (type == "quantiles") {
-        response = prediction$predictions[, which(private$.quantile == private$.quantile_response)]
+      if (self$predict_type == "quantile") {
         quantile = prediction$predictions
         attr(quantile, "probs") = private$.quantile
-        list(response = response, quantile = quantile)
-      } else {
-        list(response = prediction$predictions, se = prediction$se)
+        attr(quantile, "response") = private$.quantile_response
+        return(list(quantile = quantile))
       }
+
+      list(response = prediction$predictions, se = prediction$se)
     },
 
     .hotstart = function(task) {
-      model = self$model
+      model = self$models
       model$num.trees = self$param_set$values$num.trees
       model
     }

tests/testthat/test_regr_ranger.R

@@ -75,8 +75,12 @@ test_that("quantile prediction", {
 
   learner$train(task)
   pred = learner$predict(task)
-  expect_matrix(pred$quantile)
+
+  expect_matrix(pred$quantile, ncol = 3L)
+  expect_true(!any(apply(pred$quantile, 1L, is.unsorted)))
+  expect_equal(pred$response, pred$quantile[, 2L])
+
   tab = as.data.table(pred)
-  expect_names(names(tab), must.include = c("q0.1", "q0.5", "q0.9", "response"))
+  expect_names(names(tab), identical.to = c("row_ids", "truth", "q0.1", "q0.5", "q0.9", "response"))
   expect_equal(tab$response, tab$q0.5)
 })