wgl/data/graph.py

"""Directed graph stored as out-edge and in-edge CSR structures.

Web graphs are fundamentally directed and the in/out asymmetry is the key signal for tasks like
spam detection (see ``docs/brainstorming/edge-direction.md``). We therefore keep two CSR views and
never symmetrize implicitly:

* **out-CSR** (``out_indptr``, ``out_indices``): successors of each node (who it links to).
* **in-CSR**  (``in_indptr``,  ``in_indices``):  predecessors of each node (who links to it).

The COO ``edge_index`` of shape ``(2, num_edges)`` (rows: ``src``, ``dst``) is the canonical edge
list. Arrays are plain NumPy so they save/load with ``np.savez`` and can be memory-mapped at scale.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from numpy.typing import NDArray


def _build_csr(
    rows: NDArray[np.int64],
    cols: NDArray[np.int64],
    num_nodes: int,
    index_dtype: np.dtype = np.int64,
) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
    """Build a CSR ``(indptr, indices)`` grouping ``cols`` by ``rows``.

    Args:
        rows: Source endpoints (the CSR row of each entry).
        cols: Destination endpoints (the value stored in each row).
        num_nodes: Total number of nodes (CSR has ``num_nodes + 1`` indptr entries).
        index_dtype: dtype of the returned ``indices`` (the node-id column). ``int32`` halves the
            footprint when node ids fit (``num_nodes < 2**31``) — worth it at 10^10-edge scale.
            ``indptr`` is always ``int64`` (it accumulates to ``num_edges``, which can exceed 2^31).

    Returns:
        A tuple ``(indptr, indices)``.
    """
    order = np.argsort(rows, kind="stable")
    sorted_rows = rows[order]
    sorted_cols = cols[order]
    counts = np.bincount(sorted_rows, minlength=num_nodes)
    indptr = np.zeros(num_nodes + 1, dtype=np.int64)
    np.cumsum(counts, out=indptr[1:])
    return indptr, sorted_cols.astype(index_dtype, copy=False)


@dataclass
class DirectedGraph:
    """An immutable directed graph with cached out/in CSR adjacency.

    Attributes:
        num_nodes: Number of nodes.
        edge_index: COO edges of shape ``(2, num_edges)`` (row 0 = src, row 1 = dst).
        out_indptr: Out-CSR row pointers, shape ``(num_nodes + 1,)``.
        out_indices: Out-CSR successors.
        in_indptr: In-CSR row pointers, shape ``(num_nodes + 1,)``.
        in_indices: In-CSR predecessors.
        node_names: Optional per-node string identifiers (e.g. reversed hostnames).
        features: Optional per-node feature matrix, shape ``(num_nodes, feature_dim)`` (e.g.
            ogbn-arxiv's word2vec vectors). Persisted as ``node_features.npy``; ``None`` for
            structure-only graphs (the CC web graph has no such features).
    """

    num_nodes: int
    edge_index: NDArray[np.int64]
    out_indptr: NDArray[np.int64]
    out_indices: NDArray[np.int64]
    in_indptr: NDArray[np.int64]
    in_indices: NDArray[np.int64]
    node_names: list[str] | None = None
    features: NDArray[np.float32] | None = None

    @property
    def num_edges(self) -> int:
        """Number of directed edges."""
        return int(self.edge_index.shape[1])

    @classmethod
    def from_edges(
        cls,
        edge_index: NDArray[np.int64],
        num_nodes: int | None = None,
        node_names: list[str] | None = None,
    ) -> DirectedGraph:
        """Construct a graph from a COO edge list, building both CSR views.

        Args:
            edge_index: Array of shape ``(2, num_edges)`` with src in row 0 and dst in row 1.
            num_nodes: Node count; inferred from the max id + 1 when omitted.
            node_names: Optional per-node identifiers.

        Returns:
            A new :class:`DirectedGraph`.
        """
        edge_index = np.ascontiguousarray(edge_index, dtype=np.int64)
        if edge_index.ndim != 2 or edge_index.shape[0] != 2:
            msg = f"edge_index must have shape (2, E), got {edge_index.shape}"
            raise ValueError(msg)
        src, dst = edge_index[0], edge_index[1]
        if num_nodes is None:
            num_nodes = int(edge_index.max()) + 1 if edge_index.size else 0
        out_indptr, out_indices = _build_csr(src, dst, num_nodes)
        in_indptr, in_indices = _build_csr(dst, src, num_nodes)
        return cls(
            num_nodes=num_nodes,
            edge_index=edge_index,
            out_indptr=out_indptr,
            out_indices=out_indices,
            in_indptr=in_indptr,
            in_indices=in_indices,
            node_names=node_names,
        )

    def out_neighbors(self, node: int) -> NDArray[np.int64]:
        """Return the successors of ``node`` (nodes it links to)."""
        return self.out_indices[self.out_indptr[node] : self.out_indptr[node + 1]]

    def in_neighbors(self, node: int) -> NDArray[np.int64]:
        """Return the predecessors of ``node`` (nodes that link to it)."""
        return self.in_indices[self.in_indptr[node] : self.in_indptr[node + 1]]

    def out_degree(self) -> NDArray[np.int64]:
        """Return the out-degree of every node."""
        return np.diff(self.out_indptr)

    def in_degree(self) -> NDArray[np.int64]:
        """Return the in-degree of every node."""
        return np.diff(self.in_indptr)

    def to_undirected(self) -> DirectedGraph:
        """Return a symmetrized copy with both ``(u, v)`` and ``(v, u)`` for every edge.

        Used only to build the *undirected baseline* for direction ablations — it deliberately
        discards the in/out asymmetry that the directed models exploit, so it should never be used
        for the primary directed experiments.

        Returns:
            A new :class:`DirectedGraph` over the same nodes with reciprocated, de-duplicated edges.
        """
        both = np.concatenate([self.edge_index, self.edge_index[::-1]], axis=1)
        undirected = np.unique(both, axis=1) if both.size else both
        return DirectedGraph.from_edges(
            undirected, num_nodes=self.num_nodes, node_names=self.node_names
        )

    def subgraph_from_edge_mask(self, mask: NDArray[np.bool_]) -> DirectedGraph:
        """Return a graph over the same node set keeping only the masked edges.

        Used to build the *message* graph for link prediction (a structural view that excludes
        held-out evaluation edges).

        Args:
            mask: Boolean mask over edges of length ``num_edges``.

        Returns:
            A new :class:`DirectedGraph` with the same ``num_nodes`` and node names.
        """
        return DirectedGraph.from_edges(
            self.edge_index[:, mask], num_nodes=self.num_nodes, node_names=self.node_names
        )

    def save(self, directory: str | Path) -> None:
        """Persist the graph to ``directory`` (``graph.npz`` + ``meta.json``).

        Args:
            directory: Output directory (created if needed).
        """
        directory = Path(directory)
        directory.mkdir(parents=True, exist_ok=True)
        np.savez(
            directory / "graph.npz",
            edge_index=self.edge_index,
            out_indptr=self.out_indptr,
            out_indices=self.out_indices,
            in_indptr=self.in_indptr,
            in_indices=self.in_indices,
        )
        meta = {"num_nodes": self.num_nodes, "num_edges": self.num_edges}
        (directory / "meta.json").write_text(json.dumps(meta, indent=2))
        if self.node_names is not None:
            (directory / "node_names.txt").write_text("\n".join(self.node_names))
        if self.features is not None:
            np.save(directory / "node_features.npy", self.features)

    @classmethod
    def load(cls, directory: str | Path, mmap: bool = False) -> DirectedGraph:
        """Load a graph previously written by :meth:`save`.

        Args:
            directory: Directory containing ``graph.npz`` / ``meta.json``.
            mmap: If True, memory-map the arrays (recommended for large graphs).

        Returns:
            The loaded :class:`DirectedGraph`.
        """
        directory = Path(directory)
        data = np.load(directory / "graph.npz", mmap_mode="r" if mmap else None)
        meta = json.loads((directory / "meta.json").read_text())
        names_path = directory / "node_names.txt"
        node_names = names_path.read_text().splitlines() if names_path.exists() else None
        features_path = directory / "node_features.npy"
        features = np.load(features_path) if features_path.exists() else None
        return cls(
            num_nodes=int(meta["num_nodes"]),
            edge_index=data["edge_index"],
            out_indptr=data["out_indptr"],
            out_indices=data["out_indices"],
            in_indptr=data["in_indptr"],
            in_indices=data["in_indices"],
            node_names=node_names,
            features=features,
        )