From c59fbd144382988208b1db928d5363df37ce9290 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Mon, 30 Oct 2023 12:04:43 -0400
Subject: [PATCH] add orbital_rotation argument to slater determinant one rdm

---
 python/ffsim/states/states.py |  4 ++++
 tests/states/states_test.py   | 14 +++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py
index 2d0daedae..5e449858d 100644
--- a/python/ffsim/states/states.py
+++ b/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)
diff --git a/tests/states/states_test.py b/tests/states/states_test.py
index 0d4894ca3..ab9bc5964 100644
--- a/tests/states/states_test.py
+++ 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():