tensor-network-visualization 1.0.0__py3-none-any.whl

tensor_network_viz/tensorkrowch/draw_3d.py

@@ -0,0 +1,202 @@
+"""3D drawing for tensor networks."""
+
+from __future__ import annotations
+
+import numpy as np
+from mpl_toolkits.mplot3d.axes3d import Axes3D
+
+from ..config import PlotConfig
+from ._draw_common import _draw_scale_params
+from .curves import (
+    _ellipse_points_3d,
+    _group_contractions,
+    _quadratic_curve,
+    _require_self_endpoints,
+)
+from .graph import _GraphData
+from .layout import AxisDirections, NodePositions, _orthogonal_unit
+
+
+def _curved_edge_points_3d(
+    *,
+    start: np.ndarray,
+    end: np.ndarray,
+    offset_index: int,
+    edge_count: int,
+    scale: float = 1.0,
+) -> np.ndarray:
+    midpoint = (start + end) / 2.0
+    delta = end - start
+    distance = max(float(np.linalg.norm(delta)), 1e-6)
+    direction = delta / distance
+    reference = np.array([0.0, 0.0, 1.0], dtype=float)
+    perpendicular = np.cross(direction, reference)
+    if np.linalg.norm(perpendicular) < 1e-6:
+        perpendicular = np.cross(direction, np.array([0.0, 1.0, 0.0], dtype=float))
+    perpendicular = perpendicular / np.linalg.norm(perpendicular)
+    offset = (offset_index - (edge_count - 1) / 2.0) * 0.18 * scale * distance
+    control = midpoint + perpendicular * offset
+    return _quadratic_curve(start, control, end)
+
+
+def _draw_3d(
+    *,
+    ax: Axes3D,
+    graph: _GraphData,
+    positions: NodePositions,
+    directions: AxisDirections,
+    show_tensor_labels: bool,
+    show_index_labels: bool,
+    config: PlotConfig,
+    scale: float = 1.0,
+) -> None:
+    ax.cla()
+    pair_groups = _group_contractions(graph)
+    p = _draw_scale_params(config, scale, is_3d=True)
+
+    for edge in graph.edges:
+        if edge.kind == "dangling":
+            endpoint = edge.endpoints[0]
+            direction = directions[(endpoint.node_id, endpoint.axis_index)]
+            start = positions[endpoint.node_id] + direction * p.r
+            end = start + direction * p.stub
+            ax.plot(
+                [start[0], end[0]],
+                [start[1], end[1]],
+                [start[2], end[2]],
+                color=config.dangling_edge_color,
+                linewidth=p.lw,
+                zorder=2,
+            )
+            if show_index_labels and edge.label:
+                label_pos = end + direction * p.label_offset
+                ax.text(
+                    label_pos[0],
+                    label_pos[1],
+                    label_pos[2],
+                    edge.label,
+                    color=config.label_color,
+                    fontsize=p.font_dangling,
+                    zorder=5,
+                    ha="center",
+                    va="bottom",
+                )
+        elif edge.kind == "self":
+            endpoint_a, endpoint_b = _require_self_endpoints(edge)
+            direction_a = directions[(endpoint_a.node_id, endpoint_a.axis_index)]
+            direction_b = directions[(endpoint_b.node_id, endpoint_b.axis_index)]
+            orientation = direction_a + direction_b
+            if np.linalg.norm(orientation) < 1e-6:
+                orientation = np.array([1.0, 0.0, 0.0], dtype=float)
+            orientation = orientation / np.linalg.norm(orientation)
+            normal = _orthogonal_unit(orientation)
+            binormal = np.cross(orientation, normal)
+            binormal = binormal / np.linalg.norm(binormal)
+            center_pt = (
+                positions[endpoint_a.node_id]
+                + orientation * (p.r + p.loop_r)
+            )
+            curve = _ellipse_points_3d(
+                center_pt, normal, binormal, width=p.ellipse_w, height=p.ellipse_h
+            )
+            ax.plot(
+                curve[:, 0],
+                curve[:, 1],
+                curve[:, 2],
+                color=config.bond_edge_color,
+                linewidth=p.lw,
+                zorder=2,
+            )
+            if show_index_labels and edge.label:
+                label_pos = center_pt + binormal * p.ellipse_w
+                ax.text(
+                    label_pos[0],
+                    label_pos[1],
+                    label_pos[2],
+                    edge.label,
+                    color=config.label_color,
+                    fontsize=p.font_dangling,
+                    zorder=5,
+                    ha="center",
+                    va="bottom",
+                )
+        else:
+            key = tuple(sorted(edge.node_ids))
+            group = pair_groups[key]
+            offset_index = group.index(edge)
+            curve = _curved_edge_points_3d(
+                start=positions[edge.node_ids[0]],
+                end=positions[edge.node_ids[1]],
+                offset_index=offset_index,
+                edge_count=len(group),
+                scale=scale,
+            )
+            ax.plot(
+                curve[:, 0],
+                curve[:, 1],
+                curve[:, 2],
+                color=config.bond_edge_color,
+                linewidth=p.lw,
+                zorder=1,
+            )
+            if show_index_labels and edge.label:
+                midpoint = curve[len(curve) // 2]
+                delta = positions[edge.node_ids[1]] - positions[edge.node_ids[0]]
+                dist = max(float(np.linalg.norm(delta)), 1e-6)
+                direction = delta / dist
+                perpendicular = np.cross(direction, np.array([0.0, 0.0, 1.0], dtype=float))
+                if np.linalg.norm(perpendicular) < 1e-6:
+                    perpendicular = np.cross(direction, np.array([0.0, 1.0, 0.0], dtype=float))
+                perpendicular = perpendicular / np.linalg.norm(perpendicular)
+                if perpendicular[2] < 0:
+                    perpendicular = -perpendicular
+                label_pos = midpoint + perpendicular * p.label_offset
+                ax.text(
+                    label_pos[0],
+                    label_pos[1],
+                    label_pos[2],
+                    edge.label,
+                    color=config.label_color,
+                    fontsize=p.font_bond,
+                    zorder=5,
+                    ha="center",
+                    va="bottom",
+                )
+
+    coords = np.stack(list(positions.values()))
+    ax.scatter(
+        coords[:, 0],
+        coords[:, 1],
+        coords[:, 2],
+        s=p.scatter_s,
+        c=config.node_color,
+        edgecolors=config.node_edge_color,
+        linewidths=p.lw,
+        depthshade=False,
+    )
+
+    if show_tensor_labels:
+        for node_id, node in graph.nodes.items():
+            x, y, z = positions[node_id]
+            ax.text(
+                x, y, z,
+                node.name,
+                color=config.tensor_label_color,
+                fontsize=p.font_node,
+                ha="center",
+                va="center",
+                zorder=5,
+            )
+
+    _style_3d_axes(ax, coords)
+
+
+def _style_3d_axes(ax: Axes3D, coords: np.ndarray) -> None:
+    span = np.ptp(coords, axis=0)
+    span = np.maximum(span, 1.0)
+    center = coords.mean(axis=0)
+    ax.set_xlim(center[0] - span[0] * 0.9, center[0] + span[0] * 0.9)
+    ax.set_ylim(center[1] - span[1] * 0.9, center[1] + span[1] * 0.9)
+    ax.set_zlim(center[2] - span[2] * 0.9, center[2] + span[2] * 0.9)
+    ax.set_box_aspect(span)
+    ax.set_axis_off()

tensor_network_viz/tensorkrowch/graph.py

@@ -0,0 +1,190 @@
+"""Graph data structures and construction from tensor networks."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+EdgeKind = Literal["contraction", "dangling", "self"]
+
+
+@dataclass(frozen=True)
+class _EdgeEndpoint:
+    node_id: int
+    axis_index: int
+    axis_name: str | None
+
+
+@dataclass(frozen=True)
+class _NodeData:
+    name: str
+    axes_names: tuple[str, ...]
+    degree: int
+
+
+@dataclass(frozen=True)
+class _EdgeData:
+    name: str | None
+    kind: EdgeKind
+    node_ids: tuple[int, ...]
+    endpoints: tuple[_EdgeEndpoint, ...]
+    label: str | None
+
+
+@dataclass(frozen=True)
+class _GraphData:
+    nodes: dict[int, _NodeData]
+    edges: tuple[_EdgeData, ...]
+
+
+def _get_network_nodes(network: Any) -> list[Any]:
+    if hasattr(network, "nodes"):
+        raw_nodes = network.nodes
+    elif hasattr(network, "leaf_nodes"):
+        raw_nodes = network.leaf_nodes
+    else:
+        raise TypeError(
+            "Tensor network must expose either a 'nodes' attribute or a 'leaf_nodes' attribute."
+        )
+
+    iterable = raw_nodes.values() if isinstance(raw_nodes, dict) else raw_nodes
+
+    try:
+        items = list(iterable)
+    except TypeError as exc:
+        raise TypeError("Tensor network nodes must be iterable.") from exc
+
+    unique_nodes: list[Any] = []
+    seen: set[int] = set()
+    for node in items:
+        if node is None:
+            continue
+        node_id = id(node)
+        if node_id in seen:
+            continue
+        seen.add(node_id)
+        unique_nodes.append(node)
+    return unique_nodes
+
+
+def _iterable_attr(obj: Any, attr_name: str, object_name: str) -> list[Any]:
+    value = _require_attr(obj, attr_name, object_name)
+    if isinstance(value, dict):
+        return list(value.values())
+    try:
+        return list(value)
+    except TypeError as exc:
+        msg = f"{object_name.capitalize()} attribute '{attr_name}' must be iterable."
+        raise TypeError(msg) from exc
+
+
+def _require_attr(obj: Any, attr_name: str, object_name: str) -> Any:
+    if not hasattr(obj, attr_name):
+        raise TypeError(f"{object_name.capitalize()} is missing required attribute '{attr_name}'.")
+    return getattr(obj, attr_name)
+
+
+def _stringify(value: Any) -> str:
+    return "" if value is None else str(value)
+
+
+def _optional_string(value: Any) -> str | None:
+    if value is None:
+        return None
+    text = str(value)
+    return text or None
+
+
+def _build_edge_label(
+    kind: EdgeKind,
+    endpoints: tuple[_EdgeEndpoint, ...],
+    edge_name: str | None,
+) -> str | None:
+    axis_names = [endpoint.axis_name for endpoint in endpoints if endpoint.axis_name]
+    if kind == "dangling":
+        return axis_names[0] if axis_names else edge_name
+    if len(axis_names) >= 2:
+        return f"{axis_names[0]}<->{axis_names[1]}"
+    return edge_name
+
+
+def _build_graph(network: Any) -> _GraphData:
+    node_refs = _get_network_nodes(network)
+    if not node_refs:
+        raise ValueError("The tensor network does not expose any nodes to visualize.")
+
+    nodes: dict[int, _NodeData] = {}
+    edge_refs: dict[int, Any] = {}
+    edge_endpoints: dict[int, list[_EdgeEndpoint]] = {}
+
+    for node in node_refs:
+        name = _stringify(_require_attr(node, "name", "node"))
+        edges = tuple(_iterable_attr(node, "edges", "node"))
+        axes_names = tuple(_stringify(item) for item in _iterable_attr(node, "axes_names", "node"))
+        if len(edges) != len(axes_names):
+            raise TypeError(
+                f"Node {name!r} has {len(edges)} edges but {len(axes_names)} axes_names."
+            )
+
+        node_id = id(node)
+        nodes[node_id] = _NodeData(
+            name=name,
+            axes_names=axes_names,
+            degree=len(edges),
+        )
+
+        for axis_index, edge in enumerate(edges):
+            if edge is None:
+                continue
+            edge_id = id(edge)
+            edge_refs[edge_id] = edge
+            edge_endpoints.setdefault(edge_id, []).append(
+                _EdgeEndpoint(
+                    node_id=node_id,
+                    axis_index=axis_index,
+                    axis_name=axes_names[axis_index],
+                )
+            )
+
+    edges: list[_EdgeData] = []
+    for edge_id, edge in edge_refs.items():
+        name = _optional_string(_require_attr(edge, "name", "edge"))
+        node1 = _require_attr(edge, "node1", "edge")
+        node2 = _require_attr(edge, "node2", "edge")
+
+        node1_id = id(node1) if node1 is not None and id(node1) in nodes else None
+        node2_id = id(node2) if node2 is not None and id(node2) in nodes else None
+        endpoints = tuple(edge_endpoints.get(edge_id, ()))
+        if not endpoints:
+            continue
+        if len(endpoints) > 2:
+            raise TypeError("Edges with more than two endpoints are not supported.")
+
+        if node1_id is not None and node2_id is not None:
+            kind: EdgeKind
+            if node1_id == node2_id:
+                kind = "self"
+                node_ids = (node1_id,)
+            else:
+                kind = "contraction"
+                node_ids = (node1_id, node2_id)
+        elif node1_id is not None:
+            kind = "dangling"
+            node_ids = (node1_id,)
+        elif node2_id is not None:
+            kind = "dangling"
+            node_ids = (node2_id,)
+        else:
+            continue
+
+        edges.append(
+            _EdgeData(
+                name=name,
+                kind=kind,
+                node_ids=node_ids,
+                endpoints=endpoints,
+                label=_build_edge_label(kind=kind, endpoints=endpoints, edge_name=name),
+            )
+        )
+
+    return _GraphData(nodes=nodes, edges=tuple(edges))