How do I assign multiple labels at once in matplotlib?

Question

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:

Answer 1

You can give the labels while plotting the curves

import pylab as plt

x = [0, 1, 2, 3, 4]
y = [ [0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [9, 8, 7, 6, 5] ]
labels=['foo', 'bar', 'baz']
colors=['r','g','b']

# loop over data, labels and colors
for i in range(len(y)):
    plt.plot(x,y[i],'o-',color=colors[i],label=labels[i])

plt.legend()
plt.show()

Answer 2

You can iterate over your line objects list, so labels are individually assigned. An example with the built-in python iter function:

lineObjects = plt.plot(x, y)
plt.legend(iter(lineObjects), ('foo', 'bar', 'baz'))`

Edit: after updating to matplotlib 1.1.1, it looks like the plt.plot(x, y), with y as a list of lists (as provided by the author of the question), doesn't work anymore. The one step plotting without iteration over the y arrays is still possible thought after passing y as numpy.array (assuming (numpy)[http://numpy.scipy.org/] as been previously imported).

In this case, use plt.plot(x, y) (if the data in the 2D y array are arranged as columns [axis 1]) or plt.plot(x, y.transpose()) (if the data in the 2D y array are arranged as rows [axis 0])

Edit 2: as pointed by @pelson (see commentary below), the iter function is unnecessary and a simple plt.legend(lineObjects, ('foo', 'bar', 'baz')) works perfectly

Answer 3

I came over the same problem and now I found a solution that is most easy! Hopefully that's not too late for you. No iterator, just assign your result to a structure...

from numpy import *
from matplotlib.pyplot import *
from numpy.random import *

a = rand(4,4)
a
>>> array([[ 0.33562406,  0.96967617,  0.69730654,  0.46542408],
   [ 0.85707323,  0.37398595,  0.82455736,  0.72127002],
   [ 0.19530943,  0.4376796 ,  0.62653007,  0.77490795],
   [ 0.97362944,  0.42720348,  0.45379479,  0.75714877]])

[b,c,d,e] = plot(a)
legend([b,c,d,e], ["b","c","d","e"], loc=1)
show()

Looks like this:

Answer 4

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

Answer 5

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.

Answer 6

The best current solution is:

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