Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Accelerated Failure-Time Models › Implementation in NumPy and SciPy

Michael BrenndoerferOctober 26, 202610 min read

Machine Learning from ScratchAll chapters

Part 6 of 15Survival Analysis

Chapter 51Accelerated Failure-Time Models › Implementation in NumPy and SciPy

Reading section 7 of 10 · Implementation in NumPy and SciPy · 10 minSection 7 of 10 · 10 min

Section 7 of 10 · 10 min read

Accelerated Failure-Time Models › Implementation in NumPy and SciPySection view · Sections in this chapter7 / 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

We now fit both a log-normal and a Weibull AFT model from primitives. Working from primitives rather than calling a library routine lets each step of the derivation become a piece of code: the standardized residual $w_i$ , the density contribution of an event, the survival contribution of a censored observation, and the optimization over regression coefficients and the log-scale. The sample is deliberately small (one binary covariate, forty subjects, with sixteen censored observations in this realization) because the point is to see the likelihood machinery clearly, not to showcase a large data set. The data-generating process is fixed by a seed so the chapter is reproducible, but the fitted results are computed from the actual simulated observations rather than inserted by hand.

Setting Up the Censored LikelihoodLink Copied

The fitted output mixes several conceptually distinct quantities, and it helps to name them before looking at any numbers. The regression coefficients $\boldsymbol{\beta}$ describe the log-time location; exponentiating a slope gives a time ratio. The scale parameter $\sigma$ describes dispersion on the log-time scale; for the Weibull model, the shape parameter is $k=1/\sigma$ . The event indicator separates density contributions from survival contributions. The chosen error distribution determines $f_W$ and $S_W$ , and therefore determines the shape of the fitted hazard. Reading the code below with these roles in mind makes the connection to the formula immediate.

PythonIn [13]
CollapseExpandimport numpy as np
import pandas as pd
from scipy import optimize, stats
from scipy.special import log_ndtr

rng = np.random.default_rng(20250308)
n = 40
x = rng.integers(0, 2, size=n)
true_log_time = 2.00 - 0.45 * x + 0.55 * rng.normal(size=n)
true_time = np.exp(true_log_time)
censor_time = rng.uniform(3.0, 14.0, size=n)
observed_time = np.minimum(true_time, censor_time)
event = (true_time <= censor_time).astype(int)

X = np.column_stack([np.ones(n), x])
log_observed_time = np.log(observed_time)

The generating coefficients place the true log-time mean at 2.00 for the control group and 1.55 for the treated group, with a log-scale of 0.55. For the correctly specified log-normal fit, these generating values provide useful references, although estimates from forty subjects can differ through sampling variation. The Weibull fit uses a different error family, so its location and scale need not recover the normal generating parameters. A divergence between fit and truth here would signal either finite-sample variability, misspecification, or a problem with the numerical procedure.

The helper functions below encode the two likelihoods exactly as derived. For the log-normal model we use log_ndtr(-z) to compute $\log \Phi(-z)$ stably, rather than taking the logarithm of a potentially tiny survival probability. For the Weibull model, overflow in an extreme parameter region returns an infinite objective rather than clipping the exponent and changing the likelihood. Contributions are indexed by event status so a zero weight never multiplies an infinite log term.

PythonIn [14]
CollapseExpanddef lognormal_negll(params, X, log_time, event):
    if not np.all(np.isfinite(params)):
        return np.inf
    beta, log_sigma = params[:-1], params[-1]
    with np.errstate(over="ignore", under="ignore"):
        sigma = np.exp(log_sigma)
    if not np.isfinite(sigma) or sigma <= 0:
        return np.inf
    z = (log_time - X @ beta) / sigma
    observed = event == 1
    with np.errstate(over="ignore", invalid="ignore"):
        log_pdf = (
            -log_time[observed]
            - log_sigma
            - 0.5 * np.log(2 * np.pi)
            - 0.5 * z[observed] ** 2
        )
        log_surv = log_ndtr(-z[~observed])
        loglik = log_pdf.sum() + log_surv.sum()
    return float(-loglik) if np.isfinite(loglik) else np.inf


def weibull_negll(params, X, log_time, event):
    if not np.all(np.isfinite(params)):
        return np.inf
    beta, log_sigma = params[:-1], params[-1]
    with np.errstate(over="ignore", under="ignore"):
        sigma = np.exp(log_sigma)
    if not np.isfinite(sigma) or sigma <= 0:
        return np.inf
    z = (log_time - X @ beta) / sigma
    observed = event == 1
    with np.errstate(over="ignore", invalid="ignore"):
        exp_z = np.exp(z)
        loglik = (
            -log_time[observed] - log_sigma + z[observed]
        ).sum() - exp_z.sum()
    return float(-loglik) if np.isfinite(loglik) else np.inf

We optimize with L-BFGS-B and a small multi-start search. The bounds keep $\log\sigma$ in a numerically reasonable range. This does not replace careful modeling, but it prevents an optimizer from exploring absurd scale values during early iterations.

PythonIn [15]
CollapseExpanddef fit_aft(negll, start):
    bounds = [(None, None)] * X.shape[1] + [(-4.0, 3.0)]
    candidates = [
        start,
        start + np.array([0.0, 0.0, 0.5]),
        start - np.array([0.0, 0.0, 0.5]),
    ]
    fits = [
        optimize.minimize(
            negll,
            candidate,
            args=(X, log_observed_time, event),
            method="L-BFGS-B",
            bounds=bounds,
            options={"maxiter": 2000},
        )
        for candidate in candidates
    ]
    successful = [
        result for result in fits if result.success and np.isfinite(result.fun)
    ]
    if not successful:
        raise RuntimeError(
            "No AFT starting point produced a finite converged fit."
        )
    return min(successful, key=lambda result: result.fun)


start = np.array([np.mean(log_observed_time), 0.0, np.log(0.6)])
res_ln = fit_aft(lognormal_negll, start)
res_w = fit_aft(weibull_negll, start)

The table reports the fitted intercept, slope, log-time scale, time ratio for the binary covariate, maximized log-likelihood, and AIC. The log-likelihoods are comparable because both models are fitted to the same observations with the same censoring indicator and both are full likelihoods. Smaller AIC is better, but AIC is only one piece of evidence; a model with a better AIC can still have the wrong tail shape for the decision at hand. Reading the table alongside the survival-curve plots is therefore essential: coefficients summarize the acceleration effect, while the curves show where the distributional assumption is doing the most work.

Out[17]:

	Model	beta0	beta1	sigma	time_ratio_x1	log_likelihood	AIC	converged
0	Log-normal	2.031439	-0.331528	0.625007	0.717826	-67.866745	141.733490	True
1	Weibull	2.391341	-0.534171	0.448551	0.586155	-66.839888	139.679777	True

Checking an intercept-only referenceLink Copied

SciPy's censored univariate distribution fit provides a separate check of a simpler case. We remove the covariate from the design and fit an intercept-only log-normal model to the same observed durations and censoring indicators. This validates the two-parameter special case of the likelihood, rather than claiming a library check of the full covariate model.

PythonIn [18]
CollapseExpandX_intercept = np.ones((n, 1))
intercept_fit = optimize.minimize(
    lognormal_negll,
    np.array([np.mean(log_observed_time), np.log(0.6)]),
    args=(X_intercept, log_observed_time, event),
    method="L-BFGS-B",
    bounds=[(None, None), (-4.0, 3.0)],
    options={"maxiter": 2000},
)
if not intercept_fit.success:
    raise RuntimeError(intercept_fit.message)
censored_data = stats.CensoredData.right_censored(observed_time, event == 0)
reference_sigma, reference_location, reference_scale = stats.lognorm.fit(
    censored_data, floc=0.0
)
reference_beta0 = np.log(reference_scale)
reference_params = np.array([reference_beta0, np.log(reference_sigma)])
reference_nll = lognormal_negll(
    reference_params, X_intercept, log_observed_time, event
)
np.testing.assert_allclose(
    intercept_fit.x, reference_params, rtol=0.0, atol=3e-4
)
np.testing.assert_allclose(
    intercept_fit.fun, reference_nll, rtol=0.0, atol=1e-5
)
reference_comparison = pd.DataFrame(
    {
        "fit": ["Primitive intercept-only", "SciPy censored log-normal"],
        "beta0": [intercept_fit.x[0], reference_beta0],
        "sigma": [np.exp(intercept_fit.x[1]), reference_sigma],
        "negative_log_likelihood": [intercept_fit.fun, reference_nll],
    }
)

Out[19]:

	fit	beta0	sigma	negative_log_likelihood
0	Primitive intercept-only	1.794534	0.637115	68.778574
1	SciPy censored log-normal	1.794531	0.637113	68.778574

The assertions compare both parameters and the objective within stated numerical tolerances. Holding SciPy's location at zero matches the positive lifetime model; its scale equals exp(beta0), and its shape equals the log-time scale sigma. Agreement checks the intercept-only numerical likelihood, not the distributional adequacy of the pooled data or the precision of the covariate effect.

Fitting and Evaluating Survival CurvesLink Copied

The fitted curves use the estimated coefficients in the survival functions derived above. We compare them with group-specific Kaplan-Meier estimates, which do not impose a parametric lifetime family but still require appropriate censoring assumptions. In a forty-subject sample, a gap can reflect sampling variation as well as misspecification. The overlays are diagnostics rather than a proof that a family is correct or wrong.

PythonIn [20]
CollapseExpandt_pred = np.linspace(0.2, 15.0, 400)

eta_ln0 = res_ln.x[0]
eta_ln1 = res_ln.x[0] + res_ln.x[1]
sigma_ln = np.exp(res_ln.x[-1])
surv_ln_x0 = lognormal_survival(t_pred, eta_ln0, sigma_ln)
surv_ln_x1 = lognormal_survival(t_pred, eta_ln1, sigma_ln)

eta_w0 = res_w.x[0]
eta_w1 = res_w.x[0] + res_w.x[1]
sigma_w = np.exp(res_w.x[-1])
surv_w_x0 = weibull_survival(t_pred, eta_w0, sigma_w)
surv_w_x1 = weibull_survival(t_pred, eta_w1, sigma_w)


def kaplan_meier(time, event):
    order = np.argsort(time)
    time = time[order]
    event = event[order]
    event_times = np.unique(time[event == 1])

    step_times = [0.0]
    step_survival = [1.0]
    survival = 1.0

    for current_time in event_times:
        at_risk = np.sum(time >= current_time)
        deaths = np.sum((time == current_time) & (event == 1))
        survival *= 1.0 - deaths / at_risk
        step_times.append(current_time)
        step_survival.append(survival)

    if len(time) and step_times[-1] < time.max():
        step_times.append(time.max())
        step_survival.append(survival)
    return np.array(step_times), np.array(step_survival)


km_t0, km_s0 = kaplan_meier(observed_time[x == 0], event[x == 0])
km_t1, km_s1 = kaplan_meier(observed_time[x == 1], event[x == 1])

VisualizationOut [21]
CollapseExpand

A Kaplan-Meier step curve and fitted log-normal and Weibull survival lines for group 0. — Censored-likelihood log-normal and Weibull fits for group x = 0, compared with its Kaplan-Meier estimate. Each parametric family uses its own fitted coefficients, and the nonparametric step curve shows how well each family follows the observed survival. The Kaplan-Meier line stops at observed follow-up; the parametric curves extend to time 15 under their distributional assumptions.

VisualizationOut [22]
CollapseExpand

A Kaplan-Meier step curve and fitted log-normal and Weibull survival lines for group 1. — Censored-likelihood log-normal and Weibull fits for group x = 1, compared with its Kaplan-Meier estimate. Each parametric family uses its own fitted coefficients, so the fitted lines reflect the treated-group estimates. The Kaplan-Meier line stops at observed follow-up, while the parametric curves extend to time 15 under their distributional assumptions.

The comparison is deliberately visual and local. If one parametric curve follows the Kaplan-Meier steps closely in the region where most events occur but diverges after the last few observations, that divergence may be an extrapolation issue rather than evidence about the middle of the distribution. If the fitted curve misses the early steps systematically, the distribution or the linear log-time predictor is suspect. The printed coefficient table and the plotted curves should be read together: coefficients summarize the acceleration effect, while the curves show where the distributional assumption is doing the most work.

Key ParametersLink Copied

Below are the main parameters that affect how the AFT fits behave, with practical guidance on how to select them. The same names appear in the derivation and the implementation.

beta0: Intercept of the log-time location. It sets the baseline log time when all covariates are zero; centering covariates changes it without changing slopes.
beta1 (and other slopes): Log-time effect of each covariate. Exponentiating a slope gives a time ratio, so beta1 = -0.45 corresponds to a time ratio of about 0.64.
log_sigma (equivalently sigma): Log-scale dispersion parameterized so positive sigma is represented on the real log scale. Exponentiation makes sigma positive; the bounds restrict it to the chosen finite range, about 0.018 to 20.1; the Weibull shape is k = 1/sigma.
Error distribution: The choice of f_W and S_W. The normal error gives a log-normal model with a possibly non-monotone hazard; the minimum Gumbel error gives a Weibull model with a monotone hazard and a proportional hazards representation.
maxiter and starting values: Controls for the multi-start L-BFGS-B search. Multiple starts help with flat or non-concave likelihoods, especially under heavy censoring.

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!