Add nuclear gradient functions

gbasis/integrals/libcint.py

@@ -939,7 +939,7 @@ def int2e(notation="physicist", origin=None, inv_origin=None):
 
             # Transpose integrals in `out` array to proper notation
             if physicist:
-                out = out.transpose(0, 2, 1, 3)
+                out = out.transpose(0, 2, 1, 3, *range(4, out.ndim))
 
             # Return normalized integrals
             if self.coord_type == "cartesian":
@@ -1124,6 +1124,76 @@ def moment(self, orders, origin=None):
         # Return integrals in `out` array
         return out
 
+    def grad_overlap(self):
+        r"""
+        Compute the nuclear gradient of the overlap integrals.
+
+        Returns
+        -------
+        out : np.ndarray(Nbasis, Nbasis, N, 3, dtype=float)
+            Gradient array.
+
+        """
+        # Allocate output array
+        d_ovlp = np.zeros((self.nbfn, self.nbfn, self.natm, 3))
+        # Compute nuclear gradient
+        d_s = self._d_ovlp()
+        for iatm, icoords in enumerate(self.atcoords):
+            start, end = self._atm_offs[iatm : iatm + 2]
+            d_ovlp[start:end, :, iatm, :] -= d_s[start:end, :, :]
+            d_ovlp[:, :, iatm, :] += d_ovlp[:, :, iatm, :].transpose(1, 0, 2)
+        # Return output array
+        return d_ovlp
+
+    def grad_core_hamiltonian(self):
+        r"""
+        Compute the nuclear gradient of the core Hamiltonian (T + V) integrals.
+
+        Returns
+        -------
+        out : np.ndarray(Nbasis, Nbasis, N, 3, dtype=float)
+            Gradient array.
+
+        """
+        # Allocate output array
+        d_hcore = np.zeros((self.nbfn, self.nbfn, self.natm, 3))
+        # Compute nuclear gradient
+        d_h = self._d_kin()
+        d_h += self._d_nuc()
+        for iatm, (iz, icoords) in enumerate(zip(self.atnums, self.atcoords)):
+            start, end = self._atm_offs[iatm : iatm + 2]
+            d_rinv = -iz * self._d_rinv(inv_origin=icoords)
+            d_rinv[start:end, :, :] -= d_h[start:end, :, :]
+            d_rinv += d_rinv.transpose(1, 0, 2)
+            d_hcore[:, :, iatm, :] = d_rinv
+        # Return output array
+        return d_hcore
+
+    def grad_electron_repulsion(self):
+        r"""
+        Compute the nuclear gradient of the electron repulsion integrals.
+
+        Returns
+        -------
+        out : np.ndarray(Nbasis, Nbasis, Nbasis, Nbasis, N, 3, dtype=float)
+            Gradient array.
+
+        """
+        # Allocate output array
+        d_eri = np.zeros((self.nbfn, self.nbfn, self.nbfn, self.nbfn, self.natm, 3))
+        # Compute nuclear gradient
+        d_i = self._d_eri()
+        for iatm, icoords in enumerate(self.atcoords):
+            start, end = self._atm_offs[iatm : iatm + 2]
+            d_eri[start:end, :, :, :, iatm, :] -= d_i[start:end, :, :, :, :]
+            d_eri[:, :, :, :, iatm, :] += d_eri[:, :, :, :, iatm, :].transpose(2, 3, 0, 1, 4)
+            d_eri[:, :, :, :, iatm, :] += d_eri[:, :, :, :, iatm, :].transpose(1, 0, 2, 3, 4)
+            d_eri[:, :, :, :, iatm, :] += d_eri[:, :, :, :, iatm, :].transpose(0, 1, 3, 2, 4)
+            d_eri[:, :, :, :, iatm, :] += d_eri[:, :, :, :, iatm, :].transpose(1, 0, 3, 2, 4)
+        d_eri *= 0.25
+        # Return output array
+        return d_eri
+
 
 def normalized_coeffs(shell):
     r"""

pyproject.toml

@@ -61,7 +61,7 @@ doc = [
     "nbsphinx-link"
 ]
 "dev" = ["tox"]
-"test" = ["tox", "pytest", "pytest-cov"]
+"test" = ["tox", "pytest", "pytest-cov", "derivcheck"]
 "iodata" = ["iodata>0.1.7"]
 "pyscf" = ["pyscf>=1.6.1"]
 

tests/test_libcint.py

@@ -20,6 +20,8 @@
 
 from utils import find_datafile
 
+from derivcheck import assert_deriv
+
 
 TEST_BASIS_SETS = [
     pytest.param("data_sto6g.nwchem", id="STO-6G"),
@@ -58,6 +60,13 @@
 ]
 
 
+TEST_NUCLEAR_GRADIENTS = [
+    pytest.param("overlap", id="Overlap"),
+    pytest.param("core_hamiltonian", id="CoreHamiltonian"),
+    pytest.param("electron_repulsion", id="ElectronRepulsion"),
+]
+
+
 @pytest.mark.parametrize("integral", TEST_INTEGRALS)
 @pytest.mark.parametrize("coord_type", TEST_COORD_TYPES)
 @pytest.mark.parametrize("atsyms, atcoords", TEST_SYSTEMS)
@@ -156,3 +165,41 @@ def test_integral(basis, atsyms, atcoords, coord_type, integral):
         raise ValueError("Invalid integral name '{integral}' passed")
 
     npt.assert_allclose(lc_int, py_int, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("gradient", TEST_NUCLEAR_GRADIENTS)
+@pytest.mark.parametrize("coord_type", TEST_COORD_TYPES)
+@pytest.mark.parametrize("atsyms, atcoords", TEST_SYSTEMS)
+@pytest.mark.parametrize("basis", TEST_BASIS_SETS)
+def test_nuclear_gradient(basis, atsyms, atcoords, coord_type, gradient):
+    r"""
+    Test gbasis.integrals.libcint.CBasis nuclear gradients.
+
+    """
+    basis_dict = parse_nwchem(find_datafile(basis))
+
+    def f(x):
+        atcoords = x.reshape(len(atsyms), 3) / 0.5291772083
+        atnums = np.asarray([ELEMENTS.index(i) for i in atsyms], dtype=float)
+        py_basis = make_contractions(basis_dict, atsyms, atcoords, coord_types=coord_type)
+        lc_basis = CBasis(py_basis, atsyms, atcoords, coord_type=coord_type)
+        if gradient == "overlap":
+            return lc_basis.overlap().reshape(-1)
+        elif gradient == "core_hamiltonian":
+            return lc_basis.kinetic_energy().reshape(-1) + lc_basis.nuclear_attraction().reshape(-1)
+        elif gradient == "electron_repulsion":
+            return lc_basis.electron_repulsion().reshape(-1)
+
+    def g(x):
+        atcoords = x.reshape(len(atsyms), 3) / 0.5291772083
+        atnums = np.asarray([ELEMENTS.index(i) for i in atsyms], dtype=float)
+        py_basis = make_contractions(basis_dict, atsyms, atcoords, coord_types=coord_type)
+        lc_basis = CBasis(py_basis, atsyms, atcoords, coord_type=coord_type)
+        if gradient == "overlap":
+            return lc_basis.grad_overlap().reshape(lc_basis.nbfn**2, -1)
+        elif gradient == "core_hamiltonian":
+            return lc_basis.grad_core_hamiltonian().reshape(lc_basis.nbfn**2, -1)
+        elif gradient == "electron_repulsion":
+            return lc_basis.grad_electron_repulsion().reshape(lc_basis.nbfn**4, -1)
+
+    assert_deriv(f, g, atcoords.reshape(-1))