9.4 Survival analysis

When we study the time until something happens, some cases have not had the event yet when the study ends. Their time is cut short, not unknown: they lasted at least this long. Survival methods use that partial information instead of throwing it away.

Code: data and style used in this chapter

import matplotlib.pyplot as plt
import numpy as np
import statsmodels.api as sm
from matplotlib.axes import Axes
from scipy import optimize, stats
from statsmodels.duration.survfunc import SurvfuncRight, survdiff

from statly import style

style.use()
rng = np.random.default_rng(1)

heart = sm.datasets.heart.load_pandas().data  # 69 heart transplant patients, Stanford
days = (
    heart.survival.to_numpy()
)  # days from the transplant to death, or to the end of follow-up
died = heart.censors.to_numpy().astype(int)  # 1 = died, 0 = still alive when last seen
age = heart.age.to_numpy()


def fit_band(fit: SurvfuncRight) -> tuple[np.ndarray, np.ndarray]:
    """A 95% band from the standard errors, kept between 0 and 1."""
    half = 1.96 * fit.surv_prob_se
    return np.clip(fit.surv_prob - half, 0, 1), np.clip(fit.surv_prob + half, 0, 1)


def plot_curve(
    ax: Axes, fit: SurvfuncRight, color: str, label: str | None = None, band: bool = True
) -> None:
    """A Kaplan-Meier curve as steps, with censoring ticks and (optionally) a 95% band."""
    times = np.r_[0, fit.surv_times, fit.time.max()]  # run on to the last patient followed
    probs = np.r_[1, fit.surv_prob, fit.surv_prob[-1]]
    ax.step(times, probs, where="post", color=color, lw=2, label=label)
    if band:
        low, high = fit_band(fit)
        ax.fill_between(fit.surv_times, low, high, step="post", color=color, alpha=0.15, lw=0)
    censored_times = fit.time[fit.status == 0]
    at = probs[
        np.searchsorted(times, censored_times, side="right") - 1
    ]  # the step's height there
    ax.plot(censored_times, at, "|", ms=8, color=color)

9.4.1 Censoring

Sixty-nine patients received a heart transplant at Stanford. For each, we know how many days they lived afterwards, but 24 were still alive when they were last seen. For them we know only that they survived at least that long. Their times are censored.

Code

shown = np.sort(rng.choice(len(days), 30, replace=False))
order = shown[np.argsort(days[shown])]


def plot_timelines(ax: Axes) -> None:
    """One line per patient: deaths as filled dots, censored patients as open circles."""
    for row, i in enumerate(order):
        ax.plot([0, days[i]], [row, row], color=style.GREY, lw=1.2)
        if died[i]:
            ax.plot(days[i], row, "o", ms=6, color=style.CORAL)
        else:
            ax.plot(days[i], row, "o", ms=6, mfc="none", mec=style.BLUE, mew=1.5)
    ax.plot([], [], "o", color=style.CORAL, label="died")
    ax.plot([], [], "o", mfc="none", mec=style.BLUE, mew=1.5, label="alive when last seen")
    ax.legend(fontsize=8, loc="lower right")
    ax.set(xlabel="days after the transplant", ylabel="patient", yticks=[])


fig, ax = plt.subplots(figsize=(8, 4))
plot_timelines(ax)
plt.show()

Figure 9.4.1: Thirty of the patients, sorted by follow-up time. Each line runs from the transplant to death (filled dot) or to the last time the patient was seen alive (open circle).

The obvious summaries are both wrong. The average of all 69 times, 377 days, counts the censored patients as if they had died on their last visit: too low. The average of only those who died, 223 days, drops the patients who did best: lower still. We need methods that count a censored patient as alive for as long as we watched them, and no longer.

These patients are right-censored: the event, if it comes, comes after the last observation. Less common are left censoring (the event happened before the first observation, at an unknown time) and interval censoring (it happened between two check-ups). All methods here rest on one assumption: censoring is independent of the outcome. A patient who leaves the study must have the same outlook as those who stay; if the sickest patients drop out to a hospice, the survival estimate will look too good.

9.4.2 Survival and hazard

Two functions describe the time T until the event:

the survival function $S(t) = P(T > t)$ : the share still without the event at time t;
the hazard function $h(t)$ : the risk of the event at time t for those still without it, per unit of time. It is the instantaneous rate: roughly, the chance of dying in the next day, given alive today.

They are tied together through the cumulative hazard $H(t) = \int_0^t h(u)\,du$ , the hazard added up over time:

$S(t) = e^{-H(t)}$

The shape of the hazard is often the most telling thing about a process:

Code

t = np.linspace(0.01, 3, 300)
SHAPES = [
    ("constant (k = 1)", 1.0, style.BLUE),
    ("rising (k = 2.5)", 2.5, style.CORAL),
    ("falling (k = 0.5)", 0.5, style.GREEN),
]

fig, axes = plt.subplots(2, 3, figsize=(10.5, 4.8), sharex=True)
for col, (name, k, color) in enumerate(SHAPES):
    scale = 1 / np.log(2) ** (1 / k)  # same median, 1, for every shape
    hazard = (k / scale) * (t / scale) ** (k - 1)
    survival_curve = np.exp(-((t / scale) ** k))
    axes[0, col].plot(t, hazard, color=color, lw=2)
    axes[0, col].set(title=name, ylim=(0, 4))
    axes[1, col].plot(t, survival_curve, color=color, lw=2)
    axes[1, col].yaxis.set_major_formatter(style.PERCENT)
    axes[1, col].set(xlabel="time", ylim=(0, 1.02))
axes[0, 0].set_ylabel("hazard h(t)")
axes[1, 0].set_ylabel("survival S(t)")
fig.tight_layout()
plt.show()

Figure 9.4.2: Three hazard shapes (top) and the survival curves they produce (bottom), from Weibull distributions with the same median. Constant hazard (the exponential distribution, Chapter 2.1): risk never changes, like radioactive decay. Rising hazard: wear and ageing. Falling hazard: early risk that fades, like infant mortality or the weeks after surgery.

9.4.3 The Kaplan-Meier estimate

The Kaplan-Meier estimate builds the survival curve one death at a time. Sort the distinct times $t_j$ at which deaths occur. At each, let $n_j$ be the number still at risk (alive and still followed) and $d_j$ the number who die then. Surviving past $t_j$ , given alive just before it, has probability $1 - d_j / n_j$ , and these multiply:

$\hat S(t) = \prod_{t_j \le t} \left(1 - \frac{d_j}{n_j}\right)$

A censored patient counts in $n_j$ up to their last visit, and then quietly leaves the count. The first steps for the transplant patients:

Code

survival = SurvfuncRight(days, died)
first = survival.summary().head(8).reset_index()
first.columns = ["day", "survival", "standard error", "at risk", "deaths"]
first["survival"] = first["survival"].map("{:.3f}".format)
first["standard error"] = first["standard error"].map("{:.3f}".format)
first[["day", "at risk", "deaths"]] = first[["day", "at risk", "deaths"]].astype(int)
print(first.to_markdown(index=False, colalign=["right"] * 5))

Table 9.4.1: The first steps of the Kaplan-Meier estimate for the 69 transplant patients.

day	survival	standard error	at risk	deaths
1	0.971	0.02	69	2
3	0.956	0.025	66	1
10	0.942	0.028	65	1
12	0.927	0.031	64	1
14	0.912	0.034	62	1
15	0.897	0.037	61	1
23	0.882	0.039	60	1
25	0.867	0.041	59	1

Two patients died on day 1: 1 − 2/69 = 0.971. On day 3, one of the 66 still at risk died: 0.971 × (1 − 1/66) = 0.956. And so on. The standard errors come from Greenwood’s formula.

Code

median = survival.quantile(0.5)

fig, ax = plt.subplots(figsize=(8, 3.6))
plot_curve(ax, survival, style.BLUE)
ax.plot([0, median, median], [0.5, 0.5, 0], color=style.INK, ls="--", lw=1)
style.note(ax, f"median survival: {median:.0f} days", x=0.35, y=0.62, color=style.INK)
ax.yaxis.set_major_formatter(style.PERCENT)
ax.set(xlabel="days after the transplant", ylabel="share still alive", ylim=(0, 1.02))
plt.show()

Figure 9.4.3: The Kaplan-Meier estimate of survival after a heart transplant, with a 95% band. Each step down is a death; ticks mark censored patients. The dashed lines read off the median survival: the time by which half the patients had died.

Half the patients died within about 230 days. Most deaths came in the first months; after that the curve falls more slowly. The band widens to the right, where few patients are still being followed.

9.4.4 The cumulative hazard

The hazard itself is hard to estimate directly, but its running total is easy. The Nelson-Aalen estimate adds up the deaths per patient at risk:

$\hat H(t) = \sum_{t_j \le t} \frac{d_j}{n_j}$

Its slope is the hazard: steep where many die relative to those at risk, flat where few do.

Code

death_times = np.sort(np.unique(days[died == 1]))
at_risk = np.array([(days >= t_j).sum() for t_j in death_times])
deaths = np.array([((days == t_j) & (died == 1)).sum() for t_j in death_times])
cumulative = np.cumsum(deaths / at_risk)


def plot_nelson_aalen(ax: Axes) -> None:
    """The cumulative hazard as steps, with a constant-hazard line for comparison."""
    ax.step(np.r_[0, death_times], np.r_[0, cumulative], where="post", color=style.PURPLE, lw=2)
    by_1000 = cumulative[death_times <= 1000][-1]
    ax.plot([0, 1000], [0, by_1000], color=style.INK, ls="--", lw=1)
    ax.set(xlabel="days after the transplant", ylabel="cumulative hazard H(t)")


fig, ax = plt.subplots(figsize=(8, 3.2))
plot_nelson_aalen(ax)
plt.show()

Figure 9.4.4: The Nelson-Aalen cumulative hazard of the transplant patients. It rises steeply in the first two months and much more slowly after: the risk of death is highest just after the operation and then falls, like the falling hazard of Figure 2.5.4.2. The dashed line is what a constant hazard with the same total over the first 1,000 days would look like.

9.4.5 Comparing groups: the log-rank test

Did older patients do worse? Split the patients at the median age, 48, decided before looking at the curves, and draw one curve per group:

Code

older = (age >= 48).astype(int)
statistic, p_value = survdiff(days, died, older)


def plot_groups(ax: Axes) -> None:
    """One Kaplan-Meier curve per age group, with the log-rank p-value."""
    for group, color, label in ((0, style.GREEN, "under 48"), (1, style.CORAL, "48 or older")):
        plot_curve(
            ax,
            SurvfuncRight(days[older == group], died[older == group]),
            color,
            label,
            band=False,
        )
    style.note(ax, f"log-rank p = {p_value:.2f}", x=0.5, y=0.8, color=style.INK)
    ax.yaxis.set_major_formatter(style.PERCENT)
    ax.legend(fontsize=8)
    ax.set(
        title="Two age groups",
        xlabel="days after the transplant",
        ylabel="share still alive",
        ylim=(0, 1.02),
    )


def plot_observed_minus_expected(ax: Axes) -> None:
    """Running total of observed minus expected deaths in the older group."""
    difference = []
    for t_j in death_times:
        risk_all = days >= t_j
        risk_old = risk_all & (older == 1)
        d_all = ((days == t_j) & (died == 1)).sum()
        d_old = ((days == t_j) & (died == 1) & (older == 1)).sum()
        difference.append(d_old - d_all * risk_old.sum() / risk_all.sum())
    ax.step(
        np.r_[0, death_times],
        np.r_[0, np.cumsum(difference)],
        where="post",
        color=style.CORAL,
        lw=2,
    )
    ax.axhline(0, color=style.INK, lw=0.8)
    ax.set(
        title="Observed minus expected",
        xlabel="days after the transplant",
        ylabel="deaths, older group",
    )


fig, axes = style.subplots(2, height=3.3, width=10)
plot_groups(axes[0])
plot_observed_minus_expected(axes[1])
fig.tight_layout()
plt.show()

Figure 9.4.5: Left: Kaplan-Meier curves for patients under 48 and 48 or older at the transplant. Right: the log-rank test at work. At each death, the deaths expected in the older group if the groups had the same hazard; the line adds up observed minus expected deaths in the older group over time. It drifts upward: the older patients die more often than expected, but with 69 patients the drift is within what chance could give.

At each death time $t_j$ , if the groups had the same hazard, the deaths would split in proportion to the numbers at risk: group 1 would expect $E_j = d_j\, n_{1j} / n_j$ . The log-rank test adds up observed minus expected over all death times and compares it with its variance V:

$\chi^2 = \frac{\left(\sum_j (O_{1j} - E_{1j})\right)^2}{V}$

which follows a chi-square distribution with one degree of freedom when there is no difference. The same yardstick appears if the age labels are shuffled among the patients, which destroys any real link between age and survival:

Code

observed_chi2 = survdiff(days, died, older)[0]
null_logrank = np.array([survdiff(days, died, rng.permutation(older))[0] for _ in range(2000)])
fig, ax = plt.subplots(figsize=(8, 3.1))
style.null_panel(
    ax,
    observed_chi2,
    samples=null_logrank,
    dist=stats.chi2(1),
    tail="right",
    xlabel="log-rank statistic",
)
ax.set_xlim(0, 10)
plt.show()

Figure 9.4.6: The log-rank statistic for 2,000 shufflings of the ‘under 48’ and ‘48 or older’ labels among the 69 patients (grey), and the chi-square curve with 1 degree of freedom (blue). The observed statistic (dashed) is unremarkable among them: about one shuffle in nine gives a value at least as large, p = 0.11.

Here p = 0.11.

It is tempting to try other cut points until one gives a small p-value. Splitting at 50 instead of 48 gives p = 0.001; at 45, p = 0.27. That is the hidden multiple testing of Chapter 3.10. The better answer is not to cut at all, but to model age as a number: Cox regression, Chapter 9.5.

9.4.6 Parametric survival models

Kaplan-Meier assumes no shape. A parametric model fits a distribution for the survival times, such as the exponential (constant hazard) or the Weibull (hazard that rises or falls with its shape parameter k), by maximum likelihood. Censored cases enter the likelihood through the probability of surviving past their last visit, $S(t_i)$ , instead of the density $f(t_i)$ :

$\ell = \sum_{\text{deaths}} \log f(t_i) + \sum_{\text{censored}} \log S(t_i)$

Code

def weibull_log_likelihood(params: np.ndarray) -> float:
    """Negative log-likelihood of a Weibull with censoring; params are log k and log scale."""
    k, scale = np.exp(params)
    z = days / scale
    log_density = np.log(k / scale) + (k - 1) * np.log(z) - z**k
    log_survival = -(z**k)
    return -float(np.sum(np.where(died == 1, log_density, log_survival)))


weibull_k, weibull_scale = np.exp(
    optimize.minimize(weibull_log_likelihood, x0=[0, np.log(500)]).x
)
exponential_rate = died.sum() / days.sum()  # deaths per day at risk: the exponential's estimate
grid = np.linspace(0, 1800, 400)

fig, ax = plt.subplots(figsize=(8, 3.4))
plot_curve(ax, survival, style.BLUE, "Kaplan-Meier", band=False)
ax.plot(
    grid,
    np.exp(-exponential_rate * grid),
    color=style.INK,
    ls="--",
    lw=1.6,
    label="exponential",
)
ax.plot(
    grid,
    np.exp(-((grid / weibull_scale) ** weibull_k)),
    color=style.CORAL,
    lw=2,
    label=f"Weibull, k = {weibull_k:.2f}",
)
ax.yaxis.set_major_formatter(style.PERCENT)
ax.legend(fontsize=8)
ax.set(xlabel="days after the transplant", ylabel="share still alive", ylim=(0, 1.02))
plt.show()

Figure 9.4.7: The Kaplan-Meier estimate (steps) with an exponential fit (dashed, constant hazard) and a Weibull fit (red), both by maximum likelihood with censoring. The Weibull’s shape k is well below 1, a falling hazard, and it follows the early drop that the exponential misses.

Parametric models give smooth curves, extrapolate (with care) beyond the data, and need less data when their shape is right. The exponential’s estimate has a neat form: total deaths divided by total time at risk, here one death per 578 patient-days.

9.4.7 Beyond death

“Survival” is a historical name. The same methods handle any time until an event, when the event has not happened for everyone yet: time until a machine fails, a customer cancels, a loan defaults, a student finds a job. Wherever the question is “how long until?”, look for censoring.

9.4.8 In one picture

Censored times are lower limits; censoring must be independent of the outcome.
$S(t) = P(T > t)$ ; the hazard is the risk at t for those still at risk; $S(t) = e^{-H(t)}$ .
Kaplan-Meier: $\hat S(t) = \prod (1 - d_j/n_j)$ ; censored cases leave the risk set quietly.
Nelson-Aalen: $\hat H(t) = \sum d_j / n_j$ ; its slope is the hazard.
The log-rank test adds observed minus expected deaths over time; do not hunt for cut points.
Parametric models (exponential, Weibull) fit a shape, with censored cases entering through $S(t)$ .

Next: regression for survival times.