Repository: QCHackers/tqec
Pull request #344 (2024-10-09 14:50:51)

Introduce a `QubitMap` structure
During the circuit construction stage, we are performing a lot of operations on qubits such as remapping, changing indices, offsetting, etc. I think it would be beneficial for the long term to have a specific data-structure to represent a map* from indices to qubits.

* it is more than a map, it should be a bijection. In that regard, I think using https://github.com/jab/bidict might be a good idea, but I am a little bit reluctant to add a whole dependency just for that. Any opinion on whether or not we should use that package or keep the implementation as it is for the moment?

github-actions[bot]: A preview of  is uploaded and can be seen here:

✨ https://qchackers.github.io/tqec/pull/344/ ✨

Changes may take a few minutes to propagate.
<!-- Sticky Pull Request Commentdocumentation-preview -->
github-actions[bot]: ![Code Coverage](https://img.shields.io/badge/Code%20Coverage-81%25-success?style=flat)

Package | Line Rate | Complexity | Health
-------- | --------- | ---------- | ------
src.tqec | 97% | 0 | ✔
src.tqec.circuit | 97% | 0 | ✔
src.tqec.circuit.detectors | 94% | 0 | ✔
src.tqec.circuit.detectors.match_utils | 90% | 0 | ✔
src.tqec.compile | 34% | 0 | ❌
src.tqec.computation | 74% | 0 | ➖
src.tqec.computation.block_graph | 81% | 0 | ✔
src.tqec.noise_models | 91% | 0 | ✔
src.tqec.plaquette | 88% | 0 | ✔
src.tqec.plaquette.library | 100% | 0 | ✔
src.tqec.simulation | 0% | 0 | ❌
src.tqec.simulation.plotting | 0% | 0 | ❌
src.tqec.templates | 78% | 0 | ✔
**Summary** | **81%** (3173 / 3912) | **0** | ✔

<!-- Sticky Pull Request Commentcode-coverage -->
github-actions[bot]: A preview of 77e10701b3f31e478e9c099fdb9ffcbcc7742238 is uploaded and can be seen here:

✨ https://qchackers.github.io/tqec/pull/344/ ✨

Changes may take a few minutes to propagate.
<!-- Sticky Pull Request Commentdocumentation-preview -->
github-actions[bot]: ![Code Coverage](https://img.shields.io/badge/Code%20Coverage-81%25-success?style=flat)

Package | Line Rate | Complexity | Health
-------- | --------- | ---------- | ------
src.tqec | 97% | 0 | ✔
src.tqec.circuit | 97% | 0 | ✔
src.tqec.circuit.detectors | 94% | 0 | ✔
src.tqec.circuit.detectors.match_utils | 90% | 0 | ✔
src.tqec.compile | 34% | 0 | ❌
src.tqec.computation | 74% | 0 | ➖
src.tqec.computation.block_graph | 81% | 0 | ✔
src.tqec.noise_models | 91% | 0 | ✔
src.tqec.plaquette | 88% | 0 | ✔
src.tqec.plaquette.library | 100% | 0 | ✔
src.tqec.simulation | 0% | 0 | ❌
src.tqec.simulation.plotting | 0% | 0 | ❌
src.tqec.templates | 78% | 0 | ✔
**Summary** | **81%** (3173 / 3912) | **0** | ✔

<!-- Sticky Pull Request Commentcode-coverage -->

Patch:
diff --git a/src/tqec/circuit/qubit.py b/src/tqec/circuit/qubit.py
index d66c5ada..2aac2423 100644
--- a/src/tqec/circuit/qubit.py
+++ b/src/tqec/circuit/qubit.py
@@ -13,7 +13,6 @@ from fractions import Fraction
 
 import stim
 
-from tqec.exceptions import TQECException
 from tqec.position import Displacement, Position2D
 
 NumericType = int | float | Fraction
@@ -164,34 +163,3 @@ def get_used_qubit_indices(circuit: stim.Circuit) -> frozenset[int]:
         instruction.
     """
     return frozenset(count_qubit_accesses(circuit).keys())
-
-
-def get_final_qubits(circuit: stim.Circuit) -> dict[int, GridQubit]:
-    """Returns the existing qubits and their coordinates at the end of the
-    provided `circuit`.
-
-    Warning:
-        This function, just like
-        [`stim.Circuit.get_final_qubit_coordinates`](https://github.com/quantumlib/Stim/blob/main/doc/python_api_reference_vDev.md#stim.Circuit.get_final_qubit_coordinates),
-        returns the qubit coordinates **at the end** of the provided `circuit`.
-
-    Args:
-        circuit: instance to get qubit coordinates from.
-
-    Raises:
-        TQECException: if any of the final qubits is not defined with exactly 2
-            coordinates (we only consider qubits on a 2-dimensional grid).
-
-    Returns:
-        a mapping from qubit indices (keys) to qubit coordinates (values).
-    """
-    qubit_coordinates = circuit.get_final_qubit_coordinates()
-    qubits: dict[int, GridQubit] = {}
-    for qi, coords in qubit_coordinates.items():
-        if len(coords) != 2:
-            raise TQECException(
-                "Qubits should be defined on exactly 2 spatial dimensions. "
-                f"Found {qi} -> {coords} defined on {len(coords)} spatial dimensions."
-            )
-        qubits[qi] = GridQubit(*coords)
-    return qubits
diff --git a/src/tqec/circuit/qubit_map.py b/src/tqec/circuit/qubit_map.py
new file mode 100644
index 00000000..87da7964
--- /dev/null
+++ b/src/tqec/circuit/qubit_map.py
@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import functools
+from collections import Counter
+from dataclasses import dataclass, field
+from typing import Iterable
+
+import stim
+
+from tqec.circuit.qubit import GridQubit
+from tqec.exceptions import TQECException
+
+
+@dataclass(frozen=True)
+class QubitMap:
+    """Represent a bijection between :class:`GridQubit` instances and indices.
+
+    This class aims at representing a bidirectional mapping (hence the
+    "bijection") between qubits and their associated indices.
+
+    Raises:
+        TQECException: if the provided mapping from indices to qubits is not a
+            bijection (i.e., if at least to values represent the same qubit).
+    """
+
+    i2q: dict[int, GridQubit] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        qubit_counter = Counter(self.i2q.values())
+        if len(qubit_counter) != len(self.i2q):
+            duplicated_qubits = frozenset(
+                q for q in qubit_counter if qubit_counter[q] > 1
+            )
+            raise TQECException(
+                f"Found qubit(s) with more than one index: {duplicated_qubits}."
+            )
+
+    @staticmethod
+    def from_qubits(qubits: Iterable[GridQubit]) -> QubitMap:
+        return QubitMap(dict(enumerate(qubits)))
+
+    @functools.cached_property
+    def q2i(self) -> dict[GridQubit, int]:
+        return {q: i for i, q in self.i2q.items()}
+
+    @property
+    def indices(self) -> Iterable[int]:
+        return self.i2q.keys()
+
+    @property
+    def qubits(self) -> Iterable[GridQubit]:
+        return self.i2q.values()
+
+    def with_mapped_qubits(self, qubit_map: dict[GridQubit, GridQubit]) -> QubitMap:
+        """Change the qubits involved in `self` without changing the associated indices.
+
+        Args:
+            qubit_map: a map from qubits to qubits that should associate a qubit
+                to each of the qubits represented by `self`.
+
+        Raises:
+            KeyError: if any qubit in `self.qubits` is not present in the keys of
+                the provided `qubit_map`.
+
+        Returns:
+            a new instance representing the updated mapping.
+        """
+        return QubitMap({i: qubit_map[q] for i, q in self.i2q.items()})
+
+    def items(self) -> Iterable[tuple[int, GridQubit]]:
+        return self.i2q.items()
+
+    def filter_by_qubits(self, qubits_to_keep: Iterable[GridQubit]) -> QubitMap:
+        """Filter the qubit map to only keep qubits present in the provided
+        `qubits_to_keep`.
+
+        Args:
+            qubits_to_keep: the qubits to keep in the circuit.
+
+        Returns:
+            a copy of `self` for which the assertion
+            `set(return_value.qubits).issubset(qubits_to_keep)` is `True`.
+        """
+        kept_qubits = frozenset(qubits_to_keep)
+        return QubitMap({i: q for i, q in self.i2q.items() if q in kept_qubits})
+
+
+def get_qubit_map(circuit: stim.Circuit) -> QubitMap:
+    """Returns the existing qubits and their coordinates at the end of the
+    provided `circuit`.
+
+    Warning:
+        This function, just like
+        [`stim.Circuit.get_final_qubit_coordinates`](https://github.com/quantumlib/Stim/blob/main/doc/python_api_reference_vDev.md#stim.Circuit.get_final_qubit_coordinates),
+        returns the qubit coordinates **at the end** of the provided `circuit`.
+
+    Args:
+        circuit: instance to get qubit coordinates from.
+
+    Raises:
+        TQECException: if any of the final qubits is not defined with exactly 2
+            coordinates (we only consider qubits on a 2-dimensional grid).
+
+    Returns:
+        a mapping from qubit indices (keys) to qubit coordinates (values).
+    """
+    qubit_coordinates = circuit.get_final_qubit_coordinates()
+    qubits: dict[int, GridQubit] = {}
+    for qi, coords in qubit_coordinates.items():
+        if len(coords) != 2:
+            raise TQECException(
+                "Qubits should be defined on exactly 2 spatial dimensions. "
+                f"Found {qi} -> {coords} defined on {len(coords)} spatial dimensions."
+            )
+        qubits[qi] = GridQubit(*coords)
+    return QubitMap(qubits)
diff --git a/src/tqec/circuit/schedule.py b/src/tqec/circuit/schedule.py
index c98a3b72..ab91ac9a 100644
--- a/src/tqec/circuit/schedule.py
+++ b/src/tqec/circuit/schedule.py
@@ -35,7 +35,8 @@ from dataclasses import dataclass, field
 import stim
 
 from tqec.circuit.moment import Moment, iter_stim_circuit_without_repeat_by_moments
-from tqec.circuit.qubit import GridQubit, get_final_qubits
+from tqec.circuit.qubit import GridQubit
+from tqec.circuit.qubit_map import QubitMap, get_qubit_map
 from tqec.exceptions import TQECException, TQECWarning
 
 
@@ -139,7 +140,7 @@ class ScheduledCircuit:
         self,
         moments: list[Moment],
         schedule: Schedule | list[int] | int,
-        i2q: dict[int, GridQubit],
+        qubit_map: QubitMap,
     ) -> None:
         """Represent a quantum circuit with scheduled moments.
 
@@ -155,7 +156,7 @@ class ScheduledCircuit:
             schedule: schedule of the provided `moments`. If an integer is
                 provided, each moment of the provided `stim.Circuit` is
                 scheduled sequentially, starting by the provided integer.
-            i2q: a map from indices to qubits that is used to re-create
+            qubit_map: a map from indices to qubits that is used to re-create
                 `QUBIT_COORDS` instructions when generating a `stim.Circuit`.
 
         Raises:
@@ -183,18 +184,18 @@ class ScheduledCircuit:
             )
 
         self._moments: list[Moment] = moments
-        self._i2q: dict[int, GridQubit] = i2q
+        self._qubit_map: QubitMap = qubit_map
         self._schedule: Schedule = schedule
 
     @staticmethod
     def empty() -> ScheduledCircuit:
-        return ScheduledCircuit([], Schedule(), {})
+        return ScheduledCircuit([], Schedule(), QubitMap())
 
     @staticmethod
     def from_circuit(
         circuit: stim.Circuit,
         schedule: Schedule | list[int] | int = 0,
-        i2q: dict[int, GridQubit] | None = None,
+        i2q: QubitMap | None = None,
     ) -> ScheduledCircuit:
         """Build a :class:`ScheduledCircuit` instance from a circuit and a
         schedule.
@@ -243,7 +244,7 @@ class ScheduledCircuit:
             # Here, we know for sure that no qubit is (re)defined in another place
             # than the first moment, so we can directly get qubit coordinates from
             # that moment.
-            i2q = get_final_qubits(moments[0].circuit)
+            i2q = get_qubit_map(moments[0].circuit)
         # And because we want the cleanest possible moments, we can remove the
         # `QUBIT_COORDS` instructions from the first moment.
         moments[0].remove_all_instructions_inplace(frozenset(["QUBIT_COORDS"]))
@@ -266,7 +267,7 @@ class ScheduledCircuit:
     def get_qubit_coords_definition_preamble(self) -> stim.Circuit:
         """Get a circuit with only `QUBIT_COORDS` instructions."""
         ret = stim.Circuit()
-        for qi, qubit in sorted(self._i2q.items(), key=lambda t: t[0]):
+        for qi, qubit in sorted(self._qubit_map.items(), key=lambda t: t[0]):
             ret.append("QUBIT_COORDS", qi, (float(qubit.x), float(qubit.y)))
         return ret
 
@@ -376,7 +377,9 @@ class ScheduledCircuit:
             a modified instance of `ScheduledCircuit` (a copy if inplace is
             `True`, else `self`).
         """
-        mapped_final_qubits = {qubit_index_map[qi]: q for qi, q in self._i2q.items()}
+        mapped_final_qubits = QubitMap(
+            {qubit_index_map[qi]: q for qi, q in self._qubit_map.items()}
+        )
         mapped_moments: list[Moment] = []
         for moment in self._moments:
             mapped_moment_circuit = stim.Circuit()
@@ -399,7 +402,7 @@ class ScheduledCircuit:
                 )
             mapped_moments.append(Moment(mapped_moment_circuit))
         if inplace:
-            self._i2q = mapped_final_qubits
+            self._qubit_map = mapped_final_qubits
             self._moments = mapped_moments
             return self
         else:
@@ -430,21 +433,17 @@ class ScheduledCircuit:
             else self).
         """
         operand = self if inplace else deepcopy(self)
-        operand._i2q = {qi: qubit_map[q] for qi, q in operand._i2q.items()}
+        operand._qubit_map = operand._qubit_map.with_mapped_qubits(qubit_map)
         return operand
 
     def __copy__(self) -> ScheduledCircuit:
-        return ScheduledCircuit(
-            self._moments,
-            self._schedule,
-            self._i2q,
-        )
+        return ScheduledCircuit(self._moments, self._schedule, self._qubit_map)
 
     def __deepcopy__(self, memo: dict[ty.Any, ty.Any]) -> ScheduledCircuit:
         return ScheduledCircuit(
             deepcopy(self._moments, memo=memo),
             deepcopy(self._schedule, memo=memo),
-            deepcopy(self._i2q, memo=memo),
+            deepcopy(self._qubit_map, memo=memo),
         )
 
     @property
@@ -464,7 +463,7 @@ class ScheduledCircuit:
 
     @property
     def qubits(self) -> frozenset[GridQubit]:
-        return frozenset(self._i2q.values())
+        return frozenset(self._qubit_map.qubits)
 
     def _get_moment_index_by_schedule(self, schedule: int) -> int | None:
         """Get the index of the moment scheduled at the provided schedule.
@@ -558,7 +557,7 @@ class ScheduledCircuit:
     @property
     def q2i(self) -> dict[GridQubit, int]:
         """Return the map from qubits used in `self` their index."""
-        return {q: i for i, q in self._i2q.items()}
+        return {q: i for i, q in self._qubit_map.items()}
 
     def append_observable(
         self, index: int, targets: ty.Sequence[stim.GateTarget]
@@ -608,14 +607,12 @@ class ScheduledCircuit:
                 continue
             filtered_moments.append(filtered_moment)
             filtered_schedule.append(schedule)
-        i2q = {i: q for i, q in self._i2q.items() if i in qubits_indices_to_keep}
+        qubit_map = self._qubit_map.filter_by_qubits(qubits_to_keep)
         filtered_circuit = ScheduledCircuit(
-            filtered_moments,
-            Schedule(filtered_schedule),
-            i2q,
+            filtered_moments, Schedule(filtered_schedule), qubit_map
         )
         # The qubit indices may not be contiguous anymore, so we need to remap them.
-        indices_map = {oi: ni for ni, oi in enumerate(i2q.keys())}
+        indices_map = {oi: ni for ni, oi in enumerate(qubit_map.indices)}
         return filtered_circuit.map_qubit_indices(indices_map)
 
 
@@ -639,7 +636,7 @@ class _ScheduledCircuits:
         """
         # We might need to remap qubits to avoid index collision on several
         # circuits.
-        self._circuits, self._global_q2i = relabel_circuits_qubit_indices(circuits)
+        self._circuits, self._qubit_map = relabel_circuits_qubit_indices(circuits)
         self._iterators = [circuit.scheduled_moments for circuit in self._circuits]
         self._current_moments = [next(it, None) for it in self._iterators]
 
@@ -710,7 +707,7 @@ class _ScheduledCircuits:
 
     @property
     def q2i(self) -> dict[GridQubit, int]:
-        return self._global_q2i
+        return self._qubit_map.q2i
 
 
 def _sort_target_groups(
@@ -823,7 +820,7 @@ def merge_scheduled_circuits(
 
     all_moments: list[Moment] = []
     all_schedules = Schedule()
-    global_i2q: dict[int, GridQubit] = {i: q for q, i in scheduled_circuits.q2i.items()}
+    global_i2q = QubitMap({i: q for q, i in scheduled_circuits.q2i.items()})
 
     while scheduled_circuits.has_pending_moment():
         schedule, moments = scheduled_circuits.collect_moments_at_minimum_schedule()
@@ -855,7 +852,7 @@ def merge_scheduled_circuits(
 
 def relabel_circuits_qubit_indices(
     circuits: ty.Sequence[ScheduledCircuit],
-) -> tuple[list[ScheduledCircuit], dict[GridQubit, int]]:
+) -> tuple[list[ScheduledCircuit], QubitMap]:
     """Relabel the qubit indices of the provided circuits to avoid collision.
 
     When several :class:`ScheduledCircuit` are constructed without a global
@@ -886,12 +883,12 @@ 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 list `[qubit_map[q] for q in sorted(qubit_map)]` is exactly
-           `list(range(len(qubit_map)))`.
+        2. the returned qubit map
     """
     # First, get a global qubit index map.
     needed_qubits = frozenset.union(*[c.qubits for c in circuits])
-    global_q2i = {q: i for i, q in enumerate(sorted(needed_qubits))}
+    qubit_map = QubitMap.from_qubits(sorted(needed_qubits))
+    q2i = qubit_map.q2i
     # Then, get the remapped circuits. Note that map_qubit_indices should
     # have approximately the same runtime cost whatever the value of inplace
     # so we ask for a new instance to avoid keeping a reference to the given
@@ -899,9 +896,9 @@ def relabel_circuits_qubit_indices(
     relabeled_circuits: list[ScheduledCircuit] = []
     for circuit in circuits:
         local_indices_to_global_indices = {
-            local_index: global_q2i[q] for q, local_index in circuit.q2i.items()
+            local_index: q2i[q] for q, local_index in circuit.q2i.items()
         }
         relabeled_circuits.append(
             circuit.map_qubit_indices(local_indices_to_global_indices, inplace=False)
         )
-    return relabeled_circuits, global_q2i
+    return relabeled_circuits, qubit_map
diff --git a/src/tqec/plaquette/library/css.py b/src/tqec/plaquette/library/css.py
index 113c8f2c..7e9e3e57 100644
--- a/src/tqec/plaquette/library/css.py
+++ b/src/tqec/plaquette/library/css.py
@@ -6,6 +6,7 @@ from typing import Literal
 
 import stim
 
+from tqec.circuit.qubit_map import QubitMap
 from tqec.circuit.schedule import ScheduledCircuit
 from tqec.plaquette.enums import MeasurementBasis, ResetBasis
 from tqec.plaquette.plaquette import Plaquette
@@ -35,9 +36,10 @@ def make_css_surface_code_plaquette(
     sq = 0
     # data qubits ordered as "Z" shape
     dqs = range(1, 5)
-    i2q = {0: qubits.syndrome_qubits[0]} | {
-        i + 1: q for i, q in enumerate(qubits.data_qubits)
-    }
+    i2q = QubitMap(
+        {0: qubits.syndrome_qubits[0]}
+        | {i + 1: q for i, q in enumerate(qubits.data_qubits)}
+    )
 
     circuit = stim.Circuit()
     # 1. Initialization
diff --git a/src/tqec/plaquette/library/zxxz.py b/src/tqec/plaquette/library/zxxz.py
index 9dfdf18a..5339814e 100644
--- a/src/tqec/plaquette/library/zxxz.py
+++ b/src/tqec/plaquette/library/zxxz.py
@@ -6,6 +6,7 @@ from typing import Literal
 
 import stim
 
+from tqec.circuit.qubit_map import QubitMap
 from tqec.circuit.schedule import ScheduledCircuit
 from tqec.plaquette.enums import MeasurementBasis, ResetBasis
 from tqec.plaquette.plaquette import Plaquette
@@ -38,9 +39,10 @@ def make_zxxz_surface_code_plaquette(
     sq = 0
     # data qubits ordered as "Z" shape
     dqs = list(range(1, 5))
-    i2q = {0: qubits.syndrome_qubits[0]} | {
-        i + 1: q for i, q in enumerate(qubits.data_qubits)
-    }
+    i2q = QubitMap(
+        {0: qubits.syndrome_qubits[0]}
+        | {i + 1: q for i, q in enumerate(qubits.data_qubits)}
+    )
 
     circuit = stim.Circuit()
     # 1. Initialization


Test patch:
diff --git a/src/tqec/circuit/qubit_map_test.py b/src/tqec/circuit/qubit_map_test.py
new file mode 100644
index 00000000..d8288dfe
--- /dev/null
+++ b/src/tqec/circuit/qubit_map_test.py
@@ -0,0 +1,93 @@
+import re
+
+import pytest
+import stim
+
+from tqec.circuit.qubit import GridQubit
+from tqec.circuit.qubit_map import QubitMap, get_qubit_map
+from tqec.exceptions import TQECException
+
+
+def test_qubit_map_creation() -> None:
+    q = GridQubit(0, 0)
+    QubitMap()
+    QubitMap({0: q})
+    QubitMap({1: q})
+
+    qset_str = re.escape(str(frozenset([q])))
+    with pytest.raises(
+        TQECException,
+        match=f"^Found qubit\\(s\\) with more than one index: {qset_str}.$",
+    ):
+        QubitMap({0: q, 1: q})
+
+    QubitMap.from_qubits([])
+    QubitMap.from_qubits([GridQubit(0, 0)])
+    with pytest.raises(
+        TQECException,
+        match=f"^Found qubit\\(s\\) with more than one index: {qset_str}.$",
+    ):
+        QubitMap.from_qubits([q, q])
+
+
+def test_qubit_map_getters() -> None:
+    assert frozenset(QubitMap().indices) == frozenset()
+    assert frozenset(QubitMap().qubits) == frozenset()
+
+    qmap = QubitMap({i: GridQubit(i, -i) for i in range(4)})
+    assert frozenset(qmap.indices) == frozenset(range(4))
+    assert frozenset(qmap.qubits) == frozenset(GridQubit(i, -i) for i in range(4))
+
+
+def test_qubit_map_q2i() -> None:
+    assert QubitMap().q2i == {}
+
+    qmap = QubitMap({i: GridQubit(i, -i) for i in range(4)})
+    assert qmap.q2i == {GridQubit(i, -i): i for i in range(4)}
+
+
+def test_qubit_map_with_mapped_qubits() -> None:
+    qubit_map = {GridQubit(i, -i): GridQubit(-i, i) for i in range(4)}
+    assert QubitMap().with_mapped_qubits(qubit_map) == QubitMap()
+
+    qmap = QubitMap({i: GridQubit(i, -i) for i in range(4)})
+    assert qmap.with_mapped_qubits(qubit_map) == QubitMap(
+        {i: GridQubit(-i, i) for i in range(4)}
+    )
+
+
+def test_qubit_map_filter_by_qubits() -> None:
+    qubits = [GridQubit(i, -i) for i in range(3)]
+    assert QubitMap().filter_by_qubits(qubits) == QubitMap()
+
+    qmap = QubitMap({i: GridQubit(i, -i) for i in range(4)})
+    assert qmap.filter_by_qubits(qubits) == QubitMap(
+        {i: GridQubit(i, -i) for i in range(3)}
+    )
+
+
+def test_get_final_qubits() -> None:
+    assert get_qubit_map(stim.Circuit("QUBIT_COORDS(0, 0) 0")) == QubitMap(
+        {0: GridQubit(0, 0)}
+    )
+    assert get_qubit_map(stim.Circuit("QUBIT_COORDS(0, 0.5) 0")) == QubitMap(
+        {0: GridQubit(0, 0.5)}
+    )
+    assert get_qubit_map(stim.Circuit("QUBIT_COORDS(0, 0) 1")) == QubitMap(
+        {1: GridQubit(0, 0)}
+    )
+    # assert get_final_qubits(
+    #     stim.Circuit("QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 4) 0")
+    # ) == {0: GridQubit(1, 4)}
+    assert get_qubit_map(
+        stim.Circuit("QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 4) 1")
+    ) == QubitMap({0: GridQubit(0, 0), 1: GridQubit(1, 4)})
+
+    with pytest.raises(
+        TQECException,
+        match=re.escape(
+            "Qubits should be defined on exactly 2 spatial dimensions. "
+            f"Found {0} -> [0.0, 0.0, 1.0] defined on 3 spatial dimensions."
+        ),
+    ):
+        get_qubit_map(stim.Circuit("QUBIT_COORDS(0, 0, 1) 0"))
diff --git a/src/tqec/circuit/qubit_test.py b/src/tqec/circuit/qubit_test.py
index 080ec493..28645828 100644
--- a/src/tqec/circuit/qubit_test.py
+++ b/src/tqec/circuit/qubit_test.py
@@ -1,16 +1,9 @@
-import re
 from fractions import Fraction
 
 import pytest
 import stim
 
-from tqec.circuit.qubit import (
-    GridQubit,
-    count_qubit_accesses,
-    get_final_qubits,
-    get_used_qubit_indices,
-)
-from tqec.exceptions import TQECException
+from tqec.circuit.qubit import GridQubit, count_qubit_accesses, get_used_qubit_indices
 from tqec.position import Displacement, Position2D
 
 
@@ -109,30 +102,3 @@ def test_used_qubit_indices() -> None:
     assert get_used_qubit_indices(
         stim.Circuit("REPEAT 34{\nH 0 1 2 3\nTICK\nH 0 1 2 3\n}")
     ) == frozenset(range(4))
-
-
-def test_get_final_qubits() -> None:
-    assert get_final_qubits(stim.Circuit("QUBIT_COORDS(0, 0) 0")) == {
-        0: GridQubit(0, 0)
-    }
-    assert get_final_qubits(stim.Circuit("QUBIT_COORDS(0, 0.5) 0")) == {
-        0: GridQubit(0, 0.5)
-    }
-    assert get_final_qubits(stim.Circuit("QUBIT_COORDS(0, 0) 1")) == {
-        1: GridQubit(0, 0)
-    }
-    # assert get_final_qubits(
-    #     stim.Circuit("QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 4) 0")
-    # ) == {0: GridQubit(1, 4)}
-    assert get_final_qubits(
-        stim.Circuit("QUBIT_COORDS(0, 0) 0\nQUBIT_COORDS(1, 4) 1")
-    ) == {0: GridQubit(0, 0), 1: GridQubit(1, 4)}
-
-    with pytest.raises(
-        TQECException,
-        match=re.escape(
-            "Qubits should be defined on exactly 2 spatial dimensions. "
-            f"Found {0} -> [0.0, 0.0, 1.0] defined on 3 spatial dimensions."
-        ),
-    ):
-        get_final_qubits(stim.Circuit("QUBIT_COORDS(0, 0, 1) 0"))
diff --git a/src/tqec/circuit/schedule_test.py b/src/tqec/circuit/schedule_test.py
index 7bed9f0c..9aeb7b67 100644
--- a/src/tqec/circuit/schedule_test.py
+++ b/src/tqec/circuit/schedule_test.py
@@ -5,6 +5,7 @@ import stim
 
 from tqec.circuit.moment import Moment
 from tqec.circuit.qubit import GridQubit
+from tqec.circuit.qubit_map import QubitMap
 from tqec.circuit.schedule import (
     Schedule,
     ScheduledCircuit,
@@ -125,9 +126,9 @@ def test_scheduled_circuit_construction() -> None:
     with pytest.raises(ScheduleException):
         ScheduledCircuit.from_circuit(stim.Circuit("REPEAT 10{\nH 0 1 2\n}"), [1])
 
-    ScheduledCircuit([], 0, {})
+    ScheduledCircuit([], 0, QubitMap())
     moments = [Moment(stim.Circuit("H 0 1 2")), Moment(stim.Circuit("H 0 1 2"))]
-    i2q = {i: GridQubit(i, i) for i in range(3)}
+    i2q = QubitMap({i: GridQubit(i, i) for i in range(3)})
     ScheduledCircuit(moments, 0, i2q)
     ScheduledCircuit(moments, [0, 1], i2q)
     ScheduledCircuit(moments, Schedule.from_offsets([0, 1]), i2q)
@@ -159,7 +160,7 @@ def test_scheduled_circuit_schedule_property() -> None:
         stim.Circuit("H 0 1 2\nTICK\nH 0 1 2"), [4, 56]
     ).schedule.schedule == [4, 56]
     moments = [Moment(stim.Circuit("H 0 1 2")), Moment(stim.Circuit("H 0 1 2"))]
-    i2q = {i: GridQubit(i, i) for i in range(3)}
+    i2q = QubitMap({i: GridQubit(i, i) for i in range(3)})
     assert ScheduledCircuit(moments, 0, i2q).schedule.schedule == [0, 1]
 
 
@@ -355,7 +356,7 @@ def test_relabel_circuits_qubit_indices() -> None:
             ScheduledCircuit.from_circuit(stim.Circuit("QUBIT_COORDS(1, 1) 0\nX 0")),
         ]
     )
-    assert q2i == {GridQubit(0, 0): 0, GridQubit(1, 1): 1}
+    assert q2i == QubitMap({0: GridQubit(0, 0), 1: GridQubit(1, 1)})
     assert len(circuits) == 2
     assert circuits[0].get_circuit() == stim.Circuit("QUBIT_COORDS(0, 0) 0\nH 0")
     assert circuits[1].get_circuit() == stim.Circuit("QUBIT_COORDS(1, 1) 1\nX 1")