diff --git a/python/ffsim/contract/num_op_sum.py b/python/ffsim/contract/num_op_sum.py index d48310112..a835e3456 100644 --- a/python/ffsim/contract/num_op_sum.py +++ b/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): diff --git a/python/ffsim/hamiltonians/double_factorized_hamiltonian.py b/python/ffsim/hamiltonians/double_factorized_hamiltonian.py index 96e85ea54..c8247a45c 100644 --- a/python/ffsim/hamiltonians/double_factorized_hamiltonian.py +++ b/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,6 +225,7 @@ 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: @@ -228,6 +233,35 @@ def from_molecular_hamiltonian( 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 diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py index 87faed17b..a7bd874d2 100644 --- a/python/ffsim/hamiltonians/molecular_hamiltonian.py +++ b/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 ) diff --git a/tests/hamiltonians/double_factorized_hamiltonian_test.py b/tests/hamiltonians/double_factorized_hamiltonian_test.py index 2c25167e5..add976df2 100644 --- a/tests/hamiltonians/double_factorized_hamiltonian_test.py +++ b/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():