Support stress for Becke grids

gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py

@@ -221,7 +221,7 @@ def get_energy(cell):
 
         assert np.abs(test_energy - ref_energy) < 1e-9
         assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 3e-8
-        # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 3e-7 # TODO: Support Becke grid stress tensor
+        assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-7
 
         dm = mf.make_rdm1()
         kmesh = np.array([1,1,1])
@@ -336,7 +336,7 @@ def get_energy(cell):
 
         assert np.abs(test_energy - ref_energy) < 1e-9
         assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-7
-        # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-6 # TODO: Support Becke grid stress tensor
+        assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-6
 
         dm = mf.make_rdm1()
         kpts = cell.make_kpts(kmesh)
@@ -458,7 +458,7 @@ def get_energy(cell):
 
         assert np.abs(test_energy - ref_energy) < 1e-9
         assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 1e-6
-        # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-6 # TODO: Support Becke grid stress tensor
+        assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 2e-6
 
 if __name__ == '__main__':
     print("Full Tests for PBC GDF Hcore gradient and stress tensor")

gpu4pyscf/pbc/dft/gen_grid.py

@@ -246,7 +246,10 @@ def get_becke_weight_derivative(grids, natm, grid_range = None):
     dweight_dA_unitcell = cp.zeros([natm, 3, ngrids])
     cp.add.at(dweight_dA_unitcell, grids_supatm_to_atm_idx, dweight_dA_supercell)
 
-    return dweight_dA_unitcell
+    weight_stress = cp.einsum("Axg,Ay->xyg", dweight_dA_supercell, grids_supatm_coords)
+
+    dweight_gradient_stress = cp.vstack((dweight_dA_unitcell, weight_stress))
+    return dweight_gradient_stress
 
 class UniformGrids(lib.StreamObject):
     '''Uniform Grid class.'''

gpu4pyscf/pbc/grad/krks.py

@@ -27,60 +27,163 @@
 from gpu4pyscf.lib.cupy_helper import contract
 from gpu4pyscf.pbc.dft import multigrid, multigrid_v3, BeckeGrids
 from gpu4pyscf.pbc.dft.gen_grid import get_becke_weight_derivative
+from gpu4pyscf.pbc.dft.numint import _GTOvalOpt
+from gpu4pyscf.pbc.grad.krks_stress import _eval_ao_strain_derivatives
 
 __all__ = ['Gradients']
 
+XX, XY, XZ = 4, 5, 6
+YX, YY, YZ = 5, 7, 8
+ZX, ZY, ZZ = 6, 8, 9
+
+def get_d2mu_dr2(ao_ks):
+    assert ao_ks.ndim == 4
+    nkpts = ao_ks.shape[0]
+    ngrids = ao_ks.shape[2]
+    nao = ao_ks.shape[3]
+
+    d2mu_dr2 = cp.empty([nkpts, 3, 3, ngrids, nao], dtype = ao_ks.dtype)
+    d2mu_dr2[:,0,0,:,:] = ao_ks[:,XX,:,:]
+    d2mu_dr2[:,0,1,:,:] = ao_ks[:,XY,:,:]
+    d2mu_dr2[:,1,0,:,:] = ao_ks[:,XY,:,:]
+    d2mu_dr2[:,0,2,:,:] = ao_ks[:,XZ,:,:]
+    d2mu_dr2[:,2,0,:,:] = ao_ks[:,XZ,:,:]
+    d2mu_dr2[:,1,1,:,:] = ao_ks[:,YY,:,:]
+    d2mu_dr2[:,1,2,:,:] = ao_ks[:,YZ,:,:]
+    d2mu_dr2[:,2,1,:,:] = ao_ks[:,YZ,:,:]
+    d2mu_dr2[:,2,2,:,:] = ao_ks[:,ZZ,:,:]
+    return d2mu_dr2
+
 def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
     '''derivatives of the Exc per cell'''
     assert dm_kpts.ndim == 3
     xctype = ni._xc_type(xc_code)
     nao = cell.nao
     nkpts = len(kpts)
-    vmat = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype)
+
     if xctype == 'LDA':
+        ao_deriv = 0
+    elif xctype == 'GGA':
         ao_deriv = 1
-        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts,
-                                                   sort_grids=True):
+    elif xctype == 'MGGA':
+        ao_deriv = 1
+    else:
+        raise NotImplementedError(f"Unrecognized xctype = {xctype}")
+    eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv)
+
+    vmat = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype)
+    de_stress_rho = cp.zeros((3,3))
+    exc_sum = 0
+
+    if xctype == 'LDA':
+        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True):
             rho = ni.eval_rho(cell, ao_ks[:,0], dm_kpts, xctype=xctype, hermi=hermi)
-            vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1]
+            exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2]
             wv = weight * vxc[0]
             aow = cp.einsum('kpi,p->kpi', ao_ks[:,0], wv)
             for kn in range(nkpts):
                 vmat[kn] += _d1_dot_(ao_ks[kn,1:4], aow[kn])
 
+            del aow, vxc
+
+            exc_sum += cp.sum(weight * (rho * exc))
+
+            del rho, exc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            ao_ks_strain = ao_ks_strain[:,:,:,0]
+            ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T)
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            drho_stress = contract('kxypg,kpg->xyg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xyg,g->xy', drho_stress, wv).real
+
+            del ao_ks_strain, dm_nu, drho_stress, wv
+
     elif xctype == 'GGA':
-        ao_deriv = 2
-        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts,
-                                                   sort_grids=True):
+        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True):
             rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi)
-            vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1]
+            exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2]
             wv = weight * vxc
             wv[0] *= .5
             for kn in range(nkpts):
                 vmat[kn] += _gga_grad_sum_(ao_ks[kn], wv)
 
+            del vxc
+
+            exc_sum += cp.sum(weight * (rho[0] * exc))
+
+            del rho, exc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T)
+            d2ao = get_d2mu_dr2(ao_ks)
+            ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T)
+            del d2ao
+
+            wv[0] *= 2
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv).real
+            del dmu_stress, dm_nu
+            dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4])
+            dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real
+            del dnu_stress, dm_dmu
+
+            del ao_ks_strain, wv
+
     elif xctype == 'MGGA':
-        ao_deriv = 2
-        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts,
-                                                   sort_grids=True):
+        for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True):
             rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi)
-            vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1]
+            exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2]
             wv = weight * vxc
             wv[0] *= .5
             wv[4] *= .5  # for the factor 1/2 in tau
             for kn in range(nkpts):
                 vmat[kn] += _gga_grad_sum_(ao_ks[kn], wv[:4])
                 vmat[kn] += _tau_grad_dot_(ao_ks[kn], wv[4])
 
+            del vxc
+
+            exc_sum += cp.sum(weight * (rho[0] * exc))
+
+            del rho, exc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T)
+            d2ao = get_d2mu_dr2(ao_ks)
+            ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T)
+            del d2ao
+
+            wv[0] *= 2
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv[0:4]).real
+            del dmu_stress, dm_nu
+            dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4])
+            dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real
+            del dnu_stress, dm_dmu
+            dm_dnu = contract('kpq,kdgq->kdpg', dm_kpts, ao_ks[:,1:4].conj())
+            tau_stress = contract('kxydpg,kdpg->xyg', ao_ks_strain[:,:,:,1:4], dm_dnu)
+            de_stress_rho += 2 * contract('xyg,g->xy', tau_stress, wv[4]).real
+            del tau_stress, dm_dnu
+
+            del ao_ks_strain, wv
+
     elif xctype == 'HF':
         pass
     elif xctype == 'NLC':
         raise NotImplementedError("NLC")
     else:
-        raise NotImplementedError(xc_code)
+        raise NotImplementedError(f"Unrecognized xctype = {xctype}")
+
+    de_stress_weight = exc_sum * cp.eye(3)
 
-    exc = krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1)
-    exc *= -1.0 / nkpts
+    exc = np.zeros((cell.natm + 3, 3))
+    exc[:-3] = -krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1)
+    exc[:-3] *= 1.0 / nkpts
+    exc[-3:] = (de_stress_rho / nkpts + de_stress_weight).get()
     return exc
 
 def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
@@ -103,8 +206,9 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
         ao_deriv = 1
     else:
         raise NotImplementedError(f"Unrecognized xctype = {xctype}")
+    eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv)
 
-    de_grid_response_weight = cp.zeros((natm, 3), dtype=cp.float64)
+    de_grid_response_weight = cp.zeros((natm + 3, 3), dtype=cp.float64)
     g1 = 0
     for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts):
         g0, g1 = g1, g1 + weight.size
@@ -121,6 +225,7 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
 
     dvmat_orbital_response = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype)
     de_grid_response_rho = cp.zeros((natm, 3), dtype=dm_kpts.dtype)
+    de_stress_rho = cp.zeros((3,3), dtype=cp.float64)
 
     g1 = 0
     for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts):
@@ -140,7 +245,20 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
                 dvmat_orbital_response[kn] += vtmp_k
                 de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2
                 del vtmp_k
-            del wv, aow, rho, vxc
+
+            del aow, rho, vxc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            ao_ks_strain = ao_ks_strain[:,:,:,0]
+            associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]]
+            ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T)
+            del associated_supatm_coords
+
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            drho_stress = contract('kxypg,kpg->xyg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xyg,g->xy', drho_stress, wv).real
+
+            del ao_ks_strain, dm_nu, drho_stress, wv
 
         elif xctype == 'GGA':
             rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi)
@@ -153,7 +271,26 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
                 dvmat_orbital_response[kn] += vtmp_k
                 de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2
                 del vtmp_k
-            del wv, rho, vxc
+            del rho, vxc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]]
+            ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T)
+            d2ao = get_d2mu_dr2(ao_ks)
+            ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T)
+            del d2ao, associated_supatm_coords
+
+            wv[0] *= 2
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv).real
+            del dmu_stress, dm_nu
+            dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4])
+            dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real
+            del dnu_stress, dm_dmu
+
+            del ao_ks_strain, wv
 
         elif xctype == 'MGGA':
             rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi)
@@ -167,14 +304,39 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1):
                 dvmat_orbital_response[kn] += vtmp_k
                 de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2
                 del vtmp_k
-            del wv, rho, vxc
+            del rho, vxc
+
+            ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt)
+            associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]]
+            ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T)
+            d2ao = get_d2mu_dr2(ao_ks)
+            ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T)
+            del d2ao, associated_supatm_coords
+
+            wv[0] *= 2
+            dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj())
+            dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv[0:4]).real
+            del dmu_stress, dm_nu
+            dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4])
+            dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu)
+            de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real
+            del dnu_stress, dm_dmu
+            dm_dnu = contract('kpq,kdgq->kdpg', dm_kpts, ao_ks[:,1:4].conj())
+            tau_stress = contract('kxydpg,kdpg->xyg', ao_ks_strain[:,:,:,1:4], dm_dnu)
+            de_stress_rho += 2 * contract('xyg,g->xy', tau_stress, wv[4]).real
+            del tau_stress, dm_dnu
+
+            del ao_ks_strain, wv
 
         else:
             raise NotImplementedError(f"Unrecognized xctype = {xctype}")
     assert g1 == ngrids
 
-    exc = de_grid_response_rho.get().real
-    exc -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1)
+    exc = np.zeros((cell.natm + 3, 3), dtype=np.float64)
+    exc[:-3] = de_grid_response_rho.get().real
+    exc[:-3] -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1)
+    exc[-3:] = de_stress_rho.get()
     exc *= 1.0 / nkpts
     exc += de_grid_response_weight.get()
     return exc
@@ -233,12 +395,7 @@ def energy_ee(self, dm, kpts):
                 fn = get_vxc_full_response
             else:
                 fn = get_vxc
-            de[:-3] = fn(ni, cell, grids, xc, dm, kpts)
-            if isinstance(grids, BeckeGrids):
-                de[-3:] = np.nan
-            else:
-                de[-3:] = multigrid_v3.MultiGridNumInt(cell).energy_strain_gradient(
-                    xc, dm, kpts=kpts, spin=0, with_j=False, with_nuc=False)
+            de = fn(ni, cell, grids, xc, dm, kpts)
         t0 = log.timer_debug1('vxc', *t0)
 
         if j_factor != 0 or k_sr != 0 or k_lr != 0:

gpu4pyscf/pbc/grad/krks_stress.py

@@ -100,13 +100,6 @@ def get_vxc(ks_grad, cell, dm_kpts, kpts, with_j=False, with_nuc=False):
 
     assert kpts.ndim == 2
     assert dm_kpts.ndim == 3
-    if not cell.cart:
-        c2s = asarray(cell.cart2sph_coeff())
-        dm_kpts = sandwich_dot(dm_kpts, c2s.T)
-        # Ensure all AOs are evaluated in the Cartesian GTOs as ao_ks strain
-        # derivatives currently supports Cartesian format only
-        cell = cell.copy()
-        cell.cart = True
     nkpts, nao = dm_kpts.shape[:2]
     assert nkpts == len(kpts)
 
@@ -303,8 +296,8 @@ def _eval_ao_strain_derivatives(cell, coords, kpts=None, deriv=0, out=None,
     coords = cp.asarray(coords.T, order='C')
     bvk_ncells = opt.bvk_ncells
     comp = (deriv+1)*(deriv+2)*(deriv+3)//6
-    nao = cell.nao_nr(cart=True)
-    cart = 1
+    nao = cell.nao_nr()
+    cart = cell.cart
     out = cp.empty((3, 3, comp, bvk_ncells, nao, ngrids))
 
     drv = libpbc.PBCeval_gto_strain_tensor