add single-factorized hamiltonian (#76)

python/ffsim/__init__.py

@@ -24,7 +24,11 @@
     apply_orbital_rotation,
     apply_tunneling_interaction,
 )
-from ffsim.hamiltonians import DoubleFactorizedHamiltonian, MolecularHamiltonian
+from ffsim.hamiltonians import (
+    DoubleFactorizedHamiltonian,
+    MolecularHamiltonian,
+    SingleFactorizedHamiltonian,
+)
 from ffsim.molecular_data import MolecularData
 from ffsim.operators import (
     FermionAction,
@@ -73,6 +77,7 @@
     "HopGateAnsatzOperator",
     "MolecularData",
     "MolecularHamiltonian",
+    "SingleFactorizedHamiltonian",
     "SupportsApplyUnitary",
     "SupportsApproximateEquality",
     "SupportsFermionOperator",

python/ffsim/hamiltonians/__init__.py

@@ -12,8 +12,10 @@
 
 from ffsim.hamiltonians.double_factorized_hamiltonian import DoubleFactorizedHamiltonian
 from ffsim.hamiltonians.molecular_hamiltonian import MolecularHamiltonian
+from ffsim.hamiltonians.single_factorized_hamiltonian import SingleFactorizedHamiltonian
 
 __all__ = [
     "DoubleFactorizedHamiltonian",
     "MolecularHamiltonian",
+    "SingleFactorizedHamiltonian",
 ]

python/ffsim/hamiltonians/double_factorized_hamiltonian.py

@@ -152,28 +152,9 @@ def from_molecular_hamiltonian(
         and constant, and returns a :class:`DoubleFactorizedHamiltonian` storing the
         results.
 
-        The number of terms :math:`t` in the decomposition depends on the allowed
-        error threshold. A larger error threshold leads to a smaller number of terms.
-        Furthermore, the `max_vecs` parameter specifies an optional upper bound
-        on :math:`t`.
-
-        The default behavior of this routine is to perform a straightforward
-        "exact" factorization of the two-body tensor based on a nested
-        eigenvalue decomposition. Additionally, one can choose to optimize the
-        coefficients stored in the tensor to achieve a "compressed" factorization.
-        This option is enabled by setting the `optimize` parameter to `True`.
-        The optimization attempts to minimize a least-squares objective function
-        quantifying the error in the decomposition.
-        It uses `scipy.optimize.minimize`_, passing both the objective function
-        and its gradient. The diagonal coulomb matrices returned by the optimization
-        can be optionally constrained to have only certain elements allowed to be
-        nonzero. This is achieved by passing the `diag_coulomb_mask` parameter, which is
-        an :math:`N \times N` matrix of boolean values where :math:`N` is the number
-        of orbitals. The nonzero elements of this matrix indicate where the diagonal
-        Coulomb matrices are allowed to be nonzero. Only the upper triangular part of
-        the matrix is used because the diagonal Coulomb matrices are symmetric.
-
-        Note: Currently, only real-valued two-body tensors are supported.
+        See :class:`DoubleFactorizedHamiltonian` for a description of the
+        `z_representation` argument. See :func:`ffsim.linalg.double_factorized` for a
+        description of the rest of the arguments.
 
         Args:
             hamiltonian: The Hamiltonian whose double-factorized representation to

python/ffsim/hamiltonians/single_factorized_hamiltonian.py

@@ -0,0 +1,142 @@
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+from __future__ import annotations
+
+import dataclasses
+
+import numpy as np
+from scipy.sparse.linalg import LinearOperator
+
+from ffsim.contract.num_op_sum import num_op_sum_linop
+from ffsim.contract.one_body import one_body_linop
+from ffsim.hamiltonians.molecular_hamiltonian import MolecularHamiltonian
+from ffsim.linalg.double_factorized_decomposition import (
+    _truncated_eigh,
+    modified_cholesky,
+)
+from ffsim.states import dim
+
+
+@dataclasses.dataclass(frozen=True)
+class SingleFactorizedHamiltonian:
+    r"""A Hamiltonian in the single-factorized representation.
+
+    The single-factorized form of the molecular Hamiltonian is
+
+    .. math::
+
+        H = \sum_{\sigma, pq} \kappa_{pq} a^\dagger_{\sigma, p} a_{\sigma, q}
+        + \frac12 \sum_{t=1}^L \left(\mathcal{M}^{(t)}\right)^2
+        + \text{constant}'.
+
+    Here each :math:`\mathcal{M}^{(t)}` is a one-body operator:
+
+    .. math::
+
+        \mathcal{M}^{(t)} =
+        \sum_{\sigma, pq} M^{(t)}_{pq} a^\dagger_{\sigma, p} a_{\sigma, q}
+
+    where each :math:`M^{(t)}` is a Hermitian matrix.
+
+    Attributes:
+        one_body_tensor (np.ndarray): The one-body tensor :math:`\kappa`.
+        one_body_squares (np.ndarray): The one-body tensors :math:`M^{(t)}` whose
+            squares are summed in the Hamiltonian.
+        constant (float): The constant.
+    """
+
+    one_body_tensor: np.ndarray
+    one_body_squares: np.ndarray
+    constant: float = 0.0
+
+    @property
+    def norb(self) -> int:
+        """The number of spatial orbitals."""
+        return self.one_body_tensor.shape[0]
+
+    @staticmethod
+    def from_molecular_hamiltonian(
+        hamiltonian: MolecularHamiltonian,
+        *,
+        tol: float = 1e-8,
+        max_vecs: int | None = None,
+        cholesky: bool = True,
+    ) -> SingleFactorizedHamiltonian:
+        r"""Initialize a SingleFactorizedHamiltonian from a MolecularHamiltonian.
+
+        The number of terms in the decomposition depends on the allowed
+        error threshold. A larger error threshold leads to a smaller number of terms.
+        Furthermore, the `max_vecs` parameter specifies an optional upper bound
+        on the number of terms.
+
+        Note: Currently, only real-valued two-body tensors are supported.
+
+        Args:
+            hamiltonian: The Hamiltonian whose single-factorized representation to
+                compute.
+            tol: Tolerance for error in the decomposition.
+                The error is defined as the maximum absolute difference between
+                an element of the original tensor and the corresponding element of
+                the reconstructed tensor.
+            max_vecs: An optional limit on the number of terms to keep in the
+                decomposition of the two-body tensor. This argument overrides ``tol``.
+            cholesky: Whether to perform the factorization using a modified Cholesky
+                decomposition. If False, a full eigenvalue decomposition is used
+                instead, which can be much more expensive.
+
+        Returns:
+            The single-factorized Hamiltonian.
+        """
+        one_body_tensor = hamiltonian.one_body_tensor - 0.5 * np.einsum(
+            "prqr", hamiltonian.two_body_tensor
+        )
+
+        norb = hamiltonian.norb
+        reshaped_tensor = np.reshape(
+            hamiltonian.two_body_tensor, (norb**2, norb**2)
+        )
+
+        if cholesky:
+            cholesky_vecs = modified_cholesky(
+                reshaped_tensor, tol=tol, max_vecs=max_vecs
+            )
+            one_body_squares = cholesky_vecs.T.reshape((-1, norb, norb))
+        else:
+            eigs, vecs = _truncated_eigh(reshaped_tensor, tol=tol, max_vecs=max_vecs)
+            vecs *= np.sqrt(eigs)
+            one_body_squares = vecs.T.reshape((-1, norb, norb))
+
+        return SingleFactorizedHamiltonian(
+            one_body_tensor=one_body_tensor,
+            one_body_squares=one_body_squares,
+            constant=hamiltonian.constant,
+        )
+
+    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)
+        one_body_square_linops = [
+            0.5 * one_body_linop(one_body, norb, nelec) ** 2
+            for one_body in self.one_body_squares
+        ]
+
+        def matvec(vec: np.ndarray):
+            result = self.constant * vec
+            result += num_linop @ vec
+            for linop in one_body_square_linops:
+                result += linop @ vec
+            return result
+
+        return LinearOperator(
+            shape=(dim_, dim_), matvec=matvec, rmatvec=matvec, dtype=complex
+        )

python/ffsim/linalg/double_factorized_decomposition.py

@@ -106,8 +106,8 @@ def double_factorized(
 
     The number of terms :math:`L` in the decomposition depends on the allowed
     error threshold. A larger error threshold may yield a smaller number of terms.
-    Furthermore, the ``max_vecs`` parameter specifies an optional upper bound
-    on :math:`L`. The ``max_vecs`` parameter is always respected, so if it is
+    Furthermore, the `max_vecs` parameter specifies an optional upper bound
+    on :math:`L`. The `max_vecs` parameter is always respected, so if it is
     too small, then the error of the decomposition may exceed the specified
     error threshold.
 
@@ -119,17 +119,20 @@ def double_factorized(
     The optimization attempts to minimize a least-squares objective function
     quantifying the error in the decomposition.
     It uses `scipy.optimize.minimize`, passing both the objective function
-    and its gradient. The core tensors returned by the optimization can be optionally
-    constrained to have only certain elements allowed to be nonzero. This is achieved by
-    passing the `diag_coulomb_mask` parameter, which is an :math:`N \times N` matrix of
-    boolean values where :math:`N` is the number of orbitals. The nonzero elements of
-    this matrix indicate where the core tensors are allowed to be nonzero. Only the
-    upper triangular part of the matrix is used because the core tensors are symmetric.
+    and its gradient. The diagonal Coulomb matrices returned by the optimization can be
+    optionally constrained to have only certain elements allowed to be nonzero.
+    This is achieved by passing the `diag_coulomb_mask` parameter, which is an
+    :math:`N \times N` matrix of boolean values where :math:`N` is the number of
+    orbitals. The nonzero elements of this matrix indicate where the diagonal Coulomb
+    matrices are allowed to be nonzero. Only the upper triangular part of the matrix is
+    used because the diagonal Coulomb matrices are symmetric.
 
     References:
         - `arXiv:1808.02625`_
         - `arXiv:2104.08957`_
 
+    Note: Currently, only real-valued two-body tensors are supported.
+
     Args:
         two_body_tensor: The two-body tensor to decompose.
         tol: Tolerance for error in the decomposition.
@@ -461,8 +464,8 @@ def double_factorized_t2(
 
     The number of terms :math:`L` in the decomposition depends on the allowed
     error threshold. A larger error threshold may yield a smaller number of terms.
-    Furthermore, the ``max_vecs`` parameter specifies an optional upper bound
-    on :math:`L`. The ``max_vecs`` parameter is always respected, so if it is
+    Furthermore, the `max_vecs` parameter specifies an optional upper bound
+    on :math:`L`. The `max_vecs` parameter is always respected, so if it is
     too small, then the error of the decomposition may exceed the specified
     error threshold.
 

tests/hamiltonians/single_factorized_hamiltonian_test.py

@@ -0,0 +1,59 @@
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""Tests for single-factorized Hamiltonian."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import ffsim
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, cholesky",
+    [
+        (4, (2, 2), False),
+        (4, (2, 2), True),
+        (4, (2, 1), False),
+        (4, (2, 1), True),
+        (4, (2, 0), False),
+        (4, (2, 0), True),
+    ],
+)
+def test_linear_operator(norb: int, nelec: tuple[int, int], cholesky: bool):
+    """Test linear operator."""
+    norb = 4
+    nelec = (2, 2)
+    rng = np.random.default_rng(2474)
+
+    dim = ffsim.dim(norb, 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
+    )
+
+    df_hamiltonian = ffsim.SingleFactorizedHamiltonian.from_molecular_hamiltonian(
+        mol_hamiltonian, cholesky=cholesky
+    )
+
+    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=rng)
+
+    actual = actual_linop @ state
+    expected = expected_linop @ state
+
+    np.testing.assert_allclose(actual, expected, atol=1e-8)