Improve efficiency of merge_scheduled_circuits

src/tqec/circuit/generation.py

@@ -18,13 +18,14 @@
 
 from __future__ import annotations
 
-from copy import deepcopy
-
 import numpy
 import numpy.typing as npt
 
-from tqec.circuit.qubit import GridQubit
-from tqec.circuit.schedule import ScheduledCircuit, merge_scheduled_circuits
+from tqec.circuit.schedule import (
+    ScheduledCircuit,
+    merge_scheduled_circuits,
+    relabel_circuits_qubit_indices,
+)
 from tqec.exceptions import TQECException
 from tqec.plaquette.plaquette import Plaquette, Plaquettes
 from tqec.position import Displacement
@@ -142,13 +143,26 @@ def generate_circuit_from_instantiation(
                     plaquette.origin.x + column_index * increments.x,
                     plaquette.origin.y + row_index * increments.y,
                 )
+                # Warning: the variable `mapped_scheduled_circuit` shares with
+                #          `plaquette_circuits[plaquette_index]` a reference to
+                #          the circuit data-structure. This is not an issue here
+                #          as we never attempt to mutate that circuit before
+                #          calling `relabel_circuits_qubit_indices`, that
+                #          explicitly returns a copy of its inputs, and do not
+                #          mutate them either.
                 mapped_scheduled_circuit = scheduled_circuit.map_to_qubits(
-                    lambda q: q + qubit_offset, inplace=False
+                    lambda q: q + qubit_offset, inplace_qubit_map=False
                 )
                 all_scheduled_circuits.append(mapped_scheduled_circuit)
                 additional_mergeable_instructions |= plaquette.mergeable_instructions
 
-    # Merge everything!
+    # Merge everything, but first make sure that the circuits are compatible.
+    # Note that relabel_circuits_qubit_indices guarantees in its documentation
+    # that the input circuits are not mutated but rather copied. This allows us
+    # to not deepcopy the circuits earlier in the function.
+    all_scheduled_circuits, qubit_map = relabel_circuits_qubit_indices(
+        all_scheduled_circuits
+    )
     return merge_scheduled_circuits(
-        all_scheduled_circuits, additional_mergeable_instructions
+        all_scheduled_circuits, qubit_map, additional_mergeable_instructions
     )

src/tqec/circuit/schedule.py

@@ -30,7 +30,7 @@
 import itertools
 import typing as ty
 import warnings
-from copy import deepcopy
+from copy import copy, deepcopy
 from dataclasses import dataclass, field
 
 import stim
@@ -407,26 +407,33 @@ def map_qubit_indices(
     def map_to_qubits(
         self,
         qubit_map: ty.Callable[[GridQubit], GridQubit],
-        inplace: bool = False,
+        inplace_qubit_map: bool = False,
     ) -> ScheduledCircuit:
         """Map the qubits the `ScheduledCircuit` instance is applied on.
 
         Note:
             This method only changes the `QUBIT_COORDS` instructions at the
-            beginning of the circuit. As long as `inplace` is `True`, this
-            method is very efficient. If `inplace` is `False`, the deep copy of
-            `self` is the most costly part of this method.
+            beginning of the circuit. As such, it never has to iterate on the
+            whole quantum circuit and so this method is very efficient.
+
+        Warning:
+            The underlying quantum circuit data-structure is never copied (even
+            if `inplace_qubit_map == True`), so the returned instance should be
+            used with care, in particular if any method mutating the underlying
+            circuit is called on `self` or the returned instance after calling
+            that method.
 
         Args:
             qubit_map: the map used to modify the qubits.
-            inplace: if True, perform the modification in place and return self.
-                Else, perform the modification in a copy and return the copy.
+            inplace_qubit_map: if True, replaces the qubit map directly in
+                `self` and return `self`. Else, create a new instance from
+                `self` **without copying the underlying moments**, simply
+                replacing the qubit map.
 
         Returns:
-            a modified instance of `ScheduledCircuit` (a copy if inplace is True,
-            else self).
+            an instance of `ScheduledCircuit` with a new qubit map.
         """
-        operand = self if inplace else deepcopy(self)
+        operand = self if inplace_qubit_map else copy(self)
         operand._qubit_map = operand._qubit_map.with_mapped_qubits(qubit_map)
         return operand
 
@@ -627,26 +634,28 @@ def qubit_map(self) -> QubitMap:
 
 
 class _ScheduledCircuits:
-    def __init__(self, circuits: list[ScheduledCircuit]) -> None:
+    def __init__(
+        self, circuits: list[ScheduledCircuit], global_qubit_map: QubitMap
+    ) -> None:
         """Represents a collection of :class`ScheduledCircuit` instances.
 
-        This class aims at providing accessors for several instances of
-        :class:`ScheduledCircuit`. It allows to iterate on gates globally, for
+        This class aims at providing accessors for several compatible instances
+        of :class:`ScheduledCircuit`. It allows to iterate on gates globally, for
         all the managed instances of :class:`ScheduledCircuit`, and implement a
         few other accessor methods to help with the task of merging multiple
         :class`ScheduledCircuit` together.
 
         Args:
             circuits: the instances that should be managed. Note that the
-                instances provided here do not have to be "compatible" with each
-                other. In particular, the qubit indices of each circuit can
-                overlap. Due to the computations that are needed internally to
-                avoid overlapping indices (that will break the computation), the
-                instances provided here are copied in the `__init__` method.
+                instances provided here have to be "compatible" with each
+                other.
+            global_qubit_map: a unique qubit map that can be used to map qubits
+                to indices for all the provided `circuits`.
         """
         # We might need to remap qubits to avoid index collision on several
         # circuits.
-        self._circuits, self._qubit_map = relabel_circuits_qubit_indices(circuits)
+        self._circuits = circuits
+        self._global_qubit_map = global_qubit_map
         self._iterators = [circuit.scheduled_moments for circuit in self._circuits]
         self._current_moments = [next(it, None) for it in self._iterators]
 
@@ -717,7 +726,7 @@ def collect_moments_at_minimum_schedule(self) -> tuple[int, list[Moment]]:
 
     @property
     def q2i(self) -> dict[GridQubit, int]:
-        return self._qubit_map.q2i
+        return self._global_qubit_map.q2i
 
 
 def _sort_target_groups(
@@ -806,27 +815,33 @@ def remove_duplicate_instructions(
 
 def merge_scheduled_circuits(
     circuits: list[ScheduledCircuit],
+    global_qubit_map: QubitMap,
     mergeable_instructions: ty.Iterable[str] = (),
 ) -> ScheduledCircuit:
     """Merge several ScheduledCircuit instances into one instance.
 
-    This function takes several scheduled circuits as input and merge them,
-    respecting their schedules, into a unique `ScheduledCircuit` instance that
-    will then be returned to the caller.
+    This function takes several **compatible** scheduled circuits as input and
+    merge them, respecting their schedules, into a unique `ScheduledCircuit`
+    instance that will then be returned to the caller.
 
-    KeyError: if any of the provided circuit contains a qubit target that is
-            not defined by a `QUBIT_COORDS` instruction.
+    The provided circuits should be compatible between each other. Compatible
+    circuits are circuits that can all be described with a unique global qubit
+    map. In other words, if two circuits from the list of compatible circuits
+    use the same qubit index, that should mean that they use the same qubit.
+    You can obtain compatible circuits by using
+    :func:`relabel_circuits_qubit_indices`.
 
     Args:
-        circuits: the circuits to merge.
+        circuits: **compatible** circuits to merge.
+        qubit_map: global qubit map for all the provided `circuits`.
         mergeable_instructions: a list of instruction names that are considered
             mergeable. Duplicate instructions with a name in this list will be
             merged into a single instruction.
 
     Returns:
         a circuit representing the merged scheduled circuits given as input.
     """
-    scheduled_circuits = _ScheduledCircuits(circuits)
+    scheduled_circuits = _ScheduledCircuits(circuits, global_qubit_map)
 
     all_moments: list[Moment] = []
     all_schedules = Schedule()
@@ -872,13 +887,12 @@ def relabel_circuits_qubit_indices(
 
     Warning:
         all the qubit targets used in each of the provided circuits should have
-        a corresponding `QUBIT_COORDS([coordinates]) [qubit target]` for this
-        function to work correctly. If that is not the case, a KeyError will be
-        raised.
+        a corresponding entry in the circuit qubit map for this function to work
+        correctly. If that is not the case, a KeyError will be raised.
 
     Raises:
         KeyError: if any of the provided circuit contains a qubit target that is
-            not defined by a `QUBIT_COORDS` instruction.
+            not present in its qubit map.
 
     Args:
         circuits: circuit instances to remap. This parameter is not mutated by
@@ -890,7 +904,7 @@ def relabel_circuits_qubit_indices(
         are assigned to an index such that:
 
         1. the sequence of indices is `range(0, len(qubit_map))`.
-        2. the returned qubit map
+        2. qubits are assigned indices in sorted order.
     """
     # First, get a global qubit index map.
     # Using itertools to avoid the edge case `len(circuits) == 0`

src/tqec/circuit/schedule_test.py

@@ -269,7 +269,7 @@ def test_scheduled_circuit_map_to_qubits() -> None:
     assert circuit.get_circuit() == stim.Circuit(
         "QUBIT_COORDS(0.0, 0.0) 0\nQUBIT_COORDS(0.0, 1.0) 1\nH 0 1"
     )
-    circuit.map_to_qubits(lambda q: qubit_map[q], inplace=True)
+    circuit.map_to_qubits(lambda q: qubit_map[q], inplace_qubit_map=True)
     assert circuit.get_circuit() == mapped_circuit.get_circuit()
 
 
@@ -371,20 +371,13 @@ def test_relabel_circuits_qubit_indices() -> None:
 def test_merge_scheduled_circuits() -> None:
     # Any qubit target not defined by a QUBIT_COORDS instruction should raise
     # an exception.
-    with pytest.raises(KeyError):
-        merge_scheduled_circuits(
-            [
-                ScheduledCircuit.from_circuit(stim.Circuit("H 0 1 2")),
-                ScheduledCircuit.from_circuit(stim.Circuit("H 0 1 2")),
-            ]
-        )
-
-    circuit = merge_scheduled_circuits(
+    _circuits, _qubit_map = relabel_circuits_qubit_indices(
         [
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0")),
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(1, 1) 0\nX 0")),
         ]
     )
+    circuit = merge_scheduled_circuits(_circuits, _qubit_map)
     assert circuit.get_circuit() == stim.Circuit(
         "QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 1) 1\nH 0\nX 1"
     )
@@ -394,18 +387,20 @@ def test_merge_scheduled_circuits() -> None:
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0")),
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0")),
         ],
+        global_qubit_map=QubitMap({0: GridQubit(0, 0)}),
         mergeable_instructions=["H"],
     )
     assert circuit.get_circuit() == stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0")
 
-    circuit = merge_scheduled_circuits(
+    _circuits, _qubit_map = relabel_circuits_qubit_indices(
         [
             ScheduledCircuit.from_circuit(
                 stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0\nTICK\nM 0"), [0, 2]
             ),
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(1, 1) 0\nX 0"), 1),
         ]
     )
+    circuit = merge_scheduled_circuits(_circuits, _qubit_map)
     assert circuit.get_circuit() == stim.Circuit(
         "QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 1) 1\nH 0\nTICK\nX 1\nTICK\nM 0"
     )

src/tqec/compile/block.py

@@ -138,7 +138,7 @@ def get_shifted_circuits(self, k: int) -> list[ScheduledCircuit]:
         )
         return [
             generate_circuit(self._template, k, layer).map_to_qubits(
-                lambda q: q + offset, inplace=True
+                lambda q: q + offset, inplace_qubit_map=True
             )
             for layer in self._layers
         ]