feat: add lazy contraction scheme controls

Author: DOKOS-TAYOS

README.md

@@ -104,7 +104,7 @@ Everything below maps to real parameters—there are no hidden mode switches.
 | **Display mode** | `show=True` / `False` | If `True`: Jupyter **kernel** uses `IPython.display.display(fig)`; otherwise `plt.show()`. If `False`: neither runs—use for `savefig` / batch. |
 | **Label policy** | `PlotConfig` + overrides | Defaults: `show_tensor_labels`, `show_index_labels`. Per-call: `show_tensor_network(..., show_tensor_labels=..., show_index_labels=...)`. |
 | **Hover labels** | `PlotConfig(hover_labels=True)` | Tensor names and bond labels appear on pointer hover (2D axes hit-test; 3D screen-space distance). Needs an **interactive** Matplotlib window. |
-| **Contraction scheme** | `PlotConfig(show_contraction_scheme=True)` | **Einsum:** cumulative per-step highlights from the trace. **Other engines:** set **`contraction_scheme_by_name`**. **2D:** rounded boxes (AABB + pad); colored borders (no fill by default). **3D:** wireframe box. See **`docs/guide.md`**. |
+| **Contraction scheme** | `PlotConfig(show_contraction_scheme=True)` | **Einsum:** cumulative per-step highlights from the trace. **Other engines:** set **`contraction_scheme_by_name`**. Compatible figures now add Matplotlib toggles for **Scheme**, **Playback**, and **Cost hover**; if you start with those flags off, the scheme bundle is computed lazily on first use. **2D:** rounded boxes (AABB + pad); colored borders (no fill by default). **3D:** wireframe box. See **`docs/guide.md`**. |
 | **Einsum workflow** | `engine="einsum"` | **Auto:** `EinsumTrace` + `einsum` (binary `pair_tensor`, unary/ternary+ `einsum_trace_step`; implicit `->`, `out=`). **Manual:** `pair_tensor` / `einsum_trace_step` (ellipsis needs `metadata` shapes). See **`examples/einsum_general.py`**. |
 
 ## Minimal examples

docs/guide.md

@@ -360,7 +360,10 @@ with **`PlotConfig(show_contraction_scheme=True)`**; **2D** rounded **`FancyBbox
 axis-aligned bounding box of the tensors in that step (not per-tensor tight hulls); **later** steps
 drawn **under** earlier ones; **3D** wireframes (no fill). Other engines: **`contraction_scheme_by_name`**
 (each step should list the tensors you want in that hull; end with the full network if you want one
-global region).
+global region). When a figure has usable contraction steps, it also gets Matplotlib toggles for
+**`Scheme`**, **`Playback`**, and **`Cost hover`**. If those flags start disabled in
+**`PlotConfig`**, the scheme geometry, playback viewer, and cost tooltips are built lazily on the
+first toggle that needs them, then reused for the rest of the figure lifetime.
 
 <a id="toc-plotconfig"></a>
 
@@ -399,6 +402,8 @@ Frozen dataclass in [`src/tensor_network_viz/config.py`](../src/tensor_network_v
 | `contraction_scheme_linewidth` | `None` | → **`DEFAULT_CONTRACTION_SCHEME_LINEWIDTH`** (scaled like other strokes). |
 | `contraction_scheme_colors` | `None` | Cycle of colors; built-in categorical palette if unset. |
 | `contraction_scheme_by_name` | `None` | Override schedule: per step, tuple of **`node.name`** strings for non-virtual tensors. |
+| `contraction_playback` | `False` | Start with the step slider and **Play/Pause/Reset** visible; still requires **`show_contraction_scheme=True`** in code. |
+| `contraction_scheme_cost_hover` | `False` | Show the naive dense contraction-cost tooltip when hovering a visible scheme hull. |
 
 ### Recipe: publication palette
 
@@ -464,6 +469,11 @@ fig, ax = show_tensor_network(
 )
 ```
 
+Compatible figures show the three figure-level toggles even if you leave these flags off in
+**`PlotConfig`**. Turning on **`Playback`** or **`Cost hover`** from the UI auto-enables
+**`Scheme`**. Turning **`Scheme`** off hides the hulls and playback widgets, but remembered child
+states are restored instantly when you turn it back on.
+
 ### Recipe: custom positions with validation
 
 ```python

src/tensor_network_viz/_core/draw/graph_pipeline.py

@@ -3,7 +3,11 @@
 from typing import Any, Literal
 
 from ...config import PlotConfig
-from ...contraction_viewer import attach_playback_to_tensor_network_figure
+from ...contraction_viewer import (
+    _ContractionControls,
+    _ContractionSchemeBundle,
+    attach_playback_to_tensor_network_figure,
+)
 from ...einsum_module.contraction_cost import format_contraction_step_tooltip
 from ..contractions import _ContractionGroups
 from ..graph import _GraphData
@@ -15,12 +19,150 @@
 )
 from .disk_metrics import _tensor_disk_radius_px_3d_nominal
 from .render_prep import (
+    _apply_render_hover_state,
     _draw_edges_nodes_and_labels,
     _prepare_render_context,
     _register_render_hover,
+    _RenderPrepContext,
 )
 
 
+def _has_contraction_scheme_source(
+    graph: _GraphData,
+    config: PlotConfig,
+) -> bool:
+    if config.contraction_scheme_by_name is not None:
+        return len(config.contraction_scheme_by_name) > 0
+    steps = graph.contraction_steps
+    return steps is not None and len(steps) > 0
+
+
+def _scheme_bounds_2d(
+    artists_by_step: list[Any | None],
+) -> tuple[float, float, float, float] | None:
+    xs: list[float] = []
+    ys: list[float] = []
+    for artist in artists_by_step:
+        if artist is None:
+            continue
+        get_x = getattr(artist, "get_x", None)
+        get_y = getattr(artist, "get_y", None)
+        get_width = getattr(artist, "get_width", None)
+        get_height = getattr(artist, "get_height", None)
+        if not all(callable(fn) for fn in (get_x, get_y, get_width, get_height)):
+            continue
+        x0 = float(get_x())
+        y0 = float(get_y())
+        x1 = x0 + float(get_width())
+        y1 = y0 + float(get_height())
+        xs.extend((x0, x1))
+        ys.extend((y0, y1))
+    if not xs or not ys:
+        return None
+    return (min(xs), max(xs), min(ys), max(ys))
+
+
+def _scheme_bounds_3d(
+    scheme_aabb: list[tuple[float, float, float, float, float, float] | None],
+) -> tuple[float, float, float, float, float, float] | None:
+    boxes = [box for box in scheme_aabb if box is not None]
+    if not boxes:
+        return None
+    xmin = min(box[0] for box in boxes)
+    xmax = max(box[1] for box in boxes)
+    ymin = min(box[2] for box in boxes)
+    ymax = max(box[3] for box in boxes)
+    zmin = min(box[4] for box in boxes)
+    zmax = max(box[5] for box in boxes)
+    return (xmin, xmax, ymin, ymax, zmin, zmax)
+
+
+def _expand_axes_for_scheme_bounds(
+    *,
+    ax: Any,
+    dimensions: Literal[2, 3],
+    bounds_2d: tuple[float, float, float, float] | None,
+    bounds_3d: tuple[float, float, float, float, float, float] | None,
+) -> None:
+    if dimensions == 2 and bounds_2d is not None:
+        xmin, xmax, ymin, ymax = bounds_2d
+        cur_x0, cur_x1 = ax.get_xlim()
+        cur_y0, cur_y1 = ax.get_ylim()
+        ax.set_xlim(min(float(cur_x0), xmin), max(float(cur_x1), xmax))
+        ax.set_ylim(min(float(cur_y0), ymin), max(float(cur_y1), ymax))
+        return
+    if dimensions == 3 and bounds_3d is not None:
+        xmin, xmax, ymin, ymax, zmin, zmax = bounds_3d
+        cur_x0, cur_x1 = ax.get_xlim3d()
+        cur_y0, cur_y1 = ax.get_ylim3d()
+        cur_z0, cur_z1 = ax.get_zlim3d()
+        ax.set_xlim3d(min(float(cur_x0), xmin), max(float(cur_x1), xmax))
+        ax.set_ylim3d(min(float(cur_y0), ymin), max(float(cur_y1), ymax))
+        ax.set_zlim3d(min(float(cur_z0), zmin), max(float(cur_z1), zmax))
+
+
+def _build_contraction_scheme_bundle(
+    *,
+    ax: Any,
+    graph: _GraphData,
+    positions: NodePositions,
+    config: PlotConfig,
+    dimensions: Literal[2, 3],
+    scale: float,
+    context: _RenderPrepContext,
+    strict: bool,
+) -> _ContractionSchemeBundle:
+    scheme_steps_eff = _effective_contraction_steps(graph, config)
+    if not scheme_steps_eff:
+        raise ValueError("contraction scheme requires a non-empty contraction step sequence.")
+
+    per_step_artists, scheme_aabb = _draw_contraction_scheme(
+        ax=ax,
+        graph=graph,
+        positions=positions,
+        steps=scheme_steps_eff,
+        config=config,
+        dimensions=dimensions,
+        scale=scale,
+        p=context.params,
+    )
+    has_drawable_artists = any(artist is not None for artist in per_step_artists)
+    if strict and not has_drawable_artists:
+        raise ValueError("contraction scheme requires at least one drawable contraction step.")
+
+    metrics_row = _contraction_step_metrics_for_draw(graph, scheme_steps_eff)
+    tooltips = tuple(
+        format_contraction_step_tooltip(metric) if metric is not None else None
+        for metric in (metrics_row or ())
+    )
+    viewer = attach_playback_to_tensor_network_figure(
+        artists_by_step=per_step_artists,
+        fig=ax.figure,
+        ax=ax,
+        config=config,
+        build_ui=False,
+    )
+    bounds_2d = _scheme_bounds_2d(per_step_artists) if dimensions == 2 else None
+    bounds_3d = _scheme_bounds_3d(scheme_aabb) if dimensions == 3 else None
+    _expand_axes_for_scheme_bounds(
+        ax=ax,
+        dimensions=dimensions,
+        bounds_2d=bounds_2d,
+        bounds_3d=bounds_3d,
+    )
+    return _ContractionSchemeBundle(
+        availability="computed",
+        steps=scheme_steps_eff,
+        artists_by_step=per_step_artists,
+        scheme_aabb=scheme_aabb,
+        metrics_row=tuple(metrics_row) if metrics_row is not None else None,
+        tooltips=tooltips,
+        viewer=viewer,
+        bounds_2d=bounds_2d,
+        bounds_3d=bounds_3d,
+    )
+
+
 def _draw_graph(
     *,
     ax: Any,
@@ -54,39 +196,6 @@ def _draw_graph(
             "steps can be drawn and stepped."
         )
 
-    per_step_artists: list[Any | None] | None = None
-    scheme_steps_eff: tuple[frozenset[int], ...] | None = None
-    scheme_aabb: list[tuple[float, float, float, float, float, float] | None] | None = None
-    if config.show_contraction_scheme:
-        scheme_steps_eff = _effective_contraction_steps(graph, config)
-        if scheme_steps_eff:
-            per_step_artists, scheme_aabb = _draw_contraction_scheme(
-                ax=ax,
-                graph=graph,
-                positions=positions,
-                steps=scheme_steps_eff,
-                config=config,
-                dimensions=dimensions,
-                scale=scale,
-                p=context.params,
-            )
-            if config.contraction_playback:
-                if not any(artist is not None for artist in per_step_artists):
-                    raise ValueError(
-                        "contraction_playback requires at least one drawable "
-                        "contraction scheme step."
-                    )
-                attach_playback_to_tensor_network_figure(
-                    artists_by_step=per_step_artists,
-                    fig=ax.figure,
-                    ax=ax,
-                    config=config,
-                )
-        elif config.contraction_playback:
-            raise ValueError(
-                "contraction_playback requires a non-empty contraction step sequence on the graph."
-            )
-
     tensor_disk_radius_px_3d: float | None = None
     if dimensions == 3 and config.approximate_3d_tensor_disk_px:
         tensor_disk_radius_px_3d = _tensor_disk_radius_px_3d_nominal(ax, context.params)
@@ -100,40 +209,38 @@ def _draw_graph(
         tensor_disk_radius_px_3d=tensor_disk_radius_px_3d,
     )
 
-    metrics_row = (
-        _contraction_step_metrics_for_draw(graph, scheme_steps_eff) if scheme_steps_eff else None
-    )
-    scheme_patches_2d: list[tuple[Any, str]] = []
-    scheme_aabbs_3d: list[tuple[tuple[float, float, float, float, float, float], str, Any]] = []
-    if (
-        config.contraction_scheme_cost_hover
-        and metrics_row is not None
-        and per_step_artists is not None
-        and scheme_aabb is not None
-    ):
-        for index, artist in enumerate(per_step_artists):
-            if index >= len(metrics_row):
-                break
-            metric = metrics_row[index]
-            if metric is None or artist is None:
-                continue
-            tooltip = format_contraction_step_tooltip(metric)
-            if dimensions == 2:
-                scheme_patches_2d.append((artist, tooltip))
-                continue
-            box = scheme_aabb[index] if index < len(scheme_aabb) else None
-            if box is not None:
-                scheme_aabbs_3d.append((box, tooltip, artist))
-
-    _register_render_hover(
+    hover_state = _register_render_hover(
         ax=ax,
         context=context,
         show_tensor_labels=show_tensor_labels,
         show_index_labels=show_index_labels,
-        scheme_patches_2d=scheme_patches_2d,
-        scheme_aabbs_3d=scheme_aabbs_3d,
+        scheme_patches_2d=[],
+        scheme_aabbs_3d=[],
         tensor_disk_radius_px_3d=tensor_disk_radius_px_3d,
     )
+    if not _has_contraction_scheme_source(graph, config):
+        return
+
+    _ContractionControls(
+        fig=ax.figure,
+        ax=ax,
+        config=config,
+        bundle_builder=lambda strict: _build_contraction_scheme_bundle(
+            ax=ax,
+            graph=graph,
+            positions=positions,
+            config=config,
+            dimensions=dimensions,
+            scale=scale,
+            context=context,
+            strict=strict,
+        ),
+        refresh_hover=lambda scheme_patches_2d, scheme_aabbs_3d: _apply_render_hover_state(
+            hover_state,
+            scheme_patches_2d=scheme_patches_2d,
+            scheme_aabbs_3d=scheme_aabbs_3d,
+        ),
+    )
 
 
 __all__ = ["_draw_graph"]