9.5 Regression for survival times

The Cox model says that a predictor multiplies the hazard by a fixed factor, at every point in time, while leaving the shape of the hazard over time completely free. That factor, the hazard ratio, is the survival version of the odds ratio.

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 statsmodels.duration.survfunc import SurvfuncRight

from statly import style

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

heart = sm.datasets.heart.load_pandas().data  # 69 heart transplant patients, Stanford
days = heart.survival.to_numpy()
died = heart.censors.to_numpy().astype(int)
decades = ((heart.age.to_numpy() - 50) / 10)[:, None]  # age in decades, 0 at age 50

cox = sm.PHReg(days, decades, status=died).fit()
beta = cox.params[0]


def concordance(times: np.ndarray, events: np.ndarray, risk: np.ndarray) -> float:
    """Harrell's C: among usable pairs, the share where the higher risk had the event first."""
    agree = usable = 0.0
    for i in np.flatnonzero(events == 1):
        later = times > times[i]  # i had the event while these were still event free
        usable += later.sum()
        agree += (risk[i] > risk[later]).sum() + 0.5 * (risk[i] == risk[later]).sum()
    return agree / usable

9.5.1 Proportional hazards

The log-rank test of Chapter 9.4 compares groups, but age is a number, and cutting it into groups throws information away and invites cut-point hunting. Cox regression models the hazard of a patient with predictors $x_1, \dots, x_p$ as

$h(t \mid x) = h_0(t)\, e^{\beta_1 x_1 + \dots + \beta_p x_p}$

The baseline hazard $h_0(t)$ can have any shape at all: rising, falling, bumpy. The predictors only scale it up or down. Two patients whose predictors differ by one unit in $x_1$ have hazards in the fixed ratio $e^{\beta_1}$ at every time: the hazard ratio. On a log scale the hazard curves are parallel, shifted by $\beta_1$ :

Code

t = np.linspace(0, 10, 400)
baseline = (
    0.3 * np.exp(-0.4 * t) + 0.05 + 0.08 * np.exp(-((t - 6) ** 2))
)  # early risk, later bump

fig, axes = style.subplots(2, height=3, width=10)
for ax, scale in zip(axes, ("linear", "log")):
    ax.plot(t, baseline, color=style.BLUE, lw=2, label="baseline")
    ax.plot(t, 2 * baseline, color=style.CORAL, lw=2, label="hazard ratio 2")
    ax.set(yscale=scale, xlabel="time", ylabel="hazard", title=f"{scale} scale")
axes[0].legend(fontsize=8)
fig.tight_layout()
plt.show()

Figure 9.5.1: Proportional hazards. Left: a bumpy baseline hazard (blue) and the hazard of a patient with a hazard ratio of 2 (red): at every time, twice as high. Right: the same on a log scale, where the two curves are parallel, a constant distance log 2 apart. The model fixes this ratio, not the shape.

Because $S(t) = e^{-H(t)}$ , a hazard ratio of r turns the survival curve into $S_0(t)^r$ : a patient with twice the hazard has, at every time, the square of the baseline survival.

9.5.2 Partial likelihood

How can β be estimated without knowing $h_0(t)$ ? Cox’s idea: look only at the moments when someone dies. At such a moment, given that exactly one of the patients still at risk dies, the chance that it is patient i is their share of the total hazard:

$\frac{h_0(t)\, e^{\beta x_i}}{\sum_{j \text{ at risk}} h_0(t)\, e^{\beta x_j}} = \frac{e^{\beta x_i}}{\sum_{j \text{ at risk}} e^{\beta x_j}}$

The baseline cancels. Multiplying these shares over all deaths gives the partial likelihood, and β is chosen to maximise it, just like a likelihood in Chapter 2.3. Only the order of the events matters, not their exact times.

Code

death_order = np.argsort(np.where(died == 1, days, np.inf))
example = death_order[10]  # the eleventh death
at_risk = np.flatnonzero(days >= days[example])
ages = heart.age.to_numpy()
order = at_risk[np.argsort(ages[at_risk])]


def plot_shares(ax: Axes) -> None:
    """Hazard shares of the patients at risk at one death, under β = 0 and the fitted β."""
    for b, color, label, offset in (
        (0.0, style.GREY, "β = 0", -0.2),
        (beta, style.BLUE, f"β = {beta:.2f}", 0.2),
    ):
        weight = np.exp(b * decades[order, 0])
        share = weight / weight.sum()
        colors = [style.CORAL if i == example else color for i in order]
        ax.bar(
            np.arange(len(order)) + offset,
            share,
            width=0.4,
            color=colors,
            alpha=0.5 if b == 0 else 0.9,
            label=label,
        )
    ax.yaxis.set_major_formatter(style.PERCENT)
    ax.legend(fontsize=8)
    ax.set(
        xlabel="patients at risk, youngest to oldest",
        ylabel="share of the hazard",
        xticks=[],
        title=f"Death on day {days[example]:.0f}",
    )


def plot_profile(ax: Axes) -> None:
    """The log partial likelihood as β varies."""
    grid = np.linspace(-0.4, 1.5, 200)
    log_likelihood = np.array([cox.model.loglike(np.array([b])) for b in grid])
    ax.plot(grid, log_likelihood, color=style.PURPLE, lw=2)
    ax.axvline(beta, color=style.INK, ls="--", lw=1)
    style.note(
        ax, f"β = {beta:.2f}\nhazard ratio {np.exp(beta):.2f}", x=0.62, y=0.3, color=style.INK
    )
    ax.set(xlabel="β (per decade of age)", ylabel="log partial likelihood", title="Choosing β")


fig, axes = style.subplots(2, height=3.3, width=10.5)
plot_shares(axes[0])
plot_profile(axes[1])
fig.tight_layout()
plt.show()

Figure 9.5.2: Left: one death in the transplant data. Each bar is a patient still at risk at that moment, its height their share of the total hazard under a given β. With β = 0 everyone has the same share; with the fitted β, older patients get larger shares, and the patient who actually died (red) gets a larger one. Right: the log partial likelihood over β. Its peak is the estimate, and its curvature gives the standard error.

9.5.3 Reading the model

Code

strikes = sm.datasets.strikes.load_pandas().data
strike_fit = sm.PHReg(
    strikes.duration.to_numpy(), strikes.iprod.to_numpy()[:, None] * 100
).fit()

rows = []
for name, fit in (
    ("heart: age, per decade", cox),
    ("strikes: production, per point", strike_fit),
):
    b, se, p = fit.params[0], fit.bse[0], fit.pvalues[0]
    rows.append(
        f"| {name} | {b:.3f} | {se:.3f} | {np.exp(b):.2f} | "
        f"{np.exp(b - 1.96 * se):.2f} to {np.exp(b + 1.96 * se):.2f} | {p:.3f} |"
    )
print("| model | β | standard error | hazard ratio | 95% interval | p-value |")
print("|:--|--:|--:|--:|--:|--:|")
print("\n".join(rows))

Table 9.5.1: Two Cox models. Heart transplants: the hazard of death per decade of age at the transplant. Strikes: the hazard of a strike ending, per percentage point of industrial production (62 strikes in the United States, none censored).

model	β	standard error	hazard ratio	95% interval	p-value
heart: age, per decade	0.545	0.226	1.72	1.11 to 2.68	0.016
strikes: production, per point	0.091	0.032	1.10	1.03 to 1.17	0.005

Each extra decade of age multiplies the hazard of death by 1.7 (95% interval 1.1 to 2.7, p = 0.016): the age effect the log-rank test, cut at 48, could not show. For strikes, the “event” is the strike ending, so a hazard ratio above 1 means shorter strikes: each extra point of industrial production makes a strike end about 10% faster. As in logistic regression, read β on the log scale and $e^\beta$ as a factor; with several predictors, each hazard ratio holds the others fixed.

The model also gives a survival curve for any patient. The software estimates the baseline cumulative hazard $\hat H_0(t)$ , and then $\hat S(t \mid x) = e^{-\hat H_0(t)\, e^{\hat\beta x}}$ :

Code

times, cumulative_hazard, _ = cox.baseline_cumulative_hazard[0]

fig, ax = plt.subplots(figsize=(8, 3.4))
for years, color in ((30, style.GREEN), (45, style.BLUE), (55, style.CORAL)):
    risk = np.exp(beta * (years - 50) / 10)
    ax.step(
        np.r_[0, times],
        np.r_[1, np.exp(-cumulative_hazard * risk)],
        where="post",
        color=color,
        lw=2,
        label=f"age {years}",
    )
ax.yaxis.set_major_formatter(style.PERCENT)
ax.legend(fontsize=8)
ax.set(xlabel="days after the transplant", ylabel="predicted share alive", ylim=(0, 1.02))
plt.show()

Figure 9.5.3: Predicted survival after a heart transplant for patients aged 30, 45 and 55, from the Cox model. The curves have the same steps, at the observed death times, but older patients’ curves fall further at each step: each is the baseline curve raised to the power of its hazard ratio.

9.5.4 Checking proportional hazards

The whole model rests on the hazard ratio staying constant over time. Two pictures check it.

Log-minus-log curves. Since $\log(-\log S(t)) = \log H_0(t) + \beta x$ , plotting $\log(-\log \hat S(t))$ from Kaplan-Meier for each group should give parallel curves.

Schoenfeld residuals. At each death, compare the age of the patient who died with the average age of those at risk, weighted by their hazard shares. If the hazard ratio is constant, these residuals scatter around zero with no trend over time; a trend means the effect of age grows or fades.

Code

older = (heart.age.to_numpy() >= 48).astype(int)


def plot_log_minus_log(ax: Axes) -> None:
    """log(-log S) from Kaplan-Meier, one curve per age group, on a log time axis."""
    for group, color, label in ((0, style.GREEN, "under 48"), (1, style.CORAL, "48 or older")):
        fit = SurvfuncRight(days[older == group], died[older == group])
        keep = (fit.surv_prob > 0) & (fit.surv_prob < 1)
        ax.step(
            fit.surv_times[keep],
            np.log(-np.log(fit.surv_prob[keep])),
            where="post",
            color=color,
            lw=2,
            label=label,
        )
    ax.legend(fontsize=8)
    ax.set(
        xscale="log",
        xlabel="days after the transplant (log scale)",
        ylabel="log(−log S(t))",
        title="Log-minus-log",
    )


def plot_schoenfeld(ax: Axes) -> None:
    """Schoenfeld residuals for age against the time of death, with a lowess line."""
    residual = cox.schoenfeld_residuals[:, 0]
    keep = ~np.isnan(residual)
    ax.scatter(days[keep], residual[keep], s=18, color=style.BLUE, alpha=0.7)
    smooth = sm.nonparametric.lowess(residual[keep], days[keep], frac=0.6)
    ax.plot(smooth[:, 0], smooth[:, 1], color=style.CORAL, lw=2)
    ax.axhline(0, color=style.INK, lw=0.8)
    ax.set(xlabel="day of death", ylabel="residual (decades)", title="Schoenfeld residuals")


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

Figure 9.5.4: Two checks of proportional hazards for age in the transplant data. Left: log-minus-log Kaplan-Meier curves for patients under 48 and 48 or older; they cross in the first two months, where each group has only a few deaths, and run roughly parallel after that. Right: Schoenfeld residuals for age (in decades) at each death, with a smoothed line. The line stays near zero: no clear sign that the age effect changes with time.

If the hazards are clearly not proportional, there are remedies: stratify on the offending variable (each stratum gets its own baseline hazard, and the variable gets no coefficient), let the coefficient change with time through an interaction with time, or report hazard ratios separately for early and late follow-up.

9.5.5 Predictors that change over time

Some predictors are not known at the start. In the Stanford programme, patients were enrolled and then waited for a donor heart; only those who survived the wait received one. Comparing “patients who got a transplant” with “those who did not” as if the transplant were known at enrolment is a classic error called immortal time bias: the transplanted patients were guaranteed to survive the waiting time, so they look better even if the transplant does nothing.

The fix is to treat the transplant as a time-varying covariate. Each transplanted patient contributes two pieces of follow-up: untreated from enrolment to the transplant, and treated from then on. In the partial likelihood, at each death, a patient counts as treated only if they had been treated by that moment.

Code

n = 2000
death = rng.exponential(300, n)
wait = rng.exponential(200, n)
treated = wait < death  # the treatment only happens if the patient lives long enough

naive = sm.PHReg(death, treated[:, None].astype(float)).fit()

# Split treated patients into an untreated piece (0 to wait) and a treated piece (wait to death).
start = np.r_[np.zeros(n), wait[treated]]
stop = np.r_[np.where(treated, wait, death), death[treated]]
event = np.r_[~treated, np.ones(treated.sum(), bool)].astype(int)
covariate = np.r_[np.zeros(n), np.ones(treated.sum())]
varying = sm.PHReg(stop, covariate[:, None], status=event, entry=start).fit()


def plot_naive_curves(ax: Axes) -> None:
    """Kaplan-Meier curves for treated and untreated, grouped as if known at enrolment."""
    for group, color, label in (
        (True, style.GREEN, "treated"),
        (False, style.GREY, "not treated"),
    ):
        fit = SurvfuncRight(death[treated == group], np.ones((treated == group).sum()))
        ax.step(
            np.r_[0, fit.surv_times],
            np.r_[1, fit.surv_prob],
            where="post",
            color=color,
            lw=2,
            label=label,
        )
    ax.yaxis.set_major_formatter(style.PERCENT)
    ax.legend(fontsize=8)
    ax.set(
        xlim=(0, 1500),
        xlabel="days after enrolment",
        ylabel="share alive",
        title="Grouped by eventual treatment",
    )


def plot_ratios(ax: Axes) -> None:
    """Hazard ratio with 95% interval under the two codings."""
    for row, (fit, label, color) in enumerate(
        ((naive, "fixed at enrolment", style.CORAL), (varying, "time-varying", style.BLUE))
    ):
        b, se = fit.params[0], fit.bse[0]
        ax.plot(np.exp([b - 1.96 * se, b + 1.96 * se]), [row, row], color=color, lw=3)
        ax.plot(np.exp(b), row, "o", ms=9, color=color)
        ax.text(np.exp(b), row + 0.18, f"{np.exp(b):.2f}", ha="center", color=color, fontsize=9)
    ax.axvline(1, color=style.INK, ls="--", lw=1)
    ax.set(
        yticks=[0, 1],
        yticklabels=["fixed at\nenrolment", "time-\nvarying"],
        ylim=(-0.6, 1.6),
        xlim=(0, 1.4),
        xlabel="hazard ratio for treatment",
        title="Estimated effect",
    )


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

Figure 9.5.5: Immortal time bias in a simulation where treatment has no effect at all. Every patient has the same constant hazard of death; treatment arrives after a random wait, if the patient is still alive. Left: treated patients seem to live far longer. Right: the estimated hazard ratio for treatment. Coded as fixed at enrolment it suggests a strong benefit; coded as a time-varying covariate it is close to the true value, 1.

The same trap appears whenever a group is defined by something that happens during follow-up: “patients who completed treatment”, “employees who were promoted”, “users who reached level 10”. A related choice is the time scale. Time since enrolment is usual, but sometimes age is the natural clock, with patients entering the risk set at the age they joined the study.

9.5.6 How good are the predictions?

R² does not work with censored data. The usual summary is the concordance index, C, the survival version of the area under the ROC curve (Chapter 6.1). Take every pair of patients where we know who had the event first: the first must have had the event, and the other must still have been followed at that time. C is the share of these pairs where the model gave the first patient the higher risk:

$C = P\left(\text{risk}_i > \text{risk}_j \mid T_i < T_j,\ i \text{ had the event}\right)$

C = 0.5 is a coin toss and 1 is a perfect ranking.

Code

c_index = concordance(days, died, decades[:, 0])
first, second = [], []
for i in np.flatnonzero(died == 1):
    for j in np.flatnonzero(days > days[i]):
        first.append(ages[i])
        second.append(ages[j])
first, second = np.array(first), np.array(second)
right = first > second

fig, ax = plt.subplots(figsize=(5.2, 5))
jitter = rng.uniform(-0.3, 0.3, (2, len(first)))
ax.scatter(
    first[right] + jitter[0, right],
    second[right] + jitter[1, right],
    s=4,
    color=style.BLUE,
    alpha=0.35,
)
ax.scatter(
    first[~right] + jitter[0, ~right],
    second[~right] + jitter[1, ~right],
    s=4,
    color=style.CORAL,
    alpha=0.35,
)
ax.plot([8, 65], [8, 65], color=style.INK, lw=1)
style.note(ax, f"C = {c_index:.2f}", x=0.05, y=0.92, color=style.INK)
ax.set(
    xlabel="age of the patient who died first",
    ylabel="age of the other patient",
    aspect="equal",
)
plt.show()

Figure 9.5.6: Concordance for the age model on the transplant data. Each dot is a usable pair of patients: the one who died first (across) against the other (up), placed by age. Pairs above the diagonal (red) are ranked wrongly, the younger one died first; pairs below (blue) are ranked rightly. The share of blue pairs is C.

C = 0.61 here: age ranks patients better than chance, but far from perfectly. As with any prediction measure, it should be computed on patients the model was not fitted to (Chapter 6.3). A time-dependent version of the ROC curve, asking “who will have died by day t?”, is also common.

9.5.7 Many predictors: shrinkage and trees

With many predictors the tools of Part 6 carry over. Penalised Cox regression maximises the log partial likelihood minus a ridge or lasso penalty,

$\log PL(\beta) - \lambda \sum_k |\beta_k|$

with λ chosen by cross-validation on, for example, the partial likelihood or C.

Code

n_patients, n_predictors = 150, 20
predictors = rng.normal(size=(n_patients, n_predictors))
true_beta = np.r_[0.8, -0.6, 0.5, np.zeros(n_predictors - 3)]
event_time = rng.exponential(np.exp(-predictors @ true_beta))
censor_time = rng.exponential(1.5, n_patients)
observed = np.minimum(event_time, censor_time)
had_event = (event_time <= censor_time).astype(int)

lasso_model = sm.PHReg(observed, predictors, status=had_event)
penalties = np.geomspace(0.3, 0.002, 25)
path = np.array([lasso_model.fit_regularized(alpha=a, L1_wt=1.0).params for a in penalties])

fig, ax = plt.subplots(figsize=(8, 3.6))
for k in range(n_predictors):
    color = (style.BLUE, style.CORAL, style.GREEN)[k] if k < 3 else style.GREY
    ax.plot(penalties, path[:, k], color=color, lw=2 if k < 3 else 1, alpha=1 if k < 3 else 0.6)
ax.axhline(0, color=style.INK, lw=0.8)
ax.set(
    xscale="log",
    xlim=(penalties[0], penalties[-1]),
    xlabel="penalty λ (log scale, falling to the right)",
    ylabel="coefficient β",
)
plt.show()

Figure 9.5.7: A lasso Cox path on simulated data: 150 patients, 20 predictors, of which only the first three (coloured) affect the hazard. As the penalty λ falls (to the right), predictors enter the model one by one; the three real ones enter first and grow largest, while the others stay near zero.

Survival trees split the patients by the predictor and cut point that best separate the survival curves, for example by the largest log-rank statistic, and repeat within each part. As in Chapter 5.5, one tree is easy to read but unstable; random survival forests average many trees grown on bootstrap samples, and handle nonlinear effects and interactions without being told about them.

9.5.8 Competing risks

Sometimes a patient can experience one of several events, and the first one ends the story: death from heart failure or death from another cause, a customer who cancels or one who upgrades. Treating the other events as censoring and using 1 − Kaplan-Meier then overstates the chance of the event of interest, because it pretends that patients removed by the other event could still have had it. The cumulative incidence function counts each event only among those still free of all events:

$\text{CIF}_k(t) = \sum_{t_j \le t} \hat S(t_{j-1}) \frac{d_{kj}}{n_j}$

where $\hat S$ is the Kaplan-Meier estimate of being free of every event.

Code

n = 1000
cause_1 = rng.exponential(10, n)
cause_2 = rng.exponential(5, n)
time_of_death = np.minimum(cause_1, cause_2)
which = np.where(cause_1 < cause_2, 1, 2)

order = np.argsort(time_of_death)
sorted_times, sorted_causes = time_of_death[order], which[order]
still_at_risk = n - np.arange(n)
free_of_all = np.cumprod(1 - 1 / still_at_risk)
free_before = np.r_[1, free_of_all[:-1]]
incidence = {k: np.cumsum(free_before * (sorted_causes == k) / still_at_risk) for k in (1, 2)}
naive_fit = SurvfuncRight(time_of_death, (which == 1).astype(int))

fig, ax = plt.subplots(figsize=(8, 3.6))
ax.step(
    np.r_[0, naive_fit.surv_times],
    np.r_[0, 1 - naive_fit.surv_prob],
    where="post",
    color=style.BLUE,
    ls="--",
    lw=1.6,
    label="cause 1: 1 − Kaplan-Meier",
)
ax.step(
    sorted_times,
    incidence[1],
    where="post",
    color=style.BLUE,
    lw=2,
    label="cause 1: cumulative incidence",
)
ax.step(
    sorted_times,
    incidence[2],
    where="post",
    color=style.CORAL,
    lw=2,
    label="cause 2: cumulative incidence",
)
ax.yaxis.set_major_formatter(style.PERCENT)
ax.legend(fontsize=8, loc="center right")
ax.set(xlim=(0, 30), ylim=(0, 1.02), xlabel="time", ylabel="share who have died of the cause")
plt.show()

Figure 9.5.8: Competing risks in a simulation: two causes of death, each with a constant hazard, the second twice as strong. Treating the competing cause as censoring, 1 − Kaplan-Meier (dashed) keeps climbing, to almost 90% here, as if every patient would eventually die of cause 1. The cumulative incidence (solid) levels off near one third, the true share of patients who die of cause 1. The two cumulative incidences add up to the share who have died of anything.

Regression for competing risks either models each cause’s hazard separately (cause-specific Cox models) or models the cumulative incidence directly (the Fine-Gray model).

9.5.9 In one picture

Cox model: $h(t \mid x) = h_0(t) e^{\beta x}$ ; the baseline is free, the hazard ratio $e^\beta$ is fixed over time.
Partial likelihood: at each death, the dead patient’s share of the hazard; the baseline cancels.
Predicted survival: $S(t \mid x) = S_0(t)^{e^{\beta x}}$ .
Check proportional hazards with log-minus-log curves and Schoenfeld residuals; stratify if they fail.
Predictors that arrive during follow-up must be time-varying, or immortal time bias follows.
Concordance C is the survival AUC.
Lasso Cox and survival forests handle many predictors.
With competing risks, use cumulative incidence, not 1 − Kaplan-Meier.