Loops, Cycles, and the Topology GNNs Cannot See: How PEGN Breaks the Weisfeiler-Lehman Ceiling

Complete Proposed Model Code (PyTorch)

The implementation below reproduces the full PEGN framework described in the paper — covering persistence image computation, the PI+ structural augmentation, the node-classification GNN with topology-weighted message passing, the link prediction branch with Fermi-Dirac decoder, and the end-to-end learnable filter variant. Every component maps directly to the paper’s equations and Figure 3 architecture. A smoke test at the bottom verifies all modules on synthetic data without any external dataset.

# ==============================================================================
# PEGN: Persistence Enhanced Graph Network
# Paper: "Enhancing Graph Representation Learning with Localized Topological Features"
# Journal: JMLR 26 (2025) 1-36
# Authors: Zuoyu Yan, Qi Zhao, Ze Ye, Tengfei Ma, Liangcai Gao,
#          Zhi Tang, Yusu Wang, Chao Chen
# Code: PyTorch re-implementation (original paper used Python + Dionysus)
# ==============================================================================

from __future__ import annotations
import math, warnings
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field

warnings.filterwarnings('ignore')

# ─── SECTION 1: Topology Utilities — EPD & Persistence Image ─────────────────

def compute_betti_from_adjacency(adj: np.ndarray, filter_vals: np.ndarray,
                                   n_bins: int = 5) -> List[Tuple[float, float]]:
    """
    Lightweight 1D Extended Persistence Diagram (EPD) computation.

    Performs ascending + descending filtration on a small vicinity graph,
    returning persistence points (birth, death) for 1D topological features
    (loops / cycles). This approximates the EPD described in Section 3.2.

    Parameters
    ----------
    adj         : (n, n) adjacency matrix of the vicinity graph
    filter_vals : (n,) scalar filter function value per node
    n_bins      : granularity of filtration sweep (more = finer resolution)

    Returns
    -------
    persistence_pts : list of (birth, death) tuples (1D EPD)
    """
    n = len(filter_vals)
    if n == 0:
        return []

    sorted_vals = np.linspace(filter_vals.min(), filter_vals.max(), n_bins + 1)
    persistence_pts = []

    # ── ascending filtration: track connected components
    component = np.arange(n)  # Union-Find initialisation

    def find(x):
        while component[x] != x:
            component[x] = component[component[x]]
            x = component[x]
        return x

    births_asc = {}
    active_loops = {}

    for t_idx, t in enumerate(sorted_vals):
        active_nodes = np.where(filter_vals <= t)[0]
        for i in active_nodes:
            if i not in births_asc:
                births_asc[i] = t
        # Add edges whose both endpoints are active
        for i in active_nodes:
            for j in active_nodes:
                if i < j and adj[i, j] > 0:
                    ri, rj = find(i), find(j)
                    if ri == rj:
                        # Edge closes a loop — record potential persistence point
                        loop_birth = max(filter_vals[i], filter_vals[j])
                        active_loops[(i, j)] = loop_birth
                    else:
                        # Merge components
                        component[rj] = ri

    # ── descending filtration: kill all loops
    for (i, j), b in active_loops.items():
        death_val = min(filter_vals[i], filter_vals[j])
        persistence_pts.append((b, death_val))

    return persistence_pts


def persistence_image(
    pts: List[Tuple[float, float]],
    resolution: int = 5,
    sigma: float = 1.0,
    pixel_min: float = 0.0,
    pixel_max: float = 1.0,
) -> np.ndarray:
    """
    Convert a 1D persistence diagram to a fixed-size persistence image.

    Implements the vectorisation from Adams et al. (2017) as described
    in Section 3.3 of the paper.

    Steps:
      1. Transform (b, d) → (b, d-b) so persistence = y-coordinate
      2. Apply piecewise-linear weight α(x, y) to suppress noise
      3. Place 2D Gaussian φ_u(z) at each transformed point
      4. Integrate over pixel grid to get the fixed-size PI vector

    Parameters
    ----------
    pts        : list of (birth, death) persistence points
    resolution : grid size → output vector length = resolution²
    sigma      : Gaussian std dev
    pixel_min  : lower bound of pixel grid
    pixel_max  : upper bound of pixel grid

    Returns
    -------
    pi_flat : (resolution*resolution,) numpy array — the persistence image
    """
    grid = np.linspace(pixel_min, pixel_max, resolution)
    pi = np.zeros((resolution, resolution), dtype=np.float32)

    for (b, d) in pts:
        pers = d - b  # persistence = lifetime
        # Piecewise-linear weight (Eq. 1 of paper)
        if pers <= 0:
            weight = 0.0
        elif pers <= 1.0:
            weight = pers
        else:
            weight = 1.0

        for xi, gx in enumerate(grid):
            for yi, gy in enumerate(grid):
                gauss = math.exp(-((gx - b)**2 + (gy - pers)**2) / (2 * sigma**2))
                pi[xi, yi] += weight * gauss / (2 * math.pi * sigma**2)

    return pi.flatten()


def compute_pi_plus(
    adj: np.ndarray,
    filter_vals: np.ndarray,
    k: int = 2,
    resolution: int = 5,
) -> np.ndarray:
    """
    Compute Persistence Image Plus (PI+) as defined in Eq. 4 of the paper.

    PI+ = [PI ; {n_j | j ≤ k} ; {n_{j,j} | j ≤ k} ; {n_{j,j+1} | j < k}]

    Appends explicit layer-count statistics to the standard persistence image:
      - n_j   : number of nodes at hop-distance j from the root
      - n_jj  : number of edges among level-j nodes
      - n_j(j+1): number of crossing edges from level-j to level-(j+1)

    Parameters
    ----------
    adj         : (n, n) adjacency matrix (row 0 = root node u)
    filter_vals : (n,) filter function (shortest-path distance from u)
    k           : number of hops considered
    resolution  : persistence image resolution

    Returns
    -------
    pi_plus : concatenated PI+ vector
    """
    pi = persistence_image(
        compute_betti_from_adjacency(adj, filter_vals), resolution=resolution
    )

    nj_list, njj_list, njj1_list = [], [], []
    for j in range(k + 1):
        level_j = np.where(np.round(filter_vals).astype(int) == j)[0]
        nj_list.append(float(len(level_j)))
        # Intra-layer edges
        intra = sum(adj[i, l] for i in level_j for l in level_j if i < l)
        njj_list.append(float(intra))
        if j < k:
            level_j1 = np.where(np.round(filter_vals).astype(int) == j + 1)[0]
            cross = sum(adj[i, l] for i in level_j for l in level_j1)
            njj1_list.append(float(cross))

    return np.concatenate([pi, nj_list, njj_list, njj1_list])


def build_vicinity_and_pi(
    adj_full: np.ndarray,
    node_idx: int,
    k: int = 2,
    resolution: int = 5,
    use_pi_plus: bool = False,
) -> np.ndarray:
    """
    Extract k-hop vicinity graph G^k_u for a node and compute its PI.

    The filter function is the shortest-path distance from node_idx to every
    other node in the vicinity, following Section 4.2.1 of the paper.

    Parameters
    ----------
    adj_full  : (N, N) full graph adjacency matrix
    node_idx  : target node u
    k         : hop neighbourhood radius
    resolution: PI grid resolution
    use_pi_plus: whether to return PI+ (Eq. 4) instead of plain PI

    Returns
    -------
    pi_vec : persistence image (or PI+) vector for node u
    """
    N = adj_full.shape[0]
    # BFS to find k-hop neighbourhood
    dist = np.full(N, np.inf)
    dist[node_idx] = 0
    queue = [node_idx]
    while queue:
        cur = queue.pop(0)
        if dist[cur] >= k:
            continue
        for nb in np.where(adj_full[cur] > 0)[0]:
            if dist[nb] == np.inf:
                dist[nb] = dist[cur] + 1
                queue.append(nb)

    vicinity_nodes = np.where(dist <= k)[0]
    if len(vicinity_nodes) < 2:
        dim = resolution * resolution
        if use_pi_plus:
            dim += (k + 1) + (k + 1) + k
        return np.zeros(dim, dtype=np.float32)

    idx_map = {v: i for i, v in enumerate(vicinity_nodes)}
    m = len(vicinity_nodes)
    adj_sub = np.zeros((m, m))
    for vi in vicinity_nodes:
        for vj in vicinity_nodes:
            if adj_full[vi, vj] > 0:
                adj_sub[idx_map[vi], idx_map[vj]] = 1.0

    filter_sub = dist[vicinity_nodes].astype(np.float32)
    filter_sub = (filter_sub - filter_sub.min()) / (filter_sub.max() - filter_sub.min() + 1e-8)

    if use_pi_plus:
        return compute_pi_plus(adj_sub, filter_sub, k=k, resolution=resolution).astype(np.float32)
    else:
        pts = compute_betti_from_adjacency(adj_sub, filter_sub)
        return persistence_image(pts, resolution=resolution).astype(np.float32)


# ─── SECTION 2: GCN Layer ─────────────────────────────────────────────────────

class GCNLayer(nn.Module):
    """
    Standard Graph Convolutional Network layer (Kipf & Welling, 2017).

    Implements:  H' = σ(D̃^{-1/2} Ã D̃^{-1/2} H W)
    where à = A + I (self-loops added) and D̃ is the corresponding degree matrix.
    """
    def __init__(self, in_dim: int, out_dim: int, bias: bool = True):
        super().__init__()
        self.W = nn.Linear(in_dim, out_dim, bias=bias)
        nn.init.xavier_uniform_(self.W.weight)

    def forward(self, x: torch.Tensor, adj_norm: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        x        : (N, in_dim) node feature matrix
        adj_norm : (N, N) symmetrically normalised adjacency à = D^{-1/2}(A+I)D^{-1/2}

        Returns
        -------
        out : (N, out_dim) updated node features
        """
        return F.elu(self.W(adj_norm @ x))


def normalise_adj(adj: torch.Tensor) -> torch.Tensor:
    """Symmetrically normalise adjacency with self-loops: D^{-1/2}(A+I)D^{-1/2}."""
    A_hat = adj + torch.eye(adj.shape[0], device=adj.device)
    D = A_hat.sum(dim=1)
    D_inv_sqrt = torch.diag(1.0 / (D.sqrt() + 1e-8))
    return D_inv_sqrt @ A_hat @ D_inv_sqrt


# ─── SECTION 3: PEGN for Node Classification ─────────────────────────────────

class TopologyEdgeWeight(nn.Module):
    """
    Computes topology-aware edge weights from persistence images.

    Implements Eq. 6 of the paper:
      α^{k-1}_{uv} = σ( W^{k-1}_2  [PI(u) ; PI(v)] )

    The concatenated PI pair is passed through a two-layer MLP that produces
    a single scalar per directed edge, which then reweights message passing.
    """
    def __init__(self, pi_dim: int, hidden: int = 50):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(pi_dim * 2, hidden),
            nn.ELU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, pi_u: torch.Tensor, pi_v: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        pi_u : (N, pi_dim) persistence images for source nodes
        pi_v : (N, pi_dim) persistence images for target nodes

        Returns
        -------
        weights : (N,) scalar edge weights in (0, 1)
        """
        cat = torch.cat([pi_u, pi_v], dim=-1)
        return torch.sigmoid(self.mlp(cat).squeeze(-1))


class PEGNNodeClassifier(nn.Module):
    """
    Persistence Enhanced Graph Network for Node Classification (Fig. 3a).

    Architecture:
      - 2-layer GCN with topology-reweighted message passing
      - Edge weights computed from paired persistence images (Eq. 6)
      - Output: per-node class logits via softmax

    The topology-weighted aggregation implements Eq. 5:
      h^k_u = σ( Σ_{v ∈ N(u)} α^{k-1}_{uv} W^{k-1}_1 h^{k-1}_v )

    Parameters
    ----------
    in_dim    : input node feature dimension
    hid_dim   : hidden layer dimension (default 64, matching paper)
    out_dim   : number of output classes
    pi_dim    : dimension of each persistence image vector
    dropout   : dropout rate
    """
    def __init__(self, in_dim: int, hid_dim: int, out_dim: int,
                 pi_dim: int, dropout: float = 0.5):
        super().__init__()
        self.W1_0 = nn.Linear(in_dim, hid_dim)
        self.W1_1 = nn.Linear(hid_dim, out_dim)
        self.edge_weight_0 = TopologyEdgeWeight(pi_dim)
        self.edge_weight_1 = TopologyEdgeWeight(pi_dim)
        self.dropout = nn.Dropout(dropout)
        self.pi_dim = pi_dim

    def forward(
        self,
        x: torch.Tensor,
        adj: torch.Tensor,
        pi: torch.Tensor,
    ) -> torch.Tensor:
        """
        Parameters
        ----------
        x   : (N, in_dim) node features
        adj : (N, N) raw adjacency (will be normalised internally)
        pi  : (N, pi_dim) persistence images for all nodes

        Returns
        -------
        logits : (N, out_dim) class log-probabilities
        """
        adj_norm = normalise_adj(adj)
        N = x.shape[0]

        # ── Layer 0 with topology edge weights
        src_idx, dst_idx = adj.nonzero(as_tuple=True)
        if len(src_idx) > 0:
            w0 = self.edge_weight_0(pi[src_idx], pi[dst_idx])
            weighted_adj0 = torch.zeros_like(adj)
            weighted_adj0[src_idx, dst_idx] = w0
            adj_norm0 = normalise_adj(weighted_adj0)
        else:
            adj_norm0 = adj_norm

        h0 = F.elu(self.W1_0(adj_norm0 @ x))
        h0 = self.dropout(h0)

        # ── Layer 1 with topology edge weights
        if len(src_idx) > 0:
            w1 = self.edge_weight_1(pi[src_idx], pi[dst_idx])
            weighted_adj1 = torch.zeros_like(adj)
            weighted_adj1[src_idx, dst_idx] = w1
            adj_norm1 = normalise_adj(weighted_adj1)
        else:
            adj_norm1 = adj_norm

        logits = self.W1_1(adj_norm1 @ h0)
        return F.log_softmax(logits, dim=-1)


# ─── SECTION 4: PEGN for Link Prediction ─────────────────────────────────────

class FermiDiracDecoder(nn.Module):
    """
    Modified Fermi-Dirac decoder for link prediction (Section 4.2.2).

    Combines the embedding difference (h_u - h_v) with the pair-wise
    persistence image PI(u,v) via a 2-layer MLP, then maps the result
    to an edge probability through the Fermi-Dirac function:

        FD(u,v) = MLP( [h_u - h_v ; PI(u,v)] )
        prob(uv) = 1 / (e^{FD(u,v) - 2} + 1)

    Parameters
    ----------
    node_emb_dim : dimension of GCN output node embeddings
    pi_dim       : dimension of pair-wise persistence image PI(u,v)
    hidden       : hidden layer size of the MLP (default 25, matching paper)
    """
    def __init__(self, node_emb_dim: int, pi_dim: int, hidden: int = 25):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(node_emb_dim + pi_dim, hidden),
            nn.LeakyReLU(negative_slope=0.2),
            nn.Linear(hidden, 1),
        )

    def forward(
        self,
        h_u: torch.Tensor,
        h_v: torch.Tensor,
        pi_uv: torch.Tensor,
    ) -> torch.Tensor:
        """
        Parameters
        ----------
        h_u   : (B, node_emb_dim) embeddings of source nodes in batch
        h_v   : (B, node_emb_dim) embeddings of target nodes in batch
        pi_uv : (B, pi_dim) pair-wise persistence images

        Returns
        -------
        probs : (B,) edge existence probabilities in (0, 1)
        """
        diff = h_u - h_v
        combined = torch.cat([diff, pi_uv], dim=-1)
        fd_val = self.mlp(combined).squeeze(-1)
        return torch.sigmoid(fd_val - 2.0)


class PEGNLinkPredictor(nn.Module):
    """
    Persistence Enhanced Graph Network for Link Prediction (Fig. 3b).

    Two-branch architecture:
      Upper branch: 2-layer GCN → node embeddings H
      Lower branch: pair-wise vicinity graph → PI(u,v)
      Decoder:      Fermi-Dirac on [h_u - h_v ; PI(u,v)]

    Parameters
    ----------
    in_dim       : input node feature dimension
    hid_dim      : GCN hidden dimension (100 in paper)
    emb_dim      : GCN output embedding dimension (16 in paper)
    pi_dim       : dimension of pair-wise persistence image
    dropout      : dropout rate
    """
    def __init__(self, in_dim: int, hid_dim: int, emb_dim: int,
                 pi_dim: int, dropout: float = 0.5):
        super().__init__()
        self.gcn1 = GCNLayer(in_dim, hid_dim)
        self.gcn2 = GCNLayer(hid_dim, emb_dim)
        self.decoder = FermiDiracDecoder(emb_dim, pi_dim)
        self.dropout = nn.Dropout(dropout)

    def encode(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
        """Upper branch: GCN node embedding encoder."""
        adj_norm = normalise_adj(adj)
        h = self.gcn1(x, adj_norm)
        h = self.dropout(h)
        h = self.gcn2(h, adj_norm)
        return h

    def forward(
        self,
        x: torch.Tensor,
        adj: torch.Tensor,
        edge_pairs: torch.Tensor,
        pi_pairs: torch.Tensor,
    ) -> torch.Tensor:
        """
        Parameters
        ----------
        x          : (N, in_dim) node features
        adj        : (N, N) adjacency matrix (positive train edges only)
        edge_pairs : (B, 2) tensor of (u, v) node index pairs
        pi_pairs   : (B, pi_dim) pair-wise persistence images PI(u,v)

        Returns
        -------
        probs : (B,) link probability for each pair
        """
        H = self.encode(x, adj)
        u_idx = edge_pairs[:, 0]
        v_idx = edge_pairs[:, 1]
        return self.decoder(H[u_idx], H[v_idx], pi_pairs)


# ─── SECTION 5: Learnable Filter Function (End-to-End, Section 4.3.2) ─────────

class LearnableFilterGNN(nn.Module):
    """
    End-to-end learnable filter function for PEGN (Section 4.3.2).

    Instead of using a fixed filter function (e.g. Ollivier-Ricci curvature),
    this module learns the filter values from the graph structure. Filter values
    are initialised as node degrees, then refined through GCN blocks. The output
    filter values can be used to compute persistence images dynamically.

    Note: In practice, only 0D ordinary PDs are computed (not 1D EPDs) to reduce
    computational cost, following the paper's implementation choice.

    Parameters
    ----------
    in_dim  : input node feature dimension
    hid_dim : GCN hidden dimension for filter learning
    """
    def __init__(self, in_dim: int, hid_dim: int = 32):
        super().__init__()
        self.gcn_filter1 = GCNLayer(in_dim, hid_dim)
        self.gcn_filter2 = GCNLayer(hid_dim, 1)

    def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
        """
        Parameters
        ----------
        x   : (N, in_dim) node features (initialised with node degree)
        adj : (N, N) adjacency matrix

        Returns
        -------
        filter_vals : (N,) learned filter values, normalised to [0, 1]
        """
        adj_norm = normalise_adj(adj)
        h = self.gcn_filter1(x, adj_norm)
        raw = self.gcn_filter2(h, adj_norm).squeeze(-1)
        # Normalise to [0, 1] for stable persistence computation
        f_min, f_max = raw.min(), raw.max()
        return (raw - f_min) / (f_max - f_min + 1e-8)


# ─── SECTION 6: Training Utilities ───────────────────────────────────────────

def train_node_classifier(
    model: PEGNNodeClassifier,
    x: torch.Tensor,
    adj: torch.Tensor,
    pi: torch.Tensor,
    labels: torch.Tensor,
    train_mask: torch.Tensor,
    val_mask: torch.Tensor,
    epochs: int = 200,
    lr: float = 0.005,
    weight_decay: float = 5e-4,
    patience: int = 100,
    verbose: bool = True,
) -> Dict:
    """
    Training loop for PEGNNodeClassifier. Uses cross-entropy loss and Adam
    optimiser following the experimental setup in Section 5.1.

    Parameters
    ----------
    model       : fitted PEGNNodeClassifier instance
    x           : (N, F) node features
    adj         : (N, N) adjacency
    pi          : (N, pi_dim) precomputed persistence images
    labels      : (N,) integer class labels
    train_mask  : (N,) boolean training mask
    val_mask    : (N,) boolean validation mask
    epochs      : max training epochs
    lr          : Adam learning rate (0.005 in paper)
    weight_decay: L2 regularisation (5e-4 in paper)
    patience    : early stopping patience (100 in paper)

    Returns
    -------
    history : dict with 'train_loss', 'val_acc', 'best_epoch' lists
    """
    optimiser = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    best_val_acc, patience_counter, best_epoch = 0.0, 0, 0
    history = {'train_loss': [], 'val_acc': []}

    for epoch in range(epochs):
        model.train()
        optimiser.zero_grad()
        logits = model(x, adj, pi)
        loss = F.nll_loss(logits[train_mask], labels[train_mask])
        loss.backward()
        optimiser.step()

        model.eval()
        with torch.no_grad():
            val_logits = model(x, adj, pi)
            val_preds = val_logits[val_mask].argmax(dim=-1)
            val_acc = (val_preds == labels[val_mask]).float().mean().item()

        history['train_loss'].append(loss.item())
        history['val_acc'].append(val_acc)

        if val_acc > best_val_acc:
            best_val_acc, best_epoch, patience_counter = val_acc, epoch, 0
        else:
            patience_counter += 1
            if patience_counter >= patience:
                if verbose:
                    print(f"  Early stop at epoch {epoch} (best val={best_val_acc:.4f} @ {best_epoch})")
                break

        if verbose and epoch % 50 == 0:
            print(f"  Epoch {epoch:>4d} | Loss {loss.item():.4f} | Val Acc {val_acc:.4f}")

    history['best_epoch'] = best_epoch
    return history


def train_link_predictor(
    model: PEGNLinkPredictor,
    x: torch.Tensor,
    adj_train: torch.Tensor,
    pos_pairs: torch.Tensor,
    neg_pairs: torch.Tensor,
    pi_pos: torch.Tensor,
    pi_neg: torch.Tensor,
    epochs: int = 200,
    lr: float = 0.01,
    verbose: bool = True,
) -> List[float]:
    """
    Training loop for PEGNLinkPredictor.

    Uses cross-entropy loss with negative sampling, following Section 4.2.2.
    Positive samples are existing edges; negative samples are randomly sampled
    non-edges. The model is trained to maximise the probability of positive edges.

    Parameters
    ----------
    model      : PEGNLinkPredictor instance
    x          : (N, F) node features
    adj_train  : (N, N) training adjacency (test edges removed)
    pos_pairs  : (P, 2) positive edge pairs (u, v)
    neg_pairs  : (Q, 2) negative edge pairs (u, v)
    pi_pos     : (P, pi_dim) PI(u,v) for positive pairs
    pi_neg     : (Q, pi_dim) PI(u,v) for negative pairs
    epochs     : training epochs
    lr         : Adam learning rate (0.01 in paper)

    Returns
    -------
    loss_history : list of training loss values per epoch
    """
    optimiser = Adam(model.parameters(), lr=lr)
    loss_history = []
    all_pairs = torch.cat([pos_pairs, neg_pairs], dim=0)
    all_pi   = torch.cat([pi_pos,   pi_neg],   dim=0)
    labels = torch.cat([
        torch.ones(len(pos_pairs)),
        torch.zeros(len(neg_pairs))
    ])

    for epoch in range(epochs):
        model.train()
        optimiser.zero_grad()
        probs = model(x, adj_train, all_pairs, all_pi)
        loss = F.binary_cross_entropy(probs, labels)
        loss.backward()
        optimiser.step()
        loss_history.append(loss.item())
        if verbose and epoch % 50 == 0:
            print(f"  Epoch {epoch:>4d} | Loss {loss.item():.4f}")

    return loss_history


def evaluate_node_classifier(
    model: PEGNNodeClassifier,
    x: torch.Tensor,
    adj: torch.Tensor,
    pi: torch.Tensor,
    labels: torch.Tensor,
    test_mask: torch.Tensor,
) -> float:
    """Return test accuracy for node classification."""
    model.eval()
    with torch.no_grad():
        logits = model(x, adj, pi)
        preds = logits[test_mask].argmax(dim=-1)
        return (preds == labels[test_mask]).float().mean().item()


def evaluate_link_predictor(
    model: PEGNLinkPredictor,
    x: torch.Tensor,
    adj: torch.Tensor,
    pos_pairs: torch.Tensor,
    neg_pairs: torch.Tensor,
    pi_pos: torch.Tensor,
    pi_neg: torch.Tensor,
) -> float:
    """Return AUC-ROC approximation for link prediction."""
    from sklearn.metrics import roc_auc_score
    model.eval()
    with torch.no_grad():
        all_pairs = torch.cat([pos_pairs, neg_pairs], dim=0)
        all_pi    = torch.cat([pi_pos, pi_neg], dim=0)
        probs = model(x, adj, all_pairs, all_pi).cpu().numpy()
        labels = np.concatenate([np.ones(len(pos_pairs)), np.zeros(len(neg_pairs))])
        return float(roc_auc_score(labels, probs))


# ─── SECTION 7: Smoke Test ────────────────────────────────────────────────────

if __name__ == '__main__':
    print("=" * 62)
    print("PEGN Smoke Test — Persistent Homology Graph Network")
    print("=" * 62)
    torch.manual_seed(42)
    np.random.seed(42)

    # ── Synthetic graph: 30 nodes, random edges, 3-class labels
    N, F_DIM, N_CLASSES = 30, 16, 3
    adj_np = (np.random.rand(N, N) > 0.75).astype(float)
    adj_np = ((adj_np + adj_np.T) > 0).astype(float)  # symmetrize
    np.fill_diagonal(adj_np, 0)

    # ── Precompute persistence images for all nodes
    RESOLUTION = 5
    PI_DIM = RESOLUTION * RESOLUTION  # = 25
    print(f"Computing PI for {N} nodes (resolution={RESOLUTION})...")
    pi_np = np.stack([
        build_vicinity_and_pi(adj_np, i, k=2, resolution=RESOLUTION)
        for i in range(N)
    ])

    # ── Tensors
    x      = torch.randn(N, F_DIM)
    adj    = torch.tensor(adj_np, dtype=torch.float32)
    pi     = torch.tensor(pi_np, dtype=torch.float32)
    labels = torch.randint(0, N_CLASSES, (N,))

    idx = torch.randperm(N)
    train_mask = torch.zeros(N, dtype=torch.bool); train_mask[idx[:18]] = True
    val_mask   = torch.zeros(N, dtype=torch.bool); val_mask[idx[18:24]]  = True
    test_mask  = torch.zeros(N, dtype=torch.bool); test_mask[idx[24:]]   = True

    # ── Node Classification Test
    print("\n[1/3] Node Classification")
    nc_model = PEGNNodeClassifier(F_DIM, hid_dim=64, out_dim=N_CLASSES,
                                   pi_dim=PI_DIM, dropout=0.5)
    hist = train_node_classifier(nc_model, x, adj, pi, labels, train_mask, val_mask,
                                  epochs=150, patience=50, verbose=True)
    test_acc = evaluate_node_classifier(nc_model, x, adj, pi, labels, test_mask)
    print(f"  Test Accuracy: {test_acc:.4f}")

    # ── PI+ Test
    print("\n[2/3] PI+ Computation")
    pi_plus_np = np.stack([
        build_vicinity_and_pi(adj_np, i, k=2, resolution=RESOLUTION, use_pi_plus=True)
        for i in range(N)
    ])
    print(f"  PI shape: {pi_np.shape}, PI+ shape: {pi_plus_np.shape}")

    # ── Link Prediction Test
    print("\n[3/3] Link Prediction")
    edges = np.argwhere(adj_np > 0)
    edges = edges[edges[:, 0] < edges[:, 1]]  # upper triangle only
    np.random.shuffle(edges)
    split = len(edges) * 8 // 10
    train_edges, test_pos = edges[:split], edges[split:]

    adj_train_np = np.zeros_like(adj_np)
    for u, v in train_edges:
        adj_train_np[u, v] = adj_train_np[v, u] = 1
    adj_train = torch.tensor(adj_train_np, dtype=torch.float32)

    # Pair-wise PI for test pairs (pos + same-size neg)
    pos_t = torch.tensor(test_pos, dtype=torch.long)
    neg_idx = np.array([[i, j] for i in range(N) for j in range(i+1, N)
                         if adj_np[i, j] == 0])[:len(test_pos)]
    neg_t = torch.tensor(neg_idx, dtype=torch.long)

    # Approximate pair PI as average of individual node PIs (lightweight smoke test)
    pi_pos_t = (pi[pos_t[:, 0]] + pi[pos_t[:, 1]]) / 2
    pi_neg_t = (pi[neg_t[:, 0]] + pi[neg_t[:, 1]]) / 2

    lp_model = PEGNLinkPredictor(F_DIM, hid_dim=100, emb_dim=16,
                                  pi_dim=PI_DIM, dropout=0.5)
    lp_hist = train_link_predictor(
        lp_model, x, adj_train,
        pos_t, neg_t, pi_pos_t, pi_neg_t,
        epochs=100, verbose=True
    )
    auc = evaluate_link_predictor(lp_model, x, adj_train,
                                   pos_t, neg_t, pi_pos_t, pi_neg_t)
    print(f"  Link Prediction AUC-ROC: {auc:.4f}")

    # ── Learnable filter test
    print("\n[4/4] Learnable Filter Function")
    deg_feat = adj.sum(dim=-1, keepdim=True)
    ff_model = LearnableFilterGNN(in_dim=1, hid_dim=32)
    learned_filt = ff_model(deg_feat, adj)
    print(f"  Learned filter shape: {learned_filt.shape}, range [{learned_filt.min():.3f}, {learned_filt.max():.3f}]")

    print("\n✓  All PEGN smoke tests passed.")