QAOA Transpilation capabilities (#34)

qopt_best_practices/swap_strategies/build_circuit.py

@@ -55,7 +55,7 @@ def apply_swap_strategy(
     return pm_pre.run(circuit)
 
 
-def apply_qaoa_layers(  # pylint: disable=too-many-arguments,too-many-locals
+def apply_qaoa_layers(  # pylint: disable=too-many-arguments,too-many-locals,too-many-positional-arguments
     cost_layer: QuantumCircuit,
     meas_map: dict,
     num_layers: int,
@@ -116,7 +116,7 @@ def apply_qaoa_layers(  # pylint: disable=too-many-arguments,too-many-locals
     return new_circuit
 
 
-def create_qaoa_swap_circuit(  # pylint: disable=too-many-arguments
+def create_qaoa_swap_circuit(  # pylint: disable=too-many-arguments,too-many-positional-arguments
     cost_operator: SparsePauliOp,
     swap_strategy: SwapStrategy,
     edge_coloring: dict = None,

qopt_best_practices/transpilation/__init__.py

@@ -0,0 +1,3 @@
+"""Module with transpiler methods for QAOA like circuits."""
+
+from .preset_qaoa_passmanager import qaoa_swap_strategy_pm

qopt_best_practices/transpilation/preset_qaoa_passmanager.py

@@ -0,0 +1,48 @@
+"""Make a pass manager to transpile QAOA."""
+
+from typing import Any, Dict
+
+from qiskit.transpiler import PassManager
+from qiskit.transpiler.passes import HighLevelSynthesis, InverseCancellation
+from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import (
+    FindCommutingPauliEvolutions,
+    Commuting2qGateRouter,
+)
+from qiskit.circuit.library import CXGate
+
+from qopt_best_practices.transpilation.qaoa_construction_pass import QAOAConstructionPass
+
+
+def qaoa_swap_strategy_pm(config: Dict[str, Any]):
+    """Provide a pass manager to build the QAOA cirucit.
+
+    This function will be extended in the future.
+    """
+
+    num_layers = config.get("num_layers", 1)
+    swap_strategy = config.get("swap_strategy", None)
+    edge_coloring = config.get("edge_coloring", None)
+    basis_gates = config.get("basis_gates", ["sx", "x", "rz", "cx", "id"])
+
+    if swap_strategy is None:
+        raise ValueError("No swap_strategy provided in config.")
+
+    if edge_coloring is None:
+        raise ValueError("No edge_coloring provided in config.")
+
+    # 2. define pass manager for cost layer
+    qaoa_pm = PassManager(
+        [
+            HighLevelSynthesis(basis_gates=["PauliEvolution"]),
+            FindCommutingPauliEvolutions(),
+            Commuting2qGateRouter(
+                swap_strategy,
+                edge_coloring,
+            ),
+            HighLevelSynthesis(basis_gates=basis_gates),
+            InverseCancellation(gates_to_cancel=[CXGate()]),
+            QAOAConstructionPass(num_layers),
+        ]
+    )
+
+    return qaoa_pm

qopt_best_practices/transpilation/qaoa_construction_pass.py

@@ -0,0 +1,126 @@
+"""A pass to build a full QAOA ansatz circuit."""
+
+from typing import Optional
+
+from qiskit.converters import circuit_to_dag, dag_to_circuit
+from qiskit.circuit import QuantumCircuit, ParameterVector, Parameter
+from qiskit.dagcircuit import DAGCircuit
+from qiskit.transpiler import TranspilerError
+from qiskit.transpiler.basepasses import TransformationPass
+
+
+class QAOAConstructionPass(TransformationPass):
+    """Build the QAOAAnsatz from a transpiled cost operator.
+
+    This pass takes as input a single layer of a transpiled QAOA operator.
+    It then repeats this layer the appropriate number of time and adds (i)
+    the initial state, (ii) the mixer operator, and (iii) the measurements.
+    The measurements are added in such a fashion so as to undo the local
+    permuttion induced by the SWAP gates in the cost layer.
+    """
+
+    def __init__(
+        self,
+        num_layers: int,
+        init_state: Optional[QuantumCircuit] = None,
+        mixer_layer: Optional[QuantumCircuit] = None,
+    ):
+        """Initialize the pass
+
+        Limitations: The current implementation of the pass does not permute the mixer.
+        Therefore mixers with local bias fields, such as in warm-start methods, will not
+        result in the correct circuits.
+
+        Args:
+            num_layers: The number of QAOA layers to apply.
+            init_state: The initial state to use. This must match the anticipated number
+                of qubits otherwise an error will be raised when `run` is called. If this
+                is not given we will default to the equal superposition initial state.
+            mixer_layer: The mixer layer to use. This must match the anticipated number
+                of qubits otherwise an error will be raised when `run` is called. If this
+                is not given we default to the standard mixer made of X gates.
+        """
+        super().__init__()
+
+        self.num_layers = num_layers
+        self.init_state = init_state
+        self.mixer_layer = mixer_layer
+
+    def run(self, cost_layer_dag: DAGCircuit):
+        num_qubits = cost_layer_dag.num_qubits()
+
+        # Make the initial state and the mixer.
+        if self.init_state is None:
+            init_state = QuantumCircuit(num_qubits)
+            init_state.h(range(num_qubits))
+        else:
+            init_state = self.init_state
+
+        if self.mixer_layer is None:
+            mixer_layer = QuantumCircuit(num_qubits)
+            beta = Parameter("β")
+            mixer_layer.rx(2 * beta, range(num_qubits))
+        else:
+            mixer_layer = self.mixer_layer
+
+        # Do some sanity checks on qubit numbers.
+        if init_state.num_qubits != num_qubits:
+            raise TranspilerError(
+                "Number of qubits in the initial state does not match the number in the DAG. "
+                f"{init_state.num_qubits} != {num_qubits}"
+            )
+
+        if mixer_layer.num_qubits != num_qubits:
+            raise TranspilerError(
+                "Number of qubits in the mixer does not match the number in the DAG. "
+                f"{init_state.num_qubits} != {num_qubits}"
+            )
+
+        # Note: converting to circuit is inefficent. Update to DAG only.
+        cost_layer = dag_to_circuit(cost_layer_dag)
+        qaoa_circuit = QuantumCircuit(num_qubits, num_qubits)
+
+        # Re-parametrize the circuit
+        gammas = ParameterVector("γ", self.num_layers)
+        betas = ParameterVector("β", self.num_layers)
+
+        # Add initial state
+        qaoa_circuit.compose(init_state, inplace=True)
+
+        # iterate over number of qaoa layers
+        # and alternate cost/reversed cost and mixer
+        for layer in range(self.num_layers):
+            bind_dict = {cost_layer.parameters[0]: gammas[layer]}
+            bound_cost_layer = cost_layer.assign_parameters(bind_dict)
+
+            bind_dict = {mixer_layer.parameters[0]: betas[layer]}
+            bound_mixer_layer = mixer_layer.assign_parameters(bind_dict)
+
+            if layer % 2 == 0:
+                # even layer -> append cost
+                qaoa_circuit.compose(bound_cost_layer, range(num_qubits), inplace=True)
+            else:
+                # odd layer -> append reversed cost
+                qaoa_circuit.compose(
+                    bound_cost_layer.reverse_ops(), range(num_qubits), inplace=True
+                )
+
+            # the mixer layer is not reversed and not permuted.
+            qaoa_circuit.compose(bound_mixer_layer, range(num_qubits), inplace=True)
+
+        if self.num_layers % 2 == 1:
+            # iterate over layout permutations to recover measurements
+            if self.property_set["virtual_permutation_layout"]:
+                for cidx, qidx in (
+                    self.property_set["virtual_permutation_layout"].get_physical_bits().items()
+                ):
+                    qaoa_circuit.measure(qidx, cidx)
+            else:
+                print("layout not found, assigining trivial layout")
+                for idx in range(num_qubits):
+                    qaoa_circuit.measure(idx, idx)
+        else:
+            for idx in range(num_qubits):
+                qaoa_circuit.measure(idx, idx)
+
+        return circuit_to_dag(qaoa_circuit)

test/test_qaoa_construction.py

@@ -0,0 +1,97 @@
+"""Test the construction of the QAOA ansatz."""
+
+
+from unittest import TestCase
+
+from qiskit import QuantumCircuit, transpile
+from qiskit.circuit import Parameter
+from qiskit.circuit.library import QAOAAnsatz
+from qiskit.primitives import StatevectorEstimator
+from qiskit.quantum_info import SparsePauliOp
+from qiskit.transpiler import PassManager
+from qiskit.transpiler.passes.routing.commuting_2q_gate_routing import SwapStrategy
+
+from qopt_best_practices.transpilation.qaoa_construction_pass import QAOAConstructionPass
+from qopt_best_practices.transpilation.preset_qaoa_passmanager import qaoa_swap_strategy_pm
+
+
+class TestQAOAConstruction(TestCase):
+    """Test the construction of the QAOA ansatz."""
+
+    def setUp(self):
+        """Set up re-used variables."""
+        self.estimator = StatevectorEstimator()
+
+        gamma = Parameter("γ")
+        cost_op = QuantumCircuit(4)
+        cost_op.rzz(2 * gamma, 0, 1)
+        cost_op.rzz(2 * gamma, 2, 3)
+        cost_op.swap(0, 1)
+        cost_op.swap(2, 3)
+        cost_op.rzz(2 * gamma, 1, 2)
+
+        self.cost_op_circ = transpile(cost_op, basis_gates=["sx", "cx", "x", "rz"])
+
+        self.cost_op = SparsePauliOp.from_list([("IIZZ", 1), ("ZZII", 1), ("ZIIZ", 1)])
+
+        self.config = {
+            "swap_strategy": SwapStrategy.from_line(list(range(4))),
+            "edge_coloring": {(idx, idx + 1): (idx + 1) % 2 for idx in range(4)},
+        }
+
+    def test_depth_one(self):
+        """Compare the pass with the SWAPs and ensure the measurements are ordered properly."""
+        qaoa_pm = qaoa_swap_strategy_pm(self.config)
+
+        cost_op_circ = QAOAAnsatz(
+            self.cost_op, initial_state=QuantumCircuit(4), mixer_operator=QuantumCircuit(4)
+        ).decompose(reps=1)
+
+        ansatz = qaoa_pm.run(cost_op_circ)
+
+        # 1. Check the measurement map
+        qreg = ansatz.qregs[0]
+        creg = ansatz.cregs[0]
+
+        expected_meas_map = {0: 1, 1: 0, 2: 3, 3: 2}
+
+        for inst in ansatz.data:
+            if inst.operation.name == "measure":
+                qubit = qreg.index(inst.qubits[0])
+                cbit = creg.index(inst.clbits[0])
+                self.assertEqual(cbit, expected_meas_map[qubit])
+
+        # 2. Check the expectation value. Note that to use the estimator we need to
+        # Remove the final measurements and correspondingly permute the cost op.
+        ansatz.remove_final_measurements(inplace=True)
+        permuted_cost_op = SparsePauliOp.from_list([("IIZZ", 1), ("ZZII", 1), ("IZZI", 1)])
+        value = self.estimator.run([(ansatz, permuted_cost_op, [1, 2])]).result()[0].data.evs
+
+        library_ansatz = QAOAAnsatz(self.cost_op, reps=1)
+        library_ansatz = transpile(library_ansatz, basis_gates=["cx", "rz", "rx", "h"])
+
+        expected = self.estimator.run([(library_ansatz, self.cost_op, [1, 2])]).result()[0].data.evs
+
+        self.assertAlmostEqual(value, expected)
+
+    def test_depth_two_qaoa_pass(self):
+        """Compare the pass with the SWAPs to an all-to-all construction.
+
+        Note: this test only works as is because p is even and we don't have the previous
+        passes to give us the qubit permutations.
+        """
+        qaoa_pm = PassManager([QAOAConstructionPass(num_layers=2)])
+
+        ansatz = qaoa_pm.run(self.cost_op_circ)
+        ansatz.remove_final_measurements(inplace=True)
+
+        value = self.estimator.run([(ansatz, self.cost_op, [1, 2, 3, 4])]).result()[0].data.evs
+
+        library_ansatz = QAOAAnsatz(self.cost_op, reps=2)
+        library_ansatz = transpile(library_ansatz, basis_gates=["cx", "rz", "rx", "h"])
+
+        expected = (
+            self.estimator.run([(library_ansatz, self.cost_op, [1, 2, 3, 4])]).result()[0].data.evs
+        )
+
+        self.assertAlmostEqual(value, expected)

test/test_qubit_selection.py

@@ -49,7 +49,7 @@ def test_qubit_selection(self):
         path_finder = BackendEvaluator(self.backend)
         path, _, _ = path_finder.evaluate(len(self.mapped_graph))
 
-        expected_path = [45, 46, 47, 48, 49, 55, 68, 69, 70, 74]
+        expected_path = [33, 39, 40, 72, 41, 81, 53, 60, 61, 62]
         self.assertEqual(set(path), set(expected_path))
 
     def test_qubit_selection_v1_v2(self):
@@ -59,5 +59,5 @@ def test_qubit_selection_v1_v2(self):
         for backend in backends:
             path_finder = BackendEvaluator(backend)
             path, _, _ = path_finder.evaluate(len(self.mapped_graph))
-            expected_path = [1, 2, 4, 7, 8, 10, 11, 12, 13, 14]
+            expected_path = [8, 9, 11, 12, 13, 14, 15, 18, 21, 23]
             self.assertEqual(set(path), set(expected_path))