Nonparametric Bayesian Survival

Python · PyMC · R  ·  Download survival module

A Prior on the Hazard, Not a Formula

Parametric survival analysis makes you commit to a hazard shape before you look — exponential for constant risk, Weibull for monotone wear-out, log-normal for a hump. The nonparametric alternative refuses the commitment: chop time into intervals, give each one its own free hazard, and tie them together with a prior instead of a formula. A Gamma process gives independent interval hazards; a random-walk prior smooths them by borrowing across neighbours. Either way the shape is inferred, and a Weibull's monotone assumption is never imposed.

h(t)=λk for t∈(ak−1,ak],λk∼Gamma-process,h(t∣x)=h0(t) ex⊤βh(t)=\lambda_k \text{ for } t\in(a_{k-1},a_k],\qquad \lambda_k\sim\text{Gamma-process},\qquad h(t\mid x)=h_0(t)\,e^{x^\top\beta}

Recovering a hazard whose shape is known

The construction is validated against a hazard whose truth is known before it is trusted on real data. Simulating 500 subjects from a linearly rising hazard, the recovered interval hazards track the truth at correlation 0.979, picking up the wear-out shape a Weibull would have had to assume — and the credible band widens exactly where events thin out in the far tail.

The Gehan trial

The Gehan leukemia trial is the natural test case: 42 patients, 30 relapses, 21 per arm, from the 1963 study that established 6-mercaptopurine. The Bayesian nonparametric survival curve sits directly on the Kaplan–Meier step function, and the posterior mean cumulative hazard traces Nelson–Aalen almost exactly. The notebooks are honest about what that means — with data this clean the prior earns its keep mainly by supplying a credible band, and would matter considerably more if the data were sparser.

Adding a proportional-hazards term recovers the trial's headline result from the nonparametric side: HR = 0.19, 95% CrI [0.08, 0.39] — 6-MP cuts the relapse hazard to about a fifth of placebo — while also delivering the full baseline hazard and survival curves that a Cox fit conditions away.

A Gap That Was a Convention

The two engines reported slightly different Cox hazard ratios — 0.22 [0.10, 0.49] in Python against R's 0.21 [0.09, 0.47] — a gap that is not a disagreement about the data but about a convention, and on this dataset it deserves saying because 18 of the 42 event times are tied. statsmodels defaults to the Breslow approximation to the partial likelihood; R's coxph defaults to Efron. Fitting both settles it: Efron gives 0.208 [0.093, 0.466], reproducing R exactly, and Breslow gives 0.221 [0.099, 0.493], reproducing the published Python figure. Both notebooks report the convention alongside the number.

hazard ratio for 6-MPestimate95% intervalnote
Bayesian nonparametric PH0.19[0.08, 0.39]Gamma-process baseline
Cox, Breslow ties0.221[0.099, 0.493]statsmodels default
Cox, Efron ties0.208[0.093, 0.466]coxph default — matches R exactly
PyMC piecewise-exponential0.24[0.11, 0.53]same likelihood, different prior

What "the same model" is and is not claiming

The PyMC cross-check reports 0.24 against the from-scratch 0.19 — a 25% gap in the point estimate — and a baseline-hazard difference that turns out to be 55% of the mean baseline level. Both are worth reading off the intervals rather than the point estimates: they overlap across 64% of their combined range, because with 30 events a hazard ratio is simply not pinned down to two significant figures. What is exact is the correspondence itself — a Gamma-process baseline with a proportional-hazards term and a piecewise-exponential Poisson GLM are the same likelihood written two ways. The residual differences are Monte Carlo error and prior choice, largest in the sparse late intervals where the prior does most of the work, and not a difference of model.

Notebooks

Python Recovery of a known hazard shape, Gamma-process and random-walk baselines on the Gehan trial, Kaplan–Meier and Nelson–Aalen as the frequentist counterparts, nonparametric proportional hazards, and a PyMC piecewise-exponential cross-check (Python) Open → R Independent R implementation via survival — Kaplan–Meier and Nelson–Aalen estimators, coxph with its Efron tie correction named explicitly, and the Breslow baseline hazard (R) Open →

Downloads

bnp_npsurv.py Piecewise-exponential expansion of censored survival data; a Gibbs sampler for Gamma-process and random-walk baseline hazards with proportional-hazards covariates; Kaplan–Meier and Nelson–Aalen estimators (NumPy / SciPy) gehan.csv Gehan leukemia remission data — 42 patients, 30 relapses, 6-mercaptopurine against placebo

Survival Module — Source Code

"""
npsurv.py -- NONPARAMETRIC BAYESIAN SURVIVAL: piecewise-exponential hazard (from scratch).

Backs the notebooks in  "Nonparametric Bayesian Survival".

The survival arc fitted PARAMETRIC hazards (Weibull/exponential) and the Cox model as a
piecewise-exponential Poisson GLM. Here the baseline hazard itself is NONPARAMETRIC: chop time
into K intervals and let the hazard be a free constant lambda_k on each, with a prior that either
leaves the intervals independent or ties neighbours together.

    h(t) = lambda_k   for t in interval k,     x_i shifts it proportionally:  h_i(t) = lambda_k e^{x_i'beta}

Written on person-time, the piecewise-exponential likelihood is exactly Poisson: with exposure
e_ik (time subject i spent in interval k) and event indicator d_ik,
        d_ik ~ Poisson( lambda_k e^{x_i'beta} e_ik ).

Two priors on the baseline:
  * GAMMA PROCESS (Kalbfleisch 1978) -- independent increments of the cumulative hazard, i.e.
    lambda_k ~ Gamma(a,b) independently. This is CONJUGATE: lambda_k | . ~ Gamma(a + D_k, b + E_k),
    with D_k the events and E_k = sum_i e_ik e^{x_i'beta} the (covariate-weighted) exposure in
    interval k. The nonparametric baseline is then a direct Gibbs draw.
  * SMOOTHED -- a random walk on log lambda_k (log lambda_k - log lambda_{k-1} ~ N(0,sigma^2)),
    which borrows strength across intervals for a smooth hazard; updated by Metropolis.

The regression coefficient beta (proportional hazards) is updated by random-walk Metropolis on the
Poisson log-likelihood. Survival is S(t|x) = exp(-H_0(t) e^{x'beta}) from the sampled hazard, with
credible bands. Kaplan-Meier and Nelson-Aalen (the frequentist nonparametric estimators) are here too.
"""

import numpy as np


# --------------------------------------------------------------------------- #
#  person-time expansion                                                        #
# --------------------------------------------------------------------------- #

def expand_pwe(time, status, cuts):
    """exposure e_ik and event d_ik matrices for the piecewise-exponential model.
    cuts: increasing interval boundaries starting at 0. Returns E, D each (N, K)."""
    time = np.asarray(time, float); status = np.asarray(status, int)
    N = len(time); K = len(cuts) - 1
    E = np.zeros((N, K)); D = np.zeros((N, K))
    for i in range(N):
        for k in range(K):
            lo, hi = cuts[k], cuts[k + 1]
            if time[i] <= lo:
                break
            E[i, k] = min(time[i], hi) - lo
            if time[i] <= hi:
                if status[i] == 1:
                    D[i, k] = 1.0
                break
    return E, D


# --------------------------------------------------------------------------- #
#  Gibbs sampler                                                                #
# --------------------------------------------------------------------------- #

def npsurv_gibbs(time, status, cuts, rng, X=None, a=0.2, b=0.2, beta_sd=10.0,
                 smooth=False, draws=3000, burn=1500, step=0.3):
    """Piecewise-exponential Bayesian survival with a nonparametric baseline hazard.
    smooth=False -> independent Gamma-process prior (conjugate); True -> random-walk-smoothed
    log-hazard. X:(N,p) optional covariates (proportional hazards). Returns posterior draws of the
    interval hazards lambda:(draws,K), coefficients beta:(draws,p), and the interval widths."""
    time = np.asarray(time, float); status = np.asarray(status, int)
    E, D = expand_pwe(time, status, cuts); N, K = E.shape
    Dk = D.sum(0); di = D.sum(1); width = np.diff(cuts)
    p = 0 if X is None else X.shape[1]
    X = np.zeros((N, 0)) if X is None else np.asarray(X, float)
    beta = np.zeros(p); acc = 0
    lam = np.full(K, max(Dk.sum() / E.sum(), 1e-3))
    gamma = np.log(lam); sigma2 = 0.5
    LAM = np.empty((draws, K)); BETA = np.empty((draws, p))
    for it in range(draws + burn):
        w = np.exp(X @ beta) if p else np.ones(N)
        Ek = (E * w[:, None]).sum(0)
        # ---- baseline hazard ----
        if not smooth:
            lam = rng.gamma(a + Dk, 1.0 / (b + Ek))
        else:
            gamma = _update_gamma(gamma, Dk, Ek, sigma2, rng, step)
            # smoothing variance (InvGamma on the RW increments)
            dg = np.diff(gamma)
            sigma2 = 1.0 / rng.gamma(1e-2 + (K - 1) / 2, 1.0 / (1e-2 + 0.5 * (dg @ dg)))
            lam = np.exp(gamma)
        # ---- proportional-hazards coefficient (Metropolis) ----
        if p:
            EL = E @ lam
            def ll(bb):
                eta = X @ bb
                return np.sum(di * eta - np.exp(eta) * EL) - 0.5 * (bb @ bb) / beta_sd ** 2
            prop = beta + step * rng.standard_normal(p)
            if np.log(rng.random()) < ll(prop) - ll(beta):
                beta = prop; acc += 1
        if it >= burn:
            LAM[it - burn] = lam; BETA[it - burn] = beta
    return dict(lam=LAM, beta=BETA, width=width, cuts=np.asarray(cuts, float),
                accept=acc / (draws + burn))


def _update_gamma(gamma, Dk, Ek, sigma2, rng, step):
    """per-interval Metropolis for log-hazard under a random-walk prior."""
    K = len(gamma); g = gamma.copy()
    for k in range(K):
        prop = g[k] + step * rng.standard_normal()
        def lp(v, idx):
            like = Dk[idx] * v - np.exp(v) * Ek[idx]
            pri = 0.0
            if idx > 0: pri += -0.5 * (v - g[idx - 1]) ** 2 / sigma2
            if idx < K - 1: pri += -0.5 * (g[idx + 1] - v) ** 2 / sigma2
            return like + pri
        if np.log(rng.random()) < lp(prop, k) - lp(g[k], k):
            g[k] = prop
    return g


# --------------------------------------------------------------------------- #
#  posterior survival / hazard on a grid                                        #
# --------------------------------------------------------------------------- #

def survival_curves(res, tgrid, x=None):
    """posterior survival S(t|x) on tgrid from the sampled hazards. Returns (draws, len(tgrid))."""
    cuts = res["cuts"]; LAM = res["lam"]; draws, K = LAM.shape
    tgrid = np.asarray(tgrid, float)
    # cumulative baseline hazard at each grid time
    H = np.zeros((draws, len(tgrid)))
    for j, t in enumerate(tgrid):
        contrib = np.clip(np.minimum(t, cuts[1:]) - cuts[:-1], 0, None)   # time in each interval up to t
        H[:, j] = LAM @ contrib
    if x is not None and res["beta"].shape[1] > 0:
        H = H * np.exp(res["beta"] @ np.asarray(x, float))[:, None]
    return np.exp(-H)


def cum_hazard(res, tgrid):
    cuts = res["cuts"]; LAM = res["lam"]; tgrid = np.asarray(tgrid, float)
    H = np.zeros((LAM.shape[0], len(tgrid)))
    for j, t in enumerate(tgrid):
        H[:, j] = LAM @ np.clip(np.minimum(t, cuts[1:]) - cuts[:-1], 0, None)
    return H


# --------------------------------------------------------------------------- #
#  frequentist nonparametric estimators + simulation                            #
# --------------------------------------------------------------------------- #

def kaplan_meier(time, status):
    time = np.asarray(time, float); status = np.asarray(status, int)
    ts = np.sort(np.unique(time[status == 1])); S = np.empty(len(ts)); s = 1.0
    for i, t in enumerate(ts):
        d = np.sum((time == t) & (status == 1)); n = np.sum(time >= t)
        s *= (1 - d / n); S[i] = s
    return ts, S

def nelson_aalen(time, status):
    time = np.asarray(time, float); status = np.asarray(status, int)
    ts = np.sort(np.unique(time[status == 1])); H = np.empty(len(ts)); h = 0.0
    for i, t in enumerate(ts):
        d = np.sum((time == t) & (status == 1)); n = np.sum(time >= t)
        h += d / n; H[i] = h
    return ts, H

def simulate_survival(n, hazard, tmax, censor_time, rng):
    """simulate event times from a hazard function h(t) by inversion of the cumulative hazard on a
    fine grid, with administrative censoring at censor_time. hazard: callable t->h(t)."""
    tg = np.linspace(0, tmax, 4000); H = np.cumsum(hazard(tg)) * (tg[1] - tg[0])
    u = -np.log(rng.random(n))
    T = np.interp(u, H, tg, right=tmax)
    C = rng.uniform(0.4 * censor_time, censor_time, n)
    time = np.minimum(T, C); status = (T <= C).astype(int)
    return time, status

References