add spin_summed option to slater determinant one rdm

python/ffsim/states/states.py

@@ -131,7 +131,9 @@ def hartree_fock_state(norb: int, nelec: tuple[int, int]) -> np.ndarray:
 
 
 def slater_determinant_one_rdm(
-    norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]]
+    norb: int,
+    occupied_orbitals: tuple[Sequence[int], Sequence[int]],
+    spin_summed: bool = True,
 ) -> np.ndarray:
     """Return the one-particle reduced density matrix of a Slater determinant.
 
@@ -140,16 +142,22 @@ def slater_determinant_one_rdm(
         occupied_orbitals: A tuple of two sequences of integers. The first
             sequence contains the indices of the occupied alpha orbitals, and
             the second sequence similarly for the beta orbitals.
+        spin_summed: Whether to sum over the spin index.
 
     Returns:
         The one-particle reduced density matrix of the Slater determinant.
     """
-    one_rdm = np.zeros((2 * norb, 2 * norb), dtype=complex)
+    rdm_a = np.zeros((norb, norb), dtype=complex)
+    rdm_b = np.zeros((norb, norb), dtype=complex)
     alpha_orbitals = np.array(occupied_orbitals[0])
-    beta_orbitals = np.array(occupied_orbitals[1]) + norb
-    one_rdm[(alpha_orbitals, alpha_orbitals)] = 1
-    one_rdm[(beta_orbitals, beta_orbitals)] = 1
-    return one_rdm
+    beta_orbitals = np.array(occupied_orbitals[1])
+    if len(alpha_orbitals):
+        rdm_a[(alpha_orbitals, alpha_orbitals)] = 1
+    if len(beta_orbitals):
+        rdm_b[(beta_orbitals, beta_orbitals)] = 1
+    if spin_summed:
+        return rdm_a + rdm_b
+    return scipy.linalg.block_diag(rdm_a, rdm_b)
 
 
 def indices_to_strings(

tests/states/states_test.py

@@ -60,6 +60,33 @@ def test_hartree_fock_state():
     np.testing.assert_allclose(energy, hartree_fock_energy)
 
 
+@pytest.mark.parametrize(
+    "norb, occupied_orbitals, spin_summed",
+    [
+        (4, ([0, 1], [0, 1]), True),
+        (4, ([0, 1], [0, 1]), False),
+        (3, ([0], [1, 2]), True),
+        (3, ([0], [1, 2]), False),
+        (2, ([], [0]), True),
+        (2, ([], [0]), False),
+    ],
+)
+def test_slater_determinant_one_rdm(
+    norb: int, occupied_orbitals: tuple[list[int], list[int]], spin_summed: bool
+):
+    """Test Slater determinant 1-RDM."""
+    occ_a, occ_b = occupied_orbitals
+    nelec = len(occ_a), len(occ_b)
+
+    vec = ffsim.slater_determinant(norb, occupied_orbitals)
+    rdm = ffsim.slater_determinant_one_rdm(
+        norb, occupied_orbitals, spin_summed=spin_summed
+    )
+    expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+
+    np.testing.assert_allclose(rdm, expected)
+
+
 def test_indices_to_strings():
     """Test converting statevector indices to strings."""
     norb = 3
@@ -84,9 +111,9 @@ def test_indices_to_strings():
     "norb, nelec",
     [
         (4, (2, 2)),
-        (4, (1, 2)),
-        (4, (0, 1)),
-        (4, (0, 0)),
+        (3, (1, 2)),
+        (2, (0, 1)),
+        (1, (0, 0)),
     ],
 )
 def test_rdm_1_spin_summed(norb: int, nelec: tuple[int, int]):
@@ -103,9 +130,9 @@ def test_rdm_1_spin_summed(norb: int, nelec: tuple[int, int]):
     "norb, nelec",
     [
         (4, (2, 2)),
-        (4, (1, 2)),
-        (4, (0, 1)),
-        (4, (0, 0)),
+        (3, (1, 2)),
+        (2, (0, 1)),
+        (1, (0, 0)),
     ],
 )
 def test_rdm_1(norb: int, nelec: tuple[int, int]):