Rastrigin Function Graph Using matplotlib With Edge Colors

The Rastrigin function is a standard benchmark problem for optimization algorithms. The general form of the function in n dimensions is f(X) = Sum[xi^2 – 10*cos(2*pi*xi)] + 10n. For the specific case when n = 2, the function is f(x1, x2) = [(x1^2 – 10*cos(2*pi*x1)) + (x2^2 – 10*cos(2*pi*x2))] + 20.

I have graphed the Rastrigin function using R, MatLab, SciLab (a MatLab clone), Excel, and Python’s matplotlib library. I most often use matplotlib. When I used matplotlib a few years ago, the default surface plot had a wireframe outline that made the shape of the function easier to see. But at some point in early 2018, the default settings for the surface_plot() function no longer added a wireframe. This is nice for very simple graphs, but for complex shapes like the Rastrigin function, the lack of a wireframe made it difficult to understand the shape of the graph.

Here is the Rastrigin function for n = 2, using default no-wireframe outline (on the left), and with a wireframe outline added (on the right). To me at least, the graph on the right is clearly preferable.

My annoyance level finally hit a threshhold and I sat down to figure out how to add a wireframe outline to a matplotlib surface plot. If you aren’t familiar with graphing libraries, your first impression is likely, “Oh come on McCaffrey, how hard could it be to look at the matplotlib documentation and find the appropriate parameter(s)?” If you are familiar with graphing libraries, your first impression is something like, “Yeah, good luck with that. Let me know what you find over the next several days.”

Anyway, as always, the answer is easy once you know what to do: add the edgecolor and linewidth parameters to the plot_surface() function. The diffiulty is that these two parameters aren’t explicit parameters of plot_surface() — they are parameters to a hidden base class Poly3DCollection class.

Anyway, I don’t think there’s a significant moral to this story/post. But maybe it’s that for “us” (meaning me and you and anyone else who reads blog posts like this one), we mentally store up and remember minor technical problems, put them in our mental TODO list, and these problems don’t go away until we eventually solve them.

The matplotlib show_plot() function constructs a surface by using many small polygons. Modern photo editing software can convert a photograph to an image made of polygons. I’m not a huge fan of the low poly portraits because they lack the soul of pure-human art, but here are three nice ones that somewhat resemble the image of the Rastrigin function graph. Notice that none of the three illustrations uses a line edge.

# rastrigin_graph.py

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

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

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

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, \
  rstride=1, cstride=1, cmap=cm.jet, \
  edgecolor='darkred', linewidth=0.1)
# plt.savefig('rastrigin.png')
This entry was posted in Machine Learning. Bookmark the permalink.

Leave a Reply

Please log in using one of these methods to post your comment:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s