Move canonize_colored_graph and save_colored_graph to lib/colored_graphs

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

lib/colored_graphs.py

@@ -0,0 +1,60 @@
+"""Utilities for canonizing and saving colored graphs."""
+import base64
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, cast, TypedDict
+
+from sage.all import Graph, save, load  # type: ignore[attr-defined]  # pylint: disable=no-name-in-module
+
+DIR = Path(__file__).parent.parent
+PALETTE = ['red', 'blue', 'green', 'yellow']
+
+VertexColoring = dict[Any, Any]
+
+
+class ColoredGraphId(TypedDict):
+    """Canonical id representing a colored graph"""
+    graph_id: str
+    coloring_id: str
+
+
+def canonize_colored_graph(g: Graph, coloring: VertexColoring) -> tuple[Graph, VertexColoring, ColoredGraphId]:
+    """Return a new canonical graph, new canonical coloring, and a ColoredGraphId without mutating inputs"""
+    canonical, cert = cast(
+        tuple[Graph, dict[Any, int]],
+        g.canonical_label(certificate=True),
+    )
+    graph_id = base64.urlsafe_b64encode(
+        canonical.graph6_string().encode()
+    ).decode()
+
+    color_seq = [0] * g.order()
+    for orig_v, canon_idx in cert.items():
+        color_seq[canon_idx] = coloring[orig_v]
+
+    canonical_coloring = {canon_idx: color for canon_idx, color in enumerate(color_seq)}
+    coloring_id = base64.urlsafe_b64encode(bytes(color_seq)).decode()
+    return canonical, canonical_coloring, ColoredGraphId(graph_id=graph_id, coloring_id=coloring_id)
+
+
+def save_colored_graph(g: Graph, coloring: VertexColoring) -> tuple[Graph, VertexColoring, ColoredGraphId]:
+    """
+    Canonize g and coloring, save to disk, and return the canonical forms with id.
+    If already saved, load and return the cached graph.
+    """
+    g_canon, canonical_coloring, cid = canonize_colored_graph(g, coloring)
+    out_dir = DIR / "data" / "graphs" / cid['graph_id'] / cid['coloring_id']
+    if (out_dir / "graph.sobj").exists():
+        g_canon = cast(Graph, load(str(out_dir / 'graph')))
+        return g_canon, canonical_coloring, cid
+    g_canon.is_planar(set_embedding=True, set_pos=True)
+    vertex_colors: defaultdict[str, list[Any]] = defaultdict(list)
+    for v, c in canonical_coloring.items():
+        vertex_colors[PALETTE[c]].append(v)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    g_canon.plot(
+        vertex_colors=dict(vertex_colors),
+        title=f"graph: {cid['graph_id']} coloring: {cid['coloring_id']}",
+    ).save(out_dir / 'graph.png')
+    save(g_canon, str(out_dir / 'graph'))
+    return g_canon, canonical_coloring, cid