add diag coulomb hamiltonian diag

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

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)

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)
+    )