from mmf.solve.test_problems import TwoDimensional

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

potential = TwoDimensional.MexicanHat()
r = np.linspace(0,1.3,40)
theta = np.linspace(0,2*np.pi,40)
R, Theta = np.meshgrid(r, theta)
X = R*np.cos(Theta)
Y = R*np.sin(Theta) + 1.0
Z = potential.f(X, Y)

fig = plt.figure(figsize=[16,6])
ax = Axes3D(fig, [0,0,0.5,1])
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='b')

r = np.linspace(0,1.1,40)
theta = np.linspace(0,2*np.pi,40)
R, Theta = np.meshgrid(r, theta)
X = R*np.cos(Theta)
Y = R*np.sin(Theta) + 1.0
Z = potential.F(X, Y)

ax = Axes3D(fig, [0.5,0,0.5,1])
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='r')
ax.azim = -23
ax.elev = 65
plt.draw()