Bugfix: to_cpu, uks gradient, and grid_response (#230)

* several minor bugs

* remove print

* relax df gradient

gpu4pyscf/df/tests/test_df_rks_grad.py

@@ -60,6 +60,7 @@ def _check_grad(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-6):
     mf.grids.prune = None
     mf.nlcgrids.level = nlcgrids_level
     mf.conv_tol = 1e-14
+    mf.direct_scf_tol = 1e-20
     mf.verbose = 1
     mf.kernel()
 
@@ -141,7 +142,7 @@ def test_grad_mgga(self):
 
     def test_grad_rsh(self):
         print('--------RSH testing-------------')
-        _vs_cpu(mol_sph, xc='wb97', disp=None, tol=1e-8)
+        _vs_cpu(mol_sph, xc='wb97', disp=None, tol=1e-7)
 
     def test_grad_nlc(self):
         print('--------nlc testing-------------')

gpu4pyscf/grad/rhf.py

@@ -22,6 +22,7 @@
 import numpy
 from pyscf import lib, gto
 from pyscf.grad import rhf as rhf_grad_cpu
+from gpu4pyscf.lib import utils
 from gpu4pyscf.scf.hf import KohnShamDFT
 from gpu4pyscf.lib.cupy_helper import tag_array, contract, condense, sandwich_dot
 from gpu4pyscf.__config__ import props as gpu_specs
@@ -345,17 +346,12 @@ class GradientsBase(lib.StreamObject):
     as_scanner  = as_scanner
     _tag_rdm1   = rhf_grad_cpu.GradientsBase._tag_rdm1
 
-    # to_cpu can be reused only when __init__ still takes mf
-    def to_cpu(self):
-        mf = self.base.to_cpu()
-        from importlib import import_module
-        mod = import_module(self.__module__.replace('gpu4pyscf', 'pyscf'))
-        cls = getattr(mod, self.__class__.__name__)
-        obj = cls(mf)
-        return obj
 
 class Gradients(GradientsBase):
-    from gpu4pyscf.lib.utils import to_gpu, device
+
+    to_cpu = utils.to_cpu
+    to_gpu = utils.to_gpu
+    device = utils.device
 
     make_rdm1e = rhf_grad_cpu.Gradients.make_rdm1e
     grad_elec = grad_elec

gpu4pyscf/grad/rks.py

@@ -22,12 +22,12 @@
 import cupy
 import pyscf
 from pyscf import lib, gto
-from pyscf.dft import radi
 from pyscf.grad import rks as rks_grad
 from gpu4pyscf.lib.utils import patch_cpu_kernel
 from gpu4pyscf.grad import rhf as rhf_grad
 from gpu4pyscf.dft import numint, xc_deriv, rks
 from gpu4pyscf.dft.numint import _GDFTOpt, AO_THRESHOLD
+from gpu4pyscf.dft import radi
 from gpu4pyscf.lib.cupy_helper import (
     contract, get_avail_mem, add_sparse, tag_array, take_last2d, sandwich_dot)
 from gpu4pyscf.lib import logger
@@ -409,7 +409,7 @@ def get_vxc_full_response(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
 
     #:vmat = cupy.einsum('pi,npq,qj->nij', coeff, vmat, coeff)
     vmat = sandwich_dot(vmat, coeff)
-
+    
     # - sign because nabla_X = -nabla_x
     return excsum, -vmat
 
@@ -424,7 +424,7 @@ def grids_response_cc(grids):
     atm_dist = gto.inter_distance(mol, atm_coords)
     atm_dist = cupy.asarray(atm_dist)
     atm_coords = cupy.asarray(atm_coords)
-
+    
     def _radii_adjust(mol, atomic_radii):
         charges = mol.atom_charges()
         if grids.radii_adjust == radi.treutler_atomic_radii_adjust:
@@ -514,7 +514,6 @@ def get_du(ia, ib):  # JCP 98, 5612 (1993); (B10)
 
         for ia in range(mol.natm):
             dpbecke[:,ia] *= pbecke[ia]
-
         return pbecke, dpbecke
 
     natm = mol.natm

gpu4pyscf/grad/tests/test_grid_response.py

@@ -0,0 +1,86 @@
+# Copyright 2023 The GPU4PySCF Authors. All Rights Reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import numpy as np
+import pyscf
+import cupy
+import unittest
+import pytest
+from pyscf.dft import rks as cpu_rks
+from gpu4pyscf.dft import rks as gpu_rks
+from packaging import version
+
+atom = '''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+
+bas0='def2-tzvpp'
+grids_level = 5
+nlcgrids_level = 3
+def setUpModule():
+    global mol_sph, mol_cart
+    mol_sph = pyscf.M(atom=atom, basis=bas0, max_memory=32000,
+                      output='/dev/null', verbose=1)
+
+    mol_cart = pyscf.M(atom=atom, basis=bas0, max_memory=32000, cart=1,
+                       output='/dev/null', verbose=1)
+
+def tearDownModule():
+    global mol_sph, mol_cart
+    mol_sph.stdout.close()
+    mol_cart.stdout.close()
+    del mol_sph, mol_cart
+
+class KnownValues(unittest.TestCase):
+
+    def test_grids_response(self):
+        mf = cpu_rks.RKS(mol_sph, xc='b3lyp')
+        mf.kernel()
+
+        grids_cpu = mf.grids
+
+        coords_cpu = []
+        w0_cpu = []
+        w1_cpu = []
+        from pyscf.grad.rks import grids_response_cc
+        for coords, w0, w1 in grids_response_cc(grids_cpu):
+            coords_cpu.append(coords)
+            w0_cpu.append(w0)
+            w1_cpu.append(w1)
+
+        mf = cpu_rks.RKS(mol_sph, xc='b3lyp').to_gpu()
+        mf.kernel()
+        grids_gpu = mf.grids
+
+        coords_gpu = []
+        w0_gpu = []
+        w1_gpu = []
+        from gpu4pyscf.grad.rks import grids_response_cc
+        for coords, w0, w1 in grids_response_cc(grids_gpu):
+            coords_gpu.append(coords)
+            w0_gpu.append(w0)
+            w1_gpu.append(w1)
+
+        for w0, w1 in zip(w0_gpu, w0_cpu):
+            assert np.linalg.norm(w0.get() - w1) < 1e-10
+
+        for w0, w1 in zip(w1_gpu, w1_cpu):
+            assert np.linalg.norm(w0.get() - w1) < 1e-10
+
+if __name__ == "__main__":
+    print("Full Tests for grid response")
+    unittest.main()

gpu4pyscf/grad/tests/test_rks_grad.py

@@ -50,6 +50,7 @@ def _check_grad(mol, grid_response=False, xc='B3LYP', disp=None, tol=1e-9):
     mf.grids.level = grids_level
     mf.grids.prune = None
     mf.small_rho_cutoff = 1e-30
+    mf.direct_scf_tol = 1e-20
     if mf._numint.libxc.is_nlc(mf.xc):
         mf.nlcgrids.level = nlcgrids_level
     mf.kernel()
@@ -66,9 +67,8 @@ class KnownValues(unittest.TestCase):
 
     def test_grad_with_grids_response(self):
         print("-----testing DFT gradient with grids response----")
-        #FIXME: The difference (1e-6) between CPU and GPU is too large
         _check_grad(mol_sph, grid_response=True, tol=1e-6)
-
+    
     def test_grad_without_grids_response(self):
         print('-----testing DFT gradient without grids response----')
         _check_grad(mol_sph, grid_response=False)
@@ -108,7 +108,7 @@ def test_grad_d4(self):
     def test_grad_cart(self):
         print('------hybrid GGA Cart testing--------')
         _check_grad(mol_cart, xc='B3LYP', disp=None)
-
+    
 if __name__ == "__main__":
     print("Full Tests for RKS Gradient")
     unittest.main()

gpu4pyscf/grad/tests/test_uks_grad.py

@@ -43,14 +43,13 @@ def tearDownModule():
     mol_cart.stdout.close()
     del mol_sph, mol_cart
 
-# FIXME: Why is the difference between CPU and GPU so large?
-# tol=1e-6 is not acceptable
 def _check_grad(mol, grid_response=False, xc='B3LYP', disp=None, tol=1e-6):
     mf = uks.UKS(mol, xc=xc)
     mf.disp = disp
     mf.grids.level = grids_level
     mf.grids.prune = None
     mf.small_rho_cutoff = 1e-30
+    mf.direct_scf_tol = 1e-20
     if mf._numint.libxc.is_nlc(mf.xc):
         mf.nlcgrids.level = nlcgrids_level
     mf.kernel()

gpu4pyscf/grad/uhf.py

@@ -21,6 +21,7 @@
 from pyscf import lib, gto
 from pyscf.grad import uhf
 from pyscf.grad import rhf as rhf_grad_cpu
+from gpu4pyscf.lib import utils
 from gpu4pyscf.lib.cupy_helper import tag_array, contract
 from gpu4pyscf.df import int3c2e      #TODO: move int3c2e to out of df
 from gpu4pyscf.lib import logger
@@ -99,12 +100,14 @@ def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
         mf_grad.grad_disp = g_disp
         mf_grad.grad_mf = de
 
-    de -= cupy.sum(de, axis=0)/len(atmlst)
     return de.get()
 
 
 class Gradients(rhf_grad.GradientsBase):
-    from gpu4pyscf.lib.utils import to_gpu, device
+
+    to_cpu = utils.to_cpu
+    to_gpu = utils.to_gpu
+    device = utils.device
 
     grad_elec = grad_elec
 

gpu4pyscf/grad/uks.py

@@ -90,7 +90,7 @@ def get_veff(ks_grad, mol=None, dm=None, verbose=None):
             vxc_tmp[0] += vnlc
             vxc_tmp[1] += vnlc
     t0 = logger.timer(ks_grad, 'vxc', *t0)
-
+    
     mo_coeff_alpha = mf.mo_coeff[0]
     mo_coeff_beta = mf.mo_coeff[1]
     occ_coeff0 = cupy.asarray(mo_coeff_alpha[:, mf.mo_occ[0]>0.5], order='C')
@@ -195,8 +195,6 @@ def get_vxc(ni, mol, grids, xc_code, dms, relativity=0, hermi=1,
 
     vmat = take_last2d(vmat, opt.rev_ao_idx)
     exc = None
-    if nset == 1:
-        vmat = vmat[0]
 
     # - sign because nabla_X = -nabla_x
     return exc, -cupy.array(vmat)