# Source code for psi4.driver.diatomic

#
#
# Psi4: an open-source quantum chemistry software package
#
# Copyright (c) 2007-2019 The Psi4 Developers.
#
# The copyrights for code used from other parties are included in
# the corresponding files.
#
# This file is part of Psi4.
#
# Psi4 is free software; you can redistribute it and/or modify
# the Free Software Foundation, version 3.
#
# Psi4 is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with Psi4; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
#

import numpy as np

from psi4 import core
from psi4.driver import constants
from psi4.driver.p4util.exceptions import *

def least_squares_fit_polynomial(xvals, fvals, localization_point, no_factorials=True, weighted=True, polynomial_order=4):
"""Performs and unweighted least squares fit of a polynomial, with specified order
to an array of input function values (fvals) evaluated at given locations (xvals).
See http://dx.doi.org/10.1063/1.4862157, particularly eqn (7) for details. """
xpts = np.array(xvals) - localization_point
if weighted:
R = 1.0
p_nu = 1
epsilon = 1e-3
zvals = np.square(xpts/R)
weights = np.exp(-zvals) / (zvals**p_nu + epsilon**p_nu)
else:
weights = None
fit = np.polynomial.polynomial.polyfit(xpts, fvals, polynomial_order, w=weights)
# Remove the 1/n! coefficients
if no_factorials:
scalefac = 1.0
for n in range(2,polynomial_order+1):
scalefac *= n
fit[n] *= scalefac
return fit

[docs]def anharmonicity(rvals, energies, plot_fit='', mol = None):
"""Generates spectroscopic constants for a diatomic molecules.
Fits a diatomic potential energy curve using a weighted least squares approach
(c.f. http://dx.doi.org/10.1063/1.4862157, particularly eqn. 7), locates the minimum
energy point, and then applies second order vibrational perturbation theory to obtain spectroscopic
constants.  Any number of points greater than 4 may be provided, and they should bracket the minimum.
The data need not be evenly spaced, and can be provided in any order.  The data are weighted such that
those closest to the minimum have highest impact.

A dictionary with the following keys, which correspond to spectroscopic constants, is returned:

:type rvals: list
:param rvals: The bond lengths (in Angstrom) for which energies are
provided, of length at least 5 and equal to the length of the energies array

:type energies: list
:param energies: The energies (Eh) computed at the bond lengths in the rvals list

:type plot_fit: string
:param plot_fit: A string describing where to save a plot of the harmonic and anharmonic fits, the
inputted data points, re, r0 and the first few energy levels, if matplotlib
is available.  Set to 'screen' to generate an interactive plot on the screen instead. If a filename is
provided, the image type is determined by the extension; see matplotlib for supported file types.

:returns: (*dict*) Keys: "re", "r0", "we", "wexe", "nu", "ZPVE(harmonic)", "ZPVE(anharmonic)", "Be", "B0", "ae", "De"
corresponding to the spectroscopic constants in cm-1
"""

angstrom_to_bohr = 1.0 / constants.bohr2angstroms
angstrom_to_meter = 10e-10;

# Make sure the input is valid
if len(rvals) != len(energies):
raise ValidationError("The number of energies must match the number of distances")
npoints = len(rvals)
if npoints < 5:
raise ValidationError("At least 5 data points must be provided to compute anharmonicity")
core.print_out("\n\nPerforming a fit to %d data points\n" % npoints)

# Make sure the molecule the user provided is the active one
molecule = mol if mol is not None else core.get_active_molecule()
molecule.update_geometry()
natoms = molecule.natom()
if natoms != 2:
raise Exception("The current molecule must be a diatomic for this code to work!")
m1 = molecule.mass(0)
m2 = molecule.mass(1)

# Optimize the geometry, refitting the surface around each new geometry
core.print_out("\nOptimizing geometry based on current surface:\n\n");
re = np.mean(rvals)
maxit = 30
thres = 1.0e-9
for i in range(maxit):
derivs = least_squares_fit_polynomial(rvals,energies,localization_point=re)
e,g,H = derivs[0:3]
core.print_out("       E = %20.14f, x = %14.7f, grad = %20.14f\n" % (e, re, g))
if abs(g) < thres:
break
re -= g/H;
if i == maxit-1:
raise ConvergenceError("diatomic geometry optimization", maxit)
core.print_out(" Final E = %20.14f, x = %14.7f, grad = %20.14f\n" % (e, re, g));
if re < min(rvals):
raise Exception("Minimum energy point is outside range of points provided.  Use a lower range of r values.")
if re > max(rvals):
raise Exception("Minimum energy point is outside range of points provided.  Use a higher range of r values.")

# Convert to convenient units, and compute spectroscopic constants
d0,d1,d2,d3,d4 = derivs*constants.hartree2aJ
core.print_out("\nEquilibrium Energy %20.14f Hartrees\n" % e)
core.print_out("Quadratic Force Constant %14.7f MDYNE/A\n" % d2)
core.print_out("Cubic Force Constant     %14.7f MDYNE/A**2\n" % d3)
core.print_out("Quartic Force Constant   %14.7f MDYNE/A**3\n" % d4)

hbar = constants.h / (2.0 * np.pi)
mu = ((m1*m2)/(m1+m2))*constants.amu2kg
we = 5.3088375e-11 * np.sqrt(d2/mu)
wexe = (1.2415491e-6)*(we/d2)**2 * ((5.0*d3*d3)/(3.0*d2)-d4)

# Rotational constant: Be
I = ((m1*m2)/(m1+m2)) * constants.amu2kg * (re * angstrom_to_meter)**2
B = constants.h / (8.0 * np.pi**2 * constants.c * I)

# alpha_e and quartic centrifugal distortion constant
ae = -(6.0 * B**2 / we) * ((1.05052209e-3*we*d3)/(np.sqrt(B * d2**3))+1.0)
de = 4.0*B**3 / we**2

# B0 and r0 (plus re check using Be)
B0 = B - ae / 2.0
r0 = np.sqrt(constants.h / (8.0 * np.pi**2 * mu * constants.c * B0))
recheck = np.sqrt(constants.h / (8.0 * np.pi**2 * mu * constants.c * B))
r0 /= angstrom_to_meter;
recheck /= angstrom_to_meter;

# Fundamental frequency nu
nu = we - 2.0 * wexe;
zpve_nu = 0.5 * we - 0.25 * wexe;

# Generate pretty pictures, if requested
if(plot_fit):
try:
import matplotlib.pyplot as plt
except ImportError:
msg = "\n\tPlot not generated; matplotlib is not installed on this machine.\n\n"
print(msg)
core.print_out(msg)

# Correct the derivatives for the missing factorial prefactors
dvals = np.zeros(5)
dvals[0:5] = derivs[0:5]
dvals /= 2
dvals /= 6
dvals /= 24

# Default plot range, before considering energy levels
minE = np.min(energies)
maxE = np.max(energies)
minR = np.min(rvals)
maxR = np.max(rvals)

# Plot vibrational energy levels
we_au = we / constants.hartree2wavenumbers
wexe_au = wexe / constants.hartree2wavenumbers
coefs2 = [ dvals, dvals, dvals ]
coefs4 = [ dvals, dvals, dvals, dvals, dvals ]
for n in range(3):
Eharm = we_au*(n+0.5)
Evpt2 = Eharm - wexe_au*(n+0.5)**2
coefs2[-1] = -Eharm
coefs4[-1] = -Evpt2
roots2 = np.roots(coefs2)
roots4 = np.roots(coefs4)
xvals2 = roots2 + re
xvals4 = np.choose(np.where(np.isreal(roots4)), roots4).real + re
Eharm += dvals
Evpt2 += dvals
plt.plot(xvals2, [Eharm, Eharm], 'b', linewidth=1)
plt.plot(xvals4, [Evpt2, Evpt2], 'g', linewidth=1)
maxE = Eharm
maxR = np.max([xvals2,xvals4])
minR = np.min([xvals2,xvals4])

# Find ranges for the plot
dE = maxE - minE
minE -= 0.2*dE
maxE += 0.4*dE
dR = maxR - minR
minR -= 0.2*dR
maxR += 0.2*dR

# Generate the fitted PES
xpts = np.linspace(minR, maxR, 1000)
xrel = xpts - re
xpows = xrel[:, None] ** range(5)
fit2 = np.einsum('xd,d', xpows[:,0:3], dvals[0:3])
fit4 = np.einsum('xd,d', xpows, dvals)

# Make / display the plot
plt.plot(xpts, fit2, 'b', linewidth=2.5, label='Harmonic (quadratic) fit')
plt.plot(xpts, fit4, 'g', linewidth=2.5, label='Anharmonic (quartic) fit')
plt.plot([re, re], [minE, maxE], 'b--', linewidth=0.5)
plt.plot([r0, r0], [minE, maxE], 'g--', linewidth=0.5)
plt.scatter(rvals, energies, c='Black', linewidth=3, label='Input Data')
plt.legend()

plt.xlabel('Bond length (Angstroms)')
plt.ylabel('Energy (Eh)')
plt.xlim(minR, maxR)
plt.ylim(minE, maxE)
if plot_fit == 'screen':
plt.show()
else:
plt.savefig(plot_fit)
core.print_out("\n\tPES fit saved to %s.\n\n" % plot_fit)

core.print_out("\nre     = %10.6f A  check: %10.6f\n" % (re, recheck))
core.print_out("r0       = %10.6f A\n" % r0)
core.print_out("we       = %10.4f cm-1\n" % we)
core.print_out("wexe     = %10.4f cm-1\n" % wexe)
core.print_out("nu       = %10.4f cm-1\n" % nu)
core.print_out("ZPVE(nu) = %10.4f cm-1\n" % zpve_nu)
core.print_out("Be       = %10.4f cm-1\n" % B)
core.print_out("B0       = %10.4f cm-1\n" % B0)
core.print_out("ae       = %10.4f cm-1\n" % ae)
core.print_out("De       = %10.7f cm-1\n" % de)
results = {
"re"               :  re,
"r0"               :  r0,
"we"               :  we,
"wexe"             :  wexe,
"nu"               :  nu,
"ZPVE(harmonic)"   :  zpve_nu,
"ZPVE(anharmonic)" :  zpve_nu,
"Be"               :  B,
"B0"               :  B0,
"ae"               :  ae,
"De"               :  de
}
return results