Predicting Heart Failure with Deep Survival Analysis (DeepSurv)

Introduction

Traditional classification models predict if an event (like heart failure or death) will occur. Survival analysis predicts when it will occur, while naturally handling right-censored data (patients who drop out or haven't experienced the event yet). In this tutorial, we will use the pycox library to train DeepSurv, a deep learning extension of the Cox Proportional Hazards model.

Prerequisites

• Python 3.9+
• PyTorch
• pycox and pandas
• A clinical dataset with time-to-event and censoring indicators (e.g., the SUPPORT dataset)

Step 1: Data Preparation

In survival analysis, the target variable is a tuple of (time, event). If event=1, the patient experienced the event at time. If event=0, the patient was right-censored at time.

import pandas as pd
import numpy as np
from pycox.datasets import support
from sklearn.preprocessing import StandardScaler

# Load the SUPPORT dataset (Study to Understand Prognoses Preferences Outcomes and Risks of Treatment)
df = support.read_df()

# Handle missing values (simple imputation for this example)
df = df.fillna(df.mean())

# Separate features from targets
cols_leave = ['age', 'sex', 'dzgroup', 'dzclass', 'num.co', 'edu', 'income', 'scoma', 'charges', 'totcst', 'totmcst', 'avtisst', 'race', 'sps', 'aps', 'surv2m', 'surv6m', 'hday', 'diabetes', 'dementia', 'ca', 'prg2m', 'prg6m', 'dnr', 'dnrday', 'meanbp', 'wblc', 'hrt', 'resp', 'temp', 'pafi', 'alb', 'bili', 'crea', 'sod', 'ph', 'glucose', 'bun', 'urine', 'adlp', 'adls', 'sfd', 'adlsc']
features = ['age', 'sex', 'num.co', 'meanbp', 'hrt', 'resp', 'temp', 'wblc', 'crea', 'sod'] # subset for simplicity
X = df[features].values

# Targets: duration (time) and event indicator
get_target = lambda df: (df['duration'].values, df['event'].values)
y_time, y_event = get_target(df)

# Scale features
scaler = StandardScaler()
X = scaler.fit_transform(X)
X = X.astype('float32')
y_time = y_time.astype('float32')
y_event = y_event.astype('float32')

Step 2: Defining the DeepSurv Model

DeepSurv replaces the linear combination of features in a standard Cox model with a multi-layer perceptron (MLP), allowing it to model complex, non-linear interactions between clinical variables.

import torch
import torchtuples as tt
from pycox.models import CoxPH

in_features = X.shape[1]
num_nodes = [32, 32]
out_features = 1 # DeepSurv outputs a single risk score
batch_norm = True
dropout = 0.1

net = tt.practical.MLPVanilla(in_features, num_nodes, out_features, batch_norm, dropout)

# Initialize the CoxPH model with our network
model = CoxPH(net, tt.optim.Adam)

Step 3: Training the Model

We train the model using the Cox partial log-likelihood loss function. This function only considers the relative risk rankings of patients at the specific times events occur.

# Set learning rate
model.optimizer.set_lr(0.01)

batch_size = 256
epochs = 50
callbacks = [tt.callbacks.EarlyStopping()]

# Fit the model
log = model.fit(X, (y_time, y_event), batch_size, epochs, callbacks, val_split=0.2)

# Plot training loss
_ = log.plot()

Step 4: Predicting Survival Curves

Once trained, the model can generate a full survival curve for a new patient, showing their probability of survival across all time points.

# Compute the baseline hazard (required for predicting survival probabilities)
_ = model.compute_baseline_hazards()

# Predict survival for the first 5 patients in the dataset
surv = model.predict_surv_df(X[:5])

import matplotlib.pyplot as plt
surv.plot()
plt.ylabel('Survival Probability')
plt.xlabel('Time (days)')
plt.show()

Conclusion

Deep survival analysis is incredibly powerful for time-to-event predictions in cardiology and oncology. By predicting the entire survival curve, clinicians get a much richer view of patient risk trajectories than a simple binary readmission classifier.