Fix initial guess dimension problems when solving for multiple roots (pyscf/pyscf#2306)

PySCF >=v2.2.1 fci.direct_spin1 davidson solver issue

Hello

I have the following minimal example using FCI to solve a given Hamiltonian which produces an error when using a PySCF version greater than or equal to v2.2.1.
``` python
from pyscf import fci
import numpy as np

norb=4
nelec=(2, 2)
nroots=36 # (4 choose 2)*(4 choose 2)

e = -0.5
t = -0.1
h1 = np.array([[e, t, 0, 0],
               [t, e, t, 0],
               [0, t, e, t],
               [0, 0, t, e],
              ])

U = 0.1
h2 = np.zeros((norb, norb, norb, norb))
for i in range(norb):
    h2[i,i,i,i] = U

e, fcivec = fci.direct_spin1.kernel(h1, h2, norb, nelec, nroots=nroots, max_space=30, max_cycle=100,
                                    davidson_only=True,
                                   )

print("Energies:", e)
```

The error message:
``` bash
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 257, in kernel
    return _kfactory(FCISolver, h1e, eri, norb, nelec, ci0, level_shift,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 281, in _kfactory
    e, c = cis.kernel(h1e, eri, norb, nelec, ci0, ecore=ecore, **unknown)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 842, in kernel
    e, c = kernel_ms1(self, h1e, eri, norb, nelec, ci0, link_index,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 612, in kernel_ms1
    e, c = fci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 810, in eig
    lib.davidson1(lambda xs: [op(x) for x in xs],
  File "/lib/python3.11/site-packages/pyscf/lib/linalg_helper.py", line 388, in davidson1
    x0 = x0()
         ^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 580, in init_guess
    return fci.get_init_guess(norb, nelec, nroots, hdiag)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 829, in get_init_guess
    return get_init_guess(norb, nelec, nroots, hdiag)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 469, in get_init_guess
    return _get_init_guess(na, nb, nroots, hdiag, nelec)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/pyscf/fci/direct_spin1.py", line 444, in _get_init_guess
    addrs = numpy.argpartition(lib.pack_tril(hdiag), nroots-1)[:nroots]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/numpy/_core/fromnumeric.py", line 908, in argpartition
    return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lib/python3.11/site-packages/numpy/_core/fromnumeric.py", line 57, in _wrapfunc
    return bound(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^^
ValueError: kth(=35) out of bounds (21)
```

I have only encountered this error with versions of PySCF >=v2.2.1, the versions before that run with the correct output. I believe the issue is related to `davidson_only=True` as setting `davidson_only=False` for >=v2.2.1 also does not produce an error.

Is this a bug with the davidson solver?

Thanks!

diff --git a/pyscf/fci/direct_spin0_symm.py b/pyscf/fci/direct_spin0_symm.py
index ed6feb92f9..5e1e8f1f82 100644
--- a/pyscf/fci/direct_spin0_symm.py
+++ b/pyscf/fci/direct_spin0_symm.py
@@ -74,15 +74,20 @@ def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
     strsa = cistring.gen_strings4orblist(range(norb), neleca)
     na = len(strsa)
     airreps = direct_spin1_symm._gen_strs_irrep(strsa, orbsym)
-    hdiag = hdiag.reshape(na,na)

     sym_allowed = (airreps[:,None] ^ airreps) == wfnsym
     idx = numpy.arange(na)
     sym_allowed[idx[:,None] < idx] = False
     idx_a, idx_b = numpy.where(sym_allowed)

+    hdiag = hdiag.reshape(na,na)[idx_a,idx_b]
+    if hdiag.size <= nroots:
+        hdiag_indices = numpy.arange(hdiag.size)
+    else:
+        hdiag_indices = numpy.argpartition(hdiag, nroots-1)[:nroots]
+
     ci0 = []
-    for k in numpy.argpartition(hdiag[idx_a,idx_b], nroots-1)[:nroots]:
+    for k in hdiag_indices:
         addra, addrb = idx_a[k], idx_b[k]
         x = numpy.zeros((na, na))
         if addra == addrb:
diff --git a/pyscf/fci/direct_spin1.py b/pyscf/fci/direct_spin1.py
index 7f8d04f3d3..710acc872c 100644
--- a/pyscf/fci/direct_spin1.py
+++ b/pyscf/fci/direct_spin1.py
@@ -440,8 +440,11 @@ def _get_init_guess(na, nb, nroots, hdiag, nelec):
     ci0 = []
     neleca, nelecb = _unpack_nelec(nelec)
     if neleca == nelecb and na == nb:
-        hdiag = hdiag.reshape(na, na)
-        addrs = numpy.argpartition(lib.pack_tril(hdiag), nroots-1)[:nroots]
+        hdiag = lib.pack_tril(hdiag.reshape(na, na))
+        if hdiag.size <= nroots:
+            addrs = numpy.arange(hdiag.size)
+        else:
+            addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
         for addr in addrs:
             addra = (int)((2*addr+.25)**.5 - .5 + 1e-7)
             addrb = addr - addra*(addra+1)//2
@@ -449,7 +452,10 @@ def _get_init_guess(na, nb, nroots, hdiag, nelec):
             x[addra,addrb] = 1
             ci0.append(x.ravel().view(FCIvector))
     else:
-        addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
+        if hdiag.size <= nroots:
+            addrs = numpy.arange(hdiag.size)
+        else:
+            addrs = numpy.argpartition(hdiag, nroots-1)[:nroots]
         for addr in addrs:
             x = numpy.zeros((na*nb))
             x[addr] = 1
diff --git a/pyscf/fci/direct_spin1_symm.py b/pyscf/fci/direct_spin1_symm.py
index 017505836a..ee10487b8f 100644
--- a/pyscf/fci/direct_spin1_symm.py
+++ b/pyscf/fci/direct_spin1_symm.py
@@ -217,15 +217,20 @@ def _get_init_guess(airreps, birreps, nroots, hdiag, nelec, orbsym, wfnsym=0):
     neleca, nelecb = _unpack_nelec(nelec)
     na = len(airreps)
     nb = len(birreps)
-    hdiag = hdiag.reshape(na,nb)
     sym_allowed = airreps[:,None] == wfnsym ^ birreps
     if neleca == nelecb and na == nb:
         idx = np.arange(na)
         sym_allowed[idx[:,None] < idx] = False
     idx_a, idx_b = np.where(sym_allowed)

+    hdiag = hdiag.reshape(na,nb)[idx_a,idx_b]
+    if hdiag.size <= nroots:
+        hdiag_indices = np.arange(hdiag.size)
+    else:
+        hdiag_indices = np.argpartition(hdiag, nroots-1)[:nroots]
+
     ci0 = []
-    for k in np.argpartition(hdiag[idx_a,idx_b], nroots-1)[:nroots]:
+    for k in hdiag_indices:
         addra, addrb = idx_a[k], idx_b[k]
         x = np.zeros((na, nb))
         x[addra,addrb] = 1
@@ -248,7 +253,11 @@ def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
         return _get_init_guess(airreps, birreps, nroots, hdiag, nelec, orbsym, wfnsym)

     ci0 = []
-    for k in np.argpartition(hdiag, nroots-1)[:nroots]:
+    if hdiag.size <= nroots:
+        hdiag_indices = np.arange(hdiag.size)
+    else:
+        hdiag_indices = np.argpartition(hdiag, nroots-1)[:nroots]
+    for k in hdiag_indices:
         x = np.zeros_like(hdiag)
         x[k] = 1.
         ci0.append(x.ravel().view(direct_spin1.FCIvector))
diff --git a/pyscf/fci/direct_uhf.py b/pyscf/fci/direct_uhf.py
index e5c7cda67a..6aaea941b3 100644
--- a/pyscf/fci/direct_uhf.py
+++ b/pyscf/fci/direct_uhf.py
@@ -198,7 +198,7 @@ def pspace(h1e, eri, norb, nelec, hdiag=None, np=400):
     na = cistring.num_strings(norb, neleca)
     nb = cistring.num_strings(norb, nelecb)
     assert hdiag.size == na * nb
-    if hdiag.size < np:
+    if hdiag.size <= np:
         addr = numpy.arange(hdiag.size)
     else:
         try:
diff --git a/pyscf/scf/dispersion.py b/pyscf/scf/dispersion.py
index d63dd47cd6..4e038420ff 100644
--- a/pyscf/scf/dispersion.py
+++ b/pyscf/scf/dispersion.py
@@ -23,6 +23,7 @@
 import warnings
 from functools import lru_cache
 from pyscf.lib import logger
+from pyscf import scf

 # supported dispersion corrections
 DISP_VERSIONS = ['d3bj', 'd3zero', 'd3bjm', 'd3zerom', 'd3op', 'd4']
@@ -115,8 +116,13 @@ def check_disp(mf, disp=None):
     if disp == 0: # disp = False
         return False

-    # The disp method for both HF and MCSCF is set to 'hf'
-    method = getattr(mf, 'xc', 'hf')
+    # To prevent mf.do_disp() triggering the SCF.__getattr__ method, do not use
+    # method = getattr(mf, 'xc', 'hf').
+    if isinstance(mf, scf.hf.KohnShamDFT):
+        method = mf.xc
+    else:
+        # Set the disp method for both HF and MCSCF to 'hf'
+        method = 'hf'
     disp_version = parse_disp(method)[1]

     if disp is None: # Using the disp version decoded from the mf.xc attribute
@@ -172,6 +178,5 @@ def get_dispersion(mf, disp=None, with_3body=None, verbose=None):
         raise RuntimeError(f'dipersion correction: {disp_version} is not supported.')

 # Inject to SCF class
-from pyscf import scf
 scf.hf.SCF.do_disp = check_disp
 scf.hf.SCF.get_dispersion = get_dispersion

diff --git a/pyscf/fci/test/test_spin1.py b/pyscf/fci/test/test_spin1.py
index 420cccba9a..23c7b9a4fc 100644
--- a/pyscf/fci/test/test_spin1.py
+++ b/pyscf/fci/test/test_spin1.py
@@ -198,6 +198,17 @@ def test_pspace(self):
         h = fci.direct_spin1.pspace(h1e, g2e, norb, nelec)[1]
         self.assertAlmostEqual(abs(h-h.T).max(), 0, 12)

+    def test_many_roots(self):
+        norb = 4
+        nelec = (2, 2)
+        nroots = 36
+        h1 = numpy.eye(norb) * -.5
+        h2 = numpy.zeros((norb, norb, norb, norb))
+        for i in range(norb):
+            h2[i,i,i,i] = .1
+        e, fcivec = fci.direct_spin1.kernel(h1, h2, norb, nelec, nroots=nroots, davidson_only=True)
+        self.assertAlmostEqual(e[0], -2, 9)
+

 if __name__ == "__main__":
     print("Full Tests for spin1")
diff --git a/pyscf/fci/test/test_spin1_symm.py b/pyscf/fci/test/test_spin1_symm.py
index 4b63d48d27..58a055e9be 100644
--- a/pyscf/fci/test/test_spin1_symm.py
+++ b/pyscf/fci/test/test_spin1_symm.py
@@ -208,6 +208,20 @@ def test_incomplete_orbsym(self):
         with self.assertRaises(lib.exceptions.PointGroupSymmetryError):
             sol.kernel(h1, h2, no, ne, orbsym=orbsym)

+    def test_many_roots(self):
+        norb = 4
+        nelec = (2, 2)
+        nroots = 36
+        h1 = numpy.eye(norb) * -.5
+        h2 = numpy.zeros((norb, norb, norb, norb))
+        orbsym = numpy.array([0, 5, 3, 2])
+        for i in range(norb):
+            h2[i,i,i,i] = .1
+        obj = direct_spin1_symm.FCI()
+        e, fcivec = obj.kernel(h1, h2, norb, nelec, nroots=nroots,
+                               davidson_only=True, orbsym=orbsym)
+        self.assertAlmostEqual(e[0], -1.8, 9)
+
 if __name__ == "__main__":
     print("Full Tests for spin1-symm")
     unittest.main()


Fix the transformation matrix of Mole.decontracted_basis (pyscf/pyscf#2354)

`mole.to_uncontracted_cartesian_basis` behavior change

The behavior of `mole.to_uncontracted_cartesian_basis` has chaged since v2.1. The coefficients returned no longer have the size of the number of shells of the original Mole object. This breaks some code for some basis sets, e.g., this [line]( https://github.com/pyscf/pyscf/blob/ca8c7c1680defdfee2380eda3af3a28d9fb375cb/pyscf/pbc/dft/multigrid/multigrid.py#L903).
I think we should modify the docstring and maybe also add a warning.

### 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/gto/mole.py`: `def decontract_basis(mol, atoms, to_cart, aggregate) -> Union[list, tuple]: ...`
  **Inputs:** mol: gto.Mole, atoms: Union[list, tuple], to_cart: bool, aggregate: bool
  **Output:** Union[list, tuple]

```python
def decontract_basis(mol, atoms, to_cart, aggregate) -> Union[list, tuple]: ...
```

diff --git a/pyscf/gto/mole.py b/pyscf/gto/mole.py
index 0f2ec506f9..24319af8cc 100644
--- a/pyscf/gto/mole.py
+++ b/pyscf/gto/mole.py
@@ -572,7 +572,7 @@ def to_uncontracted_cartesian_basis(mol):
     '''
     return decontract_basis(mol, to_cart=True)

-def decontract_basis(mol, atoms=None, to_cart=False):
+def decontract_basis(mol, atoms=None, to_cart=False, aggregate=False):
     '''Decontract the basis of a Mole or a Cell.  Returns a Mole (Cell) object
     with the uncontracted basis environment and a list of coefficients that
     transform the uncontracted basis to the original basis. Each element in
@@ -584,6 +584,9 @@ def decontract_basis(mol, atoms=None, to_cart=False):
             are decontracted
         to_cart: bool
             Decontract basis and transfer to Cartesian basis
+        aggregate: bool
+            Whether to aggregate the transformation coefficients into a giant
+            transformation matrix

     Examples:

@@ -602,17 +605,27 @@ def decontract_basis(mol, atoms=None, to_cart=False):
     bas_exps = mol.bas_exps()
     def _to_full_contraction(mol, bas_idx):
         es = numpy.hstack([bas_exps[i] for i in bas_idx])
-        cs = scipy.linalg.block_diag(*[mol._libcint_ctr_coeff(ib) for ib in bas_idx])
-
-        es, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
-        cs_new = numpy.zeros((es.size, cs.shape[1]))
-        for i, j in enumerate(rev_idx):
-            cs_new[j] += cs[i]
-        return es[::-1], cs_new[::-1]
+        _, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
+        if aggregate:
+            cs = scipy.linalg.block_diag(
+                *[mol._libcint_ctr_coeff(i) for i in bas_idx])
+            if len(es) != len(e_idx):
+                cs_new = numpy.zeros((e_idx.size, cs.shape[1]))
+                for i, j in enumerate(rev_idx):
+                    cs_new[j] += cs[i]
+                es = es[e_idx][::-1]
+                cs = cs_new[::-1]
+            yield es, cs
+        else:
+            if len(es) != len(e_idx):
+                raise RuntimeError('Duplicated pGTOs across shells')
+            for i in bas_idx:
+                yield bas_exps[i], mol._libcint_ctr_coeff(i)

     _bas = []
-    _env = mol._env.copy()
+    env = [mol._env.copy()]
     contr_coeff = []
+    pexp = env[0].size

     lmax = mol._bas[:,ANG_OF].max()
     if mol.cart:
@@ -648,7 +661,7 @@ def _to_full_contraction(mol, bas_idx):
                 continue

         lmax = mol._bas[ib0:ib1,ANG_OF].max()
-        pexp = mol._bas[ib0,PTR_EXP]
+
         for l in range(lmax+1):
             bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,ANG_OF] == l)[0]
             if len(bas_idx) == 0:
@@ -656,30 +669,27 @@ def _to_full_contraction(mol, bas_idx):
             if bas_idx[0] + len(bas_idx) != bas_idx[-1] + 1:
                 raise NotImplementedError('Discontinuous bases of same angular momentum')

-            mol_exps, b_coeff = _to_full_contraction(mol, bas_idx)
-            nprim, nc = b_coeff.shape
-            bs = numpy.zeros((nprim, BAS_SLOTS), dtype=numpy.int32)
-            bs[:,ATOM_OF] = ia
-            bs[:,ANG_OF ] = l
-            bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
-            norm = gto_norm(l, mol_exps)
-            if atoms is None:
+            for mol_exps, b_coeff in _to_full_contraction(mol, bas_idx):
+                nprim, nc = b_coeff.shape
+                bs = numpy.zeros((nprim, BAS_SLOTS), dtype=numpy.int32)
+                bs[:,ATOM_OF] = ia
+                bs[:,ANG_OF ] = l
+                bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
                 bs[:,PTR_EXP] = pexp + numpy.arange(nprim)
                 bs[:,PTR_COEFF] = pexp + numpy.arange(nprim, nprim*2)
-                _env[pexp:pexp+nprim] = mol_exps
-                _env[pexp+nprim:pexp+nprim*2] = norm
+                norm = gto_norm(l, mol_exps)
+                env.append(mol_exps)
+                env.append(norm)
                 pexp += nprim * 2
-            else:
-                bs[:,PTR_EXP] = _env.size + numpy.arange(nprim)
-                bs[:,PTR_COEFF] = _env.size + numpy.arange(nprim, nprim*2)
-                _env = np.hstack([_env, mol_exps, norm])
-            _bas.append(bs)
+                _bas.append(bs)

-            c = numpy.einsum('pi,p,xm->pxim', b_coeff, 1./norm, c2s[l])
-            contr_coeff.append(c.reshape(nprim * c2s[l].shape[0], -1))
+                c = numpy.einsum('pi,p,xm->pxim', b_coeff, 1./norm, c2s[l])
+                contr_coeff.append(c.reshape(nprim * c2s[l].shape[0], -1))

     pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
-    pmol._env = _env
+    pmol._env = numpy.hstack(env)
+    if aggregate:
+        contr_coeff = scipy.linalg.block_diag(*contr_coeff)
     return pmol, contr_coeff

 def format_ecp(ecp_tab):
diff --git a/pyscf/pbc/df/ft_ao.py b/pyscf/pbc/df/ft_ao.py
index 4215251fdd..b85524abb7 100644
--- a/pyscf/pbc/df/ft_ao.py
+++ b/pyscf/pbc/df/ft_ao.py
@@ -526,8 +526,9 @@ def assign(type_code):
         assign(SMOOTH_BASIS)
         return ao_type

-    def decontract_basis(self, to_cart=True):
-        pcell, ctr_coeff = self.ref_cell.decontract_basis(to_cart=to_cart)
+    def decontract_basis(self, to_cart=True, aggregate=False):
+        pcell, ctr_coeff = self.ref_cell.decontract_basis(to_cart=to_cart,
+                                                          aggregate=aggregate)
         pcell = pcell.view(self.__class__)
         pcell.ref_cell = None

diff --git a/pyscf/pbc/dft/multigrid/multigrid.py b/pyscf/pbc/dft/multigrid/multigrid.py
index 23f2bbb83f..2bcb533302 100644
--- a/pyscf/pbc/dft/multigrid/multigrid.py
+++ b/pyscf/pbc/dft/multigrid/multigrid.py
@@ -601,16 +601,16 @@ def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), deriv=0,
             dms_ht = numpy.asarray(dms[:,:,idx_h[:,None],idx_t], order='C')
             dms_lh = numpy.asarray(dms[:,:,idx_l[:,None],idx_h], order='C')

-            t_cell = h_cell + grids_sparse.cell
+            l_cell = grids_sparse.cell
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-            t_cell, t_coeff = t_cell.to_uncontracted_cartesian_basis()

             if deriv == 0:
-                h_coeff = scipy.linalg.block_diag(*t_coeff[:h_cell.nbas])
-                l_coeff = scipy.linalg.block_diag(*t_coeff[h_cell.nbas:])
-                t_coeff = scipy.linalg.block_diag(*t_coeff)
-
                 if hermi == 1:
                     naol, naoh = dms_lh.shape[2:]
                     dms_ht[:,:,:,naoh:] += dms_lh.transpose(0,1,3,2)
@@ -635,10 +635,6 @@ def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), deriv=0,
                                          'LDA', kpts, grids_dense, ignore_imag, log)

             elif deriv == 1:
-                h_coeff = scipy.linalg.block_diag(*t_coeff[:h_cell.nbas])
-                l_coeff = scipy.linalg.block_diag(*t_coeff[h_cell.nbas:])
-                t_coeff = scipy.linalg.block_diag(*t_coeff)
-
                 pgto_dms = lib.einsum('nkij,pi,qj->nkpq', dms_ht, h_coeff, t_coeff)
                 shls_slice = (0, nshells_h, 0, nshells_t)
                 #:rho = eval_rho(t_cell, pgto_dms, shls_slice, 0, 'GGA', kpts,
@@ -895,13 +891,13 @@ def _get_j_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):

             h_cell = grids_dense.cell
             l_cell = grids_sparse.cell
-            t_cell = h_cell + l_cell
-            t_cell, coeff = t_cell.to_uncontracted_cartesian_basis()
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-
-            h_coeff = scipy.linalg.block_diag(*coeff[:h_cell.nbas])
-            t_coeff = scipy.linalg.block_diag(*coeff)
             shls_slice = (0, nshells_h, 0, nshells_t)
             vp = eval_mat(t_cell, vR, shls_slice, 1, 0, 'LDA', kpts)
             # Imaginary part may contribute
@@ -919,7 +915,6 @@ def _get_j_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):
                 vj_kpts[:,:,idx_l[:,None],idx_h] += \
                         vp[:,:,:,naoh:].transpose(0,1,3,2).conj()
             else:
-                l_coeff = scipy.linalg.block_diag(*coeff[h_cell.nbas:])
                 shls_slice = (nshells_h, nshells_t, 0, nshells_h)
                 vp = eval_mat(t_cell, vR, shls_slice, 1, 0, 'LDA', kpts)
                 # Imaginary part may contribute
@@ -996,15 +991,13 @@ def _get_gga_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):

             h_cell = grids_dense.cell
             l_cell = grids_sparse.cell
-            t_cell = h_cell + l_cell
-            t_cell, coeff = t_cell.to_uncontracted_cartesian_basis()
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-
-            h_coeff = scipy.linalg.block_diag(*coeff[:h_cell.nbas])
-            l_coeff = scipy.linalg.block_diag(*coeff[h_cell.nbas:])
-            t_coeff = scipy.linalg.block_diag(*coeff)
-
             shls_slice = (0, nshells_h, 0, nshells_t)
             vpR = eval_mat(t_cell, vR, shls_slice, 1, 0, 'GGA', kpts)
             vp = vpR = lib.einsum('nkpq,pi,qj->nkij', vpR, h_coeff, t_coeff)
diff --git a/pyscf/x2c/x2c.py b/pyscf/x2c/x2c.py
index 291cd8e40e..4bd83298d3 100644
--- a/pyscf/x2c/x2c.py
+++ b/pyscf/x2c/x2c.py
@@ -810,8 +810,7 @@ def to_gpu(self):

 def _uncontract_mol(mol, xuncontract=None, exp_drop=0.2):
     '''mol._basis + uncontracted steep functions'''
-    pmol, contr_coeff = mol.decontract_basis(atoms=xuncontract)
-    contr_coeff = scipy.linalg.block_diag(*contr_coeff)
+    pmol, contr_coeff = mol.decontract_basis(atoms=xuncontract, aggregate=True)
     return pmol, contr_coeff

diff --git a/pyscf/gto/test/test_mole.py b/pyscf/gto/test/test_mole.py
index 1a91682db0..b0ffeb993c 100644
--- a/pyscf/gto/test/test_mole.py
+++ b/pyscf/gto/test/test_mole.py
@@ -1037,14 +1037,27 @@ def test_decontract_basis(self):
         mol._bas = mol._bas[:5]
         pmol, c = mol.decontract_basis()
         self.assertEqual(pmol.nbas, 14)
+        self.assertEqual(len(c), 5)

         mol = gto.M(atom='He',
                     basis=('ccpvdz', [[0, [5, 1]], [1, [3, 1]]]))
         pmol, contr_coeff = mol.decontract_basis()
+        self.assertEqual(len(contr_coeff), 5)
         contr_coeff = scipy.linalg.block_diag(*contr_coeff)
         s = contr_coeff.T.dot(pmol.intor('int1e_ovlp')).dot(contr_coeff)
         self.assertAlmostEqual(abs(s - mol.intor('int1e_ovlp')).max(), 0, 12)

+        mol = gto.M(atom='H 0 0 0; F 0 0 1', basis=[[0, (2, .5), (1, .5)],
+                                                    [0, (2, .1), (1, .9)],
+                                                    [0, (4., 1)]])
+        with self.assertRaises(RuntimeError):
+            mol.decontract_basis(aggregate=False)
+        pmol, c = mol.decontract_basis(aggregate=True)
+        self.assertEqual(pmol.nbas, 6)
+        self.assertEqual(c.shape, (6, 6))
+        s = c.T.dot(pmol.intor('int1e_ovlp')).dot(c)
+        self.assertAlmostEqual(abs(s - mol.intor('int1e_ovlp')).max(), 0, 12)
+
     def test_ao_rotation_matrix(self):
         mol = gto.M(atom='O 0 0 0.2; H1 0 -.8 -.5; H2 0 .8 -.5', basis='ccpvdz')
         numpy.random.seed(1)


Enable spin-separated 3-RDMs (pyscf/pyscf#2357)

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.

Spin-separated 3-RDMs?

Is there an easy way to access spin-separated 3-RDMs?

Spin-separated 1 and 2 RDMs are supported through the `make_rdm12s` method in the `direct_spin1` fcisolver. However, only spin-traced 3/4-RDMs are available through `fci.rdm.make_dm123(4)`.

From what I can gather this is due to a lack of spin-separated RDM kernels for 3/4 RDMs, but spin-separated intermediates (t1/t2) seem to be available, so it might be trivial to implement the spin-separated kernels. However the RDM library is a little too opaque for me to get a PR started. Any help on this would be appreciated!

### 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_rdm123(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_rdm123s(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_rdm123(fcivec, norb, nelec, link_index, reorder) -> numpy.ndarray: ...
def make_rdm123s(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 b1980bae69..28e207b918 100644
--- a/examples/fci/14-density_matrix.py
+++ b/examples/fci/14-density_matrix.py
@@ -125,6 +125,18 @@
                                    (nelec_a,nelec_b))
 dm1, dm2, dm3 = fci.rdm.reorder_dm123(dm1, dm2, dm3)

+#
+# Spin-separated 3-particle density matrix
+#
+
+(dm1a, dm1b), (dm2aa, dm2ab, dm2bb), (dm3aaa, dm3aab, dm3abb, dm3bbb) = \
+        fci.direct_spin1.make_rdm123s(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))
+
 #
 # Spin-traced 3-particle transition density matrix
 #
diff --git a/examples/mcscf/04-density_matrix.py b/examples/mcscf/04-density_matrix.py
index 32ea6b8aab..ea756e1f48 100644
--- a/examples/mcscf/04-density_matrix.py
+++ b/examples/mcscf/04-density_matrix.py
@@ -31,4 +31,3 @@
 dm1_alpha, dm1_beta = mycas.make_rdm1s()

 print(numpy.allclose(dm1, dm1_alpha+dm1_beta))
-
diff --git a/pyscf/fci/addons.py b/pyscf/fci/addons.py
index 2c8e00178c..8119c80012 100644
--- a/pyscf/fci/addons.py
+++ b/pyscf/fci/addons.py
@@ -703,6 +703,121 @@ def transform_ci(ci, nelec, u):
     ci = lib.dot(ci, trans_ci_b)
     return ci

+def civec_spinless_repr_generator(ci0_r, norb, nelec_r):
+    '''Put CI vectors in the spinless representation; i.e., map
+        norb -> 2 * norb
+        (neleca, nelecb) -> (neleca+nelecb, 0)
+    This permits linear combinations of CI vectors with different
+    M == neleca-nelecb at the price of higher memory cost. This function
+    does NOT change the datatype.
+
+    Args:
+        ci0_r: sequence or generator of ndarray of length nprods
+            CAS-CI vectors in the spin-pure representation
+        norb: integer
+            Number of orbitals
+        nelec_r: sequence of tuple of length (2)
+            (neleca, nelecb) for each element of ci0_r
+
+    Returns:
+        ci1_r_gen: callable that returns a generator of length nprods
+            generates spinless CAS-CI vectors
+        ss2spinless: callable
+            Put a CAS-CI vector in the spinless representation
+            Args:
+                ci0: ndarray
+                    CAS-CI vector
+                ne: tuple of length 2
+                    neleca, nelecb of target Hilbert space
+            Returns:
+                ci1: ndarray
+                    spinless CAS-CI vector
+        spinless2ss: callable
+            Perform the reverse operation on a spinless CAS-CI vector
+            Args:
+                ci2: ndarray
+                    spinless CAS-CI vector
+                ne: tuple of length 2
+                    neleca, nelecb target Hilbert space
+            Returns:
+                ci3: ndarray
+                    CAS-CI vector of ci2 in the (neleca, nelecb) Hilbert space
+    '''
+    nelec_r_tot = [sum (n) for n in nelec_r]
+    if len(set(nelec_r_tot)) > 1:
+        raise NotImplementedError("Different particle-number subspaces")
+    nelec = nelec_r_tot[0]
+    addrs = {}
+    ndet_sp = {}
+    for ne in set(nelec_r):
+        neleca, nelecb = _unpack_nelec(ne)
+        ndeta = cistring.num_strings(norb, neleca)
+        ndetb = cistring.num_strings(norb, nelecb)
+        strsa = cistring.addrs2str(norb, neleca, list(range(ndeta)))
+        strsb = cistring.addrs2str(norb, nelecb, list(range(ndetb)))
+        strs = numpy.add.outer(strsa, numpy.left_shift(strsb, norb)).ravel()
+        addrs[ne] = cistring.strs2addr(2*norb, nelec, strs)
+        ndet_sp[ne] = tuple((ndeta,ndetb))
+    strs = strsa = strsb = None
+    ndet = cistring.num_strings(2*norb, nelec)
+    def ss2spinless(ci0, ne, buf=None):
+        if buf is None:
+            ci1 = numpy.empty(ndet, dtype=ci0.dtype)
+        else:
+            ci1 = numpy.asarray(buf).flat[:ndet]
+        ci1[:] = 0.0
+        ci1[addrs[ne]] = ci0[:,:].ravel ()
+        neleca, nelecb = _unpack_nelec (ne)
+        if abs(neleca*nelecb)%2: ci1[:] *= -1
+        # Sign comes from changing representation:
+        # ... a2' a1' a0' ... b2' b1' b0' |vac>
+        # ->
+        # ... b2' b1' b0' .. a2' a1' a0' |vac>
+        # i.e., strictly decreasing from left to right
+        # (the ordinality of spin-down is conventionally greater than spin-up)
+        return ci1[:,None]
+    def spinless2ss(ci2, ne):
+        ''' Generate the spin-separated CI vector in a particular M
+        Hilbert space from a spinless CI vector '''
+        ci3 = ci2[addrs[ne]].reshape(ndet_sp[ne])
+        neleca, nelecb = _unpack_nelec (ne)
+        if abs(neleca*nelecb)%2: ci3[:] *= -1
+        return ci3
+    def ci1_r_gen(buf=None):
+        if callable(ci0_r):
+            ci0_r_gen = ci0_r()
+        else:
+            ci0_r_gen = (c for c in ci0_r)
+        for ci0, ne in zip(ci0_r_gen, nelec_r):
+            # Doing this in two lines saves memory: ci0 is overwritten
+            ci0 = ss2spinless(ci0, ne)
+            yield ci0
+    return ci1_r_gen, ss2spinless, spinless2ss
+
+def civec_spinless_repr(ci0_r, norb, nelec_r):
+    '''Put CI vectors in the spinless representation; i.e., map
+        norb -> 2 * norb
+        (neleca, nelecb) -> (neleca+nelecb, 0)
+    This permits linear combinations of CI vectors with different
+    M == neleca-nelecb at the price of higher memory cost. This function
+    does NOT change the datatype.
+
+    Args:
+        ci0_r: sequence or generator of ndarray of length nprods
+            CAS-CI vectors in the spin-pure representation
+        norb: integer
+            Number of orbitals
+        nelec_r: sequence of tuple of length (2)
+            (neleca, nelecb) for each element of ci0_r
+
+    Returns:
+        ci1_r: ndarray of shape (nprods, ndet_spinless)
+            spinless CAS-CI vectors
+    '''
+    ci1_r_gen, *_ = civec_spinless_repr_generator(ci0_r, norb, nelec_r)
+    ci1_r = numpy.stack([x.copy() for x in ci1_r_gen()], axis=0)
+    return ci1_r
+

 def _unpack_nelec(nelec, spin=None):
     if spin is None:
diff --git a/pyscf/fci/direct_spin1.py b/pyscf/fci/direct_spin1.py
index 784a557561..92b92af144 100644
--- a/pyscf/fci/direct_spin1.py
+++ b/pyscf/fci/direct_spin1.py
@@ -50,7 +50,7 @@
 from pyscf.fci import spin_op
 from pyscf.fci import addons
 from pyscf.fci.spin_op import contract_ss
-from pyscf.fci.addons import _unpack_nelec
+from pyscf.fci.addons import _unpack_nelec, civec_spinless_repr
 from pyscf import __config__

 libfci = cistring.libfci
@@ -371,6 +371,46 @@ def make_rdm12(fcivec, norb, nelec, link_index=None, reorder=True):
         dm1, dm2 = rdm.reorder_rdm(dm1, dm2, inplace=True)
     return dm1, dm2

+def make_rdm123(fcivec, norb, nelec, link_index=None, reorder=True):
+    '''Spin traced 1-, 2-, and 3-particle density matrices.'''
+    dm1, dm2, dm3 = rdm.make_dm123('FCI3pdm_kern_sf', fcivec, fcivec, norb, nelec)
+    if reorder:
+        dm1, dm2, dm3 = rdm.reorder_dm123(dm1, dm2, dm3, inplace=True)
+    return dm1, dm2, dm3
+
+def make_rdm123s(fcivec, norb, nelec, link_index=None, reorder=True):
+    r'''Spin separated 1-, 2-, and 3-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 = make_rdm123(ci_spinless, norb*2, (nelec[0]+nelec[1],0))
+
+    rdm1a = rdm1[:norb,:norb]
+    rdm1b = rdm1[norb:,norb:]
+    # assert np.allclose(rdm1a+rdm1b, rdm1)
+
+    rdm2aa = rdm2[:norb,:norb,:norb,:norb]
+    rdm2ab = rdm2[:norb,:norb,norb:,norb:]
+    rdm2bb = rdm2[norb:,norb:,norb:,norb:]
+    # assert np.allclose(rdm2aa+rdm2bb+rdm2ab+rdm2ab.transpose(2,3,0,1), rdm2)
+
+    rdm3aaa = rdm3[:norb,:norb,:norb,:norb,:norb,:norb]
+    rdm3aab = rdm3[:norb,:norb,:norb,:norb,norb:,norb:]
+    rdm3abb = rdm3[:norb,:norb,norb:,norb:,norb:,norb:]
+    rdm3bbb = rdm3[norb:,norb:,norb:,norb:,norb:,norb:]
+    # 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)
+    return (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb), (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb)
+
+
 def trans_rdm1s(cibra, ciket, norb, nelec, link_index=None):
     r'''Spin separated transition 1-particle density matrices.
     The return values include two density matrices: (alpha,alpha), (beta,beta).
@@ -890,6 +930,16 @@ def make_rdm12(self, fcivec, norb, nelec, link_index=None, reorder=True):
         nelec = _unpack_nelec(nelec, self.spin)
         return make_rdm12(fcivec, norb, nelec, link_index, reorder)

+    @lib.with_doc(make_rdm123s.__doc__)
+    def make_rdm123s(self, fcivec, norb, nelec, link_index=None, reorder=True):
+        nelec = _unpack_nelec(nelec, self.spin)
+        return make_rdm123s(fcivec, norb, nelec, link_index, reorder)
+
+    @lib.with_doc(make_rdm123.__doc__)
+    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)
+
     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 e28ab38d2e..8ac3b59957 100644
--- a/pyscf/fci/test/test_rdm.py
+++ b/pyscf/fci/test/test_rdm.py
@@ -20,6 +20,7 @@
 from pyscf import scf
 from pyscf import ao2mo
 from pyscf import fci
+from pyscf import mcscf

 def setUpModule():
     global ci0, rdm1, rdm2, norb, nelec
@@ -62,6 +63,52 @@ def test_rdm3(self):
         dm3 = fci.rdm.make_dm123('FCI3pdm_kern_sf', ci1, ci1, norb, (5,3))[2]
         self.assertTrue(numpy.allclose(dm3ref, dm3))

+    def test_rdm3s(self):
+        dm1, dm2, dm3 = fci.direct_spin1.make_rdm123(ci0, norb, (nelec//2,nelec//2))
+
+        (dm1a, dm1b), (dm2aa, dm2ab, dm2bb), (dm3aaa, dm3aab, dm3abb, dm3bbb) = \
+        fci.direct_spin1.make_rdm123s(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))
+
+        (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb) = \
+        fci.direct_spin1.make_rdm12s(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))
+
     def test_dm4(self):
         dm4ref = make_dm4_o0(ci0, norb, nelec)
         dm4 = fci.rdm.make_dm1234('FCI4pdm_kern_sf', ci0, ci0, norb, nelec)[3]


Fix define_xc example (pyscf/pyscf#2364)

Update examples: 24-define_xc_functional.py and 31-xc_value_on_grid.py . (Fix issue #2092 , #2011)

Obtaining Hartree-Fock exchange energy density on grids

Hello pySCF team, I hope the new year is treating you well.

I'm trying to obtain Hartree-Fock exchange energy densities on grids (as the title would suggest). While there was an issue ( #1485 ) posted and seemingly resolved a little over a year ago, there was no input from any pySCF devs and the issue was closed the same day that it was posted with no commentary from anyone other than the author. I've tried executing this method to no success, as it doesn't seem to return values that get us the energy when combined with rho and the grid weights (as you would be able to with exc in a Kohn-Sham calculation).

In short, I'm looking for an analogue to exc from KS, but in HF (an ex, if you will).

If there is any insight on how this may be accomplished, please let me know and I look forward to your responses!

Unable to use custom xc functional

Hi PySCF developers,

I'm new to using PySCF. I wanted to plugin my own XC functional using the define_xc_() interface in the dft module. To begin with, I attempted to run the example that demonstrates the use of custom XC functional (https://github.com/pyscf/pyscf/blob/master/examples/dft/24-define_xc_functional.py). However, running it throws the following error:

```
Traceback (most recent call last):
  File "custom_xc_test.py", line 63, in <module>
    mf.kernel()
  File "<string>", line 2, in kernel
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/scf/hf.py", line 1779, in scf
    kernel(self, self.conv_tol, self.conv_tol_grad,
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/scf/hf.py", line 124, in kernel
    vhf = mf.get_veff(mol, dm)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/rks.py", line 80, in get_veff
    n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 1158, in nr_rks
    for i, ao, mask, wv in block_loop(ao_deriv):
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 1138, in block_loop
    exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype)[:2]
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/numint.py", line 2696, in eval_xc_eff
    out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/libxc.py", line 1599, in eval_xc1
    out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
  File "/home/biku/miniconda3/lib/python3.8/site-packages/pyscf/dft/libxc.py", line 1654, in _eval_xc
    rho = rho.reshape(spin+1,nvar,ngrids)
ValueError: cannot reshape array of size 120224 into shape (1,1,30056)
```

The custom_xc_test.py is just a copy of the example link provided above.

I used `pip install --prefer-binary pyscf` to install my PySCF.

I would really appreciate it if you can help me figure out how to use a custom XC functional.

Regards,
Bikash

diff --git a/examples/dft/24-define_xc_functional.py b/examples/dft/24-define_xc_functional.py
index e4efae3089..72c7f25cfd 100644
--- a/examples/dft/24-define_xc_functional.py
+++ b/examples/dft/24-define_xc_functional.py
@@ -50,11 +50,11 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
     kxc = None  # 3rd order functional derivative

     # Mix with existing functionals
-    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho, spin, relativity, deriv,
-                               verbose)
+    pbe_xc = dft.libxc.eval_xc('pbe,pbe', rho, 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

 mf = dft.RKS(mol)
@@ -70,4 +70,3 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=Non
 mf = mf.define_xc_(eval_xc, 'GGA', rsh=rsh_coeff)
 mf.verbose = 4
 mf.kernel()
-
diff --git a/examples/dft/31-xc_value_on_grid.py b/examples/dft/31-xc_value_on_grid.py
index c78592334e..d0e7a980c9 100644
--- a/examples/dft/31-xc_value_on_grid.py
+++ b/examples/dft/31-xc_value_on_grid.py
@@ -5,7 +5,7 @@
 potential on given grid coordinates.
 '''

-import numpy
+import numpy as np
 from pyscf import gto, dft, lib
 from pyscf.dft import numint
 from pyscf.dft import r_numint
@@ -37,19 +37,19 @@
 #
 ex, vx = dft.libxc.eval_xc('B88', rho)[:2]
 ec, vc = dft.libxc.eval_xc(',P86', rho)[:2]
-print('Exc = %.12f' % numpy.einsum('i,i,i->', ex+ec, rho[0], weights))
+print('Exc = %.12f' % np.einsum('i,i,i->', ex+ec, rho[0], weights))

 #
 # Evaluate XC functional together
 #
 exc, vxc = dft.libxc.eval_xc('B88,P86', rho)[:2]
-print('Exc = %.12f' % numpy.einsum('i,i,i->', exc, rho[0], weights))
+print('Exc = %.12f' % np.einsum('i,i,i->', exc, rho[0], weights))

 #
 # Evaluate XC functional for user specified functional
 #
 exc, vxc = dft.libxc.eval_xc('.2*HF + .08*SLATER + .72*B88, .81*LYP + .19*VWN', rho)[:2]
-print('Exc = %.12f  ref = -7.520014202688' % numpy.einsum('i,i,i->', exc, rho[0], weights))
+print('Exc = %.12f  ref = -7.520014202688' % np.einsum('i,i,i->', exc, rho[0], weights))

 #
 # For density matrix computed with GKS in spin-orbital basis, plain density in
@@ -82,3 +82,20 @@
 rhoS, mS = r_numint.eval_rho(mol, aoS_value, dmSS)
 rho = rhoL + rhoS
 mx, my, mz = mL + mS
+
+#
+# The values of exact exchange
+# Exx(r) = \sum_{ij} 2 (ij(r)|ji)
+#
+mf = mol.RHF().run()
+# Compute analytically v_ij(r') = \int \chi_i(r) \chi_j(r) / |r-r'| dr
+rpq = mol.intor('int1e_grids', grids=coords)
+ao = mol.eval_gto('GTOval', coords)
+nocc = np.count_nonzero(mf.mo_occ)
+occ_orb = mf.mo_coeff[:,:nocc]
+mo = ao.dot(occ_orb)
+exx = 2 * np.einsum('rpq,ri,rj,pj,qi->r', rpq, mo, mo, occ_orb, occ_orb)
+print('Semi-numerical HFX:', exx.dot(weights))
+dm = mf.make_rdm1()
+vj, vk = mf.get_jk(dm)
+print('Analytical HFX:', .5 * np.einsum('pq,qp->', vk, dm))
diff --git a/examples/fci/03-cylindrical_symmetry.py b/examples/fci/03-cylindrical_symmetry.py
index 78ad41eca7..3b3c382f88 100644
--- a/examples/fci/03-cylindrical_symmetry.py
+++ b/examples/fci/03-cylindrical_symmetry.py
@@ -26,5 +26,5 @@
 solver = direct_spin1_cyl_sym.FCI(mol)
 for sym in ['A1g', 'A1u', 'E1ux', 'E1uy', 'E1gx', 'E1gy', 'E2ux', 'E2uy', 'E2gx', 'E2gy', 'E3ux', 'E3uy', 'E3gx', 'E3gy']:
     e, v = solver.kernel(h1e, eri, c.shape[1], mol.nelec, orbsym=orbsym,
-                         wfnsym=sym, nroots=5, verbose=0)
+                         wfnsym=sym, nroots=5, ecore=mol.energy_nuc(), verbose=0)
     print(f'{sym}: {e}')
diff --git a/pyscf/dft/libxc.py b/pyscf/dft/libxc.py
index 0f248fb505..26aef161bf 100644
--- a/pyscf/dft/libxc.py
+++ b/pyscf/dft/libxc.py
@@ -1609,13 +1609,17 @@ def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
     '''
     out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
     xctype = xc_type(xc_code)
+    idx = _libxc_to_xcfun_indices(xctype, spin, deriv)
+    return out[idx]
+
+def _libxc_to_xcfun_indices(xctype, spin=0, deriv=1):
     if deriv <= 1:
-        return out
+        return slice(None)
     elif xctype == 'LDA' or xctype == 'HF':
-        return out
+        return slice(None)
     elif xctype == 'GGA':
         if spin == 0:
-            return out
+            return slice(None)
         else:
             idx = [numpy.arange(6)] # up to deriv=1
             for i in range(2, deriv+1):
@@ -1627,7 +1631,7 @@ def eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
             idx = [numpy.arange(8)] # up to deriv=1
         for i in range(2, deriv+1):
             idx.append(_MGGA_SORT[(spin, i)])
-    return out[numpy.hstack(idx)]
+    return numpy.hstack(idx)

 def _eval_xc(xc_code, rho, spin=0, deriv=1, omega=None):
     assert deriv <= max_deriv_order(xc_code)
@@ -1762,19 +1766,43 @@ def define_xc_(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
     48.8525211046668
     '''
     if isinstance(description, str):
-        ni.eval_xc = lambda xc_code, rho, *args, **kwargs: \
-                eval_xc(description, rho, *args, **kwargs)
+        def _eval_xc(xc_code, rho, *args, **kwargs):
+            return eval_xc(description, rho, *args, **kwargs)
+        ni.eval_xc = _eval_xc
         ni.hybrid_coeff = lambda *args, **kwargs: hybrid_coeff(description)
         ni.rsh_coeff = lambda *args: rsh_coeff(description)
         ni._xc_type = lambda *args: xc_type(description)

     elif callable(description):
-        ni.eval_xc = description
+        ni.eval_xc = _eval_xc = description
         ni.hybrid_coeff = lambda *args, **kwargs: hyb
         ni.rsh_coeff = lambda *args, **kwargs: rsh
         ni._xc_type = lambda *args: xctype
+
     else:
         raise ValueError('Unknown description %s' % description)
+
+    def _eval_xc1(xc_code, rho, spin=0, deriv=1, omega=None):
+        libxc_out = _eval_xc(xc_code, rho, spin, deriv=deriv, omega=omega)
+        nvar, xlen = xc_deriv._XC_NVAR[xctype, spin]
+        outlen = lib.comb(xlen+deriv, deriv)
+        exc, vxc, fxc, kxc = libxc_out[:4]
+        out = [exc]
+        if 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:
+            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)
+        return out[idx]
+    ni.eval_xc1 = _eval_xc1
     return ni

 def define_xc(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
diff --git a/pyscf/dft/numint.py b/pyscf/dft/numint.py
index 7eef303d2c..0f9dae11e5 100644
--- a/pyscf/dft/numint.py
+++ b/pyscf/dft/numint.py
@@ -2654,6 +2654,10 @@ def eval_xc(self, xc_code, rho, spin=0, relativity=0, deriv=1, omega=None,
         return self.libxc.eval_xc(xc_code, rho, spin, relativity, deriv,
                                   omega, verbose)

+    def eval_xc1(self, xc_code, rho, spin=0, deriv=1, omega=None):
+        if omega is None: omega = self.omega
+        return self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
+
     def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None,
                     verbose=None):
         r'''Returns the derivative tensor against the density parameters
@@ -2693,12 +2697,12 @@ def eval_xc_eff(self, xc_code, rho, deriv=1, omega=None, xctype=None,
         else:
             spin = 0

-        out = self.libxc.eval_xc1(xc_code, rho, spin, deriv, omega)
+        out = self.eval_xc1(xc_code, rho, spin, deriv, omega)
         evfk = [out[0]]
         for order in range(1, deriv+1):
             evfk.append(xc_deriv.transform_xc(rho, out, xctype, spin, order))
         if deriv < 3:
-            # The return has at least [e, v, f, k] terms
+            # Returns at least [e, v, f, k] terms
             evfk.extend([None] * (3 - deriv))
         return evfk

diff --git a/pyscf/dft/test/test_libxc.py b/pyscf/dft/test/test_libxc.py
index e94bfbd39e..cb4d46bbe9 100644
--- a/pyscf/dft/test/test_libxc.py
+++ b/pyscf/dft/test/test_libxc.py
@@ -22,7 +22,7 @@
 from pyscf import lib

 def setUpModule():
-    global mol, mf, ao, rho
+    global mol, mf, ao, rho, dm
     mol = gto.Mole()
     mol.verbose = 0
     mol.output = None
@@ -206,7 +206,7 @@ def test_lyp(self):
     #    e, v = dft.libxc.eval_xc('tpss,', (rho_a, rho_b), spin=1, deriv=1)[:2]

     def test_define_xc(self):
-        def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
+        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]
             gamma = (dx**2 + dy**2 + dz**2)
@@ -229,6 +229,11 @@ def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
         self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
         self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)

+        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)
+
         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]


Fix SGX initialization issue (pyscf/pyscf#2371)

SGX Fails When PJS=True

This is benzene, HF/def2-SVP -

```
import pyscf

xyz = """C                 -0.71618036    2.71883299    0.00000000
C                  0.67897964    2.71883299    0.00000000
C                  1.37651764    3.92658399    0.00000000
C                  0.67886364    5.13509299   -0.00119900
C                 -0.71596136    5.13501499   -0.00167800
C                 -1.41356236    3.92680899   -0.00068200
H                 -1.26593936    1.76651599    0.00045000
H                  1.22848764    1.76631999    0.00131500
H                  2.47619764    3.92666399    0.00063400
H                  1.22906364    6.08723599   -0.00125800
H                 -1.26608336    6.08729599   -0.00263100
H                 -2.51316636    3.92699199   -0.00086200"""

mol = pyscf.M(
    atom=xyz,
    basis="def2-SVP",
    charge=0,
    spin=0,
    symmetry=True
)

mf = pyscf.scf.HF(mol)

mf = pyscf.sgx.sgx_fit(mf, pjs=True)

energy = mf.kernel()
```

And here's what I get:

```
Traceback (most recent call last):
  File "/Users/cwagen/business_code/pyscf/sgx_fail.py", line 32, in <module>
    energy = mf.kernel()
             ^^^^^^^^^^^
  File "<string>", line 2, in kernel
  File "/opt/miniconda3/envs/aimnet2/lib/python3.11/site-packages/pyscf/scf/hf.py", line 1713, in scf
    self.build(self.mol)
  File "/opt/miniconda3/envs/aimnet2/lib/python3.11/site-packages/pyscf/sgx/sgx.py", line 141, in build
    self.with_df.build(level=self.with_df.grids_level_i)
  File "/opt/miniconda3/envs/aimnet2/lib/python3.11/site-packages/pyscf/sgx/sgx.py", line 316, in build
    self._opt = _make_opt(self.mol, pjs=self.pjs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/aimnet2/lib/python3.11/site-packages/pyscf/sgx/sgx.py", line 239, in _make_opt
    vhfopt = _vhf.SGXOpt(mol, 'int1e_grids', 'SGXnr_ovlp_prescreen',
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: SGXOpt.__init__() got an unexpected keyword argument 'direct_scf_tol'
```

This works fine when PJS = False.

diff --git a/pyscf/scf/_vhf.py b/pyscf/scf/_vhf.py
index 7b51aa5921..313407a957 100644
--- a/pyscf/scf/_vhf.py
+++ b/pyscf/scf/_vhf.py
@@ -79,10 +79,10 @@ def init_cvhf_direct(self, mol, intor, qcondname):

     @property
     def direct_scf_tol(self):
-        return self._this.contents.direct_scf_cutoff
+        return self._this.contents.direct_scf_tol
     @direct_scf_tol.setter
     def direct_scf_tol(self, v):
-        self._this.contents.direct_scf_cutoff = v
+        self._this.contents.direct_scf_tol = v

     @property
     def prescreen(self):
@@ -162,7 +162,7 @@ def __init__(self, mol, intor=None, prescreen='CVHFnoscreen',
         self.mol = mol
         self._this = cvhfopt = _CVHFOpt()
         cvhfopt.nbas = mol.nbas
-        cvhfopt.direct_scf_cutoff = direct_scf_tol
+        cvhfopt.direct_scf_tol = direct_scf_tol
         cvhfopt.fprescreen = _fpointer(prescreen)
         cvhfopt.r_vkscreen = _fpointer('CVHFr_vknoscreen')
         self._q_cond = None
@@ -207,10 +207,10 @@ def init_cvhf_direct(self, mol, intor, qcondname):

     @property
     def direct_scf_tol(self):
-        return self._this.direct_scf_cutoff
+        return self._this.direct_scf_tol
     @direct_scf_tol.setter
     def direct_scf_tol(self, v):
-        self._this.direct_scf_cutoff = v
+        self._this.direct_scf_tol = v

     @property
     def prescreen(self):
@@ -265,9 +265,9 @@ def dm_cond(self, dm_cond):

 class SGXOpt(_VHFOpt):
     def __init__(self, mol, intor=None, prescreen='CVHFnoscreen',
-                 qcondname=None, dmcondname=None, direct_scf_cutoff=1e-14):
+                 qcondname=None, dmcondname=None, direct_scf_tol=1e-14):
         _VHFOpt.__init__(self, mol, intor, prescreen, qcondname, dmcondname,
-                         direct_scf_cutoff)
+                         direct_scf_tol)
         self.ngrids = None

     def set_dm(self, dm, atm, bas, env):
@@ -287,18 +287,22 @@ def set_dm(self, dm, atm, bas, env):
             if self._dmcondname != 'SGXnr_dm_cond':
                 raise ValueError('SGXOpt only supports SGXnr_dm_cond')
             fdmcond = getattr(libcvhf, self._dmcondname)
-        dm_cond = numpy.empty((mol.nbas, self.ngrids))
+        if self.ngrids is None:
+            ngrids = int(env[gto.NGRIDS])
+        else:
+            ngrids = self.ngrids
+        dm_cond = numpy.empty((mol.nbas, ngrids))
         fdmcond(dm_cond.ctypes, dm.ctypes, ctypes.c_int(n_dm),
                 ao_loc.ctypes, mol._atm.ctypes, ctypes.c_int(mol.natm),
                 mol._bas.ctypes, ctypes.c_int(mol.nbas), mol._env.ctypes,
-                ctypes.c_int(self.ngrids))
+                ctypes.c_int(ngrids))
         self.dm_cond = dm_cond

 class _CVHFOpt(ctypes.Structure):
     _fields_ = [('nbas', ctypes.c_int),
                 ('ngrids', ctypes.c_int),
-                ('direct_scf_cutoff', ctypes.c_double),
+                ('direct_scf_tol', ctypes.c_double),
                 ('q_cond', ctypes.c_void_p),
                 ('dm_cond', ctypes.c_void_p),
                 ('fprescreen', ctypes.c_void_p),

diff --git a/pyscf/sgx/test/test_sgx.py b/pyscf/sgx/test/test_sgx.py
index bcf8fab256..b373fc2f19 100644
--- a/pyscf/sgx/test/test_sgx.py
+++ b/pyscf/sgx/test/test_sgx.py
@@ -39,11 +39,28 @@ def test_sgx_scf(self):
                     [1   , (0. , 0.757  , 0.587)] ],
             basis = 'ccpvdz',
         )
-        method = sgx.sgx_fit(scf.RHF(mol), 'weigend')
-        method.with_df.dfj = True
-        energy = method.scf()
+        mf = sgx.sgx_fit(scf.RHF(mol), 'weigend')
+        mf.with_df.dfj = True
+        energy = mf.kernel()
         self.assertAlmostEqual(energy, -76.02686422219752, 9)

+        mf = sgx.sgx_fit(scf.RHF(mol))
+        energy = mf.kernel()
+        self.assertAlmostEqual(energy, -76.02673747035047, 8)
+
+    def test_sgx_pjs(self):
+        mol = gto.Mole()
+        mol.build(
+            atom = [["O" , (0. , 0.     , 0.)],
+                    [1   , (0. , -0.757 , 0.587)],
+                    [1   , (0. , 0.757  , 0.587)] ],
+            basis = 'ccpvdz',
+        )
+        mf = sgx.sgx_fit(mol.RHF(), pjs=True)
+        mf.with_df.dfj = True
+        energy = mf.kernel()
+        self.assertAlmostEqual(energy, -76.0267979, 6)
+
 if __name__ == "__main__":
     print("Full Tests for SGX")
     unittest.main()


Fix transition dipoles of triplet tddft states (pyscf/pyscf#2373)

TDDFT reports non-zero oscillator strengths for spin-forbidden transitions

TDDFT (w/ or w/o TDA) will report non-zero oscillator strengths for triplet excited states when the ground state is a singlet:

```
from pyscf import gto

mol = gto.M(atom='''
  H      1.2194     -0.1652      2.1600
  C      0.6825     -0.0924      1.2087
  C     -0.7075     -0.0352      1.1973
  H     -1.2644     -0.0630      2.1393
  C     -1.3898      0.0572     -0.0114
  H     -2.4836      0.1021     -0.0204
  C     -0.6824      0.0925     -1.2088
  H     -1.2194      0.1652     -2.1599
  C      0.7075      0.0352     -1.1973
  H      1.2641      0.0628     -2.1395
  C      1.3899     -0.0572      0.0114
  H      2.4836     -0.1022      0.0205
''',
              basis='sto-3g')

mf = mol.RKS(xc='cam-b3lyp').run()
td = mf.TDA()
td.singlet = False
td.kernel()
td.analyze()
```

Output:
```
** Triplet excitation energies and oscillator strengths **
Excited State   1:      4.59810 eV    269.64 nm  f=0.0000
Excited State   2:      5.75758 eV    215.34 nm  f=1.0200
Excited State   3:      5.75772 eV    215.34 nm  f=1.0198
```

diff --git a/pyscf/tdscf/rhf.py b/pyscf/tdscf/rhf.py
index b1b680b69f..d2cb63e086 100644
--- a/pyscf/tdscf/rhf.py
+++ b/pyscf/tdscf/rhf.py
@@ -530,16 +530,20 @@ def _charge_center(mol):
     return numpy.einsum('z,zr->r', charges, coords)/charges.sum()

 def _contract_multipole(tdobj, ints, hermi=True, xy=None):
+    '''ints is the integral tensor of a spin-independent operator'''
     if xy is None: xy = tdobj.xy
+    nstates = len(xy)
+    pol_shape = ints.shape[:-2]
+    nao = ints.shape[-1]
+
+    if not tdobj.singlet:
+        return numpy.zeros((nstates,) + pol_shape)
+
     mo_coeff = tdobj._scf.mo_coeff
     mo_occ = tdobj._scf.mo_occ
     orbo = mo_coeff[:,mo_occ==2]
     orbv = mo_coeff[:,mo_occ==0]

-    nstates = len(xy)
-    pol_shape = ints.shape[:-2]
-    nao = ints.shape[-1]
-
     #Incompatible to old numpy version
     #ints = numpy.einsum('...pq,pi,qj->...ij', ints, orbo.conj(), orbv)
     ints = lib.einsum('xpq,pi,qj->xij', ints.reshape(-1,nao,nao), orbo.conj(), orbv)

diff --git a/pyscf/tdscf/test/test_tdrhf.py b/pyscf/tdscf/test/test_tdrhf.py
index 243d897bb2..b2b5ed6218 100644
--- a/pyscf/tdscf/test/test_tdrhf.py
+++ b/pyscf/tdscf/test/test_tdrhf.py
@@ -15,6 +15,7 @@
 #

 import unittest
+import numpy as np
 from pyscf import lib, gto, scf, tdscf

 def setUpModule():
@@ -41,6 +42,8 @@ def test_tda_singlet(self):
         e = td.kernel()[0]
         ref = [11.90276464, 11.90276464, 16.86036434]
         self.assertAlmostEqual(abs(e[:len(ref)] * 27.2114 - ref).max(), 0, 5)
+        dip = td.transition_dipole()
+        self.assertAlmostEqual(lib.fp(np.linalg.norm(dip, axis=1)), -0.65616659, 5)

     def test_tda_triplet(self):
         td = mf.TDA().set(nstates=nstates)
@@ -48,12 +51,16 @@ def test_tda_triplet(self):
         e = td.kernel()[0]
         ref = [11.01747918, 11.01747918, 13.16955056]
         self.assertAlmostEqual(abs(e[:len(ref)] * 27.2114 - ref).max(), 0, 5)
+        dip = td.transition_dipole()
+        self.assertAlmostEqual(abs(dip).max(), 0, 8)

     def test_tdhf_singlet(self):
         td = mf.TDHF().set(nstates=nstates)
         e = td.kernel()[0]
         ref = [11.83487199, 11.83487199, 16.66309285]
         self.assertAlmostEqual(abs(e[:len(ref)] * 27.2114 - ref).max(), 0, 5)
+        dip = td.transition_dipole()
+        self.assertAlmostEqual(lib.fp(np.linalg.norm(dip, axis=1)), -0.64009191, 5)

     def test_tdhf_triplet(self):
         td = mf.TDHF().set(nstates=nstates)
@@ -61,6 +68,8 @@ def test_tdhf_triplet(self):
         e = td.kernel()[0]
         ref = [10.8919234, 10.8919234, 12.63440705]
         self.assertAlmostEqual(abs(e[:len(ref)] * 27.2114 - ref).max(), 0, 5)
+        dip = td.transition_dipole()
+        self.assertAlmostEqual(abs(dip).max(), 0, 8)

 if __name__ == "__main__":


Fix bug in GHF-X2C dipole function (pyscf/pyscf#2396)

ValueError in x2c dip_moment picture change with GHF and GKS

When I apply the .x2c() decoration method for a GHF or GKS object, calculating dipole moments always seems to raise the following:
"ValueError: could not broadcast input array from shape (3,6,6) into shape (3,2,2)"

Here is the input:
`from pyscf import gto, scf, dft

mol = gto.M(
    verbose = 0,
    atom='H 0 0 0',
    basis='STO-3G',
    spin = 1)

mf = scf.GHF(mol).x2c()
#mf = dft.GKS(mol).x2c()
mf.kernel()
mf.dip_moment()`

And here is the full error message:
`Traceback (most recent call last):
  File "---/ghf_x2c_h_atom.py", line 13, in <module>
    mf.dip_moment()
  File "---/anaconda3/envs/Python/lib/python3.12/site-packages/pyscf/x2c/x2c.py", line 785, in dip_moment
    ao_dip = self.with_x2c.picture_change((ao_dip, prp))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "---/anaconda3/envs/Python/lib/python3.12/site-packages/pyscf/x2c/x2c.py", line 338, in picture_change
    pc_mat = self._picture_change(xmol, even_operator, odd_operator)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "---/anaconda3/envs/Python/lib/python3.12/site-packages/pyscf/x2c/x2c.py", line 182, in _picture_change
    mat[:,nao:,nao:] = w_op
    ~~~^^^^^^^^^^^^^
ValueError: could not broadcast input array from shape (3,6,6) into shape (3,2,2)`

diff --git a/pyscf/x2c/x2c.py b/pyscf/x2c/x2c.py
index 4bd83298d3..d1b0e34ed6 100644
--- a/pyscf/x2c/x2c.py
+++ b/pyscf/x2c/x2c.py
@@ -774,17 +774,22 @@ def dip_moment(self, mol=None, dm=None, unit='Debye', verbose=logger.NOTE,
         charges = mol.atom_charges()
         coords  = mol.atom_coords()
         nucl_dip = numpy.einsum('i,ix->x', charges, coords)
-        with mol.with_common_orig(nucl_dip):
-            r = mol.intor_symmetric('int1e_r')
-            ao_dip = numpy.array([_block_diag(x) for x in r])
+        with mol.with_common_orig((0,0,0)):
             if picture_change:
                 xmol = self.with_x2c.get_xmol()[0]
                 nao = xmol.nao
-                prp = xmol.intor_symmetric('int1e_sprsp').reshape(3,4,nao,nao)[:,0]
-                prp = numpy.array([_block_diag(x) for x in prp])
-                ao_dip = self.with_x2c.picture_change((ao_dip, prp))
+                r = xmol.intor_symmetric('int1e_r')
+                r = numpy.array([_block_diag(x) for x in r])
+                c1 = 0.5/lib.param.LIGHT_SPEED
+                prp = xmol.intor_symmetric('int1e_sprsp').reshape(3,4,nao,nao)
+                prp = numpy.array([_sigma_dot(x*c1**2) for x in prp])
+                ao_dip = self.with_x2c.picture_change((r, prp))
+            else:
+                r = mol.intor_symmetric('int1e_r')
+                ao_dip = numpy.array([_block_diag(x) for x in r])

-        mol_dip = -numpy.einsum('xij,ji->x', ao_dip, dm).real
+        el_dip = numpy.einsum('xij,ji->x', ao_dip, dm).real
+        mol_dip = nucl_dip - el_dip

         if unit.upper() == 'DEBYE':
             mol_dip *= nist.AU2DEBYE

diff --git a/pyscf/x2c/test/test_x2c.py b/pyscf/x2c/test/test_x2c.py
index f1d963f88a..ddca2b3206 100644
--- a/pyscf/x2c/test/test_x2c.py
+++ b/pyscf/x2c/test/test_x2c.py
@@ -215,6 +215,7 @@ def test_ghf(self):
         mf.kernel(dm0=dm)
         self.assertTrue(mf.converged)
         self.assertAlmostEqual(mf.e_tot, ref.e_tot, 9)
+        self.assertAlmostEqual(abs(mf.dip_moment() - ref.dip_moment()).max(), 0, 9)

     def test_undo_x2c(self):
         mf = mol.RHF().x2c().density_fit()


Fixes calculating exchange integrals when only_dfj=True, issue #2456 (pyscf/pyscf#2457)

This PR fixes an assertion error in the following minimal example:
```
from pyscf import gto, dft

mol = gto.M(
    atom="""
        O  0.00000000  0.00000000  0.00000000
        H -0.75390364  0.00000000 -0.58783729
        H  0.75390364  0.00000000 -0.58783729
    """
)

mf = dft.RKS(mol).density_fit(only_dfj=True).run()
vk = mf.get_k(mol=mol)
```

Calculating exchange integrals fails when only_dfj=True

Hi everyone,
the following minimal example fails due to an assertion error:
```
from pyscf import gto, dft

mol = gto.M(
    atom="""
        O  0.00000000  0.00000000  0.00000000
        H -0.75390364  0.00000000 -0.58783729
        H  0.75390364  0.00000000 -0.58783729
    """
)

mf = dft.RKS(mol).density_fit(only_dfj=True).run()
vk = mf.get_k(mol=mol)
```
It should still be possible to get the regular exchange integrals which are also returned for `get_jk` if `with_j=True` and `with_k=True`. The issue appears to be in the following function:
https://github.com/pyscf/pyscf/blob/410d9608a73ad3a507b3008a6f68549288ebfc33/pyscf/df/df_jk.py#L133-L159

I would propose the following solution:
 ```
   def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        assert (with_j or with_k)
        if dm is None: dm = self.make_rdm1()
        if not self.with_df:
            return super().get_jk(mol, dm, hermi, with_j, with_k, omega)

        with_dfk = with_k and not self.only_dfj
        if with_j:
            if isinstance(self, scf.ghf.GHF):
                def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
                    vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
                                                self.direct_scf_tol, omega)
                    if dm.dtype == numpy.complex128:
                        vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
                                                    self.direct_scf_tol, omega)
                        if with_j:
                            vj = vj + vjI * 1j
                        if with_k:
                            vk = vk + vkI * 1j
                    return vj, vk
                vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
                                        jkbuild, omega)
            else:
                vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
                                            self.direct_scf_tol, omega)
        else:
            vj = None
        if with_k and not with_dfk:
            vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
        return vj, vk
```
This change would still require that either the Coulomb or the exchange matrix are requested while ensuring the regular exchange matrix is returned whenever the corresponding density-fitted integrals are not available. I can open a PR if this behavior/fix is desired.

diff --git a/pyscf/df/df_jk.py b/pyscf/df/df_jk.py
index 2ac82c243c..972b7e2033 100644
--- a/pyscf/df/df_jk.py
+++ b/pyscf/df/df_jk.py
@@ -132,28 +132,31 @@ def reset(self, mol=None):

     def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
                omega=None):
+        assert (with_j or with_k)
         if dm is None: dm = self.make_rdm1()
         if not self.with_df:
             return super().get_jk(mol, dm, hermi, with_j, with_k, omega)

+        vj = vk = None
         with_dfk = with_k and not self.only_dfj
-        if isinstance(self, scf.ghf.GHF):
-            def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
-                vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
-                                             self.direct_scf_tol, omega)
-                if dm.dtype == numpy.complex128:
-                    vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
-                                                   self.direct_scf_tol, omega)
-                    if with_j:
-                        vj = vj + vjI * 1j
-                    if with_k:
-                        vk = vk + vkI * 1j
-                return vj, vk
-            vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
-                                    jkbuild, omega)
-        else:
-            vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
-                                         self.direct_scf_tol, omega)
+        if with_j or with_dfk:
+            if isinstance(self, scf.ghf.GHF):
+                def jkbuild(mol, dm, hermi, with_j, with_k, omega=None):
+                    vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k,
+                                                self.direct_scf_tol, omega)
+                    if dm.dtype == numpy.complex128:
+                        vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k,
+                                                    self.direct_scf_tol, omega)
+                        if with_j:
+                            vj = vj + vjI * 1j
+                        if with_k:
+                            vk = vk + vkI * 1j
+                    return vj, vk
+                vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk,
+                                        jkbuild, omega)
+            else:
+                vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk,
+                                            self.direct_scf_tol, omega)
         if with_k and not with_dfk:
             vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
         return vj, vk

diff --git a/pyscf/df/test/test_df.py b/pyscf/df/test/test_df.py
index 394d6d5a55..11b07b91be 100644
--- a/pyscf/df/test/test_df.py
+++ b/pyscf/df/test/test_df.py
@@ -142,6 +142,44 @@ def test_rsh_df_custom_storage(self):
         mf.run()
         self.assertAlmostEqual(mf.e_tot, -103.4965622991, 6)

+    def test_only_dfj(self):
+        mol = gto.M(atom='H 0 0 0; H 0 0 1')
+        dm = numpy.eye(mol.nao)
+        mf = mol.RHF().density_fit()
+        refj, refk = mf.get_jk(mol, dm)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        mf = mol.RHF().density_fit(only_dfj=True)
+        refk = mol.RHF().get_k(mol, dm)
+        vj, vk = mf.get_jk(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        dm = numpy.eye(mol.nao*2)
+        mf = mol.GHF().density_fit()
+        refj, refk = mf.get_jk(mol, dm)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
+        mf = mol.GHF().density_fit(only_dfj=True)
+        refk = mol.GHF().get_k(mol, dm)
+        vj, vk = mf.get_jk(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+        vk = mf.get_k(mol, dm)
+        self.assertAlmostEqual(abs(vk - refk).max(), 0, 9)
+        vj = mf.get_j(mol, dm)
+        self.assertAlmostEqual(abs(vj - refj).max(), 0, 9)
+
 if __name__ == "__main__":
     print("Full Tests for df")
     unittest.main()


Fix biased implementation for the becke radi method (pyscf/pyscf#2499)

This fixed the Becke radial integration scheme as described in #2490.

Becke radial integration scheme does not integrate numbers of electrons correctly

From the available radial integration schemes the Becke one seems to not create a proper integration grid as the number of electrons is overestimated:

````python
from pyscf import dft, gto

mol = gto.M(atom="He")

for radi_method in (
   dft.radi.becke,
   dft.radi.treutler_ahlrichs,
   dft.radi.delley,
   dft.radi.gauss_chebyshev,
   dft.radi.mura_knowles,
):
   ks = dft.KS(mol)(xc="pbe", grids=dft.Grids(mol)(radi_method=radi_method), verbose=0)
   ks.kernel()
   rho = ks._numint.get_rho(mol, ks.make_rdm1(), ks.grids)
   nel = (rho * ks.grids.weights).sum()
   print(radi_method.__name__, nel)
````

Output:

````
becke 2.188070936036473
treutler_ahlrichs 1.9999999979072856
delley 1.9999999990380877
gauss_chebyshev 1.9999999994477295
mura_knowles 1.9999999997029547
````

cc Thijs @tvogels, Derk @derkkooi

diff --git a/pyscf/dft/radi.py b/pyscf/dft/radi.py
index 718aea118c..43bfb7b71d 100644
--- a/pyscf/dft/radi.py
+++ b/pyscf/dft/radi.py
@@ -47,14 +47,21 @@
 def murray(n, *args, **kwargs):
     raise RuntimeError('Not implemented')

-# Gauss-Chebyshev of the first kind,  and the transformed interval [0,\infty)
 def becke(n, charge, *args, **kwargs):
     '''Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033'''
     if charge == 1:
         rm = BRAGG_RADII[charge]
     else:
         rm = BRAGG_RADII[charge] * .5
-    t, w = numpy.polynomial.chebyshev.chebgauss(n)
+
+    # Points and weights for Gauss-Chebyshev quadrature points of the second kind
+    # The weights are adjusted to integrate a function on the interval [-1, 1] with uniform weighting
+    # instead of weighted by sqrt(1 - t^2) = sin(i pi / (n+1)).
+    i = numpy.arange(n) + 1
+    t = numpy.cos(i * numpy.pi / (n + 1))
+    w = numpy.pi / (n + 1) * numpy.sin(i * numpy.pi / (n + 1))
+
+    # Change of variables to map the domain to [0, inf)
     r = (1+t)/(1-t) * rm
     w *= 2/(1-t)**2 * rm
     return r, w

diff --git a/pyscf/dft/test/test_grids.py b/pyscf/dft/test/test_grids.py
index ba261f2f36..33dd9c02ec 100644
--- a/pyscf/dft/test/test_grids.py
+++ b/pyscf/dft/test/test_grids.py
@@ -91,7 +91,7 @@ def test_radi(self):

         grid.radi_method = radi.becke
         grid.build(with_non0tab=False)
-        self.assertAlmostEqual(lib.fp(grid.weights), 818061.875131255, 7)
+        self.assertAlmostEqual(lib.fp(grid.weights), 780.7183109298, 9)

     def test_prune(self):
         grid = gen_grid.Grids(h2o)


Fix iterative damping error (pyscf/pyscf#2524)

Iterative damping broken in UCCSD

Damping in unrestricted CCSD doesn't work. Note that RCCSD does not have this error.

```
from pyscf import gto, scf, cc
mol = gto.M(
    atom = 'H 0 0 0; F 0 0 1.1',  # in Angstrom
    basis = 'ccpvdz',
    symmetry = True,
)
mf = scf.HF(mol).run()
mycc = cc.UCCSD(mf)
mycc.iterative_damping = 0.5
mycc.run()
print('CCSD total energy', mycc.e_tot)
```
```
TypeError
File ~/.pyenv/versions/3.12.7/envs/main_env/lib/python3.12/site-packages/pyscf/cc/ccsd.py:80, in kernel(mycc, eris, t1, t2, max_cycle, tol, tolnormt, verbose, callback)
     78 if mycc.iterative_damping < 1.0:
     79     alpha = mycc.iterative_damping
---> 80     t1new = (1-alpha) * t1 + alpha * t1new #t1new = (1-alpha) * numpy.asarray(t1) + alpha * numpy.asarray(t1new)
     81     t2new *= alpha
     82     t2new += (1-alpha) * t2

TypeError: can't multiply sequence by non-int of type 'float'
```
The solution is to modify lines 79-82 in ccsd.py to:
```
            alpha = numpy.asarray(mycc.iterative_damping)
            t1new = (1-alpha) * numpy.asarray(t1) + alpha * numpy.asarray(t1new)
            t2new *= alpha
            t2new += (1-alpha) * numpy.asarray(t2)
```

diff --git a/pyscf/cc/ccsd.py b/pyscf/cc/ccsd.py
index 383b8a8969..59157a7481 100644
--- a/pyscf/cc/ccsd.py
+++ b/pyscf/cc/ccsd.py
@@ -76,10 +76,16 @@ def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
         normt = numpy.linalg.norm(tmpvec)
         tmpvec = None
         if mycc.iterative_damping < 1.0:
-            alpha = mycc.iterative_damping
-            t1new = (1-alpha) * t1 + alpha * t1new
-            t2new *= alpha
-            t2new += (1-alpha) * t2
+            alpha = numpy.asarray(mycc.iterative_damping)
+            if isinstance(t1, tuple): # e.g. UCCSD
+                t1new = tuple((1-alpha) * numpy.asarray(t1_part) + alpha * numpy.asarray(t1new_part)
+                    for t1_part, t1new_part in zip(t1, t1new))
+                t2new = tuple((1-alpha) * numpy.asarray(t2_part) + alpha * numpy.asarray(t2new_part)
+                    for t2_part, t2new_part in zip(t2, t2new))
+            else:
+                t1new = (1-alpha) * numpy.asarray(t1) + alpha * numpy.asarray(t1new)
+                t2new *= alpha
+                t2new += (1-alpha) * numpy.asarray(t2)
         t1, t2 = t1new, t2new
         t1new = t2new = None
         t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)

diff --git a/pyscf/cc/test/test_uccsd.py b/pyscf/cc/test/test_uccsd.py
index a05f3bb4de..162ba47af2 100644
--- a/pyscf/cc/test/test_uccsd.py
+++ b/pyscf/cc/test/test_uccsd.py
@@ -727,6 +727,20 @@ def test_ao2mo(self):
         self.assertAlmostEqual(abs(eri_incore.OVvo - eri_outcore.OVvo).max(), 0, 12)
         self.assertAlmostEqual(abs(eri_incore.OVvv - eri_outcore.OVvv).max(), 0, 12)

+    def test_damping(self):
+        mol = gto.M(
+            atom = 'H 0 0 0; F 0 0 1.1',  # in Angstrom
+            basis = 'ccpvdz',
+            spin=1,
+            charge=-1,
+            symmetry = True,
+            verbose = 0
+        )
+        mf = scf.UHF(mol).run()
+        mycc = cc.UCCSD(mf)
+        mycc.iterative_damping = 0.5
+        mycc.run()
+        self.assertAlmostEqual(mycc.e_tot, -100.07710261186985, 7)

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


Fix CCSD two-particle density matrix for complex orbitals (pyscf/pyscf#2531)

* Fix #2519
* Improve the doc-strings for two-particle density matrix functions.

Incorrect conjugation in 2RDM transformations

PySCF typically uses the following conventions for DMs
```
dm1[p, q] = <q* p>
dm2[p, q, r, s] = <p* r* s q>
```

From Helgaker (3.2), when transforming operators under unitary $U$, the annihilation operator transforms under the conjugate $U^*$, and the creation operator transforms under $U$. This means that following the PySCF conventions, the 1RDM transforms as the Hamiltonian does i.e.
```python
dm1_ao = einsum("ij,pi,qj->pq", dm1_mo, c, c.conj())
```
but the 2RDM should transform as
```python
dm2_ao = einsum("ijkl,pi,qj,rk,sl->pqrs", dm2_mo, c.conj(), c, c.conj(), c)
```
whereas in PySCF the 2RDM is also assumed to transform as the Hamiltonian (the conjugations are reversed):
https://github.com/pyscf/pyscf/blob/2c7f7625fd2719b8f3e9715e40c3937cab65c06e/pyscf/cc/ccsd_rdm.py#L379-L381

This can be verified by trying to recover correlated energies using the RDMs in the AO basis via `ao_repr=True`, which fails, but works with the reverse transformation. See the following script for an example:
https://gist.github.com/obackhouse/855724a176c7f9cad580dc2d39e7c63d
We're aware in this case that the coefficients are not actually a unitary, but this is the best way I could find to show this as an issue in PySCF, and the issue seems to apply at least to the MO->AO transformation.

I am happy to fix but wanted to check if we are missing something first?

diff --git a/pyscf/cc/ccsd_rdm.py b/pyscf/cc/ccsd_rdm.py
index de1b38f2b4..9353dcf3fe 100644
--- a/pyscf/cc/ccsd_rdm.py
+++ b/pyscf/cc/ccsd_rdm.py
@@ -377,8 +377,13 @@ def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):

 def _rdm2_mo2ao(dm2, mo):
+    '''
+    Back transform the two-particle density matrices to AO representation, where
+    the dm2 is defined in accordance with the chemist's ERI notation:
+    E = einsum('pqrs,pqrs', eri, rdm2)
+    '''
     mo_C = mo.conj()
-    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo, mo_C, mo, mo_C)
+    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_C, mo, mo_C, mo)

 if __name__ == '__main__':
diff --git a/pyscf/cc/uccsd_rdm.py b/pyscf/cc/uccsd_rdm.py
index 7e096973f0..8c2bada6ee 100644
--- a/pyscf/cc/uccsd_rdm.py
+++ b/pyscf/cc/uccsd_rdm.py
@@ -625,8 +625,8 @@ def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):

 def _dm2ab_mo2ao(dm2, mo_a, mo_b):
-    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_a, mo_a.conj(),
-                      mo_b, mo_b.conj())
+    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_a.conj(), mo_a,
+                      mo_b.conj(), mo_b)

 if __name__ == '__main__':
diff --git a/pyscf/scf/hf.py b/pyscf/scf/hf.py
index 645cfee1f4..c3f30d93eb 100644
--- a/pyscf/scf/hf.py
+++ b/pyscf/scf/hf.py
@@ -853,6 +853,16 @@ def make_rdm1(mo_coeff, mo_occ, **kwargs):
 def make_rdm2(mo_coeff, mo_occ, **kwargs):
     '''Two-particle density matrix in AO representation

+    NOTE the indices of the two-particle density matrix is ordered to
+
+    dm2[p,q,r,s] = <q^+ s^+ r p>.
+
+    HF energy can be computed
+    E = einsum('pq,qp', hcore, 1pdm) + einsum('pqrs,pqrs', eri, 2pdm) / 2
+    where h1[p,q] = <p|h|q> and eri[p,q,r,s] = (pq|rs)
+to make the density matrix consistent with the density matrix obtained
+    from post-HF methods,
+
     Args:
         mo_coeff : 2D ndarray
             Orbital coefficients. Each column is one orbital.
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/cc/test/test_gccsd.py b/pyscf/cc/test/test_gccsd.py
index 393e8bbced..d0f5123642 100644
--- a/pyscf/cc/test/test_gccsd.py
+++ b/pyscf/cc/test/test_gccsd.py
@@ -285,6 +285,7 @@ def test_rdm_complex(self):
         eri = eri + eri.transpose(1,0,3,2).conj()
         eri = eri + eri.transpose(2,3,0,1)
         eri *= .1
+        mf._eri = eri

         def get_jk(mol, dm, *args,**kwargs):
             vj = numpy.einsum('ijkl,lk->ij', eri, dm)
@@ -294,17 +295,19 @@ def get_veff(mol, dm, *args, **kwargs):
             vj, vk = get_jk(mol, dm)
             return vj - vk
         def ao2mofn(mos):
-            return eri
+            c = mos
+            return lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, c.conj(), c, c.conj(), c)

         mf.get_jk = get_jk
         mf.get_veff = get_veff
         hcore = numpy.random.random((nmo,nmo)) * .2 + numpy.random.random((nmo,nmo))* .2j
-        hcore = hcore + hcore.T.conj() + numpy.diag(range(nmo))*2
+        hcore = hcore + hcore.T.conj() + numpy.diag(numpy.arange(nmo)*2)
         mf.get_hcore = lambda *args: hcore
         mf.get_ovlp = lambda *args: numpy.eye(nmo)
         orbspin = numpy.zeros(nmo, dtype=int)
         orbspin[1::2] = 1
-        mf.mo_coeff = lib.tag_array(numpy.eye(nmo) + 0j, orbspin=orbspin)
+        u = numpy.linalg.eigh(hcore)[1]
+        mf.mo_coeff = lib.tag_array(u, orbspin=orbspin)
         mf.mo_occ = numpy.zeros(nmo)
         mf.mo_occ[:nocc] = 1

@@ -313,8 +316,8 @@ def ao2mofn(mos):
         mycc.ao2mo = lambda *args, **kwargs: eris
         mycc.kernel(eris=eris)
         mycc.solve_lambda(eris=eris)
-        dm1 = mycc.make_rdm1()
-        dm2 = mycc.make_rdm2()
+        dm1 = mycc.make_rdm1(ao_repr=True)
+        dm2 = mycc.make_rdm2(ao_repr=True)

         e1 = numpy.einsum('ij,ji', hcore, dm1)
         e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5