Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Kaplan–Meier and Nelson–Aalen Estimation › Implementation in NumPy

Michael BrenndoerferOctober 26, 20269 min read

Machine Learning from ScratchAll chapters

Part 6 of 15Survival Analysis

Chapter 49Kaplan–Meier and Nelson–Aalen Estimation › Implementation in NumPy

Reading section 6 of 10 · Implementation in NumPy · 9 minSection 6 of 10 · 9 min

Section 6 of 10 · 9 min read

Kaplan–Meier and Nelson–Aalen Estimation › Implementation in NumPySection view · Sections in this chapter6 / 10 ↓

Sections in this chapter10

Reading tools Definitions, links & notes

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Notes & comments

Notes & commentsLeave notes and join the discussionSelection toolbarShow comment actions when selecting text. Turn off for the browser menu only.Comment indicatorsShow small counts beside annotated passagesPassage highlightsGently shade annotated text selections

PythonIn [5]
CollapseExpandimport matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from book_plot_style import PALETTE, polish_axes, use_book_style
from scipy.stats import norm

PythonIn [6]
CollapseExpanddef fit_km_na(times, events):
    """Return event-time Kaplan-Meier and Nelson-Aalen estimates."""
    times = np.asarray(times, dtype=float)
    events = np.asarray(events)
    if times.ndim != 1 or events.ndim != 1 or times.shape != events.shape:
        raise ValueError(
            "times and events must be one-dimensional with the same shape"
        )
    if times.size == 0:
        raise ValueError("at least one observation is required")
    if not np.all(np.isfinite(times)) or np.any(times <= 0):
        raise ValueError("times must be finite positive elapsed durations")
    if np.any((events != 0) & (events != 1)):
        raise ValueError("events must contain only 0 and 1")
    events = events.astype(int)

    event_times = np.unique(times[events == 1])
    columns = [
        "time",
        "n_at_risk",
        "events",
        "factor",
        "S_km",
        "H_na",
        "greenwood_sum",
        "var_H_sum",
    ]
    if event_times.size == 0:
        return pd.DataFrame(columns=columns)

    rows = []
    S = 1.0
    H = 0.0
    greenwood_sum = 0.0
    var_H_sum = 0.0

    for tj in event_times:
        n = int(np.sum(times >= tj))
        d = int(np.sum((times == tj) & (events == 1)))
        factor = 1.0 - d / n
        S *= factor
        H += d / n
        if n > d:
            greenwood_sum += d / (n * (n - d))
        else:
            greenwood_sum = np.nan
        var_H_sum += d * (n - d) / n**3
        rows.append(
            {
                "time": tj,
                "n_at_risk": n,
                "events": d,
                "factor": factor,
                "S_km": S,
                "H_na": H,
                "greenwood_sum": greenwood_sum,
                "var_H_sum": var_H_sum,
            }
        )

    return pd.DataFrame(rows, columns=columns)

PythonIn [7]
CollapseExpanddef step_value(times, values, t, left_value=1.0):
    """Evaluate a right-continuous step function at time t."""
    times = np.asarray(times)
    values = np.asarray(values)
    idx = np.searchsorted(times, t, side="right") - 1
    if idx < 0:
        return float(left_value)
    return float(values[idx])


def km_loglog_ci(km, level=0.95):
    """Pointwise Greenwood log-log confidence interval for Kaplan-Meier."""
    if not np.isfinite(level) or not 0 < level < 1:
        raise ValueError("level must be strictly between 0 and 1")
    z = norm.ppf(0.5 + level / 2.0)
    S = km["S_km"].to_numpy(dtype=float)
    green = km["greenwood_sum"].to_numpy(dtype=float)
    lower = np.full_like(S, np.nan)
    upper = np.full_like(S, np.nan)
    mask = (S > 0) & (S < 1)
    se = np.zeros_like(S)
    se[mask] = np.sqrt(green[mask]) / np.abs(np.log(S[mask]))
    lower[mask] = S[mask] ** np.exp(z * se[mask])
    upper[mask] = S[mask] ** np.exp(-z * se[mask])
    return lower, upper

At an event that exhausts the risk set, the point estimates remain valid but greenwood_sum and the transformed limits are missing. An all-censored nonempty sample returns an empty event-time table, so there is no estimated median to report.

We first reproduce the ten-person example from the previous section.

PythonIn [8]
CollapseExpandhand = pd.DataFrame(
    {
        "subject": np.arange(1, 11),
        "time": np.array([2, 3, 3, 5, 6, 6, 8, 9, 10, 12], dtype=float),
        "event": np.array([1, 1, 1, 0, 1, 0, 1, 0, 1, 0], dtype=int),
    }
)

km_hand = fit_km_na(hand["time"], hand["event"])

Out[9]:

	time	n_at_risk	events	factor	S_km	H_na	greenwood_sum	var_H_sum
0	2.0	10	1	0.900000	0.900000	0.100000	0.011111	0.009000
1	3.0	9	2	0.777778	0.700000	0.322222	0.042857	0.028204
2	6.0	6	1	0.833333	0.583333	0.488889	0.076190	0.051353
3	8.0	4	1	0.750000	0.437500	0.738889	0.159524	0.098228
4	10.0	2	1	0.500000	0.218750	1.238889	0.659524	0.223228

The table matches the hand calculation. The factor column shows the conditional survival factor at each event time, and the H_na column accumulates the Nelson–Aalen increments. At the final event time, the Kaplan–Meier estimate is $0.21875$ and the Nelson–Aalen estimate is $1.238889$ .

The cumulative curves hide the event-by-event arithmetic. At each event time we can compare the additive Nelson-Aalen increment with the exact decrement in log survival that the product-limit estimator uses.

PythonIn [10]
CollapseExpandincrement_times = km_hand["time"].to_numpy(dtype=float)
na_increments = km_hand["events"].to_numpy(dtype=float) / km_hand[
    "n_at_risk"
].to_numpy(dtype=float)
km_log_decrements = -np.log(km_hand["factor"].to_numpy(dtype=float))
cum_na_increments = np.cumsum(na_increments)
cum_km_decrements = np.cumsum(km_log_decrements)

VisualizationOut [11]
CollapseExpand

Grouped bars compare Nelson-Aalen increments with exact log-survival decrements at five times. — Nelson-Aalen increments d_j/n_j and exact log-survival decrements -log(1-d_j/n_j) at the five event times in the ten-person example. The exact log-survival decrement is larger at each event time, so its cumulative negative log survival exceeds the Nelson-Aalen estimate.

ConsoleOut [12]
CollapseExpandKM at final event time 10: 0.218750
Nelson-Aalen at final event time: 1.238889
exp(-H_NA): 0.289706
Greenwood SE of S at final event time: 0.177649
Greenwood-type SE of Nelson-Aalen at final event time: 0.472470
KM median: 8

Now we generate a larger right-censored sample. The true event times follow a Weibull distribution with shape $1.5$ and scale $12$ , and censoring times follow an independent exponential distribution with mean $28$ . Because we simulate both the event times and the censoring times, we know the true survival curve and can compare it with the estimates.

PythonIn [13]
CollapseExpandrng = np.random.default_rng(20240517)
n_sim = 240
shape_true = 1.5
scale_true = 12.0

true_time = scale_true * (-np.log(rng.uniform(size=n_sim))) ** (
    1.0 / shape_true
)
censor_time = rng.exponential(scale=28.0, size=n_sim)
sim_time = np.minimum(true_time, censor_time)
sim_event = (true_time <= censor_time).astype(int)

sim = pd.DataFrame({"time": sim_time, "event": sim_event})
km_sim = fit_km_na(sim["time"], sim["event"])

Out[14]:

	time	n_at_risk	events	factor	S_km	H_na	greenwood_sum	var_H_sum
0	0.366713	235	1	0.995745	0.995745	0.004255	0.000018	0.000018
1	0.803951	228	1	0.995614	0.991377	0.008641	0.000038	0.000037
2	0.816653	227	1	0.995595	0.987010	0.013047	0.000057	0.000057
3	0.883790	225	1	0.995556	0.982623	0.017491	0.000077	0.000076
4	0.923124	224	1	0.995536	0.978237	0.021955	0.000097	0.000096

ConsoleOut [15]
CollapseExpandSample size: 240
Observed events: 162
Censored observations: 78
KM median in simulation: 9.240

We compute the pointwise log-log confidence limits and the Greenwood standard errors before plotting, so the plot cells themselves only present already-computed results.

PythonIn [16]
CollapseExpandlo_sim, hi_sim = km_loglog_ci(km_sim)
se_sim = np.where(
    km_sim["S_km"].to_numpy() > 0,
    km_sim["S_km"].to_numpy() * np.sqrt(km_sim["greenwood_sum"].to_numpy()),
    np.nan,
)

For a reference comparison, we use two installed implementations of Kaplan–Meier: SciPy and statsmodels. The check includes the tied ten-person table and the seeded sample. statsmodels also exposes its event counts and risk-set sizes; applying our grouped Nelson–Aalen sum to those independently computed counts checks the bookkeeping, without calling a library Nelson–Aalen fitter.

PythonIn [17]
CollapseExpandfrom scipy.stats import CensoredData, ecdf
from statsmodels.duration.survfunc import SurvfuncRight

reference_rows = []
for case, data, fitted in [
    ("Tied hand table", hand, km_hand),
    ("Seeded sample", sim, km_sim),
]:
    durations = data["time"].to_numpy(dtype=float)
    indicators = data["event"].to_numpy(dtype=int)
    scipy_sf = ecdf(CensoredData.right_censored(durations, indicators == 0)).sf
    reference_km = SurvfuncRight(durations, indicators)
    event_grid = fitted["time"].to_numpy(dtype=float)
    np.testing.assert_array_equal(event_grid, reference_km.surv_times)
    np.testing.assert_array_equal(fitted["n_at_risk"], reference_km.n_risk)
    np.testing.assert_array_equal(fitted["events"], reference_km.n_events)
    scipy_values = scipy_sf.evaluate(event_grid)
    reference_na = np.cumsum(reference_km.n_events / reference_km.n_risk)
    np.testing.assert_allclose(
        fitted["S_km"], scipy_values, atol=1e-12, rtol=1e-12
    )
    np.testing.assert_allclose(
        fitted["S_km"], reference_km.surv_prob, atol=1e-12, rtol=1e-12
    )
    np.testing.assert_allclose(
        fitted["H_na"], reference_na, atol=1e-12, rtol=1e-12
    )
    reference_rows.append(
        {
            "Case": case,
            "KM–SciPy max error": np.max(np.abs(fitted["S_km"] - scipy_values)),
            "KM–statsmodels max error": np.max(
                np.abs(fitted["S_km"] - reference_km.surv_prob)
            ),
            "NA count-check max error": np.max(
                np.abs(fitted["H_na"] - reference_na)
            ),
        }
    )
comparison = pd.DataFrame(reference_rows)

Out[18]:

	Case	KM–SciPy max error	KM–statsmodels max error	NA count-check max error
0	Tied hand table	0.000000e+00	5.551115e-17	0.0
1	Seeded sample	1.110223e-16	3.330669e-16	0.0

The computed errors are within floating-point precision. The tied table exercises the events-before-censorings convention. These numerical checks establish agreement for the examples; they do not validate independent censoring in an observed population.

A small edge case illustrates the median problem. In this three-person example, the only event occurs early and the other two subjects are censored afterward, so survival stays at $2/3$ and the median is not reached.

PythonIn [19]
CollapseExpandno_median = fit_km_na([1.0, 2.0, 3.0], [1, 0, 0])

Out[20]:

	time	n_at_risk	events	factor	S_km	H_na	greenwood_sum	var_H_sum
0	1.0	3	1	0.666667	0.666667	0.333333	0.166667	0.074074

ConsoleOut [21]
CollapseExpandLast KM value in edge case: 0.666667
Did the curve reach 0.5? False

Key ParametersLink Copied

The times argument contains positive follow-up times in a consistent unit. It should not contain calendar dates unless they have been converted to elapsed time from a meaningful origin. The events argument is the event indicator: $1$ for an observed event and $0$ for right censoring. Missing times or missing event indicators cannot be handled by the basic estimator; they require a missing-data strategy before fitting.

The tie convention is part of the estimator, not a cosmetic detail. Our function counts every subject with $Y_i\ge t_j$ in $n_j$ , including subjects censored exactly at $t_j$ , and counts only observed events in $d_j$ . This implements events before censorings. The level argument in km_loglog_ci controls the pointwise confidence level. The usual censoring and sampling assumptions and a sufficiently supported large-sample regime are needed for the Greenwood/delta-method approximation, particularly away from risk exhaustion and transform boundaries.

In SciPy, CensoredData.right_censored receives a Boolean censoring mask, so we pass indicators == 0; ecdf(...).sf then represents the estimated survival function. statsmodels SurvfuncRight instead receives the event indicators directly, with 1 meaning an event. Keeping these opposite indicator conventions explicit prevents a reference check from silently reversing the data.

Key MethodsLink Copied

fit_km_na returns a table at the distinct event times. The columns n_at_risk and events expose the risk-set construction, which is often more informative than the final curve alone. The columns S_km and H_na are the two estimators. The columns greenwood_sum and var_H_sum are cumulative variance components, so their square roots give standard errors after multiplying by the appropriate factors.

step_value evaluates a right-continuous step function at a scalar query time, given sorted event times and an aligned value array. The reference comparison above instead evaluates the library curves on the fitted event-time grid. km_loglog_ci returns pointwise lower and upper survival limits. For a full curve, evaluate the limits at event times and shade only the finite limits, remembering that the resulting band is pointwise rather than simultaneous. Evaluating a step function beyond observed follow-up does not establish survival there.

The SciPy survival object provides evaluate for step-function values. The statsmodels arrays surv_times, surv_prob, n_risk, and n_events expose the event-time curve and the counts used in our comparison. The from-scratch table makes the same risk-set logic visible for inspection.

Comments

Add to the conversation

A question, a correction, or an example from your own work.

No comments yet. Be the first to share your thoughts!

Key ParametersLink Copied

Key MethodsLink Copied

Comments

Add to the conversation

Stay up to date