0

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.

Andrew M
  • 119
  • 6
  • 1
    Without a more complete code and test data, it is hard to guess what you are doing. Anyway, `plt.legend(handles=[], ...)` isn't supposed to draw anything. Handles can either be existing graphical elements of a plot, or dummy elements that got created without adding them to the plot. You can either assign a label or provide these labels in the call to `plt.legend()`. By default, all elements in the plot that got a label (not staring with `'_'`) will be used as handles. As `plt.legend()` removes a possible previously existing legend, the clearest way would be to call it after the `for` loop. – JohanC Sep 21 '22 at 15:24
  • Thank you @JohanC ! I was overthinking it somewhat. As you pointed out, getting rid of the plt.legend(handles=...) caused the legend to default to what I had named it in the plot call :) As you guessed, my confusion was with the utility of plt.legend() – Andrew M Sep 21 '22 at 20:07

0 Answers0