From 97f34c99360137b332afb6354daee46daf85ff37 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Fri, 24 Nov 2023 21:43:25 -0500
Subject: [PATCH] add orbital rotation for molecular hamiltonian

---
 .../hamiltonians/molecular_hamiltonian.py     | 40 +++++++++++++++++
 .../molecular_hamiltonian_test.py             | 45 +++++++++++++++++++
 2 files changed, 85 insertions(+)

diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py
index a7bd874d2..83eaa395c 100644
--- a/python/ffsim/hamiltonians/molecular_hamiltonian.py
+++ b/python/ffsim/hamiltonians/molecular_hamiltonian.py
@@ -14,6 +14,7 @@
 import itertools
 
 import numpy as np
+from opt_einsum import contract
 from pyscf.fci.direct_nosym import absorb_h1e, contract_1e, contract_2e, make_hdiag
 from scipy.sparse.linalg import LinearOperator
 
@@ -54,6 +55,45 @@ def norb(self) -> int:
         """The number of spatial orbitals."""
         return self.one_body_tensor.shape[0]
 
+    def rotated(self, orbital_rotation: np.ndarray) -> MolecularHamiltonian:
+        """Return the Hamiltonian in a rotated orbital basis.
+
+        Given an orbital rotation :math:`\mathcal{U}`, returns the operator
+
+        .. math::
+
+            \mathcal{U} H \mathcal{U}^\dagger
+
+        where :math:`H` is the original Hamiltonian.
+
+        Args:
+            orbital_rotation: The orbital rotation.
+
+        Returns:
+            The rotated Hamiltonian.
+        """
+        one_body_tensor_rotated = contract(
+            "ab,Aa,Bb->AB",
+            self.one_body_tensor,
+            orbital_rotation,
+            orbital_rotation.conj(),
+            optimize="greedy",
+        )
+        two_body_tensor_rotated = contract(
+            "abcd,Aa,Bb,Cc,Dd->ABCD",
+            self.two_body_tensor,
+            orbital_rotation,
+            orbital_rotation.conj(),
+            orbital_rotation,
+            orbital_rotation.conj(),
+            optimize="greedy",
+        )
+        return MolecularHamiltonian(
+            one_body_tensor=one_body_tensor_rotated,
+            two_body_tensor=two_body_tensor_rotated,
+            constant=self.constant,
+        )
+
     def _linear_operator_(self, norb: int, nelec: tuple[int, int]) -> LinearOperator:
         """Return a SciPy LinearOperator representing the object."""
         n_alpha, n_beta = nelec
diff --git a/tests/hamiltonians/molecular_hamiltonian_test.py b/tests/hamiltonians/molecular_hamiltonian_test.py
index 859bb0f1a..e34f62a9f 100644
--- a/tests/hamiltonians/molecular_hamiltonian_test.py
+++ b/tests/hamiltonians/molecular_hamiltonian_test.py
@@ -105,3 +105,48 @@ def test_fermion_operator(norb: int, nelec: tuple[int, int]):
     actual = linop @ vec
     expected = expected_linop @ vec
     np.testing.assert_allclose(actual, expected)
+
+
+def test_rotated():
+    """Test rotating orbitals."""
+    norb = 5
+    nelec = (3, 2)
+
+    rng = np.random.default_rng()
+
+    # generate a random molecular Hamiltonian
+    one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
+    # TODO remove dtype=float after adding support for complex
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=rng, dtype=float)
+    constant = rng.standard_normal()
+    mol_hamiltonian = ffsim.MolecularHamiltonian(
+        one_body_tensor, two_body_tensor, constant=constant
+    )
+
+    # generate a random orbital rotation
+    orbital_rotation = ffsim.random.random_orthogonal(norb, seed=rng)
+
+    # rotate the Hamiltonian
+    mol_hamiltonian_rotated = mol_hamiltonian.rotated(orbital_rotation)
+
+    # convert the original and rotated Hamiltonians to linear operators
+    linop = ffsim.linear_operator(mol_hamiltonian, norb, nelec)
+    linop_rotated = ffsim.linear_operator(mol_hamiltonian_rotated, norb, nelec)
+
+    # generate a random statevector
+    vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng)
+
+    # rotate the statevector
+    rotated_vec = ffsim.apply_orbital_rotation(vec, orbital_rotation, norb, nelec)
+
+    # test definition
+    actual = linop_rotated @ vec
+    expected = ffsim.apply_orbital_rotation(vec, orbital_rotation.T.conj(), norb, nelec)
+    expected = linop @ expected
+    expected = ffsim.apply_orbital_rotation(expected, orbital_rotation, norb, nelec)
+    np.testing.assert_allclose(actual, expected)
+
+    # test expectation is preserved
+    original_expectation = np.vdot(vec, linop @ vec)
+    rotated_expectation = np.vdot(rotated_vec, linop_rotated @ rotated_vec)
+    np.testing.assert_allclose(original_expectation, rotated_expectation)