# load arrays with polar coordinates of points on a surface, then plot them # this file is unit02_draw_surface_polar.py # import libraries from mpl_toolkits.mplot3d import Axes3D import numpy as np import matplotlib.pyplot as plt from matplotlib import cm # define parameters for our plot rmax = 3 thetamax = 2*np.pi # number of rows and columns in our grid number_r_values = 61 number_theta_values = number_r_values # 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 # get the r, theta values in the grid. r = np.linspace(0, rmax, number_r_values) theta = np.linspace(0, thetamax, number_theta_values) rgrid, thetagrid = np.meshgrid(r, theta) # transform them to cartesian system xgrid, ygrid = rgrid*np.cos(thetagrid), rgrid*np.sin(thetagrid) # 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, polar-- 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, cmap=cm.coolwarm ) # now save the plot to a png (portable network graphics) file plt.savefig("SurfacePlotPolar.png")