Cox model diagnostics

Fitting a Cox model is only half the work. Before you trust its hazard ratios, you should check that its central assumption holds, look at residuals for influential or poorly fit observations, and understand what the model predicts. This page covers the proportional hazards test, residuals, baseline hazard and survival prediction, stratification, robust standard errors, and the concordance index. Together these turn a fitted model into a defensible one.

We begin with the lung outcomes. The response y is a Surv object pairing each follow-up time with an event indicator, and it is the target every model on this page is fit against.

import greenwood as gw

# Load the bundled lung dataset as a Polars DataFrame
lung = gw.load_dataset("lung", backend="polars")

# Build a right-censored response; status == 2 marks a death
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Display the response summary
y
Surv(type=right, n=228, events=165)

The printed response records 165 events among 228 subjects, with a + marking each censored observation. We fit a Cox model to it using age and sex as covariates. This is the model we diagnose throughout the page, so we fit it once here and reuse it below.

# Fit the model we will diagnose throughout this page
cox = gw.CoxPH().fit(y, lung[["age", "sex"]])

With the response and model in hand, the sections below work through each diagnostic in turn.

Checking the proportional hazards assumption

The Cox model assumes each covariate’s effect on the hazard is constant over time. When this fails, for example if a treatment helps early but not late, the reported hazard ratio is a misleading average. The Grambsch-Therneau test checks the assumption by looking for a trend in the scaled Schoenfeld residuals against time. A small p-value is evidence that the assumption is violated for that covariate.

Calling cox_zph() returns a ZPHResult, an object that bundles the test statistics together. Displaying it directly gives a compact summary of the test.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit a Cox model with age and sex as covariates
cox = gw.CoxPH().fit(y, lung[["age", "sex"]])

# Run the Grambsch-Therneau proportional hazards test
zph = cox.cox_zph()

# Print per-covariate test statistics and p-values
zph
ZPHResult(transform='identity', age: p=0.7065, sex: p=0.0992, GLOBAL p=0.2425)

For the full per-covariate breakdown, call to_frame() on that result. This is the table you will read most often.

# Export the ZPH test table as a tidy DataFrame (one row per term plus GLOBAL)
zph.to_frame(format="polars")
PolarsRows3Columns4
term
str
chisq
f64
df
i64
p_value
f64
0 age 0.141747561345 1 0.70654984373
1 sex 2.71840108815 1 0.0991973377972
2 GLOBAL 2.8333913831 2 0.242514035592

The table has one row per covariate plus a GLOBAL row that tests all covariates jointly. Large p-values, as here, are reassuring: they give no evidence against proportional hazards. The test uses a time transform, which defaults to "identity". Other options are "log", "km" (maps time to the 1 - KM(t-) probability scale), and "rank" (uses average ranks of all exit times). The "km" transform is the default in R’s cox.zph() and is often a good choice when the time scale is skewed.

WarningA non-significant test is not proof

Failing to reject the proportional hazards assumption does not prove it holds, especially in small samples. Complement the test by plotting scaled Schoenfeld residuals or by fitting time-stratified models when you have reason to suspect a time-varying effect.

Visualizing the proportional hazards check

The plot_schoenfeld() function plots the scaled Schoenfeld residuals against time for each covariate. Under proportional hazards, the residuals should scatter randomly around the coefficient estimate (the dashed line) and the loess smooth should be flat. A trend in the smooth is visual evidence that the covariate’s effect changes over time.

# Scaled Schoenfeld residual plots with Grambsch-Therneau p-values
gw.plot_schoenfeld(cox)

Each panel shows one covariate. The blue points are scaled Schoenfeld residuals at each event time, the red curve is a loess smooth, and the dashed grey line marks the overall coefficient estimate. The p-value in each panel title is from the cox_zph() test. A flat smooth close to the reference line confirms proportional hazards; a curve that departs from it reveals where and how the effect changes.

The time axis can be transformed to spread out events more evenly, which sometimes makes trends easier to see. Pass transform="log" for log time, transform="km" for the Kaplan-Meier failure probability, or transform="rank" for time ranks:

# Log-time axis to spread out early events
gw.plot_schoenfeld(cox, transform="log")

To suppress the p-value annotation (for example, when you want a cleaner figure for a report), pass show_zph=False:

# Clean version without p-value annotations
gw.plot_schoenfeld(cox, show_zph=False, title="PH diagnostic")

Localizing violations with time windows

A global cox_zph() test tells you whether proportional hazards is violated, but not when. If a treatment helps early but fades later, you want to know where the transition happens. The breaks= parameter splits the follow-up into time windows and runs the Grambsch-Therneau test within each one.

# Test PH within three windows: (0, 180], (180, 365], and (365, inf]
zph_w = cox.cox_zph(breaks=[180, 365])
zph_w
ZPHResult(transform='identity', age: p=0.7065, sex: p=0.0992, GLOBAL p=0.2425)
  3 time windows

The printed summary shows the global result plus a count of windows. To see per-window statistics, pass detail="windows" to to_frame().

# Per-window test statistics, one row per term per window plus a GLOBAL row per window
zph_w.to_frame(detail="windows", format="polars")
PolarsRows9Columns6
window
str
n_events
i64
term
str
chisq
f64
df
i64
p_value
f64
0 (0.0, 180] 63 age 6.89146891265 1 0.00866080701024
1 (0.0, 180] 63 sex 6.53458002382e-05 1 0.99355023083
2 (0.0, 180] 63 GLOBAL 6.96432126283 2 0.0307409194304
3 (180, 365] 58 age 1.96656714823 1 0.160812711741
4 (180, 365] 58 sex 0.1024253971 1 0.748938250496
5 (180, 365] 58 GLOBAL 2.12258986012 2 0.346007464638
6 (365, inf] 44 age 0.0726370475837 1 0.787535238825
7 (365, inf] 44 sex 1.10705994029 1 0.292721946329
8 (365, inf] 44 GLOBAL 1.20170438374 2 0.548344142504

In this example the age effect shows a significant departure from proportional hazards in the first six months (p = 0.009 in the (0, 180] window) but not later. This pattern would be invisible in the overall test, where the age p-value is 0.71. The windowed view makes it actionable: you might consider a time-varying coefficient for age in the early period, or stratify the analysis at 180 days.

Windows with fewer than two events are automatically skipped. Choosing breaks= is a judgment call. Clinical milestones (30-day mortality, 1-year landmark) or quantiles of the event-time distribution are both reasonable starting points. The global test result in the ZPHResult is unchanged by breaks=. It always reflects the full follow-up.

Smooth non-linear hazard ratio curves

A Cox model reports a single coefficient per covariate, implying a linear relationship between the covariate and the log-hazard. But the true effect may be non-linear. For instance, age might have little effect in the middle of its range but increase the hazard steeply at extremes. The smooth_hr() method checks for this by refitting the model with a flexible B-spline basis for the covariate of interest and plotting the resulting log-hazard ratio curve.

# Compute a smooth log-HR curve for age, holding sex at its mean
shr = cox.smooth_hr("age")
shr
SmoothHRResult(term='age', reference=62.44736842105263, df=4, grid=200 points)

The result is a SmoothHRResult containing the evaluation grid, log-HR, HR, and confidence bands. Pass it to plot_smooth_hr() for a publication-ready visualization.

# Plot the smooth log-HR curve with 95% confidence band
gw.plot_smooth_hr(shr)

The dashed line at zero marks the null (no effect). Where the confidence band includes zero, the effect is not statistically significant at the chosen confidence level. Here the curve is roughly linear, which is consistent with the single-coefficient model. A strongly curved shape would suggest that the linear specification is inadequate.

You can also plot on the hazard ratio scale. This is often more interpretable for clinical audiences.

# Same curve on the HR scale (HR = 1 is the reference)
gw.plot_smooth_hr(shr, scale="hr")

The df parameter controls the flexibility of the spline. Higher values allow more wiggly curves. The default (df=4) is a good starting point. The reference parameter sets the covariate value where the log-HR is zero; it defaults to the weighted mean.

# A more flexible curve (df=6) with age 50 as the reference
shr6 = cox.smooth_hr("age", df=6, reference=50)
gw.plot_smooth_hr(shr6)

The tidy DataFrame is available via to_frame() for further analysis.

# Export the smooth curve as a Polars DataFrame
shr.to_frame(format="polars")
PolarsRows200Columns4
age
f64
log_hr
f64
log_hr_lower
f64
log_hr_upper
f64
0 39 -1.25066193497 -3.09508407184 0.593760201895
1 39.216080402 -1.19942650469 -2.97007560168 0.5712225923
2 39.432160804 -1.14942827631 -2.84788501739 0.549028464761
3 39.648241206 -1.10065320929 -2.72863335818 0.527326939612
4 39.864321608 -1.05308726304 -2.61433210906 0.50815758298
5 40.0804020101 -1.00671639701 -2.50268896456 0.489256170544
6 40.2964824121 -0.961526570634 -2.39554720661 0.472494065346
7 40.5125628141 -0.917503743347 -2.29249807345 0.457490586754
197 81.567839196 1.0169020714 -0.245907457412 2.27971160022
198 81.783919598 1.04038530439 -0.286270694452 2.36704130323
199 82 1.0641151679 -0.329459959805 2.4576902956

Residuals

Residuals reveal how individual observations relate to the fitted model. The residuals() method supports seven types, each answering a different diagnostic question.

Martingale and deviance residuals

Martingale residuals, one per subject, are the difference between the observed number of events and the number the model expected. Large negative values flag subjects who lived much longer than predicted. Deviance residuals are a normalized version that is more symmetrically distributed and better for spotting outliers.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit a Cox model with age and sex as covariates
cox = gw.CoxPH().fit(y, lung[["age", "sex"]])

# Compute martingale residuals (one per subject); peek at the first five
cox.residuals("martingale")[:5]
array([ 0.00438999, -0.50576203, -3.12981924,  0.53275397, -2.35065745])
# Deviance residuals are more symmetric, useful for identifying outliers
cox.residuals("deviance")[:5]
array([ 0.00439643, -0.43923326, -2.50192695,  0.67549295, -1.51096053])

Schoenfeld and scaled Schoenfeld residuals

Schoenfeld residuals have one row per event and one column per covariate. They underlie the proportional hazards test and are useful for spotting time trends in a covariate’s effect. Scaled Schoenfeld residuals (Grambsch-Therneau) are centered on the coefficient estimate, so under the PH assumption, each row’s expected value equals the true coefficient.

# Compute Schoenfeld residuals (one row per event, one column per covariate)
cox.residuals("schoenfeld", format="polars")
PolarsRows165Columns2
age
f64
sex
f64
0 0.935464711892 0.727079613965
1 10.0052299319 -0.272302937094
2 17.0052299319 -0.272302937094
3 3.00522993189 -0.272302937094
4 10.1404536933 -0.275789612537
5 12.227706252 -0.278411033642
6 1.22770625199 -0.278411033642
7 5.29450657479 -0.280979164464
162 -2.90825326766 -0.155341657484
163 0.611217430947 -0.196097246718
164 -4.65909168016 -0.171473016943
# Scaled Schoenfeld residuals: centered on beta, used by cox_zph()
cox.residuals("scaledsch", format="polars")
PolarsRows165Columns2
age
f64
sex
f64
0 0.0403858886709 2.86408456074
1 0.153658390962 -1.63265384459
2 0.251912824506 -1.53435602268
3 0.0554039574181 -1.7309516665
4 0.155507476888 -1.64688765923
5 0.184768068006 -1.6297065243
6 0.030368243866 -1.78417453017
7 0.087415203921 -1.73894893284
162 -0.0259573144402 -1.27281710834
163 0.0228708875903 -1.41196917521
164 -0.0507592164325 -1.37204245852

Score, dfbeta, and dfbetas residuals

Score residuals (efficient score contributions) have one row per observation and one column per covariate. They are the building blocks for the robust sandwich variance. Dfbeta residuals approximate the change in each coefficient if observation i were deleted. Dfbetas standardizes dfbeta by the coefficient standard error, making values comparable across covariates on different scales.

# Score residuals: one row per subject, one column per covariate
cox.residuals("score", format="polars")
PolarsRows228Columns2
age
f64
sex
f64
0 -0.033840982067 -0.0467194061563
1 -1.7737935818 0.149432139267
2 21.5560164909 0.967719839093
3 -3.58464005499 -0.175211113143
4 5.02675449886 0.864528565573
5 -47.2045902525 1.31521916507
6 2.08357635267 0.252774580847
7 2.01860354501 0.174753442068
225 -1.52826223923 -0.0945906530503
226 -0.938449449516 0.107978170568
227 1.10434148047 -0.14307750502
# Dfbeta: approximate leave-one-out influence on each coefficient
cox.residuals("dfbeta", format="polars")
PolarsRows228Columns2
age
f64
sex
f64
0 -6.85492847608e-06 -0.00131299357242
1 -0.000138176822328 0.00403944024828
2 0.00191610298024 0.0289715170777
3 -0.000319852509834 -0.00521837526568
4 0.000501196702015 0.0246710649503
5 -0.00390370311399 0.0328641884291
6 0.000198760004338 0.00726567311062
7 0.000186592753759 0.00507226154139
225 -0.000138057658692 -0.00278259189815
226 -7.06431169614e-05 0.00294807412698
227 8.17681727223e-05 -0.00391821713747
# Dfbetas: standardized dfbeta, comparable across covariates
cox.residuals("dfbetas", format="polars")
PolarsRows228Columns2
age
f64
sex
f64
0 -0.000743220776571 -0.00784073539384
1 -0.0149813211842 0.0241221151353
2 0.207746521343 0.173007700978
3 -0.0346788491779 -0.031162300032
4 0.0543404359924 0.147326914787
5 -0.423244862438 0.19625336393
6 0.0215498331296 0.04338804204
7 0.0202306430819 0.0302897603078
225 -0.0149684012989 -0.0166166592439
226 -0.00765922393365 0.0176048608588
227 0.0088654177854 -0.0233982133925

Martingale and deviance residuals are most useful for spotting outliers or a nonlinear covariate relationship. Schoenfeld residuals are the basis for the proportional hazards test. Score, dfbeta, and dfbetas residuals assess the influence of individual observations on the fitted coefficients.

Leverage

Leverage measures how much each observation influences the fitted values. It is the diagonal of the hat matrix, computed as h_i = L_i^\top V L_i where L_i is the score residual for observation i and V is the variance-covariance matrix. Values near zero have little influence on the fit. Values close to one dominate it.

# Leverage (hat-matrix diagonal): one value per observation
cox.residuals("leverage")[:10]
array([6.15742575e-05, 8.48719358e-04, 6.93398593e-02, 2.06087346e-03,
       2.38482332e-02, 2.27496316e-01, 2.25070912e-03, 1.26305196e-03,
       4.15570837e-03, 1.59745075e-03])

The sum of all leverage values approximately equals the number of covariates (here, 2).

Influence diagnostics

The influence_diagnostics() method combines leverage, residuals, and influence measures into a single table for identifying observations that are unusual, poorly fit, or disproportionately influential.

# All influence diagnostics in one table
cox.influence_diagnostics(format="polars")
PolarsRows228Columns8
leverage
f64
martingale
f64
deviance
f64
dfbeta_age
f64
dfbeta_sex
f64
dfbetas_age
f64
dfbetas_sex
f64
ld
f64
0 6.15742575023e-05 0.00438999436025 0.00439643487827 -6.85492847608e-06 -0.00131299357242 -0.000743220776571 -0.00784073539384 6.15742575023e-05
1 0.000848719358341 -0.505762029951 -0.439233258531 -0.000138176822328 0.00403944024828 -0.0149813211842 0.0241221151353 0.000848719358341
2 0.0693398592849 -3.12981924011 -2.50192695341 0.00191610298024 0.0289715170777 0.207746521343 0.173007700978 0.0693398592849
3 0.00206087345754 0.532753970241 0.675492945159 -0.000319852509834 -0.00521837526568 -0.0346788491779 -0.031162300032 0.00206087345754
4 0.0238482331693 -2.35065744728 -1.51096053409 0.000501196702015 0.0246710649503 0.0543404359924 0.147326914787 0.0238482331693
5 0.22749631643 -4.25370864739 -2.9167477256 -0.00390370311399 0.0328641884291 -0.423244862438 0.19625336393 0.22749631643
6 0.00225070912001 0.442667491621 0.532777181506 0.000198760004338 0.00726567311062 0.0215498331296 0.04338804204 0.00225070912001
7 0.00126305195764 0.290003990946 0.324012157388 0.000186592753759 0.00507226154139 0.0202306430819 0.0302897603078 0.00126305195764
225 0.000474195491434 -0.133128811149 -0.516001571992 -0.000138057658692 -0.00278259189815 -0.0149684012989 -0.0166166592439 0.000474195491434
226 0.000384622645155 -0.366567124263 -0.856232590203 -7.06431169614e-05 0.00294807412698 -0.00765922393365 0.0176048608588 0.000384622645155
227 0.000650908717076 -0.203602296262 -0.6381258438 8.17681727223e-05 -0.00391821713747 0.0088654177854 -0.0233982133925 0.000650908717076

The ld column is the likelihood displacement, which measures how much the log-likelihood would change if each observation were removed. It is the quadratic form \Delta\hat{\beta}_i^\top V^{-1} \Delta\hat{\beta}_i using the dfbeta vector. Large values of ld flag observations whose removal would materially change the fitted coefficients. Combining ld with leverage and deviance gives a multi-dimensional view of influence: high leverage means the observation is unusual in covariate space, a large deviance residual means the outcome was unexpected given the model, and a large likelihood displacement means the observation actually pulled the coefficients.

Visualizing influence

Scanning a table for outliers is tedious when you have hundreds of observations. plot_influence() turns the same diagnostics into a row of scatter plots, with the most influential observations highlighted in red and labeled by observation number.

# Default: deviance, leverage, and likelihood displacement panels
gw.plot_influence(cox)

Each panel plots one diagnostic against the linear predictor. Observations in the upper corners of the likelihood displacement panel deserve the closest attention, since removing them would shift the fitted coefficients the most.

You can choose which panels to show and how many observations to highlight.

# Show only deviance and leverage, highlighting the top 5
gw.plot_influence(cox, panels=["deviance", "leverage"], highlight=5)

When no observations stand out, the model fit is stable across the dataset. When a few points dominate the likelihood displacement, investigate whether they reflect data-entry errors, a subpopulation the model handles poorly, or genuinely extreme cases that should remain.

Baseline hazard and predicted survival

Although the Cox model does not assume a shape for the baseline hazard, it can estimate one after fitting. The baseline cumulative hazard, and the survival curve derived from it, describe a reference subject.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit a Cox model with age and sex as covariates
cox = gw.CoxPH().fit(y, lung[["age", "sex"]])

# Retrieve the Breslow cumulative baseline hazard and implied survival
cox.baseline_hazard(format="polars")
PolarsRows186Columns3
time
f64
cumhaz
f64
survival
f64
0 5 0.00295531012895 0.997049052501
1 11 0.0119063396838 0.988164260305
2 12 0.0149282203866 0.985182653095
3 13 0.0210294286393 0.979190147912
4 15 0.0241081722771 0.976180108421
5 26 0.0272054191316 0.973161315038
6 30 0.0303227228765 0.970132399105
7 31 0.0334599993722 0.967093594809
183 965 2.01305845682 0.133579502131
184 1010 2.01305845682 0.133579502131
185 1022 2.01305845682 0.133579502131

More useful in practice is predicting the survival curve for specific covariate values. First we build a small frame of the subjects we want predictions for: a 50-year-old with sex 1 and a 70-year-old with sex 2. Showing it confirms the covariate values before we predict from them.

import pandas as pd

# Two new subjects: a 50-year-old male and a 70-year-old female
newdata = pd.DataFrame({"age": [50, 70], "sex": [1, 2]})

# Preview the covariate frame before predicting
newdata
PandasRows2Columns2
age
i64
sex
i64
0 50 1
1 70 2

Given that new data and a set of times, predict returns the estimated survival probability for each subject.

# Predict survival probabilities at specific times for each subject in newdata
cox.predict(newdata, type="survival", times=[180, 365, 730], format="polars")
PolarsRows3Columns3
time
f64
subject_1
f64
subject_2
f64
0 180 0.7297626204 0.767074757103
1 365 0.417781082587 0.479674113531
2 730 0.119648956441 0.167440538104

Each column is a subject from newdata, and each row is a requested time. The type="lp" and type="risk" options instead return the linear predictor and the relative risk.

Confidence bands and conditional survival

Pass ci=True to add a pointwise confidence band around each predicted curve. Every subject then gains _lower and _upper columns, built from the standard error of the cumulative hazard and matching R’s survfit (which combines baseline-hazard and coefficient uncertainty).

# Add pointwise 95% confidence bands; each subject gains _lower and _upper columns
cox.predict(newdata, type="survival", times=[180, 365, 730], ci=True, format="polars")
PolarsRows3Columns7
time
f64
subject_1
f64
subject_1_lower
f64
subject_1_upper
f64
subject_2
f64
subject_2_lower
f64
subject_2_upper
f64
0 180 0.7297626204 0.65075626184 0.818360903092 0.767074757103 0.696725970364 0.844526697744
1 365 0.417781082587 0.312492619485 0.558544497004 0.479674113531 0.379184415457 0.6067951261
2 730 0.119648956441 0.0574578159532 0.249154489079 0.167440538104 0.0910524175713 0.307914216319

You can also predict conditional on a subject having already survived some time, which is useful for updating a prognosis partway through follow-up. With conditional_after=180, the returned probabilities are P(T > t \mid T > 180) = S(t) / S(180).

# Predict P(T > t | T > 180): survival conditional on being alive at day 180
cox.predict(newdata, type="survival", times=[365, 730], conditional_after=180, format="polars")
PolarsRows2Columns3
time
f64
subject_1
f64
subject_2
f64
0 365 0.568989927935 0.622110389604
1 730 0.162953886474 0.217160975455

Confidence bands quantify the uncertainty in a predicted survival curve, and conditional survival gives an updated estimate that is directly meaningful for a patient already some time into follow-up. Both are straightforward extensions of the base predict() call.

Median survival time

A common clinical summary is the predicted median survival time: the earliest time at which the survival curve drops to 0.5. Use predict_median() for a direct answer without having to inspect the full curve.

# Predicted median survival time for each subject
cox.predict_median(newdata, format="polars")
PolarsRows1Columns3
p
f64
subject_1
f64
subject_2
f64
0 0.5 320 361

The result is a single-row frame with one column per subject. When the survival curve never drops below 0.5 (heavy censoring or low risk), the median is NaN.

With ci=True, confidence intervals are obtained by inverting the pointwise confidence bands of the survival curve. The lower bound of the median comes from the lower survival band, and the upper bound from the upper survival band.

# Median survival with confidence intervals
cox.predict_median(newdata, ci=True, format="polars")
PolarsRows1Columns7
p
f64
subject_1
f64
subject_1_lower
f64
subject_1_upper
f64
subject_2
f64
subject_2_lower
f64
subject_2_upper
f64
0 0.5 320 268 442 361 301 519

Survival quantiles

predict_median() is really a special case of predict_quantile(), which generalizes the same inversion to any set of failure probabilities. Pass a list of values in (0, 1) via p= to get the predicted survival time at each quantile, one column per subject.

# Predicted survival-time quantiles at the 25th, 50th, and 75th percentiles
cox.predict_quantile(newdata, p=[0.25, 0.5, 0.75], format="polars")
PolarsRows3Columns3
p
f64
subject_1
f64
subject_2
f64
0 0.25 176 186
1 0.5 320 361
2 0.75 558 643

The result has one row per requested p and one column per subject. As with predict_median(), pass ci=True to add _lower and _upper columns built by inverting the pointwise confidence bands of the survival curve.

# Survival quantiles with confidence intervals
cox.predict_quantile(newdata, p=[0.25, 0.5, 0.75], ci=True, format="polars")
PolarsRows3Columns7
p
f64
subject_1
f64
subject_1_lower
f64
subject_1_upper
f64
subject_2
f64
subject_2_lower
f64
subject_2_upper
f64
0 0.25 176 144 223 186 163 268
1 0.5 320 268 442 361 301 519
2 0.75 558 442 728 643 520 814

Expected survival time

predict_expectation() computes the restricted mean survival time (RMST): the area under the predicted survival curve up to a time horizon tau. It answers a different question than the quantile methods above. Instead of “when does survival reach some probability”, it asks “how much event-free time is expected, on average, up to tau”. Unlike predict_median() and predict_quantile(), tau is required and must be specified explicitly, since RMST is only defined relative to a chosen horizon.

# Expected survival time (RMST) up to 365 days for each subject
cox.predict_expectation(newdata, tau=365, format="polars")
PolarsRows1Columns3
tau
f64
subject_1
f64
subject_2
f64
0 365 265.67633284 278.322690708

The result has a tau column alongside one column per subject. As with the other prediction methods, pass ci=True to add _lower and _upper columns.

# Expected survival time with confidence intervals
cox.predict_expectation(newdata, tau=365, ci=True, format="polars")
PolarsRows1Columns7
tau
f64
subject_1
f64
subject_1_lower
f64
subject_1_upper
f64
subject_2
f64
subject_2_lower
f64
subject_2_upper
f64
0 365 265.67633284 239.524155511 296.251450048 278.322690708 254.747557683 305.17385481

Stratification

Sometimes a variable violates the proportional hazards assumption but is not of direct interest, such as enrolling hospital. Stratification lets each level of that variable have its own baseline hazard while sharing the coefficients of the covariates you do care about. You pass the stratifying variable to strata. Here we stratify by sex while keeping age and ph.ecog as covariates.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit a stratified model: separate baselines per sex, shared age and ECOG coefficients
stratified = gw.CoxPH().fit(y, lung[["age", "ph.ecog"]], strata=lung["sex"])

As with any fitted model, we can print stratified for a summary

stratified
CoxPH (Cox proportional hazards model, ties='efron')

            coef  exp(coef)  se(coef)      z          p
age      0.01057      1.011  0.009241  1.143     0.2529
ph.ecog   0.4624      1.588    0.1148  4.029  5.591e-05

n = 227, events = 164
Likelihood ratio test = 19.48 on 2 df, p = 5.895e-05

or read its coefficient table as a DataFrame with to_frame().

# Export the coefficient table; sex does not appear (absorbed into the baselines)
stratified.to_frame(format="polars")
PolarsRows2Columns7
term
str
estimate
f64
std_error
f64
statistic
f64
p_value
f64
conf_low
f64
conf_high
f64
0 age 0.0105662546009 0.00924137389309 1.14336404123 0.252887475744 -0.00754650539718 0.0286790145991
1 ph.ecog 0.462424434358 0.114761097855 4.0294528634 5.59068243623e-05 0.237496815737 0.68735205298

Notice that the stratifying variable, sex, does not appear as a coefficient. Its effect is absorbed into the separate baselines, which is exactly the point of stratification.

Predicted survival from a stratified model

Survival predictions use the correct stratum-specific baseline for each subject. When predicting on the training data, stratum assignments are taken automatically from the fitted model.

# Predict survival from the stratified model at two time points
# Each subject gets curves from their own stratum's baseline hazard
stratified.predict(type="survival", times=[180, 365], format="polars").select(
    ["time", "subject_1", "subject_2", "subject_3"]
)
PolarsRows2Columns4
time
f64
subject_1
f64
subject_2
f64
subject_3
f64
0 180 0.617897046625 0.752344813136 0.778276506304
1 365 0.285735961821 0.476911220712 0.520872064146

When predicting for new subjects, pass strata= to specify which stratum each subject belongs to.

# Two new subjects: a 50-year-old (stratum sex=1) and a 70-year-old (stratum sex=2)
newdata = pd.DataFrame({"age": [50, 70], "ph.ecog": [1, 2]})

# strata= maps each new subject to a stratum seen at fit time
stratified.predict(
    newdata, type="survival", strata=pd.Series([1, 2]), times=[180, 365], format="polars"
)
PolarsRows2Columns3
time
f64
subject_1
f64
subject_2
f64
0 180 0.688254178057 0.747673798296
1 365 0.378288748953 0.335772254241

Each subject’s survival curve is computed against the baseline hazard of its stratum.

Robust and clustered standard errors

When observations are not independent, for example when the same subject contributes several rows or when subjects are grouped into clusters, the model-based standard errors are too small. The robust, or sandwich, variance corrects for this. Set robust=True for the Lin-Wei sandwich estimator, and pass cluster to group correlated observations.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit with the Lin-Wei sandwich variance; coefficients are unchanged, SEs are adjusted
robust = gw.CoxPH().fit(y, lung[["age", "sex"]], robust=True)

# Show only the standard errors to compare against the model-based values
robust.to_frame(format="polars")[["term", "std_error"]]
PolarsRows2Columns2
term
str
std_error
f64
0 age 0.00948922108421
1 sex 0.159919372277

The coefficients are unchanged; only the standard errors, and therefore the p-values and intervals, are adjusted. The model-based standard errors remain available as naive_std_error_ for comparison.

The concordance index

The concordance index, or C-statistic, measures how well the model’s risk ordering agrees with the observed order of events. It is the probability that, for a random comparable pair of subjects, the one who failed first had the higher predicted risk. A value of 0.5 is no better than chance and 1.0 is perfect discrimination.

# Load data and build the response
lung = gw.load_dataset("lung", backend="polars")
y = gw.Surv.right(lung["time"], event=(lung["status"] == 2))

# Fit a Cox model with age and sex as covariates
cox = gw.CoxPH().fit(y, lung[["age", "sex"]])

# Compute Harrell's C-statistic: probability of correct pairwise risk ordering
cox.concordance()
0.6028530028979714

The C-statistic summarizes discrimination across all observed pairs, but it does not say whether predicted probabilities are well calibrated. For a more complete picture of model performance, including calibration and the Brier score, see Prediction performance.

Next steps

You can now validate a Cox model and use it to predict.

  • Prediction performance extends discrimination and adds the Brier score for calibration, including for external risk scores.
  • Parametric survival models provide an alternative when the proportional hazards assumption is untenable. That page also covers AFT residuals (response, Cox-Snell, martingale, deviance, score, dfbeta, dfbetas) for diagnosing parametric models.
  • Revisit Cox regression for the modeling basics if any of this is unfamiliar.