DeepWalk & Node2Vec

"""DeepWalk / node2vec, transcribed the way the papers read.

Three stages, and keeping them separate is the whole point:

  1. bias the transition probabilities with p and q  (second order: the
     probability depends on where the walk CAME FROM, not just where it is)
  2. sample walks, which are the "sentences"
  3. run skip-gram with negative sampling on those sentences, unchanged

Linear-scan sampling and per-pair Python updates. Both are replaced later -
and note now that this profiles as a SAMPLING problem, not an arithmetic one.
"""

import math
import random

SEED = 1234


def biased_transition(graph, previous, current, return_p, inout_q):
    """Unnormalized node2vec weights over the neighbours of `current`.

    d(previous, candidate) is 0, 1 or 2 and nothing else, because a candidate
    is a neighbour of current and current is a neighbour of previous:

      d = 0  the candidate IS previous    -> weight 1/p  (backtracking)
      d = 1  the candidate neighbours it  -> weight 1     (staying local)
      d = 2  neither                      -> weight 1/q  (heading outward)
    """
    previous_neighbours = set(graph[previous]) if previous is not None else set()
    weights = []
    for candidate in graph[current]:
        if previous is None:
            weights.append(1.0)
        elif candidate == previous:
            weights.append(1.0 / return_p)
        elif candidate in previous_neighbours:
            weights.append(1.0)
        else:
            weights.append(1.0 / inout_q)
    return weights


def sample_index(weights, rng):
    """Draw from an unnormalized discrete distribution by linear scan.

    O(degree) per step. On a graph with a few hub nodes this single function
    is most of the runtime, which is the observation the fast stage acts on.
    """
    total = sum(weights)
    if total <= 0.0:
        return None
    threshold = rng.random() * total
    cumulative = 0.0
    for index, weight in enumerate(weights):
        cumulative += weight
        if cumulative >= threshold:
            return index
    return len(weights) - 1


def generate_walk(graph, start, length, return_p, inout_q, rng):
    """One second-order random walk. The sentence the skip-gram will read."""
    walk = [start]
    while len(walk) < length:
        current = walk[-1]
        if not graph[current]:
            break
        previous = walk[-2] if len(walk) >= 2 else None
        weights = biased_transition(graph, previous, current, return_p, inout_q)
        choice = sample_index(weights, rng)
        if choice is None:
            break
        walk.append(graph[current][choice])
    return walk


def build_noise_distribution(graph):
    """Degree to the 3/4 power. An empirical constant from word2vec, not a
    derived one - it damps hub nodes without removing them."""
    return [len(neighbours) ** 0.75 for neighbours in graph]


def sigmoid(value):
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


def train(graph, dimensions, walk_length, walks_per_node, window, negatives,
          return_p, inout_q, learning_rate, epochs):
    """Skip-gram with negative sampling over sampled walks."""
    rng = random.Random(SEED)
    num_nodes = len(graph)

    # TWO tables. `centre` is the one you keep; exporting `context` instead
    # produces vectors that work just well enough to hide the mistake.
    centre = [
        [(rng.random() - 0.5) / dimensions for _ in range(dimensions)]
        for _ in range(num_nodes)
    ]
    context = [[0.0] * dimensions for _ in range(num_nodes)]

    noise = build_noise_distribution(graph)

    for epoch in range(epochs):
        # The learning rate decays linearly and there is no other stopping
        # rule: training ends when the corpus is exhausted.
        rate = learning_rate * (1.0 - epoch / epochs)

        order = list(range(num_nodes))
        rng.shuffle(order)

        for _ in range(walks_per_node):
            for start in order:
                walk = generate_walk(graph, start, walk_length, return_p, inout_q, rng)

                for position, target in enumerate(walk):
                    low = max(0, position - window)
                    high = min(len(walk), position + window + 1)
                    for neighbour_position in range(low, high):
                        if neighbour_position == position:
                            continue
                        observed = walk[neighbour_position]

                        # One positive pair plus m negatives. Every update is
                        # a rank-one operation on two rows.
                        samples = [(observed, 1.0)]
                        for _ in range(negatives):
                            drawn = sample_index(noise, rng)
                            samples.append((drawn, 0.0))

                        gradient = [0.0] * dimensions
                        for node, label in samples:
                            score = sum(
                                centre[target][d] * context[node][d]
                                for d in range(dimensions)
                            )
                            error = (sigmoid(score) - label) * rate
                            for d in range(dimensions):
                                gradient[d] += error * context[node][d]
                                context[node][d] -= error * centre[target][d]
                        for d in range(dimensions):
                            centre[target][d] -= gradient[d]

    return centre


def cosine_neighbours(embeddings, query, top_k):
    """What the embedding is actually for: nearest neighbours in the space."""
    def norm(vector):
        return math.sqrt(sum(value * value for value in vector)) or 1.0

    query_vector = embeddings[query]
    query_norm = norm(query_vector)
    scored = []
    for node, vector in enumerate(embeddings):
        if node == query:
            continue
        dot = sum(a * b for a, b in zip(query_vector, vector))
        scored.append((dot / (query_norm * norm(vector)), node))
    scored.sort(reverse=True)
    return [node for _, node in scored[:top_k]]

"""The same method, with alias tables, typed boundaries, and one table kept.

What changes: sampling becomes O(1) via the alias method, the walk bias is
precomputed per (previous, current) edge rather than rebuilt on every step,
and the two embedding tables are returned in a structure that names which one
is which - because exporting the context table is a real and silent mistake.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass
from typing import NamedTuple, Sequence

Graph = Sequence[Sequence[int]]


class DisconnectedGraph(ValueError):
    """Raised when the graph cannot support walks at all."""


class InvalidWalkParameters(ValueError):
    """Raised on p or q outside the positive reals, which is meaningless."""


class AliasTable:
    """O(1) draws from a fixed discrete distribution after an O(n) build.

    The construction pairs each over-full bucket with an under-full one until
    every bucket holds exactly 1/n of the mass, split between one or two
    outcomes. A draw is then: pick a bucket, flip a biased coin. Two array
    lookups, no scan, no matter how skewed the distribution.
    """

    __slots__ = ('_probability', '_alias')

    def __init__(self, weights: Sequence[float]) -> None:
        if not weights:
            raise InvalidWalkParameters('an alias table needs at least one outcome')
        total = math.fsum(weights)
        if total <= 0.0:
            raise InvalidWalkParameters('alias weights must sum to something positive')

        size = len(weights)
        scaled = [weight * size / total for weight in weights]
        self._probability = [0.0] * size
        self._alias = [0] * size

        small = [i for i, value in enumerate(scaled) if value < 1.0]
        large = [i for i, value in enumerate(scaled) if value >= 1.0]

        while small and large:
            lesser = small.pop()
            greater = large.pop()
            self._probability[lesser] = scaled[lesser]
            self._alias[lesser] = greater
            scaled[greater] -= 1.0 - scaled[lesser]
            if scaled[greater] < 1.0:
                small.append(greater)
            else:
                large.append(greater)

        # Whatever is left is exactly 1.0 up to floating-point drift.
        for remaining in small + large:
            self._probability[remaining] = 1.0

    def draw(self, rng: random.Random) -> int:
        bucket = rng.randrange(len(self._probability))
        if rng.random() < self._probability[bucket]:
            return bucket
        return self._alias[bucket]


@dataclass(frozen=True)
class WalkConfig:
    """Frozen so a sampler's behaviour cannot change after construction."""

    walk_length: int
    walks_per_node: int
    return_p: float
    inout_q: float

    def __post_init__(self) -> None:
        # Guard clauses in the config, so every downstream method can assume
        # the parameters are sane rather than re-checking them.
        if self.return_p <= 0.0 or self.inout_q <= 0.0:
            raise InvalidWalkParameters('p and q must be strictly positive')
        if self.walk_length < 2:
            raise InvalidWalkParameters('a walk shorter than two nodes has no pairs')
        if self.walks_per_node < 1:
            raise InvalidWalkParameters('at least one walk per node is required')


class Embeddings(NamedTuple):
    """Both tables, named. `centre` is the one to export."""

    centre: list[list[float]]
    context: list[list[float]]


class Node2VecSampler:
    """Precomputes every alias table the walk will need.

    Memory is the cost: a second-order walk needs one table per DIRECTED
    EDGE, so the footprint is the sum of neighbour degrees over edges - which
    is quadratic in degree and the usual cause of an out-of-memory failure on
    a dense graph. That trade is the reason first-order DeepWalk still exists.
    """

    def __init__(self, graph: Graph, config: WalkConfig) -> None:
        if not graph:
            raise DisconnectedGraph('an empty graph admits no walks')
        if all(not neighbours for neighbours in graph):
            raise DisconnectedGraph('every node is isolated; no walk can move')

        self._graph = graph
        self._config = config
        self._first_order = [
            AliasTable([1.0] * len(neighbours)) if neighbours else None
            for neighbours in graph
        ]
        self._second_order: dict[tuple[int, int], AliasTable] = {}
        self._build_second_order()

    def _build_second_order(self) -> None:
        for previous, neighbours in enumerate(self._graph):
            previous_set = frozenset(neighbours)
            for current in neighbours:
                candidates = self._graph[current]
                if not candidates:
                    continue
                self._second_order[(previous, current)] = AliasTable(
                    [self._weight(previous, previous_set, candidate) for candidate in candidates]
                )

    def _weight(self, previous: int, previous_set: frozenset[int], candidate: int) -> float:
        """d(previous, candidate) is 0, 1 or 2 and nothing else."""
        if candidate == previous:
            return 1.0 / self._config.return_p
        if candidate in previous_set:
            return 1.0
        return 1.0 / self._config.inout_q

    def walk(self, start: int, rng: random.Random) -> list[int]:
        """One second-order walk, every step an O(1) alias draw."""
        first = self._first_order[start]
        if first is None:
            return [start]

        walk = [start, self._graph[start][first.draw(rng)]]
        while len(walk) < self._config.walk_length:
            previous, current = walk[-2], walk[-1]
            table = self._second_order.get((previous, current))
            if table is None:
                break
            walk.append(self._graph[current][table.draw(rng)])
        return walk


def sigmoid(value: float) -> float:
    """Branch on sign so exp never overflows on a large negative score."""
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


def train(
    graph: Graph,
    sampler: Node2VecSampler,
    config: WalkConfig,
    dimensions: int = 128,
    window: int = 5,
    negatives: int = 5,
    learning_rate: float = 0.025,
    epochs: int = 5,
    seed: int = 1234,
) -> Embeddings:
    """Skip-gram with negative sampling. Seeded, because a walk-based method
    is otherwise not reproducible even on identical data."""
    if dimensions < 1:
        raise InvalidWalkParameters('dimensions must be positive')

    rng = random.Random(seed)
    num_nodes = len(graph)

    centre = [
        [(rng.random() - 0.5) / dimensions for _ in range(dimensions)]
        for _ in range(num_nodes)
    ]
    context = [[0.0] * dimensions for _ in range(num_nodes)]

    # Degree to the 3/4 power, built once. Recomputing it per draw would
    # discard the entire benefit of the alias table.
    noise = AliasTable([max(len(neighbours), 1) ** 0.75 for neighbours in graph])

    for epoch in range(epochs):
        rate = learning_rate * max(1.0 - epoch / epochs, 1e-4)
        order = list(range(num_nodes))
        rng.shuffle(order)

        for _ in range(config.walks_per_node):
            for start in order:
                walk = sampler.walk(start, rng)
                _train_on_walk(walk, centre, context, noise, rng,
                               window, negatives, dimensions, rate)

    return Embeddings(centre=centre, context=context)


def _train_on_walk(
    walk: Sequence[int],
    centre: list[list[float]],
    context: list[list[float]],
    noise: AliasTable,
    rng: random.Random,
    window: int,
    negatives: int,
    dimensions: int,
    rate: float,
) -> None:
    """One walk's worth of updates. Extracted so the epoch loop stays flat."""
    for position, target in enumerate(walk):
        low = max(0, position - window)
        high = min(len(walk), position + window + 1)
        target_vector = centre[target]

        for neighbour_position in range(low, high):
            if neighbour_position == position:
                continue

            gradient = [0.0] * dimensions
            samples = [(walk[neighbour_position], 1.0)]
            samples.extend((noise.draw(rng), 0.0) for _ in range(negatives))

            for node, label in samples:
                context_vector = context[node]
                score = sum(a * b for a, b in zip(target_vector, context_vector))
                error = (sigmoid(score) - label) * rate
                for d in range(dimensions):
                    gradient[d] += error * context_vector[d]
                    context_vector[d] -= error * target_vector[d]

            for d in range(dimensions):
                target_vector[d] -= gradient[d]

"""Walks and updates in bulk. The alias draw vectorizes; the SGD batches.

Three changes:

  1. Alias sampling becomes array arithmetic. A draw is a uniform bucket
     index and a uniform coin, so N draws are two random arrays and one
     np.where - the whole walk frontier advances in a single step.
  2. Walks advance a FRONTIER rather than one node at a time: all N walks
     take their step t together, so the Python loop is over walk length
     (tens) rather than over walks (millions).
  3. Skip-gram updates batch. A block of pairs becomes two gathers, one
     row-wise dot, and a scatter-add - and the scatter must accumulate,
     because a hub node appears many times in one batch.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32
INDEX = np.int32


class VectorizedAlias:
    """Alias tables for every row of a ragged distribution, in flat arrays.

    Ragged rows are stored contiguously with an offset array - the same CSR
    idea as a sparse matrix - so a draw for any row is pure index arithmetic
    and no Python-level dispatch is needed to find the right table.
    """

    def __init__(self, offsets: NDArray[np.int32], weights: NDArray[np.float32]) -> None:
        self._offsets = offsets
        self._probability = np.empty_like(weights)
        self._alias = np.zeros(weights.shape, dtype=INDEX)

        for row in range(len(offsets) - 1):
            begin, end = int(offsets[row]), int(offsets[row + 1])
            if end <= begin:
                continue
            self._build_row(weights[begin:end], begin, end)

    def _build_row(self, weights: NDArray[np.float32], begin: int, end: int) -> None:
        size = end - begin
        scaled = weights * (size / weights.sum())
        probability = self._probability[begin:end]
        alias = self._alias[begin:end]

        small = [i for i in range(size) if scaled[i] < 1.0]
        large = [i for i in range(size) if scaled[i] >= 1.0]
        while small and large:
            lesser, greater = small.pop(), large.pop()
            probability[lesser] = scaled[lesser]
            alias[lesser] = greater
            scaled[greater] -= 1.0 - scaled[lesser]
            (small if scaled[greater] < 1.0 else large).append(greater)
        for remaining in small + large:
            probability[remaining] = 1.0

    def draw(self, rows: NDArray[np.int32], rng: np.random.Generator) -> NDArray[np.int32]:
        """One draw per entry of `rows`, all at once.

        A draw is: uniform bucket, uniform coin, take the bucket or its alias.
        Both are array operations, so there is no per-draw Python work.
        """
        begin = self._offsets[rows]
        sizes = self._offsets[rows + 1] - begin
        bucket = (rng.random(len(rows)) * sizes).astype(INDEX)
        flat = begin + bucket
        coin = rng.random(len(rows), dtype=FLOAT)
        # np.where, not a branch: taking the alias is a select over the whole
        # batch rather than a decision made one draw at a time.
        return np.where(coin < self._probability[flat], bucket, self._alias[flat]).astype(INDEX)


def generate_walks(
    indptr: NDArray[np.int32],
    indices: NDArray[np.int32],
    alias: VectorizedAlias,
    starts: NDArray[np.int32],
    length: int,
    rng: np.random.Generator,
) -> NDArray[np.int32]:
    """Advance every walk's frontier together. Returns (num_walks, length).

    The Python loop is over `length` - tens of iterations - rather than over
    the millions of walks, which is the whole structural change.
    """
    num_walks = len(starts)
    walks = np.empty((num_walks, length), dtype=INDEX)
    walks[:, 0] = starts
    frontier = starts

    for step in range(1, length):
        # One batched alias draw for every walk, then one gather to turn the
        # local neighbour offset into a global node id.
        local = alias.draw(frontier, rng)
        frontier = indices[indptr[frontier] + local]
        walks[:, step] = frontier

    return walks


def build_pairs(walks: NDArray[np.int32], window: int) -> tuple[NDArray[np.int32], NDArray[np.int32]]:
    """Every (centre, context) pair inside the window, as two flat arrays.

    Built by shifting the walk matrix rather than by iterating positions: an
    offset of k pairs column j with column j + k for every walk at once.
    """
    centres: list[NDArray[np.int32]] = []
    contexts: list[NDArray[np.int32]] = []
    length = walks.shape[1]

    for offset in range(1, window + 1):
        if offset >= length:
            break
        left, right = walks[:, :-offset], walks[:, offset:]
        centres.append(left.ravel())
        contexts.append(right.ravel())
        centres.append(right.ravel())
        contexts.append(left.ravel())

    return np.concatenate(centres), np.concatenate(contexts)


def train_batch(
    centre: NDArray[np.float32],
    context: NDArray[np.float32],
    centre_ids: NDArray[np.int32],
    context_ids: NDArray[np.int32],
    noise_table: NDArray[np.int32],
    negatives: int,
    rate: float,
    rng: np.random.Generator,
) -> None:
    """One batched SGNS step, updating both tables in place.

    Positives and negatives are concatenated into one block so the whole step
    is a single gather, a single row-wise dot, and a single scatter-add.
    """
    batch = len(centre_ids)
    drawn = noise_table[rng.integers(0, len(noise_table), size=batch * negatives)]

    rows = np.repeat(centre_ids, negatives + 1)
    columns = np.empty(batch * (negatives + 1), dtype=INDEX)
    columns[0::negatives + 1] = context_ids
    mask = np.ones(len(columns), dtype=bool)
    mask[0::negatives + 1] = False
    columns[mask] = drawn

    labels = np.zeros(len(columns), dtype=FLOAT)
    labels[0::negatives + 1] = 1.0

    # Two gathers, then one row-wise dot: einsum over matching rows avoids
    # materializing the full outer product a plain matmul would build.
    target = centre[rows]
    paired = context[columns]
    scores = np.einsum('ij,ij->i', target, paired)

    # Stable logistic in place; the gradient is (sigmoid(score) - label).
    np.negative(scores, out=scores)
    np.exp(scores, out=scores)
    scores += 1.0
    np.reciprocal(scores, out=scores)
    scores -= labels
    scores *= rate

    # Scatter-add, never a plain assignment: a hub node appears many times in
    # one batch and its updates must accumulate rather than overwrite.
    np.add.at(centre, rows, -scores[:, None] * paired)
    np.add.at(context, columns, -scores[:, None] * target)


def build_noise_table(degrees: NDArray[np.int32], resolution: int = 10_000_000) -> NDArray[np.int32]:
    """Degree^0.75 flattened into a uniform lookup table.

    Pre-expanding the distribution into a table makes a negative draw one
    uniform integer and one gather, with no per-draw search at all.
    """
    weights = np.maximum(degrees, 1).astype(np.float64) ** 0.75
    cumulative = np.cumsum(weights / weights.sum())
    positions = np.linspace(0.0, 1.0, resolution, endpoint=False)
    return np.searchsorted(cumulative, positions).astype(INDEX)