Survival Analysis
2026-11-17
Interval Censoring
dat = pd.DataFrame({
"ID": range(1, 7),
"lower_time": [7, 4, 9, 4.5, 4, 8],
"upper_time": [7, 6, 10, 4.5, 7, 8],
"interval_censored": [False, True, True, False, True, False]
}).sort_values("upper_time", ascending=False).reset_index(drop=True)
dat["ID"] = pd.Categorical(
dat["ID"].astype(str),
categories=dat["ID"].astype(str),
ordered=True
)
dat["event"] = np.where(
dat["interval_censored"], "Interval-censored", "Hatched"
)
dat["plot_time"] = np.where(
dat["interval_censored"],
(dat["lower_time"] + dat["upper_time"]) / 2,
dat["upper_time"]
)
interval_data = dat[dat["interval_censored"]]
(
ggplot(dat, aes(y="ID"))
+ geom_segment(aes(x=0, xend="lower_time", yend="ID"), color="gray")
+ geom_segment(
data=interval_data,
mapping=aes(x="lower_time", xend="upper_time", y="ID", yend="ID"),
inherit_aes=False,
color="gray",
linetype="dashed",
size=1
)
+ geom_point(aes(x="plot_time", shape="event"), size=3, color="#0072B2")
+ scale_shape_manual(values={"Hatched": "o", "Interval-censored": "s"})
+ labs(x="Weeks", y="Egg ID", shape="Event")
+ theme_bw()
)Right Censoring
dat = pd.DataFrame({
"ID": [str(value) for value in range(1, 11)],
"time": [7, 9, 10, 4, 2, 10, 8, 5, 6, 7],
"event_observed": [False, True, False, False, True,
False, False, True, False, True]
})
dat["event"] = np.where(
dat["event_observed"], "Hatched", "Right-censored"
)
(
ggplot(dat, aes(x="time", y="ID"))
+ geom_segment(aes(x=0, xend="time", yend="ID"), color="gray")
+ geom_point(aes(shape="event"), size=3, color="#0072B2")
+ geom_vline(xintercept=10, linetype="dashed")
+ scale_shape_manual(values={"Hatched": "o", "Right-censored": "s"})
+ labs(x="Weeks", y="Egg ID", shape="Event")
+ theme_bw()
)Survival Curve
time, survival_probability = kaplan_meier_estimator(
y["Status"],
y["Survival_in_days"]
)
median_time = time[np.flatnonzero(survival_probability <= 0.5)[0]]
km_data = pd.DataFrame({
"time": time,
"survival_probability": survival_probability
})
(
ggplot(km_data, aes(x="time", y="survival_probability"))
+ geom_step(color="#0072B2", size=1)
+ annotate(
"segment", x=0, xend=median_time, y=0.5, yend=0.5,
color="gray", linetype="dashed"
)
+ annotate(
"segment", x=median_time, xend=median_time, y=0, yend=0.5,
color="gray", linetype="dashed"
)
+ coord_cartesian(ylim=(0, 1.05))
+ labs(x="Time (days)", y="Estimated survival probability")
+ theme_bw()
)Veterans’ Survival Curve
- Fit a survival curve using
load_veterans_lung_cancer. Survival_in_days: time-to-death or censoringStatus: event indicator (Truemeans death was observed)- Plot the survival curve and add a line indicating the 50th percentile
time, survival_probability = kaplan_meier_estimator(
y["Status"], y["Survival_in_days"]
)
median_time = time[np.flatnonzero(survival_probability <= 0.5)[0]]
km_data = pd.DataFrame({
"time": time,
"survival_probability": survival_probability
})
(
ggplot(km_data, aes(x="time", y="survival_probability"))
+ geom_step(color="#0072B2", size=1)
+ annotate(
"segment", x=0, xend=median_time, y=0.5, yend=0.5,
color="gray", linetype="dashed"
)
+ annotate(
"segment", x=median_time, xend=median_time, y=0, yend=0.5,
color="gray", linetype="dashed"
)
+ coord_cartesian(ylim=(0, 1.05))
+ labs(x="Time (days)", y="Estimated survival probability")
+ theme_bw()
)Survival Curve by Treatment
- Fit survival curves stratified by treatment group.
Survival_in_days: time-to-death or censoringStatus: event indicatorTreatment: treatment group- Plot each survival curve and add a line indicating the 50th percentile
curves = []
for treatment in X["Treatment"].unique():
mask = (X["Treatment"] == treatment).to_numpy()
time, survival_probability = kaplan_meier_estimator(
y["Status"][mask],
y["Survival_in_days"][mask]
)
curves.append(pd.DataFrame({
"time": time,
"survival_probability": survival_probability,
"Treatment": str(treatment)
}))
treatment_km_data = pd.concat(curves, ignore_index=True)
(
ggplot(
treatment_km_data,
aes(x="time", y="survival_probability", color="Treatment")
)
+ geom_step(size=1)
+ geom_hline(yintercept=0.5, color="gray", linetype="dashed")
+ scale_color_manual(values=["#0072B2", "#D55E00"])
+ coord_cartesian(ylim=(0, 1.05))
+ labs(x="Time (days)", y="Estimated survival probability")
+ theme_bw()
)Survival Curves with Confidence Intervals
Add confidence intervals to the previous plot.
curves = []
for treatment in X["Treatment"].unique():
mask = (X["Treatment"] == treatment).to_numpy()
time, survival_probability, confidence_interval = kaplan_meier_estimator(
y["Status"][mask],
y["Survival_in_days"][mask],
conf_type="log-log"
)
curves.append(pd.DataFrame({
"time": time,
"survival_probability": survival_probability,
"lower": confidence_interval[0],
"upper": confidence_interval[1],
"Treatment": str(treatment)
}))
treatment_ci_data = pd.concat(curves, ignore_index=True)
(
ggplot(
treatment_ci_data,
aes(
x="time",
y="survival_probability",
color="Treatment",
group="Treatment"
)
)
+ geom_ribbon(aes(ymin="lower", ymax="upper", fill="Treatment"), alpha=0.2)
+ geom_step(size=1)
+ geom_hline(yintercept=0.5, color="gray", linetype="dashed")
+ scale_color_manual(values=["#0072B2", "#D55E00"])
+ scale_fill_manual(values=["#0072B2", "#D55E00"])
+ coord_cartesian(ylim=(0, 1.05))
+ labs(x="Time (days)", y="Estimated survival probability")
+ theme_bw()
)Plot the Hazard Ratios
Predicted Survival Curves
survival_functions = cox_model.predict_survival_function(
X_encoded.iloc[:3]
)
prediction_frames = []
for patient_number, survival_function in enumerate(
survival_functions,
start=1
):
prediction_frames.append(pd.DataFrame({
"time": survival_function.x,
"survival_probability": survival_function(survival_function.x),
"Patient": f"Patient {patient_number}"
}))
prediction_data = pd.concat(prediction_frames, ignore_index=True)
(
ggplot(
prediction_data,
aes(x="time", y="survival_probability", color="Patient")
)
+ geom_step(size=1)
+ coord_cartesian(ylim=(0, 1.05))
+ labs(x="Time (days)", y="Predicted survival probability")
+ theme_bw()
)