Color Graphviz and TikZ edges by emission

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
 ============
 

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,15 +272,15 @@ 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):
         initial_states = frozenset()
     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,

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 = {

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,16 +169,20 @@ 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``.
 
     Args:
         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)

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
 

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