PyDiatomic phase determination from wavefunction

phase.py

import numpy as np
import cse
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import scipy.constants as const
from scipy.special import spherical_jn, spherical_yn

au = 0.52917720859
eh = 27.21138386

def Mies(x, a, b, phase):
    global zz, k
    xa = x*1e-10  # x in metres
    kxa = k*xa
    return (a*spherical_jn(0, kxa) + b*spherical_yn(0, kxa))*np.sqrt(k)*xa*zz

E = 0.3  # a.u.

X = cse.Cse('O2', VT=['potentials/X3S-1.dat'], en=E*eh*8065.541)  # en in cm-1
wf = X.wf[:, 0, 0]

lr = X.R > 9.7  # long-range (large) part of internuclear distance

# estimates - amplitude, k - based on energy in MKSA units -------------
zz = np.sqrt(2*X.μ/np.pi)/const.hbar
edash = (E*eh - X.VT[0, 0, -1])*const.e  # dissociation energy in joules

k = np.sqrt(2*X.μ*edash)/const.hbar  # 1/metres
amp = np.sqrt(2*X.μ/np.pi/k)/const.hbar

print('estimates from energy ------------------')
print(f'      edash = {edash:g} J, amp = {amp:g}, kwave = {k:g} m-1\n')

# PyDiatomic -------------
delta = np.arctan(X.eig[0])
wfamp = wf[lr].max()
ai = 1/X.AI[0, 0]

print('PyD values ------------------------------')
print(f'   wf_amp = {wfamp:g}, phase = {delta:g}, K = {X.K[0, 0]:g}\n')
print(f'        A = {ai:g}, B = {X.B[0, 0]:g}, K = {X.K[0, 0]:g}')
print()

# wavefunction fit -------------
pamp = wfamp*np.sqrt(k)/zz
p = [pamp, pamp, delta]
par, err = curve_fit(Mies, X.R[lr], wf[lr], p)

a, b, phase = par

print('\nfit to wavefunction --------------------')
print(f'        A = {a:g}, B = {b:g}, K = {b/a:g}')
print(f'    phase = {phase:g}, delta = {np.arctan(b/a):g}')

# plots -----------------------------------
fig, (ax0, ax1) = plt.subplots(1, 2, sharey=True)
ax0.plot(X.R, wf)
ax0.axis(xmin=0.8, xmax=1.2)
ax0.set_title('inner')
ax0.set_ylabel('amplitude')

ax1.plot(X.R[lr], wf[lr], label='wavefunction')
ax1.plot(X.R[lr], Mies(X.R[lr], *par), label='asymp. fit')
ax1.set_title('long range')
ax1.legend()

plt.xlabel(r'internuclear distance ($\AA$)')
plt.show()