Repository: qiskit-community/ffsim (python)

Issue:
Change RDM interface to be more space efficient and consistent with PySCF
In the [`rdm`](https://github.com/qiskit-community/ffsim/blob/b23ba626e93328fd683df2072cf1ec69a6d7834c/python/ffsim/states/rdm.py#L31) function, setting `spin_summed=False` results in an expanded tensor with many zeros. For example, a 1-RDM will be block diagonal. I think we should change this to return, for example, a pair of RDMs, as the `make_rdm1s` function from PySCF does.

`slater_determinant_rdm` should also be updated.

Pull request description:
Add functions to return spin-separated RDMs
Fixes #246 

Patch:
diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py
index 12f7139a..91caba81 100644
--- a/python/ffsim/__init__.py
+++ b/python/ffsim/__init__.py
@@ -71,8 +71,10 @@ from ffsim.states import (
     indices_to_strings,
     one_hot,
     rdm,
+    rdms,
     slater_determinant,
     slater_determinant_rdm,
+    slater_determinant_rdms,
     spin_square,
     strings_to_indices,
 )
@@ -161,10 +163,12 @@ __all__ = [
     "qiskit",
     "random",
     "rdm",
+    "rdms",
     "simulate_qdrift_double_factorized",
     "simulate_trotter_double_factorized",
     "slater_determinant",
     "slater_determinant_rdm",
+    "slater_determinant_rdms",
     "spin_square",
     "strings_to_indices",
     "testing",
diff --git a/python/ffsim/states/__init__.py b/python/ffsim/states/__init__.py
index c85d0b21..b9e57135 100644
--- a/python/ffsim/states/__init__.py
+++ b/python/ffsim/states/__init__.py
@@ -11,7 +11,7 @@
 """States."""
 
 from ffsim.states.product_state_sum import ProductStateSum
-from ffsim.states.rdm import rdm
+from ffsim.states.rdm import rdm, rdms
 from ffsim.states.states import (
     StateVector,
     dim,
@@ -21,6 +21,7 @@ from ffsim.states.states import (
     one_hot,
     slater_determinant,
     slater_determinant_rdm,
+    slater_determinant_rdms,
     spin_square,
     strings_to_indices,
 )
@@ -37,8 +38,10 @@ __all__ = [
     "indices_to_strings",
     "one_hot",
     "rdm",
+    "rdms",
     "slater_determinant",
     "slater_determinant_rdm",
+    "slater_determinant_rdms",
     "spin_square",
     "strings_to_indices",
 ]
diff --git a/python/ffsim/states/rdm.py b/python/ffsim/states/rdm.py
index a5baa817..3189ba5e 100644
--- a/python/ffsim/states/rdm.py
+++ b/python/ffsim/states/rdm.py
@@ -24,10 +24,12 @@ from pyscf.fci.direct_spin1 import (
     trans_rdm12,
     trans_rdm12s,
 )
+from typing_extensions import deprecated
 
 from ffsim.cistring import gen_linkstr_index
 
 
+@deprecated("ffsim.rdm is deprecated. Instead, use ffsim.rdms.")
 def rdm(
     vec: np.ndarray,
     norb: int,
@@ -100,7 +102,7 @@ def rdm(
             return _rdm1(vec, norb, nelec, link_index)
     if rank == 2:
         if spin_summed:
-            return _rdm2_spin_summed(
+            return _rdm2_spin_summed_deprecated(
                 vec, norb, nelec, reordered, link_index, return_lower_ranks
             )
         else:
@@ -110,23 +112,28 @@ def rdm(
     )
 
 
-def _rdm1_spin_summed(
+def _rdm2_spin_summed_deprecated(
     vec: np.ndarray,
     norb: int,
     nelec: tuple[int, int],
+    reordered: bool,
     link_index: tuple[np.ndarray, np.ndarray] | None,
-) -> np.ndarray:
-    rdm1_real = make_rdm1(vec.real, norb, nelec, link_index=link_index)
-    rdm1_imag = make_rdm1(vec.imag, norb, nelec, link_index=link_index)
-    trans_rdm1_real_imag = trans_rdm1(
-        vec.real, vec.imag, norb, nelec, link_index=link_index
+    return_lower_ranks: bool,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+    rdm1_real, rdm2_real = make_rdm12(
+        vec.real, norb, nelec, reorder=reordered, link_index=link_index
     )
-    trans_rdm1_imag_real = trans_rdm1(
-        vec.imag, vec.real, norb, nelec, link_index=link_index
+    rdm1_imag, rdm2_imag = make_rdm12(
+        vec.imag, norb, nelec, reorder=reordered, link_index=link_index
     )
-    return _assemble_rdm1_spin_summed(
-        rdm1_real, rdm1_imag, trans_rdm1_real_imag, trans_rdm1_imag_real
+    trans_rdm1_real_imag, trans_rdm2_real_imag = trans_rdm12(
+        vec.real, vec.imag, norb, nelec, reorder=reordered, link_index=link_index
     )
+    rdm2 = _assemble_rdm2_spin_summed(rdm2_real, rdm2_imag, trans_rdm2_real_imag)
+    if not return_lower_ranks:
+        return rdm2
+    rdm1 = _assemble_rdm1_spin_summed(rdm1_real, rdm1_imag, trans_rdm1_real_imag)
+    return rdm1, rdm2
 
 
 def _rdm1(
@@ -148,35 +155,19 @@ def _rdm1(
     )
 
 
-def _rdm2_spin_summed(
-    vec: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    reordered: bool,
-    link_index: tuple[np.ndarray, np.ndarray] | None,
-    return_lower_ranks: bool,
-) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
-    rdm1_real, rdm2_real = make_rdm12(
-        vec.real, norb, nelec, reorder=reordered, link_index=link_index
-    )
-    rdm1_imag, rdm2_imag = make_rdm12(
-        vec.imag, norb, nelec, reorder=reordered, link_index=link_index
-    )
-    trans_rdm1_real_imag, trans_rdm2_real_imag = trans_rdm12(
-        vec.real, vec.imag, norb, nelec, reorder=reordered, link_index=link_index
-    )
-    trans_rdm1_imag_real, trans_rdm2_imag_real = trans_rdm12(
-        vec.imag, vec.real, norb, nelec, reorder=reordered, link_index=link_index
-    )
-    rdm2 = _assemble_rdm2_spin_summed(
-        rdm2_real, rdm2_imag, trans_rdm2_real_imag, trans_rdm2_imag_real
-    )
-    if not return_lower_ranks:
-        return rdm2
-    rdm1 = _assemble_rdm1_spin_summed(
-        rdm1_real, rdm1_imag, trans_rdm1_real_imag, trans_rdm1_imag_real
-    )
-    return rdm1, rdm2
+def _assemble_rdm1(
+    rdms1_real: tuple[np.ndarray, np.ndarray],
+    rdms1_imag: tuple[np.ndarray, np.ndarray],
+    trans_rdms1_real_imag: tuple[np.ndarray, np.ndarray],
+    trans_rdms1_imag_real: tuple[np.ndarray, np.ndarray],
+) -> np.ndarray:
+    rdms1 = np.stack(rdms1_real).astype(complex)
+    rdms1 += np.stack(rdms1_imag)
+    # use minus sign for real-imag and plus sign for imag-real because
+    # the rdm1 convention in pyscf is transposed
+    rdms1 -= 1j * np.stack(trans_rdms1_real_imag)
+    rdms1 += 1j * np.stack(trans_rdms1_imag_real)
+    return scipy.linalg.block_diag(*rdms1)
 
 
 def _rdm2(
@@ -210,40 +201,6 @@ def _rdm2(
     return rdm1, rdm2
 
 
-def _assemble_rdm1_spin_summed(
-    rdm1_real: np.ndarray,
-    rdm1_imag: np.ndarray,
-    trans_rdm1_real_imag: np.ndarray,
-    trans_rdm1_imag_real: np.ndarray,
-) -> np.ndarray:
-    # use minus sign for 1j because the rdm1 convention from pyscf is transposed
-    return rdm1_real + rdm1_imag - 1j * (trans_rdm1_real_imag - trans_rdm1_imag_real)
-
-
-def _assemble_rdm1(
-    rdms1_real: tuple[np.ndarray, np.ndarray],
-    rdms1_imag: tuple[np.ndarray, np.ndarray],
-    trans_rdms1_real_imag: tuple[np.ndarray, np.ndarray],
-    trans_rdms1_imag_real: tuple[np.ndarray, np.ndarray],
-) -> np.ndarray:
-    rdms1 = np.stack(rdms1_real).astype(complex)
-    rdms1 += np.stack(rdms1_imag)
-    # use minus sign for real-imag and plus sign for imag-real because
-    # the rdm1 convention in pyscf is transposed
-    rdms1 -= 1j * np.stack(trans_rdms1_real_imag)
-    rdms1 += 1j * np.stack(trans_rdms1_imag_real)
-    return scipy.linalg.block_diag(*rdms1)
-
-
-def _assemble_rdm2_spin_summed(
-    rdm2_real: np.ndarray,
-    rdm2_imag: np.ndarray,
-    trans_rdm2_real_imag: np.ndarray,
-    trans_rdm2_imag_real: np.ndarray,
-) -> np.ndarray:
-    return rdm2_real + rdm2_imag + 1j * (trans_rdm2_real_imag - trans_rdm2_imag_real)
-
-
 def _assemble_rdm2(
     rdms2_real: tuple[np.ndarray, np.ndarray, np.ndarray],
     rdms2_imag: tuple[np.ndarray, np.ndarray, np.ndarray],
@@ -266,3 +223,206 @@ def _assemble_rdm2(
     rdm2[norb:, norb:, :norb, :norb] = rdm_ba
     rdm2[norb:, norb:, norb:, norb:] = rdm_bb
     return rdm2
+
+
+def rdms(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    *,
+    rank: int = 1,
+    spin_summed: bool = False,
+    reorder: bool = True,
+) -> np.ndarray | tuple[np.ndarray, ...]:
+    """Return the spin-separated reduced density matrices of a state vector.
+
+    The rank 1 RDM is defined as follows:
+
+    .. code::
+
+        rdm1[p, q] = ⟨p+ q⟩
+
+    The definition of higher-rank RDMs depends on the ``reorder`` argument, which
+    defaults to True.
+
+    **reorder = True**
+
+    The reordered RDMs are defined as follows:
+
+    .. code::
+
+        rdm2[p, q, r, s] = ⟨p+ r+ s q⟩
+        rdm3[p, q, r, s, t, u] = ⟨p+ r+ t+ u s q⟩
+        rdm4[p, q, r, s, t, u, v, w] = ⟨p+ r+ t+ v+ w u s q⟩
+
+    **reorder = False**
+
+    If reorder is set to False, the RDMs are defined as follows:
+
+    .. code::
+
+        rdm2[p, q, r, s] = ⟨p+ q r+ s⟩
+        rdm3[p, q, r, s, t, u] = ⟨p+ q r+ s t+ u⟩
+        rdm4[p, q, r, s, t, u, v, w] = ⟨p+ q r+ s t+ u v+ w⟩
+
+    Note:
+        Currently, only ranks 1 and 2 are supported.
+
+    Args:
+        vec: The state vector whose reduced density matrix is desired.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+        rank: The rank of the reduced density matrix.
+        spin_summed: Whether to return the "spin-summed" RDMs.
+        reorder: Whether to reorder the indices of the reduced density matrix.
+
+    Returns:
+        The reduced density matrices.
+        All RDMs up to and including the specified rank are returned, in increasing
+        order of rank. For example, if `rank=2` then a tuple `(rdm1, rdm2)` is returned.
+        The 1-RDMs are: (alpha-alpha, beta-beta).
+        The spin-summed 1-RDM is alpha-alpha + alpha-beta.
+        The 2-RDMs are: (alpha-alpha, alpha-beta, beta-beta).
+        The spin-summed 2-RDM is alpha-alpha + alpha-beta + beta-alpha + beta-beta.
+    """
+    n_alpha, n_beta = nelec
+    link_index_a = gen_linkstr_index(range(norb), n_alpha)
+    link_index_b = gen_linkstr_index(range(norb), n_beta)
+    link_index = (link_index_a, link_index_b)
+    if rank == 1:
+        if spin_summed:
+            return _rdm1_spin_summed(vec, norb, nelec, link_index)
+        else:
+            return _rdm1s(vec, norb, nelec, link_index)
+    if rank == 2:
+        if spin_summed:
+            return _rdm2_spin_summed(vec, norb, nelec, reorder, link_index)
+        else:
+            return _rdm2s(vec, norb, nelec, reorder, link_index)
+    raise NotImplementedError(
+        f"Computing the rank {rank} reduced density matrix is currently not supported."
+    )
+
+
+def _rdm1s(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    link_index: tuple[np.ndarray, np.ndarray] | None,
+) -> np.ndarray:
+    rdm1s_real = make_rdm1s(vec.real, norb, nelec, link_index=link_index)
+    rdm1s_imag = make_rdm1s(vec.imag, norb, nelec, link_index=link_index)
+    trans_rdm1s_real_imag = trans_rdm1s(
+        vec.real, vec.imag, norb, nelec, link_index=link_index
+    )
+    return _assemble_rdm1s(rdm1s_real, rdm1s_imag, trans_rdm1s_real_imag)
+
+
+def _assemble_rdm1s(
+    rdm1s_real: tuple[np.ndarray, np.ndarray],
+    rdm1s_imag: tuple[np.ndarray, np.ndarray],
+    trans_rdm1s_real_imag: tuple[np.ndarray, np.ndarray],
+) -> np.ndarray:
+    rdm1s = np.stack(rdm1s_real).astype(complex)
+    rdm1s += np.stack(rdm1s_imag)
+    # use minus sign for real-imag and plus sign for imag-real because
+    # the rdm1 convention in pyscf is transposed
+    rdm1s -= 1j * np.stack(trans_rdm1s_real_imag)
+    rdm1s += 1j * np.stack(trans_rdm1s_real_imag).transpose(0, 2, 1)
+    return rdm1s
+
+
+def _rdm2s(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    reorder: bool,
+    link_index: tuple[np.ndarray, np.ndarray] | None,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+    rdm1s_real, rdm2s_real = make_rdm12s(
+        vec.real, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    rdm1s_imag, rdm2s_imag = make_rdm12s(
+        vec.imag, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    trans_rdm1s_real_imag, trans_rdm2s_real_imag = trans_rdm12s(
+        vec.real, vec.imag, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    rdm1 = _assemble_rdm1s(rdm1s_real, rdm1s_imag, trans_rdm1s_real_imag)
+    rdm2 = _assemble_rdm2s(rdm2s_real, rdm2s_imag, trans_rdm2s_real_imag)
+    return rdm1, rdm2
+
+
+def _assemble_rdm2s(
+    rdm2s_real: tuple[np.ndarray, np.ndarray, np.ndarray],
+    rdm2s_imag: tuple[np.ndarray, np.ndarray, np.ndarray],
+    trans_rdm2s_real_imag: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
+) -> np.ndarray:
+    rdm2s = np.stack(rdm2s_real).astype(complex)
+    rdm2s += np.stack(rdm2s_imag)
+    aa, ab, _, bb = trans_rdm2s_real_imag
+    rdm2s += 1j * np.stack([aa, ab, bb])
+    rdm2s -= 1j * np.stack(
+        [
+            aa.transpose(3, 2, 1, 0),
+            ab.transpose(1, 0, 3, 2),
+            bb.transpose(3, 2, 1, 0),
+        ]
+    )
+    return rdm2s
+
+
+def _rdm1_spin_summed(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    link_index: tuple[np.ndarray, np.ndarray] | None,
+) -> np.ndarray:
+    rdm1_real = make_rdm1(vec.real, norb, nelec, link_index=link_index)
+    rdm1_imag = make_rdm1(vec.imag, norb, nelec, link_index=link_index)
+    trans_rdm1_real_imag = trans_rdm1(
+        vec.real, vec.imag, norb, nelec, link_index=link_index
+    )
+    return _assemble_rdm1_spin_summed(rdm1_real, rdm1_imag, trans_rdm1_real_imag)
+
+
+def _assemble_rdm1_spin_summed(
+    rdm1_real: np.ndarray,
+    rdm1_imag: np.ndarray,
+    trans_rdm1_real_imag: np.ndarray,
+) -> np.ndarray:
+    # use minus sign for 1j because the rdm1 convention from pyscf is transposed
+    return rdm1_real + rdm1_imag - 1j * (trans_rdm1_real_imag - trans_rdm1_real_imag.T)
+
+
+def _rdm2_spin_summed(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    reorder: bool,
+    link_index: tuple[np.ndarray, np.ndarray] | None,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
+    rdm1_real, rdm2_real = make_rdm12(
+        vec.real, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    rdm1_imag, rdm2_imag = make_rdm12(
+        vec.imag, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    trans_rdm1_real_imag, trans_rdm2_real_imag = trans_rdm12(
+        vec.real, vec.imag, norb, nelec, reorder=reorder, link_index=link_index
+    )
+    rdm1 = _assemble_rdm1_spin_summed(rdm1_real, rdm1_imag, trans_rdm1_real_imag)
+    rdm2 = _assemble_rdm2_spin_summed(rdm2_real, rdm2_imag, trans_rdm2_real_imag)
+    return rdm1, rdm2
+
+
+def _assemble_rdm2_spin_summed(
+    rdm2_real: np.ndarray,
+    rdm2_imag: np.ndarray,
+    trans_rdm2_real_imag: np.ndarray,
+) -> np.ndarray:
+    return (
+        rdm2_real
+        + rdm2_imag
+        + 1j * (trans_rdm2_real_imag - trans_rdm2_real_imag.transpose(3, 2, 1, 0))
+    )
diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py
index 1f150032..1f8d26b3 100644
--- a/python/ffsim/states/states.py
+++ b/python/ffsim/states/states.py
@@ -197,6 +197,10 @@ def hartree_fock_state(norb: int, nelec: int | tuple[int, int]) -> np.ndarray:
     return slater_determinant(norb, occupied_orbitals=(range(n_alpha), range(n_beta)))
 
 
+@deprecated(
+    "ffsim.slater_determinant_rdm is deprecated. "
+    "Instead, use ffsim.slater_determinant_rdms."
+)
 def slater_determinant_rdm(
     norb: int,
     occupied_orbitals: tuple[Sequence[int], Sequence[int]],
@@ -260,6 +264,112 @@ def slater_determinant_rdm(
     )
 
 
+@overload
+def slater_determinant_rdms(
+    norb: int,
+    occupied_orbitals: Sequence[int],
+    orbital_rotation: np.ndarray | None = None,
+) -> np.ndarray: ...
+@overload
+def slater_determinant_rdms(
+    norb: int,
+    occupied_orbitals: tuple[Sequence[int], Sequence[int]],
+    orbital_rotation: np.ndarray
+    | tuple[np.ndarray | None, np.ndarray | None]
+    | None = None,
+) -> np.ndarray: ...
+def slater_determinant_rdms(
+    norb: int,
+    occupied_orbitals: Sequence[int] | tuple[Sequence[int], Sequence[int]],
+    orbital_rotation: np.ndarray
+    | tuple[np.ndarray | None, np.ndarray | None]
+    | None = None,
+    *,
+    rank: int = 1,
+) -> np.ndarray:
+    """Return the reduced density matrices of a `Slater determinant`_.
+
+    Note:
+        Currently, only rank 1 is supported.
+
+    Args:
+        norb: The number of spatial orbitals.
+        occupied_orbitals: The occupied orbitals in the electronic configuration.
+            This is either a list of integers specifying spinless orbitals, or a
+            pair of lists, where the first list specifies the spin alpha orbitals and
+            the second list specifies the spin beta orbitals.
+        orbital_rotation: The optional orbital rotation.
+            You can pass either a single Numpy array specifying the orbital rotation
+            to apply to both spin sectors, or you can pass a pair of Numpy arrays
+            specifying independent orbital rotations for spin alpha and spin beta.
+            If passing a pair, you can use ``None`` for one of the
+            values in the pair to indicate that no operation should be applied to
+            that spin sector.
+        rank: The rank of the reduced density matrix. I.e., rank 1 corresponds to the
+            one-particle RDM, rank 2 corresponds to the 2-particle RDM, etc.
+
+    Returns:
+        The reduced density matrices of the Slater determinant.
+        All RDMs up to and including the specified rank are returned, in increasing
+        order of rank. For example, if `rank=2` then a tuple `(rdm1, rdm2)` is returned.
+        The representation of an RDM depends on whether `occupied_orbitals` is a
+        sequence of integers (spinless case), or a pair of such sequences
+        (spinful case). In the spinless case, the full RDM is returned.
+        In the spinful case, each RDM is represented as a stacked Numpy
+        array of sub-RDMs. For example, the 1-RDMs are: (alpha-alpha, alpha-beta), and
+        the 2-RDMs are: (alpha-alpha, alpha-beta, beta-beta).
+
+    .. _Slater determinant: ffsim.html#ffsim.slater_determinant
+    """
+    if not occupied_orbitals or isinstance(occupied_orbitals[0], (int, np.integer)):
+        # Spinless case
+        occupied_orbitals = list(cast(Sequence[int], occupied_orbitals))
+        if rank == 1:
+            rdm = np.zeros((norb, norb), dtype=complex)
+            if occupied_orbitals:
+                rdm[(occupied_orbitals, occupied_orbitals)] = 1
+            if orbital_rotation is not None:
+                orbital_rotation = cast(np.ndarray, orbital_rotation)
+                rdm = orbital_rotation.conj() @ rdm @ orbital_rotation.T
+            return rdm
+        raise NotImplementedError(
+            f"Returning the rank {rank} reduced density matrix is currently not "
+            "supported."
+        )
+    else:
+        # Spinful case
+        alpha_orbitals, beta_orbitals = cast(
+            Tuple[Sequence[int], Sequence[int]], occupied_orbitals
+        )
+        alpha_orbitals = list(alpha_orbitals)
+        beta_orbitals = list(beta_orbitals)
+        if rank == 1:
+            rdm_a = np.zeros((norb, norb))
+            rdm_b = np.zeros((norb, norb))
+            if alpha_orbitals:
+                rdm_a[(alpha_orbitals, alpha_orbitals)] = 1
+            if beta_orbitals:
+                rdm_b[(beta_orbitals, beta_orbitals)] = 1
+            if orbital_rotation is not None:
+                if (
+                    isinstance(orbital_rotation, np.ndarray)
+                    and orbital_rotation.ndim == 2
+                ):
+                    orbital_rotation_a: np.ndarray | None = orbital_rotation
+                    orbital_rotation_b: np.ndarray | None = orbital_rotation
+                else:
+                    orbital_rotation_a, orbital_rotation_b = orbital_rotation
+                if orbital_rotation_a is not None:
+                    rdm_a = orbital_rotation_a.conj() @ rdm_a @ orbital_rotation_a.T
+                if orbital_rotation_b is not None:
+                    rdm_b = orbital_rotation_b.conj() @ rdm_b @ orbital_rotation_b.T
+            return np.stack([rdm_a, rdm_b])
+        raise NotImplementedError(
+            f"Returning the rank {rank} reduced density matrix is currently not "
+            "supported."
+        )
+
+
 def indices_to_strings(
     indices: Sequence[int] | np.ndarray, norb: int, nelec: int | tuple[int, int]
 ) -> list[str]:


Test patch:
diff --git a/tests/python/states/rdm_test.py b/tests/python/states/rdm_test.py
index 106384de..ab77e8ab 100644
--- a/tests/python/states/rdm_test.py
+++ b/tests/python/states/rdm_test.py
@@ -16,6 +16,7 @@ import itertools
 
 import numpy as np
 import pytest
+import scipy.linalg
 
 import ffsim
 
@@ -33,23 +34,18 @@ import ffsim
         (1, (0, 0), False),
     ],
 )
-def test_rdm_1(norb: int, nelec: tuple[int, int], spin_summed: bool):
-    """Test computing 1-RDM."""
-    func = {
-        # spin_summed: function
-        False: _rdm1,
-        True: _rdm1_spin_summed,
-    }
+def test_rdm1s(norb: int, nelec: tuple[int, int], spin_summed: bool):
+    """Test computing spin-summed 1-RDM."""
     rng = np.random.default_rng()
     vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
-    expected = func[spin_summed](vec, norb, nelec)
+    rdm1_func = _rdm1_spin_summed if spin_summed else _rdm1s
+    rdm = ffsim.rdms(vec, norb, nelec, spin_summed=spin_summed)
+    expected = rdm1_func(vec, norb, nelec)
     np.testing.assert_allclose(rdm, expected, atol=1e-12)
 
 
 @pytest.mark.parametrize(
-    "norb, nelec, spin_summed, reordered",
+    "norb, nelec, spin_summed, reorder",
     [
         (4, (2, 2), True, True),
         (4, (2, 2), True, False),
@@ -69,86 +65,38 @@ def test_rdm_1(norb: int, nelec: tuple[int, int], spin_summed: bool):
         (1, (0, 0), False, False),
     ],
 )
-def test_rdm_2(norb: int, nelec: tuple[int, int], spin_summed: bool, reordered: bool):
+def test_rdm2s(norb: int, nelec: tuple[int, int], spin_summed: bool, reorder: bool):
     """Test computing 1- and 2-RDMs."""
-    func = {
-        # (spin_summed, reorder): (rdm1_function, rdm2_function)
-        (False, False): (_rdm1, _rdm2),
-        (False, True): (_rdm1, _rdm2_reordered),
-        (True, False): (_rdm1_spin_summed, _rdm2_spin_summed),
-        (True, True): (_rdm1_spin_summed, _rdm2_spin_summed_reordered),
-    }
     rng = np.random.default_rng()
     vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
 
-    rdm1, rdm2 = ffsim.rdm(
+    rdm1, rdm2 = ffsim.rdms(
         vec,
         norb,
         nelec,
         rank=2,
         spin_summed=spin_summed,
-        reordered=reordered,
+        reorder=reorder,
     )
-    rdm1_func, rdm2_func = func[spin_summed, reordered]
+
+    func = {
+        # (spin_summed, reorder): (rdm1_func, rdm2_func)
+        (False, False): (_rdm1s, _rdm2s),
+        (False, True): (_rdm1s, _rdm2s_reordered),
+        (True, False): (_rdm1_spin_summed, _rdm2_spin_summed),
+        (True, True): (_rdm1_spin_summed, _rdm2_spin_summed_reordered),
+    }
+    rdm1_func, rdm2_func = func[spin_summed, reorder]
     expected_rdm1 = rdm1_func(vec, norb, nelec)
     expected_rdm2 = rdm2_func(vec, norb, nelec)
     np.testing.assert_allclose(rdm1, expected_rdm1, atol=1e-12)
     np.testing.assert_allclose(rdm2, expected_rdm2, atol=1e-12)
 
 
-@pytest.mark.parametrize(
-    "norb, nelec, spin_summed, reordered, rank",
-    [
-        (4, (2, 2), True, True, 2),
-    ],
-)
-def test_no_lower_ranks(
-    norb: int, nelec: tuple[int, int], spin_summed: bool, reordered: bool, rank: int
-):
-    """Test computing higher-rank RDM without returning lower ranks."""
-    func = {
-        # (spin_summed, reorder): function
-        (False, False, 2): _rdm2,
-        (False, True, 2): _rdm2_reordered,
-        (True, False, 2): _rdm2_spin_summed,
-        (True, True, 2): _rdm2_spin_summed_reordered,
-    }
-    rng = np.random.default_rng()
-    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(
-        vec,
-        norb,
-        nelec,
-        rank=rank,
-        spin_summed=spin_summed,
-        reordered=reordered,
-        return_lower_ranks=False,
-    )
-    expected = func[spin_summed, reordered, rank](vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-def _rdm1_spin_summed(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
-    """Compute spin-summed 1-RDM directly from its definition."""
-    rdm = np.zeros((norb, norb), dtype=complex)
-    for i, j in itertools.combinations_with_replacement(range(norb), 2):
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_a(i), ffsim.des_a(j)): 1,
-                (ffsim.cre_b(i), ffsim.des_b(j)): 1,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[i, j] = val
-        rdm[j, i] = val.conjugate()
-    return rdm
-
-
-def _rdm1(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+def _rdm1s(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
     """Compute 1-RDM directly from its definition."""
-    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)
     for i, j in itertools.combinations_with_replacement(range(norb), 2):
         op = ffsim.FermionOperator(
             {
@@ -157,8 +105,8 @@ def _rdm1(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[i, j] = val
-        rdm[j, i] = val.conjugate()
+        rdm_a[i, j] = val
+        rdm_a[j, i] = val.conjugate()
         op = ffsim.FermionOperator(
             {
                 (ffsim.cre_b(i), ffsim.des_b(j)): 1,
@@ -166,79 +114,76 @@ def _rdm1(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[norb + i, norb + j] = val
-        rdm[norb + j, norb + i] = val.conjugate()
-    return rdm
+        rdm_b[i, j] = val
+        rdm_b[j, i] = val.conjugate()
+    return np.array([rdm_a, rdm_b])
 
 
-def _rdm2_spin_summed_reordered(
-    vec: np.ndarray, norb: int, nelec: tuple[int, int]
-) -> np.ndarray:
-    """Compute spin-summed reordered 2-RDM directly from its definition."""
-    rdm = np.zeros((norb, norb, norb, norb), dtype=complex)
+def _rdm1_spin_summed(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute spin-summed 1-RDM directly from its definition."""
+    rdms = _rdm1s(vec, norb, nelec)
+    return np.sum(rdms, axis=0)
+
+
+def _rdm2s(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute 2-RDM directly from its definition."""
+    rdm_aa = np.zeros((norb, norb, norb, norb), dtype=complex)
+    rdm_ab = np.zeros((norb, norb, norb, norb), dtype=complex)
+    rdm_bb = np.zeros((norb, norb, norb, norb), dtype=complex)
     for p, q, r, s in itertools.product(range(norb), repeat=4):
         op = ffsim.FermionOperator(
             {
-                (ffsim.cre_a(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_a(q)): 1,
-                (ffsim.cre_a(p), ffsim.cre_b(r), ffsim.des_b(s), ffsim.des_a(q)): 1,
-                (ffsim.cre_b(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_b(q)): 1,
-                (ffsim.cre_b(p), ffsim.cre_b(r), ffsim.des_b(s), ffsim.des_b(q)): 1,
+                (ffsim.cre_a(p), ffsim.des_a(q), ffsim.cre_a(r), ffsim.des_a(s)): 1,
             }
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
-    return rdm
-
+        rdm_aa[p, q, r, s] = val
 
-def _rdm2_spin_summed(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
-    """Compute spin-summed 2-RDM directly from its definition."""
-    rdm = np.zeros((norb, norb, norb, norb), dtype=complex)
-    for p, q, r, s in itertools.product(range(norb), repeat=4):
         op = ffsim.FermionOperator(
             {
-                (ffsim.cre_a(p), ffsim.des_a(q), ffsim.cre_a(r), ffsim.des_a(s)): 1,
                 (ffsim.cre_a(p), ffsim.des_a(q), ffsim.cre_b(r), ffsim.des_b(s)): 1,
-                (ffsim.cre_b(p), ffsim.des_b(q), ffsim.cre_a(r), ffsim.des_a(s)): 1,
-                (ffsim.cre_b(p), ffsim.des_b(q), ffsim.cre_b(r), ffsim.des_b(s)): 1,
             }
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
-    return rdm
+        rdm_ab[p, q, r, s] = val
 
-
-def _rdm2_reordered(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
-    """Compute reordered 2-RDM directly from its definition."""
-    rdm = np.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=complex)
-    for p, q, r, s in itertools.product(range(norb), repeat=4):
         op = ffsim.FermionOperator(
             {
-                (ffsim.cre_a(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_a(q)): 1,
+                (ffsim.cre_b(p), ffsim.des_b(q), ffsim.cre_b(r), ffsim.des_b(s)): 1,
             }
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
+        rdm_bb[p, q, r, s] = val
+
+    return np.array([rdm_aa, rdm_ab, rdm_bb])
+
 
+def _rdm2s_reordered(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute reordered 2-RDM directly from its definition."""
+    rdm_aa = np.zeros((norb, norb, norb, norb), dtype=complex)
+    rdm_ab = np.zeros((norb, norb, norb, norb), dtype=complex)
+    rdm_bb = np.zeros((norb, norb, norb, norb), dtype=complex)
+    for p, q, r, s in itertools.product(range(norb), repeat=4):
         op = ffsim.FermionOperator(
             {
-                (ffsim.cre_a(p), ffsim.cre_b(r), ffsim.des_b(s), ffsim.des_a(q)): 1,
+                (ffsim.cre_a(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_a(q)): 1,
             }
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[p, q, norb + r, norb + s] = val
+        rdm_aa[p, q, r, s] = val
 
         op = ffsim.FermionOperator(
             {
-                (ffsim.cre_b(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_b(q)): 1,
+                (ffsim.cre_a(p), ffsim.cre_b(r), ffsim.des_b(s), ffsim.des_a(q)): 1,
             }
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[norb + p, norb + q, r, s] = val
+        rdm_ab[p, q, r, s] = val
 
         op = ffsim.FermionOperator(
             {
@@ -247,47 +192,158 @@ def _rdm2_reordered(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.nd
         )
         linop = ffsim.linear_operator(op, norb, nelec)
         val = np.vdot(vec, linop @ vec)
-        rdm[norb + p, norb + q, norb + r, norb + s] = val
-    return rdm
+        rdm_bb[p, q, r, s] = val
+    return np.array([rdm_aa, rdm_ab, rdm_bb])
+
+
+def _rdm2_spin_summed(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute spin-summed 2-RDM directly from its definition."""
+    rdm_aa, rdm_ab, rdm_bb = _rdm2s(vec, norb, nelec)
+    return rdm_aa + rdm_ab + rdm_ab.transpose(2, 3, 0, 1) + rdm_bb
+
+
+def _rdm2_spin_summed_reordered(
+    vec: np.ndarray, norb: int, nelec: tuple[int, int]
+) -> np.ndarray:
+    """Compute spin-summed reordered 2-RDM directly from its definition."""
+    rdm_aa, rdm_ab, rdm_bb = _rdm2s_reordered(vec, norb, nelec)
+    return rdm_aa + rdm_ab + rdm_ab.transpose(2, 3, 0, 1) + rdm_bb
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin_summed",
+    [
+        (4, (2, 2), True),
+        (4, (2, 2), False),
+        (3, (1, 2), True),
+        (3, (1, 2), False),
+        (2, (0, 1), True),
+        (2, (0, 1), False),
+        (1, (0, 0), True),
+        (1, (0, 0), False),
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_rdm_1(norb: int, nelec: tuple[int, int], spin_summed: bool):
+    """Test computing 1-RDM."""
+    func = {
+        # spin_summed: function
+        False: _rdm1,
+        True: _rdm1_spin_summed,
+    }
+    rng = np.random.default_rng()
+    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
+
+    rdm = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+    expected = func[spin_summed](vec, norb, nelec)
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin_summed, reordered",
+    [
+        (4, (2, 2), True, True),
+        (4, (2, 2), True, False),
+        (4, (2, 2), False, True),
+        (4, (2, 2), False, False),
+        (3, (1, 2), True, True),
+        (3, (1, 2), True, False),
+        (3, (1, 2), False, True),
+        (3, (1, 2), False, False),
+        (2, (0, 1), True, True),
+        (2, (0, 1), True, False),
+        (2, (0, 1), False, True),
+        (2, (0, 1), False, False),
+        (1, (0, 0), True, True),
+        (1, (0, 0), True, False),
+        (1, (0, 0), False, True),
+        (1, (0, 0), False, False),
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_rdm_2(norb: int, nelec: tuple[int, int], spin_summed: bool, reordered: bool):
+    """Test computing 1- and 2-RDMs."""
+    func = {
+        # (spin_summed, reorder): (rdm1_function, rdm2_function)
+        (False, False): (_rdm1, _rdm2),
+        (False, True): (_rdm1, _rdm2_reordered),
+        (True, False): (_rdm1_spin_summed, _rdm2_spin_summed),
+        (True, True): (_rdm1_spin_summed, _rdm2_spin_summed_reordered),
+    }
+    rng = np.random.default_rng()
+    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
+
+    rdm1, rdm2 = ffsim.rdm(
+        vec,
+        norb,
+        nelec,
+        rank=2,
+        spin_summed=spin_summed,
+        reordered=reordered,
+    )
+    rdm1_func, rdm2_func = func[spin_summed, reordered]
+    expected_rdm1 = rdm1_func(vec, norb, nelec)
+    expected_rdm2 = rdm2_func(vec, norb, nelec)
+    np.testing.assert_allclose(rdm1, expected_rdm1, atol=1e-12)
+    np.testing.assert_allclose(rdm2, expected_rdm2, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin_summed, reordered, rank",
+    [
+        (4, (2, 2), True, True, 2),
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_no_lower_ranks(
+    norb: int, nelec: tuple[int, int], spin_summed: bool, reordered: bool, rank: int
+):
+    """Test computing higher-rank RDM without returning lower ranks."""
+    func = {
+        # (spin_summed, reorder): function
+        (False, False, 2): _rdm2,
+        (False, True, 2): _rdm2_reordered,
+        (True, False, 2): _rdm2_spin_summed,
+        (True, True, 2): _rdm2_spin_summed_reordered,
+    }
+    rng = np.random.default_rng()
+    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
+
+    rdm = ffsim.rdm(
+        vec,
+        norb,
+        nelec,
+        rank=rank,
+        spin_summed=spin_summed,
+        reordered=reordered,
+        return_lower_ranks=False,
+    )
+    expected = func[spin_summed, reordered, rank](vec, norb, nelec)
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
+
+
+def _rdm1(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute 1-RDM directly from its definition."""
+    return scipy.linalg.block_diag(*_rdm1s(vec, norb, nelec))
 
 
 def _rdm2(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
     """Compute 2-RDM directly from its definition."""
     rdm = np.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=complex)
-    for p, q, r, s in itertools.product(range(norb), repeat=4):
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_a(p), ffsim.des_a(q), ffsim.cre_a(r), ffsim.des_a(s)): 1,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
+    aa, ab, bb = _rdm2s(vec, norb, nelec)
+    rdm[:norb, :norb, :norb, :norb] = aa
+    rdm[:norb, :norb, norb:, norb:] = ab
+    rdm[norb:, norb:, :norb, :norb] = ab.transpose(2, 3, 0, 1)
+    rdm[norb:, norb:, norb:, norb:] = bb
+    return rdm
 
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_a(p), ffsim.des_a(q), ffsim.cre_b(r), ffsim.des_b(s)): 1,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[p, q, norb + r, norb + s] = val
 
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_b(p), ffsim.des_b(q), ffsim.cre_a(r), ffsim.des_a(s)): 1,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[norb + p, norb + q, r, s] = val
-
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_b(p), ffsim.des_b(q), ffsim.cre_b(r), ffsim.des_b(s)): 1,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[norb + p, norb + q, norb + r, norb + s] = val
+def _rdm2_reordered(vec: np.ndarray, norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    """Compute reordered 2-RDM directly from its definition."""
+    rdm = np.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=complex)
+    aa, ab, bb = _rdm2s_reordered(vec, norb, nelec)
+    rdm[:norb, :norb, :norb, :norb] = aa
+    rdm[:norb, :norb, norb:, norb:] = ab
+    rdm[norb:, norb:, :norb, :norb] = ab.transpose(2, 3, 0, 1)
+    rdm[norb:, norb:, norb:, norb:] = bb
     return rdm
diff --git a/tests/python/states/states_test.py b/tests/python/states/states_test.py
index 43e0ba32..388747e4 100644
--- a/tests/python/states/states_test.py
+++ b/tests/python/states/states_test.py
@@ -109,60 +109,33 @@ def test_hartree_fock_state_spinless(norb: int, nelec: int):
     assert all(vec[1:] == 0)
 
 
-@pytest.mark.parametrize(
-    "norb, nelec, spin_summed",
-    [
-        (norb, nelec, spin_summed)
-        for (norb, nelec), spin_summed in itertools.product(
-            ffsim.testing.generate_norb_nelec(range(5)), [False, True]
-        )
-    ],
-)
-def test_slater_determinant_one_rdm_same_rotation(
-    norb: int, nelec: tuple[int, int], spin_summed: bool
-):
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5)))
+def test_slater_determinant_rdm1s_same_rotation(norb: int, nelec: tuple[int, int]):
     """Test Slater determinant 1-RDM."""
     rng = np.random.default_rng()
 
     occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng)
-    occ_a, occ_b = occupied_orbitals
-    nelec = len(occ_a), len(occ_b)
-
     orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
 
     vec = ffsim.slater_determinant(
         norb, occupied_orbitals, orbital_rotation=orbital_rotation
     )
-    rdm = ffsim.slater_determinant_rdm(
+    rdm = ffsim.slater_determinant_rdms(
         norb,
         occupied_orbitals,
         orbital_rotation=orbital_rotation,
-        spin_summed=spin_summed,
     )
-    expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+    expected = ffsim.rdms(vec, norb, nelec)
 
     np.testing.assert_allclose(rdm, expected, atol=1e-12)
 
 
-@pytest.mark.parametrize(
-    "norb, nelec, spin_summed",
-    [
-        (norb, nelec, spin_summed)
-        for (norb, nelec), spin_summed in itertools.product(
-            ffsim.testing.generate_norb_nelec(range(5)), [False, True]
-        )
-    ],
-)
-def test_slater_determinant_one_rdm_diff_rotation(
-    norb: int, nelec: tuple[int, int], spin_summed: bool
-):
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5)))
+def test_slater_determinant_rdm1s_diff_rotation(norb: int, nelec: tuple[int, int]):
     """Test Slater determinant 1-RDM."""
     rng = np.random.default_rng()
 
     occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng)
-    occ_a, occ_b = occupied_orbitals
-    nelec = len(occ_a), len(occ_b)
-
     orbital_rotation_a = ffsim.random.random_unitary(norb, seed=rng)
     orbital_rotation_b = ffsim.random.random_unitary(norb, seed=rng)
 
@@ -171,13 +144,40 @@ def test_slater_determinant_one_rdm_diff_rotation(
         occupied_orbitals,
         orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
     )
-    rdm = ffsim.slater_determinant_rdm(
+    expected = ffsim.rdms(vec, norb, nelec)
+
+    rdm = ffsim.slater_determinant_rdms(
         norb,
         occupied_orbitals,
         orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
-        spin_summed=spin_summed,
     )
-    expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
+
+    rdm = ffsim.slater_determinant_rdms(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=np.stack([orbital_rotation_a, orbital_rotation_b]),
+    )
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
+
+
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nocc(range(5)))
+def test_slater_determinant_rdm1s_spinless(norb: int, nelec: int):
+    """Test Slater determinant 1-RDM, spinless."""
+    rng = np.random.default_rng()
+
+    occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng)
+    orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
+
+    vec = ffsim.slater_determinant(
+        norb, occupied_orbitals, orbital_rotation=orbital_rotation
+    )
+    rdm = ffsim.slater_determinant_rdms(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=orbital_rotation,
+    )
+    expected = ffsim.rdms(vec, norb, (nelec, 0), spin_summed=True)
 
     np.testing.assert_allclose(rdm, expected, atol=1e-12)
 
@@ -268,3 +268,72 @@ def test_indices_and_strings_roundtrip_spinful(norb: int, nelec: tuple[int, int]
     strings = ffsim.indices_to_strings(indices, norb=norb, nelec=nelec)
     indices_again = ffsim.strings_to_indices(strings, norb=norb, nelec=nelec)
     np.testing.assert_array_equal(indices_again, indices)
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin_summed",
+    [
+        (norb, nelec, spin_summed)
+        for (norb, nelec), spin_summed in itertools.product(
+            ffsim.testing.generate_norb_nelec(range(5)), [False, True]
+        )
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_slater_determinant_one_rdm_same_rotation(
+    norb: int, nelec: tuple[int, int], spin_summed: bool
+):
+    """Test Slater determinant 1-RDM."""
+    rng = np.random.default_rng()
+
+    occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng)
+    orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
+
+    vec = ffsim.slater_determinant(
+        norb, occupied_orbitals, orbital_rotation=orbital_rotation
+    )
+    rdm = ffsim.slater_determinant_rdm(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=orbital_rotation,
+        spin_summed=spin_summed,
+    )
+    expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin_summed",
+    [
+        (norb, nelec, spin_summed)
+        for (norb, nelec), spin_summed in itertools.product(
+            ffsim.testing.generate_norb_nelec(range(5)), [False, True]
+        )
+    ],
+)
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_slater_determinant_one_rdm_diff_rotation(
+    norb: int, nelec: tuple[int, int], spin_summed: bool
+):
+    """Test Slater determinant 1-RDM."""
+    rng = np.random.default_rng()
+
+    occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng)
+    orbital_rotation_a = ffsim.random.random_unitary(norb, seed=rng)
+    orbital_rotation_b = ffsim.random.random_unitary(norb, seed=rng)
+
+    vec = ffsim.slater_determinant(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
+    )
+    rdm = ffsim.slater_determinant_rdm(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
+        spin_summed=spin_summed,
+    )
+    expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed)
+
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
diff --git a/tests/python/states/wick_test.py b/tests/python/states/wick_test.py
index 348dfabc..0ea914df 100644
--- a/tests/python/states/wick_test.py
+++ b/tests/python/states/wick_test.py
@@ -52,11 +52,10 @@ def test_expectation_product(norb: int, occupied_orbitals: tuple[list[int], list
 
     # generate a random Slater determinant
     orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
-    rdm = ffsim.slater_determinant_rdm(
-        norb,
-        occupied_orbitals,
-        orbital_rotation=orbital_rotation,
-        spin_summed=False,
+    rdm = scipy.linalg.block_diag(
+        *ffsim.slater_determinant_rdms(
+            norb, occupied_orbitals, orbital_rotation=orbital_rotation
+        )
     )
 
     # get the full statevector
@@ -93,11 +92,10 @@ def test_expectation_power():
 
     # generate a random Slater determinant
     orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
-    rdm = ffsim.slater_determinant_rdm(
-        norb,
-        occupied_orbitals,
-        orbital_rotation=orbital_rotation,
-        spin_summed=False,
+    rdm = scipy.linalg.block_diag(
+        *ffsim.slater_determinant_rdms(
+            norb, occupied_orbitals, orbital_rotation=orbital_rotation
+        )
     )
 
     # get the full statevector
diff --git a/tests/python/trotter/qdrift_test.py b/tests/python/trotter/qdrift_test.py
index f08e4a97..3fd84236 100644
--- a/tests/python/trotter/qdrift_test.py
+++ b/tests/python/trotter/qdrift_test.py
@@ -274,7 +274,9 @@ def test_simulate_qdrift_double_factorized_h_chain(
     occupied_orbitals = (range(n_alpha), range(n_beta))
     initial_state = ffsim.slater_determinant(norb, occupied_orbitals)
     original_state = initial_state.copy()
-    one_rdm = ffsim.slater_determinant_rdm(norb, occupied_orbitals, spin_summed=False)
+    one_rdm = scipy.linalg.block_diag(
+        *ffsim.slater_determinant_rdms(norb, occupied_orbitals)
+    )
 
     # compute exact state
     exact_state = scipy.sparse.linalg.expm_multiply(
@@ -372,7 +374,9 @@ def test_simulate_qdrift_double_factorized_random(
     occupied_orbitals = (range(n_alpha), range(n_beta))
     initial_state = ffsim.slater_determinant(norb, occupied_orbitals)
     original_state = initial_state.copy()
-    one_rdm = ffsim.slater_determinant_rdm(norb, occupied_orbitals, spin_summed=False)
+    one_rdm = scipy.linalg.block_diag(
+        *ffsim.slater_determinant_rdms(norb, occupied_orbitals)
+    )
 
     # compute exact state
     exact_state = scipy.sparse.linalg.expm_multiply(