Graph Neural Network

"""One round of message passing - gather, aggregate, update.

The aggregator MUST be permutation-invariant: a node's neighbours have no
canonical order, so anything order-sensitive is silently wrong.
"""


def message_passing(features, adjacency, weight, aggregator="mean"):
    """features[n][d], adjacency[n] -> list of neighbour indices."""
    n = len(features)
    d_out = len(weight[0])
    updated = []

    for node in range(n):
        # Self-loop: without it a node loses its own signal entirely after
        # the first round.
        neighbours = adjacency[node] + [node]

        # Aggregate. Sum keeps degree information; mean discards it and is
        # more stable when degrees vary wildly.
        pooled = [0.0] * len(features[0])
        for neighbour in neighbours:
            for j in range(len(pooled)):
                pooled[j] += features[neighbour][j]

        if aggregator == "mean":
            pooled = [p / len(neighbours) for p in pooled]

        # Update: a shared linear transform, then a nonlinearity.
        out = [0.0] * d_out
        for j in range(d_out):
            total = sum(pooled[i] * weight[i][j] for i in range(len(pooled)))
            out[j] = max(0.0, total)     # ReLU
        updated.append(out)

    return updated


def forward(features, adjacency, weights):
    """Stack rounds. Two or three - past that, over-smoothing collapses every
    node toward the same representation."""
    h = features
    for weight in weights:
        h = message_passing(h, adjacency, weight)
    return h

"""GCN - typed, sparse, with the normalization that makes it work."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class GnnConfig:
    in_features: int
    hidden: int = 64
    out_features: int = 2
    # Two layers. Not a shrug - over-smoothing means three is usually worse
    # and four is usually much worse.
    layers: int = 2
    dropout: float = 0.5


def normalize_adjacency(edge_index: Tensor, num_nodes: int) -> Tensor:
    """Build the symmetric normalized adjacency with self-loops.

    D^-1/2 (A + I) D^-1/2. The normalization is load-bearing: without it a
    hub with 10,000 neighbours produces activations orders of magnitude
    larger than a node with three, and the network spends capacity fighting
    degree instead of learning structure.
    """
    if edge_index.dim() != 2 or edge_index.size(0) != 2:
        raise ValueError(f"edge_index must be (2, num_edges), got {tuple(edge_index.shape)}")

    self_loops = torch.arange(num_nodes, device=edge_index.device)
    row = torch.cat([edge_index[0], self_loops])
    col = torch.cat([edge_index[1], self_loops])

    degree = torch.zeros(num_nodes, device=edge_index.device)
    degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float))
    inv_sqrt = degree.pow(-0.5)
    # An isolated node has degree 0, and 0^-0.5 is inf. Left unhandled this
    # produces NaN that spreads through the whole graph on the first matmul.
    inv_sqrt[torch.isinf(inv_sqrt)] = 0.0

    values = inv_sqrt[row] * inv_sqrt[col]
    return torch.sparse_coo_tensor(
        torch.stack([row, col]), values, (num_nodes, num_nodes)
    ).coalesce()


class GcnLayer(nn.Module):
    def __init__(self, in_features: int, out_features: int) -> None:
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=True)

    def forward(self, x: Tensor, adjacency: Tensor) -> Tensor:
        # Transform FIRST, then propagate. Doing it in this order costs
        # O(n·d_in·d_out + |E|·d_out); the reverse costs O(|E|·d_in + ...),
        # which is worse whenever d_out < d_in, i.e. almost always.
        return torch.sparse.mm(adjacency, self.linear(x))


class Gcn(nn.Module):
    def __init__(self, config: GnnConfig) -> None:
        super().__init__()
        if config.layers < 1:
            raise ValueError(f"need at least one layer, got {config.layers}")

        widths = [config.in_features] + [config.hidden] * (config.layers - 1)
        self.layers = nn.ModuleList(
            GcnLayer(widths[i], widths[i + 1] if i + 1 < len(widths) else config.hidden)
            for i in range(config.layers - 1)
        )
        self.output = GcnLayer(config.hidden if config.layers > 1 else config.in_features,
                               config.out_features)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor, adjacency: Tensor) -> Tensor:
        for layer in self.layers:
            x = self.dropout(torch.relu(layer(x, adjacency)))
        return self.output(x, adjacency)

"""GNN at scale - neighbour sampling, CSR adjacency, fused scatter."""

import numpy as np
import torch
from numpy.typing import NDArray
from torch import Tensor


def build_csr(edge_index: NDArray[np.int64], num_nodes: int) -> tuple[NDArray, NDArray]:
    """Compressed sparse row adjacency: (indptr, indices).

    A COO edge list means a node's neighbours are scattered across memory.
    CSR puts them in one contiguous run, so sampling a neighbourhood is a
    single slice rather than a filtered scan of every edge.
    """
    order = np.argsort(edge_index[0], kind="stable")
    sorted_src = edge_index[0][order]
    indices = edge_index[1][order].astype(np.int64)

    counts = np.bincount(sorted_src, minlength=num_nodes)
    indptr = np.zeros(num_nodes + 1, dtype=np.int64)
    np.cumsum(counts, out=indptr[1:])

    return indptr, indices


def sample_neighbours(
    indptr: NDArray[np.int64],
    indices: NDArray[np.int64],
    seeds: NDArray[np.int64],
    fanout: int,
    rng: np.random.Generator,
) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
    """Sample up to `fanout` neighbours per seed node.

    This is not an approximation for speed - it is what makes training
    possible at all. Three hops on a degree-50 graph reaches >100k nodes per
    example, so the exact computation graph does not fit in memory.
    """
    src_out = []
    dst_out = []

    for seed in seeds:
        start, end = indptr[seed], indptr[seed + 1]
        degree = end - start

        if degree <= fanout:
            chosen = indices[start:end]
        else:
            # Sample WITHOUT replacement so a neighbour cannot be double-counted
            # into the aggregate, which would silently reweight it.
            offsets = rng.choice(degree, size=fanout, replace=False)
            chosen = indices[start + offsets]

        src_out.append(np.full(len(chosen), seed, dtype=np.int64))
        dst_out.append(chosen)

    return np.concatenate(src_out), np.concatenate(dst_out)


def aggregate_mean(features: Tensor, src: Tensor, dst: Tensor, num_nodes: int) -> Tensor:
    """Mean-aggregate neighbour features with a fused scatter-add.

    index_add_ accumulates in ONE pass over the edge list, in place. Building
    per-node neighbour lists and looping would allocate per node and destroy
    the memory locality CSR just bought.
    """
    output = torch.zeros(num_nodes, features.size(1), device=features.device)
    output.index_add_(0, src, features[dst])

    counts = torch.zeros(num_nodes, device=features.device)
    counts.index_add_(0, src, torch.ones_like(src, dtype=torch.float))

    # clamp_min, not a mask: isolated nodes divide by 1 and keep their zeros,
    # which is correct and avoids a branch.
    return output / counts.clamp_min(1.0).unsqueeze(1)