restricted mean survival time

R/misc.R

@@ -126,6 +126,21 @@ has_units <- function(x){
 contains_vi <- function(object) {!is_empty(object$importance)}
 
 
+#' Restricted mean survival time
+#'
+#' @param probs vector of survival probs
+#' @param times vector of event times
+#'
+#' @return a double
+#' @noRd
+#'
+rmst <- function(probs, times){
+
+ diffs <- collapse::fdiff(times)
+ diffs[1] <- times[1]
+ collapse::fsum(diffs * probs)
+
+}
 
 #' beautify time
 #'

R/orsf_R6.R

@@ -3108,10 +3108,14 @@ ObliqueForest <- R6::R6Class(
                          "prob"  = 6,
                          "class" = 7,
                          "leaf"  = 8,
-                         "time"  = 9),
+                         "time"  = 2), # time=2 is not a typo
     pred_mode = .dots$pred_mode %||% FALSE,
     pred_aggregate = .dots$pred_aggregate %||% (self$pred_type != 'leaf'),
-    pred_horizon = .dots$pred_horizon %||% self$pred_horizon %||% 1,
+    pred_horizon = if(self$pred_type == 'time'){
+     self$event_times
+    } else {
+     .dots$pred_horizon %||% self$pred_horizon %||% 1
+    },
     oobag = .dots$oobag %||% self$oobag_pred_mode,
     oobag_R_function = .dots$oobag_eval_function %||% self$oobag_eval_function,
     oobag_eval_type_R = switch(
@@ -3282,7 +3286,8 @@ ObliqueForestSurvival <- R6::R6Class(
 
   leaf_min_events = NULL,
   split_min_events = NULL,
-  pred_horizon = NULL
+  pred_horizon = NULL,
+  event_times = NULL
 
  ),
 
@@ -3579,6 +3584,10 @@ ObliqueForestSurvival <- R6::R6Class(
    private$max_time <- y[last_value(private$data_row_sort), 1]
    # boundary check for event-based tree parameters
    private$n_events <- collapse::fsum(y[, 2])
+   # unique event times sorted in ascending order
+   self$event_times <- sort(collapse::funique(
+    y[[1]][collapse::whichv(x = y[[2]], value = 1)]
+   ), decreasing = FALSE)
 
    # if pred_horizon is unspecified, provide sensible default
    # if it is specified, check for correctness
@@ -3755,9 +3764,23 @@ ObliqueForestSurvival <- R6::R6Class(
    unsorted <- collapse::radixorder(private$data_row_sort)
    self$pred_oobag <- self$pred_oobag[unsorted, , drop = FALSE]
 
-   # mortality predictions should always be 1 column
+   if(self$pred_type == 'time'){
+
+    self$eval_oobag$stat_values <-
+     self$eval_oobag$stat_values[, ncol(self$pred_oobag), drop = FALSE]
+
+    self$pred_oobag <- apply(self$pred_oobag,
+                             MARGIN = 1,
+                             FUN = rmst,
+                             times = self$event_times)
+
+    self$pred_oobag <- collapse::qM(self$pred_oobag)
+
+   }
+
+   # these predictions should always be 1 column
    # b/c they do not depend on the prediction horizon
-   if(self$pred_type %in% c('mort', 'time')){
+   if(self$pred_type == 'mort'){
 
     self$eval_oobag$stat_values <-
      self$eval_oobag$stat_values[, 1L, drop = FALSE]
@@ -3769,6 +3792,16 @@ ObliqueForestSurvival <- R6::R6Class(
   },
   clean_pred_new_internal = function(preds){
 
+   if(self$pred_type == 'time'){
+
+    preds <- apply(preds,
+                   MARGIN = 1,
+                   FUN = rmst,
+                   times = self$event_times)
+    preds <- collapse::qM(preds)
+
+   }
+
    # don't let multiple pred horizon values through for mort
    if(self$pred_type %in% c('mort', 'time')){
     return(preds[, 1, drop = FALSE])

src/TreeSurvival.cpp

@@ -562,43 +562,21 @@
 
    } case PRED_TIME: {
 
-    // believe it or not this method seems to be more accurate
-    // than the traditional one commented out beneath it.
-    temp_vec.fill(median(leaf_times));
-
-    // // does the kaplan meier in this node go below 50% chance of survival?
-    // uvec prob_lt_50 = find(leaf_pred_prob[leaf_id] <= 0.5);
-    //
-    // // If yes, then find the time it crosses
-    // if(prob_lt_50.size() >= 1){
-    //
-    //  // index of the first instance where survival prob is < 50
-    //  uword first_row_50 = prob_lt_50[0];
-    //
-    //  // if the survival prob here is exactly 50, or if
-    //  // there is no predicted probability before this point,
-    //  // do nothing.
-    //
-    //  // otherwise, use the predicted time just before the survival
-    //  // probability dips below 50% (think about how kaplan meiers look)
-    //  double tmp_prob = leaf_pred_prob[leaf_id][first_row_50];
-    //
-    //  if(first_row_50 > 0 && tmp_prob < 0.5) first_row_50--;
-    //
-    //  // use the time value at this specific index
-    //  double time_value = leaf_times[first_row_50];
-    //
-    //  temp_vec.fill(time_value);
-    //
-    // } else {
-    //
-    //  // if the probability of survival never goes below 50%,
-    //  // then it is more likely that the observation's time is
-    //  // greater than the max time of this node. For simplicity,
-    //  // use the max time as the prediction.
-    //  temp_vec.fill(leaf_times[leaf_times.size()-1]);
-    //
-    // }
+    // restricted mean survival time
+
+    leaf_values = vec(leaf_pred_prob[leaf_id].begin(),
+                      leaf_pred_prob[leaf_id].size(),
+                      false);
+
+    double result = leaf_times[0] * leaf_values[0];
+
+    for(uword i = 1; i < leaf_times.size(); i++){
+
+     result += (leaf_times[i] - leaf_times[i-1]) * leaf_values[i];
+
+    }
+
+    temp_vec.fill(result);
 
     break;
 

src/orsf_oop.cpp

@@ -452,6 +452,11 @@ double compute_mse_exported(arma::vec& y,
 
   case TREE_SURVIVAL:
 
+   // if(pred_type == PRED_TIME){
+   //  pred_type = PRED_SURVIVAL;
+   //  pred_horizon = find_unique_event_times();
+   // }
+
    forest = std::make_unique<ForestSurvival>(leaf_min_events,
                                              split_min_events,
                                              pred_horizon);