"""
inverter.py
Density-to-potential inversion module
Handles the primary functions
"""
import numpy as np
from dataclasses import dataclass
from opt_einsum import contract
try:
import psi4
psi4.core.be_quiet()
psi4.core.clean()
except:
pass
from .methods.wuyang import WuYang
from .methods.zmp import ZMP
from .methods.mrks import MRKS
from .methods.oucarter import OC
from .methods.pdeco import PDECO
from .methods.direct import Direct
from .grid.grider import Grider
[docs]@dataclass
class data_bucket:
"""
Data class for storing grid attributes.
"""
pass
[docs]class Inverter(Direct, WuYang, ZMP, MRKS, OC, PDECO, Grider):
"""
Attributes:
----------
wfn : psi4.core.{RHF, UHF, RKS, UKS, Wavefunction, CCWavefuncion...}
Psi4 wavefunction object
mol : psi4.core.Molecule
Psi4 molecule object
basis : psi4.core.BasisSet
Psi4 basis set object
basis_str : str
Basis set
nbf : int
Number of basis functions for main calculation
nalpha : int
Number of alpha electrons
nbeta : int
Number of beta electrons
ref : {1,2}
Reference calculation
1 -> Restricted
2 -> Unrestricted
Dt : List
List of psi4.core.Matrix for target density matrices (on AO).
ct : List
List of psi4.core.Matrix for input occupied orbitals. This might not be correct for post-HartreeFock methods.
pbs_str: string
name of Potential basis set
pbs : psi4.core.BasisSet
Potential basis set.
npbs : int
the length of pbs
v_pbs : np.ndarray shape (npbs, ) for ref==1 and (2*npbs, ) for ref==2.
potential vector on the Potential Baiss Set.
If the potential is not represented on the basis set, this should
remain 0. It will be initialized to a 0 array. One can set this
value for initial guesses before Wu-Yang method (WY) or PDE-Constrained
Optimization method (PDE-CO). For example, if PDE-CO is ran after
a WY calculation, the initial for PDE-CO will be the result of WY
if v_pbs is not zeroed.
S2 : np.ndarray
The ao overlap matrix (i.e. S matrix)
S3 : np.ndarray
The three ao overlap matrix (ao, ao, pbs)
S4 : np.ndarray
The four ao overlap matrix, the size should be (ao, ao, ao, ao)
jk : psi4.core.JK
Psi4 jk object. Built if wfn has no jk, otherwise use wfn.jk
T : np.ndarray
kinetic matrix on ao
V : np.ndarray
external potential matrix on ao
T_pbs: np.ndarray
kinetic matrix on pbs. Useful for regularization.
guide_potential_components: list of string
guide potential components name
va, vb: np.ndarray of shape (nbasis, nbasis)
guide potential Fock matrix.
Methods:
--------
"""
def __init__(self, wfn, pbs="same", debug=False):
"""
Handles Inversion
Parameters
----------
wfn : psi4.core.{RHF, UHF, RKS, UKS, Wavefunction, CCWavefuncion...}
Psi4 wavefunction object
pbs: str. default="same". If same, then the potential basis set (pbs)
is the same as orbital basis set (i.e. ao). Notice that
pbs is not needed for some methods
"""
self.wfn = wfn
self.pbs_str = pbs
self.mol = wfn.molecule()
self.basis = wfn.basisset()
self.basis_str = wfn.basisset().name()
self.nbf = wfn.basisset().nbf()
self.nalpha = wfn.nalpha()
self.nbeta = wfn.nbeta()
self.ref = 1 if psi4.core.get_global_option("REFERENCE") == "RHF" or \
psi4.core.get_global_option("REFERENCE") == "RKS" else 2
self.jk = wfn.jk() if hasattr(wfn, "jk") == True else self.generate_jk()
self.Dt = (np.array(wfn.Da_subset("AO")), np.array(wfn.Db_subset("AO")))
self.ct = (np.array(wfn.Ca_subset("AO", "OCC")), np.array(wfn.Cb_subset("AO", "OCC")))
self.pbs = self.basis if pbs == "same" \
else psi4.core.BasisSet.build( self.mol, key='BASIS', target=self.pbs_str)
self.npbs = self.pbs.nbf()
self.v_pbs = np.zeros( (self.npbs) ) if self.ref == 1 \
else np.zeros( 2 * self.npbs )
self.generate_mints_matrices()
self.grid = data_bucket
self.cubic_grid = data_bucket
self.J0 = None
self.S4 = None # Entry to save the 4 overlap matrix.
#-------------> Basics:
[docs] def generate_mints_matrices(self):
"""
Generates matrices that are methods of a mints object
"""
mints = psi4.core.MintsHelper( self.basis )
#Overlap Matrices
self.S2 = np.array(mints.ao_overlap())
A = mints.ao_overlap()
A.power( -0.5, 1e-16 )
self.A = np.array(A)
self.S3 = np.squeeze(mints.ao_3coverlap(self.basis,self.basis,self.pbs))
#Core Matrices
self.T = np.array(mints.ao_kinetic()).copy()
self.V = np.array(mints.ao_potential()).copy()
self.T_pbs = np.array(mints.ao_kinetic(self.pbs, self.pbs)).copy()
[docs] 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))
jk.set_do_K(gen_K)
jk.initialize()
return jk
[docs] def diagonalize(self, matrix, ndocc):
"""
Diagonalizes Fock Matrix
Parameters
----------
marrix: np.ndarray
Matrix to be diagonalized
ndocc: int
Number of occupied orbitals
Returns
-------
C: np.ndarray
Orbital Matrix
Cocc: np.ndarray
Occupied Orbital Matrix
D: np.ndarray
Density Matrix
eigves: np.ndarray
Eigenvalues
"""
A = self.A
Fp = A.dot(matrix).dot(A)
eigvecs, Cp = np.linalg.eigh(Fp)
C = A.dot(Cp)
Cocc = C[:, :ndocc]
D = contract('pi,qi->pq', Cocc, Cocc)
return C, Cocc, D, eigvecs
#-------------> Inversion:
[docs] def invert(self, method,
guide_potential_components = ["fermi_amaldi"],
opt_max_iter = 50,
**keywords):
"""
Handler to all available inversion methods
Parameters
----------
method: str
Method used to invert density.
Can be chosen from {wuyang, zmp, mrks, oc}.
See documentation below for each method.
guide_potential_components: list, opt
Components added as to guide inversion.
Can be chosen from {"fermi_amandi", "svwn"}
Default: ["fermi_amaldi"]
opt_max_iter: int, opt
Maximum number of iterations inside the chosen inversion.
Default: 50
direct
------
Direct inversion of a set of Kohn-Sham equations.
$$v_{xc}(r) = \frac{1}{n(r)} \sum_i^N [\phi_i^{*} (r) \nabla^2 \phi_i(r) + \varepsilon_i | \phi_i(r)|^2] $$
Parameters:
-----------
grid: np.ndarray, opt
Grid where result will be expressed in.
If not provided, dft grid will be used instead.
wuyang
------
the Wu-Yang method:
The Journal of chemical physics 118.6 (2003): 2498-2509.
Parameters:
----------
opt_max_iter: int
maximum iteration
opt_method: string, opt
Method for scipy optimizer
Currently only used by wuyang and pdeco method.
Defaul: 'trust-krylov'
https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
reg : float, opt
Regularization constant for Wuyant Inversion.
Default: None -> No regularization is added.
Becomes attribute of inverter -> inverter.lambda_reg
tol: float
tol for scipy.optimize.minimize
gtol: float
gtol for scipy.optimize.minimize: the gradient norm for
convergence
opt: dict
options for scipy.optimize.minimize
Notice that opt has lower priorities than opt_max_iter and gtol.
return:
the result are stored in self.v_pbs
zmp
---
The Zhao-Morrison-Parr Method:
Phys. Rev. A 50, 2138
Parameters:
----------
lambda_list: list
List of Lamda parameters used as a coefficient for Hartree
difference in SCF cycle.
zmp_mixing: float, optional
mixing \in [0,1]. How much of the new potential is added in.
For example, zmp_mixing = 0 means the traditional ZMP, i.e. all the potentials from previous
smaller lambda are ignored.
Zmp_mixing = 1 means that all the potentials of previous lambdas are accumulated, the larger lambda
potential are meant to fix the wrong/inaccurate region of the potential of the sum of the previous
potentials instead of providing an entire new potentials.
default: 1
opt_max_iter: float
Maximum number of iterations for scf cycle
opt_tol: float
Convergence criteria set for Density Difference and DIIS error.
return:
The result will be stored in self.proto_density_a and self.proto_density_b
For zmp_mixing==1, restricted (ref==1):
self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0])
self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[1]);
unrestricted (ref==1):
self.proto_density_a = \sum_i lambda_i * (Da_i - Dt[0]) - 1/N * (Dt[0] + Dt[1])
self.proto_density_b = \sum_i lambda_i * (Db_i - Dt[1]) - 1/N * (Dt[0] + Dt[1]);
For restricted (ref==1):
vxc = \int dr' \frac{self.proto_density_a + self.proto_density_b}{|r-r'|}
= 2 * \int dr' \frac{self.proto_density_a}{|r-r'|};
for unrestricted (ref==2):
vxc_up = \int dr' \frac{self.proto_density_a}{|r-r'|}
vxc_down = \int dr' \frac{self.proto_density_b}{|r-r'|}.
To get potential on grid, one needs to do
vxc = self.on_grid_esp(Da=self.proto_density_a, Db=self.proto_density_b, grid=grid) for restricted;
vxc_up = self.on_grid_esp(Da=self.proto_density_a, Db=np.zeros_like(self.proto_density_a),
grid=grid) for unrestricted;
mRKS
----
the modified Ryabinkin-Kohut-Staroverov method:
Phys. Rev. Lett. 115, 083001
J. Chem. Phys. 146, 084103p
Parameters:
-----------
maxiter: int
same as opt_max_iter
vxc_grid: np.ndarray of shape (3, num_grid_points), opt
When this is given, the final result will be represented
v_tol: float, opt
convergence criteria for vxc Fock matrices.
default: 1e-4
D_tol: float, opt
convergence criteria for density matrices.
default: 1e-7
eig_tol: float, opt
convergence criteria for occupied eigenvalue spectrum.
default: 1e-4
frac_old: float, opt
Linear mixing parameter for current vxc and old vxc.
If 0, no old vxc is mixed in.
Should be in [0,1)
default: 0.5.
init: string or psi4.core.Wavefunction, opt
Initial guess method.
default: "SCAN"
1) If None, input wfn info will be used as initial guess.
2) If "continue" is given, then it will not initialize
but use the densities and orbitals stored. Meaningly,
one can run a quick WY calculation as the initial
guess. This can also be used to user speficified
initial guess by setting Da, Coca, eigvec_a.
3) If it's not continue, it would be expecting a
method name string that works for psi4. A separate psi4 calculation
would be performed.
sing: tuple of float of length 4, opt.
Singularity parameter for _vxc_hole_quadrature()
default: (1e-5, 1e-4, 1e-5, 1e-4)
[0]: atol, [1]: atol1 for dft_spherical grid calculation.
[2]: atol, [3]: atol1 for vxc_grid calculation.
return:
The result will be stored in self.grid.vxc
oc
--
Ou-Carter method
J. Chem. Theory Comput. 2018, 14, 5680−5689
Parameters:
-----------
maxiter: int
same as opt_max_iter
vxc_grid: np.ndarray of shape (3, num_grid_points)
The final result will be represented on this grid
default: 1e-4
D_tol: float, opt
convergence criteria for density matrices.
default: 1e-7
eig_tol: float, opt
convergence criteria for occupied eigenvalue spectrum.
default: 1e-4
frac_old: float, opt
Linear mixing parameter for current vxc and old vxc.
If 0, no old vxc is mixed in.
Should be in [0,1)
default: 0.5.
init: string, opt
Initial guess method.
default: "SCAN"
1) If None, input wfn info will be used as initial guess.
2) If "continue" is given, then it will not initialize
but use the densities and orbitals stored. Meaningly,
one can run a quick WY calculation as the initial
guess. This can also be used to user speficified
initial guess by setting Da, Coca, eigvec_a.
3) If it's not continue, it would be expecting a
method name string that works for psi4. A separate psi4 calculation
would be performed.
wuyang
pdeco
-----
the PDE-Constrained Optimization method:
Int J Quantum Chem. 2018;118:e25425;
Nat Commun 10, 4497 (2019).
Parameters:
----------
opt_max_iter: int
maximum iteration
opt_method: string, opt
Method for scipy optimizer
Currently only used by wuyang and pdeco method.
Defaul: 'L-BFGS-B'
Options: ['L-BFGS-B', 'BFGS']
https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
reg : float, opt
Regularization constant for Wuyant Inversion.
Default: None -> No regularization is added.
Becomes attribute of inverter -> inverter.lambda_reg
gtol: float
gtol for scipy.optimize.minimize: the gradient norm for
convergence
opt: dict
options for scipy.optimize.minimize
Notice that opt has lower priorities than opt_max_iter and gtol.
return:
the result are stored in self.v_pbs
"""
#Generate Guide Potential
if method.lower()=='mrks':
if guide_potential_components[0] != 'hartree' or len(guide_potential_components) != 1:
print("The guide potential is changed to v_hartree.")
self.generate_components(["hartree"])
elif method.lower() != 'direct':
self.generate_components(guide_potential_components)
#Invert
if method.lower() == "direct":
return self.direct_inversion(**keywords)
elif method.lower() == "wuyang":
self.wuyang(opt_max_iter, **keywords)
elif method.lower() == "zmp":
self.zmp(opt_max_iter, **keywords)
elif method.lower() == "mrks":
return self.mRKS(opt_max_iter, **keywords)
elif method.lower() == 'oc':
return self.oucarter(opt_max_iter, **keywords)
elif method.lower() == 'pdeco':
return self.pdeco(opt_max_iter, **keywords)
else:
raise ValueError(f"Inversion method not available. Methods available: {['wuyang', 'zmp', 'mrks', 'oc', 'pdeco']}")
[docs] def generate_components(self, guide_potential_components):
"""
Generates exact potential components to be added to
the Hamiltonian to aide in the inversion procedure.
Parameters:
-----------
guide_potential_components: list
Components added as to guide inversion.
Can be chosen from {"fermi_amandi", "svwn"}
"""
self.va = np.zeros_like(self.T)
self.vb = np.zeros_like(self.T)
N = self.nalpha + self.nbeta
if self.J0 is None:
if type(self.wfn) == psi4.core.CCWavefunction:
nbf = self.nbf
C_NO = psi4.core.Matrix(nbf, nbf)
eigs_NO = psi4.core.Vector(nbf)
self.wfn.Da().diagonalize(C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending)
occu = np.sqrt(eigs_NO.np)
New_Orb_a = occu * C_NO.np
assert np.allclose(New_Orb_a @ New_Orb_a.T, self.Dt[0])
if self.ref == 1:
New_Orb_b = New_Orb_a
else:
self.wfn.Db().diagonalize(C_NO, eigs_NO, psi4.core.DiagonalizeOrder.Descending)
occu = np.sqrt(eigs_NO.np)
New_Orb_b = occu * C_NO.np
self.J0 = self.form_jk(New_Orb_a, New_Orb_b)[0]
else:
self.J0, _ = self.form_jk(self.ct[0], self.ct[1])
if "fermi_amaldi" in guide_potential_components:
v_fa = (1-1/N) * (self.J0[0] + self.J0[1])
self.va += v_fa
self.vb += v_fa
elif "hartree" in guide_potential_components:
v_hartree = (self.J0[0] + self.J0[1])
self.va += v_hartree
self.vb += v_hartree
else:
raise ValueError("Hartee nor FA was included."
"Convergence will likely not be achieved")
if "svwn" in guide_potential_components:
if "svwn" in guide_potential_components:
_, wfn_0 = psi4.energy( "svwn"+"/"+self.basis_str, molecule=self.mol , return_wfn = True)
if self.ref == 1:
ntarget = psi4.core.Matrix.from_array( [ self.Dt[0] + self.Dt[1] ] )
wfn_0.V_potential().set_D( [ntarget] )
vxc_a = psi4.core.Matrix( self.nbf, self.nbf )
wfn_0.V_potential().compute_V([vxc_a])
self.va += vxc_a.np
self.vb += vxc_a.np
elif self.ref == 2:
na_target = psi4.core.Matrix.from_array( self.Dt[0] )
nb_target = psi4.core.Matrix.from_array( self.Dt[1] )
wfn_0.V_potential().set_D( [na_target, nb_target] )
vxc_a = psi4.core.Matrix( self.nbf, self.nbf )
vxc_b = psi4.core.Matrix( self.nbf, self.nbf )
wfn_0.V_potential().compute_V([vxc_a, vxc_b])
self.va += vxc_a.np
self.vb += vxc_b.np
self.guide_potential_components = guide_potential_components
[docs] def finalize_energy(self):
"""
Calculates energy contributions
"""
target_one = self.wfn.to_file()['floatvar']['ONE-ELECTRON ENERGY']
target_two = self.wfn.to_file()['floatvar']['TWO-ELECTRON ENERGY']
energy_kinetic = contract('ij,ij', self.T, (self.Da + self.Db))
energy_external = contract('ij,ij', self.V, (self.Da + self.Db))
energy_hartree_a = 0.5 * contract('ij,ji', self.J0[0] + self.J0[1], self.Da)
energy_hartree_b = 0.5 * contract('ij,ji', self.J0[0] + self.J0[1], self.Db)
print("WARNING: XC Energy is not yet properly calculated")
energies = {"One-Electron Energy" : energy_kinetic + energy_external,
"Two-Electron Energy" : energy_hartree_a + energy_hartree_b,
# "XC" : energy_ks,
# "Total Energy" : energy_kinetic + energy_external + \
# energy_hartree_a + energy_hartree_b + \
# energy_ks
}
target_energies = {"One-Electron Energy" : target_one,
"Two-Electron Energy" : target_two,
}
self.energies = energies
self.target_energies = target_energies
# print(f"Final Energies: {self.energies}")
# print(f"Target Energies: {self.target_energies}")
