From a7bca8555fd17cae67206ffe37766512ba849857 Mon Sep 17 00:00:00 2001 From: Ryan James Date: Thu, 13 Aug 2026 09:49:22 -0600 Subject: [PATCH] Color Graphviz and TikZ edges by emission by default. Shared symbols get a stable categorical palette so machine diagrams are easier to read; VPA/Dyck kind colours still take precedence. Co-authored-by: Cursor --- docs/viz.rst | 6 ++++ sofic/viz/_context.py | 79 ++++++++++++++++++++++++++++++++++++++++--- sofic/viz/graphviz.py | 11 ++++-- sofic/viz/tikz.py | 12 +++++-- tests/test_tikz.py | 11 ++++++ tests/test_viz.py | 30 +++++++++++++++- 6 files changed, 140 insertions(+), 9 deletions(-) diff --git a/docs/viz.rst b/docs/viz.rst index cade977..5b24979 100644 --- a/docs/viz.rst +++ b/docs/viz.rst @@ -28,6 +28,12 @@ installed (``pip install sofic[viz]`` plus the system ``graphviz`` binary): Models implement Jupyter display via ``_repr_mimebundle_``, so simply evaluating a model in a notebook shows its diagram when Graphviz is available. +By default, edges that share an emission (or an input/label symbol, on +automata and shifts with no emissions) are drawn in the same colour. Pass +``color_by_emission=False`` to ``draw`` / ``to_graphviz`` / ``to_tikz`` for +uncoloured edges. Visibly pushdown automata and Dyck shifts keep their +call / return / internal colours. + TikZ / LaTeX ============ diff --git a/sofic/viz/_context.py b/sofic/viz/_context.py index cd203fe..b2daf54 100644 --- a/sofic/viz/_context.py +++ b/sofic/viz/_context.py @@ -13,21 +13,25 @@ ATTR_FUTURE_SYMBOL, ATTR_KIND, ATTR_OUTPUT, + EPSILON, KIND_CALL, KIND_INTERNAL, KIND_RETURN, Transition, ) from sofic.viz._edge import ( + PART_EMISSION, PART_KIND, PART_MATCH_TAG, PART_MULTIPLICITY, PART_PROB, PART_QUASIPROB, PART_STACK, + PART_SYMBOL, EdgePart, EdgeSpec, edge_spec, + part_value, ) from sofic.viz._format import ( format_belief, @@ -123,6 +127,58 @@ def _dyck_edge_color(kind: Any) -> str | None: return None +# Tableau 10 — Graphviz hex, stable assignment by sorted ``repr`` of the symbol. +EMISSION_PALETTE: tuple[str, ...] = ( + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + "#bcbd22", + "#17becf", +) + +_NAMED_RGB: dict[str, tuple[int, int, int]] = { + "seagreen": (46, 139, 87), + "firebrick": (178, 34, 34), + "steelblue": (70, 130, 180), +} + + +def _rgb_from_graphviz_color(color: str) -> tuple[int, int, int] | None: + if color.startswith("#") and len(color) == 7: + return int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16) + return _NAMED_RGB.get(color) + + +def tikz_draw_color(color: str) -> str: + """TikZ ``draw=`` value for a Graphviz color string (hex or named).""" + rgb = _rgb_from_graphviz_color(color) + if rgb is None: + return color + red, green, blue = rgb + return f"{{rgb,255:red,{red};green,{green};blue,{blue}}}" + + +def _emission_color_key(model: StateMachine, transition: Transition) -> Any: + """Emission (else input/label symbol) used to color ``transition``, or None.""" + key = part_value(edge_spec(model, transition), PART_EMISSION, PART_SYMBOL) + if key is None or key is EPSILON: + return None + return key + + +def _emission_color_map(model: StateMachine) -> dict[Any, str]: + keys = sorted( + {key for transition in model.transitions() if (key := _emission_color_key(model, transition)) is not None}, + key=repr, + ) + return {key: EMISSION_PALETTE[index % len(EMISSION_PALETTE)] for index, key in enumerate(keys)} + + _TRANSIENT_FILL = "mistyrose" _RECURRENT_FILL = "honeydew" _RECURRENCE_ATOL = 1e-12 @@ -164,7 +220,12 @@ def _recurrence_fill_sets( return frozenset(), frozenset() -def viz_context(model: StateMachine, *, style: str = "auto") -> VizContext: +def viz_context( + model: StateMachine, + *, + style: str = "auto", + color_by_emission: bool = True, +) -> VizContext: from sofic.automata.base import LabeledAutomaton from sofic.automata.transducers import Transducer from sofic.automata.vpa import VisiblyPushdownAutomaton @@ -198,8 +259,9 @@ def edge_label(transition: Transition) -> str: def _dyck_color(transition: Transition) -> str | None: return _dyck_edge_color(transition.data.get(ATTR_KIND)) - edge_color: Callable[[Transition], str | None] = lambda _t: None + specialized_color: Callable[[Transition], str | None] = lambda _t: None edge_style: Callable[[Transition], str | None] = lambda _t: None + emission_colors = _emission_color_map(model) if color_by_emission else {} if isinstance(model, LabeledAutomaton): initial_states = model.initial_states @@ -210,7 +272,7 @@ def _dyck_color(transition: Transition) -> str | None: if model.initial_state is not None: initial_states = frozenset({model.initial_state}) accepting_states = model.accepting_states - edge_color = _dyck_color + specialized_color = _dyck_color elif isinstance(model, MixedStatePresentation): initial_states = frozenset({model.initial_mixed_state}) elif isinstance(model, BidirectionalEpsilonMachine): @@ -218,7 +280,7 @@ def _dyck_color(transition: Transition) -> str | None: elif isinstance(model, (NMachine, MooreHMM, MealyHMM, StochasticModel, QuasiStochasticModel)): initial_states = _stochastic_initials(model) elif isinstance(model, SoficDyckShift): - edge_color = _dyck_color + specialized_color = _dyck_color for state in model.states(): attrs = model.graph.state_attrs(state) @@ -257,6 +319,15 @@ def node_fillcolor(state: Hashable) -> str | None: return _RECURRENT_FILL return None + def edge_color(transition: Transition) -> str | None: + color = specialized_color(transition) + if color: + return color + key = _emission_color_key(model, transition) + if key is None: + return None + return emission_colors.get(key) + return VizContext( title=title, initial_states=initial_states, diff --git a/sofic/viz/graphviz.py b/sofic/viz/graphviz.py index 425b0b4..0fa5115 100644 --- a/sofic/viz/graphviz.py +++ b/sofic/viz/graphviz.py @@ -37,14 +37,21 @@ def model_to_graphviz( *, rankdir: str | None = None, style: str = "auto", + color_by_emission: bool = True, graph_attr: dict[str, str] | None = None, node_attr: dict[str, str] | None = None, edge_attr: dict[str, str] | None = None, ) -> graphviz.Digraph: - """Return a :class:`graphviz.Digraph` for ``model``.""" + """Return a :class:`graphviz.Digraph` for ``model``. + + Edges sharing an emission (or input/label symbol when the model has no + emissions) get a common colour from a categorical palette. Pass + ``color_by_emission=False`` for uncoloured edges. Visibly pushdown / Dyck + kind colours (call / return / internal) take precedence. + """ graphviz = _require_graphviz() model = _model_for_viz(model) - context = viz_context(model, style=style) + context = viz_context(model, style=style, color_by_emission=color_by_emission) resolved_rankdir = rankdir if rankdir is not None else (context.rankdir or "LR") attrs = { diff --git a/sofic/viz/tikz.py b/sofic/viz/tikz.py index 7894556..cc41b84 100644 --- a/sofic/viz/tikz.py +++ b/sofic/viz/tikz.py @@ -12,7 +12,7 @@ from sofic.base import StateMachine from sofic.graph import Transition -from sofic.viz._context import VizContext, viz_context +from sofic.viz._context import VizContext, tikz_draw_color, viz_context from sofic.viz._edge import ( PART_EMISSION, PART_KIND, @@ -169,6 +169,7 @@ def model_to_tikz( rankdir: str | None = None, label: str | None = None, edge_label_pos: float = 0.5, + color_by_emission: bool = True, ) -> str: """Return a Vaucanson-style TikZ picture for ``model``. @@ -176,9 +177,12 @@ def model_to_tikz( edge_label_pos: Fraction along each edge (0 = source, 1 = target) at which to place the edge label. Use ``1/3`` to keep labels clear of mid-edge crossings. + color_by_emission: Colour edges by emission (or input/label symbol). + Defaults to ``True``. Visibly pushdown / Dyck kind colours take + precedence. """ model = _model_for_viz(model) - context = viz_context(model, style=style) + context = viz_context(model, style=style, color_by_emission=color_by_emission) if layout == "circle": coords = layout_circle(model, radius=radius, positions=positions, angles=angles) @@ -249,6 +253,10 @@ def model_to_tikz( has_reverse=has_reverse, loop_style=loop_styles.get((source, target, index)), ) + color = context.edge_color(transition) + if color: + draw = f"draw={tikz_draw_color(color)}" + style_opts = f"{style_opts}, {draw}" if style_opts else draw edge_label = _tikz_edge_label(model, transition) source_name = node_name(source) target_name = node_name(target) diff --git a/tests/test_tikz.py b/tests/test_tikz.py index ec4420e..422a990 100644 --- a/tests/test_tikz.py +++ b/tests/test_tikz.py @@ -39,6 +39,17 @@ def test_dfa_tikz_symbol_only(): assert r"\Edge{" not in tikz +def test_tikz_colors_edges_by_emission_by_default(): + from sofic.viz._context import EMISSION_PALETTE, tikz_draw_color + + tikz = model_to_tikz(golden_mean_forward(0.5)) + assert f"draw={tikz_draw_color(EMISSION_PALETTE[0])}" in tikz + assert f"draw={tikz_draw_color(EMISSION_PALETTE[1])}" in tikz + + plain = model_to_tikz(golden_mean_forward(0.5), color_by_emission=False) + assert "draw=" not in plain + + def test_sofic_dyck_tikz_marks_matched_edges(): from sofic.examples import sofic_dyck_nondeterminizable_shift diff --git a/tests/test_viz.py b/tests/test_viz.py index 6c39937..747a593 100644 --- a/tests/test_viz.py +++ b/tests/test_viz.py @@ -8,7 +8,7 @@ from sofic.examples.epsilon_machines import golden_mean, golden_mean_bidirectional from sofic.generators.markov import MarkovChain from sofic.graph import ATTR_PROB -from sofic.viz._context import viz_context +from sofic.viz._context import EMISSION_PALETTE, viz_context from sofic.viz._format import ( format_belief, format_distribution, @@ -277,3 +277,31 @@ def test_format_prob_label_symbolic(): assert format_prob_label(sp.Rational(1, 2)) == "1/2" assert "a" in format_prob_label(a / (a + 1)) assert format_prob_label(0.5) == "1/2" + + +def test_edges_colored_by_emission_by_default(): + eps = golden_mean() + source = model_to_graphviz(eps).source + assert f'color="{EMISSION_PALETTE[0]}"' in source + assert f'color="{EMISSION_PALETTE[1]}"' in source + + context = viz_context(eps) + by_symbol: dict[object, set[str | None]] = {} + for transition in eps.transitions(): + by_symbol.setdefault(transition.data["emission"], set()).add(context.edge_color(transition)) + assert by_symbol[0] == {EMISSION_PALETTE[0]} + assert by_symbol[1] == {EMISSION_PALETTE[1]} + + +def test_color_by_emission_can_be_disabled(): + source = model_to_graphviz(golden_mean(), color_by_emission=False).source + assert 'color="#' not in source + + +def test_dfa_edges_colored_by_input_symbol(): + source = model_to_graphviz(_dfa()).source + assert f'color="{EMISSION_PALETTE[0]}"' in source + assert f'color="{EMISSION_PALETTE[1]}"' in source + + plain = model_to_graphviz(_dfa(), color_by_emission=False).source + assert 'color="#' not in plain