From c15f3c5f00f0953aeea7c36d003b5f1ec9d500d2 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Mon, 30 Oct 2023 11:59:34 -0400
Subject: [PATCH] add spin_summed option to slater determinant one rdm

---
 python/ffsim/states/states.py | 20 ++++++++++++------
 tests/states/states_test.py   | 39 +++++++++++++++++++++++++++++------
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py
index 0b92ad3cb..2d0daedae 100644
--- a/python/ffsim/states/states.py
+++ b/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(
diff --git a/tests/states/states_test.py b/tests/states/states_test.py
index b93e0232a..0d4894ca3 100644
--- a/tests/states/states_test.py
+++ b/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]):