initial_validation_split() stratification for Surv objects

[krz]

Hi,

I'm using tidymodels for a survial prediction task. My data has very few events so I want to make sure the validation split is stratified by the event variable. However, this seems to no be implemented currently.

Here is an example:

# example data
library(survival)
attach(kidney)

kidney <- kidney %>% 
  mutate(
    kidney_surv = Surv(time, status), 
    .keep = "unused"
  )

kidney_split <- initial_validation_split(kidney, strata = kidney_surv)

I do get the following error message: strata cannot be a Surv object. Use the time or event variable directly.

But it is unclear to me how it is supposed to work. None of these work:

kidney_split <- initial_validation_split(kidney, strata = kidney_surv)
kidney_split <- initial_validation_split(kidney, strata = kidney_surv[, "status"])
kidney_split <- initial_validation_split(kidney, strata = "kidney_surv[, "status"]")
kidney_split <- initial_validation_split(kidney, strata = kidney_surv[, 2])

[hfrick]

This is indeed a little bumpy. We have an issue to revisit this when we feel comfortable that we have a good solution: #445 .

Our current idea is to remove the original time and status variables not when you make the Surv object but rather after you've made your resampling object, i.e., in a recipe:

library(rsample)
library(recipes)
# [...]
library(survival)
attach(kidney)

kidney <- kidney %>%
  mutate(
    surv = Surv(time, status) # keep the original variables here
  )

# so you can use them for stratification
kidney_split <- initial_validation_split(kidney, strata = status)

# then remove them in your recipe
kidney_rec <- recipe(surv ~ ., data = training(kidney_split)) %>% 
  step_rm(time, status)

^{Created on 2025-01-24 with reprex v2.1.1}

[krz]

thanks, this works for me