Source code for andrey.viz

"""One visual grammar for causal graphs.

:func:`draw` renders a :class:`~andrey.core.structure.GraphStructure` as a self-contained SVG.
Marks are drawn by shape (arrowhead, hollow circle, tail), and color is kept for uncertain edges
and latent nodes, so a drawing reads without color. Graphviz places the nodes when pygraphviz (the
``viz`` extra) is installed, and a built-in layering otherwise; edges bend around nodes.

Examples
--------
>>> import numpy as np
>>> from andrey.core.structure import GraphStructure, TAIL, ARROW
>>> import andrey.viz
>>> M = np.zeros((2, 2), np.int8); M[0, 1] = TAIL; M[1, 0] = ARROW  # X -> Y
>>> svg = andrey.viz.draw(GraphStructure.from_numpy(M, kind="dag", labels=["X", "Y"]))
>>> svg.startswith("<svg") and "polygon" in svg
True
"""

from __future__ import annotations

import math
from collections.abc import Iterable, Mapping

import numpy as np

from andrey.core.structure import (
    ARROW,
    CIRCLE,
    KINDS,
    LATENT,
    TAIL,
    GraphStructure,
    TemporalStructure,
)

from . import _graphviz, _layout, _svg
from ._graphviz import to_graphviz
from .theme import DARK, LIGHT, Palette

__all__ = ["draw", "layout", "to_graphviz", "classify_edge", "SVG", "Palette", "LIGHT", "DARK"]
ENGINES = ("auto", "builtin", *_graphviz.ENGINES)


[docs] class SVG(str): """An SVG document that renders itself in notebooks and doc pages. A plain :class:`str` subclass, so it slices, concatenates, and writes to a file exactly like the markup it holds. The only addition is ``_repr_html_``, which Jupyter calls to display the graph inline instead of printing several kilobytes of markup. Examples -------- >>> from andrey.viz import SVG >>> svg = SVG("<svg></svg>") >>> isinstance(svg, str) and svg._repr_html_() == svg True """ __slots__ = () def _repr_html_(self) -> str: """Return the markup itself: an SVG document is already valid inline HTML.""" return str(self)
# The smallest canvas a default drawing gets, in pixels. _MIN_WIDTH, _MIN_HEIGHT = 360, 240 def _scale(unit: dict[int, tuple[float, float]], widths: list[float]) -> float: """Pixels per unit of ``unit`` that keep every pair of nodes clear of each other. Two nodes are clear when they are a column apart across, wide enough for both labels, or a row apart down. """ if len(unit) < 2: return 1.0 points = np.array([unit[k] for k in range(len(unit))]) dx = np.abs(points[:, None, 0] - points[None, :, 0]) dy = np.abs(points[:, None, 1] - points[None, :, 1]) w = np.asarray(widths) across = np.maximum(_layout.COLUMN, (w[:, None] + w[None, :]) / 2 + 12) with np.errstate(divide="ignore"): need = np.minimum( np.where(dx > 0, across / dx, np.inf), np.where(dy > 0, _layout.ROW / dy, np.inf) ) need = need[np.triu_indices(len(unit), 1)] need = need[np.isfinite(need)] # nodes on the same point cannot be told apart by scaling return float(need.max()) if need.size else 1.0
[docs] def classify_edge(mark_i: int, mark_j: int, kind: str) -> str: """Name the edge type from its endpoint-mark pair and the graph kind. The name selects the edge's color and mark shapes: ``directed`` is the graphite backbone, and the rest carry the palette's uncertainty colors. Parameters ---------- mark_i, mark_j : int The endpoint marks at nodes ``i`` and ``j``: ``TAIL``, ``ARROW``, or ``CIRCLE`` from :mod:`andrey.core.structure`. kind : str The graph kind (``dag`` / ``cpdag`` / ``pag`` / ``digraph``). Only ``ARROW``/``ARROW`` reads differently across kinds. Returns ------- str ``"directed"`` (``->``), ``"undirected"`` (``--``), ``"bidirected"`` (``<->`` outside a ``digraph``), ``"twocycle"`` (``<->`` in a ``digraph``: two directed edges), or ``"partial"`` (an edge with a circle end: ``o->``, ``-o``, or ``o-o``). Raises ------ ValueError If a mark is not ``TAIL``, ``ARROW``, or ``CIRCLE``, or ``kind`` is not one of the four kinds. Notes ----- ``ARROW``/``ARROW`` is kind-relative: a ``bidirected`` edge (latent confounder) under a PAG, a ``twocycle`` (feedback, two directed edges) under a directed graph. The two are drawn differently on purpose, so a reader never mistakes feedback for a hidden common cause. Examples -------- >>> from andrey.core.structure import ARROW, TAIL >>> from andrey.viz import classify_edge >>> classify_edge(TAIL, ARROW, "dag") 'directed' >>> classify_edge(ARROW, ARROW, "pag"), classify_edge(ARROW, ARROW, "digraph") ('bidirected', 'twocycle') """ if kind not in KINDS: raise ValueError(f"kind must be one of {KINDS}, got {kind!r}") marks = {int(mark_i), int(mark_j)} if not marks <= {TAIL, ARROW, CIRCLE}: raise ValueError( f"marks must be TAIL ({TAIL}), ARROW ({ARROW}), or CIRCLE ({CIRCLE}), " f"got ({mark_i}, {mark_j})" ) if marks == {TAIL, ARROW}: return "directed" if int(mark_i) == TAIL and int(mark_j) == TAIL: return "undirected" if int(mark_i) == ARROW and int(mark_j) == ARROW: return "twocycle" if kind == "digraph" else "bidirected" return "partial"
def _require_graph(structure: object) -> None: """Raise ``TypeError`` unless ``structure`` is one graph; a temporal result names its views.""" if isinstance(structure, GraphStructure): return if isinstance(structure, TemporalStructure): raise TypeError( "a TemporalStructure holds one graph per lag; draw one of them with " "structure.lag(k), or all lags merged with structure.summary_graph()" ) raise TypeError(f"expected a GraphStructure, got {type(structure).__name__}") def _endpoints(structure: GraphStructure) -> list[tuple[int, int, int, int, str]]: """List each edge as ``(i, j, mark_i, mark_j, edge_type)``.""" edges = [] for row in structure.to_edges(): i, j = int(row["i"]), int(row["j"]) mi, mj = int(row["mark_i"]), int(row["mark_j"]) edges.append((i, j, mi, mj, classify_edge(mi, mj, structure.kind))) return edges def _builtin(n: int, edges: list[tuple[int, int, int, int, str]]) -> dict[int, tuple[float, float]]: """Layered by the directed edges, or on a circle.""" directed: list[tuple[int, int]] = [] for i, j, _, mj, et in edges: if i == j: continue if et == "directed": directed.append((i, j) if mj == ARROW else (j, i)) # arrow end is the child elif et == "twocycle": directed.append((i, j)) directed.append((j, i)) placed = _layout.layered(n, directed) if directed else None return _layout.circular(n) if placed is None else placed
[docs] def layout( structure: GraphStructure, engine: str = "auto", options: Mapping[str, object] | None = None ) -> dict[int, tuple[float, float]]: """Return the node positions :func:`draw` uses when none are given. Pass the result as ``positions=`` to draw several structures over the same nodes, such as an estimate and its truth, with every node in the same place. Parameters ---------- structure : GraphStructure Any kind. engine : str, default="auto" ``"dot"`` lays the nodes out in ranks along the directed edges, causes on the left; ``"neato"``, ``"fdp"``, ``"sfdp"``, ``"circo"``, ``"twopi"``, ``"osage"``, and ``"patchwork"`` are Graphviz's other layouts. These need pygraphviz, the ``viz`` extra. ``"builtin"`` needs nothing: nodes in layers along the directed edges, or on a circle when there are none or they form a cycle. ``"auto"`` is ``"dot"`` when pygraphviz is installed, and ``"builtin"`` otherwise. options : mapping of str to object or None, default=None Graphviz graph attributes, such as ``{"rankdir": "TB"}`` to run causes top to bottom, or ``{"nodesep": 1}`` to spread nodes within a rank. The built-in layout takes none. Returns ------- dict of int to (float, float) Each node's ``(x, y)`` in the unit square, ``y`` downward, in the layout's proportions: the longer side spans the square, and the shorter one is centered. Raises ------ TypeError If ``structure`` is not a ``GraphStructure``, such as a ``TemporalStructure``. ValueError If ``engine`` is not one of the above, or the built-in layout is given ``options``. ImportError If a Graphviz engine is named and pygraphviz is not installed. Examples -------- >>> import numpy as np >>> from andrey.core.structure import GraphStructure, TAIL, ARROW >>> import andrey.viz >>> M = np.zeros((2, 2), np.int8); M[0, 1] = TAIL; M[1, 0] = ARROW # X -> Y >>> andrey.viz.layout(GraphStructure.from_numpy(M, kind="dag"), engine="builtin") {0: (0.0, 0.5), 1: (1.0, 0.5)} """ _require_graph(structure) if engine not in ENGINES: raise ValueError(f"engine must be one of {', '.join(ENGINES)}, not {engine!r}.") if engine == "auto": engine = "dot" if _graphviz.available() else "builtin" if engine == "builtin": if options: raise ValueError("the built-in layout takes no options; they are Graphviz's.") placed = _builtin(int(structure.n_nodes), _endpoints(structure)) else: placed = _graphviz.positions(structure, engine, options) return _layout.unit(placed)
[docs] def draw( structure: GraphStructure, *, dark: bool = False, palette: Palette | None = None, width: int | None = None, height: int | None = None, background: bool = True, positions: Mapping[int, tuple[float, float]] | None = None, engine: str = "auto", engine_options: Mapping[str, object] | None = None, highlight: Iterable[tuple[int, int]] = (), ) -> SVG: """Render a causal structure to an SVG string. Parameters ---------- structure : GraphStructure Any kind (``dag`` / ``cpdag`` / ``pag`` / ``digraph``). Its ``.node_types`` flag latent variables, which are drawn hollow and dashed. If every label has at most two characters, all labels sit inside nodes; otherwise, all labels sit below them. dark : bool, default=False Select the built-in dark palette instead of the light one. Ignored when ``palette`` is set. palette : Palette or None, default=None A :class:`Palette` (``andrey.viz.LIGHT`` / ``DARK`` or a custom one); overrides ``dark``. width : int or None, default=None Canvas width in pixels. ``None`` sizes the canvas to the layout, with every node and label clear of the others. Either way the layout keeps its proportions and is centered. The SVG is fluid (``width="100%"``); this sets the coordinate system. A width too narrow for the widest label and the margins is widened to the narrowest that holds them. height : int or None, default=None Canvas height; ``None`` sizes it to the layout. A height too short for the margins is raised the same way. background : bool, default=True Paint the canvas with the palette's paper color. Set ``False`` for a transparent SVG to composite onto an already-themed page. positions : mapping of int to (float, float) or None, default=None Each node's ``(x, y)`` in the unit square, ``y`` downward, as :func:`layout` returns. ``None`` uses :func:`layout` with ``engine`` and ``engine_options``. engine : str, default="auto" The layout :func:`layout` runs when ``positions`` is ``None``: ``"dot"`` and Graphviz's other layouts, or ``"builtin"``. ``"auto"`` is ``"dot"`` when pygraphviz is installed. engine_options : mapping of str to object or None, default=None Graphviz graph attributes for the layout, such as ``{"rankdir": "TB"}``. highlight : iterable of (int, int), default=() Endpoints to light in the palette's ``accent`` color. ``(i, j)`` lights the mark at node ``i`` on the edge between nodes ``i`` and ``j``, and that half of the edge, with a soft glow under it (class ``highlight``, for a page to style or animate); add ``(j, i)`` to light both ends. The default lights nothing. Returns ------- SVG A self-contained SVG document (no external dependencies or references). :class:`SVG` is a :class:`str` subclass, so it behaves as the markup everywhere a string is expected, and additionally renders inline in notebooks and doc pages. Raises ------ TypeError If ``structure`` is not a ``GraphStructure``, such as a ``TemporalStructure``. ValueError If ``positions`` misses a node or places one outside the unit square, ``engine`` is unknown, or ``highlight`` names a pair ``(i, j)`` with ``i == j`` or no edge between ``i`` and ``j``. ImportError If a Graphviz engine is named and pygraphviz is not installed. Examples -------- >>> import numpy as np >>> import andrey >>> import andrey.viz >>> rng = np.random.default_rng(0) >>> x = rng.standard_normal(500) >>> y = x + 0.5 * rng.standard_normal(500) >>> out = andrey.pc(np.column_stack([x, y])) >>> svg = andrey.viz.draw(out.structure, dark=True) >>> svg.startswith("<svg") True """ _require_graph(structure) n = int(structure.n_nodes) labels = list(structure.labels) if structure.labels is not None else [str(i) for i in range(n)] node_types = structure.node_types latent = {k for k in range(n) if node_types is not None and int(node_types[k]) == LATENT} edges = _endpoints(structure) if positions is None: unit = layout(structure, engine, engine_options) else: unit = {} for k in range(n): point = positions.get(k) if ( point is None or len(point) != 2 or not all(math.isfinite(v) and 0 <= v <= 1 for v in point) ): raise ValueError("positions must give every node an (x, y) inside the unit square.") unit[k] = (float(point[0]), float(point[1])) ends = {(i, j) for i, j, *_ in edges if i != j} | {(j, i) for i, j, *_ in edges if i != j} rings = frozenset((int(i), int(j)) for i, j in highlight) if not rings <= ends: raise ValueError( f"highlight names no endpoint of an edge between two nodes: {sorted(rings - ends)}." ) inside = all(len(label) <= 2 for label in labels) widths = [_svg.footprint(label, inside)[0] for label in labels] pad = _svg.PAD # Room for the widest label at either side, so none is cut off at the canvas edge. pad_x = max(pad, max(widths, default=0.0) / 2 + 4) xs = [unit[k][0] for k in range(n)] ys = [unit[k][1] for k in range(n)] span_x = max(xs, default=0.0) - min(xs, default=0.0) span_y = max(ys, default=0.0) - min(ys, default=0.0) fits = [] # the scale each given side allows # A side too small for its padding and one padding of drawing is enlarged, so the scale stays # positive: a long label widens the side padding past the fixed floor. if width is not None: width = max(int(width), math.ceil(2 * pad_x + pad)) fits.append((width - 2 * pad_x) / span_x if span_x else math.inf) if height is not None: height = max(int(height), 3 * int(pad)) fits.append((height - 2 * pad) / span_y if span_y else math.inf) scale = min(fits, default=math.inf) if not math.isfinite(scale): scale = _scale(unit, widths) if width is None: width = max(round(span_x * scale + 2 * pad_x), _MIN_WIDTH) if height is None: height = max(round(span_y * scale + 2 * pad), _MIN_HEIGHT) left = (width - span_x * scale) / 2 - min(xs, default=0.0) * scale top = (height - span_y * scale) / 2 - min(ys, default=0.0) * scale pos = {k: (left + xs[k] * scale, top + ys[k] * scale) for k in range(n)} resolved = palette if palette is not None else (DARK if dark else LIGHT) return SVG( _svg.render( n=n, pos=pos, edges=edges, labels=labels, latent=latent, width=width, height=height, palette=resolved, background=background, highlight=rings, ) )