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