steeldb/text/

ot.rs

//! Entropy-regularised optimal transport (Sinkhorn–Knopp) + k-means codebook — the ontology-sensing
//! numeric core, ported from `src/ingest/spo/ot.ts` (itself the port of python spo_sinkhorn.py).
//! Vectors are assumed L2-normalised, so cosine == dot. This is what turns a cloud of span embeddings
//! into a MECE facet codebook — no hand-authored schema.

pub type Vec32 = Vec<f32>;

fn dot(a: &[f32], b: &[f32]) -> f32 {
    let mut s = 0.0;
    for i in 0..a.len() {
        s += a[i] * b[i];
    }
    s
}

fn l2(v: &[f32]) -> Vec32 {
    let mut n = 0.0f32;
    for &x in v {
        n += x * x;
    }
    n = n.sqrt() + 1e-9;
    v.iter().map(|x| x / n).collect()
}

/// k-means++ over cosine (vectors L2-normalised → argmax dot = nearest). Returns k prototypes.
pub fn kmeans(x: &[Vec32], k: usize, iters: usize, seed: u32) -> Vec<Vec32> {
    let k = k.max(2).min((x.len() / 4).max(2));
    let dim = x[0].len();
    let mut s = seed;
    let mut rnd = || {
        s = s.wrapping_mul(1664525).wrapping_add(1013904223);
        s as f64 / 4294967296.0
    };
    let mut c: Vec<Vec32> = vec![x[(rnd() * x.len() as f64) as usize].clone()];
    while c.len() < k {
        let d2: Vec<f32> = x
            .iter()
            .map(|xi| {
                let best = c.iter().fold(-1.0f32, |b, cj| b.max(dot(xi, cj)));
                (1.0 - best).max(1e-6)
            })
            .collect();
        let tot: f32 = d2.iter().sum();
        let mut r = rnd() as f32 * tot;
        let mut idx = 0;
        while idx < d2.len() && {
            r -= d2[idx];
            r > 0.0
        } {
            idx += 1;
        }
        c.push(x[idx.min(x.len() - 1)].clone());
    }
    for _ in 0..iters {
        let mut sum = vec![vec![0.0f32; dim]; k];
        let mut cnt = vec![0usize; k];
        for xi in x {
            let mut a = 0;
            let mut bestv = f32::NEG_INFINITY;
            for j in 0..k {
                let d = dot(xi, &c[j]);
                if d > bestv {
                    bestv = d;
                    a = j;
                }
            }
            cnt[a] += 1;
            for i in 0..dim {
                sum[a][i] += xi[i];
            }
        }
        let mut moved = false;
        for j in 0..k {
            if cnt[j] == 0 {
                continue;
            }
            let nc = l2(&sum[j]);
            if 1.0 - dot(&nc, &c[j]) > 1e-5 {
                moved = true;
            }
            c[j] = nc;
        }
        if !moved {
            break;
        }
    }
    c
}

/// Sinkhorn OT. `cost[n][k]` (cosine distance); uniform source + target marginals. Returns (plan, cost).
///
/// This is unmodified Sinkhorn–Knopp: Gibbs kernel `exp(-C/eps)` with uniform marginals. Use
/// [`sinkhorn_weighted`] when the rows are derived from reified hyperedges, where degree is not uniform.
pub fn sinkhorn(m: &[Vec<f32>], eps: f32, iters: usize) -> (Vec<Vec<f32>>, f32) {
    let n = m.len();
    let k = m[0].len();
    let kmat: Vec<Vec<f32>> = m.iter().map(|row| row.iter().map(|v| (-v / eps).exp()).collect()).collect();
    let (a, b) = (1.0 / n as f32, 1.0 / k as f32);
    let mut u = vec![1.0f32; n];
    let mut v = vec![1.0f32; k];
    for _ in 0..iters {
        for i in 0..n {
            let mut s = 0.0;
            for j in 0..k {
                s += kmat[i][j] * v[j];
            }
            u[i] = a / (s + 1e-12);
        }
        for j in 0..k {
            let mut s = 0.0;
            for i in 0..n {
                s += kmat[i][j] * u[i];
            }
            v[j] = b / (s + 1e-12);
        }
    }
    let mut cost = 0.0;
    let pi: Vec<Vec<f32>> = (0..n)
        .map(|i| {
            (0..k)
                .map(|j| {
                    let p = u[i] * kmat[i][j] * v[j];
                    cost += p * m[i][j];
                    p
                })
                .collect()
        })
        .collect();
    (pi, cost)
}

/// Sinkhorn OT with a **non-uniform source marginal**.
///
/// NOTE ON PROVENANCE: this is an extension, not the reference behaviour. Both the reference Python
/// (`spo_sinkhorn.py`) and [`sinkhorn`] above use uniform source AND target marginals — the solver is
/// unmodified textbook Sinkhorn–Knopp in both. The reference's adaptation to reified hyperedges happens
/// UPSTREAM of the solver: spans are clustered separately per kind, and non-semantic kinds (quantities,
/// boilerplate) are routed out to numeric rules rather than entering the codebook at all.
///
/// Standard Sinkhorn gives every row the same mass `1/n`. That is correct when rows are interchangeable
/// samples, and wrong when rows are terms drawn from reified hyperedges, because hyperedge degree is heavily
/// skewed: a hub entity participates in a large fraction of all statements while a rare one appears twice.
/// Under a uniform source marginal the solver must move the same mass for both, so a supernode's term is
/// forced to spread across topics it has no real affinity for, and the topic it truly belongs to is diluted.
///
/// Weighting each row by its hyperedge support fixes the asymmetry: mass now reflects how much evidence the
/// corpus actually has for that term. The target marginal stays uniform, because that is what prevents any
/// one topic from collapsing and absorbing the others.
///
/// `support[i]` is the number of reified statements term `i` participates in. Weights are normalised to sum
/// to 1; a zero or missing support falls back to uniform so the function is total.
pub fn sinkhorn_weighted(m: &[Vec<f32>], eps: f32, iters: usize, support: &[f32]) -> (Vec<Vec<f32>>, f32) {
    let n = m.len();
    let k = m[0].len();
    let kmat: Vec<Vec<f32>> = m.iter().map(|row| row.iter().map(|v| (-v / eps).exp()).collect()).collect();

    // source marginal from hyperedge support; dampened with a square root so a hub does not take over the
    // plan entirely — the same reason term weighting uses sqrt(tf) rather than tf
    let mut a: Vec<f32> = (0..n)
        .map(|i| support.get(i).copied().unwrap_or(0.0).max(0.0).sqrt())
        .collect();
    let total: f32 = a.iter().sum();
    if total <= 0.0 {
        a = vec![1.0 / n as f32; n];
    } else {
        for x in a.iter_mut() {
            *x /= total;
        }
    }
    let b = 1.0 / k as f32;

    let mut u = vec![1.0f32; n];
    let mut v = vec![1.0f32; k];
    for _ in 0..iters {
        for i in 0..n {
            let mut s = 0.0;
            for j in 0..k {
                s += kmat[i][j] * v[j];
            }
            u[i] = a[i] / (s + 1e-12);
        }
        for j in 0..k {
            let mut s = 0.0;
            for i in 0..n {
                s += kmat[i][j] * u[i];
            }
            v[j] = b / (s + 1e-12);
        }
    }
    let mut cost = 0.0;
    let pi: Vec<Vec<f32>> = (0..n)
        .map(|i| {
            (0..k)
                .map(|j| {
                    let p = u[i] * kmat[i][j] * v[j];
                    cost += p * m[i][j];
                    p
                })
                .collect()
        })
        .collect();
    (pi, cost)
}

/// Codebook = k-means prototypes; assign each vector via the Sinkhorn plan (argmax over targets).
pub fn codebook(x: &[Vec32], k: usize, eps: f32) -> (Vec<Vec32>, Vec<usize>, f32) {
    if x.len() < 2 {
        return (x.to_vec(), vec![0; x.len()], 0.0);
    }
    let protos = kmeans(x, k, 60, 1);
    let m: Vec<Vec<f32>> = x.iter().map(|xi| protos.iter().map(|c| 1.0 - dot(xi, c)).collect()).collect();
    let (pi, cost) = sinkhorn(&m, eps, 200);
    let assign = pi
        .iter()
        .map(|row| {
            let mut a = 0;
            let mut bv = f32::NEG_INFINITY;
            for (j, &p) in row.iter().enumerate() {
                if p > bv {
                    bv = p;
                    a = j;
                }
            }
            a
        })
        .collect();
    (protos, assign, cost)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn weighted_sinkhorn_respects_hyperedge_support() {
        let cost = vec![vec![0.05f32, 0.95], vec![0.10, 0.90], vec![0.95, 0.05]];
        let support = vec![100.0f32, 50.0, 2.0];
        let (uni, _) = sinkhorn(&cost, 0.1, 80);
        let (wei, _) = sinkhorn_weighted(&cost, 0.1, 80, &support);
        let row_mass = |p: &Vec<Vec<f32>>, i: usize| p[i].iter().sum::<f32>();
        assert!((row_mass(&uni, 0) - row_mass(&uni, 2)).abs() < 1e-4);
        assert!(row_mass(&wei, 0) > row_mass(&wei, 2) * 3.0);
        assert!(wei[0][0] > wei[0][1]);
        assert!(wei[2][1] > wei[2][0]);
    }

    #[test]
    fn weighted_sinkhorn_falls_back_to_uniform_without_support() {
        let cost = vec![vec![0.1f32, 0.9], vec![0.9, 0.1]];
        let (a, _) = sinkhorn(&cost, 0.1, 50);
        let (b, _) = sinkhorn_weighted(&cost, 0.1, 50, &[0.0, 0.0]);
        for i in 0..2 { for j in 0..2 { assert!((a[i][j]-b[i][j]).abs() < 1e-5); } }
    }
}