survML: Conditional survival function estimation using machine learning

The survML package implements two methods for estimating a conditional survival function using off-the-shelf machine learning. The first, called global survival stacking and performed using the stackG() function, involves decomposing the conditional cumulative hazard function into regression functions depending only on the observed data. The second, called local survival stacking or discrete-time hazard estimation, involves discretizing time and estimating the probability of an event of interest within each discrete time period. This procedure is implemented in the stackL() function.

These functions can be used for both left-truncated, right-censored data (commonly seen in prospective studies) and right-truncated data (commonly seen in retrospective studies).

More details on each method, as well as examples, follow.

Installing survML

You can install a stable version of survML from CRAN using

install.packages("survML")

Alternatively, the development version of survML is available on GitHub. You can install it using the devtools package as follows:

## install.packages("devtools") # run only if necessary
install_github(repo = "cwolock/survML")

Bugs reports and feature requests

To submit a bug report or request a new feature, please submit a new GitHub Issue.

Global survival stacking

In a basic survival analysis setting with right-censored data (for simplicity, we don’t discuss truncation here), the ideal data for each individual consist of a covariate vector \(X\), an event time \(T\), and a censoring time \(C\). The observed data consist of \(X\), the observed follow-up time \(Y:=\text{min}(T,C)\), and the event indicator \(\Delta := I(T \leq C)\). Global survival stacking requires three components: (1) the conditional probability that \(\Delta = 1\) given \(X\), (2) the CDF of \(Y\) given \(X\) among among censored subjects, and (3) the CDF of \(Y\) given \(X\) among uncensored subjects. All three of these can be estimated using standard binary regression or classification methods.

Estimating (1) is a standard binary regression problem. We use pooled binary regression to estimate (2) and (3). In essence, at time \(t\) each on a user-specified grid, the CDF is a binary regression using the outcome \(I(Y \leq t)\). The data sets for each \(t\) are combined into a single, pooled data set, including \(t\) as a covariate. Currently, survML allows Super Learner to be used for binary regression, but more learners will be added in future versions.

The stackG function performs global survival stacking. The most important user-specified arguments are described here:

Example

Here’s a small example applying stackG to simulated data.

# This is a small simulation example
set.seed(123)
n <- 500
X <- data.frame(X1 = rnorm(n), X2 = rbinom(n, size = 1, prob = 0.5))

S0 <- function(t, x){
  pexp(t, rate = exp(-2 + x[,1] - x[,2] + .5 * x[,1] * x[,2]), lower.tail = FALSE)
}
T <- rexp(n, rate = exp(-2 + X[,1] - X[,2] + .5 *  X[,1] * X[,2]))

G0 <- function(t, x) {
  as.numeric(t < 15) *.9*pexp(t,
                              rate = exp(-2 -.5*x[,1]-.25*x[,2]+.5*x[,1]*x[,2]),
                              lower.tail=FALSE)
}
C <- rexp(n, exp(-2 -.5 * X[,1] - .25 * X[,2] + .5 * X[,1] * X[,2]))
C[C > 15] <- 15

time <- pmin(T, C)
event <- as.numeric(T <= C)

# note that this a very small library, just for demonstration
SL.library <- c("SL.mean", "SL.glm", "SL.gam")

fit <- stackG(time = time,
              event = event,
              X = X,
              newX = X,
              newtimes = seq(0, 15, .1),
              direction = "prospective",
              bin_size = 0.02,
              time_basis = "continuous",
              time_grid_approx = sort(unique(time)),
              surv_form = "exp",
              SL_control = list(SL.library = SL.library,
                                V = 5))

We can plot the fitted versus true conditional survival at various times for one particular individual in our data set:

plot_dat <- data.frame(fitted = fit$S_T_preds[1,], 
                       true = S0(t =  seq(0, 15, .1), X[1,]))

p <- ggplot(data = plot_dat, mapping = aes(x = true, y = fitted)) + 
  geom_point() + 
  geom_abline(slope = 1, intercept = 0, color = "red") + 
  theme_bw() + 
  ylab("fitted") +
  xlab("true") + 
  ggtitle("Global survival stacking example (event time distribution)")

p

The stackG function simultaneously produces estimates for the conditional censoring distribution. This may be useful, for example, for producing inverse probability of censoring (IPCW) weights.

plot_dat <- data.frame(fitted = fit$S_C_preds[1,], 
                       true = G0(t =  seq(0, 15, .1), X[1,]))

p <- ggplot(data = plot_dat, mapping = aes(x = true, y = fitted)) + 
  geom_point() + 
  geom_abline(slope = 1, intercept = 0, color = "red") + 
  theme_bw() + 
  ylab("fitted") +
  xlab("true") + 
  ggtitle("Global survival stacking example (censoring time distribution)")

p

Local survival stacking

For discrete time-to-event variables, the hazard function at a single time is a conditional probability whose estimation can be framed as a binary regression problem: among those who have not experienced the event by time \(t\), what proportion experience the outcome at that time? Local survival stacking assumes a discrete survival process and is based on estimating this conditional event probability at each time in a user-specified grid. These binary regressions are estimated jointly by “stacking” the data sets corresponding to all times in the grid. This idea dates back at least to work by Polley and van der Laan (2011) and was also recently described by Craig et al. (2021).

Example

fit <- stackL(time = time,
              event = event,
              X = X,
              newX = X,
              newtimes = seq(0, 15, .1),
              direction = "prospective",
              bin_size = 0.02,
              time_basis = "continuous",
              SL_control = list(SL.library = SL.library,
                                V = 5))
plot_dat <- data.frame(fitted = fit$S_T_preds[1,], 
                       true = S0(t =  seq(0, 15, .1), X[1,]))

p <- ggplot(data = plot_dat, mapping = aes(x = true, y = fitted)) + 
  geom_point() + 
  geom_abline(slope = 1, intercept = 0, color = "red") + 
  theme_bw() + 
  ylab("fitted") +
  xlab("true") + 
  ggtitle("Local survival stacking example")

p

References

For details of global survival stacking, please see the following preprint:

Charles J. Wolock, Peter B. Gilbert, Noah Simon and Marco Carone. “A framework for leveraging machine learning tools to estimate personalized survival curves.” arXiv: 2211.03031.

Local survival stacking is described in:

Eric C. Polley and Mark J. van der Laan. “Super Learning for Right-Censored Data” in Targeted Learning (2011).

Erin Craig, Chenyang Zhong, and Robert Tibshirani. “Survival stacking: casting survival analysis as a classification problem.” arXiv:2107.13480.