amachine.am_visualization.am_draw

View Source

  1from pathlib import Path
  2from typing import Any
  3import tempfile, webbrowser, os
  4
  5import graphviz
  6from matplotlib import colormaps
  7from matplotlib.colors import to_hex
  8
  9def view_graph(GV):
 10    svg_bytes = GV.pipe(format="svg")
 11
 12    html = """<!DOCTYPE html>
 13<html>
 14<head>
 15<style>
 16  * { margin: 0; padding: 0; box-sizing: border-box; }
 17  body { background: #ffffff; overflow: hidden; width: 100vw; height: 100vh; }
 18  #viewport { width: 100%; height: 100%; overflow: hidden; cursor: grab; }
 19  #viewport.dragging { cursor: grabbing; }
 20  #container { transform-origin: 0 0; display: inline-block; }
 21  #hint {
 22    position: fixed; bottom: 12px; right: 12px;
 23    color: rgb(180,180,180); font: 12px monospace; pointer-events: none;
 24  }
 25</style>
 26</head>
 27<body>
 28<div id="viewport">
 29  <div id="container">
 30""" + svg_bytes.decode("utf-8") + """
 31  </div>
 32</div>
 33<div id="hint">scroll: zoom &nbsp;|&nbsp; drag: pan &nbsp;|&nbsp; dbl-click: reset</div>
 34<script>
 35  const viewport = document.getElementById("viewport");
 36  const container = document.getElementById("container");
 37
 38  let scale = 1, tx = 0, ty = 0;
 39  let dragging = false, ox = 0, oy = 0;
 40
 41  function apply() {
 42    container.style.transform = `translate(${tx}px,${ty}px) scale(${scale})`;
 43  }
 44
 45  function fitToWindow() {
 46    const svg = container.querySelector("svg");
 47    const r = svg.getBoundingClientRect();
 48    const vw = viewport.clientWidth, vh = viewport.clientHeight;
 49    scale = Math.min(vw / r.width, vh / r.height) * 0.95;
 50    tx = (vw - r.width * scale) / 2;
 51    ty = (vh - r.height * scale) / 2;
 52    apply();
 53  }
 54
 55  window.addEventListener("load", fitToWindow);
 56
 57  viewport.addEventListener("wheel", e => {
 58    e.preventDefault();
 59    const factor = e.deltaY < 0 ? 1.1 : 1 / 1.1;
 60    const rect = viewport.getBoundingClientRect();
 61    const mx = e.clientX - rect.left;
 62    const my = e.clientY - rect.top;
 63    tx = mx - (mx - tx) * factor;
 64    ty = my - (my - ty) * factor;
 65    scale *= factor;
 66    apply();
 67  }, { passive: false });
 68
 69  viewport.addEventListener("mousedown", e => {
 70    dragging = true;
 71    ox = e.clientX - tx;
 72    oy = e.clientY - ty;
 73    viewport.classList.add("dragging");
 74  });
 75  window.addEventListener("mousemove", e => {
 76    if (!dragging) return;
 77    tx = e.clientX - ox;
 78    ty = e.clientY - oy;
 79    apply();
 80  });
 81  window.addEventListener("mouseup", () => {
 82    dragging = false;
 83    viewport.classList.remove("dragging");
 84  });
 85
 86  viewport.addEventListener("dblclick", fitToWindow);
 87</script>
 88</body>
 89</html>"""
 90
 91    with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False, dir=".") as f:
 92        f.write(html)
 93        path = f.name
 94
 95    webbrowser.open(f"file://{path}")
 96    input("Press Enter to continue...")
 97    os.unlink(path)
 98
 99ClassParam = str | tuple[int, str]
100
101_ENGINE_CONFIGS: dict[str, dict[str, str]] = {
102    "dot": {
103        "rankdir": "LR",
104        "ranksep": "0.5",
105        "nodesep": "1.9",
106        "splines": "spline",
107        "constraint": "true",
108        "concentrate": "false",
109        "ratio": "auto",
110    },
111    "neato": {
112        "overlap": "scale",
113        "overlap_scaling": "-4",
114        "esep": "+2.5",
115        "sep": "+1.75",
116        "model": "shortpath",
117        "damping": "0.85",
118        "epsilon": "0.00001",
119        "maxiter": "10000",
120        "start": "5",
121    },
122    "fdp": {
123        "overlap": "prism",
124        "sep": "+1.5",
125        "K": "1.0",
126        "splines": "true",
127        "len": "3.0",
128        "maxiter": "5000",
129    },
130}
131
132_DEFAULT_FILL_COLOR   = "gray"
133_DEFAULT_BORDER_COLOR = "black"
134_PLAIN_BORDER_WIDTH   = 3.0
135_THICK_BORDER_WIDTH   = 24.0
136_DEFAULT_EDGE_WIDTH   = 6.0
137_COLORED_EDGE_WIDTH   = 12.0
138_DEFAULT_ARROW_SIZE   = 1.0
139_COLORED_ARROW_SIZE   = 1.75
140
141def _as_composition_depth_cls(param: ClassParam) -> tuple[int, str]:
142    return (0, param) if isinstance(param, str) else param
143
144
145def _collect_classes(states: list[Any], composition_depth: int, cls: str) -> list[str]:
146    prefix = cls.split("_")[0]
147    found: set[str] = set()
148    for state in states:
149        for label in state.classes[composition_depth]:
150            if label.split("_")[0] == prefix:
151                found.add(label)
152    return sorted(found)
153
154
155def _palette(cmap_name: str, n_cycle: int, count: int) -> list[str]:
156    cmap = colormaps[cmap_name]
157    return [to_hex(cmap(i % n_cycle)) for i in range(count)]
158
159
160def _resolve_color(
161    node_idx: int,
162    c_by: ClassParam | None,
163    class_list: list[str] | None,
164    colors: list[str] | None,
165    sccs: list[list[int]] | None,
166    state: Any,
167    default: str,
168) -> str:
169    if class_list is None:
170        return default
171
172    if sccs and c_by == "strongly_connected_components":
173        for i, component in enumerate(sccs):
174            if node_idx in component:
175                return colors[i % len(colors)]
176        return default
177
178    if c_by is not None:
179        composition_depth, _ = _as_composition_depth_cls(c_by)
180        for cl in state.classes[composition_depth]:
181            try:
182                idx = class_list.index(cl)
183                return colors[idx % len(colors)]
184            except ValueError:
185                continue
186
187    return default
188
189
190def _state_class_prefix(state: Any, composition_depth: int, cls_prefix: str) -> str | None:
191    for cl in state.classes[composition_depth]:
192        if cl.split("_")[0] == cls_prefix:
193            return cl
194    return None
195
196
197def create_digraph(engine: str = "dot") -> graphviz.Digraph:
198
199    graph_attr = {**_ENGINE_CONFIGS.get(engine, {}), "outputorder": "edgesfirst"}
200
201    node_attr = {
202        "shape": "box",
203        "style": "rounded, filled",
204        "fillcolor": "lightblue",
205        "fontname": "Helvetica",
206    }
207
208    edge_attr = {"color": "gray40"}
209
210    return graphviz.Digraph(
211        engine=engine,
212        graph_attr=graph_attr,
213        node_attr=node_attr,
214        edge_attr=edge_attr,
215    )
216
217
218def draw_graph(
219    aM: Any,
220    output_dir: Path | str | None = None,
221    title: str = "am_graph",
222    view: bool = True,
223    format: str = "png",
224    strongly_connected_components: list[list[int]] | None = None,
225    engine: str = "dot",
226    symbols_only: bool = False,
227    no_labels: bool = False,
228    color_nodes_by: ClassParam | None = None,
229    color_borders_by: ClassParam | None = None,
230    color_edges_by: str | None = None,
231    edge_width_by_cross_class: ClassParam | None = None,
232    cluster_by: ClassParam | None = None,
233    symbol_categories: dict[str, int] | None = None,
234) -> None:
235
236    if color_edges_by == "category" and symbol_categories is None:
237        raise ValueError("color_edges_by='category' requires symbol_categories")
238
239    # Normalise optional ClassParams to (composition_depth, cls) tuples up front.
240    if edge_width_by_cross_class is not None:
241        edge_width_by_cross_class = _as_composition_depth_cls(edge_width_by_cross_class)
242    if cluster_by is not None:
243        cluster_by = _as_composition_depth_cls(cluster_by)
244
245    GV = create_digraph(engine=engine)
246
247    # Node-colour configuration
248    
249    border_width   = _PLAIN_BORDER_WIDTH
250    arrowhead      = "normal"
251    node_colors: list[str] | None = None
252    node_classes: list[str] | None = None
253
254    if strongly_connected_components and color_nodes_by == "strongly_connected_components":
255        node_colors  = _palette("tab20_r", 8, len(strongly_connected_components))
256        border_width = _THICK_BORDER_WIDTH
257        arrowhead    = "nonenonenormal"
258    elif color_nodes_by is not None:
259        composition_depth, cls   = _as_composition_depth_cls(color_nodes_by)
260        node_classes = _collect_classes(aM.states, composition_depth, cls)
261        node_colors  = _palette("tab20_r", 8, len(node_classes))
262        border_width = _THICK_BORDER_WIDTH
263        arrowhead    = "nonenonenormal"
264
265    # Border-colour configuration
266    
267    border_colors: list[str] | None = None
268    border_classes: list[str] | None = None
269
270    if strongly_connected_components and color_borders_by == "strongly_connected_components":
271        border_colors = _palette("Set1", 8, len(strongly_connected_components))
272    elif color_borders_by is not None:
273        composition_depth, cls    = _as_composition_depth_cls(color_borders_by)
274        border_classes = _collect_classes(aM.states, composition_depth, cls)
275        border_colors  = _palette("Set1", 8, len(border_classes))
276
277    # Edge-colour configuration
278    
279    default_edge_width = _DEFAULT_EDGE_WIDTH
280    default_arrow_size = _DEFAULT_ARROW_SIZE
281    edge_colors: list[str] | None = None
282    categories: list[int] | None  = None
283
284    if color_edges_by == "category":
285        categories     = list(set(symbol_categories.values()))
286        edge_colors    = _palette("Dark2", 8, len(categories))
287        default_edge_width = _COLORED_EDGE_WIDTH
288        default_arrow_size = _COLORED_ARROW_SIZE
289    elif color_edges_by == "symbol":
290        edge_colors    = _palette("tab20", 20, len(aM.alphabet))
291        default_edge_width = _COLORED_EDGE_WIDTH
292        default_arrow_size = _COLORED_ARROW_SIZE
293
294    # Add nodes
295    
296    for node_idx, state in enumerate(aM.states):
297        fill_color = _resolve_color(
298            node_idx, color_nodes_by,
299            node_classes, node_colors,
300            strongly_connected_components, state,
301            _DEFAULT_FILL_COLOR,
302        )
303        border_color = _resolve_color(
304            node_idx, color_borders_by,
305            border_classes, border_colors,
306            strongly_connected_components, state,
307            _DEFAULT_BORDER_COLOR,
308        )
309
310        node_kw: dict[str, str] = dict(
311            label=state.name or str(node_idx),
312            shape="circle",
313            style="bold,filled",
314            fillcolor=fill_color,
315            color=border_color,
316            width="3.0",
317            fontsize="36",
318            penwidth=str(border_width),
319        )
320
321        my_class: str | None = None
322        if cluster_by is not None:
323            composition_depth, cls = cluster_by
324            my_class = _state_class_prefix(state, composition_depth, cls)
325
326        if my_class is not None:
327            with GV.subgraph(name=f"cluster_{my_class}") as sub:
328                sub.attr(
329                    style="solid, filled",
330                    color="blue",
331                    fillcolor="#f0f8ff",
332                    label=my_class,
333                    margin="90",
334                )
335                sub.node(str(node_idx), **node_kw)
336        else:
337            GV.node(str(node_idx), **node_kw)
338
339    # Add edges
340
341    for tr in aM.transitions:
342        u, v   = tr.origin_state_idx, tr.target_state_idx
343        symbol = aM.alphabet[tr.symbol_idx]
344
345        pr_str     = str(tr.pq) if aM.is_q_weighted else str(round(tr.prob, 4))
346        label_text = symbol if symbols_only else f"{symbol}({pr_str})"
347        fontsize   = "70" if symbols_only else "12"
348
349        if no_labels:
350            html_label = None
351        else:
352            html_label = (
353                f'<<TABLE BORDER="1" COLOR="black" CELLBORDER="0" CELLSPACING="0"'
354                f' CELLPADDING="4" STYLE="ROUNDED" BGCOLOR="white">'
355                f"<TR><TD>{label_text}</TD></TR>"
356                f"</TABLE>>"
357            )
358
359        # Edge colour
360        edge_color = "black"
361        if color_edges_by == "category" and symbol in symbol_categories:
362            cat = symbol_categories[symbol]
363            try:
364                edge_color = edge_colors[categories.index(cat) % 8]
365            except ValueError:
366                pass
367        elif color_edges_by == "symbol":
368            edge_color = edge_colors[tr.symbol_idx % 20]
369
370        # Edge width
371        edge_width = default_edge_width
372        if edge_width_by_cross_class is not None:
373            composition_depth, cls = edge_width_by_cross_class
374            origin_class = _state_class_prefix(aM.states[u], composition_depth, cls)
375            target_class = _state_class_prefix(aM.states[v], composition_depth, cls)
376            if origin_class != target_class:
377                edge_width = default_edge_width * 1.5
378
379        GV.edge(
380            str(u), str(v),
381            label=html_label,
382            headlabel=" ",
383            fontsize=fontsize,
384            fontname="Times-Italic",
385            fontcolor="#1f1f1f",
386            color=edge_color,
387            arrowhead=arrowhead,
388            penwidth=str(edge_width),
389            arrowsize=str(default_arrow_size),
390        )
391
392    # Render
393    if view:
394        view_graph( GV )
395
396    if output_dir is not None:
397        GV.attr(bgcolor="transparent")
398        GV.render(title, directory=output_dir, view=False, format=format, cleanup=True)

def view_graph(GV): View Source

10def view_graph(GV):
11    svg_bytes = GV.pipe(format="svg")
12
13    html = """<!DOCTYPE html>
14<html>
15<head>
16<style>
17  * { margin: 0; padding: 0; box-sizing: border-box; }
18  body { background: #ffffff; overflow: hidden; width: 100vw; height: 100vh; }
19  #viewport { width: 100%; height: 100%; overflow: hidden; cursor: grab; }
20  #viewport.dragging { cursor: grabbing; }
21  #container { transform-origin: 0 0; display: inline-block; }
22  #hint {
23    position: fixed; bottom: 12px; right: 12px;
24    color: rgb(180,180,180); font: 12px monospace; pointer-events: none;
25  }
26</style>
27</head>
28<body>
29<div id="viewport">
30  <div id="container">
31""" + svg_bytes.decode("utf-8") + """
32  </div>
33</div>
34<div id="hint">scroll: zoom &nbsp;|&nbsp; drag: pan &nbsp;|&nbsp; dbl-click: reset</div>
35<script>
36  const viewport = document.getElementById("viewport");
37  const container = document.getElementById("container");
38
39  let scale = 1, tx = 0, ty = 0;
40  let dragging = false, ox = 0, oy = 0;
41
42  function apply() {
43    container.style.transform = `translate(${tx}px,${ty}px) scale(${scale})`;
44  }
45
46  function fitToWindow() {
47    const svg = container.querySelector("svg");
48    const r = svg.getBoundingClientRect();
49    const vw = viewport.clientWidth, vh = viewport.clientHeight;
50    scale = Math.min(vw / r.width, vh / r.height) * 0.95;
51    tx = (vw - r.width * scale) / 2;
52    ty = (vh - r.height * scale) / 2;
53    apply();
54  }
55
56  window.addEventListener("load", fitToWindow);
57
58  viewport.addEventListener("wheel", e => {
59    e.preventDefault();
60    const factor = e.deltaY < 0 ? 1.1 : 1 / 1.1;
61    const rect = viewport.getBoundingClientRect();
62    const mx = e.clientX - rect.left;
63    const my = e.clientY - rect.top;
64    tx = mx - (mx - tx) * factor;
65    ty = my - (my - ty) * factor;
66    scale *= factor;
67    apply();
68  }, { passive: false });
69
70  viewport.addEventListener("mousedown", e => {
71    dragging = true;
72    ox = e.clientX - tx;
73    oy = e.clientY - ty;
74    viewport.classList.add("dragging");
75  });
76  window.addEventListener("mousemove", e => {
77    if (!dragging) return;
78    tx = e.clientX - ox;
79    ty = e.clientY - oy;
80    apply();
81  });
82  window.addEventListener("mouseup", () => {
83    dragging = false;
84    viewport.classList.remove("dragging");
85  });
86
87  viewport.addEventListener("dblclick", fitToWindow);
88</script>
89</body>
90</html>"""
91
92    with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False, dir=".") as f:
93        f.write(html)
94        path = f.name
95
96    webbrowser.open(f"file://{path}")
97    input("Press Enter to continue...")
98    os.unlink(path)

ClassParam = str | tuple[int, str]

def create_digraph(engine: str = 'dot') -> graphviz.graphs.Digraph: View Source

198def create_digraph(engine: str = "dot") -> graphviz.Digraph:
199
200    graph_attr = {**_ENGINE_CONFIGS.get(engine, {}), "outputorder": "edgesfirst"}
201
202    node_attr = {
203        "shape": "box",
204        "style": "rounded, filled",
205        "fillcolor": "lightblue",
206        "fontname": "Helvetica",
207    }
208
209    edge_attr = {"color": "gray40"}
210
211    return graphviz.Digraph(
212        engine=engine,
213        graph_attr=graph_attr,
214        node_attr=node_attr,
215        edge_attr=edge_attr,
216    )

219def draw_graph(
220    aM: Any,
221    output_dir: Path | str | None = None,
222    title: str = "am_graph",
223    view: bool = True,
224    format: str = "png",
225    strongly_connected_components: list[list[int]] | None = None,
226    engine: str = "dot",
227    symbols_only: bool = False,
228    no_labels: bool = False,
229    color_nodes_by: ClassParam | None = None,
230    color_borders_by: ClassParam | None = None,
231    color_edges_by: str | None = None,
232    edge_width_by_cross_class: ClassParam | None = None,
233    cluster_by: ClassParam | None = None,
234    symbol_categories: dict[str, int] | None = None,
235) -> None:
236
237    if color_edges_by == "category" and symbol_categories is None:
238        raise ValueError("color_edges_by='category' requires symbol_categories")
239
240    # Normalise optional ClassParams to (composition_depth, cls) tuples up front.
241    if edge_width_by_cross_class is not None:
242        edge_width_by_cross_class = _as_composition_depth_cls(edge_width_by_cross_class)
243    if cluster_by is not None:
244        cluster_by = _as_composition_depth_cls(cluster_by)
245
246    GV = create_digraph(engine=engine)
247
248    # Node-colour configuration
249    
250    border_width   = _PLAIN_BORDER_WIDTH
251    arrowhead      = "normal"
252    node_colors: list[str] | None = None
253    node_classes: list[str] | None = None
254
255    if strongly_connected_components and color_nodes_by == "strongly_connected_components":
256        node_colors  = _palette("tab20_r", 8, len(strongly_connected_components))
257        border_width = _THICK_BORDER_WIDTH
258        arrowhead    = "nonenonenormal"
259    elif color_nodes_by is not None:
260        composition_depth, cls   = _as_composition_depth_cls(color_nodes_by)
261        node_classes = _collect_classes(aM.states, composition_depth, cls)
262        node_colors  = _palette("tab20_r", 8, len(node_classes))
263        border_width = _THICK_BORDER_WIDTH
264        arrowhead    = "nonenonenormal"
265
266    # Border-colour configuration
267    
268    border_colors: list[str] | None = None
269    border_classes: list[str] | None = None
270
271    if strongly_connected_components and color_borders_by == "strongly_connected_components":
272        border_colors = _palette("Set1", 8, len(strongly_connected_components))
273    elif color_borders_by is not None:
274        composition_depth, cls    = _as_composition_depth_cls(color_borders_by)
275        border_classes = _collect_classes(aM.states, composition_depth, cls)
276        border_colors  = _palette("Set1", 8, len(border_classes))
277
278    # Edge-colour configuration
279    
280    default_edge_width = _DEFAULT_EDGE_WIDTH
281    default_arrow_size = _DEFAULT_ARROW_SIZE
282    edge_colors: list[str] | None = None
283    categories: list[int] | None  = None
284
285    if color_edges_by == "category":
286        categories     = list(set(symbol_categories.values()))
287        edge_colors    = _palette("Dark2", 8, len(categories))
288        default_edge_width = _COLORED_EDGE_WIDTH
289        default_arrow_size = _COLORED_ARROW_SIZE
290    elif color_edges_by == "symbol":
291        edge_colors    = _palette("tab20", 20, len(aM.alphabet))
292        default_edge_width = _COLORED_EDGE_WIDTH
293        default_arrow_size = _COLORED_ARROW_SIZE
294
295    # Add nodes
296    
297    for node_idx, state in enumerate(aM.states):
298        fill_color = _resolve_color(
299            node_idx, color_nodes_by,
300            node_classes, node_colors,
301            strongly_connected_components, state,
302            _DEFAULT_FILL_COLOR,
303        )
304        border_color = _resolve_color(
305            node_idx, color_borders_by,
306            border_classes, border_colors,
307            strongly_connected_components, state,
308            _DEFAULT_BORDER_COLOR,
309        )
310
311        node_kw: dict[str, str] = dict(
312            label=state.name or str(node_idx),
313            shape="circle",
314            style="bold,filled",
315            fillcolor=fill_color,
316            color=border_color,
317            width="3.0",
318            fontsize="36",
319            penwidth=str(border_width),
320        )
321
322        my_class: str | None = None
323        if cluster_by is not None:
324            composition_depth, cls = cluster_by
325            my_class = _state_class_prefix(state, composition_depth, cls)
326
327        if my_class is not None:
328            with GV.subgraph(name=f"cluster_{my_class}") as sub:
329                sub.attr(
330                    style="solid, filled",
331                    color="blue",
332                    fillcolor="#f0f8ff",
333                    label=my_class,
334                    margin="90",
335                )
336                sub.node(str(node_idx), **node_kw)
337        else:
338            GV.node(str(node_idx), **node_kw)
339
340    # Add edges
341
342    for tr in aM.transitions:
343        u, v   = tr.origin_state_idx, tr.target_state_idx
344        symbol = aM.alphabet[tr.symbol_idx]
345
346        pr_str     = str(tr.pq) if aM.is_q_weighted else str(round(tr.prob, 4))
347        label_text = symbol if symbols_only else f"{symbol}({pr_str})"
348        fontsize   = "70" if symbols_only else "12"
349
350        if no_labels:
351            html_label = None
352        else:
353            html_label = (
354                f'<<TABLE BORDER="1" COLOR="black" CELLBORDER="0" CELLSPACING="0"'
355                f' CELLPADDING="4" STYLE="ROUNDED" BGCOLOR="white">'
356                f"<TR><TD>{label_text}</TD></TR>"
357                f"</TABLE>>"
358            )
359
360        # Edge colour
361        edge_color = "black"
362        if color_edges_by == "category" and symbol in symbol_categories:
363            cat = symbol_categories[symbol]
364            try:
365                edge_color = edge_colors[categories.index(cat) % 8]
366            except ValueError:
367                pass
368        elif color_edges_by == "symbol":
369            edge_color = edge_colors[tr.symbol_idx % 20]
370
371        # Edge width
372        edge_width = default_edge_width
373        if edge_width_by_cross_class is not None:
374            composition_depth, cls = edge_width_by_cross_class
375            origin_class = _state_class_prefix(aM.states[u], composition_depth, cls)
376            target_class = _state_class_prefix(aM.states[v], composition_depth, cls)
377            if origin_class != target_class:
378                edge_width = default_edge_width * 1.5
379
380        GV.edge(
381            str(u), str(v),
382            label=html_label,
383            headlabel=" ",
384            fontsize=fontsize,
385            fontname="Times-Italic",
386            fontcolor="#1f1f1f",
387            color=edge_color,
388            arrowhead=arrowhead,
389            penwidth=str(edge_width),
390            arrowsize=str(default_arrow_size),
391        )
392
393    # Render
394    if view:
395        view_graph( GV )
396
397    if output_dir is not None:
398        GV.attr(bgcolor="transparent")
399        GV.render(title, directory=output_dir, view=False, format=format, cleanup=True)