kmf for time series survival model

ts_survival_kmf.py

# loop over each department type, fit KMF and visualize
fig, ax = plt.subplots(1,1,figsize=(10,8))
kmf = KaplanMeierFitter()
depts = ['HOBBIES', 'FOODS']
for dept in depts:
  depts_list = [d for d in results.dept_id.unique()\
                 if d.startswith(dept)]
  res_depts = results[results.dept_id.isin(depts_list)]
  kmf.fit(res_depts.thresh_week[:STUDY_END_WEEKS],\
          res_depts.event[:STUDY_END_WEEKS])
  kmf.plot(ax=ax, label = dept, ci_show=False)
  ax.legend()
  print(dept)
  #print(kmf.median_survival_time_)
  print(kmf.percentile(0.95))
  print('...')
  
ax.set_ylabel('Survival Probability')
ax.set_xlabel('Weeks')
ax.set_title('Kaplan Meier Estimator')
ax.set_xticks(np.arange(0,STUDY_END_WEEKS,2))
plt.show()