Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
%matplotlib inline
import numpy as np
from lifelines import KaplanMeierFitter
from matplotlib import pyplot as plt
from pylab import rcParams
rcParams['figure.figsize'] = 20, 10
plt.style.use('ggplot')
def run_survival(data, group_by=None, groups=[]):
time_column = 'time'
observation_column = 'death'
ax = plt.subplot(111)
kmf = KaplanMeierFitter()
if group_by is None:
kmf.fit(data[time_column], data[observation_column])
print(kmf.survival_function_.head())
print('Median')
print(kmf.median_)
kmf.plot(ax=ax)
else:
kmf.fit(data[time_column], data[observation_column], label='baseline')
print(kmf.survival_function_.head())
print('Median')
print(kmf.median_)
kmf.plot(ax=ax)
grouped_data = data.groupby([group_by])
plt.title(group_by)
if len(groups) == 0:
groups = np.sort(data[group_by].unique())
for group in groups:
d = grouped_data.get_group(group)
print(group, len(d))
kmf.fit(d[time_column], d[observation_column], label=group)
print(kmf.survival_function_.head())
print('Median')
print(kmf.median_)
kmf.plot(ax=ax)
return ax
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.