matplotlib legend not working correctly with handles

Question

Matplotlib can display a legend automatically or manually with giving it handles to the plots. But somehow the latter is not working correctly for me. Take this example:

legend_handles = {}
lgh, = plt.plot([0, 1], [0, 1], '-r')
lgh, = plt.plot([0, 1], [1, 1], '-r')
legend_handles["a"] = lgh
lgh, = plt.plot([0, 1], [1, 0], '-b')
legend_handles["b"] = lgh
plt.legend(legend_handles);

This will give a legend with two red lines, rather than a blue and a red line.

How do I get it to display a legend to only a selection of the plots?

Answer 1

There is no indication that legends would support dictionaries as input. Instead the signature is either of

legend()                  ## (1)
legend(labels)            ## (2)
legend(handles, labels)   ## (3)

Here you are using (2), so of the three lines, only the first 2 are labelled with the keys of the dictionary (because the dictionary only has two keys).

If you need to use a dictionary, you need to unpack it first to obtain two lists that can be used to achieve case (3).

import matplotlib.pyplot as plt

legend_handles = {}
lgh1, = plt.plot([0, 1], [0, 1], '-r')
lgh2, = plt.plot([0, 1], [1, 1], '-r')
legend_handles["a"] = lgh1
lgh3, = plt.plot([1, 0], [1, 0], '-b')
legend_handles["b"] = lgh3

labels, handles = zip(*legend_handles.items())
plt.legend(handles, labels)

plt.show()

However, not using a dictionary at all seems even simpler:

import matplotlib.pyplot as plt

lgh1, = plt.plot([0, 1], [0, 1], '-r')
lgh2, = plt.plot([0, 1], [1, 1], '-r')
lgh3, = plt.plot([1, 0], [1, 0], '-b')

plt.legend([lgh1, lgh3], list("ab"))

plt.show()

Not to forget, the canonical solution for creating legends by supplying the label to the artists directly,

import matplotlib.pyplot as plt

lgh1, = plt.plot([0, 1], [0, 1], '-r', label="a")
lgh2, = plt.plot([0, 1], [1, 1], '-r')
lgh3, = plt.plot([1, 0], [1, 0], '-b', label="b")

plt.legend()

plt.show()

Result in all cases: