Matplotlib legend help

Question

I am writing a script that plot's several points. I am also trying to create a legend from these points. To sum up my script, I am plotting several 'types' of points (call them 'a', 'b', 'c'). These points have different colors and shapes: 'a'-'go' 'b'-'rh' 'c'-'k^'.

This is a shortened version of the relevant parts of my script:

lbl = #the type of point x,y is (a,b,c)
for x,y in coords:
   if lbl in LABELS:
      plot(x, y, color)
   else:
      LABELS.add(lbl)
      plot(x, y, color, label=lbl)
 legend()

What I am doing here is just plotting a bunch of points and assigning a label to them. However, I found out if I added a label to each point, then the legend will contain an entry for each point. I only want one entry per type of point (a, b, c). So, I changed my script to look like the above. Is there a better way to do this? If I have a million different types of points, then the data structure LABELS (a set) will take up a lot of space.

Answer 1

Group x and y according to the type of point. Plot all the points of the same type with one call to plot:

import pylab
import numpy as np

lbl=np.array(('a','b','c','c','b','a','b','c','a','c'))
x=np.random.random(10)
y=np.random.random(10)
for t,color in zip(('a','b','c'),('go','rh','k^')):
    pylab.plot(x[lbl==t],y[lbl==t],color,label=t)
pylab.legend()
pylab.show()

Answer 2

You can add the legend manually like this

labels = ['Label 1', 'Label 2', 'Label 3']
legend_location = 'upper right'
colors = ['r', 'g', 'b']
circles = [Circle((0, 0), 1, fc=colors[0]), Circle((0, 0), 1, fc=colors[1]), Circle((0, 0), 1, fc=colors[2])]

plt.legend(circles, labels, loc=legend_location)