# Crabtree Demo

print("Welcome to python! Loading dependencies will take about 5 minutes. Please wait...\n")
!mamba install psi4 -yc psi4/label/dev >/dev/null
%matplotlib inline
# -*- coding: utf-8 -*-

# NH3 Inversion Potential Calculator
# Kyle N. Crabtree
#
# This program calculates and plots the inversion potential of ammonia at a chosen
# level of theory and standard basis set. The coordinate used in this program
# is the "umbrella angle": the symmetric angle with respect to planarity (i.e.,
# theta = 0 is the planar D_3h geometry, and for other angles, all three H
# atoms are symmetrically raised above or lowered beneath the D_3h horizontal
# symmetry plane). The numbour of anngles and the max angle are set in the
# beginning of the code, and only one half of the potential is calculated.
# By specifying multiple basis sets or methods, the program will calculate
# each in turn, placing all results onto a single plot. Output files are
# generated and named according to the current calculation.
#
# Parameters:
# methods: list of methods as strings. Example: methods = ['scf','mp2']
# basis_sets: list of basis sets as strings.
#             Example: basis_sets = ['3-21g','6-31g','cc-pvdz','cc-pvtz']
# details: Print starting geometry before each calculation
# optimize_path: If true, a constrained optimization will be performed
#                at each angle to adjust the N-H bond lengths
# numangles: The number of points at which to perform the calculation
# maxangle: Maximum angle. The program will calculate numangles values between
#           -abs(maxangle) and 0, inclusive.

import psi4
import numpy as np
import matplotlib.pyplot as plt

#methods = ['scf','b3lyp','ccsd(t)']
methods = ['scf']
basis_sets = ['3-21g','6-31g*','cc-pvtz']
details = False
optimize_path = False
numangles = 41
maxangle = 40.
psi4.set_memory('4096 MB')
print("\n\nWelcome to psi4! Beginning computation\n")
#In this function, calculate the coordinates of the hydrogen
#atoms (units: A) in ammonia as a function of the angle (0 degrees = planar)
#and the N-H bond length (in A)
def calc_coordinates(angle_degrees, nhlength):
#atom = [x,y,z]
N = [0.0,0.0,0.0]

#use trig functions and geometry to calculate x,y,z coordinates
#of hydrogen atoms here
angle = angle_degrees/180.0*np.pi
z = np.sin(angle)*nhlength
x1 = np.cos(angle)*nhlength
x23 = np.cos(2.0*np.pi/3.0)*np.cos(angle)*nhlength
y = np.sin(2.0*np.pi/3.0)*np.cos(angle)*nhlength

H1 = [x1,0.0,z] #insert x, y, and z variables here
H2 = [x23,y,z]
H3 = [x23,-y,z]

return """
N {0} {1} {2}
H {3} {4} {5}
H {6} {7} {8}
H {9} {10} {11}
""".format(N[0],N[1],N[2],H1[0],H1[1],H1[2],H2[0],H2[1],H2[2],H3[0],H3[1],H3[2])

plt.figure()

count = 0
total = len(methods)*len(basis_sets)
res_angles, res_egys, res_bonds, res_barriers, res_calcs = [], [], [], [], []
for method in methods:
for basis in basis_sets:
count+=1
calc = method+'/'+basis
res_calcs.append(calc)
outfile = method+'_'+basis+'_output.dat'
outfile = outfile.replace('*','s')
psi4.core.set_output_file(outfile,False)

print('\n-------------------------------------')
print('Starting '+calc+' calculation ('+str(count)+'/'+str(total)+')')
print('-------------------------------------')

R = 1.0
A = 110.0
D = 120.0

nh3 = psi4.geometry("""
X
N 1 1.0
H 2 {0} 1 {1}
H 2 {0} 1 {1} 3 120.0
H 2 {0} 1 {1} 3 -120.0
""".format(R,A))

print('Performing initial geometry optimization.')
psi4.set_module_options('optking',{'FROZEN_BEND':''})
egyend = psi4.optimize(calc,molecule=nh3)
nh3.print_distances()
nh3.print_bond_angles()
dm = nh3.distance_matrix().to_array()
nhbond = dm[0,1]*psi4.constants.bohr2angstroms
res_bonds.append(nhbond)

#find optimized umbrella angle.
#Get geometry matrix, and translate N atom to origin
gm = nh3.geometry().to_array().transpose()
for i in range(0,3):
d = gm[i,0]
for j in range(0,4):
gm[i,j] -= d

#rotate H1 into the xz plane by rotating about z axis
theta = 0.0
if abs(gm[0,1]) > 1e-10:
theta = np.arctan(-gm[1,1]/gm[0,1])
else:
theta=np.pi/2.0

Rtheta = np.zeros((3,3))
Rtheta[0,0] = np.cos(theta)
Rtheta[0,1] = -np.sin(theta)
Rtheta[1,0] = np.sin(theta)
Rtheta[1,1] = np.cos(theta)
Rtheta[2,2] = 1.0
gmrot = np.dot(Rtheta,gm)

#rotate H1 onto z axis by rotating about y axis
theta2 = 0.0
if abs(gmrot[0,1]) > 1e-10:
theta2 = np.arctan(gmrot[2,1]/gmrot[0,1])
else:
theta2 = np.pi/2.0

Rtheta2 = np.zeros((3,3))
Rtheta2[0,0] = np.cos(theta2)
Rtheta2[0,2] = np.sin(theta2)
Rtheta2[1,1] = 1.0
Rtheta2[2,0] = -np.sin(theta2)
Rtheta2[2,2] = np.cos(theta2)
gmrot2 = np.dot(Rtheta2,gmrot)

#rotate about x axis until z(H2) = z(H3)
theta3 = 0.0
if abs(gmrot2[2,3]-gmrot2[2,2]) > 1e-10:
if abs(gmrot2[1,2]-gmrot2[1,3]) > 1e-10:
theta3 = np.arctan((gmrot2[2,3]-gmrot2[2,2])/(gmrot2[1,2]-gmrot2[1,3]))
else:
theta3 = np.pi/2.0

Rtheta3 = np.zeros((3,3))
Rtheta3[0,0] = 1.0
Rtheta3[1,1] = np.cos(theta3)
Rtheta3[1,2] = -np.sin(theta3)
Rtheta3[2,1] = np.sin(theta3)
Rtheta3[2,2] = np.cos(theta3)

gmrot3 = np.dot(Rtheta3,gmrot2)

#umbrella angle is the rotation angle about the y axis required for
#Z(H1) = Z(H2) = Z(H3)
umbangle = (np.arctan((gmrot3[2,2]-gmrot3[2,1])/(gmrot3[0,2]-gmrot3[0,1])))/np.pi*180.0
res_angles.append(umbangle)
res_egys.append(egyend)

print('Optimized energy: {:.6f} Hartree'.format(egyend))
print('Optimized N-H bond length: {:.3f} Angstrom'.format(nhbond))
print('Optimized umbrella angle: {:.3f} deg'.format(umbangle))

egyend *= psi4.constants.hartree2kcalmol

angles = []
for i in range(0,numangles):
angles.append(i*abs(maxangle)/(numangles-1.) - abs(maxangle))

angle_array = np.zeros(2*len(angles)-1)
egy_array = np.zeros(2*len(angles)-1)

idx = 0
for angle in angles:
geo = calc_coordinates(angle,nhbond)
if details:
print(geo)
nh3 = psi4.geometry(geo)
egy = 0.
if optimize_path:
psi4.set_module_options('optking',{'FROZEN_BEND':'2 1 3 3 1 4 4 1 2'})
egy = psi4.optimize(calc,molecule=nh3)*psi4.constants.hartree2kcalmol - egyend
else:
egy = psi4.energy(calc,molecule=nh3)*psi4.constants.hartree2kcalmol - egyend
print("Angle: {:5.1f} deg, Egy: {:7.4f} kcal/mol".format(angle,egy))
angle_array[idx] = angle
angle_array[-(idx+1)] = -angle
egy_array[idx] = egy
egy_array[-(idx+1)] = egy
idx+=1

barrier = egy_array[idx-1]
res_barriers.append(barrier)
print('\n-------------------------------------')
print('Barrier height: {:6.4f} kcal/mol'.format(barrier))

plt.plot(angle_array,egy_array,label=calc.upper())

filename = 'nh3_'+method+'_'+basis
if optimize_path:
filename += '_opt'
filename+='.txt'
filename = filename.replace('*','s')
print('Detailed calculation information written to '+outfile)
print('Potential energy data written to '+filename)
print('-------------------------------------')
with open(filename,'w') as f:
f.write('Angle (deg)\tEnergy (kcal/mol)')
for angle,egy in zip(angle_array,egy_array):
f.write('\n'+str(angle)+'\t'+str(egy))

plt.xlabel("Umbrella Angle (degrees)")
plt.ylabel("Relative Potential Energy (kcal/mol)")
plt.title("NH$_3$ Inversion Potential")
plt.legend()

print("\n\nSummary\n")
print("--------------------------------------------------------------------------")
print("Calculation       |Energy (au)|N-H Bond (A)|Angle (deg)|Barrier (kcal/mol)")
print("--------------------------------------------------------------------------")
for c, e, b, a, ba in zip(res_calcs, res_egys, res_bonds, res_angles, res_barriers):
print("{:18s} {:<11.4f} {:<12.4f} {:<11.3f} {:<18.4f}".format(c,e,b,a,ba))