使用plotnine绘制Lifelines KM曲线时CI ribbon显示异常求助

问题描述 #

使用lifelines自带的plot_survival_function()可以生成CI置信带与生存曲线完美对齐的KM曲线：

from lifelines.datasets import load_dd
data = load_dd()
kmf = KaplanMeierFitter()
kmf.fit(data["duration"], event_observed=data["observed"])
kmf.plot_survival_function()
plt.xlim(0, 20)
plt.show()

但改用plotnine重构代码后，CI置信带与阶梯状的生存曲线无法正确匹配对齐：

km_df = kmf.survival_function_.reset_index()
km_df.columns = ["Time", "Survival"]

ci_df = kmf.confidence_interval_.reset_index()
ci_df.columns = ["Time", "Lower", "Upper"]
km_merged = km_df.merge(ci_df, on="Time")

(
    p9.ggplot(km_merged, p9.aes(x="Time", y="Survival")) +
    p9.geom_step(size=1.2, color="blue") +
    p9.geom_ribbon(p9.aes(ymin="Lower", ymax="Upper"), alpha=0.25, fill="blue")
)

问题原因 #

Kaplan-Meier生存曲线是**阶梯式（step）**的：在两个事件时间点之间，生存率保持恒定，直到下一个事件发生才跳变。而geom_ribbon默认会用直线直接连接相邻时间点的CI上下界，没有遵循阶梯的逻辑，导致置信带的形状和生存曲线的阶梯结构不匹配。lifelines自带的绘图内部会自动为CI置信带生成阶梯式的点，而plotnine需要手动处理这部分数据。

解决方法 #

手动构造符合阶梯逻辑的CI置信带数据，让每个时间区间内的CI上下界保持恒定，和生存曲线的阶梯结构一致。具体代码如下：

from lifelines.datasets import load_dd
import plotnine as p9
import pandas as pd

# 1. 拟合KM模型
data = load_dd()
kmf = KaplanMeierFitter()
kmf.fit(data["duration"], event_observed=data["observed"])

# 2. 获取生存曲线和CI的基础数据
survival_df = kmf.survival_function_.reset_index()
survival_df.columns = ["Time", "Survival"]
ci_df = kmf.confidence_interval_.reset_index()
ci_df.columns = ["Time", "Lower", "Upper"]

# 3. 构造阶梯式的完整数据（包含每个时间点的前后状态）
step_times = [0]
step_survival = [1.0]
step_lower = [1.0]
step_upper = [1.0]

prev_surv = 1.0
prev_low = 1.0
prev_high = 1.0

for idx, row in survival_df.iterrows():
    curr_time = row["Time"]
    curr_surv = row["Survival"]
    curr_low = ci_df.loc[idx, "Lower"]
    curr_high = ci_df.loc[idx, "Upper"]
    
    # 添加前一个状态延续到当前时间点的记录
    step_times.append(curr_time)
    step_survival.append(prev_surv)
    step_lower.append(prev_low)
    step_upper.append(prev_high)
    
    # 添加当前时间点的新状态记录
    step_times.append(curr_time)
    step_survival.append(curr_surv)
    step_lower.append(curr_low)
    step_upper.append(curr_high)
    
    prev_surv = curr_surv
    prev_low = curr_low
    prev_high = curr_high

# 转换为DataFrame
km_step_data = pd.DataFrame({
    "Time": step_times,
    "Survival": step_survival,
    "Lower": step_lower,
    "Upper": step_upper
})

# 4. 用plotnine绘制对齐的KM曲线和CI置信带
(
    p9.ggplot(km_step_data, p9.aes(x="Time", y="Survival")) +
    p9.geom_step(size=1.2, color="blue") +
    p9.geom_ribbon(p9.aes(ymin="Lower", ymax="Upper"), alpha=0.25, fill="blue") +
    p9.xlim(0, 20)
)

原理说明 #

通过为每个事件时间点生成两行记录：一行保留上一个时间区间的生存率和CI值（让状态延续到当前时间点），另一行记录当前时间点更新后的状态。这样geom_ribbon绘制时，会在每个时间区间内填充恒定的CI范围，和geom_step绘制的生存曲线完全对齐，还原lifelines自带绘图的效果。

内容的提问来源于stack exchange，提问作者mjs