Support arrdist

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: mlr3proba
 Title: Probabilistic Supervised Learning for 'mlr3'
-Version: 0.5.2
+Version: 0.5.3
 Authors@R:
     c(person(given = "Raphael",
              family = "Sonabend",
@@ -43,7 +43,7 @@ Depends:
 Imports:
     checkmate,
     data.table,
-    distr6 (>= 1.6.11),
+    distr6 (>= 1.8.0),
     ggplot2,
     mlr3misc (>= 0.7.0),
     mlr3viz,
@@ -62,7 +62,7 @@ Suggests:
     param6 (>= 0.2.4),
     pracma,
     rpart,
-    set6 (>= 0.1.7),
+    set6 (>= 0.2.6),
     simsurv,
     survAUC,
     testthat,

R/MeasureRegrLogloss.R

@@ -43,7 +43,7 @@ MeasureRegrLogloss = R6::R6Class("MeasureRegrLogloss",
       distr = prediction$distr
       truth = prediction$truth
 
-      if (inherits(distr, "Matdist")) {
+      if (inherits(distr, c("Matdist", "Arrdist"))) {
         pdf = diag(distr$pdf(truth))
       } else {
         pdf = as.numeric(distr$pdf(data = matrix(truth, nrow = 1)))

R/PipeOpSurvAvg.R

@@ -101,7 +101,7 @@ delayedAssign(
       distr = map(inputs, "distr")
 
       ok = mlr3misc::map_lgl(distr, function(.x) {
-        checkmate::test_class(.x, "Matdist")
+        checkmate::test_class(.x, "Matdist") | checkmate::test_class(.x, "Arrdist")
       })
 
       if (all(ok)) {

R/PipeOpTaskSurvRegr.R

@@ -228,7 +228,7 @@ PipeOpTaskSurvRegr = R6Class("PipeOpTaskSurvRegr",
       }
 
       est = est$train(task)$predict(task)$distr
-      if (inherits(est, "Matdist")) {
+      if (inherits(est, c("Matdist", "Arrdist"))) {
         weights = diag(est$survival(task$truth()[, 1]))
       } else {
         weights = as.numeric(est$survival(data = matrix(task$truth()[, 1], nrow = 1)))

R/PredictionDataSurv.R

@@ -14,8 +14,10 @@ check_prediction_data.PredictionDataSurv = function(pdata, ...) { # nolint
   assert_numeric(pdata$lp, len = n, any.missing = FALSE, null.ok = TRUE)
   if (inherits(pdata$distr, "VectorDistribution")) {
     assert(nrow(pdata$distr$modelTable) == n)
-  } else if (inherits(pdata$distr, "Matdist")) {
+  } else if (inherits(pdata$distr, c("Matdist", "Arrdist"))) {
     assert(nrow(gprm(pdata$distr, "pdf")) == n)
+  } else if (class(pdata$distr)[1] == "array") { # from Arrdist
+    assert_array(pdata$distr, d = 3, any.missing = FALSE, null.ok = TRUE)
   } else {
     assert_matrix(pdata$distr, nrows = n, any.missing = FALSE, null.ok = TRUE)
   }
@@ -70,7 +72,7 @@ c.PredictionDataSurv = function(..., keep_duplicates = TRUE) {
   }
 
   if ("distr" %in% predict_types) {
-    if (inherits(dots[[1]], c("Matdist", "VectorDistribution"))) {
+    if (inherits(dots[[1]], c("Matdist", "VectorDistribution", "Arrdist"))) {
       result$distr = do.call(c, map(dots, "distr"))
     } else {
       result$distr = tryCatch(

R/PredictionSurv.R

@@ -13,12 +13,22 @@
 #' learner = lrn("surv.kaplan")
 #' p = learner$train(task, row_ids = 1:20)$predict(task, row_ids = 21:30)
 #' head(as.data.table(p))
+#' class(p$data$distr) # survival matrix stored
+#' p$distr # Matdist
 PredictionSurv = R6Class("PredictionSurv",
   inherit = Prediction,
   public = list(
     #' @description
     #' Creates a new instance of this [R6][R6::R6Class] class.
     #'
+    #' @details
+    #' Upon initialization, the `distr` input will be coerced to a survival matrix
+    #' or array (accessible via `$data$distr`) if it's a [Distribution][distr6::Distribution]
+    #' object. The active field `$distr` always returns a distribution
+    #' ([Matdist][distr6::Matdist] or [Arrdist][distr6::Arrdist]) depening on
+    #' the class of the stored `$data$distr`. In the case of an [Arrdist][distr6::Arrdist],
+    #' the distribution is by default initialized using `which.curve = 'mean'`.
+    #'
     #' @param task ([TaskSurv])\cr
     #'   Task, used to extract defaults for `row_ids` and `truth`.
     #'
@@ -33,10 +43,10 @@ PredictionSurv = R6Class("PredictionSurv",
     #'   observation in the test set. For a pair of continuous ranks, a higher rank indicates that
     #'   the observation is more likely to experience the event.
     #'
-    #' @param distr (`matrix()|[distr6::Matdist]|[distr6::VectorDistribution]`)\cr
-    #'   Either a matrix of predicted survival probabilities or a [distr6::VectorDistribution]
-    #'   or a [distr6::Matdist].
-    #'   If a matrix then column names must be given and correspond to survival times.
+    #' @param distr (`matrix()|[distr6::Arrdist]|[distr6::Matdist]|[distr6::VectorDistribution]`)\cr
+    #'   Either a matrix of predicted survival probabilities, a [distr6::VectorDistribution],
+    #'   a [distr6::Matdist] or an [distr6::Arrdist].
+    #'   If a matrix/array then column names must be given and correspond to survival times.
     #'   Rows of matrix correspond to individual predictions. It is advised that the
     #'   first column should be time `0` with all entries `1` and the last
     #'   with all entries `0`. If a `VectorDistribution` then each distribution in the vector
@@ -57,7 +67,7 @@ PredictionSurv = R6Class("PredictionSurv",
       distr = NULL, lp = NULL, response = NULL, check = TRUE) {
 
       if (inherits(distr, "Distribution")) {
-        # coerce to matrix if possible
+        # coerce to matrix/array if possible
         distr = private$.simplify_distr(distr)
       }
 
@@ -90,14 +100,14 @@ PredictionSurv = R6Class("PredictionSurv",
       self$data$crank %??% rep(NA_real_, length(self$data$row_ids))
     },
 
-    #' @field distr ([distr6::Matdist]|[distr6::VectorDistribution])\cr
-    #' Convert the stored survival matrix to a survival distribution.
+    #' @field distr ([distr6::Matdist]|[distr6::Arrdist]|[distr6::VectorDistribution])\cr
+    #' Convert the stored survival array or matrix to a survival distribution.
     distr = function() {
       if (inherits(self$data$distr, "Distribution")) {
         return(self$data$distr)
       }
 
-      private$.distrify_survmatrix(self$data$distr %??% NA_real_)
+      private$.distrify_survarray(self$data$distr %??% NA_real_)
     },
 
     #' @field lp (`numeric()`)\cr
@@ -117,7 +127,7 @@ PredictionSurv = R6Class("PredictionSurv",
     .censtype = NULL,
     .distr = function() self$data$distr %??% NA_real_,
     .simplify_distr = function(x) {
-      if (inherits(x, "Matdist")) {
+      if (inherits(x, c("Matdist", "Arrdist"))) {
         1 - gprm(x, "cdf")
       } else {
         if (!inherits(x, "VectorDistribution")) {
@@ -148,9 +158,16 @@ PredictionSurv = R6Class("PredictionSurv",
         surv
       }
     },
-    .distrify_survmatrix = function(x) {
-      distr6::as.Distribution(1 - x, fun = "cdf",
-        decorators = c("CoreStatistics", "ExoticStatistics"))
+    .distrify_survarray = function(x) {
+      if (inherits(x, "matrix")) {
+        # create Matdist
+        distr6::as.Distribution(1 - x, fun = "cdf",
+          decorators = c("CoreStatistics", "ExoticStatistics"))
+      } else {
+        # create Arrdist
+        distr6::Arrdist$new(cdf = 1 - x, which.curve = "mean",
+          decorators = c("CoreStatistics", "ExoticStatistics"))
+      }
     }
   )
 )

R/surv_measures.R

@@ -1,7 +1,7 @@
 surv_logloss = function(truth, distribution, eps = 1e-15, IPCW = TRUE, train = NULL, ...) {
 
   # calculate pdf at true death time
-  if (inherits(distribution, "Matdist")) {
+  if (inherits(distribution, c("Matdist", "Arrdist"))) {
     pred = diag(distribution$pdf(truth[, 1]))
   } else {
     pred = as.numeric(distribution$pdf(data = matrix(truth[, 1], nrow = 1)))
@@ -36,7 +36,7 @@ surv_logloss = function(truth, distribution, eps = 1e-15, IPCW = TRUE, train = N
   truth = truth[uncensored, 1]
   distribution = distribution[uncensored]
 
-  if (inherits(distribution, "Matdist")) {
+  if (inherits(distribution, c("Matdist", "Arrdist"))) {
     cens = diag(distribution$survival(truth))
   } else {
     cens = as.numeric(distribution$survival(data = matrix(truth, nrow = 1)))

R/surv_return.R

@@ -1,36 +1,59 @@
 #' @title Get Survival Predict Types
-#' @description Internal helper function to easily return the correct survival predict types and to
-#' automatically coerce a predicted survival probability matrix to a [distr6::Matdist].
+#'
+#' @description Internal helper function to easily return the correct survival predict types.
+#'
 #' @param times (`numeric()`) \cr Vector of survival times.
-#' @param surv (`matrix()`)\cr Matrix of predicted survival probabilities, rows are observations,
-#' columns are times. Number of columns should be equal to length of `times`.
+#' @param surv (`matrix()|array()`)\cr Matrix or array of predicted survival
+#' probabilities, rows (1st dimension) are observations, columns (2nd dimension)
+#' are times and in the case of an array there should be one more dimension.
+#' Number of columns should be equal to length of `times`. In case a `numeric()`
+#' vector is provided, it is converted to a single row (one observation) matrix.
+#' @param which.curve Which curve (3rd dimension) should the `crank` be
+#' calculated for, in case `surv` is an `array`? If between (0,1) it is taken as
+#' the quantile of the curves otherwise if greater than 1 it is taken as the
+#' curve index. It can also be 'mean' and the survival probabilities are averaged
+#' across the 3rd dimension. Default value (`NULL`) is the **0.5 quantile** which
+#' is the median across the 3rd dimension of the survival array.
 #' @param crank (`numeric()`)\cr Relative risk/continuous ranking. Higher value is associated
 #' with higher risk. If `NULL` then either set as `-response` if available or
 #' `lp` if available (this assumes that the `lp` prediction comes from a PH type
 #' model - in case of an AFT model the user should provide `-lp`).
 #' In case neither `response` or `lp` are provided, then `crank` is calculated
 #' as the sum of the cumulative hazard function (expected mortality) derived from
-#' the predicted survival function (`surv`).
+#' the predicted survival function (`surv`). In case `surv` is a 3d array, we use
+#' the `which.curve` parameter to decide which survival matrix (index in the 3rd
+#' dimension) will be chosen for the calculation of `crank`.
 #' @param lp (`numeric()`)\cr Predicted linear predictor, used to impute `crank` if `NULL`.
 #' @param response (`numeric()`)\cr Predicted survival time, passed through function without
 #' modification.
+#'
 #' @details
 #' Uses [survivalmodels::surv_to_risk] to reduce survival matrices to relative
 #' risks / rankings if `crank` is NULL.
+#'
 #' @references
 #' Sonabend, R., Bender, A., & Vollmer, S. (2022). Avoiding C-hacking when
 #' evaluating survival distribution predictions with discrimination measures.
 #' Bioinformatics. https://doi.org/10.1093/BIOINFORMATICS/BTAC451
 #' @export
-.surv_return = function(times = NULL, surv = NULL, crank = NULL, lp = NULL, response = NULL) {
+.surv_return = function(times = NULL, surv = NULL, which.curve = NULL,
+  crank = NULL, lp = NULL, response = NULL) {
 
   if (!is.null(surv)) {
     if (class(surv)[1] == "numeric") {
-      surv = matrix(surv, nrow = 1)
+      # in case of a vector (one observation) convert to matrix
+      surv = matrix(surv, nrow = 1, dimnames = list(NULL, names(surv)))
     }
-    times <- times %||% colnames(surv)
-    assert(length(times) == ncol(surv))
-    colnames(surv) <- times
+    if (class(surv)[1] == "array") {
+      if (length(dim(surv)) != 3) {
+        stop("3D survival arrays supported only")
+      }
+    }
+    times = times %||% colnames(surv)
+    if (length(times) != ncol(surv)) {
+      stop("'times' must have the same length as the 2nd dimension (columns of 'surv')")
+    }
+    colnames(surv) = times
   }
 
   if (is.null(crank)) {
@@ -42,14 +65,54 @@
       # assumes PH-type lp where high value = high risk
       crank = lp
     } else if (!is.null(surv)) {
-      crank = survivalmodels::surv_to_risk(surv)
+      if (inherits(surv, "matrix")) {
+        crank = survivalmodels::surv_to_risk(surv)
+      } else { # array
+        surv_mat = .ext_surv_mat(surv, which.curve)
+        crank = survivalmodels::surv_to_risk(surv_mat)
+      }
     }
   }
 
   list(
-    distr = surv,
+    distr = surv, # matrix or array
     crank = crank,
     lp = lp,
     response = response
   )
 }
+
+# helper function to extract a survival matrix from a 3D survival array
+.ext_surv_mat = function(arr, which.curve) {
+  # if NULL return the 'median' curve (default)
+  if (is.null(which.curve)) {
+    return(array(apply(arr, c(1, 2), quantile, 0.5), c(nrow(arr), ncol(arr)),
+      dimnames(arr)[c(1, 2)]))
+  }
+
+  # which.curve must be length 1 and either 'mean' or >0
+  ok = (length(which.curve) == 1) &&
+    ((is.character(which.curve) && which.curve == "mean") ||
+      (is.numeric(which.curve) && which.curve > 0))
+  if (!ok) {
+    stop("'which.curve' has to be a numeric between (0,1) or the index of the
+      3rd dimension or 'mean'")
+  }
+
+  if (is.numeric(which.curve) && which.curve > dim(arr)[3L]) {
+    stop(sprintf("Length is %s on third dimension but curve '%s' requested,
+          change 'which.curve' parameter.", dim(arr)[3L], which.curve))
+  }
+
+  # mean
+  if (which.curve == "mean") {
+    apply(arr, c(1, 2), mean)
+  # curve chosen based on quantile
+  } else if (which.curve < 1) {
+    array(apply(arr, c(1, 2), quantile, which.curve), c(nrow(arr), ncol(arr)),
+      dimnames(arr)[c(1, 2)])
+  # curve chosen based on index
+  } else {
+    array(arr[, , which.curve], c(nrow(arr), ncol(arr)), dimnames(arr)[c(1, 2)])
+  }
+}

inst/testthat/helper_expectations.R

@@ -31,7 +31,7 @@ expect_prediction_surv = function(p) {
   checkmate::expect_data_table(data.table::as.data.table(p), nrows  = length(p$row_ids))
   checkmate::expect_atomic_vector(p$missing)
   if ("distr" %in% p$predict_types) {
-    expect_true(class(p$distr)[[1]] %in% c("VectorDistribution", "Matdist"))
+    expect_true(class(p$distr)[[1]] %in% c("VectorDistribution", "Matdist", "Arrdist"))
   }
   expect_true(inherits(p, "PredictionSurv"))
 }