Repository: pyscf/pyscf — pull request #2357 (python)
Title: 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.

Resolved issues:
title: Spin-separated 3-RDMs?
body: 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: ...
```

Hints:

Fix patch:
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
 


Test patch:
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]