如何用for循环在Python的lifelines中绘制多组Kaplan Meier曲线？

我希望绘制包含多组的Kaplan Meier生存曲线。现有代码仅能在一张图中展示两组曲线，但分组数量固定为两组。我需要使用for循环遍历包含所有分组的列表，处理数量不固定且大于2的分组，请问如何通过lifelines库实现？

原代码如下：

from lifelines import KaplanMeierFitter
from lifelines.datasets import load_waltons

waltons = load_waltons()
ix = waltons['group'] == 'control'

ax = plt.subplot(111)

kmf_control = KaplanMeierFitter()
ax = kmf_control.fit(waltons.loc[ix]['T'], waltons.loc[ix]['E'],label='control').plot_survival_function(ax=ax)

kmf_exp = KaplanMeierFitter()
ax = kmf_exp.fit(waltons.loc[~ix]['T'], waltons.loc[~ix]['E'], label='exp').plot_survival_function(ax=ax)


from lifelines.plotting import add_at_risk_counts
add_at_risk_counts(kmf_exp, kmf_control, ax=ax)
plt.tight_layout()

要实现任意数量分组的KM生存曲线绘制，只需通过动态遍历分组、收集模型对象来适配多组场景，修改后的完整代码如下：

from lifelines import KaplanMeierFitter
from lifelines.datasets import load_waltons
from lifelines.plotting import add_at_risk_counts
import matplotlib.pyplot as plt

# 加载数据集
waltons = load_waltons()
# 获取所有唯一分组名称
all_groups = waltons['group'].unique()

# 初始化绘图轴
ax = plt.subplot(111)
# 存储每个分组拟合后的KM模型实例
kmf_instances = []

# 遍历每个分组完成拟合与绘图
for group_name in all_groups:
    # 筛选当前分组的生存数据
    subset = waltons[waltons['group'] == group_name]
    # 初始化并拟合KM模型
    kmf = KaplanMeierFitter()
    kmf.fit(subset['T'], subset['E'], label=group_name)
    # 绘制当前分组的生存曲线
    kmf.plot_survival_function(ax=ax)
    # 将拟合好的模型存入列表
    kmf_instances.append(kmf)

# 添加上风险人数统计表格，解包所有模型对象传入函数
add_at_risk_counts(*kmf_instances, ax=ax)
plt.tight_layout()
plt.show()

关键说明： #

• 动态获取分组：使用waltons['group'].unique()自动提取所有分组，无需预先指定分组数量，兼容2组及以上的任意场景
• 收集模型实例：通过列表kmf_instances保存每个分组的拟合模型，最后用*解包传入add_at_risk_counts，自动生成对应所有分组的风险人数统计
• 统一循环逻辑：每个分组的拟合和绘图逻辑在循环内统一处理，避免重复编写多组冗余代码

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