How to use custom legend label in lifelines Kaplan-Meier plot?

Using the lifelines library, you can easily plot Kaplan-Meier plots, e.g. as seen in our previous post Minimal Python Kaplan-Meier Plot example:

from lifelines.datasets import load_leukemia
from lifelines import KaplanMeierFitter
df = load_leukemia()

kmf = KaplanMeierFitter()
kmf.fit(df['t'], df['Rx']) # t = Timepoints, Rx: 0=censored, 1=event
kmf.plot()

What if you want a custom label instead of KM_estimates to appear in the legend?

Use kmf.fit(..., label='Your label'). Since we use the leukemias dataset for this example, we use the label 'Leukemia'

Full example:

from lifelines.datasets import load_leukemia
from lifelines import KaplanMeierFitter
# Load datasets
df_leukemia = load_leukemia()

# Fit & plot leukemia dataset
kmf_leukemia = KaplanMeierFitter()
kmf_leukemia.fit(df_leukemia['t'], df_leukemia['Rx'], label="Leukemia")
ax = kmf_leukemia.plot()
# Set Y axis to fixed scale
ax.set_ylim([0.0, 1.0])