CODEWITHSUNDEEP
pythonmatplotlib
Joysn
Joysn

Reputation: 1029

Consistent way of getting labels from plot, bar and other drawings with matplotlib

With line plots, I can get all labels like this and build a legend:

p1 = ax1.plot(x, 'P1', data=df)
p2 = ax1.plot(x, 'P2', data=df)
p3 = ax1.plot(x, 'P3', data=df)
p4 = ax1.plot(x, 'P4', data=df)

p = p1+p2+p3+p4
labs = [l.get_label() for l in p]
ax1.legend(p, labs, loc=0, frameon=False)

When I have bar plots, this does not work anymore. E.g.:

b1 = ax1.bar(x-2*w, 'B1', data=df, width=w, label="TP")
b2 = ax1.bar(x-w, 'B2', data=df, width=w, label="FN")
b3 = ax1.bar(x, 'B3', data=df, width=w, label="FP")
b4 = ax2.bar(x+w, 'B4', data=df, width=w, label="AP")
b5 = ax2.bar(x+2*w, 'B5', data=df, width=w, label="AR")

b1.get_label() returns a string similar to a __str__ method:

'0    87
Name: TP, dtype: object'

Why does .get_label() not behave identically?

Upvotes: 1

Views: 410

Answers (1)

JohanC
JohanC

Reputation: 80329

ax1.plot(...) returns a tuple of Line2D elements. Usually this is a tuple of just one element, but it can be longer when more lines are plotted in the same call (lines = ax1.plot(x1,y1,'r',x2,y2,'b') would return 2 Line2D elements).

When you do p1+p2+p3+p4, you append these tuples, creating a tuple of 4 elements. ax.bar, on the other hand, returns a single Bar container. These can't be concatenated via +. You need to create a tuple (b1,b2,b3,b4,b5) or a list [b1,b2,b3,b4,b5].

You'll often see a mysterious comma used in p1, = ax1.plot(...). That way, the first element of the tuple is assigned to p1.

Also note that you don't need to extract the labels. If you call ax1.legend(handles=p), matplotlib will extract and use these labels automatically.

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

x = np.arange(20)
df = pd.DataFrame({f'P{i}': np.random.randn(20).cumsum() for i in range(1, 5)})
fig, ax1 = plt.subplots()
p1 = ax1.plot(x, 'P1', data=df)
p2 = ax1.plot(x, 'P2', data=df)
p3 = ax1.plot(x, 'P3', data=df)
p4 = ax1.plot(x, 'P4', data=df)

p = p1 + p2 + p3 + p4
ax1.legend(handles=p, loc='best', frameon=False)
plt.show()

calling handles

The same can be written as follows, making it easier to combine handles from different functions:

p1, = ax1.plot(x, 'P1', data=df)
p2, = ax1.plot(x, 'P2', data=df)
p3, = ax1.plot(x, 'P3', data=df)
p4, = ax1.plot(x, 'P4', data=df)

p = [p1, p2, p3, p4]
ax1.legend(handles=p, frameon=False)
plt.show()

That makes it similar to how you would work with bars:

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

x = np.arange(5)
df = pd.DataFrame({f'B{i}': np.random.rand(5).cumsum() for i in range(1, 6)})
fig, ax1 = plt.subplots()
w = 0.19
b1 = ax1.bar(x - 2 * w, 'B1', data=df, width=w, label="TP")
b2 = ax1.bar(x - w, 'B2', data=df, width=w, label="FN")
b3 = ax1.bar(x, 'B3', data=df, width=w, label="FP")
b4 = ax1.bar(x + w, 'B4', data=df, width=w, label="AP")
b5 = ax1.bar(x + 2 * w, 'B5', data=df, width=w, label="AR")

ax1.legend(handles=[b1, b2, b3, b4, b5], frameon=False)
plt.show()

bar plot with legend from handles

Of course, in these cases, the legend can also be created automatically. However, explicit working with these handles can be interesting if you need finetuning the legend, or you want to combine two handles into one.

Upvotes: 1

Related Questions