Bugs in libxc.define_xc_ (fix #2551) (pyscf/pyscf#2574)

Fix `dft.libxc.define_xc_` and update `examples/dft/24-define_xc_functional.py` (fix #2551)

pyscf/examples/dft/24-define_xc_functional.py is not working

I'm trying to run examples of pyscf, and I figured out that dft/24-define_xc_functional.py is not working

Below is my error message.

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[87], line 63
     61 mf = mf.define_xc_(eval_xc, 'GGA', hyb=hybrid_coeff)
     62 mf.verbose = 4
---> 63 mf.kernel()

File <string>:2, in kernel(self, dm0, **kwargs)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/scf/hf.py:1793, in SCF.scf(self, dm0, **kwargs)
   1788 self.build(self.mol)
   1790 if self.max_cycle > 0 or self.mo_coeff is None:
   1791     self.converged, self.e_tot, \
   1792             self.mo_energy, self.mo_coeff, self.mo_occ = \
-> 1793             kernel(self, self.conv_tol, self.conv_tol_grad,
   1794                    dm0=dm0, callback=self.callback,
   1795                    conv_check=self.conv_check, **kwargs)
   1796 else:
   1797     # Avoid to update SCF orbitals in the non-SCF initialization
   1798     # (issue #495).  But run regular SCF for initial guess if SCF was
   1799     # not initialized.
   1800     self.e_tot = kernel(self, self.conv_tol, self.conv_tol_grad,
   1801                         dm0=dm0, callback=self.callback,
   1802                         conv_check=self.conv_check, **kwargs)[1]

File ~/miniconda3/lib/python3.11/site-packages/pyscf/scf/hf.py:126, in kernel(mf, conv_tol, conv_tol_grad, dump_chk, dm0, callback, conv_check, **kwargs)
    123     dm = dm0
    125 h1e = mf.get_hcore(mol)
--> 126 vhf = mf.get_veff(mol, dm)
    127 e_tot = mf.energy_tot(dm, h1e, vhf)
    128 logger.info(mf, 'init E= %.15g', e_tot)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/rks.py:80, in get_veff(ks, mol, dm, dm_last, vhf_last, hermi)
     78 else:
     79     max_memory = ks.max_memory - lib.current_memory()[0]
---> 80     n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
     81     logger.debug(ks, 'nelec by numeric integration = %s', n)
     82     if ks.do_nlc():

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:1158, in nr_rks(ni, mol, grids, xc_code, dms, relativity, hermi, max_memory, verbose)
   1156 elif xctype == 'GGA':
   1157     ao_deriv = 1
-> 1158     for i, ao, mask, wv in block_loop(ao_deriv):
   1159         wv[0] *= .5  # *.5 because vmat + vmat.T at the end
   1160         aow = _scale_ao_sparse(ao[:4], wv[:4], mask, ao_loc, out=aow)

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:1138, in nr_rks.<locals>.block_loop(ao_deriv)
   1136 for i in range(nset):
   1137     rho = make_rho(i, ao, mask, xctype)
-> 1138     exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype)[:2]
   1139     if xctype == 'LDA':
   1140         den = rho * weight

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/numint.py:2696, in LibXCMixin.eval_xc_eff(self, xc_code, rho, deriv, omega, xctype, verbose)
   2693 else:
   2694     spin = 0
-> 2696 out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
   2697 evfk = [out[0]]
   2698 for order in range(1, deriv+1):

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/libxc.py:1610, in eval_xc1(xc_code, rho, spin, deriv, omega)
   1606 def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
   1607     '''Similar to eval_xc.
   1608     Returns an array with the order of derivatives following xcfun convention.
   1609     '''
-> 1610     out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
   1611     xctype = xc_type(xc_code)
   1612     if deriv <= 1:

File ~/miniconda3/lib/python3.11/site-packages/pyscf/dft/libxc.py:1665, in _eval_xc(xc_code, rho, spin, deriv, omega)
   1663 nvar, xlen = xc_deriv._XC_NVAR[xctype, spin]
   1664 ngrids = rho.shape[-1]
-> 1665 rho = rho.reshape(spin+1,nvar,ngrids)
   1666 outlen = lib.comb(xlen+deriv, deriv)
   1667 out = numpy.zeros((outlen,ngrids))

ValueError: cannot reshape array of size 120224 into shape (1,1,30056)
```

And I also trying to run libxc.py's example of "define_xc_" [function](https://github.com/pyscf/pyscf/blob/c58968907adc96bc748591c9d72f1bc2c8ee66dc/pyscf/dft/libxc.py#L1123)

```python
def eval_xc(xc_code, rho, *args, **kwargs):
    exc = 0.01 * rho**2
    vrho = 0.01 * 3 * rho**2
    vxc = (vrho, None, None, None)
    fxc = None  # 2nd order functional derivative
    kxc = None  # 3rd order functional derivative
    return exc, vxc, fxc, kxc
define_xc_(mf._numint, eval_xc, xctype='LDA')
mf.kernel()

>>converged SCF energy = -75.8511902095329
```

but the code's results are same with 'SVWN' functional (default of pyscf).

```python
mf = dft.RKS(mol)
mf.xc = 'SVWN'
mf.kernel()

>>converged SCF energy = -75.8511902095329
```
I think there are some problems in running scf with self-made xc functional.

### Minimal API referenced by tests (typed signatures)
_Added only when tests call brand‑new symbols introduced by the fix; types reflect annotations present in the repo or `Unknown` if not annotated._

- Function in `examples/dft/24-define_xc_functional.py`: `def eval_mgga_xc(xc_code, rho, spin, relativity, deriv, omega, verbose) -> numpy.ndarray: ...`
  **Inputs:** xc_code: str, rho: numpy.ndarray, spin: int, relativity: int, deriv: int, omega: float, verbose: bool
  **Output:** numpy.ndarray

```python
def eval_mgga_xc(xc_code, rho, spin, relativity, deriv, omega, verbose) -> numpy.ndarray: ...
```

diff --git a/examples/dft/24-define_xc_functional.py b/examples/dft/24-define_xc_functional.py
index 72c7f25cfd..d5e2986f53 100644
--- a/examples/dft/24-define_xc_functional.py
+++ b/examples/dft/24-define_xc_functional.py
@@ -12,6 +12,7 @@
   functionals provided by Libxc or XcFun library.
 '''

+import numpy as np
 from pyscf import gto
 from pyscf import dft

@@ -25,28 +26,19 @@
 # half-half exact exchange and GGA functional
 hybrid_coeff = 0.5

-def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
+def eval_gga_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
     # A fictitious XC functional to demonstrate the usage
-    rho0, dx, dy, dz = rho[:4]
+    rho0, dx, dy, dz = rho
     gamma = (dx**2 + dy**2 + dz**2)
     exc = .01 * rho0**2 + .02 * (gamma+.001)**.5
     vrho = .01 * 3 * rho0**2 + .02 * (gamma+.001)**.5
     vgamma = .02 * .5 * (gamma+.001)**(-.5)
-    vlapl = None
-    vtau = None
-    vxc = (vrho, vgamma, vlapl, vtau)
+    vxc = (vrho, vgamma)
     v2rho2 = 0.01 * 6 * rho0
-    v2rhosigma = .02 * .5 * (gamma+.001)**(-.5)
+    v2rhosigma = np.zeros(gamma.shape)
     v2sigma2 = 0.02 * .5 * -.5 * (gamma+.001)**(-1.5)
-    v2lapl2 = None
-    vtau2 = None
-    v2rholapl = None
-    v2rhotau = None
-    v2lapltau = None
-    v2sigmalapl = None
-    v2sigmatau = None
     # 2nd order functional derivative
-    fxc = (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)
+    fxc = (v2rho2, v2rhosigma, v2sigma2)
     kxc = None  # 3rd order functional derivative

     # Mix with existing functionals
@@ -58,15 +50,42 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
     return exc, vxc, fxc, kxc

 mf = dft.RKS(mol)
-mf = mf.define_xc_(eval_xc, 'GGA', hyb=hybrid_coeff)
+mf = mf.define_xc_(eval_gga_xc, 'GGA', hyb=hybrid_coeff)
 mf.verbose = 4
 mf.kernel()

+def eval_mgga_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
+    # A fictitious XC functional to demonstrate the usage
+    rho0, dx, dy, dz, tau = rho
+    gamma = (dx**2 + dy**2 + dz**2)
+    exc = .01 * rho0**2 + .02 * (gamma+.001)**.5 +  0.01 * tau**2
+    vrho = .01 * 3 * rho0**2 + .02 * (gamma+.001)**.5
+    vgamma = .02 * .5 * (gamma+.001)**(-.5)
+    vtau = 0.02 * tau
+    vxc = (vrho, vgamma, vtau)
+    v2rho2 = 0.01 * 6 * rho0
+    v2rhosigma = np.zeros(gamma.shape)
+    v2sigma2 = 0.02 * .5 * -.5 * (gamma+.001)**(-1.5)
+    v2tau2 = np.full(tau.shape, 0.02)
+    v2rhotau = np.zeros(tau.shape)
+    v2sigmatau = np.zeros(tau.shape)
+    # 2nd order functional derivative
+    fxc = (v2rho2, v2rhosigma, v2sigma2, v2tau2, v2rhotau, v2sigmatau)
+    kxc = None  # 3rd order functional derivative
+
+    # Mix with existing functionals
+    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho[:4], spin, relativity, deriv, verbose)
+    exc += pbe_xc[0] * 0.5
+    vrho += pbe_xc[1][0] * 0.5
+    vgamma += pbe_xc[1][1] * 0.5
+    # The output follows the libxc.eval_xc API convention
+    return exc, vxc, fxc, kxc
+
 # half exact exchange in which 40% of the exchange is computed with short
 # range part of the range-separation Coulomb operator (omega = 0.8)
 beta = 0.2
 rsh_coeff = (0.8, hybrid_coeff-beta, beta)
 mf = dft.RKS(mol)
-mf = mf.define_xc_(eval_xc, 'GGA', rsh=rsh_coeff)
+mf = mf.define_xc_(eval_mgga_xc, 'MGGA', rsh=rsh_coeff)
 mf.verbose = 4
 mf.kernel()
diff --git a/pyscf/dft/libxc.py b/pyscf/dft/libxc.py
index 2d703479d3..4f67f3c652 100644
--- a/pyscf/dft/libxc.py
+++ b/pyscf/dft/libxc.py
@@ -770,7 +770,7 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
           | vtau[:,2]   = (u, d)

         * fxc for restricted case:
-          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)
+          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, v2tau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)

         * fxc for unrestricted case:
           | v2rho2[:,3]     = (u_u, u_d, d_d)
@@ -1186,16 +1186,33 @@ def _eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
         outlen = lib.comb(xlen+deriv, deriv)
         exc, vxc, fxc, kxc = libxc_out[:4]
         out = [exc]
-        if vxc is not None:
+        if deriv > 0:
+            assert vxc is not None
             out.extend([x for x in vxc if x is not None])
-        if fxc is not None:
-            out.extend([fxc[i] for i in [0, 1, 2, 6, 4, 9]])
-        if kxc is not None:
+        if deriv > 1:
+            assert fxc is not None
+            if xctype == 'GGA':
+                assert len(fxc) == 3, 'fxc for GGA should be arranged as (v2rho2, v2rhosigma, v2sigma2)'
+            elif xctype == 'MGGA':
+                if len(fxc) == 10:
+                    fxc = [fxc[i] for i in [0, 1, 2, 6, 4, 9]]
+                else:
+                    assert len(fxc) == 6, (
+                        'fxc for MGGA should be arranged as\n'
+                        '(v2rho2, v2rhosigma, v2sigma2, v2tau2, v2rhotau, v2sigmatau)\nor\n'
+                        '(v2rho2, v2rhosigma, v2sigma2, v2lapl2, v2tau2, '
+                        'v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)')
+            assert all(x is not None for x in fxc)
+            out.extend(fxc)
+        if deriv > 2:
+            assert kxc is not None
             out.extend([x for x in kxc if x is not None])
+
         if spin == 1:
             # Returns of eval_xc are structured as [grid_id,deriv_component]
             # for each term in libxc_out. Change the shape to [deriv_comp, grid_id]
             out = [x.T for x in out]
+
         out = numpy.vstack(out)[:outlen]
         assert len(out) == outlen
         idx = _libxc_to_xcfun_indices(xctype, spin, deriv)

diff --git a/pyscf/dft/test/test_libxc.py b/pyscf/dft/test/test_libxc.py
index 6dc223789d..04b28d4ded 100644
--- a/pyscf/dft/test/test_libxc.py
+++ b/pyscf/dft/test/test_libxc.py
@@ -218,7 +218,7 @@ def test_lyp(self):
     #    rho_b = numpy.array([[4.53272893e-06, 4.18968775e-06,-2.83034672e-06, 2.61832978e-06, 5.63360737e-06, 8.97541777e-07]]).T
     #    e, v = dft.libxc.eval_xc('tpss,', (rho_a, rho_b), spin=1, deriv=1)[:2]

-    def test_define_xc(self):
+    def test_define_xc_gga(self):
         def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
             # A fictitious XC functional to demonstrate the usage
             rho0, dx, dy, dz = rho[:4]
@@ -242,11 +242,19 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
         self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
         self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)

+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=2)
+
+        # Ensure the xc_code in the input has no impact to the result.
         n, exc, vxc = ni.nr_rks(mol, mf.grids, 'm06x', dm)
         self.assertAlmostEqual(n, 4, 5)
         self.assertAlmostEqual(lib.fp(exc), 0.01197588220700074, 6)
         self.assertAlmostEqual(lib.fp(vxc), -0.043974912389152986, 6)

+        ni = dft.libxc.define_xc(ni, eval_xc, 'MGGA')
+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=1)
+
         mf = mf.define_xc_('0.5*B3LYP+0.5*B3LYP')
         exc0, vxc0 = mf._numint.eval_xc(None, rho, 0, deriv=1)[:2]
         exc1, vxc1 = dft.libxc.eval_xc('0.5*B3LYP+0.5*B3LYP', rho, 0, deriv=1)[:2]
@@ -263,6 +271,33 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
         self.assertAlmostEqual(abs(vxc0[0]-vxc1[0]).max(), 0, 9)
         self.assertAlmostEqual(abs(vxc0[1]-vxc1[1]).max(), 0, 9)

+    def test_define_xc_mgga(self):
+        def eval_mgga_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
+            # A fictitious XC functional to demonstrate the usage
+            rho0, dx, dy, dz, tau = rho
+            gamma = (dx**2 + dy**2 + dz**2)
+            exc = .01 * rho0**2 + .02 * (gamma+.001)**.5 + .01 * tau**2
+            vrho = .01 * 2 * rho0
+            vgamma = .02 * .5 * (gamma+.001)**(-.5)
+            vtau = .02 * tau
+            vxc = (vrho, vgamma, vtau)
+            zeros = numpy.zeros_like(rho0)
+            fxc = (zeros, zeros, zeros)
+            kxc = None  # 3rd order functional derivative
+            return exc, vxc, fxc, kxc
+
+        mf = dft.RKS(mol)
+        ni = dft.libxc.define_xc(mf._numint, eval_mgga_xc, 'MGGA')
+        numpy.random.seed(1)
+        rho = numpy.random.random((5,10))
+        exc, vxc = ni.eval_xc(None, rho, 0, deriv=1)[:2]
+        self.assertAlmostEqual(lib.fp(exc), 0.016353861820705, 9)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)
+
+        with self.assertRaises(AssertionError):
+            ni.eval_xc_eff(None, rho, deriv=2)
+
     def test_m05x(self):
         rho =(numpy.array([1., 1., 0., 0., 0., 0.165 ]).reshape(-1,1),
               numpy.array([.8, 1., 0., 0., 0., 0.1050]).reshape(-1,1))


Fix GTH-PP integrals (fix #2575) (pyscf/pyscf#2577)

Possible bug in GTH PP

MWE
```
from pyscf.pbc import gto
import numpy as np

cell= gto.Cell(atom='''Cu .0 .0 .0''',
               a=np.array([[0.5, 0.5, 0], [0, 0.5, 0.5], [0.5, 0, 0.5]])* 3.0,
               basis='gth-dzvp-molopt-sr',
               pseudo='gth-pbe-q11',
               verbose=4
              ).build()
kpts = cell.make_kpts([2,2,2])
mf = cell.KRKS(xc='pbe0', kpts=kpts).density_fit()
mf.max_cycle=2
mf.run()
```
traceback
```
  File "/share/home/srwang/pbc/gthpp.py", line 15, in <module>
    mf.run()
  File "/share/home/srwang/pyscf2/pyscf/lib/misc.py", line 638, in run
    self.kernel(*args)
  File "<string>", line 2, in kernel
  File "/share/home/srwang/pyscf2/pyscf/scf/hf.py", line 2000, in scf
    kernel(self, self.conv_tol, self.conv_tol_grad,
  File "/share/home/srwang/pyscf2/pyscf/scf/hf.py", line 128, in kernel
    h1e = mf.get_hcore(mol)
  File "/share/home/srwang/pyscf2/pyscf/pbc/scf/khf.py", line 77, in get_hcore
    nuc = lib.asarray(mf.with_df.get_pp(kpts))
  File "/share/home/srwang/pyscf2/pyscf/pbc/df/df.py", line 393, in get_pp
    vpp = dfbuilder.get_pp()
  File "/share/home/srwang/pyscf2/pyscf/pbc/df/rsdf_builder.py", line 1031, in get_pp
    vpp += pp2builder.get_pp_loc_part2()
ValueError: operands could not be broadcast together with shapes (8,25,25) (8,) (8,25,25)
```

diff --git a/pyscf/pbc/df/aft.py b/pyscf/pbc/df/aft.py
index 1440cc7f63..87a1d323a3 100644
--- a/pyscf/pbc/df/aft.py
+++ b/pyscf/pbc/df/aft.py
@@ -329,7 +329,8 @@ def get_pp_loc_part2(self):
             else:
                 lib.logger.warn(cell, 'cell.pseudo was specified but its elements %s '
                                 'were not found in the system.', cell._pseudo.keys())
-            vpploc = [0] * nkpts
+            nao = cell.nao
+            vpploc = np.zeros((nkpts, nao, nao))
             return vpploc

         rcut = self._estimate_rcut_3c1e(rs_cell, fake_cells)
diff --git a/pyscf/tdscf/_lr_eig.py b/pyscf/tdscf/_lr_eig.py
index 28beab7ffd..0eaa964bf5 100644
--- a/pyscf/tdscf/_lr_eig.py
+++ b/pyscf/tdscf/_lr_eig.py
@@ -524,10 +524,7 @@ def real_eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=Non
     assert pick is None
     assert callable(precond)

-    if isinstance(verbose, logger.Logger):
-        log = verbose
-    else:
-        log = logger.Logger(sys.stdout, verbose)
+    log = logger.new_logger(verbose)

     assert x0.ndim == 2
     A_size = x0.shape[1] // 2

diff --git a/pyscf/pbc/df/test/test_aft.py b/pyscf/pbc/df/test/test_aft.py
index f637883643..c3cfc14e7d 100644
--- a/pyscf/pbc/df/test/test_aft.py
+++ b/pyscf/pbc/df/test/test_aft.py
@@ -352,6 +352,20 @@ def test_aft_get_pp(self):
         self.assertAlmostEqual(abs(v0 - v3).max(), 0, 8)
         self.assertAlmostEqual(lib.fp(v0) - (-5.7070290795125445-0.00038722541238732697j), 0, 8)

+    # issue 2575
+    def test_aft_get_pp1(self):
+        cell = pgto.M(atom='Cu .0 .0 .0', a=np.eye(3)*3, spin=1,
+               basis=[[1, [1, 1]]], pseudo='''Cu
+1    0   10
+0.53    0
+1
+0.26    1   1.8''')
+        kpts = cell.make_kpts([2,1,1])
+        ref = fft.FFTDF(cell, kpts).get_pp()
+        v = aft.AFTDF(cell, kpts).get_pp()
+        self.assertAlmostEqual(abs(v - ref).max(), 0, 8)
+        self.assertAlmostEqual(lib.fp(v), 0.37494551685809624, 8)
+
     def test_aft_get_nuc(self):
         v0 = fft.FFTDF(cell, kpts[0]).get_nuc()
         v1 = aft.AFTDF(cell, kpts[0]).get_nuc()


Fix dimension bug in spinor X2C code (pyscf/pyscf#2611)

spin-orbit contracted basis set

In PySCF's basis set internal format, the kappa value can be used to treat spin-orbit contracted basis set.

> kappa can have value
>  (corresponding to spinors with
> ),
>  (corresponding to spinors with
> ) or 0. When kappa is 0, both types of spinors are assumed in the basis.

```
[[angular, kappa, [[exp, c_1, c_2, ..],
                   [exp, c_1, c_2, ..],
                   ... ]],
 [angular, kappa, [[exp, c_1, c_2, ..],
                   [exp, c_1, c_2, ..]
                   ... ]]]
```
There is a possible bug in treating contracted basis set with contraction coefficients for different kappa. The following example with a fake basis set of this kind.
```
from pyscf import gto
from pyscf import scf

mol = gto.M(atom='''
O 0  0  0
H 0  1  1
H 0 -1 -1''',
            basis={
                'O':
                [[
                    0, [61420.0, 9e-05, -2e-05], [9199.0, 0.000698, -0.000159],
                    [2091.0, 0.003664, -0.000829],
                    [590.9, 0.015218, -0.003508], [192.3, 0.052423, -0.012156],
                    [69.32, 0.145921, -0.036261], [26.97, 0.305258, -0.082992],
                    [11.1, 0.398508, -0.15209], [4.682, 0.21698, -0.115331]
                ], [0, [0.5547, 1.0]],
                 [1, -2, [63.42, 0.006044], [14.66, 0.041799], [4.459, 0.1611]],
                 [1, 1, [34.46, 0.015928], [7.749, 0.09974], [2.28, 0.31049]],
                 [1, 1, [1.531, 1.0]], [1, 0, [0.5302, 1.0]]],
                'H':
                '631g'
            })
print(mol._basis)
print(mol.nao_nr())
print(mol.nao_2c())
mf = scf.X2C(mol)
mf.kernel()
```
The calculation won't run through.

diff --git a/examples/x2c/01-x2c1e.py b/examples/x2c/01-x2c1e.py
new file mode 100644
index 0000000000..a2f9061c27
--- /dev/null
+++ b/examples/x2c/01-x2c1e.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+#
+# Author: Qiming Sun <osirpt.sun@gmail.com>
+#
+
+'''
+Constructing the X matrix using one-electron approximation. This is the default
+configuration of the X2C module.
+'''
+
+import pyscf
+from pyscf.x2c import UHF as X2C_UHF
+from pyscf.x2c import dft as x2c_dft
+
+mol = pyscf.M(
+    verbose = 0,
+    atom = [["O" , (0. , 0.     , 0.)],
+            [1   , (0. , -0.757 , 0.587)],
+            [1   , (0. , 0.757  , 0.587)] ],
+    basis = 'ccpvdz-dk',
+)
+
+method = mol.RHF()
+enr = method.kernel()
+print('E(NR) = %.12g' % enr)
+
+method = X2C_UHF(mol)
+ex2c = method.kernel()
+print('E(X2C1E-UHF) = %.12g' % ex2c)
+
+# Customizing basis sets for X matrix
+method.with_x2c.basis = {'O': 'unc-ccpvqz', 'H':'unc-ccpvdz'}
+print('E(X2C1E-UHF) = %.12g' % method.kernel())
+
+# Constructing X matrix using single center approximation
+method.with_x2c.approx = 'atom1e'
+print('E(X2C1E-UHF) = %.12g' % method.kernel())
+
+method = x2c_dft.UKS(mol).run()
+print('E(X2C1E-UKS) = %.12g' % method.e_tot)
diff --git a/pyscf/hessian/uks.py b/pyscf/hessian/uks.py
index 74bdb85bac..dfb4592edc 100644
--- a/pyscf/hessian/uks.py
+++ b/pyscf/hessian/uks.py
@@ -41,7 +41,7 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     mf = hessobj.base
     ni = mf._numint
     if mf.do_nlc():
-        raise NotImplementedError('RKS Hessian for NLC functional')
+        raise NotImplementedError('UKS Hessian for NLC functional')

     if mo_energy is None: mo_energy = mf.mo_energy
     if mo_occ is None:    mo_occ = mf.mo_occ
diff --git a/pyscf/x2c/x2c.py b/pyscf/x2c/x2c.py
index d1b0e34ed6..1754f3cd7c 100644
--- a/pyscf/x2c/x2c.py
+++ b/pyscf/x2c/x2c.py
@@ -61,7 +61,10 @@ def get_xmol(self, mol=None):
             xmol.build(False, False, basis=self.basis)
             return xmol, None
         elif self.xuncontract:
-            xmol, contr_coeff = _uncontract_mol(mol, self.xuncontract)
+            if all(mol._bas[:,mole.KAPPA_OF] == 0):
+                xmol, contr_coeff = _uncontract_mol(mol, self.xuncontract)
+            else:
+                xmol, contr_coeff = _decontract_spinor(mol, self.xuncontract)
             return xmol, contr_coeff
         else:
             return mol, None
@@ -113,9 +116,13 @@ def get_hcore(self, mol=None):
             h1 = reduce(numpy.dot, (c.T.conj(), h1, c))
         elif self.xuncontract:
             np, nc = contr_coeff_nr.shape
-            contr_coeff = numpy.zeros((np*2,nc*2))
-            contr_coeff[0::2,0::2] = contr_coeff_nr
-            contr_coeff[1::2,1::2] = contr_coeff_nr
+            if np == h1.shape[0]:
+                # spinor adapted contraction matrix generated by _decontract_spinor
+                contr_coeff = contr_coeff_nr
+            else:
+                contr_coeff = numpy.zeros((np*2,nc*2))
+                contr_coeff[0::2,0::2] = contr_coeff_nr
+                contr_coeff[1::2,1::2] = contr_coeff_nr
             h1 = reduce(numpy.dot, (contr_coeff.T.conj(), h1, contr_coeff))
         return h1

@@ -986,30 +993,89 @@ def _sigma_dot(mat):
     nao = mat.shape[-1] * 2
     return lib.einsum('sxy,spq->xpyq', quaternion, mat).reshape(nao, nao)

-
-if __name__ == '__main__':
-    from pyscf.x2c import dft
-    mol = mole.Mole()
-    mol.build(
-        verbose = 0,
-        atom = [["O" , (0. , 0.     , 0.)],
-                [1   , (0. , -0.757 , 0.587)],
-                [1   , (0. , 0.757  , 0.587)] ],
-        basis = 'ccpvdz-dk',
-    )
-
-    method = hf.RHF(mol)
-    enr = method.kernel()
-    print('E(NR) = %.12g' % enr)
-
-    method = UHF(mol)
-    ex2c = method.kernel()
-    print('E(X2C1E) = %.12g' % ex2c)
-    method.with_x2c.basis = {'O': 'unc-ccpvqz', 'H':'unc-ccpvdz'}
-    print('E(X2C1E) = %.12g' % method.kernel())
-    method.with_x2c.approx = 'atom1e'
-    print('E(X2C1E) = %.12g' % method.kernel())
-
-    method = dft.UKS(mol)
-    ex2c = method.kernel()
-    print('E(X2C1E) = %.12g' % ex2c)
+def _decontract_spinor(mol, atoms=None):
+    pmol = mol.copy(deep=False)
+
+    # Some input basis may be segmented basis from a general contracted set.
+    # This may lead to duplicated pGTOs. First contract all basis to remove
+    # duplicated primitive functions.
+    bas_exps = mol.bas_exps()
+    _bas = []
+    env = [mol._env.copy()]
+    contr_coeff = []
+    pexp = env[0].size
+    ao_loc = mol.ao_loc_2c()
+
+    aoslices = mol.aoslice_by_atom()
+    for ia, (ib0, ib1) in enumerate(aoslices[:,:2]):
+        if ib0 == ib1: # No basis on atom ia
+            continue
+
+        if atoms is not None:
+            if isinstance(atoms, str):
+                to_apply = ((atoms == mol.atom_pure_symbol(ia)) or
+                            (atoms == mol.atom_symbol(ia)))
+            elif isinstance(atoms, (tuple, list)):
+                to_apply = ((mol.atom_pure_symbol(ia) in atoms) or
+                            (mol.atom_symbol(ia) in atoms) or
+                            (ia in atoms))
+            else:
+                to_apply = True
+            if not to_apply:
+                _bas.append(mol._bas[ib0:ib1])
+                contr_coeff.append(numpy.eye(ao_loc[ib1] - ao_loc[ib0]))
+                continue
+
+        lmax = mol._bas[ib0:ib1,mole.ANG_OF].max()
+        for l in range(lmax+1):
+            bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,mole.ANG_OF] == l)[0]
+            if len(bas_idx) == 0:
+                continue
+            if bas_idx[0] + len(bas_idx) != bas_idx[-1] + 1:
+                raise NotImplementedError('Discontinuous bases of same angular momentum')
+
+            es = numpy.hstack([bas_exps[i] for i in bas_idx])
+            _, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
+            cs = scipy.linalg.block_diag(
+                *[mol._libcint_ctr_coeff(i) for i in bas_idx])
+            nprim = len(e_idx)
+            if len(es) != nprim:
+                cs_new = numpy.zeros((nprim, cs.shape[1]))
+                for i, j in enumerate(rev_idx):
+                    cs_new[j] += cs[i]
+                es = es[e_idx][::-1]
+                cs = cs_new[::-1]
+
+            bs = numpy.zeros((nprim, mole.BAS_SLOTS), dtype=numpy.int32)
+            bs[:,mole.ATOM_OF] = ia
+            bs[:,mole.ANG_OF ] = l
+            bs[:,mole.NCTR_OF] = bs[:,mole.NPRIM_OF] = 1
+            bs[:,mole.PTR_EXP] = pexp + numpy.arange(nprim)
+            bs[:,mole.PTR_COEFF] = pexp + numpy.arange(nprim, nprim*2)
+            norm = mole.gto_norm(l, es)
+            env.append(es)
+            env.append(norm)
+            pexp += nprim * 2
+            _bas.append(bs)
+
+            degen = 4 * l + 2
+            c = numpy.einsum('pi,p,xm->pxim', cs, 1./norm, numpy.eye(degen))
+            c = c.reshape(nprim*degen, -1)
+            full_idx = numpy.arange(c.shape[1]).reshape(-1, degen)
+            spinor_idx = []
+            p0 = p1 = 0
+            for ib in bas_idx:
+                p0, p1 = p1, p1 + mol.bas_nctr(ib)
+                kappa = mol.bas_kappa(ib)
+                if kappa == 0:
+                    spinor_idx.append(full_idx[p0:p1].ravel())
+                elif kappa > 0:
+                    spinor_idx.append(full_idx[p0:p1,:l*2].ravel())
+                else:
+                    spinor_idx.append(full_idx[p0:p1,l*2:].ravel())
+            contr_coeff.append(c[:,numpy.hstack(spinor_idx)])
+
+    pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
+    pmol._env = numpy.hstack(env)
+    contr_coeff = scipy.linalg.block_diag(*contr_coeff)
+    return pmol, contr_coeff

diff --git a/pyscf/x2c/test/test_x2c.py b/pyscf/x2c/test/test_x2c.py
index ddca2b3206..8776d2029c 100644
--- a/pyscf/x2c/test/test_x2c.py
+++ b/pyscf/x2c/test/test_x2c.py
@@ -228,6 +228,23 @@ def test_undo_x2c(self):
         mf = mf.undo_x2c()
         self.assertEqual(mf.__class__.__name__, 'DFGHF')

+    # issue 2605
+    def test_kappa_spinor(self):
+        mol = gto.M(
+            atom='''He 0.   0.7  .1
+                    He 0.5 -0.2 -.1''',
+            basis=[[0, [0.5547, 1.]],
+                   [1, [11., 0.94, .31], [4.68, 0.22, .80]],
+                   [1, -2, [1.9, 1.]],
+                   [1, 1, [1.53, 1.]],
+                   [1, 0, [0.53, 1.]]],
+        )
+        pmol, c = x2c.SpinorX2CHelper(mol).get_xmol()
+        self.assertEqual(c.shape, (pmol.nao_2c(), mol.nao_2c()))
+        ref = mol.intor('int1e_ovlp_spinor')
+        dat = c.T.dot(pmol.intor('int1e_ovlp_spinor')).dot(c)
+        self.assertAlmostEqual(abs(ref - dat).max(), 0, 12)
+

 if __name__ == "__main__":
     print("Full Tests for x2c")


Complex density matrices for DFHF (pyscf/pyscf#2676)

[UNEXPECTED BEHAVIOR] Density-fitting complex RHF does not give correct energy

I notice a significant energy difference when testing df complex RHF. Here is my code
```
from pyscf import gto, scf
import numpy as np
mol = gto.Mole()
dista = 1.5
mol.atom = [['Be',(0.0, 0.0, 0.0)], ['H',(dista, 1.344-0.46*dista, 0.0)],['H',(dista, -1.344+0.46*dista, 0.0)]]
mol.spin = 0
mol.basis = 'cc-pvdz'
mol.build()
rrhf = scf.RHF(mol)

# real RHF
rrhf = scf.RHF(mol)
rrhf.kernel()

# cmplx RHF
crhf = scf.RHF(mol)
dm = rrhf.make_rdm1() + 0j
dm[0,:] += .1j
dm[:,0] -= .1j
crhf.kernel(dm)

# DF RHF
dfrhf = scf.RHF(mol).density_fit()
dfrhf.mo_coeff = rrhf.mo_coeff
dfrhf.mo_occ = rrhf.mo_occ
print("E(dfrrhf)=", dfrhf.energy_tot())

dfrhf.mo_coeff = crhf.mo_coeff
dfrhf.mo_occ = crhf.mo_occ
print("E(dfchf)=", dfrhf.energy_tot())
```
My output shows
```
E(rrhf) = -15.5335928065009
E(crhf) = -15.5335928064977
E(dfrhf) = -15.533610777908848
E(dfchf) = -14.751606976368922
```
The energies match for real orbitals but differ a lot for complex orbitals.

diff --git a/pyscf/df/df_jk.py b/pyscf/df/df_jk.py
index 9e5670537f..222eaceb53 100644
--- a/pyscf/df/df_jk.py
+++ b/pyscf/df/df_jk.py
@@ -258,6 +258,33 @@ def get_jk(dfobj, dm, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-13):
     vj = 0
     vk = numpy.zeros_like(dms)

+    if numpy.iscomplexobj(dms):
+        if with_j:
+            vj = numpy.zeros_like(dms)
+        max_memory = dfobj.max_memory - lib.current_memory()[0]
+        blksize = max(4, int(min(dfobj.blockdim, max_memory*.22e6/8/nao**2)))
+        buf = numpy.empty((blksize,nao,nao))
+        buf1 = numpy.empty((nao,blksize,nao))
+        for eri1 in dfobj.loop(blksize):
+            naux, nao_pair = eri1.shape
+            eri1 = lib.unpack_tril(eri1, out=buf)
+            if with_j:
+                tmp = numpy.einsum('pij,nji->pn', eri1, dms.real)
+                vj.real += numpy.einsum('pn,pij->nij', tmp, eri1)
+                tmp = numpy.einsum('pij,nji->pn', eri1, dms.imag)
+                vj.imag += numpy.einsum('pn,pij->nij', tmp, eri1)
+            buf2 = numpy.ndarray((nao,naux,nao), buffer=buf1)
+            for k in range(nset):
+                buf2[:] = lib.einsum('pij,jk->ipk', eri1, dms[k].real)
+                vk[k].real += lib.einsum('ipk,pkj->ij', buf2, eri1)
+                buf2[:] = lib.einsum('pij,jk->ipk', eri1, dms[k].imag)
+                vk[k].imag += lib.einsum('ipk,pkj->ij', buf2, eri1)
+            t1 = log.timer_debug1('jk', *t1)
+        if with_j: vj = vj.reshape(dm_shape)
+        if with_k: vk = vk.reshape(dm_shape)
+        logger.timer(dfobj, 'df vj and vk', *t0)
+        return vj, vk
+
     if with_j:
         idx = numpy.arange(nao)
         dmtril = lib.pack_tril(dms + dms.conj().transpose(0,2,1))
@@ -322,6 +349,7 @@ def get_jk(dfobj, dm, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-13):
         buf = numpy.empty((2,blksize,nao,nao))
         for eri1 in dfobj.loop(blksize):
             naux, nao_pair = eri1.shape
+            assert (nao_pair == nao*(nao+1)//2)
             if with_j:
                 # uses numpy.matmul
                 vj += dmtril.dot(eri1.T).dot(eri1)
@@ -388,6 +416,7 @@ def get_j(dfobj, dm, hermi=0, direct_scf_tol=1e-13):
     opt = dfobj._vjopt
     fakemol = opt.fakemol
     dm = numpy.asarray(dm, order='C')
+    assert dm.dtype == numpy.float64
     dm_shape = dm.shape
     nao = dm_shape[-1]
     dm = dm.reshape(-1,nao,nao)

diff --git a/pyscf/df/test/test_df_jk.py b/pyscf/df/test/test_df_jk.py
index e418a70204..d426430ae3 100644
--- a/pyscf/df/test/test_df_jk.py
+++ b/pyscf/df/test/test_df_jk.py
@@ -194,6 +194,18 @@ def test_get_j(self):
         self.assertAlmostEqual(abs(vj0-vj1).max(), 0, 12)
         self.assertAlmostEqual(lib.fp(vj0), -194.15910890730052, 9)

+    def test_df_jk_complex_dm(self):
+        mol = gto.M(atom='H 0 0 0; H 0 0 1')
+        mf = mol.RHF().run()
+        dm = mf.make_rdm1() + 0j
+        dm[0,:] += .1j
+        dm[:,0] -= .1j
+        mf.kernel(dm)
+        self.assertTrue(mf.mo_coeff.dtype == numpy.complex128)
+        dfmf = mf.density_fit()
+        self.assertAlmostEqual(dfmf.energy_tot(), -1.0661355663696201, 9)
+        self.assertAlmostEqual(dfmf.energy_tot(), mf.e_tot, 3)
+

 if __name__ == "__main__":
     print("Full Tests for df")


When building RSH K integrals, cderi is overwritten (fix #2790) (pyscf/pyscf#2797)

Causing a KeyError when reading cderi file

[BUG REPORT] RSH density fitting fails with j3c-lr KeyError

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

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

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__':


feat: Add spin-seperated 4-RDM for FCI  (pyscf/pyscf#2826)

Extends the functionality to obtain the spin-separated 4-RDM from an FCI object. Relies on the same hack as in #2357 for the 3-RDM. Implements the same tests and asserts (checking consistency with the spin-summed RDM) as in the aforementioned PR.

Enable spin-separated 3-RDMs

The spin-separated 3-RDMs (aaa, aab, abb, bbb components) have been implemented thanks to @MatthewRHermes's hack (see #2326, code here is largely lifted out of https://github.com/MatthewRHermes/mrh/blob/master/my_pyscf/lassi/op_o0.py#L54-L168). A suitable test (mcscf/04) have been expanded to assert that the spin-separated RDMs add (with correct transpositions) to the spin-summed RDMs.

### Minimal API referenced by tests (typed signatures)
_Added only when tests call brand‑new symbols introduced by the fix; types reflect annotations present in the repo or `Unknown` if not annotated._

- Function in `pyscf/fci/direct_spin1.py`: `def make_rdm1234(fcivec, norb, nelec, link_index, reorder) -> numpy.ndarray: ...`
  **Inputs:** fcivec: numpy.ndarray, norb: int, nelec: Union[int, tuple], link_index: numpy.ndarray, reorder: bool
  **Output:** numpy.ndarray

- Function in `pyscf/fci/direct_spin1.py`: `def make_rdm1234s(fcivec, norb, nelec, link_index, reorder) -> numpy.ndarray: ...`
  **Inputs:** fcivec: numpy.ndarray, norb: int, nelec: Union[int, tuple], link_index: numpy.ndarray, reorder: bool
  **Output:** numpy.ndarray

```python
def make_rdm1234(fcivec, norb, nelec, link_index, reorder) -> numpy.ndarray: ...
def make_rdm1234s(fcivec, norb, nelec, link_index, reorder) -> numpy.ndarray: ...
```

diff --git a/examples/fci/14-density_matrix.py b/examples/fci/14-density_matrix.py
index 28e207b918..52ceb78e1e 100644
--- a/examples/fci/14-density_matrix.py
+++ b/examples/fci/14-density_matrix.py
@@ -158,6 +158,37 @@
                                          (nelec_a,nelec_b))
 dm1, dm2, dm3, dm4 = fci.rdm.reorder_dm1234(dm1, dm2, dm3, dm4)

+#
+# Spin-separated 4-particle density matrix
+#
+
+(dm1a, dm1b), (dm2aa, dm2ab, dm2bb), (dm3aaa, dm3aab, dm3abb, dm3bbb), (dm4aaaa, dm4aaab, dm4aabb, dm4abbb, dm4bbbb) = (
+    fci.direct_spin1.make_rdm1234s(fcivec0, norb, (nelec_a, nelec_b))
+)
+assert(numpy.allclose(dm1a+dm1b, dm1))
+assert(numpy.allclose(dm2aa+dm2bb+dm2ab+dm2ab.transpose(2,3,0,1), dm2))
+assert(numpy.allclose(dm3aaa+dm3bbb+dm3aab+dm3aab.transpose(0,1,4,5,2,3)+\
+dm3aab.transpose(4,5,0,1,2,3)+dm3abb+dm3abb.transpose(2,3,0,1,4,5)+dm3abb.transpose(2,3,4,5,0,1), dm3))
+assert(numpy.allclose(
+    dm4aaaa
+    + dm4bbbb
+    + dm4aaab
+    + dm4aaab.transpose(0, 1, 2, 3, 6, 7, 4, 5)
+    + dm4aaab.transpose(0, 1, 6, 7, 2, 3, 4, 5)
+    + dm4aaab.transpose(6, 7, 0, 1, 2, 3, 4, 5)
+    + dm4aabb
+    + dm4aabb.transpose(0, 1, 4, 5, 2, 3, 6, 7)
+    + dm4aabb.transpose(4, 5, 0, 1, 2, 3, 6, 7)
+    + dm4aabb.transpose(0, 1, 4, 5, 6, 7, 2, 3)
+    + dm4aabb.transpose(4, 5, 0, 1, 6, 7, 2, 3)
+    + dm4aabb.transpose(4, 5, 6, 7, 0, 1, 2, 3)
+    + dm4abbb
+    + dm4abbb.transpose(2, 3, 0, 1, 4, 5, 6, 7)
+    + dm4abbb.transpose(2, 3, 4, 5, 0, 1, 6, 7)
+    + dm4abbb.transpose(2, 3, 4, 5, 6, 7, 0, 1),
+    dm4,
+))
+
 #
 # Spin-traced 4-particle transition density matrix
 #
diff --git a/pyscf/fci/direct_spin1.py b/pyscf/fci/direct_spin1.py
index 2f4d3777db..ea3bfbf9a5 100644
--- a/pyscf/fci/direct_spin1.py
+++ b/pyscf/fci/direct_spin1.py
@@ -410,6 +410,77 @@ def make_rdm123s(fcivec, norb, nelec, link_index=None, reorder=True):
     # rdm3aab.transpose(4,5,0,1,2,3)+rdm3abb+rdm3abb.transpose(2,3,0,1,4,5)+rdm3abb.transpose(2,3,4,5,0,1), rdm3)
     return (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb), (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb)

+def make_rdm1234(fcivec, norb, nelec, link_index=None, reorder=True):
+    '''Spin traced 1-, 2-, 3, 4-particle density matrices.'''
+    dm1, dm2, dm3, dm4 = rdm.make_dm1234('FCI4pdm_kern_sf', fcivec, fcivec, norb, nelec)
+    if reorder:
+        dm1, dm2, dm3, dm4 = rdm.reorder_dm1234(dm1, dm2, dm3, dm4, inplace=True)
+    return dm1, dm2, dm3, dm4
+
+def make_rdm1234s(fcivec, norb, nelec, link_index=None, reorder=True):
+    r'''Spin separated 1-, 2-, 3, 4-particle density matrices.
+
+    1pdm[p,q] = :math:`\langle q_\alpha^\dagger p_\alpha \rangle +
+                       \langle q_\beta^\dagger  p_\beta \rangle`;
+    2pdm[p,q,r,s] = :math:`\langle p_\alpha^\dagger r_\alpha^\dagger s_\alpha q_\alpha\rangle +
+                           \langle p_\beta^\dagger  r_\alpha^\dagger s_\alpha q_\beta\rangle +
+                           \langle p_\alpha^\dagger r_\beta^\dagger  s_\beta  q_\alpha\rangle +
+                           \langle p_\beta^\dagger  r_\beta^\dagger  s_\beta  q_\beta\rangle`.
+    '''
+    if (not reorder):
+        raise NotImplementedError('reorder=False not currently supported')
+    ci_spinless = civec_spinless_repr([fcivec,], norb, [nelec,])
+    rdm1, rdm2, rdm3, rdm4 = make_rdm1234(ci_spinless, norb*2, (nelec[0]+nelec[1],0))
+
+    # define slices
+    alpha = slice(0, norb)
+    beta = slice(norb, None)
+
+    rdm1a = rdm1[alpha, alpha]
+    rdm1b = rdm1[beta, beta]
+    # assert np.allclose(rdm1a+rdm1b, rdm1)
+
+    rdm2aa = rdm2[alpha, alpha, alpha, alpha]
+    rdm2ab = rdm2[alpha, alpha, beta, beta]
+    rdm2bb = rdm2[beta, beta, beta, beta]
+    # assert np.allclose(rdm2aa+rdm2bb+rdm2ab+rdm2ab.transpose(2,3,0,1), rdm2)
+
+    rdm3aaa = rdm3[alpha, alpha, alpha, alpha, alpha, alpha]
+    rdm3aab = rdm3[alpha, alpha, alpha, alpha, beta, beta]
+    rdm3abb = rdm3[alpha, alpha, beta, beta, beta, beta]
+    rdm3bbb = rdm3[beta, beta, beta, beta, beta, beta]
+    # assert np.allclose(rdm3aaa+rdm3bbb+rdm3aab+rdm3aab.transpose(0,1,4,5,2,3)+\
+    # rdm3aab.transpose(4,5,0,1,2,3)+rdm3abb+rdm3abb.transpose(2,3,0,1,4,5)+rdm3abb.transpose(2,3,4,5,0,1), rdm3)
+
+    rdm4aaaa = rdm4[alpha, alpha, alpha, alpha, alpha, alpha, alpha, alpha]
+    rdm4aaab = rdm4[alpha, alpha, alpha, alpha, alpha, alpha, beta, beta]
+    rdm4aabb = rdm4[alpha, alpha, alpha, alpha, beta, beta, beta, beta]
+    rdm4abbb = rdm4[alpha, alpha, beta, beta, beta, beta, beta, beta]
+    rdm4bbbb = rdm4[beta, beta, beta, beta, beta, beta, beta, beta]
+    # assert numpy.allclose(
+    #     rdm4aaaa
+    #     + rdm4bbbb
+    #     + rdm4aaab
+    #     + rdm4aaab.transpose(0, 1, 2, 3, 6, 7, 4, 5)
+    #     + rdm4aaab.transpose(0, 1, 6, 7, 2, 3, 4, 5)
+    #     + rdm4aaab.transpose(6, 7, 0, 1, 2, 3, 4, 5)
+    #     + rdm4aabb
+    #     + rdm4aabb.transpose(0, 1, 4, 5, 2, 3, 6, 7)
+    #     + rdm4aabb.transpose(0, 1, 4, 5, 6, 7, 2, 3)
+    #     + rdm4aabb.transpose(4, 5, 2, 3, 0, 1, 6, 7)
+    #     + rdm4aabb.transpose(4, 5, 2, 3, 6, 7, 0, 1)
+    #     + rdm4abbb
+    #     + rdm4abbb.transpose(2, 3, 0, 1, 4, 5, 6, 7)
+    #     + rdm4abbb.transpose(2, 3, 4, 5, 0, 1, 6, 7)
+    #     + rdm4abbb.transpose(2, 3, 4, 5, 6, 7, 0, 1),
+    #     rdm4,
+    # )
+    return (
+        (rdm1a, rdm1b),
+        (rdm2aa, rdm2ab, rdm2bb),
+        (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb),
+        (rdm4aaaa, rdm4aaab, rdm4aabb, rdm4abbb, rdm4bbbb),
+    )

 def trans_rdm1s(cibra, ciket, norb, nelec, link_index=None):
     r'''Spin separated transition 1-particle density matrices.
@@ -939,6 +1010,16 @@ def make_rdm123(self, fcivec, norb, nelec, link_index=None, reorder=True):
         nelec = _unpack_nelec(nelec, self.spin)
         return make_rdm123(fcivec, norb, nelec, link_index, reorder)

+    @lib.with_doc(make_rdm1234s.__doc__)
+    def make_rdm1234s(self, fcivec, norb, nelec, link_index=None, reorder=True):
+        nelec = _unpack_nelec(nelec, self.spin)
+        return make_rdm1234s(fcivec, norb, nelec, link_index, reorder)
+
+    @lib.with_doc(make_rdm1234.__doc__)
+    def make_rdm1234(self, fcivec, norb, nelec, link_index=None, reorder=True):
+        nelec = _unpack_nelec(nelec, self.spin)
+        return make_rdm1234(fcivec, norb, nelec, link_index, reorder)
+
     def make_rdm2(self, fcivec, norb, nelec, link_index=None, reorder=True):
         r'''Spin traced 2-particle density matrice

diff --git a/pyscf/fci/test/test_rdm.py b/pyscf/fci/test/test_rdm.py
index 8ac3b59957..7366783c7d 100644
--- a/pyscf/fci/test/test_rdm.py
+++ b/pyscf/fci/test/test_rdm.py
@@ -134,6 +134,103 @@ def test_dm4(self):
         self.assertTrue(numpy.allclose(dm3a, numpy.einsum('mnijppkl->mnijkl',dm4b)/5))
         self.assertTrue(numpy.allclose(dm3a, numpy.einsum('mnijklpp->mnijkl',dm4b)/5))

+    def test_rdm4s(self):
+        dm1, dm2, dm3, dm4 = fci.direct_spin1.make_rdm1234(ci0, norb, (nelec//2,nelec//2))
+
+        (dm1a, dm1b), (dm2aa, dm2ab, dm2bb), (dm3aaa, dm3aab, dm3abb, dm3bbb), (dm4aaaa, dm4aaab, dm4aabb, dm4abbb, dm4bbbb) = \
+        fci.direct_spin1.make_rdm1234s(ci0, norb, (nelec//2,nelec//2))
+
+        self.assertTrue(numpy.allclose(dm1a+dm1b, dm1))
+        self.assertTrue(numpy.allclose(dm2aa+dm2bb+dm2ab+dm2ab.transpose(2,3,0,1), dm2))
+        self.assertTrue(numpy.allclose(dm3aaa+dm3bbb+dm3aab+dm3aab.transpose(0,1,4,5,2,3)+\
+        dm3aab.transpose(4,5,0,1,2,3)+dm3abb+dm3abb.transpose(2,3,0,1,4,5)+dm3abb.transpose(2,3,4,5,0,1), dm3))
+        self.assertTrue(numpy.allclose(
+            dm4aaaa
+            + dm4bbbb
+            + dm4aaab
+            + dm4aaab.transpose(0, 1, 2, 3, 6, 7, 4, 5)
+            + dm4aaab.transpose(0, 1, 6, 7, 2, 3, 4, 5)
+            + dm4aaab.transpose(6, 7, 0, 1, 2, 3, 4, 5)
+            + dm4aabb
+            + dm4aabb.transpose(0, 1, 4, 5, 2, 3, 6, 7)
+            + dm4aabb.transpose(4, 5, 0, 1, 2, 3, 6, 7)
+            + dm4aabb.transpose(0, 1, 4, 5, 6, 7, 2, 3)
+            + dm4aabb.transpose(4, 5, 0, 1, 6, 7, 2, 3)
+            + dm4aabb.transpose(4, 5, 6, 7, 0, 1, 2, 3)
+            + dm4abbb
+            + dm4abbb.transpose(2, 3, 0, 1, 4, 5, 6, 7)
+            + dm4abbb.transpose(2, 3, 4, 5, 0, 1, 6, 7)
+            + dm4abbb.transpose(2, 3, 4, 5, 6, 7, 0, 1),
+            dm4,
+        ))
+
+        (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb), (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb) = \
+        fci.direct_spin1.make_rdm123s(ci0, norb, (nelec//2,nelec//2))
+
+        self.assertTrue(numpy.allclose(rdm1a, dm1a))
+        self.assertTrue(numpy.allclose(rdm1b, dm1b))
+
+        fac = 1. / (nelec//2-1)
+        self.assertTrue(numpy.allclose(rdm1a, numpy.einsum('ijmm->ij',dm2aa)*fac))
+        self.assertTrue(numpy.allclose(rdm1a, numpy.einsum('mmij->ij',dm2aa)*fac))
+        self.assertTrue(numpy.allclose(rdm1b, numpy.einsum('ijmm->ij',dm2bb)*fac))
+        self.assertTrue(numpy.allclose(rdm1b, numpy.einsum('mmij->ij',dm2bb)*fac))
+
+        fac = 1. / (nelec//2) # dm2ab has a prefactor N * N instead of N(N-1)
+        self.assertTrue(numpy.allclose(rdm1a, numpy.einsum('ijmm->ij',dm2ab)*fac))
+        self.assertTrue(numpy.allclose(rdm1b, numpy.einsum('mmij->ij',dm2ab)*fac))
+
+        fac = 1. / (nelec//2-2)
+        self.assertTrue(numpy.allclose(rdm2aa, numpy.einsum('ijklmm->ijkl',dm3aaa)*fac))
+        self.assertTrue(numpy.allclose(rdm2aa, numpy.einsum('ijmmkl->ijkl',dm3aaa)*fac))
+        self.assertTrue(numpy.allclose(rdm2aa, numpy.einsum('mmijkl->ijkl',dm3aaa)*fac))
+
+        self.assertTrue(numpy.allclose(rdm2bb, numpy.einsum('ijklmm->ijkl',dm3bbb)*fac))
+        self.assertTrue(numpy.allclose(rdm2bb, numpy.einsum('ijmmkl->ijkl',dm3bbb)*fac))
+        self.assertTrue(numpy.allclose(rdm2bb, numpy.einsum('mmijkl->ijkl',dm3bbb)*fac))
+
+        fac = 1. / (nelec//2) # dm3aab/abb has a prefactor N * N * (N-1) instead of N(N-1)(N-2)
+        self.assertTrue(numpy.allclose(rdm2aa, numpy.einsum('ijklmm->ijkl',dm3aab)*fac))
+        self.assertTrue(numpy.allclose(rdm2bb, numpy.einsum('mmijkl->ijkl',dm3abb)*fac))
+
+        fac = 1. / (nelec//2-1)
+        self.assertTrue(numpy.allclose(rdm2ab, numpy.einsum('ijmmkl->ijkl',dm3aab)*fac))
+        self.assertTrue(numpy.allclose(rdm2ab, numpy.einsum('mmijkl->ijkl',dm3aab)*fac))
+        self.assertTrue(numpy.allclose(rdm2ab, numpy.einsum('ijklmm->ijkl',dm3abb)*fac))
+        self.assertTrue(numpy.allclose(rdm2ab, numpy.einsum('ijmmkl->ijkl',dm3abb)*fac))
+
+        # NOTE: Tests skipeed because system only contains 3 alpha and beta electrons
+        # fac = 1. / (nelec//2-3)
+        # self.assertTrue(numpy.allclose(rdm3aaa, numpy.einsum('ijklmnoo->ijklmn',dm4aaaa)*fac))
+        # self.assertTrue(numpy.allclose(rdm3aaa, numpy.einsum('ijkloomn->ijklmn',dm4aaaa)*fac))
+        # self.assertTrue(numpy.allclose(rdm3aaa, numpy.einsum('ijooklmn->ijklmn',dm4aaaa)*fac))
+        # self.assertTrue(numpy.allclose(rdm3aaa, numpy.einsum('ooijklmn->ijklmn',dm4aaaa)*fac))
+        #
+        # self.assertTrue(numpy.allclose(rdm3bbb, numpy.einsum('ijklmnoo->ijklmn',dm4bbbb)*fac))
+        # self.assertTrue(numpy.allclose(rdm3bbb, numpy.einsum('ijkloomn->ijklmn',dm4bbbb)*fac))
+        # self.assertTrue(numpy.allclose(rdm3bbb, numpy.einsum('ijooklmn->ijklmn',dm4bbbb)*fac))
+        # self.assertTrue(numpy.allclose(rdm3bbb, numpy.einsum('ooijklmn->ijklmn',dm4bbbb)*fac))
+
+        fac = 1. / (nelec//2-2)
+        self.assertTrue(numpy.allclose(rdm3aab, numpy.einsum('ooijklmn->ijklmn',dm4aaab)*fac))
+        self.assertTrue(numpy.allclose(rdm3aab, numpy.einsum('ijooklmn->ijklmn',dm4aaab)*fac))
+        self.assertTrue(numpy.allclose(rdm3aab, numpy.einsum('ijkloomn->ijklmn',dm4aaab)*fac))
+
+        self.assertTrue(numpy.allclose(rdm3abb, numpy.einsum('ijklmnoo->ijklmn',dm4abbb)*fac))
+        self.assertTrue(numpy.allclose(rdm3abb, numpy.einsum('ijkloomn->ijklmn',dm4abbb)*fac))
+        self.assertTrue(numpy.allclose(rdm3abb, numpy.einsum('ijooklmn->ijklmn',dm4abbb)*fac))
+
+        fac = 1. / (nelec//2-1)
+        self.assertTrue(numpy.allclose(rdm3aab, numpy.einsum('ijklmnoo->ijklmn',dm4aabb)*fac))
+        self.assertTrue(numpy.allclose(rdm3aab, numpy.einsum('ijkloomn->ijklmn',dm4aabb)*fac))
+
+        self.assertTrue(numpy.allclose(rdm3abb, numpy.einsum('ijooklmn->ijklmn',dm4aabb)*fac))
+        self.assertTrue(numpy.allclose(rdm3abb, numpy.einsum('ooijklmn->ijklmn',dm4aabb)*fac))
+
+        fac = 1. / (nelec//2)
+        self.assertTrue(numpy.allclose(rdm3bbb, numpy.einsum('ooijklmn->ijklmn',dm4abbb)*fac))
+        self.assertTrue(numpy.allclose(rdm3aaa, numpy.einsum('ijklmnoo->ijklmn',dm4aaab)*fac))
+
     def test_tdm2(self):
         dm1 = numpy.einsum('ij,ijkl->lk', ci0, _trans1(ci0, norb, nelec))
         self.assertTrue(numpy.allclose(rdm1, dm1))


Enable hermi=0 for Breit and Gaunt JK integrals (pyscf/pyscf#2870)

[UNEXPECTED BEHAVIOUR] No Gaunt or Breit possible in TDDFT

**Gaunt and Breit Interactions don't work in TDDFT**
When using the later mentioned code, the regular `pyscf.dft.dks().DKS()` works, when using the options `with_gaunt = True` or `with_breit = True`. When subsequently performing a TDDFT calculation, an unimplemented Error occurs. This is not the case when `with_gaunt` and `with_breit` are turned to `False`.

**Minimal Working Example**
A minimal example, where the error occurs:
```python
from pyscf import gto, dft

mol = gto.Mole()

mol.build(
    atom = '''H 0 0 0
    F 0 0 1.5''',
    basis = '631g',
    verbose = 4,
    symmetry = True,
    )

mf = dft.dks.DKS(mol)
mf.xc = 'pbe0'
mf.with_gaunt = True
mf.with_breit = True
mf.kernel()

mytd = mf.TDDFT()

mytd.kernel()
```

**Error Message**
```bash
Traceback (most recent call last):
  File "/home/kevin/Documents/Spin-Orbit-Coupling/pyscf/HF/test.py", line 22, in <module>
    mytd.kernel()
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/dhf.py", line 567, in kernel
    self.converged, w, x1 = lr_eig(
                            ^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/_lr_eig.py", line 340, in eig
    axt = aop(xt)
          ^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/tdscf/dhf.py", line 504, in vind
    v1ao = vresp(dms) # = <mj||nb> Xjb + <mb||nj> Yjb
           ^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/_response_functions.py", line 297, in vind
    vj, vk = mf.get_jk(mol, dm1, hermi, with_j=with_j)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/dhf.py", line 679, in get_jk
    vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, False)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kevin/.local/lib/python3.12/site-packages/pyscf/scf/dhf.py", line 1038, in _call_veff_gaunt_breit
    raise NotImplementedError
NotImplementedError
```

**Version and hardware:**
- PySCF version: 2.9.0
- PySCF install type: built from source and pip, both result in the same error
- Operating system: Ubuntu 24.04

diff --git a/pyscf/scf/dhf.py b/pyscf/scf/dhf.py
index 7a9ced7eed..06395fcef7 100644
--- a/pyscf/scf/dhf.py
+++ b/pyscf/scf/dhf.py
@@ -1010,32 +1010,57 @@ def _call_veff_gaunt_breit(mol, dm, hermi=1, mf_opt=None, with_breit=False):
     vj = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex128)
     vk = numpy.zeros((n_dm,n2c*2,n2c*2), dtype=numpy.complex128)

-    jks = ('lk->s1ij', 'jk->s1il')
-    vj_ls, vk_ls = _vhf.rdirect_mapdm(intor_prefix+'ssp1ssp2_spinor', 's1', jks, dms[:n_dm], 1,
-                            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
-    vj[:,:n2c,n2c:] = vj_ls
-    vk[:,:n2c,n2c:] = vk_ls
-    t0 = log.timer_debug1('LSLS contribution', *t0)
-
-    jks = ('lk->s1ij',) * n_dm \
-        + ('li->s1kj',) * n_dm \
-        + ('jk->s1il',) * n_dm
-    vx = _vhf.rdirect_bindm(intor_prefix+'ssp1sps2_spinor', 's1', jks, dms[n_dm:], 1,
-                            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
-
-    t0 = log.timer_debug1('LSSL contribution', *t0)
-    vj[:,:n2c,n2c:]+= vx[      :n_dm  ,:,:]
-    vk[:,n2c:,n2c:] = vx[n_dm  :n_dm*2,:,:]
-    vk[:,:n2c,:n2c] = vx[n_dm*2:      ,:,:]
-
-    if hermi == 1:
-        vj[:,n2c:,:n2c] = vj[:,:n2c,n2c:].transpose(0,2,1).conj()
-        vk[:,n2c:,:n2c] = vk[:,:n2c,n2c:].transpose(0,2,1).conj()
-    elif hermi == 2:
-        vj[:,n2c:,:n2c] = -vj[:,:n2c,n2c:].transpose(0,2,1).conj()
-        vk[:,n2c:,:n2c] = -vk[:,:n2c,n2c:].transpose(0,2,1).conj()
+    if hermi != 0:
+        jks = ('lk->s1ij', 'jk->s1il')
+        vj_ls, vk_ls = _vhf.rdirect_mapdm(
+            intor_prefix+'ssp1ssp2_spinor', 's1', jks, dms[:n_dm], 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
+        vj[:,:n2c,n2c:] = vj_ls
+        vk[:,:n2c,n2c:] = vk_ls
+        t0 = log.timer_debug1('LSLS contribution', *t0)
+
+        jks = (('lk->s1ij',) * n_dm +
+               ('li->s1kj',) * n_dm +
+               ('jk->s1il',) * n_dm)
+        vx = _vhf.rdirect_bindm(
+            intor_prefix+'ssp1sps2_spinor', 's1', jks, dms[n_dm:], 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
+        vj[:,:n2c,n2c:]+= vx[      :n_dm  ,:,:]
+        vk[:,n2c:,n2c:] = vx[n_dm  :n_dm*2,:,:]
+        vk[:,:n2c,:n2c] = vx[n_dm*2:      ,:,:]
+        t0 = log.timer_debug1('LSSL contribution', *t0)
+
+        if hermi == 1:
+            vj[:,n2c:,:n2c] = vj[:,:n2c,n2c:].transpose(0,2,1).conj()
+            vk[:,n2c:,:n2c] = vk[:,:n2c,n2c:].transpose(0,2,1).conj()
+        elif hermi == 2:
+            vj[:,n2c:,:n2c] = -vj[:,:n2c,n2c:].transpose(0,2,1).conj()
+            vk[:,n2c:,:n2c] = -vk[:,:n2c,n2c:].transpose(0,2,1).conj()
     else:
-        raise NotImplementedError
+        jks = ('lk->s1ij', 'jk->s1il')
+        subj, subk = _vhf.rdirect_mapdm(
+            intor_prefix+'ssp1ssp2_spinor', 's1', jks,
+            (dms[:n_dm] + [x.conj().T for x in dms[n_dm:n_dm*2]]), 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lsls)
+        vj[:,:n2c,n2c:] = subj[      :n_dm  ,:,:]
+        vj[:,n2c:,:n2c] = subj[n_dm  :n_dm*2,:,:].conj().transpose(0,2,1)
+        vk[:,:n2c,n2c:] = subk[      :n_dm  ,:,:]
+        vk[:,n2c:,:n2c] = subk[n_dm  :n_dm*2,:,:].conj().transpose(0,2,1)
+        t0 = log.timer_debug1('LSLS contribution', *t0)
+
+        jks = (('ji->s1kl',) * n_dm +
+               ('lk->s1ij',) * n_dm +
+               ('li->s1kj',) * n_dm +
+               ('jk->s1il',) * n_dm)
+        vx = _vhf.rdirect_bindm(
+            intor_prefix+'ssp1sps2_spinor', 's1', jks, dms, 1,
+            mol._atm, mol._bas, mol._env, opt_gaunt_lssl)
+        vj[:,n2c:,:n2c]+= vx[      :n_dm  ,:,:]
+        vj[:,:n2c,n2c:]+= vx[n_dm  :n_dm*2,:,:]
+        vk[:,n2c:,n2c:] = vx[n_dm*2:n_dm*3,:,:]
+        vk[:,:n2c,:n2c] = vx[n_dm*3:n_dm*4,:,:]
+        t0 = log.timer_debug1('LSSL contribution', *t0)
+
     vj = vj.reshape(dm.shape)
     vk = vk.reshape(dm.shape)
     c1 = .5 / lib.param.LIGHT_SPEED

diff --git a/pyscf/scf/test/test_dhf.py b/pyscf/scf/test/test_dhf.py
index 07ee57229f..2db8b48017 100644
--- a/pyscf/scf/test/test_dhf.py
+++ b/pyscf/scf/test/test_dhf.py
@@ -250,11 +250,17 @@ def test_gaunt(self):
         n4c = erig.shape[0]
         numpy.random.seed(1)
         dm = numpy.random.random((2,n4c,n4c))+numpy.random.random((2,n4c,n4c))*1j
-        dm = dm + dm.transpose(0,2,1).conj()
         c1 = .5 / lib.param.LIGHT_SPEED
         vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
         vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
+        vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, hermi=0)
+        self.assertAlmostEqual(abs(vj0 - vj1).max(), 0, 12)
+        self.assertAlmostEqual(abs(vk0 - vk1).max(), 0, 12)

+        dm = dm + dm.transpose(0,2,1).conj()
+        c1 = .5 / lib.param.LIGHT_SPEED
+        vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
+        vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
         vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm)
         self.assertTrue(numpy.allclose(vj0, vj1))
         self.assertTrue(numpy.allclose(vk0, vk1))
@@ -279,11 +285,17 @@ def test_breit_high_cost(self):
         n4c = erig.shape[0]
         numpy.random.seed(1)
         dm = numpy.random.random((n4c,n4c))+numpy.random.random((n4c,n4c))*1j
+        c1 = .5 / lib.param.LIGHT_SPEED
+        vj0 = -numpy.einsum('ijkl,xlk->xij', erig, dm) * c1**2
+        vk0 = -numpy.einsum('ijkl,xjk->xil', erig, dm) * c1**2
+        vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, hermi=0)
+        self.assertAlmostEqual(abs(vj0 - vj1).max(), 0, 12)
+        self.assertAlmostEqual(abs(vk0 - vk1).max(), 0, 12)
+
         dm = dm + dm.T.conj()
         c1 = .5 / lib.param.LIGHT_SPEED
         vj0 = numpy.einsum('ijkl,lk->ij', erig, dm) * c1**2
         vk0 = numpy.einsum('ijkl,jk->il', erig, dm) * c1**2
-
         vj1, vk1 = scf.dhf._call_veff_gaunt_breit(h4, dm, with_breit=True)
         self.assertTrue(numpy.allclose(vj0, vj1))
         self.assertTrue(numpy.allclose(vk0, vk1))
@@ -325,23 +337,22 @@ def _fill_gaunt(mol, erig):
     eri0 = numpy.zeros((n4c,n4c,n4c,n4c), dtype=numpy.complex128)
     eri0[:n2c,n2c:,:n2c,n2c:] = erig # ssp1ssp2

-    eri2 = erig.take(idx,axis=0).take(idx,axis=1) # sps1ssp2
+    eri2 = erig[idx[:,None],idx] # sps1ssp2
     eri2[sign_mask,:] *= -1
     eri2[:,sign_mask] *= -1
     eri2 = -eri2.transpose(1,0,2,3)
     eri0[n2c:,:n2c,:n2c,n2c:] = eri2

-    eri2 = erig.take(idx,axis=2).take(idx,axis=3) # ssp1sps2
+    eri2 = erig[:,:,idx[:,None],idx] # ssp1sps2
     eri2[:,:,sign_mask,:] *= -1
     eri2[:,:,:,sign_mask] *= -1
     eri2 = -eri2.transpose(0,1,3,2)
     #self.assertTrue(numpy.allclose(eri0, eri2))
     eri0[:n2c,n2c:,n2c:,:n2c] = eri2

-    eri2 = erig.take(idx,axis=0).take(idx,axis=1)
-    eri2 = eri2.take(idx,axis=2).take(idx,axis=3) # sps1sps2
-    eri2 = eri2.transpose(1,0,2,3)
-    eri2 = eri2.transpose(0,1,3,2)
+    eri2 = erig[idx[:,None],idx]
+    eri2 = eri2[:,:,idx[:,None],idx] # sps1sps2
+    eri2 = eri2.transpose(1,0,3,2)
     eri2[sign_mask,:] *= -1
     eri2[:,sign_mask] *= -1
     eri2[:,:,sign_mask,:] *= -1


Add NLC response function for GKS (pyscf/pyscf#2871)

[FEATURE REQUEST] Nonlocal corrections and GKS

Hello!

I would like to run GKS with the VV10 dispersion correction. However, I get the error "NotImplementedError: NLC" from the `_gen_ghf_response` function in scf/_response_functions.py. Is there any way to include VV10 for use with the GKS module? Or has VV10 only been developed for RKS and UKS thus far?

Thank you!

diff --git a/pyscf/dft/numint2c.py b/pyscf/dft/numint2c.py
index f34bc6b892..0daadda128 100644
--- a/pyscf/dft/numint2c.py
+++ b/pyscf/dft/numint2c.py
@@ -679,8 +679,8 @@ def nr_nlc_vxc(self, mol, grids, xc_code, dm, spin=0, relativity=0, hermi=1,
         n, exc, v = ni.nr_nlc_vxc(mol, grids, xc_code, dm_a+dm_b, relativity,
                                   hermi, max_memory, verbose)
         vmat = np.zeros_like(dm)
-        vmat[:nao,:nao] = v[0]
-        vmat[nao:,nao:] = v[1]
+        vmat[:nao,:nao] = v
+        vmat[nao:,nao:] = v
         return n, exc, vmat

     @lib.with_doc(numint.nr_rks_fxc.__doc__)
diff --git a/pyscf/hessian/rks.py b/pyscf/hessian/rks.py
index 554bdd2f1f..b559c6d7e8 100644
--- a/pyscf/hessian/rks.py
+++ b/pyscf/hessian/rks.py
@@ -1849,19 +1849,20 @@ def get_vnlc_resp(mf, mol, mo_coeff, mo_occ, dm1s, max_memory):

         TODO: check the effect of different grid, using mf.nlcgrids right now
     """
-    if mo_coeff.ndim == 2:
+    if isinstance(mo_coeff, numpy.ndarray) and mo_coeff.ndim == 2:
         mocc = mo_coeff[:,mo_occ>0]
         mo_occ = mo_occ[mo_occ > 0]
         dm0 = (mocc * mo_occ) @ mocc.T
     else:
-        assert mo_coeff.ndim == 3 # unrestricted case
-        assert mo_coeff.shape[0] == 2
-        assert mo_occ.shape[0] == 2
+        assert mo_coeff[0].ndim == 2 # unrestricted case
+        assert len(mo_coeff) == 2
+        assert len(mo_occ) == 2
         mocc_a = mo_coeff[0][:, mo_occ[0] > 0]
         mocc_b = mo_coeff[1][:, mo_occ[1] > 0]
-        mo_occ_a = mo_occ[0, mo_occ[0] > 0]
-        mo_occ_b = mo_occ[1, mo_occ[1] > 0]
+        mo_occ_a = mo_occ[0][mo_occ[0] > 0]
+        mo_occ_b = mo_occ[1][mo_occ[1] > 0]
         dm0 = (mocc_a * mo_occ_a) @ mocc_a.T + (mocc_b * mo_occ_b) @ mocc_b.T
+    dm0 = numpy.asarray(dm0.real, order='C')

     output_in_2d = False
     if dm1s.ndim == 2:
diff --git a/pyscf/scf/_response_functions.py b/pyscf/scf/_response_functions.py
index fd7f458e48..df75029caa 100644
--- a/pyscf/scf/_response_functions.py
+++ b/pyscf/scf/_response_functions.py
@@ -267,8 +267,6 @@ def _gen_ghf_response(mf, mo_coeff=None, mo_occ=None,
         ni = mf._numint
         assert isinstance(ni, (numint2c.NumInt2C, r_numint.RNumInt))
         ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
-        if with_nlc and mf.do_nlc():
-            raise NotImplementedError('NLC')
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)
         hybrid = ni.libxc.is_hybrid_xc(mf.xc)

@@ -288,6 +286,20 @@ def vind(dm1):
             else:
                 v1 = ni.get_fxc(mol, mf.grids, mf.xc, dm0, dm1, 0, 0, hermi,
                                 rho0, vxc, fxc, max_memory=max_memory)
+                if with_nlc and mf.do_nlc():
+                    from pyscf.hessian.rks import get_vnlc_resp
+                    nao = mo_coeff.shape[0] // 2
+                    dm1_sf = dm1[...,:nao,:nao] + dm1[...,nao:,nao:]
+                    mo_uks = [mo_coeff[:nao], mo_coeff[nao:]]
+                    mo_occ_uks = [mo_occ, mo_occ]
+                    # The get_vnlc_resp function uses mf._numint and explicitly
+                    # calls the NumInt functions for mf._numint
+                    ni1c = mf._numint._to_numint1c()
+                    with lib.temporary_env(mf, _numint=ni1c):
+                        vxc_nlc = get_vnlc_resp(mf, mol, mo_uks, mo_occ_uks,
+                                                dm1_sf.real, max_memory)
+                        v1[...,:nao,:nao] += vxc_nlc
+                        v1[...,nao:,nao:] += vxc_nlc
             if not hybrid:
                 if with_j:
                     vj = mf.get_j(mol, dm1, hermi=hermi)

diff --git a/pyscf/dft/test/test_gks.py b/pyscf/dft/test/test_gks.py
index 4a0ca4662f..07e59ab464 100644
--- a/pyscf/dft/test/test_gks.py
+++ b/pyscf/dft/test/test_gks.py
@@ -15,6 +15,7 @@

 import unittest
 import numpy
+import scipy.linalg
 from pyscf import gto
 from pyscf import lib
 from pyscf.dft import numint, numint2c, libxc, radi
diff --git a/pyscf/dft/test/test_h2o.py b/pyscf/dft/test/test_h2o.py
index 585239bdb6..aa55453984 100644
--- a/pyscf/dft/test/test_h2o.py
+++ b/pyscf/dft/test/test_h2o.py
@@ -15,6 +15,7 @@

 import unittest
 import numpy
+import scipy.linalg
 from pyscf import gto
 from pyscf import lib
 from pyscf import dft
@@ -371,32 +372,34 @@ def test_nr_rks_nlc(self):
         method = dft.RKS(h2o)
         dm = method.get_init_guess()
         method.xc = 'wB97M_V'
+        method.nlcgrids.level = 0
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc), 22.767792068559917, 8)
+        self.assertAlmostEqual(lib.fp(vxc), 22.767790354129247, 8)

         method.xc = 'B97M_V'
         vxc = method.get_veff(h2o, dm, dm, vxc)
-        self.assertAlmostEqual(lib.fp(vxc), 23.067046560473408, 8)
+        self.assertAlmostEqual(lib.fp(vxc), 23.06704484604275, 8)

         method.nlc = False
         assert method.do_nlc() == False
         vxc = method.get_veff(h2o, dm, dm, vxc)
         self.assertAlmostEqual(lib.fp(vxc), 23.05881308880983, 8)

-    def test_nr_rks_nlc_small_memory_high_cost(self):
+    def test_nr_rks_nlc_small_memory(self):
         method = dft.RKS(h2o)
         dm = method.get_init_guess()
         method._eri = None
         method.max_memory = 0
         method.xc = 'wB97M_V'
+        method.nlcgrids.level = 0
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc), 22.767792068559917, 8)
+        self.assertAlmostEqual(lib.fp(vxc), 22.767790354129247, 8)

         method._eri = None
         method.max_memory = 0
         method.xc = 'B97M_V'
         vxc = method.get_veff(h2o, dm, dm, vxc)
-        self.assertAlmostEqual(lib.fp(vxc), 23.067046560473408, 8)
+        self.assertAlmostEqual(lib.fp(vxc), 23.06704484604275, 8)

     def test_nr_rks_rsh_cart_high_cost(self):
         mol1 = h2o.copy()
@@ -417,37 +420,39 @@ def test_nr_uks_rsh(self):
         self.assertAlmostEqual(lib.fp(vxc[0]), 23.16975737295899, 8)
         self.assertAlmostEqual(lib.fp(vxc[1]), 23.16975737295899, 8)

-    def test_nr_uks_nlc_high_cost(self):
+    def test_nr_uks_nlc(self):
         method = dft.UKS(h2o)
         dm = method.get_init_guess()
         dm = (dm[0], dm[0])
         method.xc = 'wB97M_V'
+        method.nlcgrids.level = 0
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767792068559917, 8)
-        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767792068559917, 8)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767790354129247, 8)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767790354129247, 8)

         method.xc = 'B97M_V'
         vxc = method.get_veff(h2o, dm, dm, vxc)
-        self.assertAlmostEqual(lib.fp(vxc[0]), 23.067046560473408, 8)
-        self.assertAlmostEqual(lib.fp(vxc[1]), 23.067046560473408, 8)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 23.06704484604275, 8)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 23.06704484604275, 8)

-    def test_nr_uks_nlc_small_memory_high_cost(self):
+    def test_nr_uks_nlc_small_memory(self):
         method = dft.UKS(h2o)
         dm = method.get_init_guess()
         dm = (dm[0], dm[0])
         method._eri = None
         method.max_memory = 0
         method.xc = 'wB97M_V'
+        method.nlcgrids.level = 0
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767792068559917, 8)
-        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767792068559917, 8)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767790354129247, 8)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767790354129247, 8)

         method._eri = None
         method.max_memory = 0
         method.xc = 'B97M_V'
         vxc = method.get_veff(h2o, dm, dm, vxc)
-        self.assertAlmostEqual(lib.fp(vxc[0]), 23.067046560473408, 8)
-        self.assertAlmostEqual(lib.fp(vxc[1]), 23.067046560473408, 8)
+        self.assertAlmostEqual(lib.fp(vxc[0]), 23.06704484604275, 8)
+        self.assertAlmostEqual(lib.fp(vxc[1]), 23.06704484604275, 8)

     def test_nr_gks_rsh(self):
         method = dft.GKS(h2o)
@@ -458,35 +463,26 @@ def test_nr_gks_rsh(self):
         vxc = method.get_veff(h2o, dm)
         self.assertAlmostEqual(lib.fp(vxc), 5.115622298912124+0j, 8)

-    def test_nr_gks_nlc_high_cost(self):
+    def test_nr_gks_nlc(self):
         method = dft.GKS(h2o)
         dm = method.get_init_guess()
         dm = dm + numpy.sin(dm)*.02j
         dm = dm + dm.conj().T
         method.xc = 'wB97M_V'
+        method.nlcgrids.level = 0
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc), 3.172920887028461+0j, 8)
-
+        self.assertAlmostEqual(lib.fp(vxc), 3.17466829111814+0j, 8)
         method.xc = 'B97M_V'
         vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc), 2.0041673361905317+0j, 8)
-
-    def test_nr_gks_nlc_small_memory_high_cost(self):
-        method = dft.GKS(h2o)
-        dm = method.get_init_guess()
-        dm = dm + numpy.sin(dm)*.02j
-        dm = dm + dm.conj().T
-        method._eri = None
-        method.max_memory = 0
-        method.xc = 'wB97M_V'
-        vxc = method.get_veff(h2o, dm)
-        self.assertAlmostEqual(lib.fp(vxc), 3.172920887028461+0j, 8)
-
-        method._eri = None
-        method.max_memory = 0
-        method.xc = 'B97M_V'
-        vxc = method.get_veff(h2o, dm, dm, vxc)
-        self.assertAlmostEqual(lib.fp(vxc), 2.0041673361905317+0j, 8)
+        self.assertAlmostEqual(lib.fp(vxc), 2.00591474028017+0j, 8)
+
+        mf = h2o_cation.UKS(xc='wb97mv')
+        mf.nlcgrids.level = 0
+        dm = mf.get_init_guess()
+        ref = scipy.linalg.block_diag(*mf.get_veff(h2o_cation, dm))
+        mf_gks = mf.to_gks()
+        vxc = mf_gks.get_veff(h2o_cation, scipy.linalg.block_diag(*dm))
+        self.assertAlmostEqual(abs(vxc - ref).max(), 0, 12)

     def test_nr_rks_vv10_high_cost(self):
         method = dft.RKS(h2o)
diff --git a/pyscf/scf/test/test_response_function.py b/pyscf/scf/test/test_response_function.py
new file mode 100644
index 0000000000..0874b0a035
--- /dev/null
+++ b/pyscf/scf/test/test_response_function.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+# Copyright 2025 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 scipy.linalg
+from pyscf import lib
+from pyscf import gto
+from pyscf.scf import _response_functions
+
+class KnownValues(unittest.TestCase):
+    def test_gks_nlc(self):
+        mol = gto.M(
+            verbose = 5,
+            output = '/dev/null',
+            atom = [
+            ["O" , (0. , 0.     , 0.)],
+            [1   , (0. , -0.757 , 0.587)],
+            [1   , (0. , 0.757  , 0.587)]],
+            charge = 1,
+            spin = 1,
+            basis = '631g')
+        nao = mol.nao
+        mf_uks = mol.UHF().run().to_uks()
+        mf_uks.xc = 'wb97mv'
+        mf_uks.nlcgrids.level = 0
+
+        dm = mf_uks.make_rdm1()
+        dm1 = np.random.rand(2, nao, nao)
+        vind = mf_uks.gen_response(with_nlc=True)
+        ref = scipy.linalg.block_diag(*vind(dm1))
+
+        mf_gks = mf_uks.to_gks()
+        vind = mf_gks.gen_response(with_nlc=True)
+        v = vind(scipy.linalg.block_diag(*dm1))
+        self.assertAlmostEqual(abs(v - ref).max(), 0, 12)
+
+if __name__ == "__main__":
+    print("Full Tests for response_functions")
+    unittest.main()