Making a legend / getting handles of Kaplan Meier plot objects in python lifelines

97 Views Asked by At

I'm working with the lifelines package to make Kaplan-Meier curves. I'd like to add the censored data, but also have a legend that only mentions the two lines.

I'm calling the function iteratively twice to plot two separate lines, as so:

def plot_km(col,months,dpi):
  ax = plt.subplot(111)
  clrIdx = 0
  for r in df[col].unique():
    ix = df[col] == r
    plt.rcParams['figure.dpi'] = dpi
    plt.rcParams['savefig.dpi'] = dpi
    kmf.fit(T[ix], C[ix],label=r)
    kmf.plot(ax=ax, color=colorsKM[clrIdx],show_censors=True,censor_styles={'ms': 6, 'marker': 's','label':'_out_'})
    if clrIdx == 1:
        plt.legend(handles=[],labels=['test1', 'test2'])
    clrIdx += 1

Where the output is a KM curve as well as the censored datapoints. However, I can't seem to figure out a way to interact with the handles/labels that gets the desired output. The above code seems to ignore the censored objects by using 'label':'_out_' , but it ignores my custom labels in the plt.legend() call. If I enter anything for handles, e.g.: plt.legend(handles=[line0] , it throws "NameError: name 'line0' is not defined"

enter image description here

I tried playing around with h, l = ax.get_legend_handles_labels() but this always returns empty. I believe my issue is with not understanding how each of these "artists"(?) are getting stored, and how to get them again.

0

There are 0 best solutions below