"""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])