| #!/usr/bin/env python3 | |
| ## Potential Energy Surface | |
| ## Based on Matplotlib 3D plotting documentation at | |
| ## https://matplotlib.org/stable/gallery/index.html#d-plotting | |
| ## Load the libraries | |
| import numpy as np | |
| from mpl_toolkits.mplot3d import axes3d | |
| import matplotlib.pyplot as plt | |
| ## Set up the figure | |
| fig = plt.figure(figsize=(8,4),dpi=300) | |
| ## Change the subplot space to remove white space | |
| fig.subplots_adjust(left=0, right=1.25, bottom=0, top=1) | |
| ax = fig.add_subplot(111, projection="3d") | |
| #ax = fig.gca(projection="3d") | |
| ## Create the PES gridpoints | |
| ## Need to be negative for Z so that the TS is high and the | |
| ## reactant and product are low | |
| x, y = np.mgrid[-1:1:30j, -1:1:30j] | |
| z = -np.cos(np.pi*-x)*np.cos(np.pi*-y) | |
| ## Create the line defining the reaction path | |
| ## X = zero, y = evenly spaced, z matches the PES points | |
| x_path = np.zeros(shape=100) | |
| y_path = np.linspace(-1,1, num=100) | |
| z_path = -np.cos(np.pi*x_path)*np.cos(np.pi*y_path) | |
| ## Create the chain-of-replica beads at a 5th of the reaction path interval | |
| ## Give them a slight offset so they're visible | |
| x_bead = np.zeros(shape=20) | |
| y_bead = np.linspace(-1,1, num=20) | |
| z_bead = -np.cos(np.pi*x_bead)*np.cos(np.pi*y_bead) + -0.025*np.pi | |
| ## The PES contour plot | |
| ## Break into 3 groups so that the chain-of-replica beads loook right! | |
| ax.plot_surface(-x[0:14], -y[0:14], -z[0:14], cmap="winter_r", lw=-1, | |
| rstride=1, cstride=1, alpha=0.75) | |
| ax.plot_surface(-x[13:17], -y[13:17], -z[13:17], cmap="winter_r", lw=-1, | |
| rstride=1, cstride=1, alpha=0.75) | |
| ax.plot_surface(-x[16:], -y[16:], -z[16:], cmap="winter_r", lw=-1, | |
| rstride=1, cstride=1, alpha=0.75) | |
| ## The line defining the reaction path | |
| ax.plot(-x_path, -y_path, -z_path, lw=2, c='black') | |
| ## Points for Chain-of-replica beads | |
| ax.scatter(-x_bead, -y_bead, -z_bead, zdir='z', s=55, c='darkorange', | |
| depthshade=False, alpha=1) | |
| ## Wire mesh | |
| # ax.plot_wireframe(-x, -y, -z, rstride=1, cstride=1) | |
| # Get current rotation angle and elevation | |
| # print(ax.azim) | |
| # print(ax.elev) | |
| # Set viewpoint (GUI gives values in top right corner) | |
| ax.view_init(azim=-40, elev=63.) | |
| ## Option to save a series of images at various rotations (to make a gif???) | |
| # for ii in xrange(0,360,1): | |
| # ax.view_init(elev=10., azim=ii) | |
| # savefig("movie%d.png" % ii) | |
| ## Add the axis labels using the math format for subscript | |
| ax.set_xlabel('$q_1$', labelpad=-5) | |
| ax.set_ylabel('$q_2$', labelpad=-5) | |
| ax.set_zlabel(r'Energy', labelpad=-10) | |
| ax.get_xaxis().set_ticklabels([]) | |
| ax.get_yaxis().set_ticklabels([]) | |
| ax.set_zticklabels([]) | |
| ## Show the plot so you can orient it properly (and re-save with GUI) | |
| # plt.show() | |
| ## Save the figure | |
| fig.savefig('PES_w_beads_image.png', dpi=fig.dpi, bbox_inches='tight', | |
| transparent=True) |
