From 945aa8bf05e5cbc17b6307b1d90199a124c8b58b Mon Sep 17 00:00:00 2001
From: Masato Fukushima <masa1063fuk@gmail.com>
Date: Fri, 2 Jan 2026 00:09:07 +0900
Subject: [PATCH 1/6] Add coordinate support for nodes and QUBIT_COORDS in stim
 output
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This adds the ability to store 2D/3D coordinates for nodes in GraphState
and emit QUBIT_COORDS instructions in stim compiler output.

Changes:
- Add coordinates property and set_coordinate method to GraphState
- Add coordinate parameter to add_physical_node and from_graph
- Add coordinate field to N command
- Add input_coordinates field and coordinates property to Pattern
- Inherit coordinates from GraphState to Pattern via qompiler
- Emit QUBIT_COORDS instructions in stim_compile (with emit_qubit_coords option)
- Add use_graph_coordinates option to visualizer
- Add comprehensive tests for coordinate functionality

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 graphqomb/command.py        |  7 ++-
 graphqomb/graphstate.py     | 79 ++++++++++++++++++++++++++++++--
 graphqomb/pattern.py        | 20 +++++++++
 graphqomb/qompiler.py       | 14 +++++-
 graphqomb/stim_compiler.py  | 60 +++++++++++++++++++++++--
 graphqomb/visualizer.py     | 50 ++++++++++++++++++++-
 tests/test_graphstate.py    | 69 ++++++++++++++++++++++++++++
 tests/test_stim_compiler.py | 89 +++++++++++++++++++++++++++++++++++++
 8 files changed, 377 insertions(+), 11 deletions(-)

diff --git a/graphqomb/command.py b/graphqomb/command.py
index 56ae7c901..c193b42d7 100644
--- a/graphqomb/command.py
+++ b/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}"
 
 
diff --git a/graphqomb/graphstate.py b/graphqomb/graphstate.py
index 4253ac542..e8a150554 100644
--- a/graphqomb/graphstate.py
+++ b/graphqomb/graphstate.py
@@ -169,6 +169,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 +190,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 +201,7 @@ def __init__(self) -> None:
         self.__physical_edges = {}
         self.__meas_bases = {}
         self.__local_cliffords = {}
+        self.__coordinates = {}
 
         self.__node_counter = 0
 
@@ -268,6 +281,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 +332,15 @@ 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 +350,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 +411,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 +669,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 +699,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 +780,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 +847,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 +869,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 +960,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
 
 
diff --git a/graphqomb/pattern.py b/graphqomb/pattern.py
index b9718da41..13e8b08e5 100644
--- a/graphqomb/pattern.py
+++ b/graphqomb/pattern.py
@@ -38,12 +38,17 @@ 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 +106,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
+
 
 def is_runnable(pattern: Pattern) -> None:
     """Check if the pattern is runnable.
diff --git a/graphqomb/qompiler.py b/graphqomb/qompiler.py
index cb6a708cb..16c42d9e8 100644
--- a/graphqomb/qompiler.py
+++ b/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,17 @@ 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,
     )
diff --git a/graphqomb/stim_compiler.py b/graphqomb/stim_compiler.py
index ea1f8970d..6c26478ca 100644
--- a/graphqomb/stim_compiler.py
+++ b/graphqomb/stim_compiler.py
@@ -22,10 +22,33 @@
     from graphqomb.pauli_frame import PauliFrame
 
 
+def _emit_qubit_coords(
+    stim_io: StringIO,
+    node: int,
+    coordinate: tuple[float, ...] | None,
+) -> None:
+    r"""Emit QUBIT_COORDS instruction if coordinate is available.
+
+    Parameters
+    ----------
+    stim_io : `StringIO`
+        The output stream to write to.
+    node : `int`
+        The qubit index.
+    coordinate : `tuple`\[`float`, ...\] | `None`
+        The coordinate tuple (2D or 3D), or None if no coordinate.
+    """
+    if coordinate is not None:
+        coords_str = ", ".join(str(c) for c in coordinate)
+        stim_io.write(f"QUBIT_COORDS({coords_str}) {node}\n")
+
+
 def _initialize_nodes(
     stim_io: StringIO,
     node_indices: Mapping[int, int],
     p_depol_after_clifford: float,
+    input_coordinates: Mapping[int, tuple[float, ...]] | None = None,
+    emit_qubit_coords: bool = True,
 ) -> None:
     r"""Initialize nodes in the stim circuit.
 
@@ -37,8 +60,15 @@ def _initialize_nodes(
         The node indices mapping to initialize.
     p_depol_after_clifford : `float`
         The probability of depolarization after Clifford gates.
+    input_coordinates : `collections.abc.Mapping`\[`int`, `tuple`\[`float`, ...\]\] | `None`, optional
+        Coordinates for input nodes, by default None.
+    emit_qubit_coords : `bool`, optional
+        Whether to emit QUBIT_COORDS instructions, by default True.
     """
     for node in node_indices:
+        if emit_qubit_coords and input_coordinates is not None:
+            coord = input_coordinates.get(node)
+            _emit_qubit_coords(stim_io, node, coord)
         stim_io.write(f"RX {node}\n")
         if p_depol_after_clifford > 0.0:
             stim_io.write(f"DEPOLARIZE1({p_depol_after_clifford}) {node}\n")
@@ -48,6 +78,8 @@ def _prepare_node(
     stim_io: StringIO,
     node: int,
     p_depol_after_clifford: float,
+    coordinate: tuple[float, ...] | None = None,
+    emit_qubit_coords: bool = True,
 ) -> None:
     r"""Prepare a node in |+> state (N command).
 
@@ -59,7 +91,13 @@ def _prepare_node(
         The node to prepare in |+> state.
     p_depol_after_clifford : `float`
         The probability of depolarization after Clifford gates.
+    coordinate : `tuple`\[`float`, ...\] | `None`, optional
+        The coordinate tuple (2D or 3D), by default None.
+    emit_qubit_coords : `bool`, optional
+        Whether to emit QUBIT_COORDS instruction, by default True.
     """
+    if emit_qubit_coords:
+        _emit_qubit_coords(stim_io, node, coordinate)
     stim_io.write(f"RX {node}\n")
     if p_depol_after_clifford > 0.0:
         stim_io.write(f"DEPOLARIZE1({p_depol_after_clifford}) {node}\n")
@@ -191,6 +229,7 @@ def stim_compile(
     *,
     p_depol_after_clifford: float = 0.0,
     p_before_meas_flip: float = 0.0,
+    emit_qubit_coords: bool = True,
 ) -> str:
     r"""Compile a pattern to stim format.
 
@@ -204,6 +243,9 @@ def stim_compile(
         The probability of depolarization after a Clifford gate, by default 0.0.
     p_before_meas_flip : `float`, optional
         The probability of flipping a measurement result before measurement, by default 0.0.
+    emit_qubit_coords : `bool`, optional
+        Whether to emit QUBIT_COORDS instructions for nodes with coordinates,
+        by default True.
 
     Returns
     -------
@@ -228,13 +270,25 @@ def stim_compile(
             meas_idx += 1
     total_measurements = meas_idx
 
-    # Initialize input nodes
-    _initialize_nodes(stim_io, pattern.input_node_indices, p_depol_after_clifford)
+    # Initialize input nodes (with coordinates if available)
+    _initialize_nodes(
+        stim_io,
+        pattern.input_node_indices,
+        p_depol_after_clifford,
+        input_coordinates=pattern.input_coordinates if emit_qubit_coords else None,
+        emit_qubit_coords=emit_qubit_coords,
+    )
 
     # Process pattern commands
     for cmd in pattern:
         if isinstance(cmd, N):
-            _prepare_node(stim_io, cmd.node, p_depol_after_clifford)
+            _prepare_node(
+                stim_io,
+                cmd.node,
+                p_depol_after_clifford,
+                coordinate=cmd.coordinate,
+                emit_qubit_coords=emit_qubit_coords,
+            )
         elif isinstance(cmd, E):
             _entangle_nodes(stim_io, cmd.nodes, p_depol_after_clifford)
         elif isinstance(cmd, M):
diff --git a/graphqomb/visualizer.py b/graphqomb/visualizer.py
index 65576ce34..aaeb1eb33 100644
--- a/graphqomb/visualizer.py
+++ b/graphqomb/visualizer.py
@@ -61,13 +61,14 @@ class FigureSetup(NamedTuple):
     fig_height: float
 
 
-def visualize(
+def visualize(  # noqa: PLR0913
     graph: BaseGraphState,
     *,
     ax: Axes | None = None,
     show_node_labels: bool = True,
     node_size: float = 300,
     show_legend: bool = True,
+    use_graph_coordinates: bool = True,
 ) -> Axes:
     r"""Visualize the GraphState.
 
@@ -83,13 +84,18 @@ def visualize(
         Size of nodes (scatter size), by default 300
     show_legend : `bool`, optional
         Whether to show color legend, by default True
+    use_graph_coordinates : `bool`, optional
+        Whether to use coordinates stored in the graph. If True and the graph
+        has coordinates, those coordinates are used (projected to 2D for 3D
+        coordinates). Nodes without coordinates will use auto-calculated
+        positions. By default True.
 
     Returns
     -------
     `matplotlib.axes.Axes`
         The Axes object containing the visualization
     """
-    node_pos = _calc_node_positions(graph)
+    node_pos = _get_node_positions(graph, use_graph_coordinates)
 
     node_colors = _determine_node_colors(graph)
 
@@ -223,6 +229,46 @@ def _setup_figure(node_pos: Mapping[int, tuple[float, float]]) -> FigureSetup:
     )
 
 
+def _get_node_positions(
+    graph: BaseGraphState,
+    use_graph_coordinates: bool,
+) -> dict[int, tuple[float, float]]:
+    """Get node positions, using graph coordinates if available and requested.
+
+    Parameters
+    ----------
+    graph : BaseGraphState
+        GraphState to visualize.
+    use_graph_coordinates : bool
+        Whether to use coordinates stored in the graph.
+
+    Returns
+    -------
+    dict[int, tuple[float, float]]
+        Mapping of node indices to their (x, y) positions.
+    """
+    if use_graph_coordinates and graph.coordinates:
+        # Use graph coordinates (project 3D to 2D by using x, y only)
+        node_pos: dict[int, tuple[float, float]] = {}
+        for node, coord in graph.coordinates.items():
+            # Take first two coordinates for 2D projection
+            node_pos[node] = (coord[0], coord[1])
+
+        # For nodes without coordinates, calculate positions
+        missing_nodes = graph.physical_nodes - node_pos.keys()
+        if missing_nodes:
+            # Calculate auto positions for all nodes
+            auto_pos = _calc_node_positions(graph)
+            # Add only the missing nodes' positions
+            for node in missing_nodes:
+                node_pos[node] = auto_pos[node]
+
+        return node_pos
+
+    # Fall back to auto-calculated positions
+    return _calc_node_positions(graph)
+
+
 def _calc_node_positions(graph: BaseGraphState) -> dict[int, tuple[float, float]]:
     """Calculate node positions for visualization with input/output nodes arranged vertically.
 
diff --git a/tests/test_graphstate.py b/tests/test_graphstate.py
index e649ff916..57fdf0f7c 100644
--- a/tests/test_graphstate.py
+++ b/tests/test_graphstate.py
@@ -285,3 +285,72 @@ def test_odd_neighbors(graph: GraphState) -> None:
 
     assert odd_neighbors({node1}, graph) == {node2, node3}
     assert odd_neighbors({node1, node2}, graph) == {node1, node2}
+
+
+# ---- Coordinate Tests ----
+
+
+def test_set_coordinate(graph: GraphState) -> None:
+    """Test setting coordinates for a node."""
+    node = graph.add_physical_node()
+    graph.set_coordinate(node, (1.0, 2.0))
+    assert graph.coordinates == {node: (1.0, 2.0)}
+
+
+def test_set_coordinate_3d(graph: GraphState) -> None:
+    """Test setting 3D coordinates for a node."""
+    node = graph.add_physical_node()
+    graph.set_coordinate(node, (1.0, 2.0, 3.0))
+    assert graph.coordinates == {node: (1.0, 2.0, 3.0)}
+
+
+def test_add_physical_node_with_coordinate() -> None:
+    """Test adding a node with coordinates."""
+    graph = GraphState()
+    node = graph.add_physical_node(coordinate=(1.5, 2.5))
+    assert graph.coordinates == {node: (1.5, 2.5)}
+
+
+def test_remove_physical_node_removes_coordinate() -> None:
+    """Test that removing a node also removes its coordinate."""
+    graph = GraphState()
+    node1 = graph.add_physical_node(coordinate=(1.0, 2.0))
+    node2 = graph.add_physical_node(coordinate=(3.0, 4.0))
+    graph.add_physical_edge(node1, node2)
+    graph.register_output(node2, 0)
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+
+    graph.remove_physical_node(node1)
+    assert node1 not in graph.coordinates
+    assert graph.coordinates == {node2: (3.0, 4.0)}
+
+
+def test_from_graph_with_coordinates() -> None:
+    """Test from_graph with coordinates parameter."""
+    nodes = ["a", "b", "c"]
+    edges = [("a", "b"), ("b", "c")]
+    coordinates = {"a": (0.0, 0.0), "b": (1.0, 0.0), "c": (2.0, 0.0)}
+
+    graph, node_map = GraphState.from_graph(
+        nodes, edges, inputs=["a"], outputs=["c"], coordinates=coordinates
+    )
+
+    assert graph.coordinates[node_map["a"]] == (0.0, 0.0)
+    assert graph.coordinates[node_map["b"]] == (1.0, 0.0)
+    assert graph.coordinates[node_map["c"]] == (2.0, 0.0)
+
+
+def test_from_base_graph_state_copies_coordinates() -> None:
+    """Test that from_base_graph_state copies coordinates."""
+    graph1 = GraphState()
+    node1 = graph1.add_physical_node(coordinate=(1.0, 2.0))
+    node2 = graph1.add_physical_node(coordinate=(3.0, 4.0))
+    graph1.add_physical_edge(node1, node2)
+    graph1.register_input(node1, 0)
+    graph1.register_output(node2, 0)
+    graph1.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+
+    graph2, node_map = GraphState.from_base_graph_state(graph1)
+
+    assert graph2.coordinates[node_map[node1]] == (1.0, 2.0)
+    assert graph2.coordinates[node_map[node2]] == (3.0, 4.0)
diff --git a/tests/test_stim_compiler.py b/tests/test_stim_compiler.py
index 996fae3ce..7c3d783f1 100644
--- a/tests/test_stim_compiler.py
+++ b/tests/test_stim_compiler.py
@@ -459,3 +459,92 @@ def test_stim_compile_respects_manual_entangle_time() -> None:
 
     assert cz_slice[in_node, mid_node] == 2
     assert cz_slice[mid_node, out_node] == 1
+
+
+# ---- Coordinate Tests ----
+
+
+def test_stim_compile_with_coordinates() -> None:
+    """Test that QUBIT_COORDS instructions are emitted for nodes with coordinates."""
+    graph = GraphState()
+    in_node = graph.add_physical_node(coordinate=(0.0, 0.0))
+    mid_node = graph.add_physical_node(coordinate=(1.0, 0.0))
+    out_node = graph.add_physical_node(coordinate=(2.0, 0.0))
+
+    graph.register_input(in_node, 0)
+    graph.register_output(out_node, 0)
+
+    graph.add_physical_edge(in_node, mid_node)
+    graph.add_physical_edge(mid_node, out_node)
+
+    graph.assign_meas_basis(in_node, PlannerMeasBasis(Plane.XY, 0.0))
+    graph.assign_meas_basis(mid_node, PlannerMeasBasis(Plane.XY, 0.0))
+
+    pattern = qompile(graph, {in_node: {mid_node}, mid_node: {out_node}})
+    stim_str = stim_compile(pattern)
+
+    # Check QUBIT_COORDS instructions are present
+    assert f"QUBIT_COORDS(0.0, 0.0) {in_node}" in stim_str
+    assert f"QUBIT_COORDS(1.0, 0.0) {mid_node}" in stim_str
+    assert f"QUBIT_COORDS(2.0, 0.0) {out_node}" in stim_str
+
+
+def test_stim_compile_with_3d_coordinates() -> None:
+    """Test that 3D coordinates are correctly emitted."""
+    graph = GraphState()
+    in_node = graph.add_physical_node(coordinate=(0.0, 0.0, 0.0))
+    out_node = graph.add_physical_node(coordinate=(1.0, 1.0, 1.0))
+
+    graph.register_input(in_node, 0)
+    graph.register_output(out_node, 0)
+
+    graph.add_physical_edge(in_node, out_node)
+    graph.assign_meas_basis(in_node, PlannerMeasBasis(Plane.XY, 0.0))
+
+    pattern = qompile(graph, {in_node: {out_node}})
+    stim_str = stim_compile(pattern)
+
+    assert f"QUBIT_COORDS(0.0, 0.0, 0.0) {in_node}" in stim_str
+    assert f"QUBIT_COORDS(1.0, 1.0, 1.0) {out_node}" in stim_str
+
+
+def test_stim_compile_without_coordinates() -> None:
+    """Test that no QUBIT_COORDS are emitted when emit_qubit_coords is False."""
+    graph = GraphState()
+    in_node = graph.add_physical_node(coordinate=(0.0, 0.0))
+    out_node = graph.add_physical_node(coordinate=(1.0, 0.0))
+
+    graph.register_input(in_node, 0)
+    graph.register_output(out_node, 0)
+
+    graph.add_physical_edge(in_node, out_node)
+    graph.assign_meas_basis(in_node, PlannerMeasBasis(Plane.XY, 0.0))
+
+    pattern = qompile(graph, {in_node: {out_node}})
+    stim_str = stim_compile(pattern, emit_qubit_coords=False)
+
+    assert "QUBIT_COORDS" not in stim_str
+
+
+def test_pattern_coordinates_property() -> None:
+    """Test that Pattern.coordinates aggregates coordinates from N commands and input nodes."""
+    graph = GraphState()
+    in_node = graph.add_physical_node(coordinate=(0.0, 0.0))
+    mid_node = graph.add_physical_node(coordinate=(1.0, 0.0))
+    out_node = graph.add_physical_node(coordinate=(2.0, 0.0))
+
+    graph.register_input(in_node, 0)
+    graph.register_output(out_node, 0)
+
+    graph.add_physical_edge(in_node, mid_node)
+    graph.add_physical_edge(mid_node, out_node)
+
+    graph.assign_meas_basis(in_node, PlannerMeasBasis(Plane.XY, 0.0))
+    graph.assign_meas_basis(mid_node, PlannerMeasBasis(Plane.XY, 0.0))
+
+    pattern = qompile(graph, {in_node: {mid_node}, mid_node: {out_node}})
+
+    # Check pattern coordinates
+    assert pattern.coordinates[in_node] == (0.0, 0.0)
+    assert pattern.coordinates[mid_node] == (1.0, 0.0)
+    assert pattern.coordinates[out_node] == (2.0, 0.0)

From 6e48537240778fbcbf911616e2e2b986f4bdcb75 Mon Sep 17 00:00:00 2001
From: Masato Fukushima <masa1063fuk@gmail.com>
Date: Fri, 2 Jan 2026 00:34:58 +0900
Subject: [PATCH 2/6] Fix ruff format and improve test coverage for coordinates
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Apply ruff format to fix formatting issues
- Add test_compose_copies_coordinates to verify compose function
  copies coordinates from both graphs
- Add tests/test_visualizer.py with coordinate tests for:
  - Graph coordinates usage
  - Partial coordinates handling
  - 3D to 2D projection
  - Empty coordinates fallback

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 graphqomb/graphstate.py  |   4 +-
 graphqomb/pattern.py     |   4 +-
 graphqomb/qompiler.py    |   6 +--
 tests/test_graphstate.py |  37 +++++++++++--
 tests/test_visualizer.py | 113 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 149 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_visualizer.py

diff --git a/graphqomb/graphstate.py b/graphqomb/graphstate.py
index e8a150554..74380fd90 100644
--- a/graphqomb/graphstate.py
+++ b/graphqomb/graphstate.py
@@ -332,9 +332,7 @@ def _ensure_node_exists(self, node: int) -> None:
             raise ValueError(msg)
 
     @typing_extensions.override
-    def add_physical_node(
-        self, coordinate: tuple[float, ...] | None = None
-    ) -> int:
+    def add_physical_node(self, coordinate: tuple[float, ...] | None = None) -> int:
         r"""Add a physical node to the graph state.
 
         Parameters
diff --git a/graphqomb/pattern.py b/graphqomb/pattern.py
index 13e8b08e5..47465daee 100644
--- a/graphqomb/pattern.py
+++ b/graphqomb/pattern.py
@@ -46,9 +46,7 @@ class Pattern(Sequence[Command]):
     output_node_indices: dict[int, int]
     commands: tuple[Command, ...]
     pauli_frame: PauliFrame
-    input_coordinates: dict[int, tuple[float, ...]] = dataclasses.field(
-        default_factory=dict
-    )
+    input_coordinates: dict[int, tuple[float, ...]] = dataclasses.field(default_factory=dict)
 
     def __len__(self) -> int:
         return len(self.commands)
diff --git a/graphqomb/qompiler.py b/graphqomb/qompiler.py
index 16c42d9e8..6b20bca8f 100644
--- a/graphqomb/qompiler.py
+++ b/graphqomb/qompiler.py
@@ -130,11 +130,7 @@ def _qompile(
             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
-    }
+    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,
diff --git a/tests/test_graphstate.py b/tests/test_graphstate.py
index 57fdf0f7c..70c8238f9 100644
--- a/tests/test_graphstate.py
+++ b/tests/test_graphstate.py
@@ -7,7 +7,7 @@
 
 from graphqomb.common import Plane, PlannerMeasBasis
 from graphqomb.euler import LocalClifford
-from graphqomb.graphstate import GraphState, bipartite_edges, odd_neighbors
+from graphqomb.graphstate import GraphState, bipartite_edges, compose, odd_neighbors
 
 
 @pytest.fixture
@@ -331,9 +331,7 @@ def test_from_graph_with_coordinates() -> None:
     edges = [("a", "b"), ("b", "c")]
     coordinates = {"a": (0.0, 0.0), "b": (1.0, 0.0), "c": (2.0, 0.0)}
 
-    graph, node_map = GraphState.from_graph(
-        nodes, edges, inputs=["a"], outputs=["c"], coordinates=coordinates
-    )
+    graph, node_map = GraphState.from_graph(nodes, edges, inputs=["a"], outputs=["c"], coordinates=coordinates)
 
     assert graph.coordinates[node_map["a"]] == (0.0, 0.0)
     assert graph.coordinates[node_map["b"]] == (1.0, 0.0)
@@ -354,3 +352,34 @@ def test_from_base_graph_state_copies_coordinates() -> None:
 
     assert graph2.coordinates[node_map[node1]] == (1.0, 2.0)
     assert graph2.coordinates[node_map[node2]] == (3.0, 4.0)
+
+
+def test_compose_copies_coordinates() -> None:
+    """Test that compose copies coordinates from both graphs."""
+    # Create first graph with coordinates
+    graph1 = GraphState()
+    g1_in = graph1.add_physical_node(coordinate=(0.0, 0.0))
+    g1_out = graph1.add_physical_node(coordinate=(1.0, 0.0))
+    graph1.add_physical_edge(g1_in, g1_out)
+    graph1.register_input(g1_in, 0)
+    graph1.register_output(g1_out, 0)
+    graph1.assign_meas_basis(g1_in, PlannerMeasBasis(Plane.XY, 0.0))
+
+    # Create second graph with coordinates
+    graph2 = GraphState()
+    g2_in = graph2.add_physical_node(coordinate=(2.0, 0.0))
+    g2_out = graph2.add_physical_node(coordinate=(3.0, 0.0))
+    graph2.add_physical_edge(g2_in, g2_out)
+    graph2.register_input(g2_in, 0)
+    graph2.register_output(g2_out, 0)
+    graph2.assign_meas_basis(g2_in, PlannerMeasBasis(Plane.XY, 0.0))
+
+    # Compose graphs
+    composed, node_map1, node_map2 = compose(graph1, graph2)
+
+    # Verify coordinates from graph1 (input node only, output is connected)
+    assert composed.coordinates[node_map1[g1_in]] == (0.0, 0.0)
+
+    # Verify coordinates from graph2
+    assert composed.coordinates[node_map2[g2_in]] == (2.0, 0.0)
+    assert composed.coordinates[node_map2[g2_out]] == (3.0, 0.0)
diff --git a/tests/test_visualizer.py b/tests/test_visualizer.py
new file mode 100644
index 000000000..15212c5af
--- /dev/null
+++ b/tests/test_visualizer.py
@@ -0,0 +1,113 @@
+"""Tests for the visualizer module."""
+
+from __future__ import annotations
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import pytest
+
+from graphqomb.common import Plane, PlannerMeasBasis
+from graphqomb.graphstate import GraphState
+from graphqomb.visualizer import visualize
+
+# Use non-interactive backend for tests
+mpl.use("Agg")
+
+
+@pytest.fixture
+def graph_with_coordinates() -> GraphState:
+    """Create a graph with coordinates.
+
+    Returns
+    -------
+    GraphState
+        A graph with nodes at coordinates (0,0), (1,0), (2,0).
+    """
+    graph = GraphState()
+    node1 = graph.add_physical_node(coordinate=(0.0, 0.0))
+    node2 = graph.add_physical_node(coordinate=(1.0, 0.0))
+    node3 = graph.add_physical_node(coordinate=(2.0, 0.0))
+
+    graph.add_physical_edge(node1, node2)
+    graph.add_physical_edge(node2, node3)
+
+    graph.register_input(node1, 0)
+    graph.register_output(node3, 0)
+
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+    graph.assign_meas_basis(node2, PlannerMeasBasis(Plane.XY, 0.0))
+
+    return graph
+
+
+def test_visualize_with_graph_coordinates(graph_with_coordinates: GraphState) -> None:
+    """Test that visualize uses graph coordinates when available."""
+    ax = visualize(graph_with_coordinates, use_graph_coordinates=True)
+    assert ax is not None
+    plt.close("all")
+
+
+def test_visualize_without_graph_coordinates(graph_with_coordinates: GraphState) -> None:
+    """Test that visualize can use auto-calculated positions."""
+    ax = visualize(graph_with_coordinates, use_graph_coordinates=False)
+    assert ax is not None
+    plt.close("all")
+
+
+def test_visualize_with_partial_coordinates() -> None:
+    """Test that visualize handles graphs with partial coordinates."""
+    graph = GraphState()
+    # Only some nodes have coordinates
+    node1 = graph.add_physical_node(coordinate=(0.0, 0.0))
+    node2 = graph.add_physical_node()  # No coordinate
+    node3 = graph.add_physical_node(coordinate=(2.0, 0.0))
+
+    graph.add_physical_edge(node1, node2)
+    graph.add_physical_edge(node2, node3)
+
+    graph.register_input(node1, 0)
+    graph.register_output(node3, 0)
+
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+    graph.assign_meas_basis(node2, PlannerMeasBasis(Plane.XY, 0.0))
+
+    ax = visualize(graph, use_graph_coordinates=True)
+    assert ax is not None
+    plt.close("all")
+
+
+def test_visualize_with_3d_coordinates() -> None:
+    """Test that visualize projects 3D coordinates to 2D."""
+    graph = GraphState()
+    node1 = graph.add_physical_node(coordinate=(0.0, 0.0, 0.0))
+    node2 = graph.add_physical_node(coordinate=(1.0, 1.0, 1.0))
+
+    graph.add_physical_edge(node1, node2)
+
+    graph.register_input(node1, 0)
+    graph.register_output(node2, 0)
+
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+
+    ax = visualize(graph, use_graph_coordinates=True)
+    assert ax is not None
+    plt.close("all")
+
+
+def test_visualize_empty_coordinates() -> None:
+    """Test that visualize works with empty graph coordinates."""
+    graph = GraphState()
+    node1 = graph.add_physical_node()  # No coordinate
+    node2 = graph.add_physical_node()  # No coordinate
+
+    graph.add_physical_edge(node1, node2)
+
+    graph.register_input(node1, 0)
+    graph.register_output(node2, 0)
+
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+
+    # Should fall back to auto-calculated positions
+    ax = visualize(graph, use_graph_coordinates=True)
+    assert ax is not None
+    plt.close("all")

From f2bf62ea40448ef8dab944338171e6c93af74590 Mon Sep 17 00:00:00 2001
From: Masato Fukushima <masa1063fuk@gmail.com>
Date: Sun, 4 Jan 2026 14:10:02 +0900
Subject: [PATCH 3/6] Add 1D coordinate support and improve test coverage
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Support 1D coordinates in visualizer (y defaults to 0.0)
- Add docstring note about 3D visualization being planned
- Add test for set_coordinate with invalid node
- Add tests for N command string representation
- Add test for 1D coordinates in visualizer

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 graphqomb/visualizer.py  | 12 ++++++++++--
 tests/test_graphstate.py |  6 ++++++
 tests/test_pattern.py    | 12 ++++++++++++
 tests/test_visualizer.py | 18 ++++++++++++++++++
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/graphqomb/visualizer.py b/graphqomb/visualizer.py
index aaeb1eb33..8aad2b2c8 100644
--- a/graphqomb/visualizer.py
+++ b/graphqomb/visualizer.py
@@ -26,6 +26,8 @@
 
     from graphqomb.graphstate import BaseGraphState
 
+# Minimum number of coordinate dimensions required for 2D visualization
+_MIN_2D_COORDS = 2
 
 if sys.version_info >= (3, 11):
     from enum import StrEnum
@@ -94,6 +96,12 @@ def visualize(  # noqa: PLR0913
     -------
     `matplotlib.axes.Axes`
         The Axes object containing the visualization
+
+    Notes
+    -----
+    Currently only 2D visualization is supported. For 3D coordinates, only the
+    x and y components are used; the z component is ignored. 3D visualization
+    support is planned for a future release.
     """
     node_pos = _get_node_positions(graph, use_graph_coordinates)
 
@@ -251,8 +259,8 @@ def _get_node_positions(
         # Use graph coordinates (project 3D to 2D by using x, y only)
         node_pos: dict[int, tuple[float, float]] = {}
         for node, coord in graph.coordinates.items():
-            # Take first two coordinates for 2D projection
-            node_pos[node] = (coord[0], coord[1])
+            # Take first two coordinates for 2D projection (1D uses y=0.0)
+            node_pos[node] = (coord[0], coord[1] if len(coord) >= _MIN_2D_COORDS else 0.0)
 
         # For nodes without coordinates, calculate positions
         missing_nodes = graph.physical_nodes - node_pos.keys()
diff --git a/tests/test_graphstate.py b/tests/test_graphstate.py
index 70c8238f9..8e10ffe84 100644
--- a/tests/test_graphstate.py
+++ b/tests/test_graphstate.py
@@ -297,6 +297,12 @@ def test_set_coordinate(graph: GraphState) -> None:
     assert graph.coordinates == {node: (1.0, 2.0)}
 
 
+def test_set_coordinate_invalid_node(graph: GraphState) -> None:
+    """Test that set_coordinate raises error for non-existent node."""
+    with pytest.raises(ValueError):
+        graph.set_coordinate(999, (1.0, 2.0))
+
+
 def test_set_coordinate_3d(graph: GraphState) -> None:
     """Test setting 3D coordinates for a node."""
     node = graph.add_physical_node()
diff --git a/tests/test_pattern.py b/tests/test_pattern.py
index b24dcef60..a5ad70099 100644
--- a/tests/test_pattern.py
+++ b/tests/test_pattern.py
@@ -11,6 +11,18 @@
 from graphqomb.pauli_frame import PauliFrame
 
 
+def test_n_command_str_with_coordinate() -> None:
+    """Test N command string representation with coordinate."""
+    cmd = N(node=5, coordinate=(1.0, 2.0))
+    assert str(cmd) == "N: node=5, coord=(1.0, 2.0)"
+
+
+def test_n_command_str_without_coordinate() -> None:
+    """Test N command string representation without coordinate."""
+    cmd = N(node=3)
+    assert str(cmd) == "N: node=3"
+
+
 @pytest.fixture
 def pattern_components() -> tuple[dict[int, int], dict[int, int], PauliFrame, list[int]]:
     """Create shared components for building Pattern instances.
diff --git a/tests/test_visualizer.py b/tests/test_visualizer.py
index 15212c5af..7ac950257 100644
--- a/tests/test_visualizer.py
+++ b/tests/test_visualizer.py
@@ -94,6 +94,24 @@ def test_visualize_with_3d_coordinates() -> None:
     plt.close("all")
 
 
+def test_visualize_with_1d_coordinates() -> None:
+    """Test that visualize handles 1D coordinates by using y=0."""
+    graph = GraphState()
+    node1 = graph.add_physical_node(coordinate=(0.0,))
+    node2 = graph.add_physical_node(coordinate=(1.0,))
+
+    graph.add_physical_edge(node1, node2)
+
+    graph.register_input(node1, 0)
+    graph.register_output(node2, 0)
+
+    graph.assign_meas_basis(node1, PlannerMeasBasis(Plane.XY, 0.0))
+
+    ax = visualize(graph, use_graph_coordinates=True)
+    assert ax is not None
+    plt.close("all")
+
+
 def test_visualize_empty_coordinates() -> None:
     """Test that visualize works with empty graph coordinates."""
     graph = GraphState()

From 12ab44e99e7fa9dd2d34af30bfe19af04e0a1b4e Mon Sep 17 00:00:00 2001
From: Masato Fukushima <masa1063fuk@gmail.com>
Date: Sun, 4 Jan 2026 14:19:05 +0900
Subject: [PATCH 4/6] Address PR review comments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Add coordinate parameter to BaseGraphState.add_physical_node
- Rename _get_node_positions to _determine_node_positions

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 graphqomb/graphstate.py | 11 ++++++++---
 graphqomb/visualizer.py |  4 ++--
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/graphqomb/graphstate.py b/graphqomb/graphstate.py
index 74380fd90..9405f56f7 100644
--- a/graphqomb/graphstate.py
+++ b/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
diff --git a/graphqomb/visualizer.py b/graphqomb/visualizer.py
index 8aad2b2c8..4a8446d78 100644
--- a/graphqomb/visualizer.py
+++ b/graphqomb/visualizer.py
@@ -103,7 +103,7 @@ def visualize(  # noqa: PLR0913
     x and y components are used; the z component is ignored. 3D visualization
     support is planned for a future release.
     """
-    node_pos = _get_node_positions(graph, use_graph_coordinates)
+    node_pos = _determine_node_positions(graph, use_graph_coordinates)
 
     node_colors = _determine_node_colors(graph)
 
@@ -237,7 +237,7 @@ def _setup_figure(node_pos: Mapping[int, tuple[float, float]]) -> FigureSetup:
     )
 
 
-def _get_node_positions(
+def _determine_node_positions(
     graph: BaseGraphState,
     use_graph_coordinates: bool,
 ) -> dict[int, tuple[float, float]]:

From 607d7843308297a541ac6788efea9f82b5cf3ccd Mon Sep 17 00:00:00 2001
From: Masato Fukushima <masa1063fuk@gmail.com>
Date: Sun, 4 Jan 2026 23:40:02 +0900
Subject: [PATCH 5/6] fix ruff error

---
 tests/test_graphstate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_graphstate.py b/tests/test_graphstate.py
index 8e10ffe84..f65538a63 100644
--- a/tests/test_graphstate.py
+++ b/tests/test_graphstate.py
@@ -299,7 +299,7 @@ def test_set_coordinate(graph: GraphState) -> None:
 
 def test_set_coordinate_invalid_node(graph: GraphState) -> None:
     """Test that set_coordinate raises error for non-existent node."""
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="Node does not exist"):
         graph.set_coordinate(999, (1.0, 2.0))
 
 

From 4b79df7184c1bf1665ff1ea5e1577fcb71b9be36 Mon Sep 17 00:00:00 2001
From: masa10-f <masa1063fuk@gmail.com>
Date: Sat, 10 Jan 2026 18:44:36 +0900
Subject: [PATCH 6/6] fix missing s

---
 graphqomb/pattern.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphqomb/pattern.py b/graphqomb/pattern.py
index ac00733e3..bd9ecdaae 100644
--- a/graphqomb/pattern.py
+++ b/graphqomb/pattern.py
@@ -117,7 +117,7 @@ def coordinates(self) -> dict[int, tuple[float, ...]]:
         for cmd in self.commands:
             if isinstance(cmd, N) and cmd.coordinate is not None:
                 coords[cmd.node] = cmd.coordinate
-        return coord
+        return coords
 
     @property
     def active_volume(self) -> int: