core/representation.py

"""Aurelius core — structural node embeddings + text/structure fusion.

Why structure at all: text embeddings only know what a node *says*, not
where it *sits*. Two nodes with unrelated text but heavily overlapping
neighbourhoods (a gene and a disease, two tickers that co-move, two
modules imported together) are close in structural space while far in
text space — exactly the "hidden bridge" signal discover() ranks on, and
the reason a bare-text engine misses non-obvious connections.

Implementation: DeepWalk-style node2vec — uniform random walks over the
stored edge list, then skip-gram with negative sampling (SGNS) trained
with batched numpy SGD. Pure numpy on purpose: no gensim/PyG dependency,
runs at ingest time on the same free-tier CPU as everything else, and at
demo scale (≤ ~50k nodes) finishes in minutes. The upgrade path (biased
p/q walks, GraphSAGE for inductive embeddings over node features) slots
in behind the same two functions.

Fusion: z = [ text_norm ; α · struct_norm ], α per source (config
FUSION_ALPHA). Cosine over z blends the two signals with α² relative
weight on structure.
"""

from __future__ import annotations

import numpy as np

from config import (
    N2V_DIM, N2V_WALKS_PER_NODE, N2V_WALK_LENGTH,
    N2V_WINDOW, N2V_EPOCHS, N2V_NEGATIVES, FUSION_ALPHA,
)


def random_walks(edges: list[tuple[str, str, float]],
                 walks_per_node: int = N2V_WALKS_PER_NODE,
                 walk_length: int = N2V_WALK_LENGTH,
                 seed: int = 42) -> tuple[list[str], np.ndarray]:
    """Uniform random walks over an (undirected view of an) edge list.

    Returns (vocab, walk_matrix[int32 n_walks × walk_length]) with -1
    padding for dead-end truncation. Treating the graph as undirected for
    walk purposes is standard: structural similarity cares about shared
    neighbourhoods, not edge direction.
    """
    rng = np.random.default_rng(seed)
    adj: dict[str, list[str]] = {}
    for s, d, _w in edges:
        adj.setdefault(s, []).append(d)
        adj.setdefault(d, []).append(s)

    vocab = sorted(adj)
    index = {v: i for i, v in enumerate(vocab)}
    adj_idx = [np.array([index[nb] for nb in adj[v]], dtype=np.int32)
               for v in vocab]

    n = len(vocab)
    walks = np.full((n * walks_per_node, walk_length), -1, dtype=np.int32)
    row = 0
    for start in range(n):
        for _ in range(walks_per_node):
            cur = start
            walks[row, 0] = cur
            for pos in range(1, walk_length):
                nbs = adj_idx[cur]
                if nbs.size == 0:
                    break
                cur = int(nbs[rng.integers(nbs.size)])
                walks[row, pos] = cur
            row += 1
    return vocab, walks


def sgns_train(vocab: list[str], walks: np.ndarray,
               dim: int = N2V_DIM, window: int = N2V_WINDOW,
               epochs: int = N2V_EPOCHS, negatives: int = N2V_NEGATIVES,
               lr: float = 0.025, batch: int = 8192,
               seed: int = 42) -> dict[str, np.ndarray]:
    """Skip-gram with negative sampling over the walk corpus (batched
    numpy SGD). Returns id → dim-vector."""
    rng = np.random.default_rng(seed)
    n = len(vocab)
    if n == 0:
        return {}

    # (center, context) pairs from every window position.
    centers, contexts = [], []
    for offset in range(1, window + 1):
        c = walks[:, :-offset].ravel()
        x = walks[:, offset:].ravel()
        ok = (c >= 0) & (x >= 0)
        centers.append(c[ok]); contexts.append(x[ok])
    C = np.concatenate(centers)
    X = np.concatenate(contexts)
    n_pairs = C.size
    if n_pairs == 0:
        return {v: np.zeros(dim, dtype=np.float32) for v in vocab}

    # Unigram^0.75 negative-sampling table.
    counts = np.bincount(walks[walks >= 0].ravel(), minlength=n).astype(np.float64)
    probs = counts ** 0.75
    probs /= probs.sum()

    W  = (rng.random((n, dim), dtype=np.float32) - 0.5) / dim   # target
    Cw = np.zeros((n, dim), dtype=np.float32)                    # context

    def sigmoid(z):
        return 1.0 / (1.0 + np.exp(-np.clip(z, -8, 8)))

    order = rng.permutation(n_pairs)
    for epoch in range(epochs):
        rng.shuffle(order)
        for i0 in range(0, n_pairs, batch):
            idx = order[i0:i0 + batch]
            c, x = C[idx], X[idx]
            wc = W[c]                                   # B × d
            # positive pass
            xc = Cw[x]
            g = (sigmoid((wc * xc).sum(1)) - 1.0)[:, None] * lr  # B × 1
            dwc = g * xc
            np.add.at(Cw, x, -g * wc)
            # negative pass
            neg = rng.choice(n, size=(idx.size, negatives), p=probs)
            xn = Cw[neg]                                # B × K × d
            gn = sigmoid(np.einsum("bd,bkd->bk", wc, xn)) * lr   # B × K
            dwc += np.einsum("bk,bkd->bd", gn, xn)
            np.add.at(Cw, neg.ravel(),
                      -(gn[..., None] * wc[:, None, :]).reshape(-1, dim))
            np.add.at(W, c, -dwc)
            # On small graphs a node recurs many times per batch, so the
            # summed np.add.at updates act like a huge effective lr and the
            # matrices diverge (float32 overflow). Bounding the matrices
            # keeps training stable at any graph size.
            np.clip(W, -4.0, 4.0, out=W)
            np.clip(Cw, -4.0, 4.0, out=Cw)
        print(f"[n2v] epoch {epoch + 1}/{epochs} done ({n_pairs:,} pairs)")

    W = np.nan_to_num(W, nan=0.0, posinf=0.0, neginf=0.0)
    return {v: W[i].copy() for i, v in enumerate(vocab)}


def node2vec_embeddings(edges: list[tuple[str, str, float]],
                        dim: int = N2V_DIM) -> dict[str, np.ndarray]:
    """edges (src_id, dst_id, weight) → {node_id: structural vector}."""
    if not edges:
        return {}
    vocab, walks = random_walks(edges)
    print(f"[n2v] {len(vocab):,} nodes, {len(edges):,} edges, "
          f"{walks.shape[0]:,} walks")
    return sgns_train(vocab, walks, dim=dim)


def fuse(text_emb: np.ndarray | None, struct_emb: np.ndarray | None,
         source: str, struct_dim: int = N2V_DIM) -> np.ndarray | None:
    """z = [text_norm ; α·struct_norm]. Missing halves are zero-padded so
    fused vectors of one source are always comparable with each other."""
    alpha = FUSION_ALPHA.get(source, 0.5)
    if text_emb is None and struct_emb is None:
        return None

    def _norm(v):
        v = np.asarray(v, dtype=np.float32)
        nv = np.linalg.norm(v)
        return v / nv if nv > 0 else v

    if text_emb is not None:
        t = _norm(text_emb)
    else:
        t = None
    if struct_emb is not None:
        s = alpha * _norm(struct_emb)
    else:
        s = np.zeros(struct_dim, dtype=np.float32)
    if t is None:
        # struct-only: pad an all-zero text half of unknown dim is useless —
        # return struct alone (comparisons stay within-source anyway).
        return s
    return np.concatenate([t, s])