add diagonal coulomb trotter (#267)

python/ffsim/__init__.py

@@ -83,6 +83,7 @@
 )
 from ffsim.trotter import (
     simulate_qdrift_double_factorized,
+    simulate_trotter_diag_coulomb_split_op,
     simulate_trotter_double_factorized,
 )
 from ffsim.variational import (
@@ -170,6 +171,7 @@
     "rdms",
     "sample_state_vector",
     "simulate_qdrift_double_factorized",
+    "simulate_trotter_diag_coulomb_split_op",
     "simulate_trotter_double_factorized",
     "slater_determinant",
     "slater_determinant_rdm",

python/ffsim/trotter/__init__.py

@@ -10,10 +10,12 @@
 
 """Hamiltonian simulation via Trotter-Suzuki formulas."""
 
+from ffsim.trotter.diagonal_coulomb import simulate_trotter_diag_coulomb_split_op
 from ffsim.trotter.double_factorized import simulate_trotter_double_factorized
 from ffsim.trotter.qdrift import simulate_qdrift_double_factorized
 
 __all__ = [
     "simulate_qdrift_double_factorized",
+    "simulate_trotter_diag_coulomb_split_op",
     "simulate_trotter_double_factorized",
 ]

python/ffsim/trotter/diagonal_coulomb.py

@@ -0,0 +1,140 @@
+# (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.
+
+"""Trotter simulation for diagonal Coulomb Hamiltonian."""
+
+from __future__ import annotations
+
+import numpy as np
+import scipy.linalg
+
+from ffsim.gates import (
+    apply_diag_coulomb_evolution,
+    apply_num_op_sum_evolution,
+    apply_orbital_rotation,
+)
+from ffsim.hamiltonians import DiagonalCoulombHamiltonian
+from ffsim.trotter._util import simulate_trotter_step_iterator
+
+
+def simulate_trotter_diag_coulomb_split_op(
+    vec: np.ndarray,
+    hamiltonian: DiagonalCoulombHamiltonian,
+    time: float,
+    *,
+    norb: int,
+    nelec: tuple[int, int],
+    n_steps: int = 1,
+    order: int = 0,
+    copy: bool = True,
+) -> np.ndarray:
+    """Diagonal Coulomb Hamiltonian simulation using split-operator method.
+
+    Args:
+        vec: The state vector to evolve.
+        hamiltonian: The Hamiltonian.
+        time: The evolution time.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+        n_steps: The number of Trotter steps.
+        order: The order of the Trotter decomposition.
+        copy: Whether to copy the vector before operating on it.
+
+            - If `copy=True` then this function always returns a newly allocated
+              vector and the original vector is left untouched.
+            - If `copy=False` then this function may still return a newly allocated
+              vector, but the original vector may have its data overwritten.
+              It is also possible that the original vector is returned,
+              modified in-place.
+
+    Returns:
+        The final state of the simulation.
+    """
+    if order < 0:
+        raise ValueError(f"order must be non-negative, got {order}.")
+    if n_steps < 0:
+        raise ValueError(f"n_steps must be non-negative, got {n_steps}.")
+    if copy:
+        vec = vec.copy()
+    if n_steps == 0:
+        return vec
+
+    one_body_energies, one_body_basis_change = scipy.linalg.eigh(
+        hamiltonian.one_body_tensor
+    )
+    step_time = time / n_steps
+
+    current_basis = np.eye(norb, dtype=complex)
+    for _ in range(n_steps):
+        vec, current_basis = _simulate_trotter_step_diag_coulomb_split_op(
+            vec,
+            current_basis,
+            one_body_energies,
+            one_body_basis_change,
+            hamiltonian.diag_coulomb_mats,
+            step_time,
+            norb=norb,
+            nelec=nelec,
+            order=order,
+        )
+    vec = apply_orbital_rotation(vec, current_basis, norb=norb, nelec=nelec, copy=False)
+
+    return vec
+
+
+def _simulate_trotter_step_diag_coulomb_split_op(
+    vec: np.ndarray,
+    current_basis: np.ndarray,
+    one_body_energies: np.ndarray,
+    one_body_basis_change: np.ndarray,
+    diag_coulomb_mats: np.ndarray,
+    time: float,
+    norb: int,
+    nelec: tuple[int, int],
+    order: int,
+) -> tuple[np.ndarray, np.ndarray]:
+    diag_coulomb_aa, diag_coulomb_ab = diag_coulomb_mats
+    eye = np.eye(norb)
+    for term_index, time in simulate_trotter_step_iterator(2, time, order):
+        if term_index == 0:
+            vec = apply_orbital_rotation(
+                vec,
+                one_body_basis_change.T.conj() @ current_basis,
+                norb=norb,
+                nelec=nelec,
+                copy=False,
+            )
+            vec = apply_num_op_sum_evolution(
+                vec,
+                one_body_energies,
+                time,
+                norb=norb,
+                nelec=nelec,
+                copy=False,
+            )
+            current_basis = one_body_basis_change
+        else:
+            vec = apply_orbital_rotation(
+                vec,
+                current_basis,
+                norb=norb,
+                nelec=nelec,
+                copy=False,
+            )
+            vec = apply_diag_coulomb_evolution(
+                vec,
+                (diag_coulomb_aa, diag_coulomb_ab, diag_coulomb_aa),
+                time,
+                norb=norb,
+                nelec=nelec,
+                copy=False,
+            )
+            current_basis = eye
+    return vec, current_basis

tests/python/trotter/diag_coulomb_test.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.
+
+"""Tests for diagonal Coulomb Trotter simulation."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import scipy.sparse.linalg
+
+import ffsim
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, time, n_steps, order, target_fidelity",
+    [
+        (3, (1, 1), 0.3, 20, 0, 0.999),
+        (4, (2, 1), 0.3, 10, 2, 0.999),
+        (4, (2, 2), 0.3, 10, 1, 0.999),
+    ],
+)
+def test_random(
+    norb: int,
+    nelec: tuple[int, int],
+    time: float,
+    n_steps: int,
+    order: int,
+    target_fidelity: float,
+):
+    rng = np.random.default_rng(2488)
+    # generate random Hamiltonian
+    dim = ffsim.dim(norb, nelec)
+    one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
+    diag_coulomb_mat_aa = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mat_ab = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+    diag_coulomb_mats = np.stack([diag_coulomb_mat_aa, diag_coulomb_mat_ab])
+    dc_hamiltonian = ffsim.DiagonalCoulombHamiltonian(
+        one_body_tensor, diag_coulomb_mats
+    )
+    linop = ffsim.linear_operator(dc_hamiltonian, norb=norb, nelec=nelec)
+
+    # generate initial state
+    dim = ffsim.dim(norb, nelec)
+    initial_state = ffsim.random.random_state_vector(dim, seed=rng)
+    original_state = initial_state.copy()
+
+    # compute exact state
+    exact_state = scipy.sparse.linalg.expm_multiply(
+        -1j * time * linop,
+        initial_state,
+        traceA=np.sum(np.abs(diag_coulomb_mats)),
+    )
+
+    # make sure time is not too small
+    assert abs(np.vdot(exact_state, initial_state)) < 0.98
+
+    # simulate
+    final_state = ffsim.simulate_trotter_diag_coulomb_split_op(
+        initial_state,
+        dc_hamiltonian,
+        time,
+        norb=norb,
+        nelec=nelec,
+        n_steps=n_steps,
+        order=order,
+    )
+
+    # check that initial state was not modified
+    np.testing.assert_allclose(initial_state, original_state)
+
+    # check agreement
+    np.testing.assert_allclose(np.linalg.norm(final_state), 1.0)
+    fidelity = np.abs(np.vdot(final_state, exact_state))
+    assert fidelity >= target_fidelity
+
+
+def test_hubbard():
+    rng = np.random.default_rng(2488)
+
+    hubbard_model = ffsim.fermi_hubbard_2d(
+        norb_x=3,
+        norb_y=3,
+        tunneling=1.0,
+        interaction=4.0,
+        chemical_potential=0.5,
+        nearest_neighbor_interaction=2.0,
+        periodic=False,
+    )
+    dc_hamiltonian = ffsim.DiagonalCoulombHamiltonian.from_fermion_operator(
+        hubbard_model
+    )
+    norb = dc_hamiltonian.norb
+    nelec = (norb // 2, norb // 2)
+    time = 0.1
+    n_steps = 1
+    order = 1
+    target_fidelity = 0.999
+
+    dim = ffsim.dim(norb, nelec)
+    linop = ffsim.linear_operator(dc_hamiltonian, norb=norb, nelec=nelec)
+
+    # generate initial state
+    dim = ffsim.dim(norb, nelec)
+    initial_state = ffsim.random.random_state_vector(dim, seed=rng)
+    original_state = initial_state.copy()
+
+    # compute exact state
+    exact_state = scipy.sparse.linalg.expm_multiply(
+        -1j * time * linop,
+        initial_state,
+        traceA=np.sum(np.abs(dc_hamiltonian.diag_coulomb_mats)),
+    )
+
+    # make sure time is not too small
+    assert abs(np.vdot(exact_state, initial_state)) < 0.98
+
+    # simulate
+    final_state = ffsim.simulate_trotter_diag_coulomb_split_op(
+        initial_state,
+        dc_hamiltonian,
+        time,
+        norb=norb,
+        nelec=nelec,
+        n_steps=n_steps,
+        order=order,
+    )
+
+    # check that initial state was not modified
+    np.testing.assert_allclose(initial_state, original_state)
+
+    # check agreement
+    np.testing.assert_allclose(np.linalg.norm(final_state), 1.0)
+    fidelity = np.abs(np.vdot(final_state, exact_state))
+    assert fidelity >= target_fidelity