In this case study we show how to implement an estimator for the leave-one-out (LOO) version of the expected Brier score (denoted by LOO-BS) for Bayesian survival models, with potentially right-censored event times. We combine the “inverse of probability of censoring weighted estimator” by Gerds et al. with Pareto-smoothed importance sampling (PSIS), by Vehtari et al., in order to efficiently approximate the LOO-BS. We compare the PSIS approximation to the exact LOO-BS.

As a by-product we show how to implement proportional hazard models with a time-dependent baseline hazard by using monotone splines (M-splines).

Useful References:

For PSIS we recommend:

Vehtari, Aki, Andrew Gelman, and Jonah Gabry. “Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC.” Statistics and Computing 27.5 (2017): 1413-1432. APA

For the inverse of probability of censoring weighted estimator see:

Gerds, Thomas A., and Martin Schumacher. “Consistent estimation of the expected Brier score in general survival models with right‐censored event times.” Biometrical Journal 48.6 (2006): 1029-1040.

Monotone (M-splines) were introduced here:

Ramsay, James O. “Monotone regression splines in action.” Statistical science 3.4 (1988): 425-441.

See als this Stan discourse discussion.

knitr::opts_chunk$set(echo = TRUE)
library(rstan);rstan_options(auto_write = TRUE)
library(loo)
library(survival)
library(tibble)
library(dplyr)
library(purrr)
library(bayesplot)
library(splines2)
library(pec)
library(cowplot)
library(MCMCpack)
library(RColorBrewer)
qual_col_pals = brewer.pal.info[brewer.pal.info$category == 'qual',]
col_vector = unlist(mapply(brewer.pal, qual_col_pals$maxcolors, rownames(qual_col_pals)))

sm <- stan_model("~/Desktop/Stan/Loo_Pec/survival_parametric_baseline_hazard_simplex_loo.stan")
expose_stan_functions("~/Desktop/Stan/Loo_Pec/survival_parametric_baseline_hazard_simplex_loo.stan")
df <- read.delim("~/Desktop/Stan/Loo_Pec/ncog.txt",sep=" ")

df <- df  %>% mutate(event=d,time=t*12/365) 
df <- df%>% mutate(isGroupB=as.double(arm=="B")) %>%
  dplyr::select(isGroupB, event, time)
N <- nrow(df)
times <- pull(df,time)
is_censored <- pull(df,event)==0
msk_censored <- is_censored == 1
time_range <- range(times)
time_min <- time_range[1]
time_max <- time_range[2]
X <- as.matrix(pull(df, isGroupB))

The dataset

The table below reports survival data from a randomized control trial run by the Nothern California Oncology Group, that compares two treatments for head and neck cancer: isGroupB==0 in case of chemotherapy and isGroup==1 if chemotherapy and radiation was applied. The response variable time is the survival time in months (actually days divided by 12/365) and the variable event is equal to 0 iff the survival time of that patient is only known to exceed the corresponding survival time. For details see (starting) chapter 9.2 in “Computer Age Statistical Inference” by B. Efron and T. Hastie.

Prior

We build a proportional hazard model with a baseline-hazard that is parametrized by M-Splines. Below we show prior predictive (checks) draws for the time-dependent baseline1 hazard and survival functions.

For details we refer to accompanying Stan Code in survival_parametric_baseline_hazard_simplex_loo.stan.

The functions surv_t and hazard_t are defined in the above file and are exposed above as R functions through expose_stan_functions. They allow the evaluation of the survival and hazard function, respectively, at fixed time, for different (latent) paremters and covariates. Both functions return a matrix with as many rows as covariates and as many columns as parameters. For their respective signatures, refer to the definition in the Stan file.

Posterior

Let’s obtain our Posterior…

Get a summary of the posterior:

Hazard

quantile_func <- function(x, probs) {
  x <- as.data.frame(t(x))
  map_dbl(x, ~quantile(., probs=probs))
}

get_post_plot_df <- function(eval_lst) {
  funcs <- purrr::map(eval_lst, rowMeans)
  funcs_low <- purrr::map(eval_lst, ~quantile_func(.,probs=c(0.025)))
  funcs_up <- purrr::map(eval_lst, ~quantile_func(.,probs=c(0.975)))
  dplyr::bind_rows(
    tibble(
    t=times_plot,
    isGroupB=F,
    func=map_dbl(funcs, ~.[[1]]),
    func_low=map_dbl(funcs_low,~.[[1]]),
    func_up=map_dbl(funcs_up,~.[[1]])
    
    ),
      tibble(
    t=times_plot,
    isGroupB=T,
    func=map_dbl(funcs, ~.[[2]]),
    func_low=map_dbl(funcs_low,~.[[2]]),
    func_up=map_dbl(funcs_up,~.[[2]])
    )
  ) 
}

get_post_plot <- function(plot_df, ylab) {
  ggplot(data=plot_df) + 
  geom_ribbon(mapping=aes(x=t,ymin=func_low, ymax=func_up, fill=isGroupB), alpha=.25)+
  geom_line(mapping=aes(x=t, y=func, color=isGroupB))+
  geom_vline(xintercept = c(0,knots, time_max), alpha=.5, color='gray', linetype='dashed')+
  geom_rug(data=tibble(times=times[is_censored]), aes(x=times), col="green",alpha=0.5, size=1.25)+
  geom_rug(data=tibble(times=times[!is_censored]), aes(x=times), col="red",alpha=0.5, size=1.25)+
  ylab(ylab)+
  xlab("Time")
}

intercepts <- as.vector(post[,,"intercept"])
betas <- t(matrix(post[,,sprintf("betas[%d]" ,1:stan_data$NC)], dim(post)[1]*dim(post)[2], stan_data$NC))
gammas <- t(matrix(post[,,sprintf("gammas[%d]", 1:nbasis)], dim(post)[1]*dim(post)[2], nbasis))

X_surv = as.matrix(c(0,1))

times_plot <- seq(0, time_max, length.out = nplot_points)
i_spline_basis_evals <- iSpline(times_plot, 
                                knots=knots, 
                                degree=mspline_degree,
                                intercept=FALSE, 
                                Boundary.knots = c(0, time_max))
m_spline_basis_evals <- deriv(i_spline_basis_evals)

Survival curves

Brier Score (LOO-BS)

Introduction

Denote by \(\mathbf{T}_{-i}\) the \(N-1\) dimensional vector2 of observed times (censored and uncensored) for all individuals but for individual \(i\). Furthermore denote by \(\mathbf{T}\) the \(N\) dimensional vector of observed times (censored and uncensored) for all individuals.

Now, write \(S_i(t\vert \mathbf{T}_{-i})\) for the expected survival probability of individual \(i\) at time \(t\), where the expectation is with respect to the posterior over latent parameters \(\boldsymbol{\theta}\), given \(\mathbf{T}_{-i}\).

Suppose \(\{\boldsymbol{\theta}_s\}_{s=1\dots S}\) is a (correlated) sequence of latent model parameters drawn from the posterior, given \(\mathbf{T}\). We approximate \(S_i(t\vert \mathbf{T}_{-i})\) as follows

\[ S_i(t\vert \mathbf{T}_{-i}) \approx \frac{\sum_{s=1}^S r_{s,i} S_i(t\vert \boldsymbol{\theta}_s)}{\sum_{s=1}^S r_{s,i}} \]

where \(S_i(t\vert \boldsymbol{\theta}_s)\) is the probability that individual \(i\) survives beyond time \(t\), given latent parameters \(\boldsymbol{\theta}_s\), and

\[ r_{s,i} \equiv \frac{1}{p(T_i\vert \boldsymbol{\theta}_s)} \]

Coming to the Brier score, we now define

\[ \eta_{i,t} \equiv \left\{ 1_{\{T_i > t\}} - S_i(t\vert \mathbf{T}_{-i})\right\}^2 \]

Based on the residuals \(\eta_{i}(t)\) we define, following Gerds et al., the inverse probability of censoring weighted estimator \(\Omega(t)\):

\[ \Omega(t)\equiv \frac{1}{N} \sum_{i=1}^N w_i(t)\eta_{i}(t)^2 \]

with

\[ w_i(t) \equiv \frac{\mathbf{1}_{\{T_i \leq t\}} \delta_i}{G_i(T_i)} + \frac{\mathbf{1}_{\{T_i > t\}}}{G_i(t)} \]

Here \(\delta_i\) is the event indicator, i.e. \(\delta_i = 1\) iff individual experienced the event and \(0\) otherwise. Moreover, \(G_i(t)\) is the probability that the censoring time3, for individual \(i\), denoted by \(C_i\), exceeds \(t\). Note that for both the event and censoring survival functions we indicate a potential dependence on covariates by adding the index \(i\). \(G_i(t)\) can be obtained by fitting, for example, a Cox model to the data with reversed event indicators. By introducing a weighting based on \(G_i\) one can guarantee that the bias of the estimator for \(\Omega(t)\) does not depent on the survival model (otherwise model comparison would be hard!). Finally, it can be shown that under weak conditions, \(\Omega(t)\) is a uniformly consistent estimator for the expected mean squared error in time (aka Brier score). We refer the reader for more details and proofs to Gerds et al..

Practically, this means we can estimate \(S_i(t\vert \mathbf{T}_{-i})\) for all \(i\) and required times \(t\), based on posterior samples \(S_i(t\vert \boldsymbol{\theta}_s)\), obtained from an model for all individuals! This matrix (!) \(\{S_i(t\vert \mathbf{T}_{-i})\}_{i,t}\) can then be passed to the function pec() provided by pec R Package, which calculates \(\Omega(t)\) and provides ways to choose \(G_i\):

Comparison

## 
## Computed from 8000 by 96 log-likelihood matrix
## 
##          Estimate   SE
## elpd_loo   -274.7 13.1
## p_loo         4.4  0.4
## looic       549.5 26.3
## ------
## Monte Carlo SE of elpd_loo is 0.0.
## 
## All Pareto k estimates are good (k < 0.5).
## See help('pareto-k-diagnostic') for details.

  1. Baseline here means control group isGroupB==0↩

  2. where \(N\) is the number of individuals↩

  3. Here we assume that actual event process competes with a censoring process and that the observed time is the minimum of the two; right censoring occurrs precisely when the censoring time is smaller than the event time.↩