The SUPPORT dataset comes from the Vanderbilt University study to estimate survival for seriously ill hospitalized adults. (Refer to http://biostat.mc.vanderbilt.edu/wiki/Main/SupportDesc. for the original datasource.)
In this notebook, we will apply Deep Survival Machines for survival prediction on the SUPPORT data.
The package includes helper functions to load the dataset.
X represents an np.array of features (covariates), T is the event/censoring times and, E is the censoring indicator.
from dsm import datasets
x, t, e = datasets.load_dataset('SUPPORT')
Survival predictions are issued at certain time horizons. Here we will evaluate the performance of DSM to issue predictions at the 25th, 50th and 75th event time quantile as is standard practice in Survival Analysis.
import numpy as np
horizons = [0.25, 0.5, 0.75]
times = np.quantile(t[e==1], horizons).tolist()
We will train DSM on 70% of the Data, use a Validation set of 10% for Model Selection and report performance on the remaining 20% held out test set.
n = len(x)
tr_size = int(n*0.70)
vl_size = int(n*0.10)
te_size = int(n*0.20)
x_train, x_test, x_val = x[:tr_size], x[-te_size:], x[tr_size:tr_size+vl_size]
t_train, t_test, t_val = t[:tr_size], t[-te_size:], t[tr_size:tr_size+vl_size]
e_train, e_test, e_val = e[:tr_size], e[-te_size:], e[tr_size:tr_size+vl_size]
Lets set up the parameter grid to tune hyper-parameters. We will tune the number of underlying survival distributions, ($K$), the distribution choices (Log-Normal or Weibull), the learning rate for the Adam optimizer between $1\times10^{-3}$ and $1\times10^{-4}$ and the number of hidden layers between $0, 1$ and $2$.
from sklearn.model_selection import ParameterGrid
param_grid = {'k' : [3, 4, 6],
'distribution' : ['LogNormal', 'Weibull'],
'learning_rate' : [ 1e-4, 1e-3],
'layers' : [ [], [100], [100, 100] ]
}
params = ParameterGrid(param_grid)
from dsm import DeepSurvivalMachines
models = []
for param in params:
model = DeepSurvivalMachines(k = param['k'],
distribution = param['distribution'],
layers = param['layers'])
# The fit method is called to train the model
model.fit(x_train, t_train, e_train, iters = 100, learning_rate = param['learning_rate'])
models.append([[model.compute_nll(x_val, t_val, e_val), model]])
best_model = min(models)
model = best_model[0][1]
out_risk = model.predict_risk(x_test, times)
out_survival = model.predict_survival(x_test, times)
We evaluate the performance of DSM in its discriminative ability (Time Dependent Concordance Index and Cumulative Dynamic AUC) as well as Brier Score.
from sksurv.metrics import concordance_index_ipcw, brier_score, cumulative_dynamic_auc
cis = []
brs = []
et_train = np.array([(e_train[i], t_train[i]) for i in range(len(e_train))],
dtype = [('e', bool), ('t', float)])
et_test = np.array([(e_test[i], t_test[i]) for i in range(len(e_test))],
dtype = [('e', bool), ('t', float)])
et_val = np.array([(e_val[i], t_val[i]) for i in range(len(e_val))],
dtype = [('e', bool), ('t', float)])
for i, _ in enumerate(times):
cis.append(concordance_index_ipcw(et_train, et_test, out_risk[:, i], times[i])[0])
brs.append(brier_score(et_train, et_test, out_survival, times)[1])
roc_auc = []
for i, _ in enumerate(times):
roc_auc.append(cumulative_dynamic_auc(et_train, et_test, out_risk[:, i], times[i])[0])
for horizon in enumerate(horizons):
print(f"For {horizon[1]} quantile,")
print("TD Concordance Index:", cis[horizon[0]])
print("Brier Score:", brs[0][horizon[0]])
print("ROC AUC ", roc_auc[horizon[0]][0], "\n")