I thought I’d refresh my memory of the matplotlib library, an add-on package for Python that can create plots and graphs. I regularly use the plotting functions in R, SciLab, Excel, and Python, but using these isn’t easy or intuitive so I like to stay in practice.

One of my standard practice problems is to create a 3D graph of Rastrigin’s function. Rastrigin’s function is a standard benchmark optimization problem because it has many false local minima, but only one true global minimum at (0, 0).

Here’s one way to graph Rastrigin’s function using matplotlb:

# rastrigin_3d.py

from matplotlib import cm  # color map
from mpl_toolkits.mplot3d import Axes3D
import math
import matplotlib.pyplot as plt
import numpy as np

X = np.linspace(-4, 4, 200)    
Y = np.linspace(-4, 4, 200)    
X, Y = np.meshgrid(X, Y)

Z = (X**2 - 10 * np.cos(2 * 3.14 * X)) + \
  (Y**2 - 10 * np.cos(2 * 3.14 * Y)) + 20

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, rstride=1, \
  cstride=1, cmap=cm.jet)
    
# plt.savefig('rastrigin.png')
plt.show()

I don’t enjoy creating graphs, but in many situations graphs are extremely useful to help explain a machine learning example.