scripts/dance_core.py

"""
Faithful, self-contained reimplementation of DANCE's core mechanisms
(ICML 2026, "DANCE: Dynamic, Available, Neighbor-gated Condensation for
Federated Text-Attributed Graphs", OpenReview YhEWC2HiJl / arXiv 2601.16519).

Modules (Sec. 4):
  * label_aware_node_condensation  — Sec 4.2, Eq (3)-(4): k-means prototypes per
    class, label-stratified top-s selection to a budget K = ceil(r|V|).
  * budgeted_neighbor_gating        — Sec 4.3, Eq (5)-(9): cross-modal scores,
    hard per-hop budgets B_ell over 0/1/2-hop, chunk budget B_tok.
  * self_expressive_topology        — Sec 4.4, Eq (12)-(14): candidate support
    TopK_q, sparse self-expression via proximal gradient, symmetrise + top-k
    per row -> O(Kq)-edge adjacency.
  * hard_budget_project (Pi_B)      — top-B truncation + renormalise (Eq 6/8).

Plus helpers used by the theorem checks (Sec 5).
"""
from __future__ import annotations
import numpy as np


# ---------------------------------------------------------------- Pi_B / entmax
def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=axis, keepdims=True)


def hard_budget_project(p, B):
    """Pi_B: keep the top-B entries of distribution p, zero the rest, renormalise.
    Returns pi with ||pi||_0 <= B. p: 1-D nonneg weights summing to 1."""
    p = np.asarray(p, dtype=float)
    n = p.shape[0]
    if B >= n:
        return p.copy()
    keep = np.argpartition(-p, B - 1)[:B]
    pi = np.zeros_like(p)
    pi[keep] = p[keep]
    s = pi.sum()
    if s > 0:
        pi = pi / s
    return pi


def tail_mass(p, B):
    """delta_B(p) = sum of p over entries NOT in the top-B (Def 5.3)."""
    p = np.asarray(p, dtype=float)
    n = p.shape[0]
    if B >= n:
        return 0.0
    keep = np.argpartition(-p, B - 1)[:B]
    mask = np.ones(n, dtype=bool)
    mask[keep] = False
    return float(p[mask].sum())


def topB_margin(scores, B):
    """Delta_B(omega) = r_(B) - r_(B+1) on sorted pre-entmax scores (Def 5.5)."""
    s = np.sort(np.asarray(scores, dtype=float))[::-1]
    if B < 1 or B >= len(s):
        return 0.0
    return float(s[B - 1] - s[B])


def topB_set(scores, B):
    """Index set of the top-B scores (ties broken by index)."""
    B = min(B, len(scores))
    return frozenset(np.argsort(-np.asarray(scores))[:B].tolist())


# ------------------------------------------------ Sec 4.2 node condensation
def label_aware_node_condensation(Z, y, r, P=8, tau=0.0, conf=None, seed=0):
    """Eq (3)-(4). Z: [n,d] node embeddings; y: [n] (pseudo-)labels; r: ratio.
    Returns indices of the condensed core (size K = ceil(r n)), label-stratified,
    each class summarised by <=P k-means prototypes and scored by best cosine match.
    conf: optional confidence in [0,1]; nodes with conf<tau are dropped from stats."""
    from sklearn.cluster import MiniBatchKMeans
    rng = np.random.default_rng(seed)
    n = Z.shape[0]
    K = int(np.ceil(r * n))
    Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8)
    if conf is None:
        conf = np.ones(n)
    valid = conf >= tau
    classes = [c for c in np.unique(y[valid])]
    # per-class confident counts -> quotas proportional to distribution
    counts = {c: int(np.sum((y == c) & valid)) for c in classes}
    total = sum(counts.values()) or 1
    quotas = {c: int(np.floor(K * counts[c] / total)) for c in classes}
    # distribute remainder to largest classes
    rem = K - sum(quotas.values())
    for c in sorted(classes, key=lambda c: -counts[c])[:max(rem, 0)]:
        quotas[c] += 1
    core = []
    scores_all = np.full(n, -np.inf)
    for c in classes:
        idx = np.where((y == c) & valid)[0]
        if len(idx) == 0 or quotas[c] <= 0:
            continue
        Pc = min(P, len(idx))
        km = MiniBatchKMeans(n_clusters=Pc, random_state=seed, n_init=3, batch_size=256)
        km.fit(Zn[idx])
        protos = km.cluster_centers_
        protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8)
        sv = (Zn[idx] @ protos.T).max(axis=1)   # best prototype cosine match
        scores_all[idx] = sv
        take = min(quotas[c], len(idx))
        sel = idx[np.argsort(-sv)[:take]]
        core.extend(sel.tolist())
    core = np.array(sorted(core[:K]), dtype=int)
    return core, scores_all, quotas


# ------------------------------------------------ Sec 4.3 neighbor gating
def budgeted_neighbor_gating(v, g, t, hop_neighbors, Wq, Wk, budgets):
    """Eq (5)-(7) for a single core node v.
    g:[n,d] graph embeds; t:[n,d] text embeds; hop_neighbors: dict ell->list(idx)
    (0-hop is [v]); budgets: dict ell->B_ell. Returns selected {ell: {idx: weight}}."""
    d = g.shape[1]
    selected = {}
    for ell, cand in hop_neighbors.items():
        cand = list(cand)
        if not cand:
            selected[ell] = {}
            continue
        s = (Wq @ g[v]) @ (Wk @ t[np.array(cand)].T) / np.sqrt(d)   # Eq 5
        w = softmax(s)                                              # entmax ~ softmax
        pi = hard_budget_project(w, budgets[ell])                  # Pi_{B_ell}, Eq 6
        selected[ell] = {cand[i]: float(pi[i]) for i in np.nonzero(pi)[0]}
    return selected


def budgeted_chunk_selection(v, g, chunk_embeds, chunk_owner, cand_nodes, Ws, B_tok):
    """Eq (8)-(9): from chunks owned by selected neighbors, keep <=B_tok.
    chunk_embeds:[C,d]; chunk_owner:[C] node id; cand_nodes: selected neighbor ids.
    Returns {chunk_idx: weight} with support <= B_tok, and evidence embedding."""
    d = g.shape[1]
    cand_mask = np.isin(chunk_owner, list(cand_nodes))
    cand = np.where(cand_mask)[0]
    if len(cand) == 0:
        return {}, np.zeros(d)
    a = (Ws @ g[v]) @ chunk_embeds[cand].T / np.sqrt(d)   # Eq 8
    w = softmax(a)
    pi = hard_budget_project(w, B_tok)
    sel = {int(cand[i]): float(pi[i]) for i in np.nonzero(pi)[0]}
    emb = sum(wt * chunk_embeds[ci] for ci, wt in sel.items())   # Eq 9
    return sel, emb


# ------------------------------------------------ Sec 4.4 topology reconstruction
def candidate_support(X, S_prior, q):
    """Eq (12): C(i) = TopK_q(feature sim) ∪ TopK_q(evidence prior) per node i."""
    K = X.shape[0]
    Xn = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-8)
    feat_sim = Xn @ Xn.T
    np.fill_diagonal(feat_sim, -np.inf)
    np.fill_diagonal(S_prior, -np.inf)
    cand = []
    for i in range(K):
        a = np.argsort(-feat_sim[i])[:q]
        b = np.argsort(-S_prior[i])[:q]
        cand.append(set(a.tolist()) | set(b.tolist()))
    return cand


def self_expressive_topology(X, S_prior, q=10, k=5, alpha=1.0, beta=0.1,
                             steps=100, lr=0.05):
    """Eq (13)-(14): sparse self-expression on masked support, then symmetrise +
    keep top-k per row. Returns adjacency A (K x K, 0/weight) with O(Kk) edges and
    the coefficient matrix Z. Nonzeros of Z restricted to C(i): O(Kq)."""
    K = X.shape[0]
    cand = candidate_support(X, S_prior.copy(), q)
    mask = np.zeros((K, K), dtype=bool)
    for i, cs in enumerate(cand):
        for j in cs:
            if j != i:
                mask[i, j] = True
    Z = np.zeros((K, K))
    XtX = X @ X.T
    prior_pen = (1.0 - np.clip(S_prior, 0, 1))   # (1 - S_ij) weight on |Z_ij|
    np.fill_diagonal(prior_pen, 0.0)
    for _ in range(steps):
        # grad of alpha||X - XZ||_F^2 wrt Z = -2 alpha X^T(X - XZ) = -2a(XtX - XtX Z)
        grad = -2 * alpha * (XtX - XtX @ Z)
        Z = Z - lr * grad
        # proximal step for beta||Z||_1 + sum (1-S_ij)|Z_ij|  (soft-threshold)
        thr = lr * (beta + prior_pen)
        Z = np.sign(Z) * np.maximum(np.abs(Z) - thr, 0.0)
        Z[~mask] = 0.0
        np.fill_diagonal(Z, 0.0)
    W = np.abs(Z) + np.abs(Z).T                  # Eq 14 symmetrise
    A = np.zeros((K, K))
    for i in range(K):
        row = W[i].copy()
        if k < K:
            keep = np.argsort(-row)[:k]
        else:
            keep = np.nonzero(row)[0]
        A[i, keep] = row[keep]
    A = np.maximum(A, A.T)                        # symmetric top-k
    return A, Z, mask


if __name__ == "__main__":
    # quick self-test
    rng = np.random.default_rng(0)
    p = softmax(rng.normal(size=20))
    pi = hard_budget_project(p, 5)
    print("Pi_B support:", int((pi > 0).sum()), "tail_mass:", round(tail_mass(p, 5), 4))
    X = rng.normal(size=(30, 16))
    S = rng.random((30, 30))
    A, Z, mask = self_expressive_topology(X, S, q=6, k=4)
    print("edges:", int((A > 0).sum()), "Z nnz:", int((Z != 0).sum()), "mask support:", int(mask.sum()))