194 lines
6.8 KiB
Python
194 lines
6.8 KiB
Python
|
"""
|
||
|
|
Provides interface n2v interface to Psi4
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
from .engine import Engine
|
||
|
|
import numpy as np
|
||
|
|
from opt_einsum import contract
|
||
|
|
|
||
|
|
try:
|
||
|
|
import psi4
|
||
|
|
psi4.set_options({"save_jk" : True})
|
||
|
|
has_psi4 = True
|
||
|
|
except ImportError:
|
||
|
|
has_psi4 = False
|
||
|
|
|
||
|
|
if has_psi4:
|
||
|
|
from ..grid import Psi4Grider
|
||
|
|
class Psi4Engine(Engine):
|
||
|
|
"""
|
||
|
|
Psi4 Engine Class
|
||
|
|
"""
|
||
|
|
|
||
|
|
def set_system(self, molecule, basis, ref='1', pbs='same', wfn=None):
|
||
|
|
"""
|
||
|
|
Initializes geometry and basis infromation
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
molecule: psi4.core.Molecule
|
||
|
|
Molecule of the system used
|
||
|
|
basis: str
|
||
|
|
Basis set of calculation
|
||
|
|
ref: int
|
||
|
|
Reference: Restricted -> 1
|
||
|
|
Unrestricted -> 2
|
||
|
|
pbs: str
|
||
|
|
Basis set of potential used
|
||
|
|
wfn : psi4.core.{RHF, UHF, RKS, UKS, Wavefunction, CCWavefuncion...}
|
||
|
|
Psi4 wavefunction object
|
||
|
|
"""
|
||
|
|
self.mol = molecule
|
||
|
|
|
||
|
|
#Assert units are in bohr
|
||
|
|
# units = self.mol.to_schema(dtype='psi4')['units']
|
||
|
|
# if units != "Bohr":
|
||
|
|
# raise ValueError("Units need to be set in Bohr")
|
||
|
|
self.basis_str = basis
|
||
|
|
self.ref = ref
|
||
|
|
self.pbs = pbs
|
||
|
|
self.pbs_str = basis if pbs == 'same' else pbs
|
||
|
|
|
||
|
|
self.nalpha = wfn.nalpha()
|
||
|
|
self.nbeta = wfn.nbeta()
|
||
|
|
|
||
|
|
self.wfn = wfn
|
||
|
|
|
||
|
|
def initialize(self):
|
||
|
|
"""
|
||
|
|
Initializes basic objects required for the Psi4Engine
|
||
|
|
"""
|
||
|
|
self.basis = psi4.core.BasisSet.build( self.mol, key='BASIS', target=self.basis_str)
|
||
|
|
self.pbs = psi4.core.BasisSet.build( self.mol, key='BASIS', target=self.pbs_str)
|
||
|
|
|
||
|
|
self.nbf = self.basis.nbf()
|
||
|
|
self.npbs = self.pbs.nbf()
|
||
|
|
|
||
|
|
self.mints = psi4.core.MintsHelper( self.basis )
|
||
|
|
self.jk = self.generate_jk()
|
||
|
|
|
||
|
|
self.grid = Psi4Grider(self.mol, self.basis, self.ref)
|
||
|
|
|
||
|
|
def get_T(self):
|
||
|
|
"""Kinetic Potential in ao basis"""
|
||
|
|
return np.array( self.mints.ao_kinetic() )
|
||
|
|
|
||
|
|
def get_Tpbas(self):
|
||
|
|
"""Kinetic Potential in pbs"""
|
||
|
|
return np.array( self.mints.ao_kinetic(self.pbs, self.pbs) )
|
||
|
|
|
||
|
|
def get_V(self):
|
||
|
|
"""External potential in ao basis"""
|
||
|
|
return np.array( self.mints.ao_potential() )
|
||
|
|
|
||
|
|
def get_A(self):
|
||
|
|
"""Inverse squared root of S matrix"""
|
||
|
|
A = self.mints.ao_overlap()
|
||
|
|
A.power( -0.5, 1e-16 )
|
||
|
|
return np.array( A )
|
||
|
|
|
||
|
|
def get_S(self):
|
||
|
|
"""Overlap matrix in AO basis"""
|
||
|
|
return np.array( self.mints.ao_overlap() )
|
||
|
|
|
||
|
|
def get_S3(self):
|
||
|
|
"""3 Orbitals Overlap matrix in AO basis"""
|
||
|
|
return np.array( self.mints.ao_3coverlap(self.basis,self.basis,self.pbs) )
|
||
|
|
|
||
|
|
def get_S4(self):
|
||
|
|
"""
|
||
|
|
Calculates four overlap integral with Density Fitting method.
|
||
|
|
S4_{ijkl} = \int dr \phi_i(r)*\phi_j(r)*\phi_k(r)*\phi_l(r)
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
wfn: psi4.core.Wavefunction
|
||
|
|
Wavefunction object of moleculep
|
||
|
|
Return
|
||
|
|
------
|
||
|
|
S4
|
||
|
|
"""
|
||
|
|
|
||
|
|
print(f"4-AO-Overlap tensor will take about {self.nbf **4 / 8 * 1e-9:f} GB.")
|
||
|
|
|
||
|
|
aux_basis = psi4.core.BasisSet.build(self.mol, "DF_BASIS_SCF", "", "JKFIT", self.basis_str)
|
||
|
|
S_Pmn = np.squeeze(self.mints.ao_3coverlap(aux_basis, self.basis, self.basis ))
|
||
|
|
S_PQ = np.array(self.mints.ao_overlap(aux_basis, aux_basis))
|
||
|
|
S_PQinv = np.linalg.pinv(S_PQ, rcond=1e-9)
|
||
|
|
S4 = contract('Pmn,PQ,Qrs->mnrs', S_Pmn, S_PQinv, S_Pmn)
|
||
|
|
return S4
|
||
|
|
|
||
|
|
def generate_jk(self, gen_K=False):
|
||
|
|
"""
|
||
|
|
Creates jk object for generation of Coulomb and Exchange matrices
|
||
|
|
1.0e9 B -> 1.0 GB
|
||
|
|
"""
|
||
|
|
jk = psi4.core.JK.build(self.basis)
|
||
|
|
memory = int(jk.memory_estimate() * 1.1)
|
||
|
|
jk.set_memory(int(memory))
|
||
|
|
# added by Ehsan
|
||
|
|
# .set_do_K(gen_K = False) determines if exchane matrices should be calculated or not
|
||
|
|
jk.set_do_K(gen_K)
|
||
|
|
jk.initialize()
|
||
|
|
#print("jk: ", jk)
|
||
|
|
return jk
|
||
|
|
|
||
|
|
def compute_hartree(self, Cocc_a, Cocc_b):
|
||
|
|
"""
|
||
|
|
Generates Coulomb and Exchange matrices from occupied orbitals
|
||
|
|
"""
|
||
|
|
Cocc_a = psi4.core.Matrix.from_array(Cocc_a)
|
||
|
|
Cocc_b = psi4.core.Matrix.from_array(Cocc_b)
|
||
|
|
self.jk.C_left_add(Cocc_a)
|
||
|
|
self.jk.C_left_add(Cocc_b)
|
||
|
|
self.jk.compute()
|
||
|
|
self.jk.C_clear()
|
||
|
|
J = (np.array(self.jk.J()[0]), np.array(self.jk.J()[1]))
|
||
|
|
return J
|
||
|
|
|
||
|
|
def hartree_NO(self, Dta):
|
||
|
|
"""
|
||
|
|
Computes Hartree potential in AO basis from Natural Orbitals
|
||
|
|
"""
|
||
|
|
|
||
|
|
if self.wfn is None:
|
||
|
|
raise ValueError('Please provide a wfn object to the Inverter, i.e., Inverter.eng = wfn')
|
||
|
|
|
||
|
|
if type(self.wfn) == psi4.core.CCWavefunction:
|
||
|
|
C_NO = psi4.core.Matrix(self.nbf, self.nbf)
|
||
|
|
eigs_NO = psi4.core.Vector(self.nbf)
|
||
|
|
self.wfn.Da().diagonalize( C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending )
|
||
|
|
occ = np.sqrt( np.array(eigs_NO) )
|
||
|
|
new_CA = occ * np.array(C_NO)
|
||
|
|
assert np.allclose( new_CA @ new_CA.T, Dta )
|
||
|
|
if self.ref == 1:
|
||
|
|
new_CB = np.copy( new_CA )
|
||
|
|
else:
|
||
|
|
self.wfn.Db().diagonalize( C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending )
|
||
|
|
occ_b = np.sqrt( np.array( eigs_NO ) )
|
||
|
|
new_CB = occ_b * np.array( C_NO )
|
||
|
|
J0 = self.compute_hartree(new_CA, new_CB)
|
||
|
|
return J0
|
||
|
|
|
||
|
|
def run_single_point(self, mol, basis, method):
|
||
|
|
"""
|
||
|
|
Run a standard energy calculation
|
||
|
|
"""
|
||
|
|
|
||
|
|
wfn_temp = psi4.energy(init+"/" + self.basis_str,
|
||
|
|
molecule=self.mol,
|
||
|
|
return_wfn=True)[1]
|
||
|
|
|
||
|
|
if self.ref == 1:
|
||
|
|
D = np.array(wfn_temp.Da()) + np.array(wfn_temp.Db())
|
||
|
|
C = np.array(wfn_temp.Ca())
|
||
|
|
e = np.array(wfn_temp.epsilon_a())
|
||
|
|
|
||
|
|
else:
|
||
|
|
D = np.stack( (np.array(wfn_temp.Da()), np.array(wfn_temp.Db())) )
|
||
|
|
C = np.stack( (np.array(wfn_temp.Ca()), np.array(wfn_temp.Cb())) )
|
||
|
|
e = np.stack( (np.array(wfn_temp.epsilon_a()), np.array(wfn_temp.epsilon_b())) )
|
||
|
|
|
||
|
|
return D, C, e
|
||
|
|
|