fix bug when input is a 3d survival array

mlr-org · Sep 18, 2023 · edd33f8 · edd33f8
1 parent 6df6afe
commit edd33f8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/R/integrated_scores.R b/R/integrated_scores.R
@@ -26,7 +26,11 @@ weighted_survival_score = function(loss, truth, distribution, times, t_max, p_ma
     unique_times = .c_get_unique_times(truth[, "time"], times)
   }
 
-  if (inherits(distribution, "Distribution")) {
+  is_distr_or_array =
+    (inherits(distribution, "Distribution")) ||
+    (inherits(distribution, "array") & length(dim(array)) == 3)
+
+  if (is_distr_or_array) {
     cdf = as.matrix(distribution$cdf(unique_times))
   } else {
     mtc = findInterval(unique_times, as.numeric(colnames(distribution)))