Visualizing survival

A survival curve is far easier to understand as a picture than as a table of numbers. By default Greenwood draws curves with Altair, which produces interactive charts (tooltips and zoom) and is Narwhals-native: the numbers flow straight from your Polars fit into the chart with no Pandas or Matplotlib in the way, so a pure Polars/Narwhals workflow stays that way. This page covers the survival curve plot, confidence bands, censoring marks, and the numbers-at-risk table.

Altair is an optional dependency. If you installed Greenwood with the altair extra, or with all, you already have it. We use the lung dataset throughout. First we build the response.

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)

We then fit a Kaplan-Meier estimator stratified by sex, one curve per group, using the log-log confidence interval transform. Printing the fitted estimator summarizes each stratum.

# Fit Kaplan-Meier with log-log CI transform, stratified by sex
km = gw.KaplanMeier(conf_type="log-log").fit(y, by=lung["sex"])

# Print per-stratum median survival
km
KaplanMeier (Kaplan-Meier survival estimate)

     n  events  median  0.95LCL  0.95UCL
1  138     112     270      210      306
2   90      53     426      345      524

Everything below draws from this fitted km object. Before plotting, it also helps to see the numbers the curve is built from. Calling to_frame() on the fitted estimator returns the step-function estimates in tidy form, one row per stratum and time, with the survival probability and its confidence limits. The preview shows the first and last rows along with the full dimensions of the table.

# Export the step-function estimates as a tidy DataFrame (one row per stratum and time)
km.to_frame(format="polars")
PolarsRows206Columns9
strata
obj
time
f64
n_risk
f64
n_event
f64
n_censor
f64
estimate
f64
std_error
f64
conf_low
f64
conf_high
f64
0 1 11 138 3 0 0.978260869565 0.0124139182765 0.934121582525 0.992936557511
1 1 12 135 1 0 0.971014492754 0.014281169221 0.924619139392 0.989021541626
2 1 13 134 2 0 0.95652173913 0.0173597697718 0.905787300104 0.980228948021
3 1 15 132 1 0 0.949275362319 0.0186795340253 0.896548509216 0.975490402087
4 1 26 131 1 0 0.942028985507 0.019892897278 0.887428391293 0.970579101988
5 1 30 130 1 0 0.934782608696 0.0210182948563 0.878418312188 0.965522040145
6 1 31 129 1 0 0.927536231884 0.0220691880472 0.869508906988 0.960339509261
7 1 53 128 2 0 0.913043478261 0.0239859688855 0.851957778533 0.949657700321
203 2 765 3 1 0 0.0832144435134 0.0499212744378 0.0185052512122 0.212363854832
204 2 821 2 0 1 0.0832144435134 0.0499212744378 0.0185052512122 0.212363854832
205 2 965 1 0 1 0.0832144435134 0.0499212744378 0.0185052512122 0.212363854832

Each row is a step in one of the curves. The time and estimate columns give the point at which the survival probability drops and its value after the drop, while conf_low and conf_high give the band that the plot will shade. The strata column identifies the group, so there is one set of steps per sex, and n_risk, n_event, and n_censor record how many subjects were at risk, failed, or were censored at each time. The plot turns exactly these numbers into a picture.

The survival curve plot

The main entry point is plot_survival(). Given a fitted KaplanMeier, it draws the step curve for each stratum, a shaded confidence band, and a small notch at each censoring time.

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

# Fit Kaplan-Meier stratified by sex with log-log CI
km = gw.KaplanMeier(conf_type="log-log").fit(y, by=lung["sex"])

# Draw the interactive survival curve with confidence band and censoring marks
gw.plot_survival(km)

Each element carries meaning. The steps show the estimated survival probability, the band shows its uncertainty, and the censoring notches (angled so they sit on the curve without obscuring it) show where subjects left the risk set without an event. Because the chart is interactive, you can hover a curve to read the survival probability at a time, and drag or scroll to zoom. Curves that separate and stay apart suggest a real group difference, which you would confirm with a log-rank test from Comparing groups.

TipPlotting Cox model baseline hazard with confidence intervals

If you have fit a CoxPH model, you can extract the baseline hazard with confidence intervals and plot it to visualize the estimated survival trajectory for a reference subject:

cox = gw.CoxPH().fit(y, lung[["age", "sex"]])
baseline = cox.baseline_hazard(ci=True, format="polars")

# Convert to a format suitable for Altair plotting
# The 'survival' column is S_0(t) and 'survival_lower'/'survival_upper' are the CI bounds
baseline
PolarsRows186Columns7
time
f64
cumhaz
f64
survival
f64
cumhaz_lower
f64
cumhaz_upper
f64
survival_lower
f64
survival_upper
f64
0 5 0.00295531012895 0.997049052501 0.000288674422755 0.0302550460651 0.970198056794 0.99971136724
1 11 0.0119063396838 0.988164260305 0.00244588031142 0.057959060386 0.943688580892 0.997557108417
2 12 0.0149282203866 0.985182653095 0.00325720424168 0.068418111784 0.933869929915 0.996748094693
3 13 0.0210294286393 0.979190147912 0.00493977733923 0.0895256685725 0.914364794384 0.995072403296
4 15 0.0241081722771 0.976180108421 0.00579712384423 0.100257297612 0.904604635478 0.994219647055
5 26 0.0272054191316 0.973161315038 0.00666433934671 0.111058994991 0.894885953583 0.993357818114
6 30 0.0303227228765 0.970132399105 0.00754101568953 0.121928870128 0.885211331232 0.992487346431
7 31 0.0334599993722 0.967093594809 0.00842633597775 0.132865762883 0.875582619907 0.991609066085
183 965 2.01305845682 0.133579502131 0.567558662941 7.1400625436 0.000792702516446 0.566907763585
184 1010 2.01305845682 0.133579502131 0.567558662941 7.1400625436 0.000792702516446 0.566907763585
185 1022 2.01305845682 0.133579502131 0.567558662941 7.1400625436 0.000792702516446 0.566907763585

The survival column represents the baseline survival probability (for a reference subject with all covariates at their mean values), and the survival_lower and survival_upper columns give the 95% confidence band. You can plot this as a curve with a confidence band using Altair, or integrate it into Great Tables for a publication-ready figure.

You can turn individual elements off. Removing the band and the censoring marks produces a cleaner curve for a slide or a small figure.

# Strip the confidence band and censoring marks for a cleaner figure
gw.plot_survival(km, conf_int=False, censor_marks=False)

Adding a numbers-at-risk table

A survival plot is much easier to trust when it is accompanied by the number of subjects still at risk at regular intervals. As the risk set shrinks, the curve becomes less certain, and the table makes that visible. Pass risk_table=True to stack an aligned table beneath the curve.

# Append an aligned numbers-at-risk table beneath the curve
gw.plot_survival(km, risk_table=True)

The table shares the horizontal axis with the curve above it, so the counts line up with the times they describe. And because they share one x scale, panning or zooming the curve moves the table with it. You control the times shown with the times= argument.

# Show risk counts at specific time points; the table x-axis stays synced with the curve
gw.plot_survival(km, risk_table=True, times=[0, 250, 500, 750, 1000])

If you only need the numbers, get_risk_table_frame() returns them as a tidy frame, which you can render however you like, including through Great Tables. It is backend-neutral, so it works whether or not you have a plotting library installed.

# Return the risk counts as a tidy DataFrame (backend-neutral, no plotting required)
gw.get_risk_table_frame(km, times=[0, 250, 500, 750, 1000], format="polars")
PolarsRows10Columns3
strata
str
time
f64
n_risk
f64
0 1 0 138
1 1 250 62
2 1 500 20
3 1 750 7
4 1 1000 2
5 2 0 90
6 2 250 53
7 2 500 21
8 2 750 3
9 2 1000 0

To display a publication-ready table of numbers at risk, use risk_table(). It returns a Great Tables object that you can view interactively in notebooks or export to HTML, LaTeX, or other formats.

# Render a publication-ready Great Tables object showing numbers at risk
gw.risk_table(km, times=[0, 250, 500, 750, 1000])
Numbers at Risk
Count of subjects at risk at each time point
strata 0.0 250.0 500.0 750.0 1000.0
1 138 62 20 7 2
2 90 53 21 3 0

Customizing the plot

The plot_survival() function returns an ordinary Altair chart, so you refine it with Altair’s API. Chaining properties() sets the title and dimensions, and you can keep layering or restyling from there.

# Refine the chart with Altair's .properties() — title and dimensions
gw.plot_survival(km).properties(
    title="Survival by sex", width=640, height=360
)
NoteThe plot is a chart object, not an image

The plot_survival() function returns an Altair chart (an alt.LayerChart, or an alt.VConcatChart when risk_table=True). It renders interactively in notebooks and browsers, and you keep composing it with Altair’s API.

TipSaving figures

Call save() on the returned chart. Use chart.save("survival.html") for the interactive version, or chart.save("survival.png") (also .svg, .pdf) for a static image. Static export uses vl-convert, which ships with the altair extra.

Forest plots

A forest plot displays effect estimates with their confidence intervals, side by side. They are commonly used to show:

  • hazard ratios from Cox regression: comparing the effect of a covariate on survival
  • hazard ratio contrasts: comparing risk between two groups or treatments
  • RMST differences: difference in restricted mean survival time between groups
  • Any comparison of two distributions: as long as you have a point estimate and CI

The forest plot arranges estimates vertically, one per row, with the interval shown as a line and the point estimate as a dot. A reference line (usually at 1 for ratios, 0 for differences) marks the boundary between protective and harmful effects.

Cox model forest plots with plot_forest

The easiest way to produce a forest plot from a fitted Cox model is plot_forest(). Pass the fitted CoxPH object directly; the function extracts hazard ratios, confidence intervals, and p-values automatically and draws an interactive Altair chart with a log-scaled x-axis and a dashed reference line at HR = 1.

# 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 three covariates
cox = gw.CoxPH().fit(y, lung[["age", "sex", "ph.ecog"]])

# Draw a forest plot of hazard ratios with a log-scale x-axis
gw.plot_forest(cox)

Hazard ratios below 1 indicate a reduced hazard relative to the reference (e.g. sex = 2, female, has a lower hazard than male), and the position of the CI relative to the reference line shows whether the effect is statistically distinguishable from zero at the chosen confidence level. The tooltip on each row shows the formatted “HR (lower–upper)” text and the p-value.

Renaming terms for display

Column names from the data often need to be rewritten for a publication figure. Pass term_labels= as a dict mapping internal names to display labels. Only the terms you list are renamed. The rest keep their original names.

# Rename internal covariate labels to publication-friendly display strings
gw.plot_forest(
    cox,
    term_labels={
        "age": "Age (years)",
        "sex": "Female vs. Male",
        "ph.ecog": "ECOG PS (0–4)",
    },
    title="Cox model: NCCTG Lung Cancer",
)

Subgroup forest plots

The plot_forest() function also accepts any tidy DataFrame with term, estimate, ci_lower, and ci_upper columns. This makes it straightforward to build subgroup forest plots where you have already assembled the estimates (for example, by fitting one Cox model per subgroup and calling tidy() on each, then stacking the results).

import pandas as pd

# Build a pre-assembled subgroup results table (term, estimate, ci_lower, ci_upper)
subgroups = pd.DataFrame({
    "term": ["Age < 60", "Age \u2265 60", "Male", "Female", "ECOG 0", "ECOG 1+"],
    "estimate": [0.72, 0.91, 0.88, 0.68, 0.60, 0.95],
    "ci_lower": [0.51, 0.74, 0.69, 0.50, 0.42, 0.78],
    "ci_upper": [1.01, 1.12, 1.12, 0.92, 0.86, 1.16],
})

# Plot the subgroup forest plot directly from the DataFrame
gw.plot_forest(subgroups, title="Subgroup Analysis")

Explicit scale control

When passing a DataFrame, the scale defaults to linear. Pass scale="log" for a DataFrame of hazard ratios, or scale="linear" for RMST differences and other contrasts.

# Hazard ratios in a DataFrame: use scale="log"
hr_df = pd.DataFrame({
    "term": ["Drug A", "Drug B", "Reference"],
    "estimate": [0.82, 0.95, 1.0],
    "ci_lower": [0.68, 0.76, 0.85],
    "ci_upper": [0.99, 1.18, 1.18],
})

gw.plot_forest(hr_df, scale="log", title="Hazard Ratios: Treatment vs. Reference")
# RMST differences: linear scale (the default for DataFrames)
rmst_df = pd.DataFrame({
    "term": ["Drug A", "Drug B"],
    "estimate": [45, 25],
    "ci_lower": [-10, -20],
    "ci_upper": [95, 70],
})

gw.plot_forest(rmst_df, title="RMST Differences vs. Placebo")

Cumulative incidence plots

In studies with competing risks (e.g., death from disease vs. death from other causes), the ordinary Kaplan-Meier estimator gives biased results if you analyze each event type in isolation, because it treats competing events as censored. The correct tool is the cumulative incidence function (CIF): for each cause, it estimates the probability that a subject experiences that specific event by time t, accounting for the fact that experiencing any other cause first makes the target cause impossible. The CIFs for all causes sum to the overall event probability at every time.

The plot_cif() function accepts a fitted AalenJohansen estimator and renders the result as an interactive faceted chart, one panel per competing cause. Each panel is a right-continuous step function on the [0, 1] probability scale with a shaded confidence band.

Single group

Fit the Aalen-Johansen estimator first, then pass it directly to plot_cif(). We use the mgus2 dataset, which follows patients with monoclonal gammopathy (a precursor plasma-cell disorder). Each patient either progresses to overt plasma-cell malignancy (pcm), dies from an unrelated cause (death), or is censored (two competing events).

import numpy as np

# Load the monoclonal gammopathy dataset
mg = gw.load_dataset("mgus2", backend="polars")

# Build a single event-time per subject: progression time if they progressed,
# otherwise total follow-up time.
etime = np.where(mg["pstat"] == 1, mg["ptime"], mg["futime"])

# Encode the cause: 0 = censored, 1 = pcm progression, 2 = death without progression
cause = np.where(mg["pstat"] == 1, 1, 2 * mg["death"])

# Build the multi-state response, naming the two competing causes
y = gw.Surv.multistate(etime, event=cause, states=("pcm", "death"))

# Fit the Aalen-Johansen estimator (no grouping variable — overall population)
aj = gw.AalenJohansen().fit(y)
gw.plot_cif(aj)

Each panel shows the cumulative incidence for one cause over follow-up time in months. The shaded band is the point-wise confidence interval from the Aalen-Johansen variance estimator. Notice that the pcm curve climbs slowly while the death curve rises more steeply, reflecting that death is the more common outcome in this cohort. Because the two causes compete, neither curve can reach 1: a subject who has already died cannot later progress, so the two CIFs are bounded above by the overall event probability, not by 1 individually.

Grouped by covariate

To compare cumulative incidence between subgroups, pass by= to AalenJohansen.fit(). The estimator computes separate CIF curves for each level of the grouping variable. When you then call plot_cif(), it automatically colors the curves by group and adds a legend, keeping each cause in its own panel so you can see both differences side by side.

# Stratify by sex: fit separate CIF curves for each group
aj_sex = gw.AalenJohansen().fit(y, by=mg["sex"])

# plot_cif detects the grouping and draws one colored line per group per panel
gw.plot_cif(aj_sex, title="Cumulative incidence by sex")

Each panel now shows one colored curve per sex. Reading across both panels at the same time reveals whether males and females differ in the rate of progression versus the rate of death (two separate clinical questions answered in a single view).

Numbers at risk

Pass risk_table=True to stack an aligned numbers-at-risk table beneath the faceted panels. The table shows how many subjects remain event-free in each group at each query time, using the same color coding as the curves above it.

# Add a numbers-at-risk table below the CIF panels
gw.plot_cif(aj_sex, risk_table=True)

The default query times are six evenly spaced points from zero to the largest observed follow-up time. Supply times= to use specific time points instead:

# Choose specific months for the at-risk table
gw.plot_cif(aj_sex, risk_table=True, times=[0, 24, 60, 120, 180])

Customizing the plot

The plot_cif() function returns an Altair chart object, so it composes with the rest of the Altair grammar. The most common adjustments are available directly as keyword arguments:

# Suppress the confidence band (conf_int=False) and tighten the panel dimensions
gw.plot_cif(
    aj_sex,
    conf_int=False,   # hide the shaded CI band for a cleaner comparison
    width=350,        # width of each cause panel in pixels
    height=250,       # height of each cause panel in pixels
)
TipCIF versus 1 − Kaplan-Meier

A common mistake is to estimate the cumulative incidence of cause k as 1 - \hat{S}_k(t), where \hat{S}_k is the Kaplan-Meier curve treating all other causes as censored. This overestimates the CIF because censoring competing events inflates the at-risk set. The Aalen-Johansen CIF is the correct estimate. See Competing risks for the full explanation.

NoteConnecting the plot to a formal test

If the grouped curves look different, you can confirm the difference with Gray’s test, the competing-risks analogue of the log-rank test. That is covered in Competing risks.

Next steps

You can now produce and customize publication-quality survival figures.

  • Cox regression introduces models whose coefficients you will later visualize as forest plots.
  • Competing risks produces cumulative incidence curves, which are plotted the same way.
  • Revisit Comparing groups to pair a grouped figure with a formal test.