Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Add axhline to legend

I'm creating a lineplot from a dataframe with seaborn and I want to add a horizontal line to the plot. That works fine, but I am having trouble adding the horizontal line to the legend.

Here is a minimal, verifiable example:

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

x = np.array([2, 2, 4, 4])
y = np.array([5, 10, 10, 15])
isBool = np.array([True, False, True, False])

data = pd.DataFrame(np.column_stack((x, y, isBool)), columns=["x", "y", "someBoolean"])
print(data)

ax = sns.lineplot(x="x", y="y", hue="someBoolean", data=data)

plt.axhline(y=7, c='red', linestyle='dashed', label="horizontal")

plt.legend(("some name", "some other name", "horizontal"))

plt.show()

This results in the following plot:

incorrect plot

The legends for "some name" and "some other name" show up correctly, but the "horizontal" legend is just blank. I tried simply using plt.legend() but then the legend consists of seemingly random values from the dataset.

Any ideas?

like image 465
maaland Avatar asked Nov 28 '18 09:11

maaland


1 Answers

Simply using plt.legend() tells you what data is being plotting:

enter image description here

You are using someBoolean as the hue. So you are essentially creating two lines by applying a Boolean mask to your data. One line is for values that are False (shown as 0 on the legend above), the other for values that are True (shown as 1 on the legend above).

In order to get the legend you want you need to set the handles and the labels. You can get a list of them using ax.get_legend_handles_labels(). Then make sure to omit the first handle which, as shown above, has no artist:

ax = sns.lineplot(x="x", y="y", hue="someBoolean", data=data)

plt.axhline(y=7, c='red', linestyle='dashed', label="horizontal")

labels = ["some name", "some other name", "horizontal"]
handles, _ = ax.get_legend_handles_labels()

# Slice list to remove first handle
plt.legend(handles = handles[1:], labels = labels)

This gives:

enter image description here

like image 57
DavidG Avatar answered Oct 18 '22 09:10

DavidG