Cox Proportional Hazards › Implementation in NumPy and SciPy

Michael BrenndoerferOctober 26, 202613 min read
Machine Learning from ScratchAll chapters
Part 6 of 15Survival Analysis
Chapter 50Cox Proportional Hazards › Implementation in NumPy and SciPy
Reading section 6 of 10 · Implementation in NumPy and SciPy · 13 minSection 6 of 10 · 13 min

Section 6 of 10 · 13 min read

Cox Proportional Hazards › Implementation in NumPy and SciPySection 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

We now build the partial likelihood from scratch. We use a small equipment-failure dataset with a single binary covariate: maintained equals 1 if the equipment received regular maintenance and 0 otherwise. Our goal is to estimate the effect of maintenance on the failure hazard and then to recover the implied survival curves. The dataset is deliberately small so that the connection between the formulas and the code stays visible.

The hidden setup imports the numerical libraries and shared figure helpers. The table below is a small constructed teaching dataset, with follow-up measured in arbitrary time units; it is not a collected maintenance study.

PythonIn [7]
CollapseExpand
equipment = pd.DataFrame(
    {
        "time": [5, 8, 12, 15, 20, 24, 30, 2, 3, 6, 9, 10, 14, 18, 22],
        "event": [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1],
        "maintained": [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    }
)
X = equipment[["maintained"]].to_numpy(dtype=float)
time = equipment["time"].to_numpy(dtype=float)
event = equipment["event"].to_numpy(dtype=int)

The first seven rows are maintained (x=1x=1); the last eight are not (x=0x=0). Unmaintained devices fail earlier on average, which suggests the fitted coefficient should be negative (maintenance reduces the hazard).

Out[8]:
time event maintained
0 5 1 1
1 8 1 1
2 12 0 1
3 15 1 1
4 20 1 1
5 24 1 1
6 30 1 1
7 2 1 0
8 3 1 0
9 6 1 0
10 9 1 0
11 10 0 0
12 14 1 0
13 18 1 0
14 22 1 0

Seven maintained rows precede eight unmaintained rows, and the unmaintained records are visibly concentrated at earlier times, so a negative fitted maintenance coefficient is the expected direction. We next write the log partial likelihood in a numerically stable form. The naive computation log⁡∑jexp⁡(ηj)\log \sum_j \exp(\eta_j) overflows when the linear predictor ηj=β⊤xj\eta_j = \beta^\top x_j is large. We use the log-sum-exp identity

log⁡∑jeηj=m+log⁡∑jeηj−m,m=max⁡jηj,\log \sum_j e^{\eta_j} = m + \log \sum_j e^{\eta_j - m}, \qquad m = \max_j \eta_j,

where:

  • ηj\eta_j: the linear predictor β⊤xj\beta^\top x_j for individual jj
  • mm: the largest linear predictor in the current risk set, used as a shift to keep the exponentials bounded

which shifts the exponential arguments to be nonpositive and keeps the sum bounded. In the code, the loop over np.flatnonzero(event) visits only event times; the risk set for each event is selected by time >= time[i].

PythonIn [9]
CollapseExpand
def cox_loglike_untied(beta, X, time, event):
    """Untied-event Cox partial log-likelihood using log-sum-exp."""
    eta = X @ beta
    total = 0.0
    for i in np.flatnonzero(event):
        at_risk = time >= time[i]
        eta_r = eta[at_risk]
        m = eta_r.max()
        log_denominator = m + np.log(np.exp(eta_r - m).sum())
        total += eta[i] - log_denominator
    return total
PythonIn [10]
CollapseExpand
null_loglike = cox_loglike_untied(np.zeros(X.shape[1]), X, time, event)
ConsoleOut [11]
CollapseExpand
Log partial likelihood at beta = 0: -23.623

At β=0\beta=0, every subject in a risk set has equal relative weight, so this log partial likelihood depends on the risk-set sizes and event pattern. Its raw magnitude is not a measure of covariate discrimination. Comparing a fitted value with the null objective is more informative about the covariate model.

The gradient follows directly from the score function derived in the previous section. The code computes the risk-set weighted average of the covariates and subtracts it from the failing individual's covariate. This is the same "observed minus expected" structure we saw in the Formula section, now expressed as a vector operation over the rows in the risk set.

PythonIn [12]
CollapseExpand
def cox_gradient(beta, X, time, event):
    """Score function: observed covariate minus risk-set weighted average."""
    eta = X @ beta
    p = X.shape[1]
    grad = np.zeros(p)
    for i in np.flatnonzero(event):
        at_risk = time >= time[i]
        eta_r = eta[at_risk]
        w = np.exp(eta_r - eta_r.max())
        w = w / w.sum()
        grad += X[i] - w @ X[at_risk]
    return grad

The Hessian is the negative sum of positive-semidefinite weighted covariance matrices, so the log partial likelihood is concave. We will need it later to compute standard errors. Notice that the Hessian does not depend on the event indicator except through which risk sets are visited; censored observations still enter through their weights in each risk set.

PythonIn [13]
CollapseExpand
def cox_hessian(beta, X, time, event):
    """Negative-semidefinite Hessian for the untied partial log-likelihood."""
    eta = X @ beta
    p = X.shape[1]
    H = np.zeros((p, p))
    for i in np.flatnonzero(event):
        at_risk = time >= time[i]
        eta_r = eta[at_risk]
        w = np.exp(eta_r - eta_r.max())
        w = w / w.sum()
        Xr = X[at_risk]
        mu = w @ Xr
        H -= (Xr.T * w) @ Xr - np.outer(mu, mu)
    return H

With the three functions in hand, we optimize. We minimize the negative log partial likelihood with BFGS, using its analytic gradient, and inspect the returned status and score. Concavity alone does not guarantee a finite optimum or convergence from an arbitrary starting value. We compute standard errors from the inverse information matrix,

Var⁡^(β^)=(−∂2ℓ∂β∂β⊤∣β=β^)−1,\widehat{\operatorname{Var}}(\hat\beta) = \left(-\left.\frac{\partial^2 \ell}{\partial \beta \partial \beta^\top}\right|_{\beta=\hat\beta}\right)^{-1},

where:

  • ℓ(β)\ell(\beta): the Cox partial log-likelihood
  • ∂2ℓ∂β∂β⊤\dfrac{\partial^2 \ell}{\partial \beta \partial \beta^\top}: the Hessian matrix of ℓ\ell with respect to β\beta
  • β^\hat\beta: the fitted coefficient vector

We solve against the observed information to obtain an approximate covariance matrix, then take square roots of its diagonal for standard errors. The fitted coefficient, hazard ratio, and confidence interval are assembled into a table.

PythonIn [14]
CollapseExpand
fit = optimize.minimize(
    lambda b: -cox_loglike_untied(b, X, time, event),
    x0=np.zeros(X.shape[1]),
    jac=lambda b: -cox_gradient(b, X, time, event),
    method="BFGS",
)
beta_hat = fit.x
H_hat = cox_hessian(beta_hat, X, time, event)
covariance_hat = np.linalg.solve(-H_hat, np.eye(X.shape[1]))
se_hat = np.sqrt(np.diag(covariance_hat))
assert (
    fit.success
    and np.max(np.abs(cox_gradient(beta_hat, X, time, event))) < 1e-5
)

summary = pd.DataFrame(
    {
        "covariate": ["maintained"],
        "coef": beta_hat,
        "hazard_ratio": np.exp(beta_hat),
        "se": se_hat,
        "ci_lower": np.exp(beta_hat - 1.96 * se_hat),
        "ci_upper": np.exp(beta_hat + 1.96 * se_hat),
    }
)
optimizer_success = fit.success
Out[15]:
covariate coef hazard_ratio se ci_lower ci_upper
0 maintained -0.882 0.414 0.637 0.119 1.442
ConsoleOut [16]
CollapseExpand
Optimizer success: True

The fitted coefficient is approximately -0.9 with a hazard ratio of about 0.4, which is drawn from the summary table above; the approximate model-based 95 percent Wald hazard-ratio interval spans from its ci_lower to its ci_upper column. Under the fitted fixed-covariate model, maintained equipment has about 41 percent of the reference hazard wherever the baseline hazard is positive. This describes the constructed example's conditional association. The interval includes one, and thirteen events do not give a precise maintenance-effect estimate.

To check that our implementation matches a mature library, we compare with statsmodels' PHReg, which fits the same partial likelihood. It uses default Breslow tie handling here, and we explicitly select Efron in the separate tied comparison. Agreement on this dataset is a useful numerical check. The stabilization identity, risk-set construction, and score formula still need to be checked directly.

PythonIn [17]
CollapseExpand
from statsmodels.duration.hazard_regression import PHReg

reference = PHReg(time, X, status=event).fit()
Out[18]:
our_beta statsmodels_beta our_se statsmodels_se
0 -0.8821 -0.8821 0.6368 0.6368

The two coefficient columns and the two standard-error columns should line up to several decimal places in the displayed table; small residual differences are compatible with optimizer and information-matrix-solve tolerance. Larger discrepancies call for checking the inputs, objective conventions, derivatives, and numerical settings. This is the expected outcome: PHReg and the scratch implementation maximize the same concave log partial likelihood.

Handling tied event times

Our example has no tied event times, so the untied likelihood is exact. Real data often produce ties; exact methods and approximations depend on how those recorded times are modeled. We can demonstrate the difference by coarsening the observation times into four-unit bins, which forces multiple events to land at the same recording time. We implement both the Breslow and Efron approximations and compare them. The Breslow version treats all tied failures as if they occurred with the full risk set available; the Efron version fractionally reduces the total weight of the tied failures across successive denominators.

PythonIn [19]
CollapseExpand
def cox_loglike_breslow(beta, X, time, event):
    """Breslow approximation for tied event times."""
    eta = X @ beta
    total = 0.0
    for t in np.unique(time[event == 1]):
        at_risk = time >= t
        events_t = (time == t) & (event == 1)
        d = events_t.sum()
        eta_r = eta[at_risk]
        m = eta_r.max()
        log_denominator = m + np.log(np.exp(eta_r - m).sum())
        total += eta[events_t].sum() - d * log_denominator
    return total


def cox_loglike_efron(beta, X, time, event):
    """Efron approximation for tied event times."""
    eta = X @ beta
    total = 0.0
    for t in np.unique(time[event == 1]):
        at_risk = time >= t
        events_t = (time == t) & (event == 1)
        d = int(events_t.sum())
        eta_r = eta[at_risk]
        eta_d = eta[events_t]
        m = eta_r.max()
        s0 = np.exp(eta_r - m).sum()
        s_d = np.exp(eta_d - m).sum()
        if d == 1:
            total += eta_d.sum() - (m + np.log(s0))
        else:
            log_denom = 0.0
            for ell in range(d):
                log_denom += m + np.log(s0 - (ell / d) * s_d)
            total += eta_d.sum() - log_denom
    return total

We now coarsen the observation times into bins of four units and fit both approximations.

PythonIn [20]
CollapseExpand
binned_time = np.ceil(time / 4).astype(int) * 4
X_b = X.copy()  # reuse the same covariate matrix without aliasing

fit_breslow = optimize.minimize(
    lambda b: -cox_loglike_breslow(b, X_b, binned_time, event),
    x0=np.zeros(X_b.shape[1]),
    method="BFGS",
)
fit_efron = optimize.minimize(
    lambda b: -cox_loglike_efron(b, X_b, binned_time, event),
    x0=np.zeros(X_b.shape[1]),
    method="BFGS",
)
reference_efron = PHReg(binned_time, X_b, status=event, ties="efron").fit()

tie_comparison = pd.DataFrame(
    {
        "method": [
            "Breslow (from scratch)",
            "Efron (from scratch)",
            "statsmodels Efron",
        ],
        "beta": [fit_breslow.x[0], fit_efron.x[0], reference_efron.params[0]],
        "hazard_ratio": [
            np.exp(fit_breslow.x[0]),
            np.exp(fit_efron.x[0]),
            np.exp(reference_efron.params[0]),
        ],
    }
)
Out[21]:
method beta hazard_ratio
0 Breslow (from scratch) -0.5892 0.5548
1 Efron (from scratch) -0.6364 0.5292
2 statsmodels Efron -0.6364 0.5292

The two tie approximations can also be compared over the whole coefficient path. The display grid is fixed and no additional fitting is performed in the figure cell.

PythonIn [22]
CollapseExpand
tie_beta_grid = np.linspace(-2.0, 1.0, 300)
tie_breslow_grid = np.array(
    [
        cox_loglike_breslow(np.array([b]), X_b, binned_time, event)
        for b in tie_beta_grid
    ]
)
tie_efron_grid = np.array(
    [
        cox_loglike_efron(np.array([b]), X_b, binned_time, event)
        for b in tie_beta_grid
    ]
)
tie_breslow_opt = float(fit_breslow.x[0])
tie_efron_opt = float(fit_efron.x[0])
tie_breslow_opt_value = cox_loglike_breslow(
    fit_breslow.x, X_b, binned_time, event
)
tie_efron_opt_value = cox_loglike_efron(fit_efron.x, X_b, binned_time, event)
VisualizationOut [23]
CollapseExpand
Two concave log partial likelihood curves for Breslow and Efron tied-time approximations.
Breslow and Efron log partial likelihoods for the coarsened tied-time dataset. The two curves are close but not identical, and their maxima sit at the fitted coefficients reported in the table. The gap reflects this coarsened sample rather than a universal rule about tie bias.

The executed tied-data table reports the Breslow and Efron coefficients and their hazard ratios alongside the explicitly selected statsmodels Efron fit. Our Efron coefficient matches the statsmodels Efron coefficient to within 10−510^{-5}. These are different approximations applied to this coarsened dataset, not a universal statement about the direction or size of tie bias.

Baseline hazard and survival curves

Estimating the baseline hazard lets us convert coefficients into absolute survival predictions rather than relative risk scores. We implement Breslow's estimator and predict survival for both groups at a grid of times. The estimator accumulates one increment at each event time: the number of events divided by the total at-risk weight. The prediction function then applies S^(t∣x)=exp⁡(−H^0(t)eβ^⊤x)\hat S(t \mid x) = \exp(-\hat H_0(t)e^{\hat\beta^\top x}).

PythonIn [24]
CollapseExpand
def breslow_baseline_cumhaz(beta, X, time, event):
    """Breslow estimator of the cumulative baseline hazard."""
    eta = X @ beta
    event_times = np.sort(np.unique(time[event == 1]))
    increments = []
    for t in event_times:
        at_risk = time >= t
        d = ((time == t) & (event == 1)).sum()
        increments.append(d / np.exp(eta[at_risk]).sum())
    return event_times, np.cumsum(increments)


def predict_survival(beta, t_grid, H0_times, H0_values, x):
    """Survival curve implied by the Cox model for a given covariate vector."""
    t_grid = np.asarray(t_grid, dtype=float)
    positions = np.searchsorted(H0_times, t_grid, side="right") - 1
    H0 = np.zeros_like(t_grid)
    valid = positions >= 0
    H0[valid] = H0_values[positions[valid]]
    return np.exp(-H0 * np.exp(x @ beta))

We compute the baseline and store predicted curves for the two covariate values.

PythonIn [25]
CollapseExpand
t_grid = np.unique(np.r_[0.0, time])
H0_times, H0_values = breslow_baseline_cumhaz(beta_hat, X, time, event)
surv_unmaintained = predict_survival(
    beta_hat, t_grid, H0_times, H0_values, np.array([0.0])
)
surv_maintained = predict_survival(
    beta_hat, t_grid, H0_times, H0_values, np.array([1.0])
)
event_rug_times = time[event == 1]
event_rug_y = np.full(event_rug_times.shape, -0.02)

The fitted survival curves start at one and are nonincreasing, with the maintained group at or above the unmaintained group over the displayed follow-up. Because the hazard ratio is constant under the fitted model, the vertical gap between the curves changes over time rather than remaining fixed.

We also compute Schoenfeld residuals, which we will use as a diagnostic in the Visualizing Cox Proportional Hazards section. For each event time, the residual is the failing individual's covariate minus the risk-set weighted average. These residuals are the empirical counterpart of the score contributions: they tell us how surprising each failure was relative to the risk set that produced it.

PythonIn [26]
CollapseExpand
def schoenfeld_residuals(beta, X, time, event):
    """Schoenfeld residual for each event and covariate."""
    eta = X @ beta
    p = X.shape[1]
    event_idx = np.flatnonzero(event)
    R = np.zeros((len(event_idx), p))
    for k, i in enumerate(event_idx):
        at_risk = time >= time[i]
        eta_r = eta[at_risk]
        w = np.exp(eta_r - eta_r.max())
        w = w / w.sum()
        R[k] = X[i] - w @ X[at_risk]
    return time[event_idx], R
PythonIn [27]
CollapseExpand
event_times_schoenfeld, schoenfeld = schoenfeld_residuals(
    beta_hat, X, time, event
)
trend = np.polyfit(event_times_schoenfeld, schoenfeld[:, 0], 1)
trend_times = np.sort(event_times_schoenfeld)
trend_values = np.polyval(trend, trend_times)

The fitted quantities are now stored, and the preceding tied-event demonstration uses the same scratch functions with binned times. The next section uses these quantities for survival and diagnostic figures.

Key Parameters

The main choices that affect the fitted Cox model and its predictions are:

  • Tie handling (cox_loglike_untied, cox_loglike_breslow, cox_loglike_efron): selects how tied event times are treated. Use the untied likelihood when times are distinct; use the Efron approximation for coarsened continuous times with ties; Breslow is a useful starting point but can pull estimates toward the null with heavy ties.
  • Optimizer tolerance and starting values (optimize.minimize, method="BFGS"): concavity is not a finite-estimability or convergence guarantee, so inspect the returned status and the score. Tighter gtol and sensible starting values help on harder problems.
  • Censoring/risk-set construction (time >= time[i]): censored individuals remain at risk until their censoring time and must be included in every applicable risk set; dropping them destroys the partial likelihood.
  • Baseline grid (t_grid, breslow_baseline_cumhaz): the survival grid determines where absolute survival is reported; it must cover the horizons of interest.

Key Methods

The methods we built map one-to-one onto the questions a practitioner asks of a Cox model. cox_loglike_untied plus cox_gradient implements estimation; breslow_baseline_cumhaz implements absolute risk reconstruction; schoenfeld_residuals implements diagnostics. Each can be replaced by a more sophisticated version (Efron ties, smoothed baseline, weighted residuals), and the underlying decomposition stays the same.

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!

Newsletter

Stay up to date

Get new articles, book updates, and notes on AI and technology.

Unsubscribe at any time. Privacy policy