Contour plots

The principal matplotlib routines for creating contour plots are contour and contourf. Sometimes you would like to make a contour plot of a function of two variables; other times you may wish to make
a contour plot of some data you have. Of the two, making a contour plot of a function is simpler. The following program does so:


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm  # color maps
import matplotlib


def pmgauss(x, y):
    r1 = (x - 1) ** 2 + (y - 2) ** 2
    r2 = (x - 3) ** 2 + (y - 1) ** 2
    return 2 * np.exp(-0.5 * r1) - 3 * np.exp(-2 * r2)


a, b = 4, 3

x = np.linspace(0, a, 60)
y = np.linspace(0, b, 45)

X, Y = np.meshgrid(x, y)
Z = pmgauss(X, Y)

fig, ax = plt.subplots(2, 2, figsize=(9.4, 6.5),
                       sharex=True, sharey=True,
                       gridspec_kw={'width_ratios': [4, 5]})

CS0 = ax[0, 0].contour(X, Y, Z, 8, colors='k')
ax[0, 0].clabel(CS0, fontsize=9, fmt='%0.1f')
matplotlib.rcParams['contour.negative_linestyle'] = 'dashed'
ax[0, 0].plot(X, Y, 'o', ms=1, color='lightgray', zorder=-1)

CS1 = ax[0, 1].contourf(X, Y, Z, 12, cmap=cm.gray, zorder=0)
cbar1 = fig.colorbar(CS1, shrink=0.8, ax=ax[0, 1])
cbar1.set_label(label='height', fontsize=10)
plt.setp(cbar1.ax.yaxis.get_ticklabels(), fontsize=8)

lev2 = np.arange(-3, 2, 0.3)
CS2 = ax[1, 0].contour(X, Y, Z, levels=lev2, colors='k',
                       linewidths=0.5)
ax[1, 0].clabel(CS2, lev2[1::2], fontsize=9, fmt='%0.1f')

CS3 = ax[1, 1].contour(X, Y, Z, 10, colors='gray')
ax[1, 1].clabel(CS3, fontsize=9, fmt='%0.1f')
im = ax[1, 1].imshow(Z, interpolation='bilinear',
                     origin='lower', cmap=cm.gray,
                     extent=(0, a, 0, b))
cbar2 = fig.colorbar(im, shrink=0.8, ax=ax[1, 1])
cbar2.set_label(label='height', fontsize=10)
plt.setp(cbar2.ax.yaxis.get_ticklabels(), fontsize=8)

for i in range(2):
    ax[1, i].set_xlabel(r'$x$', fontsize=14)
    ax[i, 0].set_ylabel(r'$y$', fontsize=14)
    for j in range(2):
        ax[i, j].set_aspect('equal')
        ax[i, j].set_xlim(0, a)
        ax[i, j].set_ylim(0, b)
fig.subplots_adjust(left=0.06, bottom=0.07, right=0.99,
                    top=0.99, wspace=0.06, hspace=0.09)
fig.savefig('./figures/contour4.pdf')
plt.show()

The output below shows four different contour plots. All were produced using contour except the upper left plot which was produced using contourf.

All plot the same function, which is the sum of a pair of Gaussians, one positive and the other negative:



After defining the function to be plotted, the next step is to create the x-y array of points at which the function will be evaluated using np.meshgrid. We use np.linspace rather than np.arange to define the extent of the x-y mesh because we want the x range to go precisely from 0 to a=4 and the y range to go precisely from 0 to b=3. We use np.linspace for two reasons. First, if we use np.arange, the array of data points does not include the upper bound, while np.linspace does. This is important for producing the grayscale (or color) background that extends all the way to the upper limits of the x-y ranges in the upper-right plot, produced by contourf, of Figure shown above. Second, to produce smooth-looking contours, one generally needs about 40–200 points in each direction across the plot, irrespective of the absolute magnitude of the numbers being plotted. The number of points is directly specified by np.linspace but must be calculated for np.arange. We follow the convention that the meshgrid variables are capitalized, which seems to be a standard followed by many programmers. It’s certainly not necessary.

The upper-left contour plot takes the X-Y 2D arrays made using gridspec as its first two arguments and Z as its third argument. The third argument tells contour to make approximately 5 different levels
in Z. We give the contour object a name, as it is needed by the clabel call in the next line, which sets the font size and the format of the numbers that label the contours. The line style of the negative contours is set globally to be “dashed” by a call to matplotlib’s rcparams. We also plot the location of the X-Y grid created by gridspec just for the sake of illustrating its function; normally these would not be plotted.

The upper-right contour plot is made using contourf with 12 different Z layers indicated by the different gray levels. The gray color scheme is set by the keyword argument cmap, which here is set to the matplotlib.cm color scheme cm.gray. Other color schemes can be found in the matplotlib documentation by an internet search on “matplotlib choosing colormaps.” The color bar legend on the right is created by the colorbar method, which is attached to fig. It is associated with the upper right plot by the name CS1 of the contourf method and by the keyword argument ax=ax[0, 1]. Its size relative to the plot is determined by the shrink keyword. The font size of the color bar label is set using the generic set property method setp using a somewhat arcane but compact syntax.

For the lower-left contour plot CS2, we manually specify the levels of the contours with the keyword argument levels=lev2. We specify that only every other contour will be labeled numerically with
lev2[1::2] as the second argument of the clabel call in line 38; lev2[0::2] would also label every other contour, but the even ones instead of the odd ones.

The lower-right contour plot CS3 has 10 contour levels and a continuously varying grayscale background created using imshow. The imshow method uses only the Z array to determine the gray levels. The x-y extent of the grayscale background is determined by the keyword argment extent. By default, imshow uses the upper-left corner as its origin. We override the default using the imshow keyword argument origin='lower' so that the grayscale is consistent with the data. The keyword argument iterpolation tells imshow how to interpolate the grayscale between different Z levels.