Every model I have written about so far needed a finished answer in every row. Linear regression needed a number. Logistic regression needed a label. Poisson regression needed a count.
Then someone asks "how long until the customer cancels?" and the data set falls apart. Half your customers have not cancelled. Their lifetime does not exist yet. You cannot put them in a regression because they have no y, and you cannot delete them because they are usually your best customers.
That is right-censoring, and it is not missing data. A subscriber who is still active at month 11 tells you something precise: their true cancellation time is greater than 11. The inequality is information.
The two wrong answers
Almost everyone reaches for one of these, and both are badly biased in the same direction.
# WRONG 1 — complete cases only: keeps exactly the SHORT lifetimes
km1 = kaplan_meier(rows[rows.event == 1])
# WRONG 2 — every censoring becomes an event
km2 = kaplan_meier(rows.assign(event=1))
Deleting censored rows keeps only the subjects who failed, which is a sample of short lifetimes by construction. Recoding a censoring as an event asserts that everyone who walked out of the study died on the way. Both pull survival down.
On a simulated cohort with rate λ = 0.10, the true median is ln2/λ = 6.93. With 45% censoring:
drop censored -> 4.1 wrong
censored = event -> 4.5 wrong
Kaplan-Meier -> 6.7 right
The risk set is where censoring gets handled
Everything downstream reads from one table. Walk the distinct times in order. At each time record n (how many are still under observation), d (events right now), c (censorings right now). Then n drops by exactly d + c.
def risk_table(rows):
out, n = [], len(rows)
for t in sorted({r["t"] for r in rows}):
d = sum(1 for r in rows if r["t"] == t and r["e"] == 1)
c = sum(1 for r in rows if r["t"] == t and r["e"] == 0)
out.append((t, n, d, c)) # n = #{ t_j >= t }
n -= d + c # they leave AFTER this instant
return out
A subject censored at time 9 belongs to every risk set up to 9 and none after it. They contribute to every denominator they earned and then stop counting. No deletion. No invention. That is the entire fix, and everything else inherits it for free.
The tie convention matters: someone censored exactly at t is still at risk at t, because they were alive when that instant's events happened.
Kaplan-Meier
Surviving past t means surviving every risky instant before t in turn, so chain the conditionals:
S = 1.0
for (t, n, d, c) in risk_table(rows):
if d == 0:
continue # censoring alone does NOT step the curve
S *= 1 - d / n # the product-limit estimator
steps.append((t, S))
That is it. A right-continuous staircase that steps down only at event times, and it is the non-parametric maximum-likelihood estimator of S(t) — no distribution assumed anywhere.
There is a free correctness check hiding in it. With no censoring at all, n_{i+1} = n_i - d_i, so the product telescopes to #{t_i > t} / n — the plain empirical survival curve. Any implementation that fails its own easy case is not worth trusting on the hard one.
Error bars, and a band that stays legal
Survival curves get read hardest exactly where they are weakest — the far right tail, where three people are left. Greenwood's formula is the delta method applied to the product:
gw += d / (n * (n - d)) # Greenwood's sum
se = S * sqrt(gw)
Small n in that denominator is what makes the band flare out on the right.
The naive interval S ± z·se then embarrasses you by poking above 1 near the start and below 0 in the tail. Build the interval on log(-log S) instead, which is unconstrained, and map back:
sigma = sqrt(gw) / abs(log(S))
lo, hi = S ** exp(z*sigma), S ** exp(-z*sigma) # always inside [0,1]
Median survival, and "not reached"
Do not average the follow-up times. With censoring present that number is not the mean survival time and is not anything else useful either. The median is a lookup:
def median_survival(steps):
for t, S in steps:
if S <= 0.5:
return t
return None # "not reached" — report it, do not invent a number
If the curve flattens at 0.62 and stops, the median genuinely does not exist in your data. That is a real result: the typical subject outlived your study.
Comparing two curves: the log-rank test
Freeze at each pooled event time. There are n at risk, n_B of them in group B, and d events happen now. If the groups were interchangeable, the events would be allocated like drawing balls from an urn without replacement — hypergeometric — so the expected count in B is d·n_B/n with a closed-form variance.
O = E = V = 0
for t in sorted({r.t for r in rows if r.e == 1}):
n = sum(1 for r in rows if r.t >= t)
nB = sum(1 for r in rows if r.t >= t and r.g == 1)
d = sum(1 for r in rows if r.t == t and r.e == 1)
dB = sum(1 for r in rows if r.t == t and r.e == 1 and r.g == 1)
O += dB
E += d * nB / n
V += d * (nB/n) * (1 - nB/n) * (n - d) / (n - 1) if n > 1 else 0
chi2 = (O - E) ** 2 / V
Summed over both groups, observed and expected are identical by construction. That invariant is the first thing a broken loop breaks, so assert it.
The p-value needs no library. A chi-square on 1 df is the square of a standard normal, so its upper tail is erfc(sqrt(x/2)) — and erfc is a Taylor series for small arguments plus a continued fraction for large ones, about fifteen lines for full double precision.
Hazard, and why Cox works
Survival is cumulative. Hazard is instantaneous: the failure rate right now among those who made it this far. Its cumulative estimator is even simpler than Kaplan-Meier — add d/n where KM multiplies 1 - d/n:
H += d / n # Nelson-Aalen
S_approx = exp(-H) # shadows the KM curve when hazards are small
Hazard is the right language for comparison because hazards multiply. Cox's model says h(t|x) = h0(t)·exp(βx): an unspecified baseline shared by everyone, scaled by a covariate effect. The remarkable part is that you never estimate h0. Condition on the fact that someone failed at each event time and ask only which member of the risk set it was — the baseline cancels top and bottom:
def cox_partial_loglik(rows, beta):
ll = 0.0
for r in rows:
if r.e != 1: continue
risk = sum(exp(beta * j.x) for j in rows if j.t >= r.t)
ll += beta * r.x - log(risk) # h0(t) is simply not in this expression
return ll
That is what "semi-parametric" means: fully parametric in the covariates, completely free in time. exp(β̂) is the hazard ratio. And ll(0) = -Σ log n_i, which is a free hand-checkable anchor for your implementation.
Testing it honestly
Statistical code is easy to write and easy to get subtly wrong — an off-by-one in the risk set still produces a plausible-looking curve. So the version behind this walkthrough is checked against baselines written independently, not copies of itself:
KM == empirical survival when nothing is censored (12 random datasets, 1e-12)
n_i == #{t_j >= t_i}, and sum(d) + sum(c) == n0 (40 datasets)
Greenwood == finite-difference delta method (rel err < 1e-6)
sum(O) == sum(E) over both groups (algebraic identity)
chi2 tail == Simpson-integrated chi2(1) density (< 1e-9)
null hypothesis: mean chi2 ~ 1, rejection rate ~ 5% (400 trials)
cox_loglik(0) == -sum(log n_i) (hand-derived)
grid maximiser == independent brute-force scan
78 assertions, all passing.
In practice
from lifelines import KaplanMeierFitter, CoxPHFitter
from lifelines.statistics import logrank_test
km = KaplanMeierFitter().fit(df.time, df.event)
km.median_survival_time_ # inf when the curve never reaches 0.5
cph = CoxPHFitter().fit(df, "time", "event")
cph.summary # coef = beta, exp(coef) = hazard ratio
cph.check_assumptions(df) # PH is an ASSUMPTION — test it
Four traps. Censoring must be non-informative: if sick patients leave the trial early, no estimator saves you. Crossing curves mean hazards are not proportional, which blinds the log-rank test and misleads Cox. A p-value is not an effect size — report the hazard ratio and its interval. And competing risks need Aalen-Johansen, not 1 - KM.
The most common industrial version of all of this is quietly fitting logistic regression to "did they churn in the last 90 days?" — which throws away the timing, forces an arbitrary horizon, and mislabels every customer who has not been around for 90 days yet. Kaplan-Meier and Cox exist so you do not have to make that trade.
Generate a cohort and watch the two wrong curves peel away from the truth, step through the product-limit ledger row by risk-set row, and score the real Cox partial likelihood across a grid of β: https://dev48v.infy.uk/ml/day60-survival-analysis.html
Top comments (0)