Repository: pyscf/pyscf — pull request #2826 (python)
Title: feat: Add spin-seperated 4-RDM for FCI 

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

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

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

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

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

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

Hints:

Fix patch:
diff --git a/examples/fci/14-density_matrix.py b/examples/fci/14-density_matrix.py
index 28e207b918..52ceb78e1e 100644
--- a/examples/fci/14-density_matrix.py
+++ b/examples/fci/14-density_matrix.py
@@ -158,6 +158,37 @@
                                          (nelec_a,nelec_b))
 dm1, dm2, dm3, dm4 = fci.rdm.reorder_dm1234(dm1, dm2, dm3, dm4)
 
+#
+# Spin-separated 4-particle density matrix
+#
+
+(dm1a, dm1b), (dm2aa, dm2ab, dm2bb), (dm3aaa, dm3aab, dm3abb, dm3bbb), (dm4aaaa, dm4aaab, dm4aabb, dm4abbb, dm4bbbb) = (
+    fci.direct_spin1.make_rdm1234s(fcivec0, norb, (nelec_a, nelec_b))
+)
+assert(numpy.allclose(dm1a+dm1b, dm1))
+assert(numpy.allclose(dm2aa+dm2bb+dm2ab+dm2ab.transpose(2,3,0,1), dm2))
+assert(numpy.allclose(dm3aaa+dm3bbb+dm3aab+dm3aab.transpose(0,1,4,5,2,3)+\
+dm3aab.transpose(4,5,0,1,2,3)+dm3abb+dm3abb.transpose(2,3,0,1,4,5)+dm3abb.transpose(2,3,4,5,0,1), dm3))
+assert(numpy.allclose(
+    dm4aaaa
+    + dm4bbbb
+    + dm4aaab
+    + dm4aaab.transpose(0, 1, 2, 3, 6, 7, 4, 5)
+    + dm4aaab.transpose(0, 1, 6, 7, 2, 3, 4, 5)
+    + dm4aaab.transpose(6, 7, 0, 1, 2, 3, 4, 5)
+    + dm4aabb
+    + dm4aabb.transpose(0, 1, 4, 5, 2, 3, 6, 7)
+    + dm4aabb.transpose(4, 5, 0, 1, 2, 3, 6, 7)
+    + dm4aabb.transpose(0, 1, 4, 5, 6, 7, 2, 3)
+    + dm4aabb.transpose(4, 5, 0, 1, 6, 7, 2, 3)
+    + dm4aabb.transpose(4, 5, 6, 7, 0, 1, 2, 3)
+    + dm4abbb
+    + dm4abbb.transpose(2, 3, 0, 1, 4, 5, 6, 7)
+    + dm4abbb.transpose(2, 3, 4, 5, 0, 1, 6, 7)
+    + dm4abbb.transpose(2, 3, 4, 5, 6, 7, 0, 1),
+    dm4,
+))
+
 #
 # Spin-traced 4-particle transition density matrix
 #
diff --git a/pyscf/fci/direct_spin1.py b/pyscf/fci/direct_spin1.py
index 2f4d3777db..ea3bfbf9a5 100644
--- a/pyscf/fci/direct_spin1.py
+++ b/pyscf/fci/direct_spin1.py
@@ -410,6 +410,77 @@ def make_rdm123s(fcivec, norb, nelec, link_index=None, reorder=True):
     # rdm3aab.transpose(4,5,0,1,2,3)+rdm3abb+rdm3abb.transpose(2,3,0,1,4,5)+rdm3abb.transpose(2,3,4,5,0,1), rdm3)
     return (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb), (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb)
 
+def make_rdm1234(fcivec, norb, nelec, link_index=None, reorder=True):
+    '''Spin traced 1-, 2-, 3, 4-particle density matrices.'''
+    dm1, dm2, dm3, dm4 = rdm.make_dm1234('FCI4pdm_kern_sf', fcivec, fcivec, norb, nelec)
+    if reorder:
+        dm1, dm2, dm3, dm4 = rdm.reorder_dm1234(dm1, dm2, dm3, dm4, inplace=True)
+    return dm1, dm2, dm3, dm4
+
+def make_rdm1234s(fcivec, norb, nelec, link_index=None, reorder=True):
+    r'''Spin separated 1-, 2-, 3, 4-particle density matrices.
+
+    1pdm[p,q] = :math:`\langle q_\alpha^\dagger p_\alpha \rangle +
+                       \langle q_\beta^\dagger  p_\beta \rangle`;
+    2pdm[p,q,r,s] = :math:`\langle p_\alpha^\dagger r_\alpha^\dagger s_\alpha q_\alpha\rangle +
+                           \langle p_\beta^\dagger  r_\alpha^\dagger s_\alpha q_\beta\rangle +
+                           \langle p_\alpha^\dagger r_\beta^\dagger  s_\beta  q_\alpha\rangle +
+                           \langle p_\beta^\dagger  r_\beta^\dagger  s_\beta  q_\beta\rangle`.
+    '''
+    if (not reorder):
+        raise NotImplementedError('reorder=False not currently supported')
+    ci_spinless = civec_spinless_repr([fcivec,], norb, [nelec,])
+    rdm1, rdm2, rdm3, rdm4 = make_rdm1234(ci_spinless, norb*2, (nelec[0]+nelec[1],0))
+
+    # define slices
+    alpha = slice(0, norb)
+    beta = slice(norb, None)
+
+    rdm1a = rdm1[alpha, alpha]
+    rdm1b = rdm1[beta, beta]
+    # assert np.allclose(rdm1a+rdm1b, rdm1)
+
+    rdm2aa = rdm2[alpha, alpha, alpha, alpha]
+    rdm2ab = rdm2[alpha, alpha, beta, beta]
+    rdm2bb = rdm2[beta, beta, beta, beta]
+    # assert np.allclose(rdm2aa+rdm2bb+rdm2ab+rdm2ab.transpose(2,3,0,1), rdm2)
+
+    rdm3aaa = rdm3[alpha, alpha, alpha, alpha, alpha, alpha]
+    rdm3aab = rdm3[alpha, alpha, alpha, alpha, beta, beta]
+    rdm3abb = rdm3[alpha, alpha, beta, beta, beta, beta]
+    rdm3bbb = rdm3[beta, beta, beta, beta, beta, beta]
+    # assert np.allclose(rdm3aaa+rdm3bbb+rdm3aab+rdm3aab.transpose(0,1,4,5,2,3)+\
+    # rdm3aab.transpose(4,5,0,1,2,3)+rdm3abb+rdm3abb.transpose(2,3,0,1,4,5)+rdm3abb.transpose(2,3,4,5,0,1), rdm3)
+
+    rdm4aaaa = rdm4[alpha, alpha, alpha, alpha, alpha, alpha, alpha, alpha]
+    rdm4aaab = rdm4[alpha, alpha, alpha, alpha, alpha, alpha, beta, beta]
+    rdm4aabb = rdm4[alpha, alpha, alpha, alpha, beta, beta, beta, beta]
+    rdm4abbb = rdm4[alpha, alpha, beta, beta, beta, beta, beta, beta]
+    rdm4bbbb = rdm4[beta, beta, beta, beta, beta, beta, beta, beta]
+    # assert numpy.allclose(
+    #     rdm4aaaa
+    #     + rdm4bbbb
+    #     + rdm4aaab
+    #     + rdm4aaab.transpose(0, 1, 2, 3, 6, 7, 4, 5)
+    #     + rdm4aaab.transpose(0, 1, 6, 7, 2, 3, 4, 5)
+    #     + rdm4aaab.transpose(6, 7, 0, 1, 2, 3, 4, 5)
+    #     + rdm4aabb
+    #     + rdm4aabb.transpose(0, 1, 4, 5, 2, 3, 6, 7)
+    #     + rdm4aabb.transpose(0, 1, 4, 5, 6, 7, 2, 3)
+    #     + rdm4aabb.transpose(4, 5, 2, 3, 0, 1, 6, 7)
+    #     + rdm4aabb.transpose(4, 5, 2, 3, 6, 7, 0, 1)
+    #     + rdm4abbb
+    #     + rdm4abbb.transpose(2, 3, 0, 1, 4, 5, 6, 7)
+    #     + rdm4abbb.transpose(2, 3, 4, 5, 0, 1, 6, 7)
+    #     + rdm4abbb.transpose(2, 3, 4, 5, 6, 7, 0, 1),
+    #     rdm4,
+    # )
+    return (
+        (rdm1a, rdm1b),
+        (rdm2aa, rdm2ab, rdm2bb),
+        (rdm3aaa, rdm3aab, rdm3abb, rdm3bbb),
+        (rdm4aaaa, rdm4aaab, rdm4aabb, rdm4abbb, rdm4bbbb),
+    )
 
 def trans_rdm1s(cibra, ciket, norb, nelec, link_index=None):
     r'''Spin separated transition 1-particle density matrices.
@@ -939,6 +1010,16 @@ def make_rdm123(self, fcivec, norb, nelec, link_index=None, reorder=True):
         nelec = _unpack_nelec(nelec, self.spin)
         return make_rdm123(fcivec, norb, nelec, link_index, reorder)
 
+    @lib.with_doc(make_rdm1234s.__doc__)
+    def make_rdm1234s(self, fcivec, norb, nelec, link_index=None, reorder=True):
+        nelec = _unpack_nelec(nelec, self.spin)
+        return make_rdm1234s(fcivec, norb, nelec, link_index, reorder)
+
+    @lib.with_doc(make_rdm1234.__doc__)
+    def make_rdm1234(self, fcivec, norb, nelec, link_index=None, reorder=True):
+        nelec = _unpack_nelec(nelec, self.spin)
+        return make_rdm1234(fcivec, norb, nelec, link_index, reorder)
+
     def make_rdm2(self, fcivec, norb, nelec, link_index=None, reorder=True):
         r'''Spin traced 2-particle density matrice
 


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