We start by importing the Axes3D class from the mpl_toolkits.mplot3d library, which is the Matplotlib object used for creating three-dimensional plots. We also import the cm class, which represents a color map. We then define a function to be plotted, with the following line of code:
f = lambda x,y: x**3 - 3*x*y**2
The next step is to define the Figure object and an Axes object with a 3D projection, as done in the following lines of code:
fig = plt.figure(figsize=(12,6))
ax = fig.add_subplot(1,2,1,projection='3d')
Notice that the approach used here is somewhat different than the other recipes in this chapter. We are assigning the output of the figure() function call to the fig variable and then adding the subplot by calling the add_subplot() method from the fig object.
The next few lines of code, shown as follows, compute the data for the plot:
xvalues = np.linspace(-2,2,100)
yvalues = np.linspace(-2,2,100)
xgrid, ygrid = np.meshgrid(xvalues, yvalues)
zvalues = f(xgrid, ygrid)
The most important feature of this code is the call to meshgrid(). This is a NumPy convenience function that constructs grids suitable for three-dimensional surface plots. To understand how this function works, run the following code:
xvec = np.arange(0, 4)
yvec = np.arange(0, 3)
xgrid, ygrid = np.meshgrid(xvec, yvec)
After running this code, the xgrid array will contain the following values:
array([[0, 1, 2, 3],
[0, 1, 2, 3],
[0, 1, 2, 3]])
The ygrid array will contain the following values:
array([[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2]])
Notice that the two arrays have the same dimensions. Each grid point is represented by a pair of the (xgrid[i,j],ygrid[i,j]) type. This convention makes the computation of a vectorized function on a grid easy and efficient, with the f(xgrid, ygrid) expression.
The next step is to generate the surface plot, which is done with the following function call:
surf = ax.plot_surface(xgrid, ygrid, zvalues,
rstride=5, cstride=5,
linewidth=0, cmap=cm.plasma)
The first three arguments, xgrid, ygrid, and zvalues, specify the data to be plotted. We then use the rstride and cstride options to select a subset of the grid points. Notice that the xvalues and yvalues arrays both have length 100, so that xgrid and ygrid will have 10,000 entries each. Using all grid points would be inefficient and produce a poor plot from the visualization point of view. Thus, we set rstride=5 and cstride=5, which results in a plot containing every fifth point across each row and column of the grid.
The next option, linewidth=0, sets the line width of the plot to zero, preventing the display of a wireframe. The final argument, cmap=cm.plasma, specifies the color map for the plot. We use the cm.plasma color map, which has the effect of plotting higher functional values with a hotter color. Matplotlib offer as large number of built-in color maps, listed at https://matplotlib.org/examples/color/colormaps_reference.html.
Next, we add the filled contour plot with the following code:
ax = fig.add_subplot(1,2,2)
ax.contourf(xgrid, ygrid, zvalues, 30,
cmap=cm.plasma)
Notice that, when selecting the subplot, we do not specify the projection option, which is not necessary for two-dimensional plots. The contour plot is generated with the contourf() method. The first three arguments, xgrid, ygrid, zvalues, specify the data points, and the fourth argument, 30, sets the number of contours. Finally, we set the color map to be the same one used for the surface plot.
The final component of the plot is a color bar, which provides a visual representation of the value associated with each color in the plot, with the fig.colorbar(surf, aspect=18) method call.
Notice that we have to specify in the first argument which plot the color bar is associated to. The aspect=18 option is used to adjust the aspect ratio of the bar. Larger values will result in a narrower bar.
To finish the plot, we call the tight_layout() function. This adjusts the sizes of each plot, so that axis labels are displayed correctly.