diff --git a/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py b/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py index b2ab966a8..e571928c9 100644 --- a/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py +++ b/python/ffsim/hamiltonians/diagonal_coulomb_hamiltonian.py @@ -15,6 +15,7 @@ import numpy as np import scipy.linalg +from pyscf.fci.direct_uhf import make_hdiag from scipy.sparse.linalg import LinearOperator from ffsim._lib import FermionOperator @@ -174,3 +175,23 @@ def from_fermion_operator(op: FermionOperator) -> DiagonalCoulombHamiltonian: diag_coulomb_mats=diag_coulomb_mats, constant=constant, ) + + def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray: + """Return the diagonal entries of the Hamiltonian.""" + if np.iscomplexobj(self.one_body_tensor): + raise NotImplementedError( + "Computing diagonal of complex diagonal Coulomb Hamiltonian is not yet " + "supported." + ) + two_body_tensor_aa = np.zeros((self.norb, self.norb, self.norb, self.norb)) + two_body_tensor_ab = np.zeros((self.norb, self.norb, self.norb, self.norb)) + diag_coulomb_mat_aa, diag_coulomb_mat_ab = self.diag_coulomb_mats + for p, q in itertools.product(range(self.norb), repeat=2): + two_body_tensor_aa[p, p, q, q] = diag_coulomb_mat_aa[p, q] + two_body_tensor_ab[p, p, q, q] = diag_coulomb_mat_ab[p, q] + one_body_tensor = self.one_body_tensor + 0.5 * np.einsum( + "prqr", two_body_tensor_aa + ) + h1e = (one_body_tensor, one_body_tensor) + h2e = (two_body_tensor_aa, two_body_tensor_ab, two_body_tensor_aa) + return make_hdiag(h1e, h2e, norb=norb, nelec=nelec) + self.constant diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py index d8ea415a7..9a65584f9 100644 --- a/python/ffsim/hamiltonians/molecular_hamiltonian.py +++ b/python/ffsim/hamiltonians/molecular_hamiltonian.py @@ -158,8 +158,9 @@ def _diag_(self, norb: int, nelec: tuple[int, int]) -> np.ndarray: if np.iscomplexobj(self.two_body_tensor) or np.iscomplexobj( self.one_body_tensor ): - raise ValueError( - "Computing diagonal of complex molecular Hamiltonian is not supported." + raise NotImplementedError( + "Computing diagonal of complex molecular Hamiltonian is not yet " + "supported." ) return ( make_hdiag(self.one_body_tensor, self.two_body_tensor, norb, nelec) diff --git a/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py b/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py index 703bea3de..e483ac3ae 100644 --- a/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py +++ b/tests/python/hamiltonians/diagonal_coulomb_hamiltonian_test.py @@ -273,3 +273,24 @@ def test_from_fermion_operator_fermi_hubbard_2d(norb: int, nelec: tuple[int, int actual_periodic = actual_linop_periodic @ vec expected_periodic = expected_linop_periodic @ vec np.testing.assert_allclose(actual_periodic, expected_periodic) + + +def test_diag(): + """Test computing diagonal.""" + rng = np.random.default_rng(2222) + norb = 5 + nelec = (3, 2) + # TODO test complex one-body after adding support for it + one_body_tensor = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + diag_coulomb_mat_a = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + diag_coulomb_mat_b = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + diag_coulomb_mats = np.stack([diag_coulomb_mat_a, diag_coulomb_mat_b]) + constant = rng.standard_normal() + hamiltonian = ffsim.DiagonalCoulombHamiltonian( + one_body_tensor, diag_coulomb_mats, constant=constant + ) + linop = ffsim.linear_operator(hamiltonian, norb=norb, nelec=nelec) + hamiltonian_dense = linop @ np.eye(ffsim.dim(norb, nelec)) + np.testing.assert_allclose( + ffsim.diag(hamiltonian, norb=norb, nelec=nelec), np.diag(hamiltonian_dense) + )