Functions of Two Variables

Functions of two variables can be visualized in many ways: heat maps in 2D, contour curves in 2D, wireframes in 3D, surfaces in 3D.

We begin by defining a function of two variables, building an array of points in the domain using meshgrid, then constructing an array of output values. This is the data we need for graphing the function.

def f(x,y):
    return 2.7 - 2 * np.cos(x) * np.cos(y) - 0.7*np.cos(np.pi - 2*x)

x = np.linspace(0, 2*np.pi, 100)
y = np.linspace(0, 2*np.pi, 100)
X,Y = np.meshgrid(x, y)
Z = f(X, Y)

Colormap and Contour Figures

Colormaps and contour figures are useful for plotting functions of two variables. In most of these functions we will use a colormap to encode one dimension of the data. There are a number of predefined colormaps. It is relatively straightforward to define custom colormaps. For a list of pre-defined colormaps, see:
http://www.scipy.org/Cookbook/Matplotlib/Show_colormaps

Using pcolor

fig, ax = plt.subplots()

p = ax.pcolor(X/(2*np.pi), Y/(2*np.pi), Z, cmap=matplotlib.cm.RdBu,
    vmin=abs(Z).min(), vmax=abs(Z).max())
cb = fig.colorbar(p, ax=ax)

Using imshow

fig, ax = plt.subplots()

im = ax.imshow(Z, cmap=matplotlib.cm.RdBu, vmin=abs(Z).min(),
     vmax=abs(Z).max(), extent=[0, 1, 0, 1])
im.set_interpolation('bilinear')

cb = fig.colorbar(im, ax=ax)

3D Figures

To use 3D graphics in matplotlib, we first need to create an instance of the Axes3D class. 3D axes can be added to a matplotlib figure canvas in exactly the same way as 2D axes; or, more conveniently, by passing a projection='3d' keyword argument to the add_axes or add_subplot methods.

from mpl_toolkits.mplot3d.axes3d import Axes3D

Surface Plots

fig = plt.figure(figsize=(14,6))

# `ax` is a 3D-aware axis instance because of the projection='3d'
#  keyword argument to add_subplot
ax = fig.add_subplot(1, 2, 1, projection='3d')

p = ax.plot_surface(X, Y, Z, rstride=4, cstride=4, linewidth=0)

# surface_plot with color grading and color bar
ax = fig.add_subplot(1, 2, 2, projection='3d')
p = ax.plot_surface(X, Y, Z, rstride=1, cstride=1,
    cmap=matplotlib.cm.coolwarm, linewidth=0, antialiased=False)
cb = fig.colorbar(p, shrink=0.5)

Wireframe Plots

fig = plt.figure(figsize=(8,6))

ax = fig.add_subplot(1, 1, 1, projection='3d')

p = ax.plot_wireframe(X, Y, Z, rstride=4, cstride=4)

Contour Plots on Surfaces

g = lambda x,y : x**2-y**2

x = np.linspace(-5,5,128)
y = np.linspace(-5,5,128)
X,Y = np.meshgrid(x,y)
Z = g(X,Y)

fig = plt.figure(figsize=(12,6))

# `ax` is a 3D-aware axis instance because of the projection='3d'
# keyword argument to add_subplot
ax = fig.add_subplot(1, 2, 1, projection='3d')
ax.plot_surface(X, Y, Z, rstride=4, cstride=4, linewidth=0, alpha=0.75)
ax.contour(X,Y,Z, colors='red', zdir='z', vmin=abs(Z).min(),
    vmax=abs(Z).max(), extent=[-5, 5, -5, 5], linewidths=2)
ax.contour(X,Y,Z, colors='green', zdir='y', vmin=abs(Z).min(),
    vmax=abs(Z).max(), extent=[-5, 5, -5, 5], linewidths=2)

Contour Plots with Projections

fig = plt.figure(figsize=(8,6))

ax = fig.add_subplot(1,1,1, projection='3d')

ax.plot_surface(X, Y, Z, rstride=4, cstride=4, alpha=0.25)
ax.plot_wireframe(X, Y, Z, rstride=4, cstride=4, lw=0.5,
    color="black", alpha=0.5)
cset = ax.contour(X, Y, Z, zdir='z', offset=-np.pi,
    cmap=matplotlib.cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='x', offset=-np.pi,
    cmap=matplotlib.cm.coolwarm)
cset = ax.contour(X, Y, Z, zdir='y', offset=3*np.pi,
    cmap=matplotlib.cm.coolwarm)

ax.set_xlim3d(-np.pi, 2*np.pi);
ax.set_ylim3d(0, 3*np.pi);
ax.set_zlim3d(-np.pi, 2*np.pi);