Geometry for 3D cubic structures

Author: DOKOS-TAYOS

CHANGELOG.md

@@ -15,6 +15,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `show_tensor_network` package export: `TYPE_CHECKING` import preserves lazy runtime import while exposing the real signature, annotations, and docstring to IDEs.
 
+### Changed
+- Refactor of the `_core` module into tiny pieces.
+
 ## [1.4.0] — 2026-03-28
 
 ### Added

CONTRIBUTING.md

@@ -51,6 +51,72 @@ With the project venv (Windows):
 
 Add tests for new features or bug fixes. All tests must pass before opening a PR.
 
+### Optional: manual example smoke checks
+
+Automated tests do not open interactive Matplotlib windows. After **`pytest`** passes, you can sanity-check **layout and drawing** by running the examples below from the **repository root** (with **`pip install -e ".[dev]"`** or the matching optional extras). Run **one command at a time** (each line is a separate invocation).
+
+**2D (default labels):**
+
+```bash
+python examples/tensorkrowch_demo.py mps 2d
+python examples/tensorkrowch_demo.py disconnected 2d
+python examples/tensornetwork_demo.py weird 2d
+python examples/mera_tree_demo.py 2d
+python examples/cubic_peps_demo.py 2d
+python examples/quimb_demo.py hyper 2d
+python examples/tenpy_demo.py imps 2d
+python examples/tenpy_demo.py impo 2d
+python examples/einsum_demo.py peps 2d
+python examples/tn_tsp.py -n 4 --view 2d
+```
+
+**2D with `--hover-labels`** (interactive window only):
+
+```bash
+python examples/tensorkrowch_demo.py mps 2d --hover-labels
+python examples/tensorkrowch_demo.py disconnected 2d --hover-labels
+python examples/tensornetwork_demo.py weird 2d --hover-labels
+python examples/mera_tree_demo.py 2d --hover-labels
+python examples/cubic_peps_demo.py 2d --hover-labels
+python examples/quimb_demo.py hyper 2d --hover-labels
+python examples/tenpy_demo.py imps 2d --hover-labels
+python examples/tenpy_demo.py impo 2d --hover-labels
+python examples/einsum_demo.py peps 2d --hover-labels
+python examples/tn_tsp.py -n 4 --view 2d --hover-labels
+```
+
+**3D (default labels):**
+
+```bash
+python examples/tensorkrowch_demo.py mps 3d
+python examples/tensorkrowch_demo.py weird 3d
+python examples/tensorkrowch_demo.py disconnected 3d
+python examples/mera_tree_demo.py 3d --mera-log2 5 --tree-depth 4
+python examples/cubic_peps_demo.py 3d --lx 3 --ly 3 --lz 4
+python examples/quimb_demo.py hyper 3d
+python examples/tenpy_demo.py imps 3d
+python examples/tenpy_demo.py impo 3d
+python examples/einsum_demo.py peps 3d
+python examples/tn_tsp.py -n 5 --view 3d
+```
+
+**3D with `--hover-labels`** (interactive window only):
+
+```bash
+python examples/tensorkrowch_demo.py mps 3d --hover-labels
+python examples/tensorkrowch_demo.py weird 3d --hover-labels
+python examples/tensorkrowch_demo.py disconnected 3d --hover-labels
+python examples/mera_tree_demo.py 3d --mera-log2 5 --tree-depth 4 --hover-labels
+python examples/cubic_peps_demo.py 3d --lx 3 --ly 3 --lz 4 --hover-labels
+python examples/quimb_demo.py hyper 3d --hover-labels
+python examples/tenpy_demo.py imps 3d --hover-labels
+python examples/tenpy_demo.py impo 3d --hover-labels
+python examples/einsum_demo.py peps 3d --hover-labels
+python examples/tn_tsp.py -n 5 --view 3d --hover-labels
+```
+
+See **`examples/README.md`** for per-script options and dependencies.
+
 ## Lint and Type Checks
 
 **Ruff** (lint and format):

examples/README.md

@@ -149,6 +149,9 @@ other scripts.
 
 ## Run-all cheat sheet
 
+For a **short pre-PR smoke checklist** (2D / 3D, default labels vs `--hover-labels`), see
+[CONTRIBUTING.md](../CONTRIBUTING.md#optional-manual-example-smoke-checks).
+
 ```bash
 python examples/tensorkrowch_demo.py mps 2d
 python examples/tensorkrowch_demo.py mps 3d

src/tensor_network_viz/_core/layout/body.py

@@ -168,8 +168,12 @@ def _lift_component_layout_3d(
         node_id: np.array([coords[0], coords[1], 0.0], dtype=float)
         for node_id, coords in positions_2d.items()
     }
+    if component.structure_kind == "grid3d" and component.grid3d_mapping is not None:
+        for node_id, (i, j, k) in component.grid3d_mapping.items():
+            positions[node_id] = np.array([float(i), float(j), float(k)], dtype=float)
     _place_trimmed_leaf_nodes_3d(component, positions)
-    _promote_3d_layers(graph, component, positions)
+    if component.structure_kind != "grid3d" or component.grid3d_mapping is None:
+        _promote_3d_layers(graph, component, positions)
     return positions
 
 

src/tensor_network_viz/_core/layout_structure.py

@@ -12,7 +12,10 @@
 from .contractions import _iter_contractions
 from .graph import _GraphData
 
-StructureKind = Literal["chain", "grid", "tree", "planar", "generic"]
+StructureKind = Literal["chain", "grid", "grid3d", "tree", "planar", "generic"]
+
+# Skew in 2D projection (i,j,k) → layout xy so sites with same (i,j) and different k stay separated.
+_GRID3D_PROJECTION_K: float = 0.4
 
 
 @dataclass(frozen=True)
@@ -29,6 +32,7 @@ class _LayoutComponent:
     structure_kind: StructureKind
     chain_order: tuple[int, ...]
     grid_mapping: dict[int, tuple[int, int]] | None
+    grid3d_mapping: dict[int, tuple[int, int, int]] | None
     tree_root: int | None
 
 
@@ -61,7 +65,9 @@ def _analyze_layout_components(graph: _GraphData) -> tuple[_LayoutComponent, ...
             ),
             proxy_visible_graph=proxy_visible_graph,
         )
-        structure_kind, chain_order, grid_mapping, tree_root = _classify_anchor_graph(anchor_graph)
+        structure_kind, chain_order, grid_mapping, grid3d_mapping, tree_root = (
+            _classify_anchor_graph(anchor_graph)
+        )
         components.append(
             _LayoutComponent(
                 node_ids=tuple(sorted(component_node_ids)),
@@ -76,6 +82,7 @@ def _analyze_layout_components(graph: _GraphData) -> tuple[_LayoutComponent, ...
                 structure_kind=structure_kind,
                 chain_order=chain_order,
                 grid_mapping=grid_mapping,
+                grid3d_mapping=grid3d_mapping,
                 tree_root=tree_root,
             )
         )
@@ -199,23 +206,33 @@ def _sorted_connected_components(nx_graph: nx.Graph) -> list[tuple[int, ...]]:
 
 def _classify_anchor_graph(
     anchor_graph: nx.Graph,
-) -> tuple[StructureKind, tuple[int, ...], dict[int, tuple[int, int]] | None, int | None]:
+) -> tuple[
+    StructureKind,
+    tuple[int, ...],
+    dict[int, tuple[int, int]] | None,
+    dict[int, tuple[int, int, int]] | None,
+    int | None,
+]:
     chain_order = _detect_chain(anchor_graph)
     if chain_order is not None:
-        return "chain", chain_order, None, None
+        return "chain", chain_order, None, None, None
 
     grid_mapping = _detect_grid(anchor_graph)
     if grid_mapping is not None:
-        return "grid", (), grid_mapping, None
+        return "grid", (), grid_mapping, None, None
+
+    grid3d_mapping = _detect_grid_3d(anchor_graph)
+    if grid3d_mapping is not None:
+        return "grid3d", (), None, grid3d_mapping, None
 
     if anchor_graph.number_of_nodes() > 0 and nx.is_tree(anchor_graph):
-        return "tree", (), None, _tree_root(anchor_graph)
+        return "tree", (), None, None, _tree_root(anchor_graph)
 
     is_planar, _ = nx.check_planarity(anchor_graph)
     if anchor_graph.number_of_nodes() > 0 and is_planar:
-        return "planar", (), None, None
+        return "planar", (), None, None, None
 
-    return "generic", (), None, None
+    return "generic", (), None, None, None
 
 
 def _specialized_anchor_positions(component: _LayoutComponent) -> dict[int, np.ndarray]:
@@ -225,6 +242,8 @@ def _specialized_anchor_positions(component: _LayoutComponent) -> dict[int, np.n
         return _layout_chain(component.chain_order)
     if component.structure_kind == "grid" and component.grid_mapping is not None:
         return _layout_grid(component.grid_mapping)
+    if component.structure_kind == "grid3d" and component.grid3d_mapping is not None:
+        return _layout_grid3d_projection_2d(component.grid3d_mapping)
     if component.structure_kind == "tree" and component.tree_root is not None:
         return _layout_tree(component.anchor_graph, component.tree_root)
     if component.structure_kind == "planar":
@@ -308,6 +327,59 @@ def _layout_grid(grid_mapping: dict[int, tuple[int, int]]) -> dict[int, np.ndarr
     }
 
 
+def _expected_edges_3d_grid(lx: int, ly: int, lz: int) -> int:
+    return (lx - 1) * ly * lz + lx * (ly - 1) * lz + lx * ly * (lz - 1)
+
+
+def _detect_grid_3d(nx_graph: nx.Graph) -> dict[int, tuple[int, int, int]] | None:
+    """If *nx_graph* is a 3D rectangular grid (all factors >= 1), return node → (i,j,k)."""
+    n = nx_graph.number_of_nodes()
+    if n <= 1 or not nx.is_connected(nx_graph):
+        return None
+    m = nx_graph.number_of_edges()
+    best_mn = -1
+    best_shape: tuple[int, int, int] = (-1, -1, -1)
+    best_mapping: dict[int, tuple[int, int, int]] | None = None
+
+    for lx in range(1, n + 1):
+        if n % lx != 0:
+            continue
+        rest = n // lx
+        for ly in range(1, rest + 1):
+            if rest % ly != 0:
+                continue
+            lz = rest // ly
+            if _expected_edges_3d_grid(lx, ly, lz) != m:
+                continue
+            template = nx.grid_graph((lx, ly, lz))
+            mapping = nx.vf2pp_isomorphism(nx_graph, template)
+            if mapping is None:
+                continue
+            mn = min(lx, ly, lz)
+            shape = (lx, ly, lz)
+            if mn > best_mn or (mn == best_mn and shape > best_shape):
+                best_mn = mn
+                best_shape = shape
+                best_mapping = {
+                    int(node_id): (int(a), int(b), int(c)) for node_id, (a, b, c) in mapping.items()
+                }
+
+    return best_mapping
+
+
+def _layout_grid3d_projection_2d(
+    grid3d_mapping: dict[int, tuple[int, int, int]],
+) -> dict[int, np.ndarray]:
+    alpha = _GRID3D_PROJECTION_K
+    return {
+        node_id: np.array(
+            [float(i) + alpha * float(k), float(j) + alpha * float(k)],
+            dtype=float,
+        )
+        for node_id, (i, j, k) in grid3d_mapping.items()
+    }
+
+
 def _layout_planar(nx_graph: nx.Graph) -> dict[int, np.ndarray]:
     positions = nx.planar_layout(nx_graph)
     return {node_id: np.array(positions[node_id], dtype=float) for node_id in nx_graph.nodes}

tests/test_layout_core.py

@@ -168,6 +168,65 @@ def test_resolve_draw_scale_heuristic_when_no_contraction_edges() -> None:
     assert lo <= s <= hi
 
 
+def _build_3d_grid_graph(lx: int, ly: int, lz: int) -> _GraphData:
+    """3D nearest-neighbor cubic lattice; same bond topology as cubic PEPS (no physical legs)."""
+
+    def node_index(i: int, j: int, k: int) -> int:
+        return i * ly * lz + j * lz + k
+
+    nodes = {}
+    edge_specs: list[tuple[int, int, str, str]] = []
+
+    for i in range(lx):
+        for j in range(ly):
+            for k in range(lz):
+                axes_names: list[str] = []
+                if i > 0:
+                    axes_names.append("xm")
+                if i < lx - 1:
+                    axes_names.append("xp")
+                if j > 0:
+                    axes_names.append("ym")
+                if j < ly - 1:
+                    axes_names.append("yp")
+                if k > 0:
+                    axes_names.append("zm")
+                if k < lz - 1:
+                    axes_names.append("zp")
+                nid = node_index(i, j, k)
+                nodes[nid] = _make_node(f"P{i}_{j}_{k}", tuple(axes_names))
+
+    axis_lookup = {
+        node_id: {name: index for index, name in enumerate(node.axes_names)}
+        for node_id, node in nodes.items()
+    }
+
+    for i in range(lx):
+        for j in range(ly):
+            for k in range(lz):
+                nid = node_index(i, j, k)
+                if i < lx - 1:
+                    rid = node_index(i + 1, j, k)
+                    edge_specs.append((nid, rid, "xp", "xm"))
+                if j < ly - 1:
+                    rid = node_index(i, j + 1, k)
+                    edge_specs.append((nid, rid, "yp", "ym"))
+                if k < lz - 1:
+                    rid = node_index(i, j, k + 1)
+                    edge_specs.append((nid, rid, "zp", "zm"))
+
+    edges = [
+        _make_contraction_edge(
+            _EdgeEndpoint(left_id, axis_lookup[left_id][left_name], left_name),
+            _EdgeEndpoint(right_id, axis_lookup[right_id][right_name], right_name),
+            name=f"{left_id}_{right_id}",
+            label=None,
+        )
+        for left_id, right_id, left_name, right_name in edge_specs
+    ]
+    return _GraphData(nodes=nodes, edges=tuple(edges))
+
+
 def _build_grid_graph(rows: int, cols: int) -> _GraphData:
     nodes = {}
     edge_specs: list[tuple[int, int, str, str]] = []
@@ -554,6 +613,32 @@ def test_compute_layout_grid_2d_places_nodes_on_regular_lattice() -> None:
     )
 
 
+def test_compute_layout_3d_grid_spans_three_axes() -> None:
+    graph = _build_3d_grid_graph(2, 2, 2)
+
+    positions = _compute_layout(graph, dimensions=3, seed=0)
+    coords = np.stack([positions[node_id] for node_id in sorted(graph.nodes)])
+
+    assert float(coords[:, 0].std()) > 1e-6
+    assert float(coords[:, 1].std()) > 1e-6
+    assert float(coords[:, 2].std()) > 1e-6
+
+
+def test_compute_layout_3d_grid_uniform_nearest_neighbor_spacing() -> None:
+    graph = _build_3d_grid_graph(2, 3, 2)
+
+    positions = _compute_layout(graph, dimensions=3, seed=0)
+    lengths: list[float] = []
+    for edge in graph.edges:
+        left_ep, right_ep = edge.endpoints
+        delta = positions[left_ep.node_id] - positions[right_ep.node_id]
+        lengths.append(float(np.linalg.norm(delta)))
+
+    assert lengths
+    mean_len = float(np.mean(lengths))
+    assert all(math.isclose(L, mean_len, rel_tol=1e-5, abs_tol=1e-8) for L in lengths)
+
+
 def test_compute_layout_planar_cycle_3d_stays_on_single_plane() -> None:
     graph = _build_planar_cycle_graph()