GNN-Slides

import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

import copy
import hashlib
import json
import time

import altair as alt
import marimo as mo
import networkx as nx
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.manifold import TSNE
from sklearn.metrics import roc_auc_score, roc_curve
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GATConv, GCNConv, MessagePassing, SAGEConv
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.transforms import NormalizeFeatures, RandomLinkSplit
from torch_geometric.utils import (
    add_self_loops,
    degree,
    homophily,
    negative_sampling,
    scatter,
    to_dense_adj,
)

# several notebooks share this CPU; more threads only add contention
torch.set_num_threads(4)
DATA_DIR = mo.notebook_dir().parents[1] / "data"  # repo-level data/, shared with cs224w
# every expensive result (training runs, t-SNE, depth sweep) is cached here
CACHE_DIR = DATA_DIR / "gnn_slides_cache"
CACHE_DIR.mkdir(parents=True, exist_ok=True)


def stable_hash(obj) -> str:
    """Short, order-independent key for a config dict."""
    return hashlib.sha1(json.dumps(obj, sort_keys=True).encode()).hexdigest()[:10]

_FONTS = (
    "https://fonts.googleapis.com/css2?"
    "family=Space+Grotesk:wght@500;600;700"
    "&family=Inter:wght@400;500;600;700"
    "&family=JetBrains+Mono:wght@400;600&display=swap"
)
CSS = (
    f"@import url('{_FONTS}');"
    """
.gd { --bg:#f3f8f5; --ink:#07160f; --muted:#4d6a5c; --line:#d2e3d9;
  --g9:#052e1f; --g8:#065f46; --g7:#047857; --g6:#059669; --g5:#10b981; --g3:#6ee7b7; --g1:#d1fae5;
  --lime:#a3e635; --amber:#f59e0b; --rose:#f43f5e; --violet:#8b5cf6; --sky:#0ea5e9;
  font-family:'Inter',system-ui,sans-serif; color:var(--ink); background:var(--bg);
  border:1px solid var(--line); border-radius:6px; padding:30px 44px 46px;
      min-height:max(600px, calc(56.25cqw - 6px)); }
    @supports (zoom: calc(100cqw / 1320px)) {
      /* fixed 1320×750 design canvas, zoomed to whatever stage marimo gives the slide */
      .gd { width:1320px; min-height:750px; zoom:calc(100cqw / 1320px); }
  /* Vega draws on a canvas and ignores CSS zoom when hit-testing clicks: undo the zoom for clickable charts */
  .gd .unzoom { zoom:calc(1320px / 100cqw); }
    }
    .gd {
  position:relative; overflow:hidden; box-sizing:border-box; text-align:left; }
.gd::before { content:''; position:absolute; inset:0 0 auto 0; height:4px;
  background:linear-gradient(90deg,var(--g7),var(--g5) 45%,var(--lime)); }
.gd.dark { color:#e3f5ea; border-color:#0d3b28;
  background:
    radial-gradient(900px 420px at 88% -8%, rgba(16,185,129,.28), transparent 62%),
    radial-gradient(700px 380px at -6% 112%, rgba(163,230,53,.14), transparent 60%),
    linear-gradient(rgba(16,185,129,.06) 1px, transparent 1px) 0 0/34px 34px,
    linear-gradient(90deg, rgba(16,185,129,.06) 1px, transparent 1px) 0 0/34px 34px,
    #03130c; }
.gd .kicker { font:600 12px/1 'JetBrains Mono',monospace; letter-spacing:.22em; text-transform:uppercase;
  color:var(--g6); margin:6px 0 14px; display:flex; align-items:center; gap:10px; }
.gd .kicker::before { content:''; width:10px; height:10px; background:var(--g5); display:inline-block; }
.gd.dark .kicker { color:var(--g3); }
.gd h1, .gd h2, .gd h3 { font-family:'Space Grotesk','Inter',sans-serif; letter-spacing:-.025em; margin:0; }
.gd h1 { font-size:64px; line-height:1.0; font-weight:700; margin-bottom:18px; }
.gd h2 { font-size:34px; line-height:1.1; font-weight:700; margin-bottom:12px; }
.gd h3 { font-size:18px; font-weight:600; margin-bottom:6px; }
.gd .hl { background:linear-gradient(90deg,var(--g6),var(--g5) 40%,var(--lime)); -webkit-background-clip:text;
  background-clip:text; color:transparent; }
.gd.dark .hl { background:linear-gradient(90deg,var(--g3),var(--lime)); -webkit-background-clip:text; background-clip:text; }
.gd .lead { font-size:17px; line-height:1.55; color:var(--muted); max-width:70ch; margin:0 0 14px; }
.gd.dark .lead { color:#a7cdb9; }
.gd p, .gd li { font-size:16.5px; line-height:1.55; }
    .gd .codehilite pre, .gd pre code { font-size:14.5px !important; line-height:1.6 !important; }
.gd ul { margin:6px 0; padding-left:20px; }
.gd b, .gd strong { color:inherit; }
.gd code { font-family:'JetBrains Mono',monospace; font-size:.86em; background:var(--g1); color:var(--g9);
  padding:1px 6px; border-radius:3px; }
.gd.dark code { background:rgba(110,231,183,.12); color:#bbf7d0; }
.gd .cols { display:grid; gap:26px; align-items:start; }
.gd .card { background:#fff; border:1px solid var(--line); border-left:3px solid var(--c,var(--g5));
  border-radius:4px; padding:14px 16px; }
.gd.dark .card { background:rgba(255,255,255,.035); border-color:rgba(110,231,183,.18); border-left-color:var(--c,var(--g5)); }
.gd .card h3 { color:var(--c,var(--g7)); }
.gd.dark .card h3 { color:var(--c,var(--g3)); }
.gd .card p { margin:0; font-size:14.5px; color:var(--muted); }
.gd.dark .card p { color:#a7cdb9; }
.gd .eq { background:#fff; border:1px solid var(--line); border-left:3px solid var(--g5); border-radius:4px;
  padding:2px 18px; margin:10px 0; overflow-x:auto; }
.gd.dark .eq { background:rgba(255,255,255,.04); border-color:rgba(110,231,183,.2); }
.gd .stat { font:700 40px/1 'Space Grotesk',sans-serif; color:var(--g7); letter-spacing:-.03em; }
.gd.dark .stat { color:var(--lime); }
.gd .stat small { font:500 12px 'JetBrains Mono',monospace; color:var(--muted); letter-spacing:.06em;
  display:block; margin-top:6px; text-transform:uppercase; }
.gd .tag { display:inline-block; font:600 11.5px 'JetBrains Mono',monospace; padding:4px 10px; margin:0 6px 6px 0;
  border:1px solid var(--line); border-radius:3px; background:#fff; color:var(--g8); }
.gd.dark .tag { background:rgba(16,185,129,.08); border-color:rgba(110,231,183,.28); color:#bbf7d0; }
.gd .num { font:700 92px/0.9 'Space Grotesk',sans-serif; color:transparent; -webkit-text-stroke:1.5px var(--g5); }
.gd .foot { position:absolute; left:44px; right:44px; bottom:16px; display:flex; justify-content:space-between;
  font:500 11px 'JetBrains Mono',monospace; letter-spacing:.14em; text-transform:uppercase; color:var(--muted); opacity:.75; }
.gd.dark .foot { color:#7fb79a; }
.gd .muted { color:var(--muted); }
.gd .badge { display:inline-flex; gap:6px; align-items:center; font:600 12px 'JetBrains Mono',monospace;
  padding:4px 10px; border-radius:3px; background:var(--g1); color:var(--g8); }
.gd .badge.warn { background:#fef3c7; color:#92400e; }
.gd table.gt { border-collapse:collapse; width:100%; font-size:14px; }
.gd table.gt th { text-align:left; font:600 11.5px 'JetBrains Mono',monospace; text-transform:uppercase;
  letter-spacing:.08em; color:var(--muted); border-bottom:2px solid var(--g5); padding:8px 10px; }
.gd table.gt td { border-bottom:1px solid var(--line); padding:8px 10px; vertical-align:top; }
.gd.dark table.gt td { border-color:rgba(110,231,183,.15); }
.gd .step { display:flex; gap:14px; align-items:flex-start; margin:10px 0; }
.gd .step i { font:700 13px 'JetBrains Mono',monospace; font-style:normal; min-width:30px; height:30px;
  display:grid; place-items:center; background:var(--g7); color:#fff; border-radius:3px; }
.gd.dark .step i { background:var(--g5); color:#03130c; }
.gd .ctl { display:flex; flex-wrap:wrap; gap:18px; align-items:center; padding:10px 14px; margin:6px 0 12px;
  background:#fff; border:1px dashed var(--g3); border-radius:4px; }
.gd .ctl::before { content:'TRY IT'; font:700 10.5px 'JetBrains Mono',monospace; letter-spacing:.18em;
  color:#fff; background:var(--g6); padding:3px 7px; border-radius:2px; }
.gd pre { margin:0; }
@keyframes gdDash { to { stroke-dashoffset:-28; } }
@keyframes gdPulse { 0%,100% { opacity:.15; } 50% { opacity:.65; } }
.gd .flow { stroke-dasharray:3 11; animation:gdDash 1.1s linear infinite; }
.gd .halo { animation:gdPulse 2.6s ease-in-out infinite; }
"""
)

def H(part) -> str:
    """HTML text of a string or any marimo-renderable object (UI elements stay live)."""
    return part if isinstance(part, str) else mo.as_html(part).text

def md(text: str) -> str:
    """Markdown/LaTeX → HTML string (KaTeX rendered by marimo)."""
    return mo.md(text).text

def cols(*parts, t="1fr 1fr", gap=26, align="start") -> str:
    _inner = "".join(
        f"<div style='min-width:0'>{H(p)}</div>" for p in parts
    )
    return f"<div class='cols' style='grid-template-columns:{t};gap:{gap}px;align-items:{align}'>{_inner}</div>"

def card(title, body, color="var(--g5)") -> str:
    return f"<div class='card' style='--c:{color}'><h3>{title}</h3><p>{body}</p></div>"

def ctl(*widgets) -> str:
    return "<div class='ctl'>" + "".join(H(w) for w in widgets) + "</div>"

def slide(*parts, kicker="", title="", lead="", tone="light", chapter=""):
    """Wrap content in the themed slide frame."""
    _head = ""
    if kicker:
        _head += f"<div class='kicker'>{kicker}</div>"
    if title:
        _head += f"<h2>{title}</h2>"
    if lead:
        _head += f"<p class='lead'>{lead}</p>"
    _body = "".join(H(p) for p in parts)
    _foot = (
        f"<div class='foot'><span>Graph Neural Networks</span><span>{chapter}</span></div>"
        if chapter
        else ""
    )
    return mo.Html(
        f"<style>{CSS}</style><div style='container-type:inline-size;width:100%'>"
        f"<div class='gd {tone}'>{_head}{_body}{_foot}</div></div>"
    )

def tidy(chart):
    """Theme an Altair chart to match the deck."""
    return (
        chart.configure(background="transparent", font="Inter")
        .configure_view(stroke=None)
        .configure_axis(
            labelColor="#4d6a5c",
            titleColor="#07160f",
            gridColor="#dcebe2",
            domainColor="#a9c7b6",
            tickColor="#a9c7b6",
            labelFontSize=11,
            titleFontSize=12,
            titleFontWeight=600,
        )
        .configure_legend(
            labelColor="#07160f",
            titleColor="#4d6a5c",
            labelFontSize=11,
            titleFontSize=11,
        )
        .configure_title(
            font="Space Grotesk",
            fontSize=15,
            color="#07160f",
            anchor="start",
        )
    )

CLASS_NAMES = [
    "Case-Based",
    "Genetic Algorithms",
    "Neural Networks",
    "Probabilistic Methods",
    "Reinforcement Learning",
    "Rule Learning",
    "Theory",
]
CLASS_COLORS = [
    "#f59e0b",
    "#a3e635",
    "#10b981",
    "#0ea5e9",
    "#8b5cf6",
    "#f43f5e",
    "#64748b",
]
CLASS_SCALE = alt.Scale(domain=CLASS_NAMES, range=CLASS_COLORS)

def graph_chart(
    pos,
    edges,
    nodes,
    color,
    *,
    size=300,
    width=420,
    height=320,
    title="",
    labels=True,
    edge_color="#8fb5a1",
    tooltip=None,
    stroke=None,
):
    """Node-link drawing. `edges`: DataFrame with u, v (and optional w = stroke width);
    `nodes`: DataFrame with a `node` column plus fields used by `color` / `tooltip`."""
    _e = edges.copy()
    _e["x"] = [pos[u][0] for u in _e["u"]]
    _e["y"] = [pos[u][1] for u in _e["u"]]
    _e["x2"] = [pos[v][0] for v in _e["v"]]
    _e["y2"] = [pos[v][1] for v in _e["v"]]
    _n = nodes.assign(
        x=[pos[n][0] for n in nodes["node"]],
        y=[pos[n][1] for n in nodes["node"]],
    )
    _x = alt.X("x:Q", axis=None, scale=alt.Scale(zero=False, padding=26))
    _y = alt.Y("y:Q", axis=None, scale=alt.Scale(zero=False, padding=26))
    _enc = dict(x=_x, y=_y, x2="x2:Q", y2="y2:Q")
    if "w" in _e:
        _edge = (
            alt.Chart(_e)
            .mark_rule(color=edge_color, opacity=0.8, strokeCap="round")
            .encode(
                **_enc,
                strokeWidth=alt.StrokeWidth(
                    "w:Q", legend=None, scale=None
                ),
            )
        )
    else:
        _edge = (
            alt.Chart(_e)
            .mark_rule(color=edge_color, opacity=0.55, strokeWidth=1.4)
            .encode(**_enc)
        )
    _tt = tooltip or [c for c in nodes.columns if c not in ("x", "y")]
    _node_enc = dict(x=_x, y=_y, color=color, tooltip=_tt)
    if stroke is not None:
        _node_enc["strokeWidth"] = stroke
    _layers = [
        _edge,
        alt.Chart(_n)
        .mark_circle(size=size, opacity=1, stroke="#07160f", strokeWidth=0)
        .encode(**_node_enc),
    ]
    if labels:
        _layers.append(
            alt.Chart(_n)
            .mark_text(fontSize=10.5, fontWeight="bold", color="white")
            .encode(x=_x, y=_y, text="node:N")
        )
    return alt.layer(*_layers).properties(
        width=width, height=height, title=title
    )

def tree_svg(root, *, width=640, level_h=78, title_left=True):
    """Draw a rooted tree {'name','color','children'} with the root on top (output layer)
    and leaves at the bottom (input features)."""
    _leaves = []
    _depth = [0]

    def _count(n, d):
        _depth[0] = max(_depth[0], d)
        if not n.get("children"):
            _leaves.append(n)
        for c in n.get("children", []):
            _count(c, d + 1)

    _count(root, 0)
    _L = max(len(_leaves), 1)
    _D = _depth[0]
    _pad = 70 if title_left else 20
    _r = max(5, min(15, (width - _pad - 20) / _L / 2.4))
    _font = max(7, min(12, _r * 0.95))
    _h = _D * level_h + 2 * _r + 30
    _items = []
    _counter = [0]

    def _place(n, d):
        if not n.get("children"):
            _x = _pad + (_counter[0] + 0.5) * (width - _pad - 10) / _L
            _counter[0] += 1
        else:
            _xs = [_place(c, d + 1) for c in n["children"]]
            _x = sum(_xs) / len(_xs)
        _y = 18 + _r + d * level_h
        n["_xy"] = (_x, _y)
        return _x

    _place(root, 0)

    def _draw(n):
        _x, _y = n["_xy"]
        for c in n.get("children", []):
            _cx, _cy = c["_xy"]
            _items.append(
                f"<path d='M{_cx:.1f},{_cy - _r:.1f} C{_cx:.1f},{(_cy + _y) / 2:.1f} {_x:.1f},{(_cy + _y) / 2:.1f} {_x:.1f},{_y + _r:.1f}' "
                f"stroke='#8fb5a1' stroke-width='1.3' fill='none'/>"
            )
            _draw(c)
        if n.get("children"):
            _items.append(
                f"<rect x='{_x - _r - 4:.1f}' y='{_y - _r - 4:.1f}' width='{2 * _r + 8:.1f}' height='{2 * _r + 8:.1f}' "
                f"rx='2' fill='none' stroke='#10b981' stroke-width='1.4' stroke-dasharray='3 2'/>"
            )
        _items.append(
            f"<circle cx='{_x:.1f}' cy='{_y:.1f}' r='{_r:.1f}' fill='{n.get('color', '#10b981')}'/>"
            f"<text x='{_x:.1f}' y='{_y + _font * 0.36:.1f}' font-size='{_font:.1f}' text-anchor='middle' "
            f"fill='white' font-weight='700' font-family='JetBrains Mono'>{n['name']}</text>"
        )

    _draw(root)
    if title_left:
        for _d in range(_D + 1):
            _lab = (
                "output"
                if _d == 0
                else ("input x" if _d == _D else f"layer {_D - _d}")
            )
            _lab = f"layer {_D}" if _d == 0 else _lab
            _items.append(
                f"<text x='2' y='{18 + _r + _d * level_h + 4:.1f}' font-size='11' fill='#4d6a5c' "
                f"font-family='JetBrains Mono'>{_lab}</text>"
            )
    return (
        f"<svg viewBox='0 0 {width} {_h:.0f}' width='100%' style='max-height:{_h:.0f}px'>"
        + "".join(_items)
        + "</svg>"
    )

def edges_df(pairs, **extra):
    _df = pd.DataFrame(list(pairs), columns=["u", "v"])
    for k, v in extra.items():
        _df[k] = v
    return _df

dataset = Planetoid(
    root=str(DATA_DIR / "Planetoid"),
    name="Cora",
    transform=NormalizeFeatures(),
)
cora = dataset[0]
cora_y = cora.y.numpy()
cora_split = np.full(cora.num_nodes, "unlabelled", dtype=object)
cora_split[cora.train_mask.numpy()] = "train"
cora_split[cora.val_mask.numpy()] = "val"
cora_split[cora.test_mask.numpy()] = "test"
G_cora = nx.Graph()
G_cora.add_nodes_from(range(cora.num_nodes))
G_cora.add_edges_from(cora.edge_index.t().tolist())
cora_deg = degree(cora.edge_index[0], cora.num_nodes).long().numpy()
cora_stats = {
    "nodes": cora.num_nodes,
    "edges": cora.num_edges // 2,
    "features": cora.num_features,
    "classes": dataset.num_classes,
    "train": int(cora.train_mask.sum()),
    "val": int(cora.val_mask.sum()),
    "test": int(cora.test_mask.sum()),
    "homophily": homophily(cora.edge_index, cora.y, method="edge"),
    "mean_degree": float(cora_deg.mean()),
    "avg_words": float((cora.x > 0).sum(1).float().mean()),
}
_ = torch.manual_seed(0)

_nodes = [
    (330, 70), (430, 40), (520, 110), (390, 160), (470, 200), (560, 230),
    (300, 210), (380, 270), (480, 310), (590, 330), (330, 340), (420, 380),
]
_edges = [(0, 1), (1, 2), (0, 3), (1, 3), (3, 4), (2, 4), (4, 5), (3, 6), (6, 7),
          (4, 7), (7, 8), (5, 8), (8, 9), (7, 10), (10, 11), (8, 11), (2, 5)]
_svg = ["<svg viewBox='260 10 360 400' width='100%' style='max-width:420px'>"]
for _i, (_a, _b) in enumerate(_edges):
    (_x1, _y1), (_x2, _y2) = _nodes[_a], _nodes[_b]
    _svg.append(f"<line x1='{_x1}' y1='{_y1}' x2='{_x2}' y2='{_y2}' stroke='rgba(110,231,183,.25)' stroke-width='1.5'/>")
    _svg.append(
        f"<line class='flow' x1='{_x1}' y1='{_y1}' x2='{_x2}' y2='{_y2}' stroke='#a3e635' "
        f"stroke-width='2.2' style='animation-delay:-{_i * 0.13:.2f}s'/>"
    )
for _i, (_x, _y) in enumerate(_nodes):
    _c = "#10b981" if _i % 3 else "#a3e635"
    _svg.append(
        f"<circle class='halo' cx='{_x}' cy='{_y}' r='22' fill='{_c}' style='animation-delay:-{_i * 0.31:.2f}s'/>"
        f"<circle cx='{_x}' cy='{_y}' r='11' fill='{_c}' stroke='#03130c' stroke-width='3'/>"
    )
_svg.append("</svg>")

slide(
    f"""
    <div class='cols' style='grid-template-columns:1.25fr 1fr;align-items:center;min-height:520px'>
      <div>
        <div class='kicker'>Machine Learning with Graphs · CS224W L2 → L5</div>
        <h1>Graph Neural<br/><span class='hl'>Networks</span></h1>
        <p class='lead' style='font-size:20px'>Learning by <b style='color:#e3f5ea'>listening to your neighbours</b>:
        from lookup-table embeddings to message passing, GCN, training, and the GNN design space.</p>
        <div style='margin-top:26px'>
          <span class='tag'>node embeddings</span><span class='tag'>message passing</span>
          <span class='tag'>GCN · GraphSAGE · GAT</span><span class='tag'>PyTorch Geometric</span>
          <span class='tag'>live on Cora</span>
        </div>
      </div>
      <div style='text-align:center'>{"".join(_svg)}</div>
    </div>
    """,
    tone="dark",
)

_items = [
    ("01", "Recap · node embeddings", "Shallow lookup tables, and the four things they can't do.", "#10b981"),
    ("02", "Message passing", "The core idea: a node is computed from its neighbourhood.", "#a3e635"),
    ("03", "GCN", "The simplest useful layer, its matrix form, and its PyG code.", "#0ea5e9"),
    ("04", "Learning & prediction", "Pipeline, losses, splits. Train live on Cora and inspect the predictions.", "#f59e0b"),
    ("05", "Design space", "GraphSAGE, GAT, deep stacks, over-smoothing, augmentation, sampling.", "#8b5cf6"),
    ("06", "PyG in practice", "A six-step recipe and the tips that matter.", "#f43f5e"),
]
_grid = "".join(
    f"<div class='card' style='--c:{c}'><div style='font:700 28px Space Grotesk;color:{c}'>{n}</div>"
    f"<h3 style='color:var(--ink)'>{t}</h3><p>{d}</p></div>"
    for n, t, d, c in _items
)
slide(
    f"<div class='cols' style='grid-template-columns:repeat(3,1fr);gap:16px'>{_grid}</div>"
    "<p class='lead' style='margin-top:22px'>Every slide marked <span class='badge'>TRY IT</span> is live: "
    "the numbers and plots come from code running on the real <b>Cora</b> citation graph.</p>",
    kicker="Roadmap",
    title="Where we are going",
)

slide(
    """
    <div style='min-height:520px;display:flex;flex-direction:column;justify-content:center'>
      <div class='num'>01</div>
      <h1 style='margin-top:10px'>Recap: <span class='hl'>node embeddings</span></h1>
      <p class='lead' style='font-size:20px'>Lecture 2 in two slides. What a shallow embedding learns, and where it breaks.</p>
    </div>
    """,
    tone="dark",
)

karate = nx.karate_club_graph()
karate_pos = nx.kamada_kawai_layout(karate)
walk_start = mo.ui.slider(
    0, 33, value=0, label="start node", show_value=True
)
walk_len = mo.ui.slider(
    2, 16, value=8, label="walk length", show_value=True
)
walk_seed = mo.ui.button(
    value=0, on_click=lambda v: v + 1, label="🎲 new walk"
)

_rng = np.random.default_rng(walk_seed.value)
_walk = [walk_start.value]
for _ in range(walk_len.value):
    _walk.append(int(_rng.choice(list(karate.neighbors(_walk[-1])))))
_wedges = set(tuple(sorted(p)) for p in zip(_walk[:-1], _walk[1:]))
_all = edges_df([tuple(sorted(e)) for e in karate.edges()])
_all["w"] = [
    4.5 if (u, v) in _wedges else 1.0 for u, v in zip(_all.u, _all.v)
]
_nodes = pd.DataFrame(
    {
        "node": list(karate.nodes()),
        "role": [
            "start"
            if n == walk_start.value
            else ("on walk" if n in _walk else "other")
            for n in karate.nodes()
        ],
        "club": [karate.nodes[n]["club"] for n in karate.nodes()],
    }
)
_chart = graph_chart(
    karate_pos,
    _all,
    _nodes,
    alt.Color(
        "role:N",
        scale=alt.Scale(
            domain=["start", "on walk", "other"],
            range=["#047857", "#10b981", "#b9d4c5"],
        ),
        legend=alt.Legend(orient="bottom", title=None),
    ),
    size=380,
    width=540,
    height=400,
    edge_color="#5fae86",
)
# positive (u, v) pairs = nodes within a window of 2 on the walk
_pairs = sorted(
    {
        (a, b)
        for i, a in enumerate(_walk)
        for b in _walk[i + 1 : i + 3]
        if a != b
    }
)[:8]
_seq = " → ".join(
    f"<b>{n}</b>" if i == 0 else str(n) for i, n in enumerate(_walk)
)
slide(
    cols(
        md(
            r"""
**Encoder** — a lookup table, one free column per node:
$\;\mathrm{ENC}(u) = \mathbf{Z}\,\mathbf{1}_u,\; \mathbf{Z}\in\mathbb{R}^{d\times|V|}$

**Decoder** — a dot product: $\;\mathrm{DEC}(\mathbf{z}_u,\mathbf{z}_v)=\mathbf{z}_u^\top\mathbf{z}_v$

**Similarity** — $u,v$ co-occur on short **random walks** (DeepWalk, node2vec)

$$\mathcal{L}=\sum_{u}\sum_{v\in N_R(u)}-\log\frac{e^{\mathbf{z}_u^\top\mathbf{z}_v}}{\sum_{n\in V}e^{\mathbf{z}_u^\top\mathbf{z}_n}}$$
"""
        )
        + f"<div class='card' style='margin-top:6px'><h3>Your walk</h3><p style='font-family:JetBrains Mono;font-size:13px'>{_seq}</p>"
        f"<p style='margin-top:8px;font-size:13px'>positive pairs (window 2): "
        + ", ".join(f"<code>({a},{b})</code>" for a, b in _pairs)
        + " …</p></div>",
        ctl(walk_start, walk_len, walk_seed) + H_chart(_chart, tidy),
        t="1fr 1.05fr",
    ),
    kicker="01 · recap",
    title="Shallow embeddings: <span class='hl'>close in the graph → close in space</span>",
    chapter="01 recap",
)

def H_chart(chart, tidy):
    """Render a themed Altair chart as HTML text for embedding in a slide."""
    return mo.as_html(tidy(chart)).text

def cached_tsne(name, X):
    """t-SNE is ~10 s on Cora, so every projection is cached on disk by name."""
    _path = CACHE_DIR / f"tsne_{name}.npy"
    if _path.exists():
        return np.load(_path)
    _xy = TSNE(n_components=2, init="pca", random_state=0, perplexity=30).fit_transform(X)
    np.save(_path, _xy)
    return _xy


_dw_path = CACHE_DIR / "cora_deepwalk_p1_q1_d64.npy"
# DeepWalk embedding trained in notebooks/cs224w/L02_node_embeddings.py (10 walks/node, length 20, window 5)
dw_Z = np.load(_dw_path) if _dw_path.exists() else None
_tr, _te = cora.train_mask.numpy(), cora.test_mask.numpy()
_X = cora.x.numpy()
raw_acc = LogisticRegression(max_iter=3000).fit(_X[_tr], cora_y[_tr]).score(_X[_te], cora_y[_te])
dw_acc = (
    LogisticRegression(max_iter=3000).fit(dw_Z[_tr], cora_y[_tr]).score(dw_Z[_te], cora_y[_te])
    if dw_Z is not None
    else float("nan")
)
dw_xy = cached_tsne("deepwalk", dw_Z) if dw_Z is not None else None
raw_xy = cached_tsne("raw_features", _X)

mo.stop(dw_xy is None, mo.md("DeepWalk embedding missing from the cache folder."))
_df = pd.DataFrame({"x": dw_xy[:, 0], "y": dw_xy[:, 1], "topic": [CLASS_NAMES[c] for c in cora_y]})
_chart = (
    alt.Chart(_df)
    .mark_circle(size=14, opacity=0.8)
    .encode(
        x=alt.X("x:Q", axis=None),
        y=alt.Y("y:Q", axis=None),
        color=alt.Color("topic:N", scale=CLASS_SCALE, legend=alt.Legend(orient="right", title="paper topic")),
    )
    .properties(width=470, height=360, title="DeepWalk on Cora (t-SNE of z_u), colour = true topic")
)
slide(
    cols(
        H_chart(_chart, tidy),
        f"""
        <div class='stat'>{dw_acc:.1%}<small>test accuracy · DeepWalk z + logistic regression</small></div>
        <div class='stat' style='margin-top:18px;color:#4d6a5c'>{raw_acc:.1%}<small>same classifier on the 1,433 word features</small></div>
        <p style='margin-top:20px'>Topics form clusters <b>without ever seeing a label or a word</b>: the citation
        structure alone carries the signal. But the table has
        <code>{cora_stats["nodes"]:,} × 64 = {cora_stats["nodes"] * 64:,}</code> free parameters, one row per paper.</p>
        """,
        t="1.35fr 1fr",
        align="center",
    ),
    kicker="01 · recap · real data",
    title="It works: random walks recover the topics",
    chapter="01 recap",
)

_svg = """
<svg viewBox='0 0 420 330' width='100%'>
  <text x='10' y='18' font-family='JetBrains Mono' font-size='12' font-weight='700' fill='#047857'>SHALLOW · lookup table Z</text>
  <g font-family='JetBrains Mono' font-size='12'>
    <text x='10' y='52' fill='#4d6a5c'>u₁</text><rect x='40' y='40' width='150' height='16' fill='#10b981'/>
    <text x='10' y='76' fill='#4d6a5c'>u₂</text><rect x='40' y='64' width='150' height='16' fill='#a3e635'/>
    <text x='10' y='100' fill='#4d6a5c'>u₃</text><rect x='40' y='88' width='150' height='16' fill='#0ea5e9'/>
    <text x='10' y='124' fill='#4d6a5c'>⋮</text><rect x='40' y='112' width='150' height='16' fill='#d1fae5'/>
    <text x='10' y='148' fill='#f43f5e' font-weight='700'>new</text>
    <rect x='40' y='136' width='150' height='16' fill='none' stroke='#f43f5e' stroke-dasharray='4 3' stroke-width='2'/>
    <text x='115' y='149' text-anchor='middle' fill='#f43f5e' font-weight='700'>no row → ✗</text>
    <text x='205' y='100' fill='#4d6a5c'>|V| × d free</text><text x='205' y='116' fill='#4d6a5c'>parameters</text>
  </g>
  <line x1='10' y1='172' x2='410' y2='172' stroke='#d2e3d9'/>
  <text x='10' y='196' font-family='JetBrains Mono' font-size='12' font-weight='700' fill='#047857'>GNN · one shared function f</text>
  <g stroke='#8fb5a1' stroke-width='2'><line x1='40' y1='225' x2='110' y2='265'/><line x1='40' y1='305' x2='110' y2='265'/><line x1='60' y1='265' x2='110' y2='265'/></g>
  <circle cx='40' cy='225' r='11' fill='#10b981'/><circle cx='40' cy='305' r='11' fill='#a3e635'/><circle cx='60' cy='265' r='11' fill='#0ea5e9'/>
  <circle cx='110' cy='265' r='14' fill='#fff' stroke='#f43f5e' stroke-width='3'/>
  <text x='110' y='295' text-anchor='middle' font-family='JetBrains Mono' font-size='11' fill='#f43f5e'>new</text>
  <line x1='126' y1='265' x2='196' y2='265' stroke='#10b981' stroke-width='2.5' marker-end='url(#aw)'/>
  <defs><marker id='aw' viewBox='0 0 10 10' refX='9' refY='5' markerWidth='7' markerHeight='7' orient='auto'><path d='M0,0L10,5L0,10z' fill='#10b981'/></marker></defs>
  <rect x='200' y='240' width='90' height='50' rx='3' fill='#047857'/>
  <text x='245' y='271' text-anchor='middle' font-family='Space Grotesk' font-weight='700' font-size='17' fill='#fff'>f(·), W</text>
  <line x1='292' y1='265' x2='330' y2='265' stroke='#10b981' stroke-width='2.5' marker-end='url(#aw)'/>
  <text x='340' y='270' font-family='JetBrains Mono' font-size='13' font-weight='700' fill='#047857'>z_new ✓</text>
</svg>"""
slide(
    cols(
        f"""<div class='cols' style='grid-template-columns:1fr 1fr;gap:14px'>
          {card("① O(|V|·d) parameters", "Every node owns a private vector. ogbn-papers100M: 111M nodes × 128 dims ≈ 14 billion parameters. Nothing learned about one node helps another.", "#f43f5e")}
          {card("② Transductive only", "A new paper has no row in Z: re-run walks and re-train. Twitter, Reddit and Scholar get new nodes every second.", "#f59e0b")}
          {card("③ Ignores node features", "Cora's 1,433 words per paper, a user's bio, an atom's type: the lookup table never sees them.", "#8b5cf6")}
          {card("④ Proximity, not role", "Random walks say 'u and v are near'. Two hubs in different corners never share a walk, so roles are lost.", "#0ea5e9")}
        </div>
        <div class='eq' style='margin-top:16px;border-left-color:var(--lime)'>
          <p style='margin:10px 0'><b>What we want:</b> one <b>shared</b> function
          <code>f(own features, neighbours' features)</code>: parameters independent of |V|, works on unseen nodes and graphs,
          uses features. That function is a <span class='hl' style='font-weight:700'>Graph Neural Network</span>.</p>
        </div>""",
        _svg,
        t="1.5fr 1fr",
        align="center",
    ),
    kicker="01 · recap",
    title="Four cracks in the lookup table",
    chapter="01 recap",
)

slide(
    """
    <div style='min-height:520px;display:flex;flex-direction:column;justify-content:center'>
      <div class='num'>02</div>
      <h1 style='margin-top:10px'>Message <span class='hl'>passing</span></h1>
      <p class='lead' style='font-size:20px'>"Tell me who your neighbours are and I'll tell you who you are."<br/>
      The one idea behind every GNN.</p>
    </div>
    """,
    tone="dark",
)

_nb = [(70, 60, "#a3e635"), (330, 70, "#10b981"), (60, 260, "#f59e0b"), (330, 270, "#0ea5e9")]
_svg = "<svg viewBox='0 0 400 330' width='100%' style='max-width:420px'>"
for _i, (_x, _y, _c) in enumerate(_nb):
    _svg += f"<line x1='{_x}' y1='{_y}' x2='200' y2='165' stroke='#cfe3d7' stroke-width='2'/>"
    _svg += (
        f"<line class='flow' x1='{_x}' y1='{_y}' x2='200' y2='165' stroke='{_c}' stroke-width='3.5' "
        f"style='animation-delay:-{_i * .27:.2f}s'/>"
    )
for _i, (_x, _y, _c) in enumerate(_nb):
    _svg += (
        f"<circle cx='{_x}' cy='{_y}' r='24' fill='{_c}'/><text x='{_x}' y='{_y + 5}' text-anchor='middle' "
        f"fill='white' font-family='JetBrains Mono' font-weight='700' font-size='14'>u{_i + 1}</text>"
        f"<text x='{_x}' y='{_y + (42 if _y > 165 else -34)}' text-anchor='middle' fill='#4d6a5c' "
        f"font-family='JetBrains Mono' font-size='12'>m = W·h_u{_i + 1}</text>"
    )
_svg += (
    "<circle class='halo' cx='200' cy='165' r='52' fill='#10b981'/>"
    "<circle cx='200' cy='165' r='34' fill='#047857' stroke='#fff' stroke-width='3'/>"
    "<text x='200' y='171' text-anchor='middle' fill='white' font-family='Space Grotesk' font-weight='700' font-size='20'>v</text>"
    "<text x='200' y='232' text-anchor='middle' fill='#047857' font-family='JetBrains Mono' font-size='12' font-weight='700'>AGG ⊕ → UPDATE</text>"
    "</svg>"
)
slide(
    cols(
        """
        <div class='step'><i>1</i><div><h3>Message</h3><p>every neighbour <code>u</code> prepares a vector,
        usually a learned linear map of its current state: <code>m_u = W h_u</code>.</p></div></div>
        <div class='step'><i>2</i><div><h3>Aggregate</h3><p><code>v</code> pools the incoming messages with an
        <b>order-independent</b> function: sum, mean, max, attention. Neighbours have no natural order!</p></div></div>
        <div class='step'><i>3</i><div><h3>Update</h3><p>combine with <code>v</code>'s own state and apply a
        non-linearity. That's one <b>layer</b>. Repeat K times → <code>v</code> sees its K-hop neighbourhood.</p></div></div>
        """
        + md(
            r"""
$$\mathbf{h}_v^{(k+1)}=\sigma\Big(\mathbf{W}_k\!\!\sum_{u\in N(v)}\!\frac{\mathbf{h}_u^{(k)}}{|N(v)|}+\mathbf{B}_k\mathbf{h}_v^{(k)}\Big),\qquad \mathbf{h}_v^{(0)}=\mathbf{x}_v$$
"""
        ),
        _svg,
        t="1.15fr 1fr",
        align="center",
    ),
    kicker="02 · the core idea",
    title="Every node <span class='hl'>listens</span> to its neighbours",
    lead="The same small network runs at every node, so the weights are <b>shared</b>: parameters don't grow with |V|, and any new node can be embedded.",
    chapter="02 message passing",
)

PLAYGROUND = r"""
<!doctype html><html><head><meta charset="utf-8">
<link href="https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@600;700&family=Inter:wght@400;600&family=JetBrains+Mono:wght@500;700&display=swap" rel="stylesheet">
<style>
*{box-sizing:border-box} body{margin:0;background:#03130c;color:#e3f5ea;font-family:Inter,system-ui,sans-serif;overflow:hidden}
.wrap{display:grid;grid-template-columns:1.35fr 1fr;gap:18px;padding:14px 18px;height:100vh}
svg{width:100%;height:100%;background:
  linear-gradient(rgba(16,185,129,.06) 1px,transparent 1px) 0 0/30px 30px,
  linear-gradient(90deg,rgba(16,185,129,.06) 1px,transparent 1px) 0 0/30px 30px;border:1px solid #0d3b28;border-radius:4px}
.panel{display:flex;flex-direction:column;gap:10px;min-width:0}
h3{font:700 13px 'JetBrains Mono';letter-spacing:.18em;text-transform:uppercase;color:#6ee7b7;margin:4px 0 0}
.row{display:flex;gap:8px;flex-wrap:wrap;align-items:center}
button{font:600 13px 'JetBrains Mono';background:#0b2a1d;color:#d1fae5;border:1px solid #1f6b4a;padding:7px 11px;border-radius:3px;cursor:pointer}
button:hover{background:#10b981;color:#03130c} button.on{background:#a3e635;color:#03130c;border-color:#a3e635}
.seg button{border-radius:0;margin-right:-1px}
label{font:500 12px 'JetBrains Mono';color:#a7cdb9}
.big{font:700 34px 'Space Grotesk';color:#a3e635;line-height:1}
.card{border:1px solid #1a4d36;background:rgba(255,255,255,.03);padding:10px 12px;border-radius:4px;font-size:13px;line-height:1.5}
.bars{display:flex;gap:6px;align-items:flex-end;height:54px;margin-top:6px}
.bars div{flex:1;border-radius:2px 2px 0 0;transition:height .5s}
.mono{font-family:'JetBrains Mono';font-size:12px;color:#a7cdb9}
canvas{width:100%;height:70px;border:1px solid #1a4d36;border-radius:4px}
circle.node{transition:fill .55s, r .4s;cursor:pointer;stroke:#03130c;stroke-width:2.5}
.lbl{font:700 11px 'JetBrains Mono';fill:#03130c;pointer-events:none}
.ring{fill:none;stroke-width:2;transition:opacity .3s}
</style></head><body>
<div class="wrap">
  <svg id="g" viewBox="0 0 600 470"></svg>
  <div class="panel">
    <h3>Message-passing playground</h3>
    <div class="mono">Each node holds a 3-d "topic" vector, drawn as a colour mix of
      <b style="color:#10b981">■ A</b> <b style="color:#f59e0b">■ B</b> <b style="color:#a78bfa">■ C</b>.
      Click a node to inspect it.</div>
    <div class="row">
      <button id="step">▶ 1 layer</button><button id="auto">⏵ auto</button>
      <button id="reset">↺ reset</button><button id="noise">🎲 noisy labels</button>
    </div>
    <div class="row seg"><label style="margin-right:8px">AGG</label>
      <button data-agg="mean" class="on">mean</button><button data-agg="sum">sum</button><button data-agg="max">max</button>
      <label style="margin-left:12px"><input type="checkbox" id="self" checked> self-loop</label>
    </div>
    <div class="row"><label>receptive field K = <b id="kv">2</b></label>
      <input type="range" id="k" min="0" max="5" value="2" style="flex:1"></div>
    <div class="row" style="justify-content:space-between">
      <div><div class="mono">layer</div><div class="big" id="layer">0</div></div>
      <div><div class="mono">nodes in K-hop field</div><div class="big" id="field">–</div></div>
      <div><div class="mono">spread</div><div class="big" id="spread">–</div></div>
    </div>
    <div class="card" id="info">Click a node.</div>
    <div class="mono">spread = mean distance of node vectors to their average. → 0 means over-smoothing</div>
    <canvas id="cv" width="420" height="70"></canvas>
  </div>
</div>
<script>
const COL=[[16,185,129],[245,158,11],[167,139,250]];
const P=[[95,95],[60,185],[150,150],[170,60],[110,260],[205,225],
         [415,70],[505,55],[470,145],[545,175],[385,165],[525,250],
         [255,345],[330,300],[210,420],[320,420],[395,385],[290,215]];
const COMM=[0,0,0,0,0,0,1,1,1,1,1,1,2,2,2,2,2,2];
const E=[[0,1],[0,2],[0,3],[1,2],[1,4],[2,3],[2,5],[4,5],[2,4],
         [6,7],[6,8],[7,8],[8,9],[8,10],[6,10],[9,11],[8,11],
         [12,13],[12,14],[13,15],[14,15],[15,16],[13,16],[12,17],[13,17],
         [5,17],[10,17],[3,6],[11,16]];
const n=P.length, NB=[...Array(n)].map(()=>[]);
E.forEach(([a,b])=>{NB[a].push(b);NB[b].push(a)});
let H0, H, layer=0, agg='mean', self=true, sel=null, K=2, timer=null, hist=[], busy=false;
const svg=document.getElementById('g'), NS='http://www.w3.org/2000/svg';
function el(t,a){const e=document.createElementNS(NS,t);for(const k in a)e.setAttribute(k,a[k]);return e}
function onehot(c){const v=[0,0,0];v[c]=1;return v}
function init(noisy){
  H0=COMM.map((c,i)=>onehot(noisy && Math.random()<0.3 ? (c+1+Math.floor(Math.random()*2))%3 : c));
  H=H0.map(v=>v.slice()); layer=0; hist=[spread()]; draw(); info(); plot();
}
function norm(v){const s=v[0]+v[1]+v[2];return s>0?v.map(x=>x/s):[1/3,1/3,1/3]}
function color(v){const p=norm(v);const c=[0,1,2].map(j=>Math.round(p[0]*COL[0][j]+p[1]*COL[1][j]+p[2]*COL[2][j]));return `rgb(${c})`}
function mag(v){return Math.sqrt(v[0]**2+v[1]**2+v[2]**2)}
function spread(){const N=H.map(norm);const m=[0,1,2].map(j=>N.reduce((s,v)=>s+v[j],0)/n);
  return N.reduce((s,v)=>s+Math.hypot(v[0]-m[0],v[1]-m[1],v[2]-m[2]),0)/n}
function hops(src){const d=Array(n).fill(Infinity);d[src]=0;const q=[src];
  while(q.length){const u=q.shift();for(const w of NB[u])if(d[w]===Infinity){d[w]=d[u]+1;q.push(w)}}return d}
let gE, gD, gN;
function draw(){
  svg.innerHTML=''; gE=el('g',{}); gD=el('g',{}); gN=el('g',{}); svg.append(gE,gD,gN);
  const d = sel===null?null:hops(sel);
  E.forEach(([a,b])=>{const on=d&&d[a]<=K&&d[b]<=K&&!(d[a]===K&&d[b]===K);
    gE.append(el('line',{x1:P[a][0],y1:P[a][1],x2:P[b][0],y2:P[b][1],stroke:on?'#6ee7b7':'#1f5a40','stroke-width':on?2.6:1.6,opacity:d&&!on?.35:1}))});
  P.forEach(([x,y],i)=>{
    if(d&&d[i]<=K&&i!==sel){gN.append(el('circle',{class:'ring',cx:x,cy:y,r:22,stroke:['#a3e635','#6ee7b7','#34d399','#10b981','#059669','#047857'][d[i]],opacity:.9}));}
    if(i===sel){gN.append(el('circle',{class:'ring',cx:x,cy:y,r:26,stroke:'#fff','stroke-width':3}))}
    const r = agg==='sum' ? Math.min(13+4*Math.log2(1+mag(H[i])),30) : 15;
    const c=el('circle',{class:'node',cx:x,cy:y,r:r,fill:color(H[i]),opacity:d&&d[i]>K?.3:1});
    c.onclick=()=>{sel=(sel===i?null:i);draw();info()}; gN.append(c);
    const t=el('text',{class:'lbl',x:x,y:y+4,'text-anchor':'middle'}); t.textContent=i; gN.append(t);
  });
  document.getElementById('layer').textContent=layer;
  document.getElementById('spread').textContent=spread().toFixed(2);
  document.getElementById('field').textContent= d? d.filter(x=>x<=K).length+'/'+n : '–';
}
function info(){
  const box=document.getElementById('info');
  if(sel===null){box.innerHTML='Click a node to see its vector, its neighbours, and the receptive field of a K-layer GNN.';return}
  const v=H[sel], p=norm(v), nb=NB[sel];
  const src = self? [sel,...nb] : nb;
  box.innerHTML=`<b>node ${sel}</b> · degree ${nb.length} · |h| = ${mag(v).toFixed(2)}
    <div class="bars">${p.map((x,j)=>`<div style="height:${Math.max(2,x*54)}px;background:rgb(${COL[j]})"></div>`).join('')}</div>
    <div class="mono" style="margin-top:6px">next layer: h<sub>${sel}</sub> ← ${agg}( ${src.map(u=>'h<sub>'+u+'</sub>').join(', ')} )</div>`;
}
function update(){
  const out=H.map((v,i)=>{const src=self?[i,...NB[i]]:NB[i];
    if(agg==='max') return [0,1,2].map(j=>Math.max(...src.map(u=>H[u][j])));
    const s=[0,1,2].map(j=>src.reduce((a,u)=>a+H[u][j],0));
    return agg==='mean'? s.map(x=>x/src.length) : s});
  H=out; layer++; hist.push(spread());
}
function step(){
  if(busy)return; busy=true;
  const dots=[]; E.forEach(([a,b])=>{[[a,b],[b,a]].forEach(([s,t])=>{const c=el('circle',{r:4.5,fill:color(H[s]),stroke:'#fff','stroke-width':1});gD.append(c);dots.push([c,s,t])})});
  const t0=performance.now(), T=650;
  function frame(now){const k=Math.min(1,(now-t0)/T), e=k<.5?2*k*k:1-Math.pow(-2*k+2,2)/2;
    dots.forEach(([c,s,t])=>{c.setAttribute('cx',P[s][0]+(P[t][0]-P[s][0])*e);c.setAttribute('cy',P[s][1]+(P[t][1]-P[s][1])*e)});
    if(k<1)requestAnimationFrame(frame); else {update();draw();info();plot();busy=false}}
  requestAnimationFrame(frame);
}
function plot(){const cv=document.getElementById('cv'),x=cv.getContext('2d');const W=cv.width,Hh=cv.height;x.clearRect(0,0,W,Hh);
  const mx=Math.max(...hist,0.01),N=Math.max(hist.length,12);
  x.strokeStyle='#1a4d36';x.beginPath();x.moveTo(0,Hh-6);x.lineTo(W,Hh-6);x.stroke();
  x.strokeStyle='#a3e635';x.lineWidth=2.5;x.beginPath();
  hist.forEach((s,i)=>{const px=8+i*(W-16)/(N-1),py=Hh-6-(s/mx)*(Hh-14);i?x.lineTo(px,py):x.moveTo(px,py)});x.stroke();
  hist.forEach((s,i)=>{const px=8+i*(W-16)/(N-1),py=Hh-6-(s/mx)*(Hh-14);x.fillStyle='#10b981';x.beginPath();x.arc(px,py,3,0,7);x.fill()});}
document.getElementById('step').onclick=step;
document.getElementById('reset').onclick=()=>{stop();init(false)};
document.getElementById('noise').onclick=()=>{stop();init(true)};
function stop(){if(timer){clearInterval(timer);timer=null;document.getElementById('auto').classList.remove('on')}}
document.getElementById('auto').onclick=e=>{if(timer){stop()}else{e.target.classList.add('on');step();timer=setInterval(step,900)}};
document.querySelectorAll('[data-agg]').forEach(b=>b.onclick=()=>{document.querySelectorAll('[data-agg]').forEach(x=>x.classList.remove('on'));b.classList.add('on');agg=b.dataset.agg;stop();init(false)});
document.getElementById('self').onchange=e=>{self=e.target.checked;stop();init(false)};
document.getElementById('k').oninput=e=>{K=+e.target.value;document.getElementById('kv').textContent=K;draw()};
sel=2; init(true);
</script></body></html>
"""

slide(
    mo.iframe(PLAYGROUND, height="520px"),
    kicker="02 · try it · playground",
    title="Watch messages flow, layer by layer",
    lead="<b>▶ 1 layer</b> a few times: noisy nodes get fixed by their community. Keep going: one colour everywhere = <b>over-smoothing</b>. Try <b>sum</b>: vectors blow up.",
    tone="dark",
    chapter="02 message passing",
)

LECTURE_EDGES = [("A", "B"), ("A", "C"), ("A", "D"), ("B", "C"), ("C", "E"), ("C", "F"), ("E", "F")]
tree_root = mo.ui.dropdown(list("ABCDEF"), value="A", label="target node")
tree_depth = mo.ui.slider(1, 3, value=2, label="layers K", show_value=True)
tree_self = mo.ui.switch(value=False, label="self-loops")

_G = nx.Graph(LECTURE_EDGES)
_PAL = dict(
    zip(
        "ABCDEF",
        ["#047857", "#10b981", "#0ea5e9", "#f59e0b", "#8b5cf6", "#f43f5e"],
    )
)

def _build(v, k):
    _node = {"name": v, "color": _PAL[v]}
    if k > 0:
        _nb = sorted(_G.neighbors(v)) + ([v] if tree_self.value else [])
        _node["children"] = [_build(u, k - 1) for u in _nb]
    return _node

_tree = _build(tree_root.value, tree_depth.value)

def _size(t):
    return 1 + sum(_size(c) for c in t.get("children", []))

# small drawing of the input graph, hop distance from the target as ring colour
_pos = {
    "A": (60, 90),
    "B": (150, 40),
    "C": (170, 140),
    "D": (40, 180),
    "E": (260, 110),
    "F": (250, 200),
}
_hop = nx.single_source_shortest_path_length(_G, tree_root.value)
_g = "<svg viewBox='0 0 300 230' width='100%' style='max-width:380px'>"
for _a, _b in LECTURE_EDGES:
    _g += f"<line x1='{_pos[_a][0]}' y1='{_pos[_a][1]}' x2='{_pos[_b][0]}' y2='{_pos[_b][1]}' stroke='#8fb5a1' stroke-width='2'/>"
for _v, (_x, _y) in _pos.items():
    _in = _hop[_v] <= tree_depth.value
    _g += (
        f"<circle cx='{_x}' cy='{_y}' r='19' fill='{_PAL[_v]}' opacity='{1 if _in else 0.25}' "
        f"stroke='{'#07160f' if _v == tree_root.value else 'none'}' stroke-width='3'/>"
        f"<text x='{_x}' y='{_y + 5}' text-anchor='middle' fill='white' font-weight='700' font-family='JetBrains Mono'>{_v}</text>"
    )
_g += "</svg>"
slide(
    ctl(tree_root, tree_depth, tree_self),
    cols(
        _g
        + f"""
        <p><b>{_size(_tree)}</b> boxes in the computation graph of <b>{tree_root.value}</b>, built from
        <b>{sum(1 for d in _hop.values() if d <= tree_depth.value)}</b> distinct graph nodes.</p>
        <p class='muted' style='font-size:14px'>Dashed squares = "aggregate + neural net". All squares in a row
        share the same weights <code>W_k, B_k</code>. Nodes repeat: the tree is an <b>unrolling</b>, not the graph.
        A GNN can only tell two nodes apart if their trees differ (→ Lecture 6, WL test).</p>
        """,
        tree_svg(_tree, width=680, level_h=120),
        t="0.8fr 1.6fr",
    ),
    kicker="02 · computation graphs",
    title="Every node gets its <span class='hl'>own</span> neural network",
    chapter="02 message passing",
)

toy_edges_text = mo.ui.text(
    value="0-1, 1-2, 2-0, 2-3, 3-4, 4-5, 5-3",
    label="edges",
    full_width=False,
)

def _parse(text):
    _pairs = []
    for _tok in text.replace("\n", ",").split(","):
        if "-" in _tok:
            _a, _b = _tok.split("-", 1)
            _pairs.append((int(_a), int(_b)))
    return _pairs

try:
    toy_pairs = _parse(toy_edges_text.value)
except ValueError:
    toy_pairs = []
mo.stop(not toy_pairs, mo.callout(mo.md("Couldn't parse edges, use `0-1, 1-2`"), kind="danger"))
toy_n = max(max(p) for p in toy_pairs) + 1
# each undirected edge is stored in both directions: this is PyG's edge_index convention
toy_ei = torch.tensor(
    [[a for a, b in toy_pairs] + [b for a, b in toy_pairs],
     [b for a, b in toy_pairs] + [a for a, b in toy_pairs]]
)
toy_x = torch.randint(0, 6, (toy_n, 2), generator=torch.Generator().manual_seed(1)).float()
toy_pos = nx.spring_layout(nx.Graph(toy_pairs), seed=7)
for _i in range(toy_n):
    toy_pos.setdefault(_i, (0.0, 0.0))
toy_deg = degree(toy_ei[1], toy_n).long()

_nodes = pd.DataFrame(
    {"node": range(toy_n), "degree": toy_deg.numpy(), "x": [str(v.int().tolist()) for v in toy_x]}
)
_chart = graph_chart(
    toy_pos,
    edges_df(toy_pairs),
    _nodes,
    alt.Color("degree:O", scale=alt.Scale(scheme="greens"), legend=alt.Legend(orient="bottom")),
    size=520,
    width=330,
    height=260,
)
slide(
    ctl(toy_edges_text),
    cols(
        H_chart(_chart, tidy),
        md(
            f"""
```python
from torch_geometric.data import Data

data = Data(x=x, edge_index=edge_index)
# {f"Data(x=[{toy_n}, 2], edge_index=[2, {toy_ei.size(1)}])"}

x          = {toy_x.int().tolist()}
edge_index = [{toy_ei[0].tolist()},   # source j
              {toy_ei[1].tolist()}]   # target i
```
"""
        )
        + f"""<p style='font-size:14.5px'><b>x</b>: node features, one row per node · <b>edge_index</b>: a
        <code>[2, E]</code> list of (source → target) pairs, <b>{len(toy_pairs)}</b> undirected edges stored as
        <b>{toy_ei.size(1)}</b> directed ones. Sparse, so memory is O(E), not O(N²).</p>""",
        t="0.85fr 1.3fr",
        align="center",
    ),
    kicker="02 · PyG · graphs as tensors",
    title="A graph in PyTorch Geometric is two tensors",
    lead="Edit the edge list: everything on this and the next two slides re-computes.",
    chapter="02 message passing",
)

mp_node = mo.ui.dropdown({str(i): i for i in range(toy_n)}, value="2", label="target node")
mp_aggr = mo.ui.radio(["sum", "mean", "max"], value="mean", label="aggregation", inline=True)
mp_self = mo.ui.switch(value=False, label="self-loops")
mp_shuffle = mo.ui.button(value=0, on_click=lambda v: v + 1, label="🔀 shuffle neighbour order")

_ei = toy_ei
if mp_self.value:
    _ei, _ = add_self_loops(_ei, num_nodes=toy_n)
_i = mp_node.value
_src = _ei[0][_ei[1] == _i]
# shuffling the neighbour order changes nothing: the aggregation is permutation invariant
_perm = torch.randperm(
    len(_src), generator=torch.Generator().manual_seed(mp_shuffle.value)
)
_src = _src[_perm]
_msgs = toy_x[_src]
_agg = (
    torch.zeros(2)
    if len(_src) == 0
    else {
        "sum": lambda m: m.sum(0),
        "mean": lambda m: m.mean(0),
        "max": lambda m: m.max(0).values,
    }[mp_aggr.value](_msgs)
)
_all = scatter(
    toy_x[_ei[0]], _ei[1], dim=0, dim_size=toy_n, reduce=mp_aggr.value
)
_role = [
    "target"
    if n == _i
    else ("neighbour" if n in _src.tolist() else "other")
    for n in range(toy_n)
]
_chart = graph_chart(
    toy_pos,
    edges_df(toy_pairs),
    pd.DataFrame(
        {
            "node": range(toy_n),
            "role": _role,
            "x": [str(v.int().tolist()) for v in toy_x],
        }
    ),
    alt.Color(
        "role:N",
        scale=alt.Scale(
            domain=["target", "neighbour", "other"],
            range=["#047857", "#a3e635", "#cfe3d7"],
        ),
        legend=alt.Legend(orient="bottom", title=None),
    ),
    size=700,
    width=340,
    height=320,
)
_rows = "".join(
    f"<tr><td>{int(j)}</td><td><code>{m.int().tolist()}</code></td></tr>"
    for j, m in zip(_src, _msgs)
)
_out_rows = "\n".join(
    f"node {n}: {[round(v, 2) for v in r]}"
    + ("   ← target" if n == _i else "")
    for n, r in enumerate(_all.tolist())
)
slide(
    ctl(mp_node, mp_aggr, mp_self, mp_shuffle),
    cols(
        H_chart(_chart, tidy),
        f"<table class='gt'><tr><th>from j (in this order)</th><th>message x_j</th></tr>{_rows}</table>"
        + f"<div class='eq'><p style='margin:10px 0;font-size:16px'><b>{mp_aggr.value}</b> → "
        f"<code>x'_{_i} = {[round(v, 2) for v in _agg.tolist()]}</code></p></div>",
        md(
            f"""
**All nodes at once**, the way PyG runs it:

```python
from torch_geometric.utils import scatter
out = scatter(x[src], dst, dim=0,
              reduce="{mp_aggr.value}")
```
```
{_out_rows}
```
Row **{_i}** equals the hand computation. Press **shuffle**: the order of messages changes,
the result does not. Sum / mean / max are **permutation invariant**, so the whole layer is
**permutation equivariant**: relabel the nodes and the outputs are relabelled the same way.
"""
        ),
        t="1.05fr 0.8fr 1.1fr",
    ),
    kicker="02 · PyG · message passing by hand",
    title="Gather → <span class='hl'>scatter</span>: one line for the whole graph",
    chapter="02 message passing",
)

slide(
    """
    <div style='min-height:520px;display:flex;flex-direction:column;justify-content:center'>
      <div class='num'>03</div>
      <h1 style='margin-top:10px'>Graph Convolutional <span class='hl'>Network</span></h1>
      <p class='lead' style='font-size:20px'>Kipf & Welling, ICLR 2017. A degree-normalised average, a linear map, a ReLU.
      Still the baseline to beat.</p>
    </div>
    """,
    tone="dark",
)

slide(
    md(
        r"""
$$
\mathbf{h}_v^{(k+1)} \;=\; \sigma\Big(\; \underbrace{\sum_{u\in N(v)\,\cup\,\{v\}} \frac{1}{\sqrt{\tilde d_u\,\tilde d_v}}}_{\text{fixed weight from degrees}}\;\; \underbrace{\mathbf{W}^{(k)}\mathbf{h}_u^{(k)}}_{\text{message}} \Big)
\qquad\Longleftrightarrow\qquad
\mathbf{H}^{(k+1)} = \sigma\big(\hat{\mathbf{A}}\,\mathbf{H}^{(k)}\mathbf{W}^{(k)}\big),\;\;
\hat{\mathbf{A}}=\tilde{\mathbf{D}}^{-1/2}(\mathbf{A}+\mathbf{I})\tilde{\mathbf{D}}^{-1/2}
$$
"""
    ),
    cols(
        card(
            "Self-loops  A + I",
            "A node's own state enters as just another neighbour. Without it, the node forgets itself after one layer.",
        ),
        card(
            "Symmetric normalisation",
            "Divide by √(dᵤ dᵥ). A plain sum explodes at hubs; this tones down messages from very popular nodes.",
            "#a3e635",
        ),
        card(
            "Sparse matrix product",
            "Â is sparse, so a layer costs O(|E|·d). The whole graph in one SpMM call on a GPU.",
            "#0ea5e9",
        ),
        card(
            "Why 'convolution'?",
            "A CNN is a GNN on a pixel grid, with a separate weight per relative position. Graphs have no positions, so GCN shares one W.",
            "#f59e0b",
        ),
        t="repeat(4,1fr)",
        gap=14,
    ),
    md(
        r"""
<p class='muted' style='font-size:14px;margin-top:14px'>Lecture-3 variant: mean over neighbours + separate self weight,
$\mathbf{h}_v'=\sigma(\mathbf{W}\,\mathrm{mean}_{u\in N(v)}\mathbf{h}_u+\mathbf{B}\mathbf{h}_v)$, i.e.
$\mathbf{H}'=\sigma(\mathbf{D}^{-1}\mathbf{A}\mathbf{H}\mathbf{W}^\top+\mathbf{H}\mathbf{B}^\top)$. Both are "average-like".
The symmetric one (above) is what <code>GCNConv</code> and the homework use.</p>
"""
    ),
    kicker="03 · GCN",
    title="One layer = <span class='hl'>normalised average</span> + linear map + ReLU",
    chapter="03 GCN",
)

gcn_view = mo.ui.radio(
    ["A", "A + I", "Â", "X  →  ÂX"], value="Â", label="matrix", inline=True
)

_A = to_dense_adj(toy_ei, max_num_nodes=toy_n)[0]
_At = _A + torch.eye(toy_n)
_d = _At.sum(1)
_Ahat = _At / torch.sqrt(_d[:, None] * _d[None, :])

def _heat(mat, title, fmt, scheme="greens"):
    _r, _c = mat.shape
    _df = pd.DataFrame(
        [(i, j, float(mat[i, j])) for i in range(_r) for j in range(_c)],
        columns=["row", "col", "value"],
    )
    _base = alt.Chart(_df).encode(
        x=alt.X("col:O", title=None, axis=alt.Axis(orient="top")),
        y=alt.Y("row:O", title=None),
    )
    return (
        _base.mark_rect(stroke="white", strokeWidth=1.5).encode(
            color=alt.Color(
                "value:Q", scale=alt.Scale(scheme=scheme), legend=None
            )
        )
        + _base.mark_text(fontSize=11, font="JetBrains Mono").encode(
            text=alt.Text("value:Q", format=fmt),
            color=alt.condition(
                f"datum.value > {float(mat.max()) * 0.6}",
                alt.value("white"),
                alt.value("#07160f"),
            ),
        )
    ).properties(
        width=min(56 * _c, 420), height=min(50 * _r, 360), title=title
    )

if gcn_view.value == "X  →  ÂX":
    _chart = alt.hconcat(
        _heat(toy_x, "X (features)", ".0f"),
        _heat(_Ahat @ toy_x, "ÂX (one smoothing step)", ".2f"),
    )
    _note = (
        "Each row of <code>ÂX</code> is a weighted average of the node's own row and its neighbours' rows. "
        "Add <code>W</code> and a ReLU and you have a GCN layer: <code>H' = σ(ÂXW)</code>."
    )
else:
    _mat, _fmt = {
        "A": (_A, ".0f"),
        "A + I": (_At, ".0f"),
        "Â": (_Ahat, ".2f"),
    }[gcn_view.value]
    _chart = _heat(_mat, gcn_view.value, _fmt)
    _note = {
        "A": "Adjacency: 1 if i–j is an edge. Row sums = degrees.",
        "A + I": "Self-loops on the diagonal: the node keeps its own information. Degrees become d̃ = d + 1.",
        "Â": "Entry (i, j) = 1/√(d̃ᵢ d̃ⱼ). An edge between two hubs gets a small weight; a leaf's edge a large one.",
    }[gcn_view.value]
# the same  drawn on the graph: edge width = normalised weight
_edges = edges_df(toy_pairs)
_edges["w"] = [
    1 + 18 * float(_Ahat[u, v]) for u, v in zip(_edges.u, _edges.v)
]
_g = graph_chart(
    toy_pos,
    _edges,
    pd.DataFrame({"node": range(toy_n), "degree d̃": _d.int().tolist()}),
    alt.Color(
        "degree d̃:O",
        scale=alt.Scale(scheme="greens"),
        legend=alt.Legend(orient="bottom"),
    ),
    size=700,
    width=320,
    height=250,
    edge_color="#10b981",
    title="edge width = Âᵢⱼ",
)
slide(
    ctl(gcn_view, toy_edges_text),
    cols(
        H_chart(_chart, tidy),
        H_chart(_g, tidy)
        + f"<div class='card' style='margin-top:8px'><h3>{gcn_view.value}</h3><p style='font-size:15px'>{_note}</p></div>",
        t="1.25fr 1fr",
    ),
    kicker="03 · GCN · try it",
    title="Inside Â: who gets <span class='hl'>how much</span> say",
    chapter="03 GCN",
)

MY_GCN_SRC = '''class MyGCNConv(MessagePassing):
    """x_i' = sum_{j in N(i) ∪ {i}} 1/sqrt(d_i d_j) · W x_j + b"""

    def __init__(self, in_channels, out_channels):
        super().__init__(aggr="add")                       # ⊕ = sum
        self.lin = torch.nn.Linear(in_channels, out_channels, bias=False)
        self.bias = torch.nn.Parameter(torch.zeros(out_channels))

    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))  # A + I
        x = self.lin(x)                                    # W x  (transform first: cheaper)
        src, dst = edge_index
        deg = degree(dst, x.size(0), dtype=x.dtype)
        norm = deg.pow(-0.5)[src] * deg.pow(-0.5)[dst]     # 1/sqrt(d_i d_j) per edge
        return self.propagate(edge_index, x=x, norm=norm) + self.bias

    def message(self, x_j, norm):                          # x_j: source node of every edge
        return norm.view(-1, 1) * x_j
'''
# the displayed source is exactly the code that runs
_ns = dict(
    MessagePassing=MessagePassing,
    add_self_loops=add_self_loops,
    degree=degree,
    torch=torch,
)
exec(MY_GCN_SRC, _ns)
MyGCNConv = _ns["MyGCNConv"]

torch.manual_seed(0)
_ref = GCNConv(cora.num_features, 16)
_mine = MyGCNConv(cora.num_features, 16)
with torch.no_grad():
    _mine.lin.weight.copy_(_ref.lin.weight)
    _mine.bias.copy_(_ref.bias)
    _diff = (
        (_ref(cora.x, cora.edge_index) - _mine(cora.x, cora.edge_index))
        .abs()
        .max()
        .item()
    )
_src = MY_GCN_SRC
slide(
    cols(
        md(f"```python\n{_src}\n```"),
        f"""
        <div class='step'><i>1</i><div><p><code>aggr="add"</code> picks ⊕. Also <code>"mean"</code>, <code>"max"</code>.</p></div></div>
        <div class='step'><i>2</i><div><p><code>message()</code> describes <b>one edge</b>. Suffix <code>_j</code> = source,
        <code>_i</code> = target; PyG gathers the rows for you.</p></div></div>
        <div class='step'><i>3</i><div><p><code>propagate()</code> = message on every edge → scatter-aggregate per target → <code>update()</code>.</p></div></div>
        <div class='eq' style='margin-top:18px'><p style='margin:12px 0'>Check on <b>full Cora</b> vs PyG's <code>GCNConv</code>
        with the same weights:<br/><span class='stat' style='font-size:30px'>max |Δ| = {_diff:.1e}</span></p></div>
        <span class='badge'>✓ identical up to float rounding</span>
        """,
        t="1.75fr 1fr",
    ),
    kicker="03 · GCN · PyG code",
    title="Write your own layer with <code style='font-size:.8em'>MessagePassing</code>",
    chapter="03 GCN",
)

slide(
    """
    <div style='min-height:520px;display:flex;flex-direction:column;justify-content:center'>
      <div class='num'>04</div>
      <h1 style='margin-top:10px'>Learning <span class='hl'>& prediction</span></h1>
      <p class='lead' style='font-size:20px'>From embeddings to predictions, losses and honest evaluation.
      Then we train it, live, on 2,708 real papers.</p>
    </div>
    """,
    tone="dark",
)

_box = lambda x, y, w, t, s, c="#047857": (
    f"<rect x='{x}' y='{y}' width='{w}' height='58' rx='3' fill='#fff' stroke='{c}' stroke-width='2'/>"
    f"<text x='{x + w / 2}' y='{y + 25}' text-anchor='middle' font-family='Space Grotesk' font-weight='700' font-size='15' fill='#07160f'>{t}</text>"
    f"<text x='{x + w / 2}' y='{y + 44}' text-anchor='middle' font-family='JetBrains Mono' font-size='10.5' fill='#4d6a5c'>{s}</text>"
)
_arrow = lambda x1, y1, x2, y2: (
    f"<line x1='{x1}' y1='{y1}' x2='{x2}' y2='{y2}' stroke='#10b981' stroke-width='2.2' marker-end='url(#ah)'/>"
)
_svg = (
    "<svg viewBox='0 0 1000 200' width='100%'><defs><marker id='ah' viewBox='0 0 10 10' refX='9' refY='5' "
    "markerWidth='7' markerHeight='7' orient='auto'><path d='M0,0L10,5L0,10z' fill='#10b981'/></marker></defs>"
    + _box(0, 20, 130, "input graph", "A, X")
    + _arrow(132, 49, 160, 49)
    + _box(162, 20, 130, "GNN", "K layers, shared W")
    + _arrow(294, 49, 322, 49)
    + _box(324, 20, 140, "embeddings", "h_v ∈ ℝᵈ")
    + _arrow(466, 49, 494, 49)
    + _box(496, 20, 150, "prediction head", "node / edge / graph")
    + _arrow(648, 49, 676, 49)
    + _box(678, 20, 130, "predictions", "ŷ")
    + _arrow(810, 49, 838, 49)
    + _box(840, 20, 150, "loss  ·  metric", "CE / MSE · acc / AUC", "#f59e0b")
    + _box(678, 120, 130, "labels", "y (or from graph)", "#f59e0b")
    + _arrow(810, 149, 900, 82)
    + "<path d='M915,82 C915,190 230,190 227,82' stroke='#a3e635' stroke-width='2' fill='none' stroke-dasharray='5 4' marker-end='url(#ah)'/>"
    + "<text x='560' y='194' text-anchor='middle' font-family='JetBrains Mono' font-size='11' fill='#4d6a5c'>backprop ∂loss/∂W</text>"
    + "</svg>"
)
slide(
    _svg,
    cols(
        card(
            "Node head",
            "ŷᵥ = W⁽ᴴ⁾ hᵥ. One linear layer on top of the node's embedding, k outputs for k classes.",
        ),
        card(
            "Edge head",
            "ŷᵤᵥ = Linear(hᵤ ‖ hᵥ) or the dot product hᵤᵀhᵥ: 'does this edge exist?' (link prediction).",
            "#0ea5e9",
        ),
        card(
            "Graph head",
            "Pool all nodes: mean / max / sum → h_G. Simple but lossy on big graphs → hierarchical pooling (DiffPool).",
            "#8b5cf6",
        ),
        card(
            "Where labels come from",
            "Supervised: topic, fraud, toxicity. Self-supervised: hide an edge and predict it, or predict a node statistic.",
            "#f59e0b",
        ),
        t="repeat(4,1fr)",
        gap=14,
    ),
    kicker="04 · the training pipeline",
    title="Same pipeline for every task: only the <span class='hl'>head</span> changes",
    chapter="04 learning",
)

ego_node = mo.ui.number(0, 2707, value=1701, label="paper id")
ego_hops = mo.ui.slider(1, 3, value=2, label="hops", show_value=True)

def cora_ego(center, hops, *, max_nodes=90, extra=None, color=None, width=430, height=330, edge_w=None, size=170):
    """k-hop ego network of a Cora paper (closest nodes first, capped for readability)."""
    _dist = nx.single_source_shortest_path_length(G_cora, center, cutoff=hops)
    _keep = sorted(_dist, key=lambda n: (_dist[n], n))[:max_nodes]
    _sub = G_cora.subgraph(_keep)
    _pos = nx.spring_layout(_sub, seed=1, k=1.2 / np.sqrt(max(len(_keep), 1)))
    _nodes = pd.DataFrame(
        {
            "node": _keep,
            "topic": [CLASS_NAMES[cora_y[n]] for n in _keep],
            "split": [cora_split[n] for n in _keep],
            "hop": [_dist[n] for n in _keep],
        }
    )
    if extra is not None:
        _nodes = _nodes.assign(**{k: [f(n) for n in _keep] for k, f in extra.items()})
    _edges = edges_df(list(_sub.edges()))
    if edge_w is not None:
        _edges["w"] = [edge_w(u, v) for u, v in zip(_edges.u, _edges.v)]
    _chart = graph_chart(
        _pos,
        _edges,
        _nodes,
        color or alt.Color("topic:N", scale=CLASS_SCALE, legend=alt.Legend(orient="bottom", columns=4, title=None)),
        size=size,
        width=width,
        height=height,
        labels=False,
        stroke=alt.condition(f"datum.node == {center}", alt.value(4), alt.value(0)),
    )
    return _chart, _nodes, len(_dist)

_c = int(ego_node.value)
_chart, _nodes, _total = cora_ego(
    _c, ego_hops.value, width=520, height=400, size=210
)
_same = (_nodes[_nodes.hop == 1].topic == CLASS_NAMES[cora_y[_c]]).mean()
_bars = "".join(
    f"<div style='display:flex;align-items:center;gap:8px;font-size:12.5px;margin:3px 0'>"
    f"<span style='width:150px'>{n}</span><span style='height:10px;width:{(cora_y == i).sum() / 4.5:.0f}px;"
    f"background:{CLASS_COLORS[i]};display:inline-block'></span><span class='muted'>{(cora_y == i).sum()}</span></div>"
    for i, n in enumerate(CLASS_NAMES)
)
slide(
    cols(
        f"""
        <div class='cols' style='grid-template-columns:1fr 1fr;gap:12px'>
          <div class='stat'>{cora_stats["nodes"]:,}<small>papers (nodes)</small></div>
          <div class='stat'>{cora_stats["edges"]:,}<small>citations (edges)</small></div>
          <div class='stat'>{cora_stats["features"]:,}<small>word features · {cora_stats["avg_words"]:.0f} words/paper</small></div>
          <div class='stat'>{cora_stats["homophily"]:.0%}<small>edges link same topic</small></div>
        </div>
        <p style='margin:14px 0 6px'><b>Task:</b> predict the topic of a paper. Only
        <b>{cora_stats["train"]}</b> labelled (20 per class) · {cora_stats["val"]} val · {cora_stats["test"]:,} test.</p>
        {_bars}
        """,
        ctl(ego_node, ego_hops)
        + H_chart(
            _chart.properties(
                title=f"paper {_c} ({CLASS_NAMES[cora_y[_c]]}, {cora_split[_c]}): {_total} papers within {ego_hops.value} hops"
            ),
            tidy,
        )
        + f"<p class='muted' style='font-size:13.5px'>{_same:.0%} of its direct neighbours share its topic. "
        "This <b>homophily</b> is why message passing helps.</p>",
        t="1fr 1.15fr",
    ),
    kicker="04 · the data · try it",
    title="Cora: a citation graph with <span class='hl'>140 labels</span>",
    chapter="04 learning",
)