Repository: qiskit-community/ffsim (python)

Issue:
Update MergeOrbitalRotations to merge with Slater determinant and Hartree-Fock gates

Pull request description:
update MergeOrbitalRotations to handle Slater determinants
fixes #154 

Patch:
diff --git a/python/ffsim/qiskit/merge_orbital_rotations.py b/python/ffsim/qiskit/merge_orbital_rotations.py
index e6ce9008..2dfba7ea 100644
--- a/python/ffsim/qiskit/merge_orbital_rotations.py
+++ b/python/ffsim/qiskit/merge_orbital_rotations.py
@@ -17,12 +17,14 @@ from qiskit.dagcircuit import DAGCircuit
 from qiskit.transpiler.basepasses import TransformationPass
 
 from ffsim.qiskit.orbital_rotation import OrbitalRotationJW
+from ffsim.qiskit.slater_determinant import PrepareSlaterDeterminantJW
 
 
 class MergeOrbitalRotations(TransformationPass):
     """Merge consecutive orbital rotation gates."""
 
     def run(self, dag: DAGCircuit) -> DAGCircuit:
+        # Merge orbital rotation gates
         for run in dag.collect_runs(["orb_rot_jw"]):
             node = run[0]
             qubits = node.qargs
@@ -38,4 +40,30 @@ class MergeOrbitalRotations(TransformationPass):
                 {q: i for i, q in enumerate(qubits)},
                 cycle_check=False,
             )
+
+        # Merge Slater determinant preparation followed by orbital rotation
+        # into a single Slater determinant preparation gate
+        for node in dag.named_nodes("slater_jw"):
+            successors = list(dag.successors(node))
+            if len(successors) == 1 and successors[0].op.name == "orb_rot_jw":
+                successor_node = successors[0]
+
+                combined_mat_a = (
+                    successor_node.op.orbital_rotation_a @ node.op.orbital_rotation_a
+                )
+                combined_mat_b = (
+                    successor_node.op.orbital_rotation_b @ node.op.orbital_rotation_b
+                )
+
+                dag.substitute_node(
+                    node,
+                    PrepareSlaterDeterminantJW(
+                        node.op.norb,
+                        node.op.occupied_orbitals,
+                        orbital_rotation=(combined_mat_a, combined_mat_b),
+                    ),
+                    inplace=True,
+                )
+                dag.remove_op_node(successor_node)
+
         return dag
diff --git a/python/ffsim/qiskit/slater_determinant.py b/python/ffsim/qiskit/slater_determinant.py
index a8ae75c6..03be69d7 100644
--- a/python/ffsim/qiskit/slater_determinant.py
+++ b/python/ffsim/qiskit/slater_determinant.py
@@ -8,7 +8,7 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
-"""Slater determinant preparation gate."""
+"""Slater determinant and Hartree-Fock preparation gates."""
 
 from __future__ import annotations
 


Test patch:
diff --git a/tests/python/qiskit/merge_orbital_rotations_test.py b/tests/python/qiskit/merge_orbital_rotations_test.py
index 0a4c8de9..478113c9 100644
--- a/tests/python/qiskit/merge_orbital_rotations_test.py
+++ b/tests/python/qiskit/merge_orbital_rotations_test.py
@@ -8,24 +8,29 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
+from __future__ import annotations
+
 import numpy as np
+import pytest
 from qiskit.circuit import QuantumCircuit, QuantumRegister
-from qiskit.quantum_info import Operator
+from qiskit.quantum_info import Operator, Statevector
 from qiskit.transpiler import PassManager
+from qiskit.transpiler.passes import Decompose
 
 import ffsim
 
 
-def test_yields_equivalent_circuit():
+# TODO handle norb = 0
+@pytest.mark.parametrize("norb", range(1, 5))
+def test_yields_equivalent_circuit(norb: int):
     """Test merging orbital rotations results in an equivalent circuit."""
     rng = np.random.default_rng()
-    norb = 3
     qubits = QuantumRegister(2 * norb)
     circuit = QuantumCircuit(qubits)
     for _ in range(3):
         circuit.append(
             ffsim.qiskit.OrbitalRotationJW(
-                norb, ffsim.random.random_unitary(3, seed=rng)
+                norb, ffsim.random.random_unitary(norb, seed=rng)
             ),
             qubits,
         )
@@ -34,8 +39,8 @@ def test_yields_equivalent_circuit():
             ffsim.qiskit.OrbitalRotationJW(
                 norb,
                 (
-                    ffsim.random.random_unitary(3, seed=rng),
-                    ffsim.random.random_unitary(3, seed=rng),
+                    ffsim.random.random_unitary(norb, seed=rng),
+                    ffsim.random.random_unitary(norb, seed=rng),
                 ),
             ),
             qubits,
@@ -47,3 +52,137 @@ def test_yields_equivalent_circuit():
     np.testing.assert_allclose(
         np.array(Operator(circuit)), np.array(Operator(transpiled))
     )
+
+
+# TODO handle norb = 0
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(1, 5)))
+def test_merge_slater(norb: int, nelec: tuple[int, int]):
+    """Test merging orbital rotations into Slater determinant preparation."""
+    rng = np.random.default_rng()
+
+    qubits = QuantumRegister(2 * norb)
+    circuit = QuantumCircuit(qubits)
+
+    circuit.append(
+        ffsim.qiskit.PrepareSlaterDeterminantJW(
+            norb,
+            ffsim.testing.random_occupied_orbitals(norb, nelec),
+            ffsim.random.random_unitary(norb, seed=rng),
+        ),
+        qubits,
+    )
+    for _ in range(3):
+        circuit.append(
+            ffsim.qiskit.OrbitalRotationJW(
+                norb, ffsim.random.random_unitary(norb, seed=rng)
+            ),
+            qubits,
+        )
+    for _ in range(4):
+        circuit.append(
+            ffsim.qiskit.OrbitalRotationJW(
+                norb,
+                (
+                    ffsim.random.random_unitary(norb, seed=rng),
+                    ffsim.random.random_unitary(norb, seed=rng),
+                ),
+            ),
+            qubits,
+        )
+    pass_manager = PassManager([ffsim.qiskit.MergeOrbitalRotations()])
+    transpiled = pass_manager.run(circuit)
+    assert circuit.count_ops() == {"slater_jw": 1, "orb_rot_jw": 7}
+    assert transpiled.count_ops() == {"slater_jw": 1}
+
+    ffsim.testing.assert_allclose_up_to_global_phase(
+        np.array(Statevector(circuit)), np.array(Statevector(transpiled))
+    )
+
+
+# TODO handle norb = 0
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(1, 5)))
+def test_merge_hartree_fock(norb: int, nelec: tuple[int, int]):
+    """Test merging orbital rotations into Hartree-Fock state preparation."""
+    rng = np.random.default_rng()
+
+    qubits = QuantumRegister(2 * norb)
+    circuit = QuantumCircuit(qubits)
+
+    circuit.append(
+        ffsim.qiskit.PrepareHartreeFockJW(norb, nelec),
+        qubits,
+    )
+    for _ in range(3):
+        circuit.append(
+            ffsim.qiskit.OrbitalRotationJW(
+                norb, ffsim.random.random_unitary(norb, seed=rng)
+            ),
+            qubits,
+        )
+    for _ in range(4):
+        circuit.append(
+            ffsim.qiskit.OrbitalRotationJW(
+                norb,
+                (
+                    ffsim.random.random_unitary(norb, seed=rng),
+                    ffsim.random.random_unitary(norb, seed=rng),
+                ),
+            ),
+            qubits,
+        )
+    pass_manager = PassManager(
+        [Decompose(["hartree_fock_jw", "ucj_jw"]), ffsim.qiskit.MergeOrbitalRotations()]
+    )
+    transpiled = pass_manager.run(circuit)
+    assert circuit.count_ops() == {"hartree_fock_jw": 1, "orb_rot_jw": 7}
+    assert transpiled.count_ops() == {"slater_jw": 1}
+
+    ffsim.testing.assert_allclose_up_to_global_phase(
+        np.array(Statevector(circuit)), np.array(Statevector(transpiled))
+    )
+
+
+# TODO handle norb = 0
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(1, 5)))
+def test_merge_ucj(norb: int, nelec: tuple[int, int]):
+    """Test merging orbital rotations in UCJ operator."""
+    rng = np.random.default_rng()
+
+    qubits = QuantumRegister(2 * norb)
+    circuit = QuantumCircuit(qubits)
+
+    circuit.append(
+        ffsim.qiskit.PrepareHartreeFockJW(norb, nelec),
+        qubits,
+    )
+
+    n_reps = 3
+    ucj_op = ffsim.UCJOperator(
+        diag_coulomb_mats_alpha_alpha=np.stack(
+            [
+                ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+                for _ in range(n_reps)
+            ]
+        ),
+        diag_coulomb_mats_alpha_beta=np.stack(
+            [
+                ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
+                for _ in range(n_reps)
+            ]
+        ),
+        orbital_rotations=np.stack(
+            [ffsim.random.random_unitary(norb, seed=rng) for _ in range(n_reps)]
+        ),
+        final_orbital_rotation=ffsim.random.random_unitary(norb, seed=rng),
+    )
+    circuit.append(ffsim.qiskit.UCJOperatorJW(ucj_op), qubits)
+    pass_manager = PassManager(
+        [Decompose(["hartree_fock_jw", "ucj_jw"]), ffsim.qiskit.MergeOrbitalRotations()]
+    )
+    transpiled = pass_manager.run(circuit)
+    assert circuit.count_ops() == {"hartree_fock_jw": 1, "ucj_jw": 1}
+    assert transpiled.count_ops()["slater_jw"] == 1
+
+    ffsim.testing.assert_allclose_up_to_global_phase(
+        np.array(Statevector(circuit)), np.array(Statevector(transpiled))
+    )