How to create a legend for a contour plot in matplotlib

In this tutorial, we are going to see how to create a legend for a contour plot in matplotlib in Python. To do this we need to know about matplotlib library. Matplotlib is a data visualization library in python which is used to represent data in the form of plots and graphs.

What is a contour plot?

A contour plot is a technique of plotting constant z slices called a contour on a 2-Dimensional format to represent a 3-dimensional surface. Here, by slices, we mean that a 3-D graph is cut into equal partitions along an axis and the boundaries are merged together in a 2-D plane. The main purpose of a contour plot is to display a 3-Dimensional surface in a 2-Dimensional plot.

Create a contour plot in matplotlib

import matplotlib.pyplot as plt
import numpy as np
x=np.linspace(-2,2)
y=np.linspace(-2,2)
X,Y=np.meshgrid(x,y)
Z=X**2+Y**2
plt.contour(X,Y,Z)
plt.show()

Output:

Explanation:

Import matplotlib module to plot the contour plot and NumPy to number the axis and create grid layout.
The linspace() is a function in NumPy that is used to specify the range of the axis.
The meshgrid() function is used to create a grid layout for the specified values.
The value of z specifies how to slice the 3-D graph into a 2-D plane.
Contour() is a function that is used to generate the contour plot for the given values.
Finally, the show() function is used to display the graph.

Creating a legend in matplotlib

import matplotlib.pyplot as plt
import numpy as np
x=np.linspace(-2,2)
y=np.linspace(-2,2)
X,Y=np.meshgrid(x,y)
Z1=X**2+Y**2
Z2=(X-1)**2+(Y-1)**2
c1 = plt.contour(X, Y, Z1)
c2 = plt.contour(X, Y, Z2)
h1,l1 = c1.legend_elements()
h2,l1 = c2.legend_elements()
plt.legend([h1[0], h2[0]], ['Contour 1', 'Contour 2'])
plt.show()