diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py index 99ac7f48d..dc20c867e 100644 --- a/python/ffsim/__init__.py +++ b/python/ffsim/__init__.py @@ -57,7 +57,7 @@ one_hot, rdm, slater_determinant, - slater_determinant_one_rdm, + slater_determinant_rdm, ) from ffsim.trotter import ( simulate_qdrift_double_factorized, @@ -113,7 +113,7 @@ "simulate_qdrift_double_factorized", "simulate_trotter_double_factorized", "slater_determinant", - "slater_determinant_one_rdm", + "slater_determinant_rdm", "testing", "trace", ] diff --git a/python/ffsim/states/__init__.py b/python/ffsim/states/__init__.py index f0f68cc78..d7d70820a 100644 --- a/python/ffsim/states/__init__.py +++ b/python/ffsim/states/__init__.py @@ -18,7 +18,7 @@ one_hot, rdm, slater_determinant, - slater_determinant_one_rdm, + slater_determinant_rdm, ) from ffsim.states.wick import expectation_one_body_power, expectation_one_body_product @@ -32,5 +32,5 @@ "one_hot", "rdm", "slater_determinant", - "slater_determinant_one_rdm", + "slater_determinant_rdm", ] diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py index 5e449858d..3aaa22838 100644 --- a/python/ffsim/states/states.py +++ b/python/ffsim/states/states.py @@ -130,38 +130,50 @@ def hartree_fock_state(norb: int, nelec: tuple[int, int]) -> np.ndarray: ) -def slater_determinant_one_rdm( +def slater_determinant_rdm( norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | None = None, + rank: int = 1, spin_summed: bool = True, ) -> np.ndarray: - """Return the one-particle reduced density matrix of a Slater determinant. + """Return the reduced density matrix of a Slater determinant. + + Note: + Currently, only rank 1 is supported. Args: norb: The number of spatial orbitals. 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. + orbital_rotation: An optional orbital rotation to apply to the + electron configuration. In other words, this is a unitary matrix that + describes the orbitals of the Slater determinant. + rank: The rank of the reduced density matrix. spin_summed: Whether to sum over the spin index. Returns: - The one-particle reduced density matrix of the Slater determinant. + The reduced density matrix of the Slater determinant. """ - 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]) - if len(alpha_orbitals): - rdm_a[(alpha_orbitals, alpha_orbitals)] = 1 - if len(beta_orbitals): - rdm_b[(beta_orbitals, beta_orbitals)] = 1 - if orbital_rotation is not None: - rdm_a = orbital_rotation.conj() @ rdm_a @ orbital_rotation.T - rdm_b = orbital_rotation.conj() @ rdm_b @ orbital_rotation.T - if spin_summed: - return rdm_a + rdm_b - return scipy.linalg.block_diag(rdm_a, rdm_b) + if rank == 1: + 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]) + if len(alpha_orbitals): + rdm_a[(alpha_orbitals, alpha_orbitals)] = 1 + if len(beta_orbitals): + rdm_b[(beta_orbitals, beta_orbitals)] = 1 + if orbital_rotation is not None: + rdm_a = orbital_rotation.conj() @ rdm_a @ orbital_rotation.T + rdm_b = orbital_rotation.conj() @ rdm_b @ orbital_rotation.T + if spin_summed: + return rdm_a + rdm_b + return scipy.linalg.block_diag(rdm_a, rdm_b) + raise NotImplementedError( + f"Returning the rank {rank} reduced density matrix is currently not supported." + ) def indices_to_strings( diff --git a/tests/states/states_test.py b/tests/states/states_test.py index ab9bc5964..aeae0a165 100644 --- a/tests/states/states_test.py +++ b/tests/states/states_test.py @@ -84,7 +84,7 @@ def test_slater_determinant_one_rdm( vec = ffsim.slater_determinant( norb, occupied_orbitals, orbital_rotation=orbital_rotation ) - rdm = ffsim.slater_determinant_one_rdm( + rdm = ffsim.slater_determinant_rdm( norb, occupied_orbitals, orbital_rotation=orbital_rotation, diff --git a/tests/trotter/qdrift_test.py b/tests/trotter/qdrift_test.py index 47d9403e1..3a0619fb8 100644 --- a/tests/trotter/qdrift_test.py +++ b/tests/trotter/qdrift_test.py @@ -274,9 +274,7 @@ 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_one_rdm( - norb, occupied_orbitals, spin_summed=False - ) + one_rdm = ffsim.slater_determinant_rdm(norb, occupied_orbitals, spin_summed=False) # compute exact state exact_state = scipy.sparse.linalg.expm_multiply( @@ -374,9 +372,7 @@ 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_one_rdm( - norb, occupied_orbitals, spin_summed=False - ) + one_rdm = ffsim.slater_determinant_rdm(norb, occupied_orbitals, spin_summed=False) # compute exact state exact_state = scipy.sparse.linalg.expm_multiply(