Merge pull request #1084 from qiboteam/sabre_mesurements

Sabre with measurements

src/qibo/gates/measurements.py

@@ -211,3 +211,37 @@ def load(cls, payload):
         args["basis"] = [getattr(gates, g) for g in args["basis"]]
         args.update(args.pop("init_kwargs"))
         return cls(*qubits, **args)
+
+    # Overload on_qubits to copy also gate.result, controlled by can be removed for measurements
+    def on_qubits(self, qubit_map) -> "Gate":
+        """Creates the same measurement gate targeting different qubits
+        and preserving the measurement result register.
+
+        Args:
+            qubit_map (int): Dictionary mapping original qubit indices to new ones.
+
+        Returns:
+            A :class:`qibo.gates.Gate.M` object of the original gate
+            type targeting the given qubits.
+
+        Example:
+
+            .. testcode::
+
+                from qibo import models, gates
+                measurement = gates.M(0, 1)
+                c = models.Circuit(3)
+                c.add(measurement.on_qubits({0: 0, 1: 2}))
+                assert c.queue[0].result is measurement.result
+                print(c.draw())
+            .. testoutput::
+
+                q0: ─M─
+                q1: ─|─
+                q2: ─M─
+        """
+
+        qubits = (qubit_map.get(q) for q in self.qubits)
+        gate = self.__class__(*qubits, **self.init_kwargs)
+        gate.result = self.result
+        return gate

src/qibo/transpiler/blocks.py

@@ -78,7 +78,6 @@ def on_qubits(self, new_qubits: tuple):
         """
         qubits_dict = dict(zip(self.qubits, new_qubits))
         new_gates = [gate.on_qubits(qubits_dict) for gate in self.gates]
-
         return Block(qubits=new_qubits, gates=new_gates, name=self.name)
 
     # TODO: use real QM properties to check commutation
@@ -102,7 +101,7 @@ def kak_decompose(self):  # pragma: no cover
         This should be done only if the block is entangled and the number of
         two qubit gates is higher than the number after the decomposition.
         """
-        raise_error(NotImplementedError, "")
+        raise_error(NotImplementedError, "KAK decomposition is not available yet.")
 
 
 class CircuitBlocks:
@@ -145,9 +144,16 @@ def add_block(self, block: "Block"):
             )
         self.block_list.append(block)
 
-    def circuit(self):
-        """Return the quantum circuit."""
-        circuit = Circuit(self.qubits)
+    def circuit(self, circuit_kwargs=None):
+        """Return the quantum circuit.
+
+        Args:
+            circuit_kwargs (dict): original circuit init_kwargs.
+        """
+        if circuit_kwargs is None:
+            circuit = Circuit(self.qubits)
+        else:
+            circuit = Circuit(**circuit_kwargs)
         for block in self.block_list:
             for gate in block.gates:
                 circuit.add(gate)

src/qibo/transpiler/router.py

@@ -351,15 +351,24 @@ class CircuitMap:
     Args:
         initial_layout (dict): initial logical-to-physical qubit mapping.
         circuit (Circuit): circuit to be routed.
+        blocks (CircuitBlocks): circuit blocks representation, if None the blocks will be computed from the circuit.
     """
 
-    def __init__(self, initial_layout: dict, circuit: Circuit):
-        self.circuit_blocks = CircuitBlocks(circuit, index_names=True)
+    def __init__(self, initial_layout: dict, circuit: Circuit, blocks=None):
+        if blocks is not None:
+            self.circuit_blocks = blocks
+        else:
+            self.circuit_blocks = CircuitBlocks(circuit, index_names=True)
+        self.initial_layout = initial_layout
         self._circuit_logical = list(range(len(initial_layout)))
         self._physical_logical = list(initial_layout.values())
         self._routed_blocks = CircuitBlocks(Circuit(circuit.nqubits))
         self._swaps = 0
 
+    def set_circuit_logical(self, circuit_logical_map: list):
+        """Set the current circuit to logical qubit mapping."""
+        self._circuit_logical = circuit_logical_map
+
     def blocks_qubits_pairs(self):
         """Returns a list containing the qubit pairs of each block."""
         return [block.qubits for block in self.circuit_blocks()]
@@ -371,9 +380,13 @@ def execute_block(self, block: Block):
         self._routed_blocks.add_block(block.on_qubits(self.get_physical_qubits(block)))
         self.circuit_blocks.remove_block(block)
 
-    def routed_circuit(self):
-        """Returns circuit of the routed circuit."""
-        return self._routed_blocks.circuit()
+    def routed_circuit(self, circuit_kwargs=None):
+        """Return the routed circuit.
+
+        Args:
+            circuit_kwargs (dict): original circuit init_kwargs.
+        """
+        return self._routed_blocks.circuit(circuit_kwargs=circuit_kwargs)
 
     def final_layout(self):
         """Returns the final physical-circuit qubits mapping."""
@@ -455,6 +468,7 @@ def __init__(
         self._front_layer = None
         self.circuit = None
         self._memory_map = None
+        self._final_measurements = None
         random.seed(seed)
 
     def __call__(self, circuit: Circuit, initial_layout: dict):
@@ -475,7 +489,13 @@ def __call__(self, circuit: Circuit, initial_layout: dict):
             else:
                 self._find_new_mapping()
 
-        return self.circuit.routed_circuit(), self.circuit.final_layout()
+        routed_circuit = self.circuit.routed_circuit(circuit_kwargs=circuit.init_kwargs)
+        if self._final_measurements is not None:
+            routed_circuit = self._append_final_measurements(
+                routed_circuit=routed_circuit
+            )
+
+        return routed_circuit, self.circuit.final_layout()
 
     @property
     def added_swaps(self):
@@ -485,22 +505,60 @@ def added_swaps(self):
     def _preprocessing(self, circuit: Circuit, initial_layout: dict):
         """The following objects will be initialised:
         - circuit: class to represent circuit and to perform logical-physical qubit mapping.
+        - _final_measurements: measurement gates at the end of the circuit.
         - _dist_matrix: matrix reporting the shortest path lengh between all node pairs.
         - _dag: direct acyclic graph of the circuit based on commutativity.
         - _memory_map: list to remember previous SWAP moves.
         - _front_layer: list containing the blocks to be executed.
         - _delta_register: list containing the special weigh added to qubits
             to prevent overlapping swaps.
         """
-        self.circuit = CircuitMap(initial_layout, circuit)
+        copy_circuit = self._copy_circuit(circuit)
+        self._final_measurements = self._detach_final_measurements(copy_circuit)
+        self.circuit = CircuitMap(initial_layout, copy_circuit)
         self._dist_matrix = nx.floyd_warshall_numpy(self.connectivity)
         self._dag = _create_dag(self.circuit.blocks_qubits_pairs())
         self._memory_map = []
         self._update_dag_layers()
         self._update_front_layer()
         self._delta_register = [1.0 for _ in range(circuit.nqubits)]
 
+    @staticmethod
+    def _copy_circuit(circuit: Circuit):
+        """Return a copy of the circuit to avoid altering the original circuit.
+        This copy conserves the registers of the measurement gates."""
+        new_circuit = Circuit(circuit.nqubits)
+        for gate in circuit.queue:
+            new_circuit.add(gate)
+        return new_circuit
+
+    def _detach_final_measurements(self, circuit: Circuit):
+        """Detach measurement gates at the end of the circuit for separate handling."""
+        final_measurements = []
+        for gate in circuit.queue[::-1]:
+            if isinstance(gate, gates.M):
+                final_measurements.append(gate)
+                circuit.queue.remove(gate)
+            else:
+                break
+        if not final_measurements:
+            return None
+        return final_measurements[::-1]
+
+    def _append_final_measurements(self, routed_circuit: Circuit):
+        """Append the final measurment gates on the correct qubits conserving the measurement register."""
+        for measurement in self._final_measurements:
+            original_qubits = measurement.qubits
+            routed_qubits = (
+                self.circuit.circuit_to_physical(qubit) for qubit in original_qubits
+            )
+            routed_circuit.add(
+                measurement.on_qubits(dict(zip(original_qubits, routed_qubits)))
+            )
+        return routed_circuit
+
     def _update_dag_layers(self):
+        """Update dag layers and put them in topological order."""
         for layer, nodes in enumerate(nx.topological_generations(self._dag)):
             for node in nodes:
                 self._dag.nodes[node]["layer"] = layer
@@ -530,7 +588,12 @@ def _find_new_mapping(self):
 
     def _compute_cost(self, candidate):
         """Compute the cost associated to a possible SWAP candidate."""
-        temporary_circuit = deepcopy(self.circuit)
+        temporary_circuit = CircuitMap(
+            initial_layout=self.circuit.initial_layout,
+            circuit=Circuit(len(self.circuit.initial_layout)),
+            blocks=self.circuit.circuit_blocks,
+        )
+        temporary_circuit.set_circuit_logical(deepcopy(self.circuit._circuit_logical))
         temporary_circuit.update(candidate)
         if temporary_circuit._circuit_logical in self._memory_map:
             return float("inf")

tests/test_transpiler_router.py

@@ -298,28 +298,33 @@ def test_sabre_simple(seed):
     )
 
 
-@pytest.mark.parametrize("gates", [10, 40])
+@pytest.mark.parametrize("n_gates", [10, 40])
 @pytest.mark.parametrize("look", [0, 5])
 @pytest.mark.parametrize("decay", [0.5, 1.0])
 @pytest.mark.parametrize("placer", [Trivial, Random])
 @pytest.mark.parametrize("connectivity", [star_connectivity(), grid_connectivity()])
-def test_sabre_random_circuits(gates, look, decay, placer, connectivity):
+def test_sabre_random_circuits(n_gates, look, decay, placer, connectivity):
     placer = placer(connectivity=connectivity)
     layout_circ = Circuit(5)
     initial_layout = placer(layout_circ)
     router = Sabre(connectivity=connectivity, lookahead=look, decay_lookahead=decay)
-    circuit = generate_random_circuit(nqubits=5, ngates=gates)
+    circuit = generate_random_circuit(nqubits=5, ngates=n_gates)
+    measurement = gates.M(*range(5))
+    circuit.add(measurement)
     transpiled_circuit, final_qubit_map = router(circuit, initial_layout)
     assert router.added_swaps >= 0
     assert_connectivity(connectivity, transpiled_circuit)
     assert_placement(transpiled_circuit, final_qubit_map)
-    assert gates + router.added_swaps == transpiled_circuit.ngates
+    assert n_gates + router.added_swaps + 1 == transpiled_circuit.ngates
     assert_circuit_equivalence(
         original_circuit=circuit,
         transpiled_circuit=transpiled_circuit,
         final_map=final_qubit_map,
         initial_map=initial_layout,
     )
+    circuit_result = transpiled_circuit.execute(nshots=100)
+    assert circuit_result.frequencies() == measurement.result.frequencies()
+    assert transpiled_circuit.queue[-1].result is measurement.result
 
 
 def test_sabre_memory_map():
@@ -331,3 +336,19 @@ def test_sabre_memory_map():
     router._memory_map = [[1, 0, 2, 3, 4]]
     value = router._compute_cost((0, 1))
     assert value == float("inf")
+
+
+def test_sabre_intermediate_measurements():
+    measurement = gates.M(1)
+    circ = Circuit(3, density_matrix=True)
+    circ.add(gates.H(0))
+    circ.add(measurement)
+    circ.add(gates.CNOT(0, 2))
+    connectivity = nx.Graph()
+    connectivity.add_nodes_from([0, 1, 2])
+    connectivity.add_edges_from([(0, 1), (1, 2)])
+    router = Sabre(connectivity=connectivity)
+    initial_layout = {"q0": 0, "q1": 1, "q2": 2}
+    routed_circ, final_layout = router(circuit=circ, initial_layout=initial_layout)
+    circuit_result = routed_circ.execute(nshots=100)
+    assert routed_circ.queue[3].result is measurement.result