DFT Hessian grid response (#533)

* DFT Hessian Fock 1st derivative grid response. Only correct result, with horrible memory usage and performance

* LDA Hessian energy 2nd derivative grid response. Only correct result, with horrible memory usage and performance

* GGA Hessian energy 2nd derivative grid response. Only correct result, with horrible memory usage and performance

* MGGA Hessian energy 2nd derivative grid response + 1st round of cleanup. Still only correct result

* VV10 Hessian energy 2nd derivative grid response. All grid responses are correct now. Need optimization

* Add tests for individual terms of hessian grid response. There's some problem with overall hessian compared to numerical hessian

* Second round of cleanup, remove natm*natm*ngrids memory requirement for rho, nabla rho and tau second derivative, remove nao*ngrids operation and some redundant operations in vv10, fix some linter errors

* Third round of cleanup, remove all einsum with more than two matrix inputs in energy second derivative terms

* Fourth round of cleanup, remove all einsum with more than two matrix inputs in Fock first derivative terms

* Split the memory estimation of Becke weight second derivative from other terms, because that is the only remaining natm*natm*ngrids term

* Bug fix to remove zero-sized array operations in grid response

* Count the maximum amount of memory per grid used in each operation. Certainly overcounted.

* Fix xctype==HF test

* Turn on large_exponent rks tests, confirm the problem is from grid response

* Small bug

* Move numerical derivative routine to test files, add log info for grid response

* Fix multigpu illegal memory bug

* Remove extra computation in d2rho/dAdB grid response

* Merge common computation in d2rho/dAdB grid response

* Apply similar optimization to d2e and dFock terms

* Bug fix on two H test when HOMO LUMO are degenerate

* Minor bug fix in previous commit

* Minor bug fix in previous commit, again

* Accelerate Becke weight second derivative kernel, by moving P_B computation out from the kernel

* Another minor improvement for Becke weight second derivative

Author: henryw7

gpu4pyscf/df/hessian/rks.py

@@ -64,22 +64,7 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
 
     max_memory = None
     t1 = log.timer_debug1('computing ej, ek', *t1)
-    veff_diag = rks_hess._get_vxc_diag(hessobj, mo_coeff, mo_occ, max_memory)
-    t1 = log.timer_debug1('computing veff_diag', *t1)
-    aoslices = mol.aoslice_by_atom()
-    vxc_dm = rks_hess._get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
-    t1 = log.timer_debug1('computing veff_deriv2', *t1)
-    for i0, ia in enumerate(atmlst):
-        shl0, shl1, p0, p1 = aoslices[ia]
-        veff_dm = vxc_dm[ia]
-        de2[i0,i0] += contract('xypq,pq->xy', veff_diag[:,:,p0:p1], dm0[p0:p1])*2
-        for j0, ja in enumerate(atmlst[:i0+1]):
-            q0, q1 = aoslices[ja][2:]
-            #:contract('xypq,pq->xy', veff[:,:,q0:q1], dm0[q0:q1])*2
-            de2[i0,j0] += 2.0*cupy.sum(veff_dm[:,:,q0:q1], axis=2)
-        for j0 in range(i0):
-            de2[j0,i0] = de2[i0,j0].T
-
+    de2 += rks_hess._get_exc_deriv2(hessobj, mo_coeff, mo_occ, dm0, max_memory)
     if mf.do_nlc():
         de2 += _get_enlc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
 

gpu4pyscf/dft/numint.py

@@ -1871,7 +1871,7 @@ def _block_loop(ni, mol, grids, nao=None, deriv=0, max_memory=2000,
         non0tab: dummy argument for compatibility with PySCF
         blksize: if not given, it will be estimated with avail GPU memory.
         buf: dummy argument for compatibility with PySCF
-        grid_range: loop [grid_start, grid_end] in grids only.
+        grid_range: loop [grid_start, grid_end] in grids only. TODO: Henry 20251006 believes these parameters are not respected.
     '''
     log = logger.new_logger(mol)
     if grids.coords is None:

gpu4pyscf/hessian/rhf.py

@@ -629,6 +629,9 @@ def solve_mo1(mf, mo_energy, mo_coeff, mo_occ, h1mo,
     e_ai = 1 / (e_a[:,None] + level_shift - e_i)
     nvir, nocc = e_ai.shape
 
+    if cupy.any(cp.isinf(e_ai)) or cupy.any(cp.isnan(e_ai)):
+        raise ValueError(f"e_ai = {e_ai} contains inf or nan, likely because HOMO-LUMO gap is zero.")
+
     mocc = mo_coeff[:,occidx]
     nao, nmo = mo_coeff.shape
     natm = mol.natm

gpu4pyscf/hessian/rks.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 
 import unittest
+import pytest
 import numpy as np
 import cupy as cp
 import pyscf
 from gpu4pyscf import scf, dft
 
 def setUpModule():
-    global mol_minimal_one_atom, mol_minimal_two_atom, auxbasis_minimal_good, auxbasis_minimal_bad #, mol_real
+    global mol_minimal_one_atom, mol_minimal_two_atom, auxbasis_minimal_good, auxbasis_minimal_bad
     mol_minimal_one_atom = pyscf.M(
         atom = """
         He 100 200 300
@@ -52,27 +53,12 @@ def setUpModule():
             2.0 1.0
     """
 
-    # mol_real = pyscf.M(
-    #     atom = """
-    #         C      0.000000    0.000000    0.000000
-    #         Br     0.000000    0.000000    1.940000
-    #         H      1.027662    0.000000   -0.363333
-    #         H     -0.513831    0.889981   -0.363333
-    #         H     -0.513831   -0.889981   -0.363333
-    #     """,
-    #     basis = "def2-svp",
-    #     verbose = 4,
-    #     output='/dev/null',
-    # )
-
 def tearDownModule():
-    global mol_minimal_one_atom, mol_minimal_two_atom #, mol_real
+    global mol_minimal_one_atom, mol_minimal_two_atom
     mol_minimal_one_atom.stdout.close()
     del mol_minimal_one_atom
     mol_minimal_two_atom.stdout.close()
     del mol_minimal_two_atom
-    # mol_real.stdout.close()
-    # del mol_real
 
 class KnownValues(unittest.TestCase):
     def test_hessian_large_exp_one_atom_rhf(self):
@@ -170,64 +156,81 @@ def test_hessian_large_exp_two_atom_rhf_density_fit_bad_auxbasis(self):
 
         assert np.max(np.abs(translation_invariance)) < 1e-4
 
-    # TODO: Fix DFT Hessian
-
-    # def test_hessian_large_exp_one_atom_rks(self):
-    #     mf = dft.RKS(mol_minimal_one_atom, xc = "LDA")
-    #     mf.grids.atom_grid = (200,194)
-    #     mf.conv_tol = 1e-10
+    def test_hessian_large_exp_one_atom_rks(self):
+        mf = dft.RKS(mol_minimal_one_atom, xc = "r2SCAN")
+        mf.grids.atom_grid = (50,194)
+        mf.conv_tol = 1e-10
 
-    #     mf.kernel()
-    #     assert mf.converged
+        mf.kernel()
+        assert mf.converged
 
-    #     hobj = mf.Hessian()
-    #     hessian = hobj.kernel()
+        hobj = mf.Hessian()
+        hobj.grid_response = True
+        hessian = hobj.kernel()
 
-    #     natm = mf.mol.natm
-    #     hessian = hessian.transpose(0,2,1,3)
-    #     hessian = hessian.reshape((natm * 3, natm * 3))
-    #     translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
+        natm = mf.mol.natm
+        hessian = hessian.transpose(0,2,1,3)
+        hessian = hessian.reshape((natm * 3, natm * 3))
+        translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
 
-    #     # Zero hessian for only one atom
-    #     assert np.max(np.abs(hessian)) < 1e-3
-    #     assert np.max(np.abs(translation_invariance)) < 1e-3
+        # Zero hessian for only one atom
+        assert np.max(np.abs(hessian)) < 1e-4
+        assert np.max(np.abs(translation_invariance)) < 1e-4
 
-    # def test_hessian_large_exp_two_atom_rks(self):
-    #     mf = dft.RKS(mol_minimal_two_atom, xc = "LDA")
-    #     mf.grids.atom_grid = (200,194)
-    #     mf.conv_tol = 1e-10
+    def test_hessian_large_exp_two_atom_rks(self):
+        mf = dft.RKS(mol_minimal_two_atom, xc = "wB97X")
+        mf.grids.atom_grid = (50,194)
+        mf.conv_tol = 1e-10
+        mf.level_shift = 0.001
 
-    #     mf.kernel()
-    #     assert mf.converged
+        mf.kernel()
+        assert mf.converged
 
-    #     hobj = mf.Hessian()
-    #     hessian = hobj.kernel()
+        hobj = mf.Hessian()
+        hobj.grid_response = True
+        hessian = hobj.kernel()
 
-    #     natm = mf.mol.natm
-    #     hessian = hessian.transpose(0,2,1,3)
-    #     hessian = hessian.reshape((natm * 3, natm * 3))
-    #     translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
+        natm = mf.mol.natm
+        hessian = hessian.transpose(0,2,1,3)
+        hessian = hessian.reshape((natm * 3, natm * 3))
+        translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
 
-    #     assert np.max(np.abs(translation_invariance)) < 1e-3
+        assert np.max(np.abs(translation_invariance)) < 1e-4
 
-    # def test_hessian_large_exp_methylbromide_rks(self):
-    #     mf = dft.RKS(mol_real, xc = "wB97X")
-    #     mf.grids.atom_grid = (99,590)
-    #     mf = mf.density_fit(auxbasis = "def2-universal-jkfit")
-    #     mf.conv_tol = 1e-10
+    @pytest.mark.skip("Too slow, functionality covered by corresponding tests above")
+    def test_hessian_large_exp_methylbromide_rks(self):
+        mol_real = pyscf.M(
+            atom = """
+                C      0.000000    0.000000    0.000000
+                Br     0.000000    0.000000    1.940000
+                H      1.027662    0.000000   -0.363333
+                H     -0.513831    0.889981   -0.363333
+                H     -0.513831   -0.889981   -0.363333
+            """,
+            basis = "def2-svp",
+            verbose = 4,
+            output='/dev/null',
+        )
+
+        mf = dft.RKS(mol_real, xc = "wB97X")
+        mf.grids.atom_grid = (99,590)
+        mf = mf.density_fit(auxbasis = "def2-universal-jkfit")
+        mf.conv_tol = 1e-10
 
-    #     mf.kernel()
-    #     assert mf.converged
+        mf.kernel()
+        assert mf.converged
 
-    #     hobj = mf.Hessian()
-    #     hessian = hobj.kernel()
+        hobj = mf.Hessian()
+        hobj.auxbasis_response = 2
+        hobj.grid_response = True
+        hessian = hobj.kernel()
 
-    #     natm = mf.mol.natm
-    #     hessian = hessian.transpose(0,2,1,3)
-    #     hessian = hessian.reshape((natm * 3, natm * 3))
-    #     translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
+        natm = mf.mol.natm
+        hessian = hessian.transpose(0,2,1,3)
+        hessian = hessian.reshape((natm * 3, natm * 3))
+        translation_invariance = np.sum(hessian, axis = 0).reshape(natm, 3)
 
-    #     assert np.max(np.abs(translation_invariance)) < 1e-10
+        assert np.max(np.abs(translation_invariance)) < 1e-7
 
 if __name__ == "__main__":
     print("Edge Case Tests for Hessian Calculation with Large Exponent in Atomic Orbitals")