Fixed tensor inspector

Author: DOKOS-TAYOS

examples/einsum_demo.py

@@ -68,6 +68,19 @@ def _keep_trace_tensors_alive(trace: Any, *tensors: Any) -> None:
     trace._example_keepalive = keepalive
 
 
+def _cumulative_group_contraction_scheme(
+    groups: tuple[tuple[str, ...], ...],
+) -> tuple[tuple[str, ...], ...]:
+    if not groups:
+        return ()
+    running_names: list[str] = []
+    steps: list[tuple[str, ...]] = []
+    for group in groups:
+        running_names.extend(group)
+        steps.append(tuple(running_names))
+    return tuple(steps)
+
+
 def _site_bond_dims(n_sites: int) -> list[int]:
     return [2 + (index % 3) for index in range(max(n_sites - 1, 1))]
 
@@ -586,16 +599,13 @@ def _renderable_trace(trace: Any, args: ExampleCliArgs) -> Any:
 
 def _scheme_steps(name: str, args: ExampleCliArgs) -> tuple[tuple[str, ...], ...] | None:
     if name == "mps":
-        return cumulative_prefix_contraction_scheme(
-            tuple(f"A{i}" for i in range(args.n_sites))
-            + tuple(f"x{i}" for i in range(args.n_sites))
+        return _cumulative_group_contraction_scheme(
+            tuple((f"A{i}", f"x{i}") for i in range(args.n_sites))
         )
     if name == "mpo":
-        mpo_names = tuple(f"W{i}" for i in range(args.n_sites))
-        vectors = tuple(f"d{i}" for i in range(args.n_sites)) + tuple(
-            f"u{i}" for i in range(args.n_sites)
+        return _cumulative_group_contraction_scheme(
+            tuple((f"W{i}", f"d{i}", f"u{i}") for i in range(args.n_sites))
         )
-        return cumulative_prefix_contraction_scheme(mpo_names + vectors)
     if name == "peps":
         names = tuple(f"P{i}_{j}" for i in range(args.lx) for j in range(args.ly))
         return cumulative_prefix_contraction_scheme(names)

src/tensor_network_viz/_interaction/controller.py

@@ -33,7 +33,7 @@
 )
 from .._tensor_elements_support import _TensorRecord
 from .._typing import root_figure
-from .._ui_utils import _set_axes_visible, _set_figure_bottom_reserved
+from .._ui_utils import _set_axes_visible, _set_figure_bottom_reserved, _style_control_tray_axes
 from ..config import EngineName, PlotConfig, ViewName
 from ..contraction_viewer import _MAIN_FIGURE_BOTTOM_RESERVED, _PLAYBACK_DETAILS_TOP
 from ..einsum_module.trace import EinsumTrace
@@ -73,22 +73,27 @@
 _INTERACTIVE_CHECK_FRAME_PROPS: dict[str, float] = {"s": 44.0, "linewidth": 0.9}
 _INTERACTIVE_CHECK_MARK_PROPS: dict[str, float] = {"s": 34.0, "linewidth": 1.0}
 _INTERACTIVE_RADIO_PROPS: dict[str, float] = {"s": 38.0, "linewidth": 0.9}
-_CONTROL_TRAY_FACE: tuple[float, float, float] = (0.97, 0.97, 0.99)
-_CONTROL_TRAY_FRAME: tuple[float, float, float] = (0.78, 0.78, 0.82)
 
 
-def _style_interactive_control_axes(ax: Axes) -> None:
-    ax.set_xticks([])
-    ax.set_yticks([])
-    ax.set_navigate(False)
-    ax.patch.set_facecolor(_CONTROL_TRAY_FACE)
-    ax.patch.set_alpha(0.88)
-    ax.patch.set_edgecolor(_CONTROL_TRAY_FRAME)
-    ax.patch.set_linewidth(0.6)
-    for spine in ax.spines.values():
-        spine.set_visible(True)
-        spine.set_linewidth(0.6)
-        spine.set_color(_CONTROL_TRAY_FRAME)
+def _reveal_auxiliary_figure(figure: Figure) -> None:
+    manager = getattr(figure.canvas, "manager", None)
+    manager_show = getattr(manager, "show", None)
+    if callable(manager_show):
+        with suppress(AttributeError, RuntimeError, TypeError, ValueError):
+            manager_show()
+    else:
+        figure_show = getattr(figure, "show", None)
+        if callable(figure_show):
+            with suppress(AttributeError, RuntimeError, TypeError, ValueError):
+                figure_show()
+    draw_idle = getattr(figure.canvas, "draw_idle", None)
+    if callable(draw_idle):
+        with suppress(AttributeError, RuntimeError, TypeError, ValueError):
+            draw_idle()
+    flush_events = getattr(figure.canvas, "flush_events", None)
+    if callable(flush_events):
+        with suppress(AttributeError, RuntimeError, TypeError, ValueError):
+            flush_events()
 
 
 class _LinkedTensorInspectorController:
@@ -111,6 +116,10 @@ def __init__(
         self._closing_programmatically: bool = False
         self._close_cid: int | None = None
 
+    @property
+    def is_enabled(self) -> bool:
+        return self._enabled
+
     def bind_viewer(self, viewer: Any) -> None:
         if self._viewer is viewer:
             if self._enabled and self._viewer is not None:
@@ -128,9 +137,11 @@ def bind_viewer(self, viewer: Any) -> None:
                 call_immediately=self._enabled,
             )
 
-    def set_enabled(self, enabled: bool) -> None:
+    def set_enabled(self, enabled: bool, *, reveal: bool = False) -> None:
         target = bool(enabled)
         if target == self._enabled:
+            if target and reveal and self._figure is not None:
+                _reveal_auxiliary_figure(self._figure)
             if target and self._viewer is not None:
                 self._sync_to_step(int(self._viewer.current_step))
             return
@@ -139,6 +150,8 @@ def set_enabled(self, enabled: bool) -> None:
             self._close_figure()
             return
         self._ensure_figure()
+        if reveal and self._figure is not None:
+            _reveal_auxiliary_figure(self._figure)
         if self._viewer is not None:
             self._sync_to_step(int(self._viewer.current_step))
         else:
@@ -295,6 +308,7 @@ def __init__(
         self.figure: Figure | None = None
         self._tensor_inspector: _LinkedTensorInspectorController | None = None
         self._figure_close_cid: int | None = None
+        self._initialized: bool = False
         if self.tensor_inspector_available:
             self._tensor_inspector = _LinkedTensorInspectorController(
                 trace=cast(EinsumTrace, network),
@@ -314,6 +328,7 @@ def initialize(self) -> tuple[Figure, RenderedAxes]:
             return figure, ax
         self._build_controls()
         self._apply_scene_state(self.current_scene)
+        self._initialized = True
         set_interactive_controls(figure, self)
         set_active_axes(figure, ax)
         figure._tensor_network_viz_tensor_inspector = self._tensor_inspector  # type: ignore[attr-defined]
@@ -436,7 +451,7 @@ def _build_controls(self) -> None:
         )
         cb_bottom = float(cb_bounds[1])
         check_ax = self.figure.add_axes(cb_bounds)
-        _style_interactive_control_axes(check_ax)
+        _style_control_tray_axes(check_ax)
         self._check_ax = check_ax
         if not self._external_ax:
             radio_bounds: tuple[float, float, float, float] = (
@@ -446,7 +461,7 @@ def _build_controls(self) -> None:
                 _VIEW_SELECTOR_HEIGHT,
             )
             radio_ax = self.figure.add_axes(radio_bounds)
-            _style_interactive_control_axes(radio_ax)
+            _style_control_tray_axes(radio_ax)
             self._radio_ax = radio_ax
             active_index = 0 if self.current_view == "2d" else 1
             self._radio = RadioButtons(
@@ -569,7 +584,15 @@ def _apply_scene_state(self, scene: _InteractiveSceneState) -> None:
             if self._tensor_inspector is not None:
                 self._tensor_inspector.bind_viewer(controls._viewer)
         if self._tensor_inspector is not None:
-            self._tensor_inspector.set_enabled(self.tensor_inspector_on)
+            reveal_inspector = bool(
+                self._initialized
+                and self.tensor_inspector_on
+                and not self._tensor_inspector.is_enabled
+            )
+            self._tensor_inspector.set_enabled(
+                self.tensor_inspector_on,
+                reveal=reveal_inspector,
+            )
         _apply_scene_hover_state(scene, hover_on=self.hover_on)
         self._sync_checkbuttons()
         if not self._external_ax:

src/tensor_network_viz/_tensor_elements_controller.py

@@ -22,7 +22,7 @@
     _TextSummaryPayload,
     _valid_group_modes_for_record,
 )
-from ._ui_utils import _reserve_figure_bottom
+from ._ui_utils import _reserve_figure_bottom, _style_control_tray_axes
 from .tensor_elements_config import TensorElementsConfig, TensorElementsMode
 
 TensorElementsGroup = Literal["basic", "complex", "diagnostic"]
@@ -32,9 +32,9 @@
 ]
 
 _GROUP_OPTIONS: Final[tuple[TensorElementsGroup, ...]] = ("basic", "complex", "diagnostic")
-_GROUP_SELECTOR_BOUNDS: Final[tuple[float, float, float, float]] = (0.02, 0.048, 0.12, 0.12)
-_MODE_SELECTOR_BOUNDS: Final[tuple[float, float, float, float]] = (0.16, 0.028, 0.18, 0.18)
-_TENSOR_SLIDER_BOUNDS: Final[tuple[float, float, float, float]] = (0.42, 0.058, 0.38, 0.03)
+_GROUP_SELECTOR_BOUNDS: Final[tuple[float, float, float, float]] = (0.02, 0.048, 0.15, 0.12)
+_MODE_SELECTOR_BOUNDS: Final[tuple[float, float, float, float]] = (0.19, 0.028, 0.21, 0.16)
+_TENSOR_SLIDER_BOUNDS: Final[tuple[float, float, float, float]] = (0.48, 0.052, 0.38, 0.05)
 _TENSOR_ELEMENTS_CONTROLS_BOTTOM: Final[float] = 0.24
 _INTERACTIVE_LABEL_PROPS: Final[dict[str, Sequence[Any]]] = {"fontsize": [9.5]}
 _INTERACTIVE_RADIO_PROPS: Final[dict[str, float]] = {"s": 38.0, "linewidth": 0.9}
@@ -209,6 +209,7 @@ def initialize(self, *, show_controls: bool) -> None:
         if show_controls:
             _reserve_figure_bottom(self._figure, _TENSOR_ELEMENTS_CONTROLS_BOTTOM)
             self._group_radio_ax = self._figure.add_axes(_GROUP_SELECTOR_BOUNDS)
+            _style_control_tray_axes(self._group_radio_ax)
             self._group_radio = RadioButtons(
                 self._group_radio_ax,
                 _GROUP_OPTIONS,
@@ -221,6 +222,7 @@ def initialize(self, *, show_controls: bool) -> None:
 
             if len(self._records) > 1:
                 self._slider_ax = self._figure.add_axes(_TENSOR_SLIDER_BOUNDS)
+                _style_control_tray_axes(self._slider_ax)
                 self._slider = Slider(
                     self._slider_ax,
                     "Tensor",
@@ -246,6 +248,7 @@ def _rebuild_mode_radio(self) -> None:
         if self._mode_radio_ax is not None:
             self._mode_radio_ax.remove()
         self._mode_radio_ax = self._figure.add_axes(_MODE_SELECTOR_BOUNDS)
+        _style_control_tray_axes(self._mode_radio_ax)
         mode_options = self._current_group_modes()
         active_index = mode_options.index(self._mode)
         self._mode_radio = RadioButtons(

src/tensor_network_viz/_ui_utils.py

@@ -3,10 +3,14 @@
 from contextlib import suppress
 from typing import Any
 
+from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 
 from ._matplotlib_state import get_reserved_bottom, set_reserved_bottom
 
+_CONTROL_TRAY_FACE: tuple[float, float, float] = (0.97, 0.97, 0.99)
+_CONTROL_TRAY_FRAME: tuple[float, float, float] = (0.78, 0.78, 0.82)
+
 
 def _reserve_figure_bottom(fig: Figure, bottom: float) -> None:
     current = get_reserved_bottom(fig)
@@ -40,9 +44,26 @@ def _set_widget_active(widget: Any, active: bool) -> None:
             setter(bool(active))
 
 
+def _style_control_tray_axes(ax: Axes) -> None:
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.set_navigate(False)
+    ax.patch.set_facecolor(_CONTROL_TRAY_FACE)
+    ax.patch.set_alpha(0.88)
+    ax.patch.set_edgecolor(_CONTROL_TRAY_FRAME)
+    ax.patch.set_linewidth(0.6)
+    for spine in ax.spines.values():
+        spine.set_visible(True)
+        spine.set_linewidth(0.6)
+        spine.set_color(_CONTROL_TRAY_FRAME)
+
+
 __all__ = [
+    "_CONTROL_TRAY_FACE",
+    "_CONTROL_TRAY_FRAME",
     "_reserve_figure_bottom",
     "_set_figure_bottom_reserved",
     "_set_axes_visible",
     "_set_widget_active",
+    "_style_control_tray_axes",
 ]

tests/test_examples.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pytest
 
-from tensor_network_viz import EinsumTrace
+from tensor_network_viz import EinsumTrace, show_tensor_network
 from tensor_network_viz._tensor_elements_data import _extract_einsum_playback_step_records
 
 _EXAMPLES = Path(__file__).resolve().parent.parent / "examples"
@@ -256,6 +256,67 @@ def test_einsum_auto_examples_keep_tensors_alive_for_tensor_inspector(example_na
     assert all(step.record is not None for step in step_records)
 
 
+@pytest.mark.parametrize("example_name", ["mps", "mpo"])
+def test_einsum_auto_examples_keep_inspector_and_costs_aligned_to_real_trace_steps(
+    example_name: str,
+) -> None:
+    _require_torch()
+    run_demo = _load_example_module(
+        Path("examples/run_demo.py"),
+        f"run_demo_{example_name}_inspector_alignment",
+    )
+    demo_cli = importlib.import_module("demo_cli")
+    einsum_demo = importlib.import_module("einsum_demo")
+
+    args = run_demo.parse_args(
+        ["einsum", example_name, "--view", "2d", "--tensor-inspector", "--no-show"]
+    )
+    trace = einsum_demo._trace_steps_for(example_name, args)
+    scheme_steps = einsum_demo._scheme_steps(example_name, args)
+
+    assert isinstance(trace, EinsumTrace)
+    assert scheme_steps is not None
+    assert len(scheme_steps) == len(trace)
+
+    config = demo_cli.finalize_demo_plot_config(
+        args,
+        engine="einsum",
+        scheme_tensor_names=scheme_steps,
+    )
+    fig, _ax = show_tensor_network(
+        trace,
+        engine="einsum",
+        view=args.view,
+        config=config,
+        show=False,
+    )
+
+    controls = getattr(fig, "_tensor_network_viz_interactive_controls", None)
+    assert controls is not None
+    assert controls.current_scene.contraction_controls is not None
+    viewer = controls.current_scene.contraction_controls._viewer
+    assert viewer is not None
+    assert viewer.current_step == len(trace)
+
+    inspector = getattr(fig, "_tensor_network_viz_tensor_inspector", None)
+    assert inspector is not None
+    assert inspector._figure is not None
+    inspector_controls = getattr(
+        inspector._figure,
+        "_tensor_network_viz_tensor_elements_controls",
+        None,
+    )
+    assert inspector_controls is not None
+    assert f"r{len(trace) - 1}" in inspector_controls._panel.main_ax.get_title()
+
+    controls.cost_hover_on = True
+    controls._apply_scene_state(controls.current_scene)
+    assert viewer._cost_panel_ax is not None
+    assert viewer._cost_panel_ax.get_visible()
+    assert viewer._cost_text_artist is not None
+    assert "Contraction:" in viewer._cost_text_artist.get_text()
+
+
 def test_tensornetwork_mera_ttn_saves_figure_without_showing() -> None:
     _require_tensornetwork()
     module = _load_example_module(Path("examples/run_demo.py"), "run_demo_tn_mera_ttn")

tests/test_plotting.py

@@ -16,6 +16,7 @@
 from matplotlib.collections import LineCollection
 
 import tensor_network_viz._core.renderer as core_renderer_module
+import tensor_network_viz._interaction.controller as interaction_controller_module
 import tensor_network_viz.tensorkrowch.graph as tk_graph_module
 import tensor_network_viz.tensorkrowch.renderer as tk_renderer_module
 import tensor_network_viz.tensornetwork.graph as tn_graph_module
@@ -851,6 +852,46 @@ def test_show_tn_einsum_trace_inspector_checkbox_auto_enables_playback() -> None
     assert getattr(fig, "_tensor_network_viz_tensor_inspector", None) is not None
 
 
+def test_show_tn_reenabling_tensor_inspector_reveals_auxiliary_window(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    trace = _build_einsum_trace_for_inspector()
+    revealed: list[matplotlib.figure.Figure] = []
+
+    monkeypatch.setattr(
+        interaction_controller_module,
+        "_reveal_auxiliary_figure",
+        lambda figure: revealed.append(figure),
+    )
+
+    fig, _ax = show_tensor_network(
+        trace,
+        config=PlotConfig(
+            contraction_tensor_inspector=False,
+        ),
+        show=False,
+    )
+
+    controls = getattr(fig, "_tensor_network_viz_interactive_controls", None)
+    assert controls is not None
+    assert controls._checkbuttons is not None
+
+    _click_checkbutton(controls._checkbuttons, 6)
+
+    inspector = getattr(fig, "_tensor_network_viz_tensor_inspector", None)
+    assert inspector is not None
+    assert inspector._figure is not None
+    assert revealed == [inspector._figure]
+
+    _click_checkbutton(controls._checkbuttons, 6)
+    assert inspector._figure is None
+
+    _click_checkbutton(controls._checkbuttons, 6)
+
+    assert inspector._figure is not None
+    assert revealed == [revealed[0], inspector._figure]
+
+
 def test_show_tensor_network_non_einsum_inputs_do_not_expose_tensor_inspector_checkbox() -> None:
     left = DummyTensorKrowchNode("A", ["left"])
     right = DummyTensorKrowchNode("B", ["right"])