add ipcw and initial performance function

Author: topepo

R/grid_helpers.R

@@ -8,6 +8,12 @@ predict_model <- function(split, workflow, grid, metrics, submodels = NULL,
   x_vals <- forged$predictors
   y_vals <- forged$outcomes
 
+  # TODO patch since parsnip does not record the column names when Surv objects
+  # are used with fit_xy()
+  if (model$spec$mode == "censored regression") {
+    model$preproc$y_var <- names(y_vals)
+  }
+
   orig_rows <- as.integer(split, data = "assessment")
 
   if (length(orig_rows) != nrow(x_vals)) {
@@ -92,9 +98,28 @@ predict_model <- function(split, workflow, grid, metrics, submodels = NULL,
     }
   }
 
-  tibble::as_tibble(res)
+  # TODO do we need this?
+  # res <- tibble::as_tibble(res)
+  maybe_add_ipcw(res, model, eval_time, types)
 }
 
+maybe_add_ipcw <- function(.data, model, eval_time, types) {
+  if (!any(types == "survival")) {
+    return(.data)
+  }
+  res <-
+    tidyr::unnest(.data, cols = .pred) %>%
+    dplyr::rename(eval_time = .time) %>%
+    dplyr::full_join(
+      # TODO is the outcome name enforced or the original name?
+      parsnip::.censoring_weights_graf(model, .data, eval_time = eval_time),
+      by = c(".row", "eval_time")
+    )
+  res
+}
+
+
+
 predict_wrapper <- function(model, new_data, type, eval_time, subgrid = NULL) {
   if (is.null(subgrid)) {
     fn <- "predict.model_fit"

R/grid_performance.R

@@ -83,7 +83,7 @@ metrics_info <- function(x) {
   } else if (all(types == "class" | types == "prob")) {
     estimate_class_prob(dat, metric, param_names, outcome_name, case_weights, types, event_level)
   } else if (all(types == "time" | types == "survival")) {
-    # estimate_surv(dat, metric, param_names, outcome_name, case_weights, types, eval_time)
+    estimate_surv(dat, metric, param_names, outcome_name, case_weights, types)
   } else {
     rlang::abort("Metric type not yet supported by tune.")
   }
@@ -133,14 +133,22 @@ estimate_class_prob <- function(dat, metric, param_names, outcome_name,
     )
 }
 
-estimate_surv <- function(dat, metric, param_names, outcome_name, case_weights, eval_time) {
-  # IPCW should already be computed, un-nested and have .time
-  types <- NULL
+estimate_surv <- function(dat, metric, param_names, outcome_name, case_weights, types) {
+  # TODO mixed sets?
   if (any(types == "survival")) {
-
+    res <-
+      dat %>%
+      dplyr::group_by(!!!rlang::syms(param_names), eval_time) %>%
+      metric(
+        truth = surv,
+        estimate = .pred_survival,
+        censoring_weights = .weight_cens,
+        case_weights = !!case_weights,
+        eval_time = eval_time
+      )
   } else {
-    # pad with .time = NA
+    # pad with .time = NA?
   }
-  # metric(estimate = .pred, truth = !!sym(outcome_name), case_weights = !!case_weights)
+  res
 }