Nuclear Gradients

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

@@ -217,6 +217,14 @@ 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
+        if self.method.level == 0:
+            return self.excitation_energy + self.reference_state.energy_scf
+        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

@@ -43,6 +43,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,
@@ -63,6 +64,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.17"

adcc/adc_pp/matrix.py

@@ -21,7 +21,7 @@
 ##
 ## ---------------------------------------------------------------------
 from math import sqrt
-from collections import namedtuple
+from dataclasses import dataclass
 
 from adcc import block as b
 from adcc.functions import direct_sum, einsum, zeros_like
@@ -42,12 +42,16 @@
 #
 # Dispatch routine
 #
-"""
-`apply` is a function mapping an AmplitudeVector to the contribution of this
-block to the result of applying the ADC matrix. `diagonal` is an `AmplitudeVector`
-containing the expression to the diagonal of the ADC matrix from this block.
-"""
-AdcBlock = namedtuple("AdcBlock", ["apply", "diagonal"])
+@dataclass
+class AdcBlock:
+    """
+    `apply` is a function mapping an AmplitudeVector to the contribution of this
+    block to the result of applying the ADC matrix. `diagonal` is an
+    `AmplitudeVector` containing the expression to the diagonal of the ADC matrix
+    from this block.
+    """
+    apply: callable
+    diagonal: AmplitudeVector
 
 
 def block(ground_state, spaces, order, variant=None, intermediates=None):

adcc/backends/psi4.py

@@ -24,6 +24,7 @@
 
 from libadcc import HartreeFockProvider
 from adcc.misc import cached_property
+from adcc.gradients import GradientComponents
 
 import psi4
 
@@ -32,6 +33,72 @@
 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
+                )
+        ret = GradientComponents(
+            natoms, Gradient["N"], Gradient["S"],
+            Gradient["T"] + Gradient["V"], Gradient["TEI"]
+        )
+        return ret
+
+
 class Psi4OperatorIntegralProvider:
     def __init__(self, wfn):
         self.wfn = wfn
@@ -116,6 +183,7 @@ def __init__(self, wfn):
                                           wfn.nalpha(), wfn.nbeta(),
                                           self.restricted)
         self.operator_integral_provider = Psi4OperatorIntegralProvider(self.wfn)
+        self.gradient_provider = Psi4GradientProvider(self.wfn)
 
         self.environment = None
         self.environment_implementation = None

adcc/backends/pyscf.py

@@ -24,15 +24,76 @@
 
 from libadcc import HartreeFockProvider
 from adcc.misc import cached_property
+from adcc.gradients import GradientComponents
 
 from .EriBuilder import EriBuilder
 from ..exceptions import InvalidReference
 from ..ExcitedStates import EnergyCorrection
 
-from pyscf import ao2mo, gto, scf
+from pyscf import ao2mo, gto, scf, grad
 from pyscf.solvent import ddcosmo
 
 
+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
+            )
+        ret = GradientComponents(
+            natoms, gradient.grad_nuc(), Gradient["S"],
+            Gradient["T+V"], Gradient["TEI"]
+        )
+        return ret
+
+
 class PyScfOperatorIntegralProvider:
     def __init__(self, scfres):
         self.scfres = scfres
@@ -111,7 +172,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.
@@ -138,6 +199,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]