Repository: qiskit-community/ffsim (python)

Issue:
add return_lower_ranks argument to rdm
For example, the function `make_rdm12s` from pyscf returns not only the 2-RDM but also the 1-RDM. Currently, we just discard the 1-RDM. If `ffsim.rdm` is called with `return_lower_ranks=True`, then it should return all of the RDMs up to the given rank, not just the single one.

Pull request description:
Return lower ranks of RDMs
Fixes #75 

Adds the `return_lower_ranks` argument and sets its default value to `True`, which is why this is a breaking change (the previous behavior was equivalent to it being `False`).

Patch:
diff --git a/python/ffsim/states/__init__.py b/python/ffsim/states/__init__.py
index d7d70820..c436d66d 100644
--- a/python/ffsim/states/__init__.py
+++ b/python/ffsim/states/__init__.py
@@ -10,13 +10,13 @@
 
 """States."""
 
+from ffsim.states.rdm import rdm
 from ffsim.states.states import (
     dim,
     dims,
     hartree_fock_state,
     indices_to_strings,
     one_hot,
-    rdm,
     slater_determinant,
     slater_determinant_rdm,
 )
diff --git a/python/ffsim/states/rdm.py b/python/ffsim/states/rdm.py
new file mode 100644
index 00000000..a5baa817
--- /dev/null
+++ b/python/ffsim/states/rdm.py
@@ -0,0 +1,268 @@
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Functions to compute reduced density matrices."""
+
+from __future__ import annotations
+
+import numpy as np
+import scipy.linalg
+from pyscf.fci.direct_spin1 import (
+    make_rdm1,
+    make_rdm1s,
+    make_rdm12,
+    make_rdm12s,
+    trans_rdm1,
+    trans_rdm1s,
+    trans_rdm12,
+    trans_rdm12s,
+)
+
+from ffsim.cistring import gen_linkstr_index
+
+
+def rdm(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    rank: int = 1,
+    spin_summed: bool = True,
+    reordered: bool = True,
+    return_lower_ranks: bool = True,
+) -> np.ndarray | tuple[np.ndarray, ...]:
+    """Return the reduced density matrix (RDM) or 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 ``reordered`` argument, which
+    defaults to True.
+
+    **reordered = 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⟩
+
+    **reordered = False**
+
+    If reordered 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 sum over the spin index.
+        reordered: Whether to reorder the indices of the reduced density matrix.
+        return_lower_ranks: Whether to return lower rank RDMs in addition to the
+            specified rank. If True, then this function returns all RDMs up to and
+            including the specified rank, in increasing order of rank. For example,
+            if `rank=2` then a tuple `(rdm1, rdm2)` is returned.
+
+    Returns:
+        The reduced density matrix or matrices. If `return_lower_ranks` is False,
+        then a single matrix is returned. If `return_lower_ranks` is True, then
+        a `rank`-length tuple of matrices is returned, containing the RDMs up to
+        the specified rank in increasing order of rank.
+    """
+    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 _rdm1(vec, norb, nelec, link_index)
+    if rank == 2:
+        if spin_summed:
+            return _rdm2_spin_summed(
+                vec, norb, nelec, reordered, link_index, return_lower_ranks
+            )
+        else:
+            return _rdm2(vec, norb, nelec, reordered, link_index, return_lower_ranks)
+    raise NotImplementedError(
+        f"Computing the rank {rank} reduced density matrix is currently not supported."
+    )
+
+
+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
+    )
+    trans_rdm1_imag_real = trans_rdm1(
+        vec.imag, vec.real, norb, nelec, link_index=link_index
+    )
+    return _assemble_rdm1_spin_summed(
+        rdm1_real, rdm1_imag, trans_rdm1_real_imag, trans_rdm1_imag_real
+    )
+
+
+def _rdm1(
+    vec: np.ndarray,
+    norb: int,
+    nelec: tuple[int, int],
+    link_index: tuple[np.ndarray, np.ndarray] | None,
+) -> np.ndarray:
+    rdms1_real = make_rdm1s(vec.real, norb, nelec, link_index=link_index)
+    rdms1_imag = make_rdm1s(vec.imag, norb, nelec, link_index=link_index)
+    trans_rdms1_real_imag = trans_rdm1s(
+        vec.real, vec.imag, norb, nelec, link_index=link_index
+    )
+    trans_rdms1_imag_real = trans_rdm1s(
+        vec.imag, vec.real, norb, nelec, link_index=link_index
+    )
+    return _assemble_rdm1(
+        rdms1_real, rdms1_imag, trans_rdms1_real_imag, trans_rdms1_imag_real
+    )
+
+
+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 _rdm2(
+    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]:
+    rdms1_real, rdms2_real = make_rdm12s(
+        vec.real, norb, nelec, reorder=reordered, link_index=link_index
+    )
+    rdms1_imag, rdms2_imag = make_rdm12s(
+        vec.imag, norb, nelec, reorder=reordered, link_index=link_index
+    )
+    trans_rdms1_real_imag, trans_rdms2_real_imag = trans_rdm12s(
+        vec.real, vec.imag, norb, nelec, reorder=reordered, link_index=link_index
+    )
+    trans_rdms1_imag_real, trans_rdms2_imag_real = trans_rdm12s(
+        vec.imag, vec.real, norb, nelec, reorder=reordered, link_index=link_index
+    )
+    rdm2 = _assemble_rdm2(
+        rdms2_real, rdms2_imag, trans_rdms2_real_imag, trans_rdms2_imag_real
+    )
+    if not return_lower_ranks:
+        return rdm2
+    rdm1 = _assemble_rdm1(
+        rdms1_real, rdms1_imag, trans_rdms1_real_imag, trans_rdms1_imag_real
+    )
+    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],
+    trans_rdms2_real_imag: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
+    trans_rdms2_imag_real: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
+) -> np.ndarray:
+    rdms2 = np.stack(rdms2_real).astype(complex)
+    rdms2 += np.stack(rdms2_imag)
+    # rdms2 is currently [rdm_aa, rdm_ab, rdm_bb]
+    # the following line transforms it into [rdm_aa, rdm_ab, rdm_ba, rdm_bb]
+    rdms2 = np.insert(rdms2, 2, rdms2[1].transpose(2, 3, 0, 1), axis=0)
+    rdms2 += 1j * np.stack(trans_rdms2_real_imag)
+    rdms2 -= 1j * np.stack(trans_rdms2_imag_real)
+
+    norb, _, _, _ = rdms2_real[0].shape
+    rdm2 = np.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=complex)
+    rdm_aa, rdm_ab, rdm_ba, rdm_bb = rdms2
+    rdm2[:norb, :norb, :norb, :norb] = rdm_aa
+    rdm2[:norb, :norb, norb:, norb:] = rdm_ab
+    rdm2[norb:, norb:, :norb, :norb] = rdm_ba
+    rdm2[norb:, norb:, norb:, norb:] = rdm_bb
+    return rdm2
diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py
index 7ecc72fb..261fc91d 100644
--- a/python/ffsim/states/states.py
+++ b/python/ffsim/states/states.py
@@ -8,6 +8,8 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
+"""Fermionic quantum states."""
+
 from __future__ import annotations
 
 from collections.abc import Sequence
@@ -15,19 +17,8 @@ from collections.abc import Sequence
 import numpy as np
 import scipy.linalg
 from pyscf.fci import cistring
-from pyscf.fci.direct_spin1 import (
-    make_rdm1,
-    make_rdm1s,
-    make_rdm12,
-    make_rdm12s,
-    trans_rdm1,
-    trans_rdm1s,
-    trans_rdm12,
-    trans_rdm12s,
-)
 from scipy.special import comb
 
-from ffsim.cistring import gen_linkstr_index
 from ffsim.gates.orbital_rotation import apply_orbital_rotation
 
 
@@ -199,180 +190,3 @@ def indices_to_strings(
         f"{string_a:0{norb}b}{string_b:0{norb}b}"
         for string_a, string_b in zip(strings_a, strings_b)
     ]
-
-
-def rdm(
-    vec: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    rank: int = 1,
-    spin_summed: bool = True,
-    reordered: bool = True,
-) -> np.ndarray:
-    """Return the reduced density matrix (RDM) 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 ``reordered`` argument, which
-    defaults to True.
-
-    **reordered = 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⟩
-
-    **reordered = False**
-
-    If reordered 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 sum over the spin index.
-        reordered: Whether to reorder the indices of the reduced density matrix.
-
-    Returns:
-        The reduced density matrix.
-    """
-    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)
-    vec_real = np.real(vec)
-    vec_imag = np.imag(vec)
-    if rank == 1:
-        if spin_summed:
-            return _rdm1_spin_summed(vec_real, vec_imag, norb, nelec, link_index)
-        else:
-            return _rdm1(vec_real, vec_imag, norb, nelec, link_index)
-    if rank == 2:
-        if spin_summed:
-            return _rdm2_spin_summed(
-                vec_real, vec_imag, norb, nelec, reordered, link_index
-            )
-        else:
-            return _rdm2(vec_real, vec_imag, norb, nelec, reordered, link_index)
-    raise NotImplementedError(
-        f"Computing the rank {rank} reduced density matrix is currently not supported."
-    )
-
-
-def _rdm1_spin_summed(
-    vec_real: np.ndarray,
-    vec_imag: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    link_index: tuple[np.ndarray, np.ndarray] | None,
-) -> np.ndarray:
-    result = make_rdm1(vec_real, norb, nelec, link_index=link_index).astype(complex)
-    result += make_rdm1(vec_imag, norb, nelec, link_index=link_index)
-    # the rdm1 convention from pyscf is tranposed
-    result -= 1j * trans_rdm1(vec_real, vec_imag, norb, nelec, link_index=link_index)
-    result += 1j * trans_rdm1(vec_imag, vec_real, norb, nelec, link_index=link_index)
-    return result
-
-
-def _rdm1(
-    vec_real: np.ndarray,
-    vec_imag: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    link_index: tuple[np.ndarray, np.ndarray] | None,
-) -> np.ndarray:
-    result = np.stack(make_rdm1s(vec_real, norb, nelec, link_index=link_index)).astype(
-        complex
-    )
-    result += np.stack(make_rdm1s(vec_imag, norb, nelec, link_index=link_index))
-    # the rdm1 convention from pyscf is tranposed
-    result -= 1j * np.stack(
-        trans_rdm1s(vec_real, vec_imag, norb, nelec, link_index=link_index)
-    )
-    result += 1j * np.stack(
-        trans_rdm1s(vec_imag, vec_real, norb, nelec, link_index=link_index)
-    )
-    return scipy.linalg.block_diag(*result)
-
-
-def _rdm2_spin_summed(
-    vec_real: np.ndarray,
-    vec_imag: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    reordered: bool,
-    link_index: tuple[np.ndarray, np.ndarray] | None,
-) -> np.ndarray:
-    result = make_rdm12(
-        vec_real, norb, nelec, reorder=reordered, link_index=link_index
-    )[1].astype(complex)
-    result += make_rdm12(
-        vec_imag, norb, nelec, reorder=reordered, link_index=link_index
-    )[1]
-    result += (
-        1j
-        * trans_rdm12(
-            vec_real, vec_imag, norb, nelec, reorder=reordered, link_index=link_index
-        )[1]
-    )
-    result -= (
-        1j
-        * trans_rdm12(
-            vec_imag, vec_real, norb, nelec, reorder=reordered, link_index=link_index
-        )[1]
-    )
-    return result
-
-
-def _rdm2(
-    vec_real: np.ndarray,
-    vec_imag: np.ndarray,
-    norb: int,
-    nelec: tuple[int, int],
-    reordered: bool,
-    link_index: tuple[np.ndarray, np.ndarray] | None,
-) -> np.ndarray:
-    rdms = np.stack(
-        make_rdm12s(vec_real, norb, nelec, reorder=reordered, link_index=link_index)[1]
-    ).astype(complex)
-    rdms += np.stack(
-        make_rdm12s(vec_imag, norb, nelec, reorder=reordered, link_index=link_index)[1]
-    )
-    # result is currently [rdm_aa, rdm_ab, rdm_bb]
-    # the following line transforms it into [rdm_aa, rdm_ab, rdm_ba, rdm_bb]
-    rdms = np.insert(rdms, 2, rdms[1].transpose(2, 3, 0, 1), axis=0)
-    rdms += 1j * np.stack(
-        trans_rdm12s(
-            vec_real, vec_imag, norb, nelec, reorder=reordered, link_index=link_index
-        )[1]
-    )
-    rdms -= 1j * np.stack(
-        trans_rdm12s(
-            vec_imag, vec_real, norb, nelec, reorder=reordered, link_index=link_index
-        )[1]
-    )
-    result = np.zeros((2 * norb, 2 * norb, 2 * norb, 2 * norb), dtype=complex)
-    rdm_aa, rdm_ab, rdm_ba, rdm_bb = rdms
-    result[:norb, :norb, :norb, :norb] = rdm_aa
-    result[:norb, :norb, norb:, norb:] = rdm_ab
-    result[norb:, norb:, :norb, :norb] = rdm_ba
-    result[norb:, norb:, norb:, norb:] = rdm_bb
-    return result


Test patch:
diff --git a/tests/states/rdm_test.py b/tests/states/rdm_test.py
new file mode 100644
index 00000000..106384de
--- /dev/null
+++ b/tests/states/rdm_test.py
@@ -0,0 +1,293 @@
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Test reduced density matrices."""
+
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+import pytest
+
+import ffsim
+
+
+@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),
+    ],
+)
+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),
+    ],
+)
+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),
+    ],
+)
+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:
+    """Compute 1-RDM directly from its definition."""
+    rdm = np.zeros((2 * norb, 2 * 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,
+            }
+        )
+        linop = ffsim.linear_operator(op, norb, nelec)
+        val = np.vdot(vec, linop @ vec)
+        rdm[i, j] = val
+        rdm[j, i] = val.conjugate()
+        op = ffsim.FermionOperator(
+            {
+                (ffsim.cre_b(i), ffsim.des_b(j)): 1,
+            }
+        )
+        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
+
+
+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)
+    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,
+            }
+        )
+        linop = ffsim.linear_operator(op, norb, nelec)
+        val = np.vdot(vec, linop @ vec)
+        rdm[p, q, r, s] = val
+    return rdm
+
+
+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
+
+
+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,
+            }
+        )
+        linop = ffsim.linear_operator(op, norb, nelec)
+        val = np.vdot(vec, linop @ vec)
+        rdm[p, q, r, s] = val
+
+        op = ffsim.FermionOperator(
+            {
+                (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[p, q, norb + r, norb + s] = val
+
+        op = ffsim.FermionOperator(
+            {
+                (ffsim.cre_b(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_b(q)): 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.cre_b(r), ffsim.des_b(s), ffsim.des_b(q)): 1,
+            }
+        )
+        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
+
+
+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
+
+        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
+    return rdm
diff --git a/tests/states/states_test.py b/tests/states/states_test.py
index 647889f8..eb9e0eec 100644
--- a/tests/states/states_test.py
+++ b/tests/states/states_test.py
@@ -12,8 +12,6 @@
 
 from __future__ import annotations
 
-import itertools
-
 import numpy as np
 import pyscf
 import pytest
@@ -113,281 +111,3 @@ def test_indices_to_strings():
         "110010",
         "110100",
     ]
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_1_spin_summed(norb: int, nelec: tuple[int, int]):
-    """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)
-    expected = _rdm1_spin_summed(vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_1(norb: int, nelec: tuple[int, int]):
-    """Test computing 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=False)
-    expected = _rdm1(vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_2_spin_summed_reordered(norb: int, nelec: tuple[int, int]):
-    """Test computing spin-summed reordered 2-RDM."""
-    rng = np.random.default_rng()
-    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(vec, norb, nelec, rank=2)
-    expected = _rdm2_spin_summed_reordered(vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_2_reordered(norb: int, nelec: tuple[int, int]):
-    """Test computing reordered 2-RDM."""
-    rng = np.random.default_rng()
-    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(vec, norb, nelec, rank=2, spin_summed=False)
-    expected = _rdm2_reordered(vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_2_spin_summed(norb: int, nelec: tuple[int, int]):
-    """Test computing spin-summed 2-RDM."""
-    rng = np.random.default_rng()
-    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(vec, norb, nelec, rank=2, reordered=False)
-    expected = _rdm2_spin_summed(vec, norb, nelec)
-    np.testing.assert_allclose(rdm, expected, atol=1e-12)
-
-
-@pytest.mark.parametrize(
-    "norb, nelec",
-    [
-        (4, (2, 2)),
-        (3, (1, 2)),
-        (2, (0, 1)),
-        (1, (0, 0)),
-    ],
-)
-def test_rdm_2(norb: int, nelec: tuple[int, int]):
-    """Test computing 2-RDM."""
-    rng = np.random.default_rng()
-    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
-
-    rdm = ffsim.rdm(vec, norb, nelec, rank=2, spin_summed=False, reordered=False)
-    expected = _rdm2(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:
-    """Compute 1-RDM directly from its definition."""
-    rdm = np.zeros((2 * norb, 2 * 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,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[i, j] = val
-        rdm[j, i] = val.conjugate()
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_b(i), ffsim.des_b(j)): 1,
-            }
-        )
-        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
-
-
-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)
-    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,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
-    return rdm
-
-
-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
-
-
-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,
-            }
-        )
-        linop = ffsim.linear_operator(op, norb, nelec)
-        val = np.vdot(vec, linop @ vec)
-        rdm[p, q, r, s] = val
-
-        op = ffsim.FermionOperator(
-            {
-                (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[p, q, norb + r, norb + s] = val
-
-        op = ffsim.FermionOperator(
-            {
-                (ffsim.cre_b(p), ffsim.cre_a(r), ffsim.des_a(s), ffsim.des_b(q)): 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.cre_b(r), ffsim.des_b(s), ffsim.des_b(q)): 1,
-            }
-        )
-        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
-
-
-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
-
-        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
-    return rdm