Home > Enterprise >  Manually add legend entry to Seaborn Legend
Manually add legend entry to Seaborn Legend

Time:06-09

I create overlaying distribution plots by group memberships using

gfg = sns.displot(hist_dat, x="Outcome", hue="Group", kind="kde", fill = True,
                  palette="tab10", height=10.5, aspect=18.5/10.5, legend=True)
gfg._legend.set_title("Group")
plt.setp(gfg._legend.get_title(), fontsize=20) 

Now, I add a vertical line to the plot and would like to add an entry for this line to the legend. The solution I could make work is

plt.axvline(x = 0, color = 'b', label = 'ATE')
plt.legend(bbox_to_anchor= (1.1,0.4), frameon=False, title="Effect", 
              title_fontsize=20)

To get aligned legends, I have to adjust them manually. Is there a more straightforward way to add the ATE legend entry - e.g. in the existing legend? Many thanks!

Result:

Result

CodePudding user response:

Edit: when you use fill, you don't get lines. My bad. Changed the code slightly.

To do this, you need to specify the handles and labels to the legend argument. The way I found to get the PolyCollection artists is with ax.get_children(). Then, you call plt.legend(handles, labels). Here's a toy example:

sns.displot([0, 1, 2, 2, 3, 3, 4, 4, 5], legend=True, kind='kde', label='test', fill=True)
children = plt.gca().get_children()
l = plt.axvline(3.5, c='r')
plt.legend([children[0], l], ['curve', 'line'] )

Example plot

  • Related