add qiskit and orbital rotation circuit

pyproject.toml

@@ -13,7 +13,13 @@ classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
 ]
-dependencies = ["numpy", "opt_einsum", "pyscf>=2.3.0", "scipy"]
+dependencies = [
+    "numpy",
+    "opt_einsum",
+    "pyscf>=2.3.0",
+    "qiskit >= 1.0.0",
+    "scipy",
+]
 
 [project.urls]
 Homepage = "https://github.com/qiskit-community/ffsim"

python/ffsim/__init__.py

@@ -10,7 +10,7 @@
 
 """ffsim is a software library for fast simulation of fermionic quantum circuits."""
 
-from ffsim import contract, linalg, optimize, random, testing
+from ffsim import contract, linalg, optimize, qiskit, random, testing
 from ffsim.cistring import init_cache
 from ffsim.gates import (
     apply_diag_coulomb_evolution,
@@ -129,6 +129,7 @@
     "number_operator",
     "one_hot",
     "optimize",
+    "qiskit",
     "random",
     "rdm",
     "simulate_qdrift_double_factorized",

python/ffsim/qiskit/__init__.py

@@ -0,0 +1,17 @@
+# (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.
+
+"""Code that uses Qiskit, e.g. for constructing quantum circuits."""
+
+from ffsim.qiskit.orbital_rotation import OrbitalRotationJW
+
+__all__ = [
+    "OrbitalRotationJW",
+]

python/ffsim/qiskit/orbital_rotation.py

@@ -0,0 +1,107 @@
+# (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.
+
+"""Orbital rotation gate."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+
+import numpy as np
+from qiskit.circuit import (
+    CircuitInstruction,
+    Gate,
+    QuantumCircuit,
+    QuantumRegister,
+    Qubit,
+)
+from qiskit.circuit.library import PhaseGate, XXPlusYYGate
+
+from ffsim.linalg import givens_decomposition, is_unitary
+from ffsim.spin import Spin
+
+
+class OrbitalRotationJW(Gate):
+    """Orbital rotation under the Jordan-Wigner transformation."""
+
+    def __init__(
+        self,
+        orbital_rotation: np.ndarray,
+        spin: Spin = Spin.ALPHA_AND_BETA,
+        label: str | None = None,
+        validate: bool = True,
+        rtol: float = 1e-5,
+        atol: float = 1e-8,
+    ):
+        """Create new orbital rotation gate.
+
+        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.
+
+        Args:
+            orbital_rotation: The matrix describing the orbital rotation.
+            spin: Choice of spin sector(s) to act on.
+
+                - To act on only spin alpha, pass :const:`ffsim.Spin.ALPHA`.
+                - To act on only spin beta, pass :const:`ffsim.Spin.BETA`.
+                - To act on both spin alpha and spin beta, pass
+                :const:`ffsim.Spin.ALPHA_AND_BETA` (this is the default value).
+            label: The label of the gate.
+            validate: Whether to check that the input matrix is unitary and raise an
+                error if it isn't.
+            rtol: Relative numerical tolerance for input validation.
+            atol: Absolute numerical tolerance for input validation.
+
+        Raises:
+            ValueError: The input matrix is not unitary.
+        """
+        if validate and not is_unitary(orbital_rotation, rtol=rtol, atol=atol):
+            raise ValueError("The input orbital rotation matrix is not unitary.")
+        self.orbital_rotation = orbital_rotation
+        self.spin = spin
+        norb, _ = orbital_rotation.shape
+        super().__init__("orb_rot_jw", 2 * norb, [], label=label)
+
+    def _define(self):
+        """Gate decomposition."""
+        qubits = QuantumRegister(self.num_qubits)
+        circuit = QuantumCircuit(qubits, name=self.name)
+        norb = len(qubits) // 2
+        alpha_qubits = qubits[:norb]
+        beta_qubits = qubits[norb:]
+        if self.spin & Spin.ALPHA:
+            for instruction in _orbital_rotation_jw(
+                alpha_qubits, self.orbital_rotation
+            ):
+                circuit.append(instruction)
+        if self.spin & Spin.BETA:
+            for instruction in _orbital_rotation_jw(beta_qubits, self.orbital_rotation):
+                circuit.append(instruction)
+        self.definition = circuit
+
+    def inverse(self):
+        """Inverse gate."""
+        return OrbitalRotationJW(self.orbital_rotation.T.conj(), spin=self.spin)
+
+
+def _orbital_rotation_jw(
+    qubits: Sequence[Qubit], orbital_rotation: np.ndarray
+) -> Iterator[CircuitInstruction]:
+    givens_rotations, phase_shifts = givens_decomposition(orbital_rotation)
+    for c, s, i, j in givens_rotations:
+        angle = np.arccos(c)
+        phase_angle = np.angle(s)
+        yield CircuitInstruction(
+            XXPlusYYGate(2 * angle, phase_angle - np.pi / 2),
+            (qubits[i], qubits[j]),
+        )
+    for i, phase_shift in enumerate(phase_shifts):
+        yield CircuitInstruction(PhaseGate(np.angle(phase_shift)), (qubits[i],))

python/ffsim/qiskit/util.py

@@ -0,0 +1,59 @@
+# (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.
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+import numpy as np
+
+from ffsim.cistring import make_strings
+from ffsim.states import dim
+
+
+def qiskit_vec_to_ffsim_vec(
+    vec: np.ndarray, norb: int, nelec: tuple[int, int]
+) -> np.ndarray:
+    """Convert a Qiskit statevector to an ffsim statevector.
+
+    Args:
+        vec: A statevector in Qiskit format. It should be a one-dimensional vector
+            of length ``2 ** (2 * norb)``.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+    """
+    assert vec.shape == (1 << (2 * norb),)
+    return vec[_ffsim_indices(norb, nelec)]
+
+
+def ffsim_vec_to_qiskit_vec(
+    vec: np.ndarray, norb: int, nelec: tuple[int, int]
+) -> np.ndarray:
+    """Convert an ffsim statevector to a Qiskit statevector.
+
+    Args:
+        vec: A statevector in ffsim/pySCF format. It should be a one-dimensional vector
+            of length ``comb(norb, n_alpha) * comb(norb, n_beta)``.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+    """
+    assert vec.shape == (dim(norb, nelec),)
+    qiskit_vec = np.zeros(1 << (2 * norb), dtype=vec.dtype)
+    qiskit_vec[_ffsim_indices(norb, nelec)] = vec
+    return qiskit_vec
+
+
+@lru_cache(maxsize=None)
+def _ffsim_indices(norb: int, nelec: tuple[int, int]) -> np.ndarray:
+    n_alpha, n_beta = nelec
+    strings_a = make_strings(range(norb), n_alpha)
+    strings_b = make_strings(range(norb), n_beta) << norb
+    # Compute [a + b for a, b in product(strings_a, strings_b)]
+    return (strings_a.reshape(-1, 1) + strings_b).reshape(-1).copy()

tests/python/qiskit/__init__.py

@@ -0,0 +1,9 @@
+# (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/python/qiskit/orbital_rotation_test.py

@@ -0,0 +1,63 @@
+# (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 orbital rotation circuit."""
+
+import numpy as np
+import pytest
+from qiskit.quantum_info import Statevector
+
+import ffsim
+import ffsim.qiskit
+from ffsim.qiskit.util import ffsim_vec_to_qiskit_vec, qiskit_vec_to_ffsim_vec
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin", ffsim.testing.generate_norb_nelec_spin(range(5))
+)
+def test_random_orbital_rotation(norb: int, nelec: tuple[int, int], spin: ffsim.Spin):
+    """Test random orbital rotation circuit gives correct output state."""
+    rng = np.random.default_rng()
+    dim = ffsim.dim(norb, nelec)
+    for _ in range(3):
+        mat = ffsim.random.random_unitary(norb, seed=rng)
+        gate = ffsim.qiskit.OrbitalRotationJW(mat, spin=spin)
+
+        small_vec = ffsim.random.random_statevector(dim, seed=rng)
+        big_vec = ffsim_vec_to_qiskit_vec(small_vec, norb=norb, nelec=nelec)
+
+        statevec = Statevector(big_vec).evolve(gate)
+        result = qiskit_vec_to_ffsim_vec(np.array(statevec), norb=norb, nelec=nelec)
+
+        expected = ffsim.apply_orbital_rotation(
+            small_vec, mat, norb=norb, nelec=nelec, spin=spin
+        )
+
+        np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize(
+    "norb, nelec, spin", ffsim.testing.generate_norb_nelec_spin(range(5))
+)
+def test_inverse(norb: int, nelec: tuple[int, int], spin: ffsim.Spin):
+    """Test inverse."""
+    rng = np.random.default_rng()
+    dim = ffsim.dim(norb, nelec)
+    for _ in range(3):
+        mat = ffsim.random.random_unitary(norb, seed=rng)
+        gate = ffsim.qiskit.OrbitalRotationJW(mat, spin=spin)
+
+        vec = ffsim_vec_to_qiskit_vec(
+            ffsim.random.random_statevector(dim, seed=rng), norb=norb, nelec=nelec
+        )
+
+        statevec = Statevector(vec).evolve(gate)
+        statevec = statevec.evolve(gate.inverse())
+        np.testing.assert_allclose(np.array(statevec), vec)

tests/python/qiskit/util_test.py

@@ -0,0 +1,31 @@
+# (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 Qiskit utilities."""
+
+import numpy as np
+
+import ffsim
+from ffsim.qiskit.util import ffsim_vec_to_qiskit_vec, qiskit_vec_to_ffsim_vec
+
+
+def test_ffsim_to_qiskit_roundtrip():
+    """Test converting statevector between ffsim and Qiskit gives consistent results."""
+    norb = 5
+    nelec = 3, 2
+    big_dim = 2 ** (2 * norb)
+    small_dim = ffsim.dim(norb, nelec)
+    rng = np.random.default_rng(9940)
+    ffsim_vec = ffsim.random.random_statevector(small_dim, seed=rng)
+    qiskit_vec = ffsim_vec_to_qiskit_vec(ffsim_vec, norb=norb, nelec=nelec)
+    assert qiskit_vec.shape == (big_dim,)
+    ffsim_vec_again = qiskit_vec_to_ffsim_vec(qiskit_vec, norb=norb, nelec=nelec)
+    assert ffsim_vec_again.shape == (small_dim,)
+    np.testing.assert_array_equal(ffsim_vec, ffsim_vec_again)