Support X2C1e for GHF/GKS and implement Density Fitting for GHF/GKS

.github/workflows/unittest.yml

@@ -33,7 +33,7 @@ jobs:
           -v $GITHUB_WORKSPACE:/workspace \
           -v ~/.cache/pip:/root/.cache/pip \
           pyscf/gpu4pyscf-devel:pyscf-2.8 \
-          /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache"
+          /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && pip3 install MCFun && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache"
 
   multi-gpu:
     runs-on: [self-hosted, Linux, X64, 2T4]

gpu4pyscf/df/df_jk.py

@@ -25,11 +25,13 @@
 from gpu4pyscf.lib import logger
 from gpu4pyscf.lib.cupy_helper import (
     contract, transpose_sum, reduce_to_device, tag_array, CPArrayWithTag)
-from gpu4pyscf.dft import rks, uks, numint
+from gpu4pyscf.dft import rks, uks, numint, gks
 from gpu4pyscf.scf import hf, uhf, rohf
 from gpu4pyscf.scf.jk import _check_rsh_factors
 from gpu4pyscf.df import df, int3c2e
 from gpu4pyscf.__config__ import num_devices
+from gpu4pyscf.scf import ghf
+from gpu4pyscf.lib.cupy_helper import asarray
 
 def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
     '''For the given SCF object, update the J, K matrix constructor with
@@ -112,7 +114,7 @@ def reset(self, mol=None):
         return super().reset(mol)
 
     def get_j(self, mol=None, dm=None, hermi=1, omega=None):
-        return self.with_df.get_jk(dm, hermi, True, False, self.direct_scf_tol, omega)[0]
+        return self.get_jk(mol, dm, hermi, with_j=True, with_k=False, omega=omega)[0]
 
     def get_k(self, mol=None, dm=None, hermi=1, omega=None,
               lr_factor=None, sr_factor=None):
@@ -130,18 +132,51 @@ def get_k(self, mol=None, dm=None, hermi=1, omega=None,
     def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
                omega=None):
         if dm is None: dm = self.make_rdm1()
+
         if self.with_df and self.only_dfj:
             vj = vk = None
+            # 1. Calculate DF J (Density Fitting J)
             if with_j:
-                vj = self.get_j(mol, dm, hermi, omega)
+                if isinstance(self, ghf.GHF):
+                    # GHF: Define local strategy and call GHF adapter for J only
+                    def jkbuild(mol_obj, dm_obj, hermi, omega=None):
+                        nao = mol_obj.nao
+                        dm_obj = dm_obj.reshape(-1, nao, nao)
+                        return self.with_df.get_jk(dm_obj, hermi,
+                                                   direct_scf_tol=self.direct_scf_tol,
+                                                   omega=omega)
+                    # Pass with_k=False to get only DF J
+                    vj, _ = ghf.get_jk(mol, dm, hermi, True, False,
+                                       jkbuild=jkbuild, omega=omega)
+                else:
+                    # Standard RHF/UHF: Use Master's get_j logic
+                    vj = self.get_j(mol, dm, hermi, omega)
+
+            # 2. Calculate Exact K (because only_dfj is True, we don't use DF for K)
             if with_k:
                 vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
+
         elif self.with_df:
-            vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_k,
-                                         self.direct_scf_tol, omega)
+            # Full DF mode (DF J + DF K)
+            if isinstance(self, ghf.GHF):
+                # GHF: Define local strategy and call GHF adapter
+                def jkbuild(mol_obj, dm_obj, hermi, omega=None):
+                    nao = mol_obj.nao
+                    dm_obj = dm_obj.reshape(-1, nao, nao)
+                    return self.with_df.get_jk(dm_obj, hermi,
+                                               direct_scf_tol=self.direct_scf_tol,
+                                               omega=omega)
+                vj, vk = ghf.get_jk(mol, dm, hermi, with_j, with_k,
+                                    jkbuild=jkbuild, omega=omega)
+            else:
+                # Standard RHF/UHF: Use Master's direct call
+                vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_k,
+                                             self.direct_scf_tol, omega)
+
         else:
+            # Error handling from Master
             raise ValueError(f"with_df field not found in a df object (type = {type(self)}) during a get_jk() call.")
-            # vj, vk = super().get_jk(mol, dm, hermi, with_j, with_k, omega)
+
         return vj, vk
 
     def Gradients(self):
@@ -283,9 +318,73 @@ def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=None, hermi=1):
                         vklr *= (alpha - hyb)
                         vk += vklr
                     vxc -= vk * .5
-                    exc -= float(cupy.einsum('ij,ji->', dm, vk).real.get()) * .25
+                exc -= float(cupy.einsum('ij,ji->', dm, vk).real.get()) * .25
                 ecoul = float(cupy.einsum('ij,ji->', dm, vj).real.get()) * .5
+            elif isinstance(self, ghf.GHF):
+                ground_state = isinstance(dm, cupy.ndarray) and dm.ndim == 2
+
+                if hermi == 2:  # because rho = 0
+                    n, exc, vxc = 0, 0, 0
+                else:
+                    max_memory = self.max_memory - lib.current_memory()[0]
+                    if ni.collinear[0].lower() != 'm':
+                        raise NotImplementedError('Only multi-colinear GKS is implemented for DF')
+                    if self.grids.coords is None:
+                        self.initialize_grids(mol, dm)
+
+                    if self.grids.coords is None:
+                        self.initialize_grids(mol, dm)
+
+                    n, exc, vxc = ni.get_vxc(mol, self.grids, self.xc, dm,
+                                             hermi=hermi, max_memory=max_memory)
+                    log.debug('nelec by numeric integration = %s', n)
+                    t0 = log.timer('vxc', *t0)
+
+                    if self.do_nlc():
+                        if ni.libxc.is_nlc(self.xc):
+                            xc = self.xc
+                        else:
+                            assert ni.libxc.is_nlc(self.nlc)
+                            xc = self.nlc
+                        n_nlc, enlc, vnlc = ni.nr_nlc_vxc(mol, self.nlcgrids, xc, dm,
+                                                      hermi=hermi, max_memory=max_memory)
+                        exc += enlc
+                        vxc += vnlc
+                        log.debug('nelec with nlc grids = %s', n_nlc)
 
+                if not ni.libxc.is_hybrid_xc(self.xc):
+                    vk = None
+                    vj = self.get_j(mol, dm, hermi)
+                    vxc += vj
+                else:
+                    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(self.xc, spin=mol.spin)
+                    if omega == 0:
+                        vj, vk = self.get_jk(mol, dm, hermi)
+                        vk *= hyb
+                    elif alpha == 0:
+                        vj = self.get_j(mol, dm, hermi)
+                        vk = self.get_k(mol, dm, hermi, omega=-omega)
+                        vk *= hyb
+                    elif hyb == 0:
+                        vj = self.get_j(mol, dm, hermi)
+                        vk = self.get_k(mol, dm, hermi, omega=omega)
+                        vk *= alpha
+                    else:
+                        vj, vk = self.get_jk(mol, dm, hermi)
+                        vk *= hyb
+                        vklr = self.get_k(mol, dm, hermi, omega=omega)
+                        vklr *= (alpha - hyb)
+                        vk += vklr
+
+                    vxc += vj - vk
+
+                    if ground_state:
+                        exc -= cupy.einsum('ij,ji', dm, vk).real * .5
+
+                if ground_state:
+                    ecoul = cupy.einsum('ij,ji', dm, vj).real * .5
+                else:
+                    ecoul = None
             else:
                 raise NotImplementedError("DF only supports R/U/RO KS.")
             t0 = log.timer('veff', *t0)
@@ -302,12 +401,22 @@ def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=None, hermi=1):
             vhf = vj - vk * .5
             ecoul = float(cp.einsum('ij,ji->', dm, vj).real.get()) * .5
             return tag_array(vhf, ecoul=ecoul)
+        elif isinstance(self, ghf.GHF): # (New) GHF branch
+            vj, vk = self.get_jk(mol, dm, hermi=hermi)
+            vhf = vj - vk
+            ecoul = float(cp.einsum('ij,ji->', dm, vj).real.get()) * .5
+            return tag_array(vhf, ecoul=ecoul)
         else:
-            raise NotImplementedError("DF only supports R/U/RO HF.")
+            raise NotImplementedError("DF only supports R/U/RO/G HF.")
 
     def to_cpu(self):
         obj = self.undo_df().to_cpu().density_fit()
-        return utils.to_cpu(self, obj)
+        obj = utils.to_cpu(self, obj)
+        if hasattr(self, 'collinear'):
+            obj.collinear = self.collinear
+        if hasattr(self, 'spin_samples'):
+            obj.spin_samples = self.spin_samples
+        return obj
 
 def _jk_task_with_mo(dfobj, dms, mo_coeff, mo_occ,
                      with_j=True, with_k=True, hermi=0, device_id=0):

gpu4pyscf/df/tests/test_df_ghf.py

@@ -0,0 +1,79 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import numpy as np
+import cupy
+import pyscf
+from pyscf import scf as cpu_scf
+from pyscf.df import df_jk as cpu_df_jk
+from gpu4pyscf import scf as gpu_scf
+from gpu4pyscf.df import df_jk as gpu_df_jk
+
+
+atom = '''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+
+
+bas='def2tzvpp'
+
+
+def setUpModule():
+    global mol
+    mol = pyscf.M(atom=atom, basis=bas, charge=1, spin=1, max_memory=32000,
+                      output='/dev/null', verbose=1)
+
+
+def tearDownModule():
+    global mol
+    mol.stdout.close()
+    del mol
+
+
+class KnownValues(unittest.TestCase):
+    '''
+    known values are obtained by PySCF
+    '''
+    def test_ghf(self):
+        mf_gpu = gpu_scf.GHF(mol).density_fit(auxbasis='def2-tzvpp-jkfit')
+        e_tot = mf_gpu.kernel()
+        mf_cpu = cpu_scf.GHF(mol).density_fit(auxbasis='def2-tzvpp-jkfit')
+        e_pyscf = mf_cpu.kernel()
+
+        assert np.abs(e_tot - e_pyscf) < 1e-5
+
+    def test_to_cpu(self):
+        mf = gpu_scf.GHF(mol).density_fit()
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert isinstance(mf, cpu_df_jk._DFHF)
+        assert np.abs(e_gpu - e_cpu) < 1e-5
+
+    @unittest.skip("skip test_to_gpu")
+    def test_to_gpu(self):
+        mf = cpu_scf.GHF(mol).density_fit()
+        e_cpu = mf.kernel()
+        mf = mf.to_gpu()
+        e_gpu = mf.kernel()
+        assert isinstance(mf, gpu_df_jk._DFHF)
+        assert np.abs(e_gpu - e_cpu) < 1e-5
+
+
+if __name__ == "__main__":
+    print("Full Tests for Generalized Hartree-Fock")
+    unittest.main()

gpu4pyscf/df/tests/test_df_gks.py

@@ -0,0 +1,96 @@
+# Copyright 2021-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import numpy as np
+import cupy
+import pyscf
+from pyscf import lib
+from pyscf import dft as cpu_dft
+from pyscf.df import df_jk as cpu_df_jk
+from gpu4pyscf import dft as gpu_dft
+from gpu4pyscf.df import df_jk as gpu_df_jk
+try:
+    import mcfun
+except ImportError:
+    mcfun = None
+
+
+atom = '''
+O       0.0000000000    -0.0000000000     0.1174000000
+H      -0.7570000000    -0.0000000000    -0.4696000000
+H       0.7570000000     0.0000000000    -0.4696000000
+'''
+
+
+bas='def2tzvpp'
+
+
+def setUpModule():
+    global mol
+    mol = pyscf.M(atom=atom, basis=bas, charge=1, spin=1, max_memory=32000,
+                      output='/dev/null', verbose=1)
+
+
+def tearDownModule():
+    global mol
+    mol.stdout.close()
+    del mol
+
+
+class KnownValues(unittest.TestCase):
+    '''
+    known values are obtained by PySCF
+    '''
+    def test_gks(self):
+        mf_gpu = gpu_dft.GKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+        mf_gpu.collinear = 'm'
+        mf_gpu._numint.spin_samples = 6
+        e_tot = mf_gpu.kernel()
+        if mcfun is not None:
+            mf_cpu = cpu_dft.GKS(mol, xc='b3lyp').density_fit(auxbasis='def2-tzvpp-jkfit')
+            mf_cpu.collinear = 'm'
+            mf_cpu._numint.spin_samples = 6
+            e_pyscf = mf_cpu.kernel()
+            assert np.abs(e_tot - e_pyscf) < 1e-5
+            assert np.abs(lib.fp(mf_cpu.mo_energy) - lib.fp(mf_gpu.mo_energy.get())) < 1e-5
+        assert np.abs(e_tot - -75.99882822956384) < 1e-5
+        assert np.abs(-96.56444462841858 - lib.fp(mf_gpu.mo_energy.get())) < 1e-5
+
+    @unittest.skipIf(mcfun is None, "mcfun library not found.")
+    def test_to_cpu(self):
+        mf = gpu_dft.GKS(mol, xc='b3lyp').density_fit()
+        mf.collinear = 'm'
+        mf._numint.spin_samples = 6
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert isinstance(mf, cpu_df_jk._DFHF)
+        assert np.abs(e_gpu - e_cpu) < 1e-5
+
+    @unittest.skip("skip test_to_gpu")
+    def test_to_gpu(self):
+        mf = cpu_dft.GKS(mol, xc='b3lyp').density_fit()
+        mf.collinear = 'm'
+        mf._numint.spin_samples = 6
+        e_cpu = mf.kernel()
+        mf = mf.to_gpu()
+        e_gpu = mf.kernel()
+        assert isinstance(mf, gpu_df_jk._DFHF)
+        assert np.abs(e_gpu - e_cpu) < 1e-5
+
+
+if __name__ == "__main__":
+    print("Full Tests for Generalized Hartree-Fock")
+    unittest.main()

gpu4pyscf/dft/gks.py

@@ -130,7 +130,9 @@ def get_veff(ks, mol=None, dm=None, dm_last=None, vhf_last=None, hermi=1):
 
 
 class GKS(rks.KohnShamDFT, GHF):
-    to_gpu = utils.to_gpu
+
+    # to_gpu = utils.to_gpu
+    to_gpu = NotImplemented
     device = utils.device
 
     def __init__(self, mol, xc='LDA,VWN'):
@@ -169,6 +171,8 @@ def spin_samples(self, val):
     to_hf = NotImplemented
 
     def to_cpu(self):
-        mf = gks.GKS(self.mol)
+        mf = gks.GKS(self.mol, xc=self.xc)
         utils.to_cpu(self, out=mf)
+        mf.collinear = self.collinear
+        mf.spin_samples = self.spin_samples
         return mf