Add coordinate support for nodes and QUBIT_COORDS in stim output

graphqomb/command.py

@@ -23,17 +23,22 @@
 
 @dataclasses.dataclass
 class N:
-    """Preparation command.
+    r"""Preparation command.
 
     Attributes
     ----------
     node : `int`
         The node index to be prepared.
+    coordinate : `tuple`\[`float`, ...\] | `None`
+        Optional coordinate for the node (2D or 3D).
     """
 
     node: int
+    coordinate: tuple[float, ...] | None = None
 
     def __str__(self) -> str:
+        if self.coordinate is not None:
+            return f"N: node={self.node}, coord={self.coordinate}"
         return f"N: node={self.node}"
 
 

graphqomb/graphstate.py

@@ -93,13 +93,18 @@ def meas_bases(self) -> dict[int, MeasBasis]:
         """
 
     @abc.abstractmethod
-    def add_physical_node(self) -> int:
-        """Add a physical node to the graph state.
+    def add_physical_node(self, coordinate: tuple[float, ...] | None = None) -> int:
+        r"""Add a physical node to the graph state.
+
+        Parameters
+        ----------
+        coordinate : `tuple`\[`float`, ...\] | `None`, optional
+            coordinate tuple (2D or 3D), by default None
 
         Returns
         -------
         `int`
-            The node index intenally generated
+            The node index internally generated
         """
 
     @abc.abstractmethod
@@ -169,6 +174,17 @@ def neighbors(self, node: int) -> set[int]:
     def check_canonical_form(self) -> None:
         r"""Check if the graph state is in canonical form."""
 
+    @property
+    @abc.abstractmethod
+    def coordinates(self) -> dict[int, tuple[float, ...]]:
+        r"""Return node coordinates.
+
+        Returns
+        -------
+        `dict`\[`int`, `tuple`\[`float`, ...\]\]
+            mapping from node index to coordinate tuple (2D or 3D)
+        """
+
 
 class GraphState(BaseGraphState):
     """Minimal implementation of GraphState."""
@@ -179,6 +195,7 @@ class GraphState(BaseGraphState):
     __physical_edges: dict[int, set[int]]
     __meas_bases: dict[int, MeasBasis]
     __local_cliffords: dict[int, LocalClifford]
+    __coordinates: dict[int, tuple[float, ...]]
 
     __node_counter: int
 
@@ -189,6 +206,7 @@ def __init__(self) -> None:
         self.__physical_edges = {}
         self.__meas_bases = {}
         self.__local_cliffords = {}
+        self.__coordinates = {}
 
         self.__node_counter = 0
 
@@ -268,6 +286,31 @@ def local_cliffords(self) -> dict[int, LocalClifford]:
         """
         return self.__local_cliffords.copy()
 
+    @property
+    @typing_extensions.override
+    def coordinates(self) -> dict[int, tuple[float, ...]]:
+        r"""Return node coordinates.
+
+        Returns
+        -------
+        `dict`\[`int`, `tuple`\[`float`, ...\]\]
+            mapping from node index to coordinate tuple (2D or 3D)
+        """
+        return self.__coordinates.copy()
+
+    def set_coordinate(self, node: int, coord: tuple[float, ...]) -> None:
+        r"""Set coordinate for a node.
+
+        Parameters
+        ----------
+        node : `int`
+            node index
+        coord : `tuple`\[`float`, ...\]
+            coordinate tuple (2D or 3D)
+        """
+        self._ensure_node_exists(node)
+        self.__coordinates[node] = coord
+
     def _check_meas_basis(self) -> None:
         """Check if the measurement basis is set for all physical nodes except output nodes.
 
@@ -294,8 +337,13 @@ def _ensure_node_exists(self, node: int) -> None:
             raise ValueError(msg)
 
     @typing_extensions.override
-    def add_physical_node(self) -> int:
-        """Add a physical node to the graph state.
+    def add_physical_node(self, coordinate: tuple[float, ...] | None = None) -> int:
+        r"""Add a physical node to the graph state.
+
+        Parameters
+        ----------
+        coordinate : `tuple`\[`float`, ...\] | `None`, optional
+            coordinate tuple (2D or 3D), by default None
 
         Returns
         -------
@@ -305,6 +353,8 @@ def add_physical_node(self) -> int:
         node = self.__node_counter
         self.__physical_nodes |= {node}
         self.__physical_edges[node] = set()
+        if coordinate is not None:
+            self.__coordinates[node] = coordinate
         self.__node_counter += 1
 
         return node
@@ -364,6 +414,7 @@ def remove_physical_node(self, node: int) -> None:
             del self.__output_node_indices[node]
         self.__meas_bases.pop(node, None)
         self.__local_cliffords.pop(node, None)
+        self.__coordinates.pop(node, None)
 
     def remove_physical_edge(self, node1: int, node2: int) -> None:
         """Remove a physical edge from the graph state.
@@ -621,13 +672,14 @@ def _expand_output_local_cliffords(self) -> dict[int, LocalCliffordExpansion]:
         return node_index_addition_map
 
     @classmethod
-    def from_graph(  # noqa: C901, PLR0912
+    def from_graph(  # noqa: C901, PLR0912, PLR0913, PLR0917
         cls,
         nodes: Iterable[NodeT],
         edges: Iterable[tuple[NodeT, NodeT]],
         inputs: Sequence[NodeT] | None = None,
         outputs: Sequence[NodeT] | None = None,
         meas_bases: Mapping[NodeT, MeasBasis] | None = None,
+        coordinates: Mapping[NodeT, tuple[float, ...]] | None = None,
     ) -> tuple[GraphState, dict[NodeT, int]]:
         r"""Create a graph state from nodes and edges with arbitrary node types.
 
@@ -650,6 +702,8 @@ def from_graph(  # noqa: C901, PLR0912
         meas_bases : `collections.abc.Mapping`\[NodeT, `MeasBasis`\] | `None`, optional
             Measurement bases for nodes. Nodes not specified can be set later.
             Default is None (no bases assigned initially).
+        coordinates : `collections.abc.Mapping`\[NodeT, `tuple`\[`float`, ...\]\] | `None`, optional
+            Coordinates for nodes (2D or 3D). Default is None (no coordinates).
 
         Returns
         -------
@@ -729,6 +783,12 @@ def from_graph(  # noqa: C901, PLR0912
                 if node in node_set:
                     graph_state.assign_meas_basis(node_map[node], meas_basis)
 
+        # Assign coordinates
+        if coordinates is not None:
+            for node, coord in coordinates.items():
+                if node in node_set:
+                    graph_state.set_coordinate(node_map[node], coord)
+
         return graph_state, node_map
 
     @classmethod
@@ -790,6 +850,10 @@ def from_base_graph_state(
                 # Access private attribute to copy local cliffords
                 graph_state.apply_local_clifford(node_map[node], lc)
 
+        # Copy coordinates
+        for node, coord in base.coordinates.items():
+            graph_state.set_coordinate(node_map[node], coord)
+
         return graph_state, node_map
 
 
@@ -808,7 +872,7 @@ class ExpansionMaps(NamedTuple):
     output_node_map: dict[int, LocalCliffordExpansion]
 
 
-def compose(  # noqa: C901
+def compose(  # noqa: C901, PLR0912
     graph1: BaseGraphState, graph2: BaseGraphState
 ) -> tuple[GraphState, dict[int, int], dict[int, int]]:
     r"""Compose two graph states sequentially.
@@ -899,6 +963,16 @@ def compose(  # noqa: C901
     for u, v in graph2.physical_edges:
         composed_graph.add_physical_edge(node_map2[u], node_map2[v])
 
+    # Copy coordinates from graph1
+    for node, coord in graph1.coordinates.items():
+        if node in node_map1:
+            composed_graph.set_coordinate(node_map1[node], coord)
+
+    # Copy coordinates from graph2
+    for node, coord in graph2.coordinates.items():
+        if node in node_map2:
+            composed_graph.set_coordinate(node_map2[node], coord)
+
     return composed_graph, node_map1, node_map2
 
 

graphqomb/pattern.py

@@ -38,12 +38,15 @@ class Pattern(Sequence[Command]):
         Commands of the pattern
     pauli_frame : `PauliFrame`
         Pauli frame of the pattern to track the Pauli state of each node
+    input_coordinates : `dict`\[`int`, `tuple`\[`float`, ...\]\]
+        Coordinates for input nodes (2D or 3D)
     """
 
     input_node_indices: dict[int, int]
     output_node_indices: dict[int, int]
     commands: tuple[Command, ...]
     pauli_frame: PauliFrame
+    input_coordinates: dict[int, tuple[float, ...]] = dataclasses.field(default_factory=dict)
 
     def __len__(self) -> int:
         return len(self.commands)
@@ -101,6 +104,21 @@ def depth(self) -> int:
         """
         return sum(1 for cmd in self.commands if isinstance(cmd, TICK))
 
+    @functools.cached_property
+    def coordinates(self) -> dict[int, tuple[float, ...]]:
+        r"""Get all node coordinates from N commands and input coordinates.
+
+        Returns
+        -------
+        `dict`\[`int`, `tuple`\[`float`, ...\]\]
+            mapping from node index to coordinate tuple (2D or 3D)
+        """
+        coords = dict(self.input_coordinates)
+        for cmd in self.commands:
+            if isinstance(cmd, N) and cmd.coordinate is not None:
+                coords[cmd.node] = cmd.coordinate
+        return coords
+
     @property
     def active_volume(self) -> int:
         """Calculate tha active volume, summation of space for each timeslice.

graphqomb/qompiler.py

@@ -94,6 +94,7 @@ def _qompile(
         compiled pattern
     """
     meas_bases = graph.meas_bases
+    graph_coords = graph.coordinates
 
     dag = dag_from_flow(graph, xflow=pauli_frame.xflow, zflow=pauli_frame.zflow)
     topo_order = list(TopologicalSorter(dag).static_order())
@@ -110,7 +111,10 @@ def _qompile(
         prepare_nodes, entangle_edges, measure_nodes = timeline[time_idx]
 
         # Order within time slice: N -> E -> M
-        commands.extend(N(node) for node in prepare_nodes)
+        # N commands include coordinates if available
+        for node in prepare_nodes:
+            coord = graph_coords.get(node)
+            commands.append(N(node, coordinate=coord))
         for edge in entangle_edges:
             a, b = edge
             commands.append(E(nodes=(a, b)))
@@ -125,9 +129,13 @@ def _qompile(
         else:
             commands.extend((X(node=node), Z(node=node)))
 
+    # Collect input node coordinates
+    input_coords = {node: graph_coords[node] for node in graph.input_node_indices if node in graph_coords}
+
     return Pattern(
         input_node_indices=graph.input_node_indices,
         output_node_indices=graph.output_node_indices,
         commands=tuple(commands),
         pauli_frame=pauli_frame,
+        input_coordinates=input_coords,
     )