steeldb/

mece.rs

//! **MECE gate** — the information-gain test that decides whether a facet earns its place
//! (`design.py` §"MECE gate: Collectively-Exhaustive (coverage) × Mutually-Exclusive (orthogonality)").
//!
//! The reference implementation measures orthogonality as cosine between SPLADE *decoder weight*
//! vectors. Over a roaring index the same semantics are available directly from the postings — a facet's
//! incidence vector over situations *is* its activation — so:
//!
//! * **coverage** (Collectively Exhaustive): `|situations where the facet fires| / N`
//! * **maxcos** (Mutually Exclusive): the largest cosine between this facet's incidence vector and any
//!   other facet's, i.e. how redundant it is
//! * **gain** = `coverage × (1 − maxcos)` — high only when a facet covers a lot *and* covers something
//!   nothing else does. A candidate is kept iff `gain ≥ threshold` (an MDL-style split criterion).
//!
//! Computing this on the index rather than on model weights is both cheaper and more honest: it scores
//! the facets as the corpus actually instantiates them, not as the heads were trained.

use crate::bitmap::Postings;
use crate::index::InfonIndex;
use serde::Serialize;
use std::collections::BTreeMap;

#[derive(Debug, Clone, Serialize)]
pub struct FacetScore {
    pub facet: String,
    /// distinct tokens under this facet
    pub tokens: usize,
    /// fraction of situations where the facet fires (Collectively Exhaustive)
    pub coverage: f64,
    /// largest cosine against another facet's incidence vector (Mutually Exclusive; lower is better)
    pub maxcos: f64,
    /// the facet most redundant with this one
    pub nearest: String,
    /// coverage × (1 − maxcos)
    pub gain: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct MeceReport {
    pub situations: u32,
    pub facets: Vec<FacetScore>,
    /// mean coverage across facets — the set's Collectively-Exhaustive score
    pub mean_coverage: f64,
    /// mean maxcos — the set's (inverse) Mutually-Exclusive score
    pub mean_maxcos: f64,
    pub mean_gain: f64,
}

impl MeceReport {
    /// Facets that fail the gain threshold — redundant or too sparse to earn a retrieval head.
    pub fn failing(&self, threshold: f64) -> Vec<&FacetScore> {
        self.facets.iter().filter(|f| f.gain < threshold).collect()
    }
}

/// Union of every posting set under `facet` — the facet's incidence vector over situations (coverage).
fn facet_incidence<B: Postings>(ix: &InfonIndex<B>, facet: &str) -> B {
    let mut acc = B::empty();
    for tok in facet_tokens(ix, facet) {
        acc.or_inplace(&ix.post(&tok));
    }
    acc
}

/// The tokens belonging to a facet (exact name or `facet/...` prefix).
fn facet_tokens<B: Postings>(ix: &InfonIndex<B>, facet: &str) -> Vec<String> {
    let prefix = format!("{facet}/");
    ix.tokens().filter(|t| *t == facet || t.starts_with(&prefix)).cloned().collect()
}

/// Redundancy between two facets, measured in **term space** (as `design.py` compares decoder weight
/// vectors) rather than over union-incidence: the largest cosine between any token of `a` and any token
/// of `b`. Union-incidence degenerates in a relational corpus — every facet fires on every row, so all
/// union vectors are all-ones and cosine saturates to 1 regardless of redundancy. Token-level max cosine
/// instead answers the question that matters: *does some token here duplicate a token there?*
/// Minimum postings a token needs before it may drive the redundancy estimate. Without a floor, two rare
/// terms that happen to fire on the *same single document* score cosine 1.0 — noise reading as perfect
/// redundancy, which silently zeroes a candidate facet's gain.
pub const MIN_SUPPORT: usize = 3;

fn facet_maxcos<B: Postings>(ix: &InfonIndex<B>, a: &[String], b: &[String], cap: usize) -> f64 {
    let mut best = 0.0f64;
    for ta in a.iter().take(cap) {
        let pa = ix.post(ta);
        let la = pa.len() as f64;
        if pa.len() < MIN_SUPPORT {
            continue; // too rare to say anything about redundancy
        }
        for tb in b.iter().take(cap) {
            let pb = ix.post(tb);
            let lb = pb.len() as f64;
            if pb.len() < MIN_SUPPORT {
                continue;
            }
            let cos = pa.and(&pb).len() as f64 / (la * lb).sqrt();
            if cos > best {
                best = cos;
            }
        }
    }
    best
}

/// Score every facet in the index for MECE fitness. `skip` names structural facets that shouldn't be
/// judged (e.g. `src`, the per-file provenance tag).
pub fn report<B: Postings>(ix: &InfonIndex<B>, skip: &[&str]) -> MeceReport {
    let n = ix.situations();
    // facet → distinct token count
    let mut counts: BTreeMap<String, usize> = BTreeMap::new();
    for tok in ix.tokens() {
        let f = tok.split('/').next().unwrap_or(tok).to_string();
        *counts.entry(f).or_default() += 1;
    }
    let names: Vec<String> = counts.keys().filter(|f| !skip.contains(&f.as_str())).cloned().collect();
    let inc: Vec<B> = names.iter().map(|f| facet_incidence(ix, f)).collect();
    let lens: Vec<f64> = inc.iter().map(|b| b.len() as f64).collect();
    let toks: Vec<Vec<String>> = names.iter().map(|f| facet_tokens(ix, f)).collect();

    let mut facets = Vec::with_capacity(names.len());
    for i in 0..names.len() {
        let mut maxcos = 0.0f64;
        let mut nearest = String::new();
        for j in 0..names.len() {
            if i == j || lens[i] == 0.0 || lens[j] == 0.0 {
                continue;
            }
            let cos = facet_maxcos(ix, &toks[i], &toks[j], 64);
            if cos > maxcos {
                maxcos = cos;
                nearest = names[j].clone();
            }
        }
        let coverage = if n == 0 { 0.0 } else { lens[i] / n as f64 };
        facets.push(FacetScore {
            facet: names[i].clone(),
            tokens: counts.get(&names[i]).copied().unwrap_or(0),
            coverage: round3(coverage),
            maxcos: round3(maxcos),
            nearest,
            gain: round3(coverage * (1.0 - maxcos)),
        });
    }
    facets.sort_by(|a, b| b.gain.partial_cmp(&a.gain).unwrap_or(std::cmp::Ordering::Equal));
    let k = facets.len().max(1) as f64;
    let mean_coverage = round3(facets.iter().map(|f| f.coverage).sum::<f64>() / k);
    let mean_maxcos = round3(facets.iter().map(|f| f.maxcos).sum::<f64>() / k);
    let mean_gain = round3(facets.iter().map(|f| f.gain).sum::<f64>() / k);
    MeceReport { situations: n, facets, mean_coverage, mean_maxcos, mean_gain }
}

fn round3(v: f64) -> f64 {
    (v * 1000.0).round() / 1000.0
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::RoarPostings;
    use std::collections::HashMap;

    #[test]
    fn rare_terms_cannot_fake_perfect_redundancy() {
        // two facets whose only shared signal is a single co-occurring doc must NOT read as redundant
        let raw = HashMap::from([
            ("alpha/rare".to_string(), vec![0u32]),
            ("beta/rare".to_string(), vec![0u32]),
            // give each facet a well-supported, genuinely disjoint token so coverage is meaningful
            ("alpha/common".to_string(), vec![0u32, 1, 2, 3]),
            ("beta/common".to_string(), vec![4, 5, 6, 7]),
        ]);
        let ix: InfonIndex<RoarPostings> = InfonIndex::from_postings(raw, 8);
        let r = report(&ix, &["src"]);
        let by = |n: &str| r.facets.iter().find(|f| f.facet == n).unwrap().clone();
        // the single-doc overlap is below MIN_SUPPORT, so it is ignored; the supported tokens are disjoint
        assert!(by("alpha").maxcos < 1.0, "rare overlap must not saturate: {:?}", by("alpha"));
        assert!(by("alpha").gain > 0.0, "both facets earn gain: {:?}", by("alpha"));
        assert!(by("beta").gain > 0.0);
    }

    #[test]
    fn gain_rewards_coverage_and_punishes_redundancy() {
        // `country` and `nation` are near-duplicates (same rows) → both should score poorly on maxcos.
        // `powertrain` is disjoint-ish and well-covered → best gain. `rare` covers almost nothing.
        let raw = HashMap::from([
            ("country/japan".to_string(), vec![0u32, 1, 2, 3]),
            ("nation/japan".to_string(), vec![0u32, 1, 2, 3]),
            ("powertrain/electric".to_string(), vec![0, 2]),
            ("powertrain/diesel".to_string(), vec![1, 3]),
            ("rare/thing".to_string(), vec![0]),
        ]);
        let ix: InfonIndex<RoarPostings> = InfonIndex::from_postings(raw, 4);
        let r = report(&ix, &["src"]);
        eprintln!("{}", serde_json::to_string_pretty(&r).unwrap());
        let by = |n: &str| r.facets.iter().find(|f| f.facet == n).unwrap().clone();
        // duplicates are maximally redundant → zero gain despite full coverage
        assert_eq!(by("country").maxcos, 1.0);
        assert_eq!(by("country").nearest, "nation");
        assert_eq!(by("country").gain, 0.0);
        // full coverage, and its overlap with country is partial → positive gain, the best in the set
        assert_eq!(by("powertrain").coverage, 1.0);
        assert!(by("powertrain").gain > 0.0, "disjoint well-covered facet must earn gain");
        assert_eq!(r.facets[0].facet, "powertrain", "best gain = broad coverage + low redundancy");
        assert!(by("powertrain").maxcos < 1.0, "term-space cosine must not saturate for co-extensive facets");
        // sparse facet earns little
        assert!(by("rare").coverage <= 0.25);
        assert!(by("rare").gain < by("powertrain").gain);
        // the gate flags the redundant pair
        let failing: Vec<&str> = r.failing(0.1).iter().map(|f| f.facet.as_str()).collect();
        assert!(failing.contains(&"country") && failing.contains(&"nation"));
    }
}