20 Survival Analysis with Causal Inference

using DataFrames
using Distributions
using GLM
using Statistics
using Random
using LinearAlgebra
using Printf
using CairoMakie
using CSV

When the outcome is time-to-event, such as death, readmission, job exit, or machine failure, ordinary regression is not enough. Some observations are censored. We know the event had not happened by the end of follow-up, but we do not know the event time.

This Julia chapter implements Kaplan-Meier curves, IPW-adjusted survival curves, and RMST differences. For Cox models and doubly robust survival estimators, R’s survival and riskRegression packages are still the more practical tools; see the R companion chapter.

20.1 The censoring problem

We observe \(\min(T_i, C_i)\) and \(\delta_i = \mathbb{1}\{T_i \le C_i\}\), where \(T_i\) is the event time and \(C_i\) is the censoring time. The main quantity of interest is the survival function

\[ S(t) = P(T > t) \]

the probability of surviving past time \(t\). Causal contrasts are usually reported as:

Hazard ratio (HR) — easy to estimate but hard to interpret causally (Hernán 2010).
Restricted mean survival difference (RMST) — \(\int_0^\tau [S_1(t) - S_0(t)] dt\) — interpretable as “years gained” up to time \(\tau\).
Survival probability difference at t — direct probability statement, requires picking a follow-up time.

20.2 Simulated survival data

df = CSV.read("data/survival_sim.csv", DataFrame)
n  = nrow(df)
@printf("n = %d, events = %d (%.1f%%), censored = %d\n",
        n, Int(sum(df.event)), 100 * mean(df.event), Int(sum(1 .- df.event)))

n = 1500, events = 1010 (67.3%), censored = 490

20.3 Kaplan-Meier estimator

The Kaplan-Meier estimator computes \(\hat S(t)\) nonparametrically. It is short enough to implement directly:

"""
    kaplan_meier(time, event)
Compute Kaplan-Meier survival estimates at the unique event times.
Returns (times, survival).
"""
function kaplan_meier(time, event)
    df_sort = sort(DataFrame(t = time, d = event), :t)
    event_times = unique(df_sort.t[df_sort.d .== 1])
    survival = Float64[]
    s = 1.0
    for t in event_times
        n_at_risk = sum(df_sort.t .>= t)
        n_events  = sum((df_sort.t .== t) .& (df_sort.d .== 1))
        s *= 1 - n_events / n_at_risk
        push!(survival, s)
    end
    return event_times, survival
end

t_trt, s_trt = kaplan_meier(df.time[df.trt .== 1], df.event[df.trt .== 1])
t_ctl, s_ctl = kaplan_meier(df.time[df.trt .== 0], df.event[df.trt .== 0])

fig = Figure(size = (700, 400))
ax = Axis(fig[1, 1], xlabel = "Time (years)", ylabel = "S(t)",
          title = "Unadjusted Kaplan-Meier survival curves")
stairs!(ax, t_ctl, s_ctl, color = :steelblue, linewidth = 2,
        label = "Control")
stairs!(ax, t_trt, s_trt, color = :firebrick, linewidth = 2,
        label = "Treatment")
axislegend(ax, position = :rb, framevisible = false)
fig

The unadjusted KM curves are descriptive. They include both the treatment effect and baseline imbalance.

20.4 IPW-adjusted survival

To adjust for confounding, weight observations by the inverse propensity score before computing KM:

ps_fit = glm(@formula(trt ~ age + sex), df, Binomial(), LogitLink())
df.ps  = predict(ps_fit)
df.ipw = ifelse.(df.trt .== 1, 1 ./ df.ps, 1 ./ (1 .- df.ps))
df.ipw .= min.(df.ipw, quantile(df.ipw, 0.99))

"""
    weighted_kaplan_meier(time, event, weights)
KM estimate with sample weights — used for IPW-adjusted survival.
"""
function weighted_kaplan_meier(time, event, weights)
    df_sort = sort(DataFrame(t = time, d = event, w = weights), :t)
    event_times = unique(df_sort.t[df_sort.d .== 1])
    survival = Float64[]
    s = 1.0
    for t in event_times
        w_at_risk = sum(df_sort.w[df_sort.t .>= t])
        w_events  = sum(df_sort.w[(df_sort.t .== t) .& (df_sort.d .== 1)])
        s *= 1 - w_events / w_at_risk
        push!(survival, s)
    end
    return event_times, survival
end

t_trt_w, s_trt_w = weighted_kaplan_meier(df.time[df.trt .== 1],
                                          df.event[df.trt .== 1],
                                          df.ipw[df.trt .== 1])
t_ctl_w, s_ctl_w = weighted_kaplan_meier(df.time[df.trt .== 0],
                                          df.event[df.trt .== 0],
                                          df.ipw[df.trt .== 0])

fig2 = Figure(size = (700, 400))
ax2 = Axis(fig2[1, 1], xlabel = "Time (years)", ylabel = "S(t)",
           title = "IPW-adjusted Kaplan-Meier survival curves")
stairs!(ax2, t_ctl_w, s_ctl_w, color = :steelblue, linewidth = 2,
        label = "Control (IPW-adj)")
stairs!(ax2, t_trt_w, s_trt_w, color = :firebrick, linewidth = 2,
        label = "Treatment (IPW-adj)")
axislegend(ax2, position = :rb, framevisible = false)
fig2

The IPW-adjusted curves remove the age and sex imbalance that is in the propensity score model.

20.5 Restricted Mean Survival Time

The RMST up to time \(\tau\) is

\[ \text{RMST}(\tau) = E[\min(T, \tau)] = \int_0^\tau S(t) \, dt. \]

The contrast \(\Delta_{\text{RMST}}(\tau)=\text{RMST}_1(\tau)-\text{RMST}_0(\tau)\) is interpretable as years gained up to \(\tau\). It does not require proportional hazards.

Because the KM curve is a step function, its integral is computed exactly by a left rectangle sum (width times the survival value at the left endpoint of each segment):

function rmst(times, survival, tau)
    # Add time 0 and the truncation point, with appropriate S values
    t_vec = [0.0; times; tau]
    s_vec = [1.0; survival; survival[end]]
    # Restrict to times ≤ tau
    keep = t_vec .<= tau
    t_use = t_vec[keep]
    s_use = s_vec[keep]
    # Add the tau endpoint, unless it's already the last element (t_vec was
    # constructed with tau appended, so `keep` already retains it -- pushing
    # again would duplicate it and create a zero-width trailing interval).
    if t_use[end] != tau
        push!(t_use, tau)
        push!(s_use, s_use[end])
    end
    # Exact integral of the KM step function: width × left-height
    sum(diff(t_use) .* s_use[1:end-1])
end

τ = 5.0
rmst_trt_unadj = rmst(t_trt, s_trt, τ)
rmst_ctl_unadj = rmst(t_ctl, s_ctl, τ)
rmst_trt_adj   = rmst(t_trt_w, s_trt_w, τ)
rmst_ctl_adj   = rmst(t_ctl_w, s_ctl_w, τ)

@printf("Unadjusted RMST(τ=5):\n")
@printf("  Treatment: %.3f years    Control: %.3f years    Diff: %.3f\n",
        rmst_trt_unadj, rmst_ctl_unadj, rmst_trt_unadj - rmst_ctl_unadj)
@printf("IPW-adjusted RMST(τ=5):\n")
@printf("  Treatment: %.3f years    Control: %.3f years    Diff: %.3f\n",
        rmst_trt_adj, rmst_ctl_adj, rmst_trt_adj - rmst_ctl_adj)

Unadjusted RMST(τ=5):
  Treatment: 4.351 years    Control: 4.147 years    Diff: 0.203
IPW-adjusted RMST(τ=5):
  Treatment: 4.397 years    Control: 4.122 years    Diff: 0.275

The IPW-adjusted RMST difference is the adjusted years-gained estimate up to 5 years.

20.6 Cox proportional hazards (regression adjustment)

Julia’s survival-regression tools are less mature than R’s survival::coxph. In Julia, the choices are:

Use Survival.jl for basic Cox PH (limited adjustment capabilities).
Use a parametric Weibull/exponential model via Distributions and Optim directly.
Call out to R for serious survival modelling.

For applied work where the Cox model with covariate adjustment is the target, R remains the practical choice. See the companion R chapter for the full Cox + IPCW + doubly-robust workflow.

20.7 Doubly-robust survival estimation

R’s riskRegression::ate combines a Cox outcome model with a propensity score and censoring model. Julia does not currently have an equivalent.

For now, the practical Julia workflow is:

Description: unadjusted KM curves (this chapter).
Adjusted survival: IPW-weighted KM (this chapter).
Causal contrasts: RMST differences from the IPW-adjusted KM (this chapter).
Cox HR with adjustment + doubly-robust estimators: call out to R via RCall.jl or use the R companion chapter.

20.8 When to use these methods

Use these methods when:

The outcome is time-to-event with censoring.
The treatment effect on survival probabilities or years-of-life is the causal question.

For non-censored outcomes, the standard Estimation and Heterogeneous Effects methods are sufficient.

20.9 Summary

Survival outcomes require methods that handle censoring.
Kaplan-Meier estimates the survival curve nonparametrically.
IPW-adjusted KM handles baseline confounding through weights.
RMST differences are often the clearest causal summary.
For Cox regression and doubly robust survival estimators, use R for now.