Fix Hessian atmlst and example (pyscf#2293)

* add atmlst to examples; fix examples/hessian/01-scf_hessian.py * fix atmlst for hessian * fix lint issue * fix scipy.sparse.linalg.cg issue * fix scipy.sparse.linalg.cg issue * fix scipy.sparse.linalg.cg issue * fix scipy.sparse.linalg.cg issue
sunqm · Jul 8, 2024 · 6d39f69 · 6d39f69
1 parent d488cb7
commit 6d39f69
Show file tree

Hide file tree

Showing 10 changed files with 83 additions and 25 deletions.
diff --git a/examples/grad/01-scf_grad.py b/examples/grad/01-scf_grad.py
@@ -21,7 +21,13 @@
 mf = scf.RHF(mol)
 mf.kernel()
 g = mf.nuc_grad_method()
-g.kernel()
+grad = g.kernel()
+
+# Use atmlst to specify the atoms to calculate the gradients
+atmlst = [0, 1]
+g.verbose = 0
+err = abs(grad[atmlst] - g.kernel(atmlst=atmlst)).max()
+assert err < 1e-6
 
 mf = scf.UHF(mol).x2c()
 mf.kernel()

diff --git a/examples/hessian/01-scf_hessian.py b/examples/hessian/01-scf_hessian.py
@@ -7,7 +7,7 @@
 SCF analytical nuclear hessian calculation.
 '''
 
-from pyscf import gto
+from pyscf import gto, scf, hessian
 
 mol = gto.M(
     atom = [
@@ -22,5 +22,10 @@
 h = mf.Hessian().kernel()
 print(h.shape)
 
+# Use atmlst to specify the atoms to calculate the hessian
+atmlst = [0, 1]
+err = abs(h[atmlst][:, atmlst] - mf.Hessian().kernel(atmlst=atmlst)).max()
+assert err < 1e-6
+
 mf = mol.apply('UKS').x2c().run()
 h = mf.Hessian().kernel()
diff --git a/pyscf/grad/lagrange.py b/pyscf/grad/lagrange.py
@@ -127,7 +127,7 @@ def my_Lvec_last ():
                                                  dtype=bvec.dtype)
         x0_guess = self.get_init_guess (bvec, Adiag, Aop, precond)
         Lvec, info_int = sparse_linalg.cg(Aop_obj, -bvec, x0=x0_guess,
-                                          tol=self.conv_rtol, atol=self.conv_atol,
+                                          atol=self.conv_atol,
                                           maxiter=self.max_cycle, callback=my_call, M=prec_obj)
         logger.info (self, ('Lagrange multiplier determination {} after {} iterations\n'
                             '   |geff| = {}, |Lvec| = {}').format (

diff --git a/pyscf/grad/test/test_casscf.py b/pyscf/grad/test/test_casscf.py
@@ -268,12 +268,12 @@ def test_state_average_mix_scanner(self):
         e2_0 = mcs.e_states[0]
         e2_1 = mcs.e_states[1]
 
-        self.assertAlmostEqual(e_avg, -1.083838462141992e+02, 9)
-        self.assertAlmostEqual(lib.fp(de_avg), -1.034392760319145e-01, 7)
-        self.assertAlmostEqual(e_0, -1.083902661656155e+02, 9)
-        self.assertAlmostEqual(lib.fp(de_0), -6.398921123988113e-02, 7)
-        self.assertAlmostEqual(e_1, -1.083774262627830e+02, 9)
-        self.assertAlmostEqual(lib.fp(de_1), -1.428891618903179e-01, 7)
+        self.assertAlmostEqual(e_avg, -1.083838462141992e+02, 6)
+        self.assertAlmostEqual(lib.fp(de_avg), -1.034392760319145e-01, 6)
+        self.assertAlmostEqual(e_0, -1.083902661656155e+02, 6)
+        self.assertAlmostEqual(lib.fp(de_0), -0.0639891145578155, 6)
+        self.assertAlmostEqual(e_1, -1.083774262627830e+02, 6)
+        self.assertAlmostEqual(lib.fp(de_1), -1.428891618903179e-01, 6)
         self.assertAlmostEqual(de_avg[1,2], (e1_avg-e2_avg)/0.002*lib.param.BOHR, 4)
         self.assertAlmostEqual(de_0[1,2], (e1_0-e2_0)/0.002*lib.param.BOHR, 4)
         self.assertAlmostEqual(de_1[1,2], (e1_1-e2_1)/0.002*lib.param.BOHR, 4)

diff --git a/pyscf/hessian/rhf.py b/pyscf/hessian/rhf.py
@@ -136,9 +136,12 @@ def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     ip1ip2_opt = _make_vhfopt(mol, dm0, 'ip1ip2', 'int2e_ip1ip2')
     ipvip1_opt = _make_vhfopt(mol, dm0, 'ipvip1', 'int2e_ipvip1ipvip2')
     aoslices = mol.aoslice_by_atom()
-    e1 = numpy.zeros((mol.natm,mol.natm,3,3))
-    ej = numpy.zeros((mol.natm,mol.natm,3,3))
-    ek = numpy.zeros((mol.natm,mol.natm,3,3))
+
+    natm = len(atmlst)
+    e1 = numpy.zeros((natm, natm, 3, 3))
+    ej = numpy.zeros((natm, natm, 3, 3))
+    ek = numpy.zeros((natm, natm, 3, 3))
+
     for i0, ia in enumerate(atmlst):
         shl0, shl1, p0, p1 = aoslices[ia]
         shls_slice = (shl0, shl1) + (0, mol.nbas)*3
@@ -159,24 +162,24 @@ def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
         vk1 = vk1.reshape(3,3,nao,nao)
         t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
 
-        ej[i0,i0] += numpy.einsum('xypq,pq->xy', vj1_diag[:,:,p0:p1], dm0[p0:p1])*2
-        ek[i0,i0] += numpy.einsum('xypq,pq->xy', vk1_diag[:,:,p0:p1], dm0[p0:p1])
-        e1[i0,i0] -= numpy.einsum('xypq,pq->xy', s1aa[:,:,p0:p1], dme0[p0:p1])*2
+        ej[i0, i0] += numpy.einsum('xypq,pq->xy', vj1_diag[:,:,p0:p1], dm0[p0:p1])*2
+        ek[i0, i0] += numpy.einsum('xypq,pq->xy', vk1_diag[:,:,p0:p1], dm0[p0:p1])
+        e1[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:,:,p0:p1], dme0[p0:p1])*2
 
         for j0, ja in enumerate(atmlst[:i0+1]):
             q0, q1 = aoslices[ja][2:]
             # *2 for +c.c.
-            ej[i0,j0] += numpy.einsum('xypq,pq->xy', vj1[:,:,q0:q1], dm0[q0:q1])*4
-            ek[i0,j0] += numpy.einsum('xypq,pq->xy', vk1[:,:,q0:q1], dm0[q0:q1])
-            e1[i0,j0] -= numpy.einsum('xypq,pq->xy', s1ab[:,:,p0:p1,q0:q1], dme0[p0:p1,q0:q1])*2
+            ej[i0, j0] += numpy.einsum('xypq,pq->xy', vj1[:,:,q0:q1], dm0[q0:q1])*4
+            ek[i0, j0] += numpy.einsum('xypq,pq->xy', vk1[:,:,q0:q1], dm0[q0:q1])
+            e1[i0, j0] -= numpy.einsum('xypq,pq->xy', s1ab[:,:,p0:p1,q0:q1], dme0[p0:p1,q0:q1])*2
 
             h1ao = hcore_deriv(ia, ja)
-            e1[i0,j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)
+            e1[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)
 
         for j0 in range(i0):
-            e1[j0,i0] = e1[i0,j0].T
-            ej[j0,i0] = ej[i0,j0].T
-            ek[j0,i0] = ek[i0,j0].T
+            e1[j0, i0] = e1[i0, j0].T
+            ej[j0, i0] = ej[i0, j0].T
+            ek[j0, i0] = ek[i0, j0].T
 
     log.timer('RHF partial hessian', *time0)
     return e1, ej, ek

diff --git a/pyscf/hessian/test/test_rhf.py b/pyscf/hessian/test/test_rhf.py
@@ -52,6 +52,15 @@ def test_rhf_hess(self):
         hess = hobj.kernel()
         self.assertAlmostEqual(lib.fp(hess), -0.7816352153153946, 4)
 
+    def test_rhf_hess_atmlst(self):
+        mf = scf.RHF(mol)
+        e0 = mf.kernel()
+
+        atmlst = [0, 1]
+        hess_1 = mf.Hessian().kernel()[atmlst][:, atmlst]
+        hess_2 = mf.Hessian().kernel(atmlst=atmlst)
+        self.assertAlmostEqual(abs(hess_1-hess_2).max(), 0.0, 4)
+
     def test_finite_diff_x2c_rhf_hess(self):
         mf = scf.RHF(mol).x2c()
         mf.conv_tol = 1e-14

diff --git a/pyscf/hessian/test/test_rks.py b/pyscf/hessian/test/test_rks.py
@@ -93,6 +93,17 @@ def grad_partial_R(ia, inc):
     return e2ref
 
 class KnownValues(unittest.TestCase):
+    def test_rks_hess_atmlst(self):
+        mf = dft.RKS(mol)
+        mf.xc = 'pbe0'
+        mf.conv_tol = 1e-14
+        e0 = mf.kernel()
+
+        atmlst = [0, 1]
+        hess_1 = mf.Hessian().kernel()[atmlst][:, atmlst]
+        hess_2 = mf.Hessian().kernel(atmlst=atmlst)
+        self.assertAlmostEqual(abs(hess_1-hess_2).max(), 0.0, 4)
+
     def test_finite_diff_lda_hess(self):
         mf = dft.RKS(mol)
         mf.conv_tol = 1e-14

diff --git a/pyscf/hessian/test/test_uhf.py b/pyscf/hessian/test/test_uhf.py
@@ -50,6 +50,16 @@ def test_uhf_hess(self):
         hess = hobj.kernel()
         self.assertAlmostEqual(lib.fp(hess), -0.20243405976628576, 4)
 
+    def test_uhf_hess_atmlst(self):
+        mf = scf.UHF(mol)
+        mf.conv_tol = 1e-14
+        e0 = mf.kernel()
+
+        atmlst = [0, 1]
+        hess_1 = mf.Hessian().kernel()[atmlst][:, atmlst]
+        hess_2 = mf.Hessian().kernel(atmlst=atmlst)
+        self.assertAlmostEqual(abs(hess_1-hess_2).max(), 0.0, 4)
+
     def test_finite_diff_uhf_hess(self):
         mf = scf.UHF(mol)
         mf.conv_tol = 1e-14

diff --git a/pyscf/hessian/test/test_uks.py b/pyscf/hessian/test/test_uks.py
@@ -96,6 +96,17 @@ def grad_partial_R(ia, inc):
     return e2ref
 
 class KnownValues(unittest.TestCase):
+    def test_uks_hess_atmlst(self):
+        mf = dft.UKS(mol)
+        mf.xc = 'pbe0'
+        mf.conv_tol = 1e-14
+        e0 = mf.kernel()
+
+        atmlst = [0, 1]
+        hess_1 = mf.Hessian().kernel()[atmlst][:, atmlst]
+        hess_2 = mf.Hessian().kernel(atmlst=atmlst)
+        self.assertAlmostEqual(abs(hess_1-hess_2).max(), 0.0, 4)
+
     def test_finite_diff_lda_hess(self):
         mf = dft.UKS(mol)
         mf.conv_tol = 1e-14

diff --git a/pyscf/hessian/uhf.py b/pyscf/hessian/uhf.py
@@ -159,9 +159,12 @@ def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     ip1ip2_opt = _make_vhfopt(mol, dm0, 'ip1ip2', 'int2e_ip1ip2')
     ipvip1_opt = _make_vhfopt(mol, dm0, 'ipvip1', 'int2e_ipvip1ipvip2')
     aoslices = mol.aoslice_by_atom()
-    e1 = numpy.zeros((mol.natm,mol.natm,3,3))  # (A,B,dR_A,dR_B)
-    ej = numpy.zeros((mol.natm,mol.natm,3,3))
-    ek = numpy.zeros((mol.natm,mol.natm,3,3))
+
+    natm = len(atmlst)
+    e1 = numpy.zeros((natm, natm, 3, 3))  # (A,B,dR_A,dR_B)
+    ej = numpy.zeros((natm, natm, 3, 3))
+    ek = numpy.zeros((natm, natm, 3, 3))
+
     for i0, ia in enumerate(atmlst):
         shl0, shl1, p0, p1 = aoslices[ia]
         shls_slice = (shl0, shl1) + (0, mol.nbas)*3