UKS hessian with grid response in progress, each individual term is correct, DF hessian is correct, direct hessian is wrong

Author: henryw7

gpu4pyscf/df/hessian/uks.py

@@ -46,10 +46,8 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
 
     mocca = mo_coeff[0][:,mo_occ[0]>0]
     moccb = mo_coeff[1][:,mo_occ[1]>0]
-    dm0a = numpy.dot(mocca, mocca.T)
-    dm0b = numpy.dot(moccb, moccb.T)
-    if mf.do_nlc():
-        raise NotImplementedError("2nd derivative of NLC is not implemented.")
+    dm0a = mocca.dot(mocca.T)
+    dm0b = moccb.dot(moccb.T)
 
     omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
     with_k = mf._numint.libxc.is_hybrid_xc(mf.xc)
@@ -68,24 +66,11 @@ 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)
-    veffa_diag, veffb_diag = uks_hess._get_vxc_diag(hessobj, mo_coeff, mo_occ, max_memory)
-
-    t1 = log.timer_debug1('computing veff_diag', *t1)
-    aoslices = mol.aoslice_by_atom()
-    vxca_dm, vxcb_dm = uks_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]
-        veffa_dm = vxca_dm[ia]
-        veffb_dm = vxcb_dm[ia]
-        de2[i0,i0] += contract('xypq,pq->xy', veffa_diag[:,:,p0:p1], dm0a[p0:p1])*2
-        de2[i0,i0] += contract('xypq,pq->xy', veffb_diag[:,:,p0:p1], dm0b[p0:p1])*2
-        for j0, ja in enumerate(atmlst[:i0+1]):
-            q0, q1 = aoslices[ja][2:]
-            de2[i0,j0] += 2.0*cupy.sum(veffa_dm[:,:,q0:q1], axis=2)
-            de2[i0,j0] += 2.0*cupy.sum(veffb_dm[:,:,q0:q1], axis=2)
-        for j0 in range(i0):
-            de2[j0,i0] = de2[i0,j0].T
+    de2 += uks_hess._get_exc_deriv2(hessobj, mo_coeff, mo_occ, (dm0a, dm0b), max_memory)
+    if mf.do_nlc():
+        raise NotImplementedError("2nd derivative of NLC is not implemented.")
+    #     de2 += _get_enlc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
+
     log.timer('RKS partial hessian', *time0)
     return de2
 

gpu4pyscf/grad/tests/test_vv10_grid.py

@@ -101,6 +101,7 @@ def numerical_denlc(mf, dm, denlc_only = True):
                 energy_p = mf.enlc
             else:
                 energy_p = mf.kernel()
+                assert mf.converged
 
             xyz_m = mol.atom_coords()
             xyz_m[i_atom, i_xyz] -= dx
@@ -114,6 +115,7 @@ def numerical_denlc(mf, dm, denlc_only = True):
                 energy_m = mf.enlc
             else:
                 energy_m = mf.kernel()
+                assert mf.converged
 
             numerical_gradient[i_atom, i_xyz] = (energy_p - energy_m) / (2 * dx)
     mf.reset(mol)

gpu4pyscf/hessian/rks.py

@@ -87,11 +87,13 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
             de2 += rhf_hess._partial_ejk_ip2(
                 mol, dm0, vhfopt, j_factor, k_factor, verbose=verbose)
 
-    t1 = log.timer_debug1('hessian of 2e part', *t1)
+    t1 = log.timer_debug1('hessian of JK part', *t1)
+
     de2 += _get_exc_deriv2(hessobj, mo_coeff, mo_occ, dm0, max_memory, atmlst, log)
     if mf.do_nlc():
         de2 += _get_enlc_deriv2(hessobj, mo_coeff, mo_occ, max_memory, log)
 
+    t1 = log.timer_debug1('hessian of XC part', *t1)
     log.timer('RKS partial hessian', *time0)
     return de2
 
@@ -1086,8 +1088,20 @@ def _get_enlc_deriv2(hessobj, mo_coeff, mo_occ, max_memory, log = None):
     mol = hessobj.mol
     mf = hessobj.base
 
-    mocc = mo_coeff[:,mo_occ>0]
-    dm0 = 2 * mocc @ mocc.T
+    if mo_coeff.ndim == 3:
+        assert mo_coeff.shape[0] == 2
+        assert mo_occ.ndim == 2 and mo_occ.shape[0] == 2
+        mocc0 = mo_coeff[0][:, mo_occ[0]>0]
+        mocc1 = mo_coeff[1][:, mo_occ[1]>0]
+        dm0 = mocc0 @ mocc0.T + mocc1 @ mocc1.T
+        mocc0 = None
+        mocc1 = None
+    else:
+        mocc = mo_coeff[:,mo_occ>0]
+        dm0 = 2 * mocc @ mocc.T
+        mocc = None
+    mo_coeff = None
+    mo_occ = None
 
     grids = mf.nlcgrids
     if grids.coords is None:
@@ -1845,7 +1859,7 @@ def _get_vxc_deriv1_task(hessobj, grids, mo_coeff, mo_occ, max_memory, device_id
 
                     # rho = rho_drho[0]
                     # drho_dr = rho_drho[1:4]
-                    rho_drho_tau = None
+                    rho_drho = None
 
                     depsilon_drho = vxc[0]
                     depsilon_dnablarho = vxc[1:4]
@@ -2368,6 +2382,8 @@ def _get_vnlc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
     mf = hessobj.base
     natm = mol.natm
 
+    if mo_coeff.ndim != 2:
+        raise NotImplementedError("")
     mocc = mo_coeff[:,mo_occ>0]
     dm0 = 2 * mocc @ mocc.T
 
@@ -3428,6 +3444,7 @@ def _get_exc_deriv2_grid_response(hessobj, mo_coeff, mo_occ, max_memory):
     aoslices = mol.aoslice_by_atom()
 
     dm0 = mf.make_rdm1(mo_coeff, mo_occ)
+    assert dm0.ndim == 2
 
     d2e = cupy.zeros([mol.natm, mol.natm, 3, 3])
 

gpu4pyscf/hessian/tests/test_rks_hessian_grid_response.py

@@ -49,14 +49,16 @@ def _get_exc_deriv2_numerical(hessobj, mo_coeff, mo_occ, max_memory):
     """
     mol = hessobj.mol
     mf = hessobj.base
-    mocc = mo_coeff[:,mo_occ>0]
-    dm0 = np.dot(mocc, mocc.T) * 2
+    dm0 = mf.make_rdm1(mo_coeff, mo_occ)
 
     de2 = np.empty([mol.natm, mol.natm, 3, 3])
 
     def get_xc_de(grad_obj, dm):
         assert grad_obj.grid_response
-        from gpu4pyscf.grad.rks import get_exc_full_response
+        if dm.ndim == 2:
+            from gpu4pyscf.grad.rks import get_exc_full_response
+        else:
+            from gpu4pyscf.grad.uks import get_exc_full_response
         mol = grad_obj.mol
         ni = mf._numint
         mf.grids.build()
@@ -95,16 +97,21 @@ def _get_vxc_deriv1_numerical(hessobj, mo_coeff, mo_occ, max_memory):
     """
     mol = hessobj.mol
     mf = hessobj.base
-    mocc = mo_coeff[:,mo_occ>0]
-    dm0 = np.dot(mocc, mocc.T) * 2
+    dm0 = mf.make_rdm1(mo_coeff, mo_occ)
 
     nao = mol.nao
-    vmat = cp.empty([mol.natm, 3, nao, nao])
+    if dm0.ndim == 2:
+        vmat = cp.empty([mol.natm, 3, nao, nao])
+    else:
+        vmat = cp.empty([mol.natm, 3, 2, nao, nao])
 
     def get_vxc_vmat(mol, mf, dm):
         ni = mf._numint
         mf.grids.build()
-        n, exc, vxc = ni.nr_rks(mol, mf.grids, mf.xc, dm)
+        if dm.ndim == 2:
+            n, exc, vxc = ni.nr_rks(mol, mf.grids, mf.xc, dm)
+        else:
+            n, exc, vxc = ni.nr_uks(mol, mf.grids, mf.xc, dm)
         return vxc
 
     dx = 1e-5
@@ -125,12 +132,22 @@ def get_vxc_vmat(mol, mf, dm):
             mf.reset(mol_copy)
             vmat_m = get_vxc_vmat(mol_copy, mf, dm0)
 
-            vmat[i_atom, i_xyz, :, :] = (vmat_p - vmat_m) / (2 * dx)
+            vmat[i_atom, i_xyz] = (vmat_p - vmat_m) / (2 * dx)
     mf.reset(mol)
 
-    vmat = cp.einsum('Adij,jq->Adiq', vmat, mocc)
-    vmat = cp.einsum('Adiq,ip->Adpq', vmat, mo_coeff)
-    return vmat
+    if dm0.ndim == 3:
+        mocc0 = mo_coeff[0][:, mo_occ[0]>0]
+        mocc1 = mo_coeff[1][:, mo_occ[1]>0]
+        vmat0 = cp.einsum('Adij,jq->Adiq', vmat[:,:,0,:,:], mocc0)
+        vmat0 = cp.einsum('Adiq,ip->Adpq', vmat0, mo_coeff[0])
+        vmat1 = cp.einsum('Adij,jq->Adiq', vmat[:,:,1,:,:], mocc1)
+        vmat1 = cp.einsum('Adiq,ip->Adpq', vmat1, mo_coeff[1])
+        return vmat0, vmat1
+    else:
+        mocc = mo_coeff[:,mo_occ>0]
+        vmat = cp.einsum('Adij,jq->Adiq', vmat, mocc)
+        vmat = cp.einsum('Adiq,ip->Adpq', vmat, mo_coeff)
+        return vmat
 
 class KnownValues(unittest.TestCase):
     # All reference results from the same calculation with mf.level_shift = 0