How do I assign multiple labels at once in matplotlib?

Viewed 126753

I have the following dataset:

x = [0, 1, 2, 3, 4]
y = [ [0, 1, 2, 3, 4],
      [5, 6, 7, 8, 9],
      [9, 8, 7, 6, 5] ]

Now I plot it with:

import matplotlib.pyplot as plt
plt.plot(x, y)

However, I want to label the 3 y-datasets with this command, which raises an error when .legend() is called:

lineObjects = plt.plot(x, y, label=['foo', 'bar', 'baz'])
plt.legend()

File "./plot_nmos.py", line 33, in <module>
  plt.legend()
...
AttributeError: 'list' object has no attribute 'startswith'

When I inspect the lineObjects:

>>> lineObjects[0].get_label()
['foo', 'bar', 'baz']
>>> lineObjects[1].get_label()
['foo', 'bar', 'baz']
>>> lineObjects[2].get_label()
['foo', 'bar', 'baz']

Question

Is there an elegant way to assign multiple labels by just using the .plot() method?

9 Answers

The best current solution is:

lineObjects = plt.plot(x, y)  # y describes 3 lines
plt.legend(['foo', 'bar', 'baz'])

In case of numpy matrix plot assign multiple legends at once for each column

I would like to answer this question based on plotting a matrix that has two columns.

Say you have a 2 column matrix Ret

then one may use this code to assign multiple labels at once

import pandas as pd, numpy as np, matplotlib.pyplot as plt
pd.DataFrame(Ret).plot()

plt.xlabel('time')
plt.ylabel('Return')
plt.legend(['Bond Ret','Equity Ret'], loc=0)
plt.show()

I hope this helps

This problem comes up for me often when I have a single set of x values and multiple y values in the columns of an array. I really don't want to plot the data in a loop, and multiple calls to ax.legend/plt.legend are not really an option, since I want to plot other stuff, usually in an equally annoying format.

Unfortunately, plt.setp is not helpful here. In newer versions of matplotlib, it just converts your entire list/tuple into a string, and assigns the whole thing as a label to all the lines.

I've therefore made a utility function to wrap calls to ax.plot/plt.plot in:

def set_labels(artists, labels):
    for artist, label in zip(artists, labels):
        artist.set_label(label)

You can call it something like

x = np.arange(5)
y = np.random.ranint(10, size=(5, 3))

fig, ax = plt.subplots()
set_labels(ax.plot(x, y), 'ABC')

This way you get to specify all your normal artist parameters to plot, without having to see the loop in your code. An alternative is to put the whole call to plot into a utility that just unpacks the labels, but that would require a lot of duplication to figure out how to parse multiple datasets, possibly with different numbers of columns, and spread out across multiple arguments, keyword or otherwise.

I used the following to show labels for a dataframe without using the dataframe plot:

lines_ = plot(df)
legend(lines_, df.columns) # df.columns is a list of labels

If you're using a DataFrame, you can also iterate over the columns of the data you want to plot:

# Plot figure
fig, ax = plt.subplots(figsize=(5,5))
# Data
data = data
# Plot
for i in data.columns:
    _ = ax.plot(data[i], label=i)
    _ = ax.legend() 
plt.show()
Related