Kaplan-Meier curve explained | by Zolzaya Luvsandorj | Towards Data Science

Kaplan-Meier curve, a popular survival analysis tool, is useful in understanding survival probability over time in the presence of incomplete data. In this post, we will learn how to build a Kaplan-Meier curve from scratch to gain a better understanding, then look at two ways to build it using survival analysis libraries in Python.
Let’s imagine we have data collected from September 2020 to June 2021 for a clinical study to understand the survival-time among cancer patients. Here’s a timeline for the people who participated in this hypothetical study:
We kept the dataset for our imaginary study small so that it’s easier to closely monitor the process. Black solid circles signifies the last data entry for a record. A cross inside the circle means the record had the event (i.e. death) whereas a question mark means the record haven’t had the event. To be more specific, there are three possible scenarios.
Had event: Had event during the study (e.g. record a and c)No event: Didn’t have event by the end of study (e.g. record b and d).No event: Dropped out of study before the endpoint and didn’t have event while enrolled in the study (e.g. record f and h).
Had event: Had event during the study (e.g. record a and c)
No event: Didn’t have event by the end of study (e.g. record b and d).
No event: Dropped out of study before the endpoint and didn’t have event while enrolled in the study (e.g. record f and h).
The second and third types of records are called censored, more specifically right-censored, records as we have incomplete information regarding the event.
In the chart above, we see a calendar time in the horizontal axis. Now, we need to convert a calendar time to a survival time, a time duration measured relative to the study enrolment date. Survival time is commonly referred to as duration in the Python libraries we are going to later familiarise; hence, we will be using these terms interchangeably here onwards.
Using the following script, let’s import the libraries that we need and create a small DataFrame containing duration and event for the 9 records:
# Data manipulationimport numpy as npimport pandas as pdfrom string import ascii_lowercase# Visualisationimport matplotlib.pyplot as pltimport seaborn as snssns.set(style='dark', context='talk')# Kaplan-Meier curvefrom lifelines import KaplanMeierFitter from sksurv.nonparametric import kaplan_meier_estimator# Create a toy dataframedf = pd.DataFrame(data={'duration': [3,9,5,8,7,2,1,3,4],                        'event': [1,0,1,0,1,0,1,0,0]},                   index=list(ascii_lowercase[:9]))df
Now let’s familiarise ourselves with the conceptual knowledge. The Kaplan-Meier curve is the visual representation of the Kaplan-Meier estimator which makes the following assumptions:
Censorship: Censored observations have the same survival prospects as uncensored observations.Study entry time: Survival probabilities are the same for observations recruited early and later in the study.Event time: Event happened at the defined time
Censorship: Censored observations have the same survival prospects as uncensored observations.
Study entry time: Survival probabilities are the same for observations recruited early and later in the study.
Event time: Event happened at the defined time
The Kaplan-Meier estimator is defined by the following formula:
With this formula, we can find the survival probability at a given time. Whether i starts from 1 or 0 doesn’t make a difference as survival probability at time 0 is 1. If we look at this formula closely, we will notice that what’s inside bracket is capturing the proportion survived (i.e. 1- proportion died).
Let’s find the survival probability for the first three durations with this formula:
At one month into the study, we had 9 participants at risk and one participant had died. Hence, d1 = 1 and n1 = 9.
At two months into the study, we had 8 participants at risk and no one had died. Hence, d2 = 0 and n = 8. As you can see, when duration has no event, the survival probability remains the same as the previous duration. Therefore, if a duration has no event, we can skip the calculation and use survival probability from the previous survival time.
At three months into the study, we had 7 records at risk and 1 death, the survival probability drops to 0.76. The equation is getting longer and more tedious now. We can rewrite the formula to the following simpler form:
Using this formula, calculations are more succinct:
These three examples have illustrated how the probabilities are estimated. Instead of repetitively calculating the rest, let’s translate the logic into Python scripts and calculate the survival probability for all durations:
# Prepare unique durations in ascending orderdurations = df.sort_values('duration')['duration'].unique()# Initialise the tablecolumns = ['duration', 'n_at_risk', 'n_events',            'survival_probability']km = pd.DataFrame(columns=columns, dtype=np.number)km = km.append(pd.DataFrame([[0, df.shape[0], 0, 1]],                             columns=columns))# Calculate survival probability for each durationfor i, t in enumerate(durations):    n = np.sum(df['duration']>=t)    d = np.sum((df['duration']==t) & (df['event']==1))    s = (1 - d / n) * km.loc[i, 'survival_probability']    km = km.append(pd.DataFrame([[t, n, d, s]],                                 index=[i+1],                                columns=columns))km
The Kaplan-Meier curve needs the survival probability and duration columns. Let’s plot the curve now since we have the inputs ready.
plt.figure(figsize=(8,4))sns.lineplot(data=km, x='duration', y='survival_probability',              drawstyle='steps-post')plt.ylim(0,1.1) plt.title("Kaplan-Meier curve");
Voila, we just plotted the curve ourselves. Constructing the curve from scratch hopefully helped you to understand the underlying logic. Having learned how to construct the curve ourselves, let’s learn more practical ways to plot the Kaplan-Meier curve. We will familiarise with two libraries that provide the solution out-of-the-box.
Lifeline is an easy-to-use survival analysis library in Python. Kaplan-Meier curve can be plotted using the KaplanMeierFitter object in a handful of lines:
kmf = KaplanMeierFitter() kmf.fit(df['duration'], df['event'])plt.figure(figsize=(8,4))kmf.plot()plt.title("Kaplan-Meier curve");
From the fitted object, we can also see the survival probabilities by accessing the survival_function_ attribute:
kmf.survival_function_
We can also extract the event table, which is similar to a table we built in Section 1.
kmf.event_table
With lifelines, plotting the curve and inspecting details regarding the estimator becomes easy.
Another useful survival analysis library is scikit-survival. As the name suggests, it works harmoniously with scikit-learn. We can get the duration and survival probability by utilising kaplan_meier_estimator(). Once we have the required inputs for the curve, we can plot it using any visualisation library. Since we used seaborn in Section 1, we will look at an alternative way to plot using matplotlib:
duration, survival_probability = kaplan_meier_estimator(df['event']==1, df['duration'])plt.figure(figsize=(8,4))plt.step(x=duration, y=survival_probability, where="post")plt.ylim(0,1.1)plt.xlabel("Survival time (months)")plt.ylabel("Survival probability")plt.title("Kaplan-Meier curve");
That was it, you have learned underlying logic for preparing the inputs for the Kaplan-Meier curve and familiarised with a few different ways to plot the curve!
Would you like to access more content like this? Medium members get unlimited access to any articles on Medium. If you become a member using my referral link, a portion of your membership fee will directly go to support me.
Thank you for reading this article. If you are interested, here are links to some of my other posts:◼️ ROC curve explained◼️ Useful IPython magic commands◼️ 5 tips to learn Python from zero◼️ Organise your Jupyter Notebook with these tips◼️️ Writing 5 common SQL queries in pandas
Bye for now 🏃💨