add orbital_rotation argument to slater determinant one rdm (#66)

* add orbital_rotation argument to slater determinant one rdm

* fix tests

python/ffsim/states/states.py

@@ -133,6 +133,7 @@ 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]],
+    orbital_rotation: np.ndarray | None = None,
     spin_summed: bool = True,
 ) -> np.ndarray:
     """Return the one-particle reduced density matrix of a Slater determinant.
@@ -155,6 +156,9 @@ def slater_determinant_one_rdm(
         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)

tests/states/states_test.py

@@ -78,13 +78,21 @@ def test_slater_determinant_one_rdm(
     occ_a, occ_b = occupied_orbitals
     nelec = len(occ_a), len(occ_b)
 
-    vec = ffsim.slater_determinant(norb, occupied_orbitals)
+    rng = np.random.default_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_one_rdm(
-        norb, occupied_orbitals, spin_summed=spin_summed
+        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)
+    np.testing.assert_allclose(rdm, expected, atol=1e-12)
 
 
 def test_indices_to_strings():

tests/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_one_rdm(norb, occupied_orbitals)
+    one_rdm = ffsim.slater_determinant_one_rdm(
+        norb, occupied_orbitals, spin_summed=False
+    )
 
     # 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_one_rdm(norb, occupied_orbitals)
+    one_rdm = ffsim.slater_determinant_one_rdm(
+        norb, occupied_orbitals, spin_summed=False
+    )
 
     # compute exact state
     exact_state = scipy.sparse.linalg.expm_multiply(