Repository: pyscf/pyscf — pull request #2797 (python)
Title: When building RSH K integrals, cderi is overwritten (fix #2790)

Causing a KeyError when reading cderi file

Resolved issues:
title: [BUG REPORT] RSH density fitting fails with j3c-lr KeyError
body: This example using the range-separated hybrid functional HSE06 fails when trying to read from the `j3c-lr` density fitting object:
```
import numpy as np
from pyscf.pbc import gto, dft, df

print(df.df.LONGRANGE_AFT_TURNOVER_THRESHOLD)

cell = gto.M(
	a = np.array([[0.000, 2.715, 2.715],
			[2.715, 0.000, 2.715],
			[2.715, 2.715, 0.000]]),
	atom = '''Si 0.0000 0.0000 0.0000
			Si 1.3575 1.3575 1.3575''',
	basis = 'cc-pvdz',
)

kpts = cell.make_kpts([2,2,2])

kmf = dft.KRKS(cell, kpts, xc='HYB_GGA_XC_HSE06').density_fit()
kmf.kernel()
```

Error:
```
Traceback (most recent call last):
  File "/gpfs/home/mhott/examples_pyscf/pbc/hse_bug.py", line 23, in <module>
    kmf.kernel()
  File "<string>", line 2, in kernel
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/scf/hf.py", line 2016, in scf
    kernel(self, self.conv_tol, self.conv_tol_grad,
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/scf/hf.py", line 129, in kernel
    vhf = mf.get_veff(mol, dm)
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/dft/krks.py", line 119, in get_veff
    vk = ks.get_k(cell, dm, hermi, kpts, kpts_band, omega=-omega)
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/scf/khf.py", line 576, in get_k
    return self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band,
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/scf/khf.py", line 589, in get_jk
    vj, vk = self.with_df.get_jk(dm_kpts, hermi, kpts, kpts_band,
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df.py", line 438, in get_jk
    return rsh_df.get_jk(dm, hermi, kpts, kpts_band, with_j, with_k,
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df.py", line 448, in get_jk
    vk = df_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df_jk.py", line 611, in get_k_kpts
    make_kpt(ki, kj, True)
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df_jk.py", line 444, in make_kpt
    for LpqR, LpqI, sign in mydf.sr_loop((kpti,kptj), max_memory, False):
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df.py", line 360, in sr_loop
    with _load3c(self._cderi, self._dataname, kpti_kptj) as j3c:
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df.py", line 699, in __enter__
    return self.getitem(*self.kpti_kptj)
  File "/gpfs/home/mhott/.venv/lib/python3.9/site-packages/pyscf/pbc/df/df.py", line 753, in getitem
    raise KeyError(f'Key {key} not found')
KeyError: 'Key j3c-lr/-0.110000/8 not found'
```
The index of the k-point it's trying to reference is always the size of the k-grid (2x2x2 = 8 in this case, for 4x4x4 is 64, etc.) which is not included in the list of kpt_ij_idx, which for the 2x2x2 k-grid is `kpt_ij_idx = [ 0  9 18 27 36 45 54 63]`. I noticed that in every case, the index N_k+1 was included in this list (2x2x2 + 1 = 9, 4x4x4 + 1 = 65, etc.), but I'm not sure if this is relevant since I'm not entirely sure how these indices are generated.  

As discussed in this issue: [1193](https://github.com/pyscf/pyscf/issues/1193), I have set `LONGRANGE_AFT_TURNOVER_THRESHOLD=0.0` in `pyscf.df.df` to prevent it from switching to AFTDF which requires many TB of memory.

pyscf version: 2.8.0, 2.9.0
install type: pip install
operating system: Rocky Linux 9.4


Hints:

Fix patch:
diff --git a/pyscf/pbc/df/aft.py b/pyscf/pbc/df/aft.py
index 53124621b5..e21c478ef0 100644
--- a/pyscf/pbc/df/aft.py
+++ b/pyscf/pbc/df/aft.py
@@ -536,8 +536,6 @@ def range_coulomb(self, omega):
             rsh_df = self._rsh_df[key]
         else:
             rsh_df = self._rsh_df[key] = self.copy().reset()
-            if hasattr(self, '_dataname'):
-                rsh_df._dataname = f'{self._dataname}-lr/{key}'
             logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
 
         cell = self.cell
diff --git a/pyscf/pbc/df/df.py b/pyscf/pbc/df/df.py
index 3e3f1456c7..d6cbdae319 100644
--- a/pyscf/pbc/df/df.py
+++ b/pyscf/pbc/df/df.py
@@ -60,8 +60,6 @@
 from pyscf.pbc.df.rsdf_builder import _RSGDFBuilder, _RSNucBuilder, LINEAR_DEP_THR
 from pyscf import __config__
 
-LONGRANGE_AFT_TURNOVER_THRESHOLD = 2.5
-
 
 def make_modrho_basis(cell, auxbasis=None, drop_eta=None):
     r'''Generate a cell object using the density fitting auxbasis as
@@ -259,6 +257,7 @@ def build(self, j_only=None, with_j3c=True, kpts_band=None):
                 if self._cderi == cderi and os.path.isfile(cderi):
                     logger.warn(self, 'File %s (specified by ._cderi) is '
                                 'overwritten by GDF initialization.', cderi)
+                    os.remove(cderi)
                 else:
                     logger.warn(self, 'Value of ._cderi is ignored. '
                                 'DF integrals will be saved in file %s .', cderi)
@@ -416,18 +415,16 @@ def get_nuc(self, kpts=None):
     # post-HF methods.
     def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
                with_j=True, with_k=True, omega=None, exxdiv=None):
-        if omega is not None:  # J/K for RSH functionals
+        if omega is not None and omega != 0:  # J/K for RSH functionals
             cell = self.cell
             # * AFT is computationally more efficient than GDF if the Coulomb
             #   attenuation tends to the long-range role (i.e. small omega).
             # * Note: changing to AFT integrator may cause small difference to
-            #   the GDF integrator. If a very strict GDF result is desired,
-            #   we can disable this trick by setting
-            #   LONGRANGE_AFT_TURNOVER_THRESHOLD to 0.
+            #   the GDF integrator.
             # * The sparse mesh is not appropriate for low dimensional systems
             #   with infinity vacuum since the ERI may require large mesh to
             #   sample density in vacuum.
-            if (abs(omega) < LONGRANGE_AFT_TURNOVER_THRESHOLD and
+            if (omega > 0 and
                 cell.dimension >= 2 and cell.low_dim_ft_type != 'inf_vacuum'):
                 mydf = aft.AFTDF(cell, self.kpts)
                 ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, omega)
@@ -471,6 +468,43 @@ def prange(self, start, stop, step):
         '''
         return lib.prange(start, stop, step)
 
+    @contextlib.contextmanager
+    def range_coulomb(self, omega):
+        '''Creates a temporary density fitting object for RSH-DF integrals.
+        In this context, only LR or SR integrals for mol and auxmol are computed.
+        '''
+        cell = self.cell
+        if cell.dimension != 0:
+            assert omega < 0
+
+        key = '%.6f' % omega
+        if key in self._rsh_df:
+            rsh_df = self._rsh_df[key]
+        else:
+            rsh_df = self._rsh_df[key] = self.copy().reset()
+            rsh_df._dataname = f'{self._dataname}-sr/{key}'
+            logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
+
+        auxcell = getattr(self, 'auxcell', None)
+
+        cell_omega = cell.omega
+        cell.omega = omega
+        auxcell_omega = None
+        if auxcell is not None:
+            auxcell_omega = auxcell.omega
+            auxcell.omega = omega
+
+        assert rsh_df.cell.omega == omega
+        if getattr(rsh_df, 'auxcell', None) is not None:
+            assert rsh_df.auxcell.omega == omega
+
+        try:
+            yield rsh_df
+        finally:
+            cell.omega = cell_omega
+            if auxcell_omega is not None:
+                auxcell.omega = auxcell_omega
+
 ################################################################################
 # With this function to mimic the molecular DF.loop function, the pbc gamma
 # point DF object can be used in the molecular code
diff --git a/pyscf/pbc/df/df_jk.py b/pyscf/pbc/df/df_jk.py
index 88da4d703e..b6f51d1dcb 100644
--- a/pyscf/pbc/df/df_jk.py
+++ b/pyscf/pbc/df/df_jk.py
@@ -280,8 +280,12 @@ def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), kpts_band=None,
             log.warn('DF integrals for band k-points were not found %s. '
                      'DF integrals will be rebuilt to include band k-points.',
                      mydf._cderi)
-        mydf.build(kpts_band=kpts_band)
+        mydf.build(j_only=False, kpts_band=kpts_band)
         t0 = log.timer_debug1('Init get_k_kpts', *t0)
+    elif mydf._j_only:
+        log.warn('DF integrals for HF exchange were not initialized. '
+                 'df.j_only cannot be used with hybrid functional. DF integrals will be rebuilt.')
+        mydf.build(j_only=False, kpts_band=kpts_band)
 
     mo_coeff = getattr(dm_kpts, 'mo_coeff', None)
     if mo_coeff is not None:
diff --git a/pyscf/pbc/df/mdf.py b/pyscf/pbc/df/mdf.py
index 15fd9d8b36..da497232fb 100644
--- a/pyscf/pbc/df/mdf.py
+++ b/pyscf/pbc/df/mdf.py
@@ -153,14 +153,12 @@ def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
             # * AFT is computationally more efficient than MDF if the Coulomb
             #   attenuation tends to the long-range role (i.e. small omega).
             # * Note: changing to AFT integrator may cause small difference to
-            #   the MDF integrator. If a very strict MDF result is desired,
-            #   we can disable this trick by setting
-            #   LONGRANGE_AFT_TURNOVER_THRESHOLD to 0.
+            #   the MDF integrator.
             # * The sparse mesh is not appropriate for low dimensional systems
             #   with infinity vacuum since the ERI may require large mesh to
             #   sample density in vacuum.
-            if (omega < df.LONGRANGE_AFT_TURNOVER_THRESHOLD and
-                cell.dimension >= 2 and cell.low_dim_ft_type != 'inf_vacuum'):
+            if (omega > 0 and
+                (cell.dimension >= 2 or cell.low_dim_ft_type != 'inf_vacuum')):
                 mydf = aft.AFTDF(cell, self.kpts)
                 ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, omega)
                 mydf.mesh = cell.cutoff_to_mesh(ke_cutoff)
diff --git a/pyscf/pbc/df/rsdf.py b/pyscf/pbc/df/rsdf.py
index a11c79a073..a9e7b13100 100644
--- a/pyscf/pbc/df/rsdf.py
+++ b/pyscf/pbc/df/rsdf.py
@@ -561,6 +561,20 @@ def make_j3c(self, cderi_file, intor='int3c2e', aosym='s2', comp=None,
             ijlst_mapping[kpti_idx * nkpts + kptj_idx] = np.arange(len(kptij_lst))
             kk_idx = kpti_idx * nkpts + kptj_idx
 
+        # TODO: Store rs_density_fit cderi tensor in v1 format for the moment.
+        # It should be changed to 'v2' format in the future.
+        data_version = 'v1'
+        if h5py.is_hdf5(cderi_file):
+            feri = lib.H5FileWrap(cderi_file, 'a')
+            if 'kpts' in feri:
+                del feri['j3c-kptij']
+            if dataname in feri:
+                log.warn(f'Overwritting {dataname} in {cderi_file}.')
+                del feri[dataname]
+        else:
+            feri = lib.H5FileWrap(cderi_file, 'w')
+        feri['j3c-kptij'] = kptij_lst
+
         fswap = self.outcore_auxe2(cderi_file, intor, aosym, comp,
                                    kptij_lst, j_only, 'j3c-junk', shls_slice)
         cpu1 = log.timer_debug1('3c2e', *cpu1)
@@ -577,16 +591,6 @@ def make_j3c(self, cderi_file, intor='int3c2e', aosym='s2', comp=None,
         # Add (1) short-range G=0 (i.e., charge) part and (2) long-range part
         tspans = np.zeros((3,2))    # lr, j2c_inv, j2c_cntr
         tspannames = ["ftaop+pw", "j2c_inv", "j2c_cntr"]
-        feri = lib.H5FileWrap(cderi_file, 'w')
-
-        # TODO: Store rs_density_fit cderi tensor in v1 format for the moment.
-        # It should be changed to 'v2' format in the future.
-        data_version = 'v1'
-        if data_version == 'v1':
-            feri['j3c-kptij'] = kptij_lst
-        else:
-            feri['kpts'] = kpts
-            ijlst_mapping = None
         def make_cderi(kpt, kpt_ij_idx, j2c):
             log.debug1('make_cderi for %s', kpt)
             kptjs = kpts[kpt_ij_idx % nkpts]
diff --git a/pyscf/pbc/df/rsdf_builder.py b/pyscf/pbc/df/rsdf_builder.py
index 29c3a839b2..10a0a98667 100644
--- a/pyscf/pbc/df/rsdf_builder.py
+++ b/pyscf/pbc/df/rsdf_builder.py
@@ -924,14 +924,23 @@ def make_j3c(self, cderi_file, intor='int3c2e', aosym='s2', comp=None,
         else:
             kk_idx = None
 
+        if h5py.is_hdf5(cderi_file):
+            feri = lib.H5FileWrap(cderi_file, 'a')
+            if 'kpts' in feri:
+                del feri['kpts']
+                del feri['aosym']
+            if dataname in feri:
+                log.warn(f'Overwritting {dataname} in {cderi_file}.')
+                del feri[dataname]
+        else:
+            feri = lib.H5FileWrap(cderi_file, 'w')
+        feri['kpts'] = kpts
+        feri['aosym'] = aosym
+
         fswap = self.outcore_auxe2(cderi_file, intor, aosym, comp, j_only,
                                    'j3c', shls_slice, kk_idx=kk_idx)
         cpu1 = log.timer('pass1: real space int3c2e', *cpu0)
 
-        feri = lib.H5FileWrap(cderi_file, 'w')
-        feri['kpts'] = kpts
-        feri['aosym'] = aosym
-
         if aosym == 's2':
             nao_pair = nao*(nao+1)//2
         else:
diff --git a/pyscf/pbc/dft/krks.py b/pyscf/pbc/dft/krks.py
index 5c4387c5c3..0705df56b5 100644
--- a/pyscf/pbc/dft/krks.py
+++ b/pyscf/pbc/dft/krks.py
@@ -105,23 +105,17 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
         vxc += vj
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=cell.spin)
-        if getattr(ks.with_df, '_j_only', False):  # for GDF and MDF
-            logger.warn(ks, 'df.j_only cannot be used with hybrid functional')
-            ks.with_df._j_only = False
-            # Rebuild df object due to the change of parameter _j_only
-            if ks.with_df._cderi is not None:
-                ks.with_df.build()
         if omega == 0:
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
         elif alpha == 0: # LR=0, only SR exchange
-            vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
             vk = ks.get_k(cell, dm, hermi, kpts, kpts_band, omega=-omega)
             vk *= hyb
-        elif hyb == 0: # SR=0, only LR exchange
             vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
+        elif hyb == 0: # SR=0, only LR exchange
             vk = ks.get_k(cell, dm, hermi, kpts, kpts_band, omega=omega)
             vk *= alpha
+            vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
         else: # SR and LR exchange with different ratios
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
diff --git a/pyscf/pbc/dft/krks_ksymm.py b/pyscf/pbc/dft/krks_ksymm.py
index 0c6adf10a1..c072f52a9f 100644
--- a/pyscf/pbc/dft/krks_ksymm.py
+++ b/pyscf/pbc/dft/krks_ksymm.py
@@ -92,12 +92,6 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
         vxc += vj
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=cell.spin)
-        if getattr(ks.with_df, '_j_only', False):  # for GDF and MDF
-            logger.warn(ks, 'df.j_only cannot be used with hybrid functional')
-            ks.with_df._j_only = False
-            # Rebuild df object due to the change of parameter _j_only
-            if ks.with_df._cderi is not None:
-                ks.with_df.build()
         if omega == 0:
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
diff --git a/pyscf/pbc/dft/kuks.py b/pyscf/pbc/dft/kuks.py
index 3a6f4f9995..8ddeeb40ab 100644
--- a/pyscf/pbc/dft/kuks.py
+++ b/pyscf/pbc/dft/kuks.py
@@ -89,23 +89,17 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
         vxc += vj
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=cell.spin)
-        if getattr(ks.with_df, '_j_only', False):  # for GDF and MDF
-            logger.warn(ks, 'df.j_only cannot be used with hybrid functional')
-            ks.with_df._j_only = False
-            # Rebuild df object due to the change of parameter _j_only
-            if ks.with_df._cderi is not None:
-                ks.with_df.build()
         if omega == 0:
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
         elif alpha == 0: # LR=0, only SR exchange
-            vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
             vk = ks.get_k(cell, dm, hermi, kpts, kpts_band, omega=-omega)
             vk *= hyb
-        elif hyb == 0: # SR=0, only LR exchange
             vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
+        elif hyb == 0: # SR=0, only LR exchange
             vk = ks.get_k(cell, dm, hermi, kpts, kpts_band, omega=omega)
             vk *= alpha
+            vj = ks.get_j(cell, dm, hermi, kpts, kpts_band)
         else: # SR and LR exchange with different ratios
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
diff --git a/pyscf/pbc/dft/kuks_ksymm.py b/pyscf/pbc/dft/kuks_ksymm.py
index 39404b6127..3c272458ed 100644
--- a/pyscf/pbc/dft/kuks_ksymm.py
+++ b/pyscf/pbc/dft/kuks_ksymm.py
@@ -92,12 +92,6 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
         vxc += vj
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=cell.spin)
-        if getattr(ks.with_df, '_j_only', False):  # for GDF and MDF
-            logger.warn(ks, 'df.j_only cannot be used with hybrid functional')
-            ks.with_df._j_only = False
-            # Rebuild df object due to the change of parameter _j_only
-            if ks.with_df._cderi is not None:
-                ks.with_df.build()
         if omega == 0:
             vj, vk = ks.get_jk(cell, dm, hermi, kpts, kpts_band)
             vk *= hyb
diff --git a/pyscf/pbc/dft/rks.py b/pyscf/pbc/dft/rks.py
index 8f485dfa63..a08ecf3b05 100644
--- a/pyscf/pbc/dft/rks.py
+++ b/pyscf/pbc/dft/rks.py
@@ -37,7 +37,6 @@
 from pyscf.pbc.dft import numint
 from pyscf.dft import rks as mol_ks
 from pyscf.pbc.dft import multigrid
-from pyscf.pbc.df.df import GDF
 from pyscf.pbc.lib.kpts import KPoints
 from pyscf import __config__
 
@@ -111,13 +110,13 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
             vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpts_band)
             vk *= hyb
         elif alpha == 0: # LR=0, only SR exchange
-            vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
             vk = ks.get_k(cell, dm, hermi, kpt, kpts_band, omega=-omega)
             vk *= hyb
-        elif hyb == 0: # SR=0, only LR exchange
             vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
+        elif hyb == 0: # SR=0, only LR exchange
             vk = ks.get_k(cell, dm, hermi, kpt, kpts_band, omega=omega)
             vk *= alpha
+            vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
         else: # SR and LR exchange with different ratios
             vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpts_band)
             vk *= hyb
@@ -231,15 +230,6 @@ def build(self, cell=None):
             if not numpy.all(self.rsjk.kpts == self.kpt):
                 self.rsjk = self.rsjk.__class__(cell, kpts)
 
-        # for GDF and MDF
-        with_df = self.with_df
-        if (isinstance(with_df, GDF) and
-            self._numint.libxc.is_hybrid_xc(self.xc) and
-            len(kpts) > 1 and getattr(with_df, '_j_only', False)):
-            logger.warn(self, 'df.j_only cannot be used with hybrid functional')
-            self.with_df._j_only = False
-            self.with_df.reset()
-
         if self.verbose >= logger.WARN:
             self.check_sanity()
         return self
diff --git a/pyscf/pbc/dft/uks.py b/pyscf/pbc/dft/uks.py
index b1a61cb41b..313cd78c86 100644
--- a/pyscf/pbc/dft/uks.py
+++ b/pyscf/pbc/dft/uks.py
@@ -94,23 +94,17 @@ def get_veff(ks, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
         vxc += vj
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=cell.spin)
-        if getattr(ks.with_df, '_j_only', False):  # for GDF and MDF
-            logger.warn(ks, 'df.j_only cannot be used with hybrid functional')
-            ks.with_df._j_only = False
-            # Rebuild df object due to the change of parameter _j_only
-            if ks.with_df._cderi is not None:
-                ks.with_df.build()
         if omega == 0:
             vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpts_band)
             vk *= hyb
         elif alpha == 0: # LR=0, only SR exchange
-            vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
             vk = ks.get_k(cell, dm, hermi, kpt, kpts_band, omega=-omega)
             vk *= hyb
-        elif hyb == 0: # SR=0, only LR exchange
             vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
+        elif hyb == 0: # SR=0, only LR exchange
             vk = ks.get_k(cell, dm, hermi, kpt, kpts_band, omega=omega)
             vk *= alpha
+            vj = ks.get_j(cell, dm, hermi, kpt, kpts_band)
         else: # SR and LR exchange with different ratios
             vj, vk = ks.get_jk(cell, dm, hermi, kpt, kpts_band)
             vk *= hyb


Test patch:
diff --git a/pyscf/pbc/df/test/test_df.py b/pyscf/pbc/df/test/test_df.py
index 20036dee82..72dfdbf59c 100644
--- a/pyscf/pbc/df/test/test_df.py
+++ b/pyscf/pbc/df/test/test_df.py
@@ -15,6 +15,7 @@
 import unittest
 import tempfile
 import numpy
+import numpy as np
 from pyscf import lib
 import pyscf.pbc
 from pyscf import ao2mo, gto
@@ -266,6 +267,23 @@ def test_kpoints_input(sef):
         mydf = mdf.MDF(cell, kpts=kpts)
         assert mydf.kpts.shape == (8,3)
 
+    # issue 2790: build full-range J and SR K separately
+    def test_rsh_df(self):
+        cell = pgto.M(
+            a = np.eye(3) * 3,
+            atom = '''H 0.0000 0.0000 0.0000; H 1.3575 1.3575 1.3575''',
+            basis = [[0, [.5, 1]]],
+        )
+        kpts = cell.make_kpts([2,1,1])
+        kmf = cell.KRKS(kpts=kpts, xc='hse06').density_fit().run()
+        assert abs(kmf.e_tot - -0.687018457218418) < 1e-8
+
+        kmf = cell.KRKS(kpts=kpts, xc='camb3lyp').density_fit().run()
+        assert abs(kmf.e_tot - -0.674692142275221) < 1e-8
+
+        kmf = cell.KRKS(kpts=kpts, xc='wb97').density_fit().run()
+        assert abs(kmf.e_tot - -0.678851816639354) < 1e-8
+
 if __name__ == '__main__':
     print("Full Tests for df")
     unittest.main()
diff --git a/pyscf/pbc/dft/test/test_uks.py b/pyscf/pbc/dft/test/test_uks.py
index b78b408307..b50a7499c9 100644
--- a/pyscf/pbc/dft/test/test_uks.py
+++ b/pyscf/pbc/dft/test/test_uks.py
@@ -82,12 +82,12 @@ def test_rsh_df(self):
         mf = pbcdft.UKS(cell).density_fit()
         mf.xc = 'hse06'
         mf.kernel()
-        self.assertAlmostEqual(mf.e_tot, -2.4824271111708662, 6)
+        self.assertAlmostEqual(mf.e_tot, -2.482426252, 7)
 
         mf.xc = 'camb3lyp'
         mf.omega = .15
         mf.kernel()
-        self.assertAlmostEqual(mf.e_tot, -2.4766238116030683, 6)
+        self.assertAlmostEqual(mf.e_tot, -2.476623986, 7)
 
 
 if __name__ == '__main__':