Survival analysis in Python: Cox proportional hazards, stcox and streg

How long does an unemployed worker take to find a job, a new firm to exit, a loan to default? Duration data have two features that ordinary regression handles badly. Durations are positive and skewed. More importantly, many are censored: the spell is still running when observation stops, so we know only that it lasted at least this long. Dropping censored spells, or treating the censoring time as the end of the spell, biases the estimates towards short durations. Survival analysis models the hazard of the spell ending and uses censored spells for exactly the information they carry.

What survival models estimate

The hazard h(t | x) is the rate at which spells end at time t among those still at risk just before t. The survivor function S(t) is the probability that a spell lasts beyond t. The Kaplan–Meier estimator (Kaplan and Meier 1958) estimates S(t) without any model, treating each censored spell as at risk until its censoring time, and the log-rank test compares survivor curves between groups.

The Cox proportional hazards model (Cox 1972) adds covariates without specifying how the hazard depends on time:

h(t | x) = h₀(t) · exp(x′β)

The baseline hazard h₀(t) is left unrestricted. The coefficients are estimated from the partial likelihood, which uses only the order in which spells end, and exp(βₖ) is a hazard ratio: the proportional change in the hazard for a one-unit change in xₖ, at every point in time. That is the key assumption. If benefits lower the job-finding hazard by 40% in the first month but not at all after a year, the hazards are not proportional. Grambsch and Therneau (1994) test the assumption with scaled Schoenfeld residuals; stratifying on the offending variable, or letting its coefficient vary with time, relaxes it.

Parametric models specify the baseline hazard as well. The exponential has a constant hazard; the Weibull, h₀(t) = exp(β₀)·p·tᵖ⁻¹, a hazard that rises (p > 1) or falls (p < 1) monotonically; Gompertz, lognormal, loglogistic and generalized gamma allow other shapes. A parametric model is more efficient and can predict durations beyond the observed follow-up, but only if its shape is right. Models are written either in the proportional-hazards (PH) metric, in which coefficients are log hazard ratios, or in the accelerated failure time (AFT) metric, in which they are log time ratios that stretch or shrink the duration. The Weibull is both, with b_AFT = −b_PH / p.

All of these models assume that censoring is unrelated to the outcome once the covariates are accounted for. Follow-up that ends on a fixed date satisfies it; workers who drop out of a survey because they found a job do not.

In OpenEconometrics

The survival data are declared in each call, much as Stata's stset declares them: time= is the exit time, failure= the 0/1 event indicator (every record fails when it is omitted), entry= a delayed entry time or the start of a (start, stop] record, and id= the subject of multiple records. The Cox model below corresponds to stset weeks, failure(found_job) followed by stcox age benefit:

• oe.stcox(data=df, time="weeks", failure="found_job", x=["age", "benefit"]). Coefficients are log hazard ratios and there is no constant; extra["hazard_ratios"] holds exp(β) with standard errors and intervals.
• ties="breslow" (default), "efron" or "exactp". strata= gives each stratum its own baseline hazard, tvc= interacts covariates with a function of time chosen by texp="identity" or "log", and offset= is available.
• covariance="robust" is the Lin and Wei (1989) sandwich, clustered on id when one is given; cluster= takes one column.
• One fit reports the likelihood-ratio and score tests of β = 0, the Grambsch–Therneau tests of each covariate and the global test (tests["ph_<term>"], tests["ph_global"], as Stata's estat phtest, detail) and Harrell's C (metrics["concordance"], as estat concordance). phtest= chooses the time function of the tests.
• oe.stcurve(cox, data=df, at={"benefit": 1}) returns the fitted survivor and cumulative hazard functions, with unnamed covariates at their sample means.

oe.streg(..., distribution="weibull") corresponds to streg age benefit, distribution(weibull). distribution may be "exponential", "weibull" (default), "gompertz", "lognormal", "loglogistic" or "ggamma", and metric="aft" is Stata's time option. ancillary= lets covariates shift the shape parameter, and strata= adds stratum indicators to both equations. Results carry extra["hazard_ratios"] or extra["time_ratios"] and the transformed shape parameter, such as extra["p"] for the Weibull. For Kaplan–Meier curves and the log-rank family, oe.sts(df, "weeks", failure="found_job", by="benefit", test="logrank") mirrors sts test benefit. The survival guide gives every formula and convention.

Supported options

Options of oe.stcox

Options of oe.streg

oe.sts

Stata command: sts.

Generated from oe.capabilities() of openecon 0.3.20a1. Stata names identify the corresponding workflow, not numerical parity.

Example

The script simulates 2,000 unemployment spells from a Weibull proportional-hazards model with shape p = 1.3, so the job-finding hazard rises with time. The index is −4.0 − 0.02·age − 0.5·benefit, where age is normal with mean 40 and standard deviation 10 and half the workers receive unemployment benefits; the true hazard ratios are exp(−0.02) = 0.980 per year of age and exp(−0.5) = 0.607 for benefits. Follow-up ends at 52 weeks, and independent exponential censoring with mean 80 weeks ends some spells earlier.

import numpy as np
import pandas as pd
import openecon as oe

BETA = {"Intercept": -4.0, "age": -0.02, "benefit": -0.5}
P = 1.3                                    # Weibull shape: the job-finding hazard rises over time


def simulate(n=2000, seed=21):
    """Unemployment spells from a Weibull proportional-hazards model, with censoring."""
    rng = np.random.default_rng(seed)
    age = rng.normal(40, 10, size=n)
    benefit = rng.integers(0, 2, size=n)
    index = BETA["Intercept"] + BETA["age"] * age + BETA["benefit"] * benefit
    spell = (-np.log(rng.uniform(size=n)) / np.exp(index)) ** (1 / P)   # S(t) = exp(-exp(index) t^P)
    censor = np.minimum(52.0, rng.exponential(80.0, size=n))             # follow-up ends after a year
    return pd.DataFrame({"weeks": np.minimum(spell, censor), "found_job": (spell <= censor).astype(int),
                         "age": age, "benefit": benefit})


if __name__ == "__main__":
    df = simulate()
    print(f"{len(df)} spells; {df.found_job.mean():.0%} end in a job, the rest are censored")
    km = oe.sts(df, "weeks", failure="found_job", by="benefit", test="logrank")
    medians = dict(zip(km["summary"]["group"], km["summary"]["median"].round(1)))
    print(f"Kaplan-Meier median weeks by benefit: {medians}; log-rank chi2 = {km.attrs['statistic']:.1f}")

    cox = oe.stcox(data=df, time="weeks", failure="found_job", x=["age", "benefit"])
    print(cox.summary())
    weibull = oe.streg(data=df, time="weeks", failure="found_job", x=["age", "benefit"],
                       distribution="weibull")
    print("\nHazard ratio      true   stcox   streg Weibull")
    for term in ["age", "benefit"]:
        cox_hr = cox.extra["hazard_ratios"][term]["hazard_ratio"]
        weibull_hr = weibull.extra["hazard_ratios"][term]["ratio"]
        print(f"  {term:<12}{np.exp(BETA[term]):8.4f}{cox_hr:8.4f}{weibull_hr:10.4f}")
    p = weibull.extra["p"]
    print(f"Weibull shape p = {p['estimate']:.3f} (95% CI {p['ci_low']:.3f} to {p['ci_high']:.3f}), true {P}")

With openecon 0.3.20a1 from PyPI the script prints:

2000 spells; 50% end in a job, the rest are censored
Kaplan-Meier median weeks by benefit: {0: 29.1, 1: 46.2}; log-rank chi2 = 84.9
Cox proportional hazards regression — weeks
Observations: 2000  |  Covariance: nonrobust  |  Confidence: 95%

Term       Estimate  Std. error         z     P>|stat|    CI lower     CI upper
-------  ----------  ----------  --------  -----------  ----------  -----------
age      -0.0161665  0.00335245   -4.8223  1.41913e-06  -0.0227372  -0.00959583
benefit   -0.585539   0.0645212  -9.07514  1.13527e-19   -0.711999     -0.45908

log_likelihood: -6919.87  |  log_likelihood_null: -6973.78  |  aic: 13843.7  |  bic: 13854.9
df_model: 2  |  n_subjects: 2000  |  n_failures: 1004  |  time_at_risk: 53679
concordance: 0.596206
LR chi2 test against b = 0: chi2(2) = 107.806, p = 3.893e-24
Score test of b = 0: chi2(2) = 108.051, p = 3.443e-24
PH test (scaled Schoenfeld residuals) for age: chi2(1) = 0.515217, p = 0.4729
PH test (scaled Schoenfeld residuals) for benefit: chi2(1) = 0.378188, p = 0.5386
Global test of proportional hazards (Grambsch-Therneau): chi2(2) = 0.893869, p = 0.6396

Hazard ratio      true   stcox   streg Weibull
  age           0.9802  0.9840    0.9838
  benefit       0.6065  0.5568    0.5561
Weibull shape p = 1.321 (95% CI 1.254 to 1.392), true 1.3

Reading the output

• Half the spells are censored. 1,004 of 2,000 spells end in a job within the follow-up. The Kaplan–Meier medians, 29.1 weeks without benefits and 46.2 with them, use the censored spells correctly; the log-rank statistic of 84.9 says the two curves differ.
• The Cox model recovers the effects. The benefit coefficient is −0.5855 with an interval from −0.712 to −0.459 that contains the true −0.5; as a hazard ratio, 0.557 against a true 0.607, receiving benefits lowers the weekly job-finding rate by about 44%. Each year of age lowers it by about 1.6% (hazard ratio 0.984, true 0.980).
• The proportional-hazards tests do not reject. The Grambsch–Therneau tests give p-values of 0.47 for age, 0.54 for benefits and 0.64 globally, as they should for data generated with proportional hazards.
• The Weibull model agrees with Cox and adds the shape. Its hazard ratios, 0.9838 and 0.5561, almost equal the Cox ones, because the data are Weibull. The estimated shape p = 1.321, with an interval from 1.254 to 1.392, shows the hazard rising with time, which the Cox model leaves unestimated.
• Harrell's C of 0.596 means that in about 59.6% of comparable pairs the worker with the higher predicted hazard found a job first. Age and benefits explain only a modest part of who leaves unemployment early.

Limitations

• The common oe.predict and oe.margins adapters do not cover oe.stcox; use its own result tables and family helpers.
• Cluster-robust covariance for oe.stcox allows at most 1 clustering dimension.
• The common oe.predict and oe.margins adapters do not cover oe.streg; use its own result tables and family helpers.
• Cluster-robust covariance for oe.streg allows at most 1 clustering dimension.
• Not implemented: Stata's exactm ties, competing-risks regression (stcrreg), interval-censored survival regression (stintreg), stcurve after streg, and a standard error for Harrell's C.
• stcox weights require ties="breslow"; ties="exactp" allows neither robust or cluster covariance nor tvc.
• After tvc, the baseline functions and proportional-hazards tests are not computed. Harrell's C is not computed with weights, delayed entry, multiple records per subject or tvc.
• Shared frailty and recurrent-event models are separate functions (oe.stcox_frailty, oe.streg_frailty, oe.stcox_recurrent), not options of oe.stcox and oe.streg.
• No numerical parity with Stata is claimed: the inventory's stata_parity_validated flag is False.

References

• Cox, D. R. (1972). Regression models and life-tables. Journal of the Royal Statistical Society Series B: Statistical Methodology, 34(2), 187–202. doi:10.1111/j.2517-6161.1972.tb00899.x
• Grambsch, P. M. and Therneau, T. M. (1994). Proportional hazards tests and diagnostics based on weighted residuals. Biometrika, 81(3), 515–526. doi:10.1093/biomet/81.3.515
• Kaplan, E. L. and Meier, P. (1958). Nonparametric estimation from incomplete observations. Journal of the American Statistical Association, 53(282), 457–481. doi:10.1080/01621459.1958.10501452
• Lin, D. Y. and Wei, L. J. (1989). The robust inference for the Cox proportional hazards model. Journal of the American Statistical Association, 84(408), 1074–1078. doi:10.1080/01621459.1989.10478874