Uks dev (#78)

* Add the UHF module.

* Add the UHF and unit test.

* Reuse the _kernel in hf.py, and correct the unit test for UHF.

* Add the UKS and a unit test for it.

* Fix typos in test_uks.py

gpu4pyscf/dft/tests/test_uks.py

@@ -0,0 +1,95 @@
+# gpu4pyscf is a plugin to use Nvidia GPU in PySCF package
+#
+# Copyright (C) 2022 Qiming Sun
+#
+# 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 unittest
+import pyscf
+from gpu4pyscf.dft import uks
+
+atom = '''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+bas='def2-qzvpp'
+grids_level = 3
+nlcgrids_level = 1
+
+def setUpModule():
+    global mol
+    mol = pyscf.M(
+        atom=atom,
+        basis=bas,
+        max_memory=32000,
+        verbose = 1,
+        spin = 1,
+        charge = 1,
+        output = '/dev/null'
+    )
+
+def tearDownModule():
+    global mol
+    mol.stdout.close()
+    del mol
+
+def run_dft(xc):
+    mf = uks.UKS(mol, xc=xc)
+    mf.grids.level = grids_level
+    mf.nlcgrids.level = nlcgrids_level
+    e_dft = mf.kernel()
+    return e_dft
+
+class KnownValues(unittest.TestCase):
+    '''
+    known values are obtained by pyscf
+    '''
+    def test_uks_lda(self):
+        print('------- LDA ----------------')
+        e_tot = run_dft("LDA, vwn5")
+        assert np.allclose(e_tot, -75.42821982483972)
+
+    def test_uks_pbe(self):
+        print('------- PBE ----------------')
+        e_tot = run_dft('PBE')
+        assert np.allclose(e_tot, -75.91732813416843)
+
+    def test_uks_b3lyp(self):
+        print('-------- B3LYP -------------')
+        e_tot = run_dft('B3LYP')
+        assert np.allclose(e_tot, -76.00306439862237)
+
+    def test_uks_m06(self):
+        print('--------- M06 --------------')
+        e_tot = run_dft("M06")
+        assert np.allclose(e_tot, -75.96551006522827)
+
+    def test_uks_wb97(self):
+        print('-------- wB97 --------------')
+        e_tot = run_dft("HYB_GGA_XC_WB97")
+        assert np.allclose(e_tot, -75.987601337562)
+
+    def test_uks_vv10(self):
+        print("------- wB97m-v -------------")
+        e_tot = run_dft('HYB_MGGA_XC_WB97M_V')
+        assert np.allclose(e_tot, -75.97363094678428)
+
+    #TODO: add test cases for D3/D4 and gradient
+
+if __name__ == "__main__":
+    print("Full Tests for dft")
+    unittest.main()
+

gpu4pyscf/dft/uks.py

@@ -15,16 +15,121 @@
 # 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 cupy
 from pyscf.dft import uks
-from gpu4pyscf.dft import numint
-from gpu4pyscf.scf.uhf import UHF
+from pyscf import lib
+from gpu4pyscf import scf
+from gpu4pyscf.lib import logger
+from gpu4pyscf.dft import numint, gen_grid, rks
+from gpu4pyscf.lib.cupy_helper import tag_array
 
-class UKS(uks.UKS):
+
+def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
+    '''Coulomb + XC functional for UKS.  See pyscf/dft/rks.py
+    :func:`get_veff` fore more details.
+    '''
+    if mol is None: mol = ks.mol
+    if dm is None: dm = ks.make_rdm1()
+    t0 = logger.init_timer(ks)
+    rks.initialize_grids(ks, mol, dm)
+
+    if hasattr(ks, 'screen_tol') and ks.screen_tol is not None:
+        ks.direct_scf_tol = ks.screen_tol
+    ground_state = (isinstance(dm, cupy.ndarray) and dm.ndim == 3)
+
+    ni = ks._numint
+    if hermi == 2:  # because rho = 0
+        n, exc, vxc = (0,0), 0, 0
+    else:
+        max_memory = ks.max_memory - lib.current_memory()[0]
+        n, exc, vxc = ni.nr_uks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
+        logger.debug(ks, 'nelec by numeric integration = %s', n)
+        if ks.nlc or ni.libxc.is_nlc(ks.xc):
+            if ni.libxc.is_nlc(ks.xc):
+                xc = ks.xc
+            else:
+                assert ni.libxc.is_nlc(ks.nlc)
+                xc = ks.nlc
+            n, enlc, vnlc = ni.nr_nlc_vxc(mol, ks.nlcgrids, xc, dm[0]+dm[1],
+                                          max_memory=max_memory)
+            exc += enlc
+            vxc += vnlc
+            logger.debug(ks, 'nelec with nlc grids = %s', n)
+        t0 = logger.timer(ks, 'vxc', *t0)
+
+    if not ni.libxc.is_hybrid_xc(ks.xc):
+        vk = None
+        if (ks._eri is None and ks.direct_scf and
+            getattr(vhf_last, 'vj', None) is not None):
+            ddm = cupy.asarray(dm) - cupy.asarray(dm_last)
+            vj = ks.get_j(mol, ddm[0]+ddm[1], hermi)
+            vj += vhf_last.vj
+        else:
+            vj = ks.get_j(mol, dm[0]+dm[1], hermi)
+        vxc += vj
+    else:
+        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)
+        if (ks._eri is None and ks.direct_scf and
+            getattr(vhf_last, 'vk', None) is not None):
+            ddm = cupy.asarray(dm) - cupy.asarray(dm_last)
+            vj, vk = ks.get_jk(mol, ddm, hermi)
+            vk *= hyb
+            if abs(omega) > 1e-10:  # For range separated Coulomb operator
+                vklr = ks.get_k(mol, ddm, hermi, omega)
+                vklr *= (alpha - hyb)
+                vk += vklr
+            vj = vj[0] + vj[1] + vhf_last.vj
+            vk += vhf_last.vk
+        else:
+            vj, vk = ks.get_jk(mol, dm, hermi)
+            vj = vj[0] + vj[1]
+            vk *= hyb
+            if abs(omega) > 1e-10:
+                vklr = ks.get_k(mol, dm, hermi, omega=omega)
+                vklr *= (alpha - hyb)
+                vk += vklr
+        vxc += vj - vk
+
+        if ground_state:
+            exc -=(cupy.einsum('ij,ji', dm[0], vk[0]).real +
+                   cupy.einsum('ij,ji', dm[1], vk[1]).real) * .5
+    if ground_state:
+        ecoul = cupy.einsum('ij,ji', dm[0]+dm[1], vj).real * .5
+    else:
+        ecoul = None
+    t0 = logger.timer_debug1(ks, 'jk total', *t0)
+    vxc = tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
+    return vxc
+
+
+def energy_elec(ks, dm=None, h1e=None, vhf=None):
+    if dm is None: dm = ks.make_rdm1()
+    if h1e is None: h1e = ks.get_hcore()
+    if vhf is None or getattr(vhf, 'ecoul', None) is None:
+        vhf = ks.get_veff(ks.mol, dm)
+    if not (isinstance(dm, cupy.ndarray) and dm.ndim == 2):
+        dm = dm[0] + dm[1]
+    return rks.energy_elec(ks, dm, h1e, vhf)
+
+
+class UKS(scf.uhf.UHF, uks.UKS):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+    _keys = {'disp', 'screen_tol'}
 
-    def __init__(self, mol, xc='LDA,VWN'):
+    def __init__(self, mol, xc='LDA,VWN', disp=None):
         super().__init__(mol, xc)
+        self.disp = disp
         self._numint = numint.NumInt()
+        self.screen_tol = 1e-14
+
+        grids_level = self.grids.level
+        self.grids = gen_grid.Grids(mol)
+        self.grids.level = grids_level
 
-    get_jk = UHF.get_jk
-    _eigh = UHF._eigh
+        nlcgrids_level = self.nlcgrids.level
+        self.nlcgrids = gen_grid.Grids(mol)
+        self.nlcgrids.level = nlcgrids_level
+
+    energy_elec = energy_elec
+    get_veff = get_veff
+

gpu4pyscf/scf/diis.py

@@ -15,7 +15,7 @@
 #
 # Author: Qiming Sun <[email protected]>
 #
-# modified by Xiaojie Wu <[email protected]>
+# modified by Xiaojie Wu <[email protected]>; Zhichen Pu <[email protected]>
 
 """
 DIIS
@@ -64,9 +64,13 @@ def get_num_vec(self):
 
 def get_err_vec(s, d, f):
     '''error vector = SDF - FDS'''
-    if isinstance(f, cupy.ndarray):
+    if isinstance(f, cupy.ndarray) and f.ndim == 2:
         sdf = reduce(cupy.dot, (s,d,f))
-        errvec = (sdf.T.conj() - sdf)
-        return errvec
+        errvec = (sdf.conj().T - sdf).ravel()
+    elif f.ndim == s.ndim+1 and f.shape[0] == 2:  # for UHF
+        errvec = cupy.hstack([
+            get_err_vec(s, d[0], f[0]).ravel(),
+            get_err_vec(s, d[1], f[1]).ravel()])
     else:
-        cpu_diis.get_err_vec(s, d, f)
+        raise RuntimeError('Unknown SCF DIIS type')
+    return errvec