Merge 365cd5a into e4ca2ed

adcc/Excitation.py

@@ -66,6 +66,8 @@ def __init__(self, parent_state, index, method):
         self.__parent_state = parent_state
         self.index = index
         self.method = method
+        self.ground_state = parent_state.ground_state
+        self.reference_state = parent_state.reference_state
         for key in self.parent_state.excitation_property_keys:
             fget = getattr(type(self.parent_state), key).fget
             # Extract the kwargs passed to mark_excitation_property

adcc/ExcitedStates.py

@@ -218,6 +218,12 @@ def __add__(self, other):
         ret += other
         return ret
 
+    @property
+    @mark_excitation_property()
+    def total_energy(self):
+        # TODO: excitation_energy_uncorrected for PE-ADC with postSCF
+        return self.excitation_energy + self.ground_state.energy(self.method.level)
+
     @cached_property
     @mark_excitation_property(transform_to_ao=True)
     @timed_member_call(timer="_property_timer")

adcc/ReferenceState.py

@@ -156,7 +156,9 @@ def __init__(self, hfdata, core_orbitals=None, frozen_core=None,
         )
 
         self.environment = None  # no environment attached by default
-        for name in ["excitation_energy_corrections", "environment"]:
+        for name in ["excitation_energy_corrections",
+                     "environment",
+                     "gradient_provider"]:
             if hasattr(hfdata, name):
                 setattr(self, name, getattr(hfdata, name))
 

adcc/__init__.py

@@ -41,6 +41,7 @@
 from .AmplitudeVector import AmplitudeVector
 from .OneParticleOperator import OneParticleOperator
 from .opt_einsum_integration import register_with_opt_einsum
+from .gradients import nuclear_gradient
 
 # This has to be the last set of import
 from .guess import (guess_symmetries, guess_zero, guesses_any, guesses_singlet,
@@ -60,6 +61,7 @@
            "guess_symmetries", "guesses_spin_flip", "guess_zero", "LazyMp",
            "adc0", "cis", "adc1", "adc2", "adc2x", "adc3",
            "cvs_adc0", "cvs_adc1", "cvs_adc2", "cvs_adc2x", "cvs_adc3",
+           "nuclear_gradient",
            "banner"]
 
 __version__ = "0.15.9"

adcc/backends/psi4.py

@@ -32,6 +32,70 @@
 from ..ExcitedStates import EnergyCorrection
 
 
+class Psi4GradientProvider:
+    def __init__(self, wfn):
+        self.wfn = wfn
+        self.mol = self.wfn.molecule()
+        self.backend = "psi4"
+        self.mints = psi4.core.MintsHelper(self.wfn)
+
+    def correlated_gradient(self, g1_ao, w_ao, g2_ao_1, g2_ao_2):
+        """
+        g1_ao: relaxed one-particle density matrix
+        w_ao:  energy-weighted density matrix
+        g2_ao_1, g2_ao_2: relaxed two-particle density matrices
+        """
+        natoms = self.mol.natom()
+        Gradient = {}
+        Gradient["N"] = np.zeros((natoms, 3))
+        Gradient["S"] = np.zeros((natoms, 3))
+        Gradient["T"] = np.zeros((natoms, 3))
+        Gradient["V"] = np.zeros((natoms, 3))
+        Gradient["OEI"] = np.zeros((natoms, 3))
+        Gradient["TEI"] = np.zeros((natoms, 3))
+        Gradient["Total"] = np.zeros((natoms, 3))
+
+        # 1st Derivative of Nuclear Repulsion
+        Gradient["N"] = psi4.core.Matrix.to_array(
+            self.mol.nuclear_repulsion_energy_deriv1([0, 0, 0])
+        )
+        # Build Integral Derivatives
+        cart = ['_X', '_Y', '_Z']
+        oei_dict = {"S": "OVERLAP", "T": "KINETIC", "V": "POTENTIAL"}
+
+        deriv1_mat = {}
+        deriv1_np = {}
+        # 1st Derivative of OEIs
+        for atom in range(natoms):
+            for key in oei_dict:
+                string = key + str(atom)
+                deriv1_mat[string] = self.mints.ao_oei_deriv1(oei_dict[key], atom)
+                for p in range(3):
+                    map_key = string + cart[p]
+                    deriv1_np[map_key] = np.asarray(deriv1_mat[string][p])
+                    if key == "S":
+                        Gradient["S"][atom, p] = np.sum(w_ao * deriv1_np[map_key])
+                    else:
+                        Gradient[key][atom, p] = np.sum(g1_ao * deriv1_np[map_key])
+        # Build Total OEI Gradient
+        Gradient["OEI"] = Gradient["T"] + Gradient["V"] + Gradient["S"]
+
+        # Build TE contributions
+        for atom in range(natoms):
+            deriv2 = self.mints.ao_tei_deriv1(atom)
+            for p in range(3):
+                deriv2_np = np.asarray(deriv2[p])
+                Gradient["TEI"][atom, p] += np.einsum(
+                    'pqrs,prqs->', g2_ao_1, deriv2_np, optimize=True
+                )
+                Gradient["TEI"][atom, p] -= np.einsum(
+                    'pqrs,psqr->', g2_ao_2, deriv2_np, optimize=True
+                )
+        # Build total gradient
+        Gradient["Total"] = Gradient["OEI"] + Gradient["TEI"] + Gradient["N"]
+        return Gradient
+
+
 class Psi4OperatorIntegralProvider:
     def __init__(self, wfn):
         self.wfn = wfn
@@ -99,6 +163,7 @@ def __init__(self, wfn):
                                           wfn.nalpha(), wfn.nbeta(),
                                           self.restricted)
         self.operator_integral_provider = Psi4OperatorIntegralProvider(self.wfn)
+        self.gradient_provider = Psi4GradientProvider(self.wfn)
 
     def pe_energy(self, dm, elec_only=True):
         density_psi = psi4.core.Matrix.from_array(dm.to_ndarray())

adcc/backends/pyscf.py

@@ -29,7 +29,66 @@
 from ..exceptions import InvalidReference
 from ..ExcitedStates import EnergyCorrection
 
-from pyscf import ao2mo, gto, scf, solvent
+from pyscf import ao2mo, gto, scf, grad, solvent
+
+
+class PyScfGradientProvider:
+    def __init__(self, scfres):
+        self.scfres = scfres
+        self.mol = self.scfres.mol
+        self.backend = "pyscf"
+
+    def correlated_gradient(self, g1_ao, w_ao, g2_ao_1, g2_ao_2):
+        natoms = self.mol.natm
+        Gradient = {}
+        Gradient["N"] = np.zeros((natoms, 3))
+        Gradient["S"] = np.zeros((natoms, 3))
+        Gradient["T+V"] = np.zeros((natoms, 3))
+        Gradient["OEI"] = np.zeros((natoms, 3))
+        Gradient["TEI"] = np.zeros((natoms, 3))
+        Gradient["Total"] = np.zeros((natoms, 3))
+
+        # TODO: does RHF/UHF matter here?
+        gradient = grad.RHF(self.scfres)
+        hcore_deriv = gradient.hcore_generator()
+        Sx = -1.0 * self.mol.intor('int1e_ipovlp', aosym='s1')
+        ERIx = -1.0 * self.mol.intor('int2e_ip1', aosym='s1')
+
+        ao_slices = self.mol.aoslice_by_atom()
+        for ia in range(natoms):
+            # TODO: only contract/compute with specific slices
+            # of density matrices (especially TPDM)
+            # this requires a lot of work however...
+            k0, k1 = ao_slices[ia, 2:]
+
+            # derivative of the overlap matrix
+            Sx_a = np.zeros_like(Sx)
+            Sx_a[:, k0:k1] = Sx[:, k0:k1]
+            Sx_a += Sx_a.transpose(0, 2, 1)
+            Gradient["S"][ia] += np.einsum("xpq,pq->x", Sx_a, w_ao)
+
+            # derivative of the core Hamiltonian
+            Hx_a = hcore_deriv(ia)
+            Gradient["T+V"][ia] += np.einsum("xpq,pq->x", Hx_a, g1_ao)
+
+            # derivatives of the ERIs
+            ERIx_a = np.zeros_like(ERIx)
+            ERIx_a[:, k0:k1] = ERIx[:, k0:k1]
+            ERIx_a += (
+                + ERIx_a.transpose(0, 2, 1, 4, 3)
+                + ERIx_a.transpose(0, 3, 4, 1, 2)
+                + ERIx_a.transpose(0, 4, 3, 2, 1)
+            )
+            Gradient["TEI"][ia] += np.einsum(
+                "pqrs,xprqs->x", g2_ao_1, ERIx_a, optimize=True
+            )
+            Gradient["TEI"][ia] -= np.einsum(
+                "pqrs,xpsqr->x", g2_ao_2, ERIx_a, optimize=True
+            )
+        Gradient["N"] = gradient.grad_nuc()
+        Gradient["OEI"] = Gradient["T+V"] + Gradient["S"]
+        Gradient["Total"] = Gradient["OEI"] + Gradient["TEI"] + Gradient["N"]
+        return Gradient
 
 
 class PyScfOperatorIntegralProvider:
@@ -89,7 +148,7 @@ def coefficients(self):
 
     def compute_mo_eri(self, blocks, spins):
         coeffs = tuple(self.coefficients[blocks[i] + spins[i]] for i in range(4))
-        # TODO Pyscf usse HDF5 internal to do the AO2MO here we read it all
+        # TODO Pyscf uses HDF5 internal to do the AO2MO here we read it all
         #      into memory. This wastes memory and could be avoided if temporary
         #      files were used instead. These could be deleted on the call
         #      to `flush_cache` automatically.
@@ -116,6 +175,7 @@ def __init__(self, scfres):
         self.operator_integral_provider = PyScfOperatorIntegralProvider(
             self.scfres
         )
+        self.gradient_provider = PyScfGradientProvider(self.scfres)
         if not self.restricted:
             assert self.scfres.mo_coeff[0].shape[1] == \
                 self.scfres.mo_coeff[1].shape[1]