Matplotlib from Part 16

'''
======================
3D surface (color map)
======================

Demonstrates plotting a 3D surface colored with the coolwarm color map.
The surface is made opaque by using antialiased=False.

Also demonstrates using the LinearLocator and custom formatting for the
z axis tick labels.
'''

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np

plt.rcParams['axes.facecolor'] = '#ffff97ff'

fig = plt.figure()
ax = fig.gca(projection='3d')

fig.patch.set_facecolor('#ffff97ff')

# Make data.
X = np.arange(-8, 8, 0.02)
Y = np.arange(-5, 5, 0.02)
X, Y = np.meshgrid(X, Y)
Z = 1/2*Y**2*np.cos(X) # This is is the function from the pendulum example

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)

# Customize the z axis.
ax.set_zlim(-2.01, 2.01)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()