"""Generate the graph of the function $f(x,y) = -(\cos^2x + \cos^2y)^2$."""
import numpy as np
from pylab import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
# Thanks to CT Zhu. See the link as follows
# http://stackoverflow.com/questions/22867620/putting-arrowheads-on-vectors-in-matplotlibs-3d-plot
class Arrow3D(FancyArrowPatch):
def __init__(self, xs, ys, zs, *args, **kwargs):
FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs)
self._verts3d = xs, ys, zs
def draw(self, renderer):
xs3d, ys3d, zs3d = self._verts3d
xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
self.set_positions((xs[0],ys[0]),(xs[1],ys[1]))
FancyArrowPatch.draw(self, renderer)
# draw the frames
ax = Axes3D(figure(), azim = -135, elev = 45)
X = arange(-pi/2, pi/2, pi/100) # set ranges
Y = arange(-pi/2, pi/2, pi/100)
X, Y = meshgrid(X, Y)
Z = -(cos(X)**2+cos(Y)**2)**2 # give the function
ax.plot_wireframe(X, Y, Z, rstride = 10, cstride = 10)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel('f(x,y)')
ax.set_xlim([-pi/2,pi/2])
ax.set_ylim([-pi/2,pi/2])
ax.set_zlim([-4,0])
# draw the contoured quiver
X1 = arange(-pi/2, pi/2, pi/20) # start of the arrows
Y1 = arange(-pi/2, pi/2, pi/20)
for x in X1:
for y in Y1:
ax.add_artist(Arrow3D([x,x+4*sin(x)*cos(x)*(cos(x)**2+cos(y)**2)*0.06],\
[y,y+4*sin(y)*cos(y)*(cos(x)**2+cos(y)**2)*0.06],\
[-3.9999,-4.0000],lw=1, mutation_scale=5,arrowstyle="->", color="r"))
plt.savefig("gradient_vis.svg")