| """ |
| Export a PyG dataset (Planetoid or Heterophilous) to a minimal Graphviz DOT. |
| - Nodes are colored by class. |
| - If --filter {train|val|test} is set, nodes in that split are colored red, |
| other nodes are left uncolored. |
| - Undirected edges are deduplicated; directed edges are written as‑is. |
| - For HeterophilousGraphDataset with multiple splits (e.g., Amazon‑ratings has 10), |
| use --split-index (default 0). |
| """ |
|
|
| import argparse |
| from torch_geometric.datasets import Planetoid, HeterophilousGraphDataset |
|
|
| PALETTE = [ |
| "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", |
| "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", |
| "#bcbd22", "#17becf" |
| ] |
| HIGHLIGHT_RED = "#ff0000" |
|
|
| def _extend_palette(base, need): |
| if need <= len(base): |
| return base[:need] |
| def brighten(hex_color, factor): |
| c = int(hex_color[1:], 16) |
| r = (c >> 16) & 255 |
| g = (c >> 8) & 255 |
| b = c & 255 |
| r = int(min(255, r + (255 - r) * factor)) |
| g = int(min(255, g + (255 - g) * factor)) |
| b = int(min(255, b + (255 - b) * factor)) |
| return f"#{r:02x}{g:02x}{b:02x}" |
| out = [] |
| for i in range(need): |
| base_hex = base[i % len(base)] |
| factor = 0.18 * (i // len(base)) |
| out.append(brighten(base_hex, factor)) |
| return out |
|
|
| def _infer_num_classes(dataset, data): |
| num_classes = getattr(dataset, "num_classes", None) |
| if not isinstance(num_classes, int) or num_classes <= 0: |
| num_classes = int(data.y.max().item()) + 1 |
| return num_classes |
|
|
| def _get_split_mask(data, which): |
| name = {"train": "train_mask", "val": "val_mask", "test": "test_mask"}[which] |
| m = getattr(data, name, None) |
| if m is None and which == "val": |
| m = getattr(data, "valid_mask", None) or getattr(data, "validation_mask", None) |
| return m |
|
|
| def load_graph(root: str, use_hetero: bool, name: str, split_index: int): |
| if use_hetero: |
| ds = HeterophilousGraphDataset(root=root, name=name) |
| idx = max(0, min(split_index, len(ds) - 1)) |
| data = ds[idx] |
| num_classes = _infer_num_classes(ds, data) |
| return data, num_classes |
| else: |
| ds = Planetoid(root=f"{root}/Planetoid", name=name) |
| data = ds[0] |
| num_classes = _infer_num_classes(ds, data) |
| return data, num_classes |
|
|
| def write_dot(path: str, data, num_classes: int, directed: bool, filter_split: str | None): |
| y = data.y |
| edge_index = data.edge_index |
| colors = _extend_palette(PALETTE, num_classes) |
| highlight_mask = None |
| if filter_split is not None: |
| m = _get_split_mask(data, filter_split) |
| if m is not None: |
| highlight_mask = m.bool() |
| gtype = "digraph" if directed else "graph" |
| eop = "->" if directed else "--" |
| with open(path, "w", encoding="utf-8") as f: |
| f.write(f"{gtype} {{\n") |
| for i in range(data.num_nodes): |
| cls = int(y[i]) |
| if cls < 0 or cls >= num_classes: |
| cls = cls % num_classes |
| base_col = colors[cls] |
| if filter_split is not None: |
| if highlight_mask is not None and bool(highlight_mask[i]): |
| col = HIGHLIGHT_RED |
| f.write(f' {i} [color="{col}", style="filled", fillcolor="{col}", fontcolor="white"];\n') |
| else: |
| f.write(f' {i} ;\n') |
| else: |
| col = base_col |
| f.write(f' {i} [color="{col}", style="filled", fillcolor="{col}", fontcolor="white"];\n') |
| if directed: |
| for s, t in edge_index.t().tolist(): |
| if s != t: |
| f.write(f" {s} {eop} {t};\n") |
| else: |
| seen = set() |
| for s, t in edge_index.t().tolist(): |
| if s == t: |
| continue |
| a, b = (s, t) if s <= t else (t, s) |
| if (a, b) in seen: |
| continue |
| seen.add((a, b)) |
| f.write(f" {a} {eop} {b};\n") |
| f.write("}\n") |
|
|
| def main(): |
| parser = argparse.ArgumentParser( |
| description="Export PyG datasets to minimal DOT with class colors and optional split highlighting." |
| ) |
| parser.add_argument("-o", "--output", default="graph.dot", help="Output .dot file (default: graph.dot)") |
| parser.add_argument("--directed", action="store_true", help="Write directed edges (default: undirected)") |
| parser.add_argument("--heterophilous", action="store_true", help="Use HeterophilousGraphDataset") |
| parser.add_argument("--name", default=None, help="Dataset name (Planetoid: Cora/CiteSeer/PubMed; Heterophilous: Amazon-ratings, Roman-empire, etc.)") |
| parser.add_argument("--root", default="data", help="Root folder for datasets (default: data)") |
| parser.add_argument("--split-index", type=int, default=0, help="Split index for heterophilous datasets (default: 0)") |
| parser.add_argument("--filter", choices=["train", "val", "test"], default=None, help="Highlight nodes in the selected split as red (others keep class colors)") |
| args = parser.parse_args() |
| dataset_name = args.name if args.name is not None else ("Amazon-ratings" if args.heterophilous else "Cora") |
| data, num_classes = load_graph(args.root, args.heterophilous, dataset_name, args.split_index) |
| write_dot(args.output, data, num_classes, args.directed, args.filter) |
| suffix = f", highlight={args.filter}" if args.filter else "" |
| print(f"Wrote {args.output} using {'HeterophilousGraphDataset' if args.heterophilous else 'Planetoid'}('{dataset_name}') | nodes={data.num_nodes}, edges={data.edge_index.size(1)}, classes={num_classes}{suffix}") |
|
|
| if __name__ == "__main__": |
| main() |
