# load arrays with coordinates of points on a surface, then plot them
# this file is unit02_draw_surface.py 

# import libraries
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
# let's also import a color map so we can make the picture prettier
from matplotlib import cm

# define parameters for our plot
xmin = -3
xmax = -xmin
ymin = xmin
ymax = -ymin

# number of rows and columns in our grid
nrows = 61
ncolumns = 61

# number of rows and columns per grid line to be drawn
rowsPerGrid = 2
columnsPerGrid = 2

# define the coefficients in the potential
xcoeff = 2
ycoeff = 5

# create the arrays. 

# first get the x values of the "columns" in the grid.
x = np.linspace(xmin, xmax, ncolumns)

# now get the y values of the "rows" in the grid.
y = np.linspace(ymin, ymax, nrows)

# now generate the x and y coordinates of all nrows * ncolumns points in the grid.
xgrid, ygrid = np.meshgrid(x,y)

# now generate the potential for all points on our grid. 
zsurface = xcoeff * xgrid**2 + ycoeff * ygrid**2

# now create a (blank) figure so we can set some of its attributes.
fig = plt.figure()

# "gca" is "get current axes." set the projection attribute to 3D. 
ax = fig.gca(projection='3d')

# set labels and title
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("potential energy")
ax.set_title("Potential energy surface-- George Gollin") 

# now put the graph into the blank figure. Note the line-continuation character
# colormap (cmap) argument is optional.
surf = ax.plot_surface(xgrid, ygrid, zsurface, rstride=rowsPerGrid, \
# cstride=columnsPerGrid, color="#FF0000" )
cstride=columnsPerGrid, cmap=cm.coolwarm)

# now save the plot to a png (portable network graphics) file
plt.savefig("SurfacePlotCartesian.png")