Optimize interactive render hot paths

Author: DOKOS-TAYOS

src/tensor_network_viz/_core/draw/contraction_edges.py

@@ -13,6 +13,7 @@
 from .constants import _CURVE_NEAR_PAIR_REF, _CURVE_OFFSET_FACTOR
 from .edge_labels import _plot_contraction_index_captions
 from .fonts_and_scale import _DrawScaleParams
+from .label_descriptors import _TextLabelDescriptor
 from .labels_misc import _contraction_hover_label_text
 from .plotter import _PlotAdapter
 from .scene_state import _RenderedEdgeGeometry
@@ -113,6 +114,7 @@ def _draw_contraction_edge_labels(
     ax: Any,
     scale: float,
     zorder_label: float | None = None,
+    label_sink: list[_TextLabelDescriptor] | None = None,
 ) -> None:
     hover_targets = getattr(plotter, "_hover_edge_targets", None)
     if config.hover_labels and hover_targets is not None:
@@ -139,6 +141,7 @@ def _draw_contraction_edge_labels(
         ax=ax,
         scale=scale,
         zorder_label=zorder_label,
+        label_sink=label_sink,
     )
 
 

src/tensor_network_viz/_core/draw/dangling_self_edges.py

@@ -22,6 +22,7 @@
     _node_label_clearance,
     _self_loop_hover_label_text,
 )
+from .label_descriptors import _TextLabelDescriptor
 from .plotter import _PlotAdapter
 from .scene_state import _RenderedEdgeGeometry
 from .vectors import _perpendicular_2d
@@ -119,6 +120,7 @@ def _draw_dangling_edge_labels(
     ax: Any,
     scale: float,
     zorder_label: float | None = None,
+    label_sink: list[_TextLabelDescriptor] | None = None,
 ) -> None:
     if not edge.label and not config.hover_labels:
         return
@@ -151,6 +153,11 @@ def _draw_dangling_edge_labels(
         return
 
     raw_label = edge.label
+    stub_segment = np.stack([np.asarray(start, dtype=float), np.asarray(end, dtype=float)], axis=0)
+    stub_length = _polyline_arc_length_total(stub_segment)
+    distance_from_tip = float(_PHYS_DANGLING_2D_FRAC_FROM_TIP) * stub_length
+    point, tangent = _point_tangent_along_polyline_from_end(stub_segment, distance_from_tip)
+
     fontsize = _edge_index_fontsize_for_bond(
         raw_label,
         bond_start=start,
@@ -166,10 +173,6 @@ def _draw_dangling_edge_labels(
         bbox_pad=p.index_bbox_pad,
         zorder=zorder_label,
     )
-    stub_segment = np.stack([np.asarray(start, dtype=float), np.asarray(end, dtype=float)], axis=0)
-    stub_length = _polyline_arc_length_total(stub_segment)
-    distance_from_tip = float(_PHYS_DANGLING_2D_FRAC_FROM_TIP) * stub_length
-    point, tangent = _point_tangent_along_polyline_from_end(stub_segment, distance_from_tip)
     if dimensions == 2:
         start_2d = np.asarray(start[:2], dtype=float)
         end_2d = np.asarray(end[:2], dtype=float)
@@ -203,11 +206,18 @@ def _draw_dangling_edge_labels(
             scale=scale,
             fontsize_pt=float(fontsize),
         )
-    plotter.plot_text(
-        label_pos,
-        format_tensor_node_label(raw_label),
-        **{**text_kwargs, **align_kwargs},
-    )
+    formatted = format_tensor_node_label(raw_label)
+    kwargs = {**text_kwargs, **align_kwargs}
+    if label_sink is not None:
+        label_sink.append(
+            _TextLabelDescriptor(
+                position=np.asarray(label_pos, dtype=float).copy(),
+                text=formatted,
+                kwargs=dict(kwargs),
+            )
+        )
+        return
+    plotter.plot_text(label_pos, formatted, **kwargs)
 
 
 def _draw_self_loop_edge(
@@ -302,6 +312,7 @@ def _draw_self_loop_edge_labels(
     ax: Any,
     scale: float,
     zorder_label: float | None = None,
+    label_sink: list[_TextLabelDescriptor] | None = None,
 ) -> None:
     hover_targets = getattr(plotter, "_hover_edge_targets", None)
     if config.hover_labels and hover_targets is not None:
@@ -391,11 +402,22 @@ def _draw_self_loop_edge_labels(
                 dimensions=dimensions,
             ),
         }
-        plotter.plot_text(
-            np.asarray(q_a, dtype=float) + offset_a,
-            format_tensor_node_label(caption_a),
-            **text_kwargs_a,
-        )
+        position_a = np.asarray(q_a, dtype=float) + offset_a
+        formatted_a = format_tensor_node_label(caption_a)
+        if label_sink is not None:
+            label_sink.append(
+                _TextLabelDescriptor(
+                    position=np.asarray(position_a, dtype=float).copy(),
+                    text=formatted_a,
+                    kwargs=dict(text_kwargs_a),
+                )
+            )
+        else:
+            plotter.plot_text(
+                position_a,
+                formatted_a,
+                **text_kwargs_a,
+            )
     if caption_b:
         offset_b = (
             -direction_unit
@@ -433,9 +455,20 @@ def _draw_self_loop_edge_labels(
                 dimensions=dimensions,
             ),
         }
+        formatted_b = format_tensor_node_label(caption_b)
+        position_b = np.asarray(q_b, dtype=float) + offset_b
+        if label_sink is not None:
+            label_sink.append(
+                _TextLabelDescriptor(
+                    position=np.asarray(position_b, dtype=float).copy(),
+                    text=formatted_b,
+                    kwargs=dict(text_kwargs_b),
+                )
+            )
+            return
         plotter.plot_text(
-            np.asarray(q_b, dtype=float) + offset_b,
-            format_tensor_node_label(caption_b),
+            position_b,
+            formatted_b,
             **text_kwargs_b,
         )
 

src/tensor_network_viz/_core/draw/edge_labels.py

@@ -16,6 +16,7 @@
 )
 from ..layout import NodePositions
 from .fonts_and_scale import _DrawScaleParams
+from .label_descriptors import _TextLabelDescriptor
 from .labels_misc import _edge_index_text_kwargs
 from .plotter import _PlotAdapter
 from .viewport_geometry import (
@@ -46,6 +47,7 @@ def _plot_contraction_index_captions(
     ax: Any,
     scale: float,
     zorder_label: float | None = None,
+    label_sink: list[_TextLabelDescriptor] | None = None,
 ) -> None:
     ep_l, ep_r = _require_contraction_endpoints(edge)
     cap_l: str | None = _endpoint_index_caption(ep_l, edge, graph)
@@ -129,11 +131,18 @@ def _plot_contraction_index_captions(
                 scale=scale,
                 fontsize_pt=float(fontsize),
             )
-        plotter.plot_text(
-            position,
-            format_tensor_node_label(cap),
-            **{**text_kwargs, **align_kwargs},
-        )
+        formatted = format_tensor_node_label(cap)
+        kwargs = {**text_kwargs, **align_kwargs}
+        if label_sink is not None:
+            label_sink.append(
+                _TextLabelDescriptor(
+                    position=np.asarray(position, dtype=float).copy(),
+                    text=formatted,
+                    kwargs=dict(kwargs),
+                )
+            )
+            continue
+        plotter.plot_text(position, formatted, **kwargs)
 
 
 __all__ = ["_plot_contraction_index_captions"]

src/tensor_network_viz/_core/draw/fonts_and_scale.py

@@ -152,16 +152,17 @@ def _on_2d_limits_changed(ax: Axes) -> None:
 
 
 def _register_2d_zoom_font_scaling(ax: Axes) -> None:
-    old_cids = get_zoom_cids(ax)
-    for cid in old_cids:
-        with suppress(ValueError, KeyError):
-            ax.callbacks.disconnect(cid)
     x0, x1 = ax.get_xlim()
     y0, y1 = ax.get_ylim()
     ref_span = max(float(x1 - x0), float(y1 - y0), 1e-9)
     sizes = {t: float(t.get_fontsize()) for t in ax.texts}
     set_zoom_font_state(ax, ref_span=ref_span, sizes=sizes)
 
+    old_cids = get_zoom_cids(ax)
+    if old_cids:
+        _on_2d_limits_changed(ax)
+        return
+
     def _cb(_: object) -> None:
         _on_2d_limits_changed(ax)
 

src/tensor_network_viz/_core/draw/label_descriptors.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal, TypeAlias
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class _TextLabelDescriptor:
+    position: np.ndarray
+    text: str
+    kwargs: dict[str, Any]
+    node_id: int | None = None
+
+
+@dataclass(frozen=True)
+class _DeferredBondLabelDescriptor:
+    text: str
+    point: np.ndarray
+    tangent_geom: np.ndarray
+    tangent_align: np.ndarray
+    bond_start: np.ndarray
+    bond_end: np.ndarray
+    text_endpoint: Literal["left", "right"]
+    stub_kind: Literal["bond", "dangling"]
+    is_physical: bool
+    peer_captions_for_width: tuple[str, ...] | None = None
+    zorder: float | None = None
+
+
+@dataclass(frozen=True)
+class _DeferredSelfLoopLabelDescriptor:
+    text: str
+    point: np.ndarray
+    tangent: np.ndarray
+    bond_start: np.ndarray
+    bond_end: np.ndarray
+    offset_direction: np.ndarray
+    offset_scale: float
+    text_endpoint: Literal["left", "right"]
+    peer_captions_for_width: tuple[str, ...] | None = None
+    zorder: float | None = None
+
+
+_AnyLabelDescriptor: TypeAlias = (
+    _TextLabelDescriptor | _DeferredBondLabelDescriptor | _DeferredSelfLoopLabelDescriptor
+)
+
+
+__all__ = [
+    "_AnyLabelDescriptor",
+    "_DeferredBondLabelDescriptor",
+    "_DeferredSelfLoopLabelDescriptor",
+    "_TextLabelDescriptor",
+]

src/tensor_network_viz/_core/draw/scene_state.py

@@ -12,6 +12,7 @@
 from ..layout import AxisDirections, NodePositions
 from .fonts_and_scale import _DrawScaleParams
 from .hover import _RenderHoverState
+from .label_descriptors import _AnyLabelDescriptor, _TextLabelDescriptor
 from .plotter import _PlotAdapter
 
 
@@ -40,6 +41,8 @@ class _InteractiveSceneState:
     tensor_disk_radius_px_3d: float | None
     tensor_label_artists: list[Artist] = field(default_factory=list)
     edge_label_artists: list[Artist] = field(default_factory=list)
+    tensor_label_descriptors: tuple[_AnyLabelDescriptor, ...] | None = None
+    edge_label_descriptors: tuple[_AnyLabelDescriptor, ...] | None = None
     tensor_hover_payload: dict[int, tuple[str, float]] | None = None
     edge_hover_payload: tuple[tuple[np.ndarray, str], ...] | None = None
     contraction_controls: Any = None
@@ -48,4 +51,6 @@ class _InteractiveSceneState:
 __all__ = [
     "_InteractiveSceneState",
     "_RenderedEdgeGeometry",
+    "_AnyLabelDescriptor",
+    "_TextLabelDescriptor",
 ]

src/tensor_network_viz/_core/draw/tensors.py

@@ -22,6 +22,7 @@
 )
 from .disk_metrics import _tensor_disk_radius_px
 from .fonts_and_scale import _DrawScaleParams
+from .label_descriptors import _TextLabelDescriptor
 from .plotter import _PlotAdapter, _visible_degree_one_mask
 from .viewport_geometry import _stack_visible_tensor_coords
 
@@ -97,8 +98,9 @@ def _draw_labels(
     visible_draw_order: list[int] | None = None,
     tensor_label_zorder_by_node: dict[int, float] | None = None,
     tensor_disk_radius_px_3d: float | None = None,
+    label_sink: list[_TextLabelDescriptor] | None = None,
 ) -> None:
-    if show_tensor_labels or tensor_hover_by_node is not None:
+    if show_tensor_labels or tensor_hover_by_node is not None or label_sink is not None:
         fig = ax.figure
         ordered_ids: list[int]
         if visible_draw_order is not None:
@@ -124,14 +126,30 @@ def _draw_labels(
             if dimensions == 3:
                 cap_tensor = float(p.font_tensor_label_max) * _LABEL_FONT_3D_SCALE
                 fs = min(float(fs) * _LABEL_FONT_3D_SCALE, cap_tensor)
-            if tensor_hover_by_node is not None:
-                tensor_hover_by_node[node_id] = (display_name, float(fs))
-            if not show_tensor_labels:
-                continue
             if tensor_label_zorder_by_node is None:
                 z_lbl = float(_ZORDER_TENSOR_NAME)
             else:
                 z_lbl = float(tensor_label_zorder_by_node.get(node_id, _ZORDER_TENSOR_NAME))
+            if tensor_hover_by_node is not None:
+                tensor_hover_by_node[node_id] = (display_name, float(fs))
+            if label_sink is not None:
+                label_sink.append(
+                    _TextLabelDescriptor(
+                        position=np.asarray(pos, dtype=float).copy(),
+                        text=display_name,
+                        kwargs={
+                            "color": config.tensor_label_color,
+                            "ha": "center",
+                            "va": "center",
+                            "fontsize": float(fs),
+                            "zorder": float(z_lbl),
+                            "gid": _TENSOR_LABEL_GID,
+                        },
+                        node_id=int(node_id),
+                    )
+                )
+            if not show_tensor_labels:
+                continue
             plotter.plot_text(
                 pos,
                 display_name,