def toy_molecules(n, atoms=6, seed=SEED):
    g = torch.Generator().manual_seed(seed)
    pos = torch.randn(n, atoms, 3, generator=g) * 1.2
    d = torch.cdist(pos, pos)
    iu = torch.triu_indices(atoms, atoms, 1)
    energy = ((1 - torch.exp(-(d[:, iu[0], iu[1]] - 1.5))) ** 2).sum(1)
    return pos, energy

def random_rigid_motion(pos, seed):
    Q = torch.tensor(
        special_ortho_group.rvs(3, random_state=seed), dtype=torch.float
    )
    b = torch.randn(3, generator=torch.Generator().manual_seed(seed))
    return pos @ Q.T + b, Q

class SchNetLite(nn.Module):
    """Continuous-filter convolution on pairwise distances only → invariant by construction."""

    def __init__(self, d=32, n_rbf=16):
        super().__init__()
        self.centers = nn.Parameter(
            torch.linspace(0, 6, n_rbf), requires_grad=False
        )
        self.filter = nn.Sequential(
            nn.Linear(n_rbf, d), nn.SiLU(), nn.Linear(d, d)
        )
        self.embed = nn.Parameter(torch.randn(d) * 0.1)
        self.out = nn.Sequential(
            nn.Linear(d, d), nn.SiLU(), nn.Linear(d, 1)
        )

    def forward(self, pos):
        d = torch.cdist(
            pos, pos
        )  # [B, n, n] — the ONLY use of coordinates
        W = self.filter(
            torch.exp(-((d[..., None] - self.centers) ** 2) / 0.5)
        )  # filter W(||r_i - r_j||)
        mask = 1 - torch.eye(pos.size(1))[None, :, :, None]
        h = (W * mask * self.embed).sum(2)  # h_i = Σ_j W(r_ij) ⊙ h_j
        return self.out(h).sum((1, 2))

class CoordMLP(nn.Module):
    def __init__(self, atoms=6, d=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(atoms * 3, d),
            nn.SiLU(),
            nn.Linear(d, d),
            nn.SiLU(),
            nn.Linear(d, 1),
        )

    def forward(self, pos):
        return self.net(pos.flatten(1)).squeeze(1)

def fit_energy(model, pos, energy, epochs=400):
    torch.manual_seed(SEED)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for _ in range(epochs):
        for b in torch.randperm(len(pos)).split(256):
            opt.zero_grad()
            F.mse_loss(model(pos[b]), energy[b]).backward()
            opt.step()
    return model

geo_run = mo.ui.run_button(
    label="Train SchNet-lite and a coordinate MLP (≈ 30 s)"
)
geo_run

mo.stop(not geo_run.value, mo.md("Press the button."))
_tr_pos, _tr_E = toy_molecules(3000, seed=0)
_te_pos, _te_E = toy_molecules(500, seed=1)
_te_rot, _ = random_rigid_motion(_te_pos, seed=7)
_rows = []
for _name, _m in [
    ("SchNet-lite (distances only)", SchNetLite()),
    ("MLP on raw coordinates", CoordMLP()),
]:
    _t = time.time()
    _m = fit_energy(_m, _tr_pos, _tr_E)
    with torch.no_grad():
        _p, _p_rot = _m(_te_pos), _m(_te_rot)
    _rows.append(
        {
            "model": _name,
            "test MAE": (_p - _te_E).abs().mean().item(),
            "test MAE, rotated + shifted molecules": (_p_rot - _te_E)
            .abs()
            .mean()
            .item(),
            "mean |f(Qr+b) − f(r)|": (_p_rot - _p).abs().mean().item(),
            "seconds": time.time() - _t,
        }
    )

# PaiNN-style equivariant vector feature: v_i = Σ_j a(r_ij) · r_ij / |r_ij| with invariant scalar weights a
_a = nn.Sequential(nn.Linear(1, 16), nn.SiLU(), nn.Linear(16, 1))

def painn_vectors(pos):
    rel = pos[:, :, None, :] - pos[:, None, :, :]
    dist = rel.norm(dim=-1, keepdim=True).clamp(min=1e-9)
    w = _a(dist) * (1 - torch.eye(pos.size(1)))[None, :, :, None]
    return (w * rel / dist).sum(2)

_x = _te_pos[:50]
_xr, _Q = random_rigid_motion(_x, seed=3)
with torch.no_grad():
    _equiv_gap = (
        (painn_vectors(_xr) - painn_vectors(_x) @ _Q.T).abs().max().item()
    )
mo.vstack(
    [
        pd.DataFrame(_rows).round(4),
        mo.md(
            f"PaiNN-style vector feature: max $|\\mathbf v(Q\\mathbf r+\\mathbf b) - Q\\,\\mathbf v(\\mathbf r)|$ = **{_equiv_gap:.1e}** (equivariant)."
        ),
    ]
)

_rev = reveal(
    r"""
Measured (500 test molecules; energies have mean 7.9 and std 2.2): **SchNet-lite MAE 0.088**, unchanged on rotated and shifted molecules,
and its predictions move by exactly 0 under rigid motion. The **coordinate MLP** has MAE 2.33, *worse* than always predicting the mean
(1.62): with 3,000 examples it cannot learn that 18 raw numbers describe the same molecule in infinitely many orientations, and a rotation
changes its prediction by 2.9 on average. The PaiNN-style vector feature rotates with the input to float precision ($2\times10^{-7}$).
**Build the symmetry in; don't hope to learn it.**

""",
    predict="On rotated and shifted test molecules, will a coordinate MLP even beat predicting the mean energy? And SchNet-lite?",
)
_post = mo.md(r"""
## 2 · Network motifs

A **motif** is a small pattern that occurs much more often than in a random network with the same degrees. Its significance is
$Z=(N^{\text{real}}-\bar N^{\text{rand}})/\sigma^{\text{rand}}$; the vector of Z-scores over all 3-node patterns is the **significance profile**.
The notes' retweet example: edges $x\to y$ = "$y$ retweets $x$": a→b, b→c, a→c, c→d, a→d, d→e. Feed-forward loops (FFL: $x\to y\to z$ and
$x\to z$) are triad type **030T** in the triad census.
""")
mo.vstack([_rev, _post])

TRIADS = [
    "021D",
    "021U",
    "021C",
    "111D",
    "111U",
    "030T",
    "030C",
    "201",
    "120D",
    "120U",
    "120C",
    "210",
    "300",
]

def significance_profile(g, n_null=30, seed=SEED):
    """Z-score of each connected triad type vs. degree-preserving rewired networks (directed edge swaps)."""
    real = nx.triadic_census(g)
    null = []
    for k in range(n_null):
        h = g.copy()
        nx.directed_edge_swap(
            h,
            nswap=3 * h.number_of_edges(),
            max_tries=10**6,
            seed=seed + k,
        )
        null.append(nx.triadic_census(h))
    z = {}
    for t in TRIADS:
        vals = np.array([c[t] for c in null], dtype=float)
        z[t] = (
            (real[t] - vals.mean()) / (vals.std() + 1e-9)
            if vals.std() > 0
            else 0.0
        )
    return pd.Series(z), real

def synthetic_retweets(coordinated, n=150, seed=SEED):
    """Organic: each account retweets a few accounts, preferring popular ones. Coordinated: plus a 12-account group whose
    members all retweet a common source and each other (an amplification ring)."""
    r = np.random.default_rng(seed)
    g = nx.DiGraph()
    g.add_nodes_from(range(n))
    popularity = r.pareto(1.5, n) + 1
    for v in range(n):
        for u in r.choice(
            n,
            r.integers(1, 4),
            replace=False,
            p=popularity / popularity.sum(),
        ):
            if u != v:
                g.add_edge(int(u), v)  # edge u → v: v retweets u
    if coordinated:
        group = r.choice(n, 12, replace=False)
        source = int(group[0])
        for v in group[1:]:
            g.add_edge(source, int(v))
            for w in r.choice(group[1:], 3, replace=False):
                if w != v:
                    g.add_edge(int(v), int(w))
    return g

motif_run = mo.ui.run_button(
    label="Compute significance profiles (≈ 20 s)"
)
motif_run

def as_data(g):
    d = from_networkx(nx.convert_node_labels_to_integers(g))
    d.x = torch.ones(d.num_nodes, 1)
    return d

def subgraph_pairs(n_pairs=600, seed=SEED):
    """Half positives (a connected induced subgraph of the target), half negatives (a random small graph that is NOT one)."""
    r = np.random.default_rng(seed)
    pairs = []
    while len(pairs) < n_pairs:
        t = nx.gnp_random_graph(
            int(r.integers(10, 16)), 0.25, seed=int(r.integers(1e9))
        )
        if not nx.is_connected(t):
            continue
        k = int(r.integers(4, 7))
        if len(pairs) % 2 == 0:  # positive: BFS-grown induced subgraph
            nodes = list(
                nx.bfs_tree(t, int(r.integers(t.number_of_nodes())))
            )[:k]
            q, y = t.subgraph(nodes).copy(), 1
        else:  # negative candidate
            q = nx.gnp_random_graph(k, 0.6, seed=int(r.integers(1e9)))
            if (
                not nx.is_connected(q)
                or nx.algorithms.isomorphism.GraphMatcher(
                    t, q
                ).subgraph_is_isomorphic()
            ):
                continue
            y = 0
        pairs.append((as_data(q), as_data(t), y))
    return pairs

class OrderEncoder(nn.Module):
    def __init__(self, d=32, n_layers=3):
        super().__init__()
        self.convs = nn.ModuleList(
            GINConv(
                nn.Sequential(
                    nn.Linear(1 if i == 0 else d, d),
                    nn.ReLU(),
                    nn.Linear(d, d),
                )
            )
            for i in range(n_layers)
        )
        self.out = nn.Linear(d * n_layers, d)

    def forward(self, b):
        h, pooled = b.x, []
        for conv in self.convs:
            h = F.relu(conv(h, b.edge_index))
            pooled.append(global_add_pool(h, b.batch))
        # non-negative embedding (softplus rather than ReLU: no dead output units)
        return F.softplus(self.out(torch.cat(pooled, 1)))

def violation(zq, zt):
    return F.relu(zq - zt).pow(2).sum(1)

subgraph_run = mo.ui.run_button(
    label="Build 600 query/target pairs and train the matcher (≈ 30 s)"
)
subgraph_run

_rev = reveal(
    r"""
Measured on 150 held-out (query, target) pairs: the order-violation score separates subgraphs from non-subgraphs with **ROC AUC ≈ 0.94**
(accuracy 0.86 at a fixed threshold), after a few seconds of training and with no subgraph search at test time. Training was sensitive to
optimisation details: with ReLU outputs, a learning rate of 3e-3 and a large margin, some seeds collapsed to "every embedding = 0"
(AUC 0.5). A softplus output, lr 1e-3, margin 2 and gradient clipping were stable across seeds (0.93–0.95). SPMiner then grows frequent
motifs greedily in this embedding space.

""",
    predict="How well will the order-embedding violation score separate subgraph from non-subgraph pairs (ROC AUC)?",
)
_post = mo.md(r"""
## 4 · Explainability: GNNExplainer and fidelity

**GNNExplainer** learns a soft mask over edges so that the masked graph keeps the model's prediction while the mask stays small (a
mutual-information objective made differentiable). Benchmark **BA-Shapes**: a Barabási–Albert graph with 80 attached "house" motifs;
house nodes get labels by their position in the house (classes 1–3), base nodes class 0. Node features are constant, so a correct model
*must* use the house structure, and the ground-truth explanation is known.

**Fidelity:** $\text{fid}^+$ = how often the prediction *changes* when the explanation is removed (necessity); $\text{fid}^-$ = how often it
changes when *only* the explanation is kept (sufficiency, lower is better). The **characterisation score** combines them like an F1.
""")
mo.vstack([_rev, _post])

mo.stop(not explain_run.value, mo.md("Press the button."))
ba_data = ba_shapes()
torch.manual_seed(SEED)
ba_model = NodeGCN()
_opt = torch.optim.Adam(ba_model.parameters(), lr=0.01, weight_decay=5e-4)
for _ in range(400):
    ba_model.train()
    _opt.zero_grad()
    F.cross_entropy(
        ba_model(ba_data.x, ba_data.edge_index)[ba_data.train_mask],
        ba_data.y[ba_data.train_mask],
    ).backward()
    _opt.step()
ba_model.eval()
with torch.no_grad():
    _pred = ba_model(ba_data.x, ba_data.edge_index).argmax(1)
_test_acc = (
    (_pred[ba_data.test_mask] == ba_data.y[ba_data.test_mask])
    .float()
    .mean()
    .item()
)

explainer = Explainer(
    model=ba_model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type="model",
    node_mask_type=None,
    edge_mask_type="object",
    model_config=dict(
        mode="multiclass_classification",
        task_level="node",
        return_type="raw",
    ),
)
_house_nodes = torch.nonzero(
    (ba_data.y > 0) & (_pred == ba_data.y)
).squeeze()[:30]
_rows = []
for _v in _house_nodes.tolist():
    torch.manual_seed(SEED)
    _exp = explainer(ba_data.x, ba_data.edge_index, index=_v)
    _m = _exp.edge_mask.detach()
    _local = _m > 0  # edges in the node's receptive field
    _auc = (
        roc_auc_score(
            ba_data.edge_mask[_local].numpy(), _m[_local].numpy()
        )
        if ba_data.edge_mask[_local].unique().numel() > 1
        else np.nan
    )
    _fp, _fm = fidelity(explainer, _exp)
    _rows.append(
        {
            "node": _v,
            "explanation AUC vs ground-truth house edges": _auc,
            "fid+": _fp,
            "fid−": _fm,
        }
    )
explain_results = pd.DataFrame(_rows)
_fp_mean, _fm_mean = (
    explain_results["fid+"].mean(),
    explain_results["fid−"].mean(),
)
mo.vstack(
    [
        mo.md(
            f"GCN test accuracy on BA-Shapes: **{_test_acc:.3f}**. Explained {len(explain_results)} correctly classified house nodes:"
        ),
        pd.DataFrame(
            [
                {
                    "mean explanation AUC": explain_results.iloc[
                        :, 1
                    ].mean(),
                    "mean fid+": _fp_mean,
                    "mean fid−": _fm_mean,
                    "characterisation score": float(
                        characterization_score(
                            torch.tensor(_fp_mean), torch.tensor(_fm_mean)
                        )
                    ),
                }
            ]
        ).round(3),
    ]
)

_rev = reveal(
    r"""
Measured (one seed): the GCN classifies BA-Shapes nodes with **0.84** test accuracy using only structure. GNNExplainer's edge masks
recover the ground-truth house edges with a **median AUC of 0.82** (mean 0.65): many explanations are near-perfect, but some score low,
because GNNExplainer is a per-instance optimisation and can settle on a different subgraph that also preserves the prediction.
**fid+ = 1.0**: removing the explanation always changes the prediction (it is *necessary*). **fid− ≈ 0.47**: keeping *only* the explanation
preserves the prediction about half the time, so it is often not *sufficient* on its own. Characterisation score ≈ **0.70**. Report both
fidelities, not a single number, and compare with ground truth when you have it (you rarely do).

""",
    predict="Will GNNExplainer's explanation be *necessary* (fid+) and *sufficient* (fid−) for the GCN's prediction? Guess both.",
)
_post = mo.md(r"""
---

## 5 · Exercises (the notes' practice problems)

**Exercise 1 (order embeddings).** $\mathbf z_1=(1,2)$, $\mathbf z_2=(2,3)$, $\mathbf z_3=(3,1)$. Which ordered pairs $(i,j)$ are predicted as
"$G_i\subseteq G_j$"? → `ex1_pairs` (a set of tuples like `{(1, 2)}`).

**Exercise 2 (fidelity from a table).** Four instances, all target class 1. Full-graph predictions 1, 1, 1, 0; without the explanation 0, 1, 0, 0;
with only the explanation 1, 1, 0, 0. Give `(fid_plus, fid_minus, char_score)` with $w_\pm=1$ → `ex2`.

**Exercise 3 (code).** Write `count_ffl(g)`: the number of feed-forward loops $x\to y,\ y\to z,\ x\to z$ in a `nx.DiGraph`, **without**
`triadic_census`. It is checked against triad 030T on random graphs. (Careful: 030T has no reciprocated edges.)
""")
mo.vstack([_rev, _post])

ex1_pairs = None
ex2 = None

def count_ffl(g):
    return None

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

def _check_ffl():
    try:
        if count_ffl(retweet_toy) is None:
            return _verdict(False, False)
        ok = count_ffl(retweet_toy) == 2
        for s in range(5):
            g = nx.gnp_random_graph(30, 0.12, seed=s, directed=True)
            ok &= count_ffl(g) == nx.triadic_census(g)["030T"]
        return _verdict(ok, True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(ex1_pairs == {(1, 2)}, ex1_pairs is not None),
        mo.md("**Exercise 2**"),
        _verdict(
            ex2 is not None and np.allclose(ex2, (0.5, 0.25, 0.6)),
            ex2 is not None,
        ),
        mo.md("**Exercise 3**"),
        _check_ffl(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"Only $\mathbf z_1\le\mathbf z_2$ coordinate-wise → **{(1, 2)}**. $z_3$ is incomparable with both. Note $\min(z_2,z_3)=(2,1)$ lies below both: the model's 'common subgraph'."
                ),
                "Solution 2": mo.md(
                    r"Correct on full graph: 1,1,1,0. Without $S$: 0,1,0,0 → differences 1,0,1,0 → $\text{fid}^+=0.5$. Only $S$: 1,1,0,0 → differences 0,0,1,0 → $\text{fid}^-=0.25$. $c=\frac{2\cdot0.5\cdot0.75}{0.5+0.75}=0.6$."
                ),
                "Solution 3": mo.md("""```python
def count_ffl(g):
    n = 0
    for x, z in g.edges:
        if g.has_edge(z, x):
            continue
        for y in set(g.successors(x)) & set(g.predecessors(z)):
            if y not in (x, z) and not g.has_edge(y, x) and not g.has_edge(z, y):
                n += 1                       # x→y→z plus shortcut x→z, no reciprocal edges
    return n
```
Each FFL is counted once, from its unique shortcut edge x→z."""),
            }
        ),
    ]
)