Add double-factorized Trotter simulation Qiskit gate (#275)

* Add double-factorized Trotter simulation Qiskit gate

* mypy

python/ffsim/qiskit/__init__.py

@@ -28,6 +28,7 @@
     PrepareHartreeFockSpinlessJW,
     PrepareSlaterDeterminantJW,
     PrepareSlaterDeterminantSpinlessJW,
+    SimulateTrotterDoubleFactorizedJW,
     UCJOperatorJW,
     UCJOpSpinBalancedJW,
     UCJOpSpinlessJW,
@@ -67,6 +68,7 @@
     "PrepareHartreeFockSpinlessJW",
     "PrepareSlaterDeterminantJW",
     "PrepareSlaterDeterminantSpinlessJW",
+    "SimulateTrotterDoubleFactorizedJW",
     "UCJOperatorJW",
     "UCJOpSpinBalancedJW",
     "UCJOpSpinUnbalancedJW",

python/ffsim/qiskit/gates/__init__.py

@@ -14,6 +14,9 @@
     DiagCoulombEvolutionJW,
     DiagCoulombEvolutionSpinlessJW,
 )
+from ffsim.qiskit.gates.double_factorized_trotter import (
+    SimulateTrotterDoubleFactorizedJW,
+)
 from ffsim.qiskit.gates.givens_ansatz import (
     GivensAnsatzOperatorJW,
     GivensAnsatzOperatorSpinlessJW,
@@ -58,6 +61,7 @@
     "PrepareHartreeFockSpinlessJW",
     "PrepareSlaterDeterminantJW",
     "PrepareSlaterDeterminantSpinlessJW",
+    "SimulateTrotterDoubleFactorizedJW",
     "UCJOperatorJW",
     "UCJOpSpinBalancedJW",
     "UCJOpSpinUnbalancedJW",

python/ffsim/qiskit/gates/double_factorized_trotter.py

@@ -0,0 +1,158 @@
+# (C) Copyright IBM 2024.
+#
+# 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.
+
+"""Double-factorized Trotter evolution gate."""
+
+from __future__ import annotations
+
+from collections.abc import Generator, Iterator, Sequence
+
+import numpy as np
+import scipy.linalg
+from qiskit.circuit import (
+    CircuitInstruction,
+    Gate,
+    QuantumCircuit,
+    QuantumRegister,
+    Qubit,
+)
+
+from ffsim.hamiltonians import DoubleFactorizedHamiltonian
+from ffsim.qiskit.gates.diag_coulomb import DiagCoulombEvolutionJW
+from ffsim.qiskit.gates.num_op_sum import NumOpSumEvolutionJW
+from ffsim.qiskit.gates.orbital_rotation import OrbitalRotationJW
+from ffsim.trotter._util import simulate_trotter_step_iterator
+
+
+class SimulateTrotterDoubleFactorizedJW(Gate):
+    r"""Trotter time evolution of double-factorized Hamiltonian, Jordan-Wigner.
+
+    This gate assumes that qubits are ordered such that the first `norb` qubits
+    correspond to the alpha orbitals and the last `norb` qubits correspond to the
+    beta orbitals.
+    """
+
+    def __init__(
+        self,
+        hamiltonian: DoubleFactorizedHamiltonian,
+        time: float,
+        *,
+        n_steps: int = 1,
+        order: int = 0,
+        label: str | None = None,
+    ):
+        r"""Create double-factorized Trotter evolution gate.
+
+        Args:
+            norb: The number of spatial orbitals.
+            hamiltonian: The Hamiltonian.
+            time: The evolution time.
+            n_steps: The number of Trotter steps.
+            order: The order of the Trotter decomposition.
+            label: The label of the gate.
+        """
+        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}.")
+        self.hamiltonian = hamiltonian
+        self.time = time
+        self.n_steps = n_steps
+        self.order = order
+        super().__init__("df_trotter_jw", 2 * self.hamiltonian.norb, [], label=label)
+
+    def _define(self):
+        """Gate decomposition."""
+        qubits = QuantumRegister(self.num_qubits)
+        self.definition = QuantumCircuit.from_instructions(
+            _simulate_trotter_double_factorized(
+                qubits,
+                self.hamiltonian,
+                time=self.time,
+                n_steps=self.n_steps,
+                order=self.order,
+            ),
+            qubits=qubits,
+        )
+
+
+def _simulate_trotter_double_factorized(
+    qubits: Sequence[Qubit],
+    hamiltonian: DoubleFactorizedHamiltonian,
+    time: float,
+    n_steps: int = 1,
+    order: int = 0,
+) -> Iterator[CircuitInstruction]:
+    if n_steps == 0:
+        return
+
+    one_body_energies, one_body_basis_change = scipy.linalg.eigh(
+        hamiltonian.one_body_tensor
+    )
+    step_time = time / n_steps
+
+    current_basis = np.eye(hamiltonian.norb, dtype=complex)
+    for _ in range(n_steps):
+        current_basis = yield from _simulate_trotter_step_double_factorized(
+            qubits,
+            current_basis,
+            one_body_energies,
+            one_body_basis_change,
+            hamiltonian.diag_coulomb_mats,
+            hamiltonian.orbital_rotations,
+            step_time,
+            norb=hamiltonian.norb,
+            order=order,
+            z_representation=hamiltonian.z_representation,
+        )
+    yield CircuitInstruction(OrbitalRotationJW(hamiltonian.norb, current_basis), qubits)
+
+
+def _simulate_trotter_step_double_factorized(
+    qubits: Sequence[Qubit],
+    current_basis: np.ndarray,
+    one_body_energies: np.ndarray,
+    one_body_basis_change: np.ndarray,
+    diag_coulomb_mats: np.ndarray,
+    orbital_rotations: np.ndarray,
+    time: float,
+    norb: int,
+    order: int,
+    z_representation: bool,
+) -> Generator[CircuitInstruction, None, np.ndarray]:
+    for term_index, time in simulate_trotter_step_iterator(
+        1 + len(diag_coulomb_mats), time, order
+    ):
+        if term_index == 0:
+            yield CircuitInstruction(
+                OrbitalRotationJW(norb, one_body_basis_change.T.conj() @ current_basis),
+                qubits,
+            )
+            yield CircuitInstruction(
+                NumOpSumEvolutionJW(norb, coeffs=one_body_energies, time=time), qubits
+            )
+            current_basis = one_body_basis_change
+        else:
+            orbital_rotation = orbital_rotations[term_index - 1]
+            yield CircuitInstruction(
+                OrbitalRotationJW(norb, orbital_rotation.T.conj() @ current_basis),
+                qubits,
+            )
+            yield CircuitInstruction(
+                DiagCoulombEvolutionJW(
+                    norb,
+                    diag_coulomb_mats[term_index - 1],
+                    time,
+                    z_representation=z_representation,
+                ),
+                qubits,
+            )
+            current_basis = orbital_rotation
+    return current_basis

tests/python/qiskit/gates/double_factorized_trotter_test.py

@@ -0,0 +1,71 @@
+# (C) Copyright IBM 2024.
+#
+# 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 double-factorized Trotter evolution gate."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+from qiskit.quantum_info import Statevector
+
+import ffsim
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, time, n_steps, order, z_representation",
+    [
+        (4, (2, 2), 0.1, 1, 0, False),
+        (4, (2, 2), 0.1, 2, 0, False),
+        (4, (2, 2), 0.1, 1, 1, False),
+        (4, (2, 2), 0.1, 1, 2, False),
+    ],
+)
+def test_random(
+    norb: int,
+    nelec: tuple[int, int],
+    time: float,
+    n_steps: int,
+    order: int,
+    z_representation: bool,
+):
+    """Test random gate gives correct output state."""
+    rng = np.random.default_rng()
+    dim = ffsim.dim(norb, nelec)
+    time = 1.0
+    for _ in range(3):
+        hamiltonian = ffsim.random.random_double_factorized_hamiltonian(
+            norb, rank=norb, z_representation=z_representation, seed=rng
+        )
+        gate = ffsim.qiskit.SimulateTrotterDoubleFactorizedJW(
+            hamiltonian, time, n_steps=n_steps, order=order
+        )
+
+        small_vec = ffsim.random.random_state_vector(dim, seed=rng)
+        big_vec = ffsim.qiskit.ffsim_vec_to_qiskit_vec(
+            small_vec, norb=norb, nelec=nelec
+        )
+
+        statevec = Statevector(big_vec).evolve(gate)
+        result = ffsim.qiskit.qiskit_vec_to_ffsim_vec(
+            np.array(statevec), norb=norb, nelec=nelec
+        )
+
+        expected = ffsim.simulate_trotter_double_factorized(
+            small_vec,
+            hamiltonian,
+            time=time,
+            norb=norb,
+            nelec=nelec,
+            n_steps=n_steps,
+            order=order,
+        )
+
+        np.testing.assert_allclose(result, expected)