Add dict utilities to Layout (#138)

MarcSerraPeralta · Nov 3, 2024 · 08370da · 08370da
1 parent 645cfeb
commit 08370da
Show file tree

Hide file tree

Showing 12 changed files with 95 additions and 142 deletions.
diff --git a/README.md b/README.md
@@ -34,7 +34,7 @@ from surface_sim.experiments.rot_surface_code_css import memory_experiment
 # prepare the layout, model, and detectors objects
 layout = rot_surface_code(distance=3)
 
-qubit_inds = {q: layout.get_inds([q])[0] for q in layout.get_qubits()}
+qubit_inds = layout.qubit_inds()
 anc_qubits = layout.get_qubits(role="anc")
 data_qubits = layout.get_qubits(role="data")
 

diff --git a/surface_sim/circuit_blocks/util.py b/surface_sim/circuit_blocks/util.py
@@ -16,7 +16,7 @@ def qubit_coords(model: Model, layout: Layout) -> Circuit:
     """Returns a stim circuit that sets up the coordinates
     of the qubits.
     """
-    coord_dict = {q: layout.get_coords([q])[0] for q in layout.get_qubits()}
+    coord_dict = layout.qubit_coords()
     circuit = Circuit()
 
     circuit += model.qubit_coords(coord_dict)

diff --git a/surface_sim/layouts/layout.py b/surface_sim/layouts/layout.py
@@ -29,12 +29,15 @@ class Layout:
     ---------------
     - ``param``
     - ``get_inds``
+    - ``qubit_inds``
     - ``get_max_ind``
     - ``get_min_ind``
     - ``get_qubits``
     - ``get_logical_qubits``
     - ``get_neighbors``
     - ``get_coords``
+    - ``qubit_coords``
+    - ``anc_coords``
 
     Set information
     ---------------
@@ -86,6 +89,7 @@ def __init__(self, setup: dict[str, object]) -> None:
         self.log_x = setup.get("log_x", {})
         self.description = setup.get("description")
         self.interaction_order = setup.get("interaction_order", {})
+        self._qubit_inds = {}
 
         self.graph = nx.DiGraph()
         self._load_layout(setup)
@@ -155,6 +159,10 @@ def get_inds(self, qubits: Iterable[str]) -> list[int]:
         """
         return [self._qubit_inds[qubit] for qubit in qubits]
 
+    def qubit_inds(self) -> dict[str, int]:
+        """Returns a dictionary mapping all the qubits to their indices."""
+        return {k: v for k, v in self._qubit_inds.items()}
+
     def get_max_ind(self) -> int:
         """Returns the largest qubit index in the layout."""
         return max(self._qubit_inds.values())
@@ -258,6 +266,16 @@ def get_coords(self, qubits: Iterable[str]) -> list[list[float | int]]:
 
         return [all_coords[q] for q in qubits]
 
+    def qubit_coords(self) -> dict[str, list[float | int]]:
+        """Returns a dictionary mapping all the qubits to their coordinates."""
+        return nx.get_node_attributes(self.graph, "coords")
+
+    def anc_coords(self) -> dict[str, list[float | int]]:
+        """Returns a dictionary mapping all ancilla qubits to their coordinates."""
+        anc_qubits = self.get_qubits(role="anc")
+        anc_coords = self.get_coords(anc_qubits)
+        return {q: c for q, c in zip(anc_qubits, anc_coords)}
+
     def adjacency_matrix(self) -> DataArray:
         """Returns the adjaceny matrix corresponding to the layout.
 

diff --git a/tests/experiments/test_rot_surface_code_css.py b/tests/experiments/test_rot_surface_code_css.py
@@ -14,11 +14,9 @@
 
 def test_memory_experiment():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
 
     for rot_basis in [True, False]:
@@ -39,7 +37,7 @@ def test_memory_experiment():
         dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
         num_coords = 0
-        anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+        anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
         for dem_instr in dem:
             if dem_instr.type == "detector":
                 assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -53,11 +51,9 @@ def test_memory_experiment():
 def test_repeated_s_experiment():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
 
     for rot_basis in [True, False]:
@@ -79,7 +75,7 @@ def test_repeated_s_experiment():
         dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
         num_coords = 0
-        anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+        anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
         for dem_instr in dem:
             if dem_instr.type == "detector":
                 assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -103,14 +99,10 @@ def test_repeated_cnot_experiment():
     )
     set_trans_cnot(layout_c, layout_t)
     set_trans_cnot(layout_t, layout_c)
-    qubit_inds = {q: layout_c.get_inds([q])[0] for q in layout_c.get_qubits()}
-    qubit_inds.update({q: layout_t.get_inds([q])[0] for q in layout_t.get_qubits()})
-    anc_coords = {
-        q: layout_c.get_coords([q])[0] for q in layout_c.get_qubits(role="anc")
-    }
-    anc_coords.update(
-        {q: layout_t.get_coords([q])[0] for q in layout_t.get_qubits(role="anc")}
-    )
+    qubit_inds = layout_c.qubit_inds()
+    qubit_inds.update(layout_t.qubit_inds())
+    anc_coords = layout_c.anc_coords()
+    anc_coords.update(layout_t.anc_coords())
     model = NoiselessModel(qubit_inds)
     detectors = Detectors(
         layout_c.get_qubits(role="anc") + layout_t.get_qubits(role="anc"),
@@ -156,8 +148,7 @@ def test_repeated_cnot_experiment():
 
 def test_memory_experiment_anc_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -187,8 +178,7 @@ def test_memory_experiment_anc_detectors():
 def test_repeated_s_experiment_anc_detectors():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = repeated_s_experiment(
         model=model,
@@ -229,14 +219,10 @@ def test_repeated_cnot_experiment_anc_detectors():
     )
     set_trans_cnot(layout_c, layout_t)
     set_trans_cnot(layout_t, layout_c)
-    qubit_inds = {q: layout_c.get_inds([q])[0] for q in layout_c.get_qubits()}
-    qubit_inds.update({q: layout_t.get_inds([q])[0] for q in layout_t.get_qubits()})
-    anc_coords = {
-        q: layout_c.get_coords([q])[0] for q in layout_c.get_qubits(role="anc")
-    }
-    anc_coords.update(
-        {q: layout_t.get_coords([q])[0] for q in layout_t.get_qubits(role="anc")}
-    )
+    qubit_inds = layout_c.qubit_inds()
+    qubit_inds.update(layout_t.qubit_inds())
+    anc_coords = layout_c.anc_coords()
+    anc_coords.update(layout_t.anc_coords())
     model = NoiselessModel(qubit_inds)
     detectors = Detectors(
         layout_c.get_qubits(role="anc") + layout_t.get_qubits(role="anc"),
@@ -280,8 +266,7 @@ def test_repeated_cnot_experiment_anc_detectors():
 
 def test_memory_experiment_gauge_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -311,8 +296,7 @@ def test_memory_experiment_gauge_detectors():
 def test_repeated_s_experiment_gauge_detectors():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = repeated_s_experiment(
         model=model,
@@ -353,14 +337,10 @@ def test_repeated_cnot_experiment_gauge_detectors():
     )
     set_trans_cnot(layout_c, layout_t)
     set_trans_cnot(layout_t, layout_c)
-    qubit_inds = {q: layout_c.get_inds([q])[0] for q in layout_c.get_qubits()}
-    qubit_inds.update({q: layout_t.get_inds([q])[0] for q in layout_t.get_qubits()})
-    anc_coords = {
-        q: layout_c.get_coords([q])[0] for q in layout_c.get_qubits(role="anc")
-    }
-    anc_coords.update(
-        {q: layout_t.get_coords([q])[0] for q in layout_t.get_qubits(role="anc")}
-    )
+    qubit_inds = layout_c.qubit_inds()
+    qubit_inds.update(layout_t.qubit_inds())
+    anc_coords = layout_c.anc_coords()
+    anc_coords.update(layout_t.anc_coords())
     model = NoiselessModel(qubit_inds)
     detectors = Detectors(
         layout_c.get_qubits(role="anc") + layout_t.get_qubits(role="anc"),

diff --git a/tests/experiments/test_rot_surface_code_css_pipelined.py b/tests/experiments/test_rot_surface_code_css_pipelined.py
@@ -13,11 +13,9 @@
 
 def test_memory_experiment():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
     circuit = memory_experiment(
         model=model,
@@ -36,7 +34,7 @@ def test_memory_experiment():
     dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
     num_coords = 0
-    anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+    anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
     for dem_instr in dem:
         if dem_instr.type == "detector":
             assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -50,11 +48,9 @@ def test_memory_experiment():
 def test_repeated_s_experiment():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
     circuit = repeated_s_experiment(
         model=model,
@@ -74,7 +70,7 @@ def test_repeated_s_experiment():
     dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
     num_coords = 0
-    anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+    anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
     for dem_instr in dem:
         if dem_instr.type == "detector":
             assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -87,8 +83,7 @@ def test_repeated_s_experiment():
 
 def test_memory_experiment_anc_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -118,8 +113,7 @@ def test_memory_experiment_anc_detectors():
 def test_repeated_s_experiment_anc_detectors():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = repeated_s_experiment(
         model=model,
@@ -149,8 +143,7 @@ def test_repeated_s_experiment_anc_detectors():
 
 def test_memory_experiment_gauge_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -180,8 +173,7 @@ def test_memory_experiment_gauge_detectors():
 def test_repeated_s_experiment_gauge_detectors():
     layout = rot_surface_code_rectangle(distance_z=4, distance_x=3)
     set_trans_s(layout, "D1")
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = repeated_s_experiment(
         model=model,

diff --git a/tests/experiments/test_rot_surface_code_xzzx.py b/tests/experiments/test_rot_surface_code_xzzx.py
@@ -9,11 +9,9 @@
 
 def test_memory_experiment():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
     circuit = memory_experiment(
         model=model,
@@ -32,7 +30,7 @@ def test_memory_experiment():
     dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
     num_coords = 0
-    anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+    anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
     for dem_instr in dem:
         if dem_instr.type == "detector":
             assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -45,8 +43,7 @@ def test_memory_experiment():
 
 def test_memory_experiment_anc_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -75,8 +72,7 @@ def test_memory_experiment_anc_detectors():
 
 def test_memory_experiment_gauge_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,

diff --git a/tests/experiments/test_rot_surface_code_xzzx_google.py b/tests/experiments/test_rot_surface_code_xzzx_google.py
@@ -9,11 +9,9 @@
 
 def test_memory_experiment():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    anc_coords = {q: layout.get_coords([q])[0] for q in layout.get_qubits(role="anc")}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(
-        layout.get_qubits(role="anc"), frame="1", anc_coords=anc_coords
+        layout.get_qubits(role="anc"), frame="1", anc_coords=layout.anc_coords()
     )
     circuit = memory_experiment(
         model=model,
@@ -31,7 +29,7 @@ def test_memory_experiment():
     dem = circuit.detector_error_model(allow_gauge_detectors=True)
 
     num_coords = 0
-    anc_coords = {k: list(map(float, v)) for k, v in anc_coords.items()}
+    anc_coords = {k: list(map(float, v)) for k, v in layout.anc_coords().items()}
     for dem_instr in dem:
         if dem_instr.type == "detector":
             assert dem_instr.args_copy()[:-1] in anc_coords.values()
@@ -44,8 +42,7 @@ def test_memory_experiment():
 
 def test_memory_experiment_anc_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,
@@ -73,8 +70,7 @@ def test_memory_experiment_anc_detectors():
 
 def test_memory_experiment_gauge_detectors():
     layout = rot_surface_code(distance=3)
-    qubit_ids = {q: i for i, q in enumerate(layout.get_qubits())}
-    model = NoiselessModel(qubit_ids)
+    model = NoiselessModel(layout.qubit_inds())
     detectors = Detectors(layout.get_qubits(role="anc"), frame="1")
     circuit = memory_experiment(
         model=model,