Metrics_Hannah
Metrics_Hannah

Reputation: 45

Manually add legend entry to Seaborn Legend

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

Upvotes: 2

Views: 4463

Answers (1)

K.Cl
K.Cl

Reputation: 1783

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

Upvotes: 4

Related Questions