implement linear operator protocol for df hamiltonian (#73)

* implement linear operator protocol for df hamiltonian

* lint

python/ffsim/contract/num_op_sum.py

@@ -103,7 +103,6 @@ def num_op_sum_linop(
         A LinearOperator that implements the action of the linear combination of number
         operators.
     """
-    n_alpha, n_beta = nelec
     dim_ = dim(norb, nelec)
 
     def matvec(vec):

python/ffsim/hamiltonians/double_factorized_hamiltonian.py

@@ -13,9 +13,13 @@
 import dataclasses
 
 import numpy as np
+from scipy.sparse.linalg import LinearOperator
 
+from ffsim.contract.diag_coulomb import diag_coulomb_linop
+from ffsim.contract.num_op_sum import num_op_sum_linop
 from ffsim.hamiltonians.molecular_hamiltonian import MolecularHamiltonian
 from ffsim.linalg import double_factorized
+from ffsim.states import dim
 
 
 @dataclasses.dataclass(frozen=True)
@@ -221,13 +225,43 @@ def from_molecular_hamiltonian(
             one_body_tensor=one_body_tensor,
             diag_coulomb_mats=diag_coulomb_mats,
             orbital_rotations=orbital_rotations,
+            constant=hamiltonian.constant,
         )
 
         if z_representation:
             df_hamiltonian = df_hamiltonian.to_z_representation()
 
         return df_hamiltonian
 
+    def _linear_operator_(self, norb: int, nelec: tuple[int, int]) -> LinearOperator:
+        """Return a SciPy LinearOperator representing the object."""
+        dim_ = dim(norb, nelec)
+        eigs, vecs = np.linalg.eigh(self.one_body_tensor)
+        num_linop = num_op_sum_linop(eigs, norb, nelec, orbital_rotation=vecs)
+        diag_coulomb_linops = [
+            diag_coulomb_linop(
+                diag_coulomb_mat,
+                norb,
+                nelec,
+                orbital_rotation=orbital_rotation,
+                z_representation=self.z_representation,
+            )
+            for diag_coulomb_mat, orbital_rotation in zip(
+                self.diag_coulomb_mats, self.orbital_rotations
+            )
+        ]
+
+        def matvec(vec: np.ndarray):
+            result = self.constant * vec
+            result += num_linop @ vec
+            for linop in diag_coulomb_linops:
+                result += linop @ vec
+            return result
+
+        return LinearOperator(
+            shape=(dim_, dim_), matvec=matvec, rmatvec=matvec, dtype=complex
+        )
+
 
 def _df_z_representation(
     diag_coulomb_mats: np.ndarray, orbital_rotations: np.ndarray

python/ffsim/hamiltonians/molecular_hamiltonian.py

@@ -14,7 +14,6 @@
 import itertools
 
 import numpy as np
-import scipy.sparse.linalg
 from pyscf.fci.direct_nosym import absorb_h1e, contract_1e, contract_2e, make_hdiag
 from scipy.sparse.linalg import LinearOperator
 
@@ -82,7 +81,7 @@ def matvec(vec: np.ndarray):
             )
             return result
 
-        return scipy.sparse.linalg.LinearOperator(
+        return LinearOperator(
             shape=(dim_, dim_), matvec=matvec, rmatvec=matvec, dtype=complex
         )
 

tests/hamiltonians/double_factorized_hamiltonian_test.py

@@ -20,64 +20,35 @@
 
 
 @pytest.mark.parametrize("z_representation", [False, True])
-def test_double_factorized_hamiltonian(z_representation: bool):
-    # set parameters
+def test_linear_operator(z_representation: bool):
+    """Test linear operator."""
     norb = 4
     nelec = (2, 2)
+    rng = np.random.default_rng(2474)
 
-    # generate random Hamiltonian
     dim = ffsim.dim(norb, nelec)
-    one_body_tensor = ffsim.random.random_hermitian(norb, seed=2474)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=7054)
-    mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
-    hamiltonian = ffsim.linear_operator(
-        mol_hamiltonian,
-        norb=norb,
-        nelec=nelec,
+    one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
+    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=rng)
+    constant = rng.standard_normal()
+    mol_hamiltonian = ffsim.MolecularHamiltonian(
+        one_body_tensor, two_body_tensor, constant
     )
 
-    # perform double factorization
     df_hamiltonian = ffsim.DoubleFactorizedHamiltonian.from_molecular_hamiltonian(
         mol_hamiltonian,
         z_representation=z_representation,
     )
 
-    # generate random state
+    actual_linop = ffsim.linear_operator(df_hamiltonian, norb, nelec)
+    expected_linop = ffsim.linear_operator(mol_hamiltonian, norb, nelec)
+
     dim = ffsim.dim(norb, nelec)
-    state = ffsim.random.random_statevector(dim, seed=1360)
-
-    # apply Hamiltonian terms
-    result = df_hamiltonian.constant * state
-
-    eigs, vecs = np.linalg.eigh(df_hamiltonian.one_body_tensor)
-    tmp = ffsim.apply_orbital_rotation(state, vecs.T.conj(), norb=norb, nelec=nelec)
-    tmp = ffsim.contract.contract_num_op_sum(tmp, eigs, norb=norb, nelec=nelec)
-    tmp = ffsim.apply_orbital_rotation(tmp, vecs, norb=norb, nelec=nelec)
-    result += tmp
-
-    for diag_coulomb_mat, orbital_rotation in zip(
-        df_hamiltonian.diag_coulomb_mats, df_hamiltonian.orbital_rotations
-    ):
-        tmp = ffsim.apply_orbital_rotation(
-            state, orbital_rotation.T.conj(), norb=norb, nelec=nelec
-        )
-        tmp = ffsim.contract.contract_diag_coulomb(
-            tmp,
-            diag_coulomb_mat,
-            norb=norb,
-            nelec=nelec,
-            z_representation=z_representation,
-        )
-        tmp = ffsim.apply_orbital_rotation(
-            tmp, orbital_rotation, norb=norb, nelec=nelec
-        )
-        result += tmp
-
-    # apply Hamiltonian directly
-    expected = hamiltonian @ state
-
-    # check agreement
-    np.testing.assert_allclose(result, expected, atol=1e-8)
+    state = ffsim.random.random_statevector(dim, seed=rng)
+
+    actual = actual_linop @ state
+    expected = expected_linop @ state
+
+    np.testing.assert_allclose(actual, expected, atol=1e-8)
 
 
 def test_z_representation_round_trip():