How do you create a legend for a contour plot in matplotlib?
You could also do it directly with the lines of the contour, without using proxy artists.
import matplotlibimport numpy as npimport matplotlib.cm as cmimport matplotlib.mlab as mlabimport matplotlib.pyplot as pltmatplotlib.rcParams['xtick.direction'] = 'out'matplotlib.rcParams['ytick.direction'] = 'out'delta = 0.025x = np.arange(-3.0, 3.0, delta)y = np.arange(-2.0, 2.0, delta)X, Y = np.meshgrid(x, y)Z1 = mlab.bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)Z2 = mlab.bivariate_normal(X, Y, 1.5, 0.5, 1, 1)# difference of GaussiansZ = 10.0 * (Z2 - Z1)# Create a simple contour plot with labels using default colors. The# inline argument to clabel will control whether the labels are draw# over the line segments of the contour, removing the lines beneath# the labelplt.figure()CS = plt.contour(X, Y, Z)plt.clabel(CS, inline=1, fontsize=10)plt.title('Simplest default with labels')labels = ['line1', 'line2','line3','line4', 'line5', 'line6']for i in range(len(labels)): CS.collections[i].set_label(labels[i])plt.legend(loc='upper left')
Will produce:
However, you might also want to look into annotations for your own need. In my opinion it will give you a more fine grained control on where and what you write on the image, here is the same example with some annotation:
### better with annotation, more flexibleplt.figure(2)CS = plt.contour(X, Y, Z)plt.clabel(CS, inline=1, fontsize=10)plt.title('Simplest default with labels')plt.annotate('some text here',(1.4,1.6))plt.annotate('some text there',(-2,-1.5))
You can create proxy artists to make the legend:
import numpy as npimport matplotlib.pyplot as pltx, y = np.meshgrid(np.arange(10),np.arange(10))z = np.sqrt(x**2 + y**2)cs = plt.contourf(x,y,z,levels=[2,3,4,6])proxy = [plt.Rectangle((0,0),1,1,fc = pc.get_facecolor()[0]) for pc in cs.collections]plt.legend(proxy, ["range(2-3)", "range(3-4)", "range(4-6)"])plt.show()
I had a similar question but needed to go a bit beyond HYRY's answer. To make a package user friendly I wanted ax.legend()
to work without requiring users to pass any handles, which can be achieved by passing the label on to the proxy
proxy = plt.Rectangle((0, 0), 1, 1, fc='red', label='some label')
and then adding the proxy to the axis' patches:
ax.patches += [proxy]
(do ax = plt.gca()
to get the current axis)
This is described in more detail in this answer.