update mlr3proba to 0.5.3 + refactoring

mlr-org · Sep 11, 2023 · f162a31 · f162a31
1 parent c8041b4
commit f162a31
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 12 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -67,7 +67,7 @@ Suggests:
     mgcv,
     mlr3cluster,
     mlr3learners (>= 0.4.2),
-    mlr3proba,
+    mlr3proba (>= 0.5.3),
     mlr3pipelines,
     mvtnorm,
     nnet,

diff --git a/R/learner_BART_surv_bart.R b/R/learner_BART_surv_bart.R
@@ -200,17 +200,9 @@ delayedAssign(
           c(3, 2, 1)
         )
 
-        # Mean posterior survival matrix (N obs x K times)
-        surv = matrix(pred$surv.test.mean, nrow = N, ncol = K, byrow = TRUE)
-
-        # get crank as expected mortality using mean posterior
-        pred_list = mlr3proba::.surv_return(times = times, surv = surv)
-
-        # replace with the full survival posterior
-        pred_list$distr = surv.array
-
-        # return list with crank and distr
-        pred_list
+        # distr => full survival array
+        # crank => expected mortality using the mean posterior survival matrix
+        mlr3proba::.surv_return(times = times, surv = surv.array, which.curve = 'mean')
       }
     )
   )