add orbital rotation for molecular hamiltonian

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

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)