car-topology 0.55.0

//! The topology codebook — a short, fixed list of graphs to select among.
//!
//! ## Why a codebook instead of a decoder
//!
//! The published designers treat topology design as conditional graph
//! generation: a variational, autoregressive, or diffusion decoder searches the
//! `N×N` adjacency space and a graph network ranks what it sampled. The paper
//! this module implements measures that the space does not need a decoder —
//! topologies that survive a reward filter collapse to about six distinct
//! graphs even as the quantizer's capacity grows from 8 to 64, and the best
//! fixed family stays within 1.4 accuracy points of every generated topology
//! they measured. Capacity spent modelling the adjacency manifold is spent
//! where the problem does not live.
//!
//! ## What this implementation does differently, and why
//!
//! The paper fits a vector-quantized autoencoder (encoder MLP → nearest
//! codebook entry → decoder MLP → edge logits, EMA codebook updates,
//! straight-through gradients). CAR has no autograd framework in the workspace
//! and this crate is a deterministic leaf, so the quantizer here is a
//! **deterministic k-medoids over Hamming distance on `vec(A)`**: farthest-first
//! initialization, Lloyd iterations to a fixed point, medoids chosen from the
//! observed survivors themselves.
//!
//! That substitution is not a shortcut past the paper's finding — it is the
//! finding taken at face value. A learned decoder earns its keep by
//! interpolating a manifold; when the reward-surviving set is six points, an
//! index over those points loses nothing, and it gains three things a VQ-VAE
//! does not have: codes that decode to topologies that were actually executed
//! (no sigmoid threshold that can emit a graph nobody ran), a fit that is
//! reproducible bit-for-bit, and no dead-code reset heuristic to tune.
//!
//! What is genuinely lost: the paper's codebook can emit a graph *between* two
//! observed ones, and this one cannot. [`Codebook::used_codes`] is how you check
//! whether that matters on your records — if it saturates well below capacity,
//! as the paper measures and [`crate::diagnostics::design_space_collapse`]
//! re-measures on yours, there was no manifold to interpolate.

use serde::{Deserialize, Serialize};

use crate::error::TopologyError;
use crate::record::{RecordSet, DEFAULT_SURVIVOR_THRESHOLD};
use crate::topology::Topology;

/// How the codebook is fitted.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CodebookConfig {
    /// Maximum number of codes `K`. The paper's default is 16, and its capacity
    /// ablation finds at most six in use for every `K ≥ 8`.
    pub capacity: usize,
    /// Records with utility strictly above this enter the codebook. The paper's
    /// `u > 0.5`: the codebook stores topologies that solved their task.
    pub survivor_threshold: f32,
    /// Cap on Lloyd iterations. The fit stops earlier on a fixed point; this is
    /// only a guard against a pathological oscillation.
    pub max_iterations: usize,
}

impl Default for CodebookConfig {
    fn default() -> Self {
        Self {
            capacity: 16,
            survivor_threshold: DEFAULT_SURVIVOR_THRESHOLD,
            max_iterations: 32,
        }
    }
}

impl CodebookConfig {
    fn validate(&self) -> Result<(), TopologyError> {
        if self.capacity == 0 {
            return Err(TopologyError::BadConfig {
                field: "capacity",
                expected: "at least 1",
                found: "0".into(),
            });
        }
        if !self.survivor_threshold.is_finite() {
            return Err(TopologyError::BadConfig {
                field: "survivor_threshold",
                expected: "finite",
                found: format!("{}", self.survivor_threshold),
            });
        }
        Ok(())
    }
}

/// A fitted, query-independent index of reward-surviving topologies.
///
/// Query dependence lives entirely in [`crate::CodePredictor`]; the codebook is
/// the same for every query, exactly as in the paper (neither its encoder nor
/// its decoder is conditioned on `c`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "CodebookWire")]
pub struct Codebook {
    n: usize,
    codes: Vec<Topology>,
    /// How many distinct survivor topologies quantize to each code.
    occupancy: Vec<usize>,
    capacity: usize,
    distinct_survivors: usize,
}

/// On-disk shape, validated on the way in — see [`Codebook::validate`].
#[derive(Deserialize)]
struct CodebookWire {
    n: usize,
    codes: Vec<Topology>,
    occupancy: Vec<usize>,
    capacity: usize,
    distinct_survivors: usize,
}

impl TryFrom<CodebookWire> for Codebook {
    type Error = TopologyError;

    fn try_from(wire: CodebookWire) -> Result<Self, Self::Error> {
        let book = Codebook {
            n: wire.n,
            codes: wire.codes,
            occupancy: wire.occupancy,
            capacity: wire.capacity,
            distinct_survivors: wire.distinct_survivors,
        };
        book.validate()?;
        Ok(book)
    }
}

impl Codebook {
    /// Fit a codebook to the reward-surviving records of `records`.
    ///
    /// Deterministic: the same records and config always produce the same codes
    /// in the same order, because every stage breaks ties on
    /// [`Topology::key`] rather than on iteration order.
    pub fn fit(records: &RecordSet, config: &CodebookConfig) -> Result<Self, TopologyError> {
        config.validate()?;

        let survivor_indices = records.survivors(config.survivor_threshold);
        if survivor_indices.is_empty() {
            return Err(TopologyError::NoRecords { kind: "surviving" });
        }

        // Distinct survivor topologies, in a key-sorted order so the fit does
        // not depend on the order records were logged in.
        let mut distinct: Vec<Topology> = Vec::new();
        for &i in &survivor_indices {
            let t = &records.records()[i].topology;
            if !distinct.iter().any(|d| d == t) {
                distinct.push(t.clone());
            }
        }
        distinct.sort_by(|a, b| a.key().cmp(&b.key()));
        let distinct_survivors = distinct.len();

        let n = records.team_size();

        // The collapse case, and the common one: fewer distinct survivors than
        // capacity, so every one gets its own code and quantization is exact.
        if distinct.len() <= config.capacity {
            let occupancy = vec![1; distinct.len()];
            return Ok(Self {
                n,
                codes: distinct,
                occupancy,
                capacity: config.capacity,
                distinct_survivors,
            });
        }

        let codes = k_medoids(&distinct, config.capacity, config.max_iterations)?;
        let mut occupancy = vec![0usize; codes.len()];
        for t in &distinct {
            occupancy[nearest(&codes, t)?] += 1;
        }
        Ok(Self {
            n,
            codes,
            occupancy,
            capacity: config.capacity,
            distinct_survivors,
        })
    }

    /// Build a codebook directly from a list of topologies, bypassing record
    /// collection.
    ///
    /// This is the cold-start path: an operator with no execution history yet
    /// can index [`crate::CoordinationShape::ALL`] and still get one-pass
    /// selection, at the cost of a predictor and proxy that have nothing
    /// measured to learn from. Duplicates are dropped.
    pub fn from_topologies(topologies: Vec<Topology>) -> Result<Self, TopologyError> {
        if topologies.is_empty() {
            return Err(TopologyError::EmptyCodebook);
        }
        let n = topologies[0].n();
        let mut codes: Vec<Topology> = Vec::new();
        for t in topologies {
            if t.n() != n {
                return Err(TopologyError::SizeMismatch {
                    expected: n,
                    found: t.n(),
                });
            }
            if !codes.iter().any(|c| c == &t) {
                codes.push(t);
            }
        }
        let occupancy = vec![1; codes.len()];
        let capacity = codes.len();
        let distinct_survivors = codes.len();
        Ok(Self {
            n,
            codes,
            occupancy,
            capacity,
            distinct_survivors,
        })
    }

    /// Check the invariants every constructor establishes.
    ///
    /// Serde bypasses constructors, so a persisted codebook can come back with
    /// an occupancy vector that no longer lines up with its codes, or codes
    /// over mixed team sizes. Neither panics — they make `used_codes` report a
    /// number that means nothing, and `decode` hand out a topology the proxy
    /// cannot score. Failing the load is better than either.
    pub fn validate(&self) -> Result<(), TopologyError> {
        if self.codes.is_empty() {
            return Err(TopologyError::EmptyCodebook);
        }
        if self.occupancy.len() != self.codes.len() {
            return Err(TopologyError::BadConfig {
                field: "occupancy",
                expected: "one entry per code",
                found: format!(
                    "{} entries for {} codes",
                    self.occupancy.len(),
                    self.codes.len()
                ),
            });
        }
        for code in &self.codes {
            if code.n() != self.n {
                return Err(TopologyError::SizeMismatch {
                    expected: self.n,
                    found: code.n(),
                });
            }
        }
        if self.capacity < self.codes.len() {
            return Err(TopologyError::BadConfig {
                field: "capacity",
                expected: "at least the number of codes",
                found: format!("{} for {} codes", self.capacity, self.codes.len()),
            });
        }
        Ok(())
    }

    /// Team size every code is defined over.
    pub fn team_size(&self) -> usize {
        self.n
    }

    /// The codes, in index order.
    pub fn codes(&self) -> &[Topology] {
        &self.codes
    }

    /// Number of codes actually emitted — `min(capacity, distinct survivors)`
    /// in the collapse case.
    pub fn len(&self) -> usize {
        self.codes.len()
    }

    /// Whether the codebook holds no codes. Never true for a fitted codebook.
    pub fn is_empty(&self) -> bool {
        self.codes.is_empty()
    }

    /// Capacity `K` the codebook was fitted at.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Number of codes with at least one survivor assigned.
    ///
    /// This is the paper's headline measurement: on their records it saturates
    /// near six for every `K ≥ 8`, which is what makes searching the full `N×N`
    /// adjacency at test time unjustifiable. Re-measure it on your own records
    /// via [`crate::diagnostics::design_space_collapse`] before believing it
    /// holds for your workload.
    pub fn used_codes(&self) -> usize {
        self.occupancy.iter().filter(|&&o| o > 0).count()
    }

    /// Number of distinct reward-surviving topologies the fit saw.
    pub fn distinct_survivors(&self) -> usize {
        self.distinct_survivors
    }

    /// Decode a code index to its topology.
    pub fn decode(&self, code: usize) -> Option<&Topology> {
        self.codes.get(code)
    }

    /// Quantize a topology to its nearest code under Hamming distance.
    pub fn encode(&self, topology: &Topology) -> Result<usize, TopologyError> {
        if topology.n() != self.n {
            return Err(TopologyError::SizeMismatch {
                expected: self.n,
                found: topology.n(),
            });
        }
        nearest(&self.codes, topology)
    }
}

/// Nearest code under Hamming distance, ties broken toward the lower index.
fn nearest(codes: &[Topology], topology: &Topology) -> Result<usize, TopologyError> {
    let mut best = None;
    for (idx, code) in codes.iter().enumerate() {
        let d = code.hamming(topology)?;
        match best {
            None => best = Some((idx, d)),
            Some((_, bd)) if d < bd => best = Some((idx, d)),
            _ => {}
        }
    }
    best.map(|(idx, _)| idx).ok_or(TopologyError::EmptyCodebook)
}

/// Deterministic k-medoids over Hamming distance.
///
/// Farthest-first initialization (start from the key-lowest point, repeatedly
/// add the point farthest from the current set), then Lloyd iterations that
/// re-elect each cluster's medoid as the member minimizing total in-cluster
/// distance. Every tie — in assignment and in medoid election — breaks toward
/// the lower index of a key-sorted input, so the result is a pure function of
/// the input set.
fn k_medoids(
    points: &[Topology],
    k: usize,
    max_iterations: usize,
) -> Result<Vec<Topology>, TopologyError> {
    debug_assert!(points.len() > k, "caller handles the exact-index case");

    let mut medoids: Vec<usize> = vec![0];
    while medoids.len() < k {
        let mut best: Option<(usize, usize)> = None;
        for (i, p) in points.iter().enumerate() {
            if medoids.contains(&i) {
                continue;
            }
            let mut d_min = usize::MAX;
            for &m in &medoids {
                d_min = d_min.min(points[m].hamming(p)?);
            }
            if best.is_none_or(|(_, bd)| d_min > bd) {
                best = Some((i, d_min));
            }
        }
        match best {
            Some((i, _)) => medoids.push(i),
            None => break,
        }
    }
    medoids.sort_unstable();

    for _ in 0..max_iterations {
        // Assign every point to its nearest medoid.
        let mut clusters: Vec<Vec<usize>> = vec![Vec::new(); medoids.len()];
        for (i, p) in points.iter().enumerate() {
            let mut best: Option<(usize, usize)> = None;
            for (slot, &m) in medoids.iter().enumerate() {
                let d = points[m].hamming(p)?;
                if best.is_none_or(|(_, bd)| d < bd) {
                    best = Some((slot, d));
                }
            }
            clusters[best.expect("at least one medoid").0].push(i);
        }

        // Re-elect each medoid as the in-cluster point of minimum total
        // distance. An emptied cluster keeps its old medoid rather than being
        // re-seeded: a dropped code would shrink `k` silently, and
        // `used_codes` is supposed to report emptiness, not hide it.
        let mut next = medoids.clone();
        for (slot, members) in clusters.iter().enumerate() {
            if members.is_empty() {
                continue;
            }
            let mut best: Option<(usize, usize)> = None;
            for &cand in members {
                let mut total = 0usize;
                for &other in members {
                    total += points[cand].hamming(&points[other])?;
                }
                if best.is_none_or(|(_, bt)| total < bt) {
                    best = Some((cand, total));
                }
            }
            next[slot] = best.expect("non-empty cluster").0;
        }
        next.sort_unstable();
        next.dedup();
        if next == medoids {
            break;
        }
        medoids = next;
    }

    Ok(medoids.into_iter().map(|i| points[i].clone()).collect())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::record::ExecutionRecord;
    use crate::topology::CoordinationShape;

    fn set(records: Vec<ExecutionRecord>) -> RecordSet {
        RecordSet::new(records).unwrap()
    }

    fn rec(task: &str, t: Topology, u: f32, tokens: u64) -> ExecutionRecord {
        ExecutionRecord::new(task, vec![0.5, 0.5], t, u, tokens)
    }

    fn six_family_records(n: usize) -> Vec<ExecutionRecord> {
        let mut out = Vec::new();
        for (ti, task) in ["t1", "t2", "t3"].iter().enumerate() {
            for (fi, topo) in Topology::collection_protocol(n)
                .unwrap()
                .into_iter()
                .enumerate()
            {
                // Every family solves every task here, so all six survive.
                out.push(rec(task, topo, 1.0, 100 + (ti * 10 + fi) as u64));
            }
        }
        out
    }

    #[test]
    fn capacity_beyond_the_survivor_count_stays_idle() {
        let records = set(six_family_records(4));
        let mut used = Vec::new();
        for capacity in [8, 16, 32, 64] {
            let book = Codebook::fit(
                &records,
                &CodebookConfig {
                    capacity,
                    ..Default::default()
                },
            )
            .unwrap();
            used.push(book.used_codes());
        }
        // The paper's finding, on this fixture: the used-code count is the
        // number of distinct survivors, not the capacity.
        assert_eq!(used, vec![6, 6, 6, 6]);
    }

    #[test]
    fn only_reward_surviving_topologies_enter_the_codebook() {
        let n = 4;
        let records = set(vec![
            rec("t", Topology::complete(n).unwrap(), 1.0, 100),
            rec("t", Topology::chain(n).unwrap(), 0.0, 100),
            rec("t", Topology::star(n, 0).unwrap(), 1.0, 100),
        ]);
        let book = Codebook::fit(&records, &CodebookConfig::default()).unwrap();
        assert_eq!(book.len(), 2);
        assert!(book
            .codes()
            .iter()
            .all(|c| c != &Topology::chain(n).unwrap()));
    }

    #[test]
    fn a_set_with_no_survivors_is_an_error_not_an_empty_book() {
        let records = set(vec![rec("t", Topology::chain(4).unwrap(), 0.0, 100)]);
        assert!(matches!(
            Codebook::fit(&records, &CodebookConfig::default()),
            Err(TopologyError::NoRecords { kind: "surviving" })
        ));
    }

    #[test]
    fn fitting_is_deterministic_under_record_shuffling() {
        let mut a = six_family_records(4);
        let book_a = Codebook::fit(&set(a.clone()), &CodebookConfig::default()).unwrap();
        a.reverse();
        let book_b = Codebook::fit(&set(a), &CodebookConfig::default()).unwrap();
        assert_eq!(book_a.codes(), book_b.codes());
    }

    #[test]
    fn quantization_compresses_when_survivors_exceed_capacity() {
        // 12 distinct survivors, capacity 4.
        let n = 5;
        let mut records = Vec::new();
        for seed in 0..12u64 {
            records.push(rec(
                "t",
                Topology::erdos_renyi(n, 0.5, seed + 1).unwrap(),
                1.0,
                100,
            ));
        }
        let book = Codebook::fit(
            &set(records),
            &CodebookConfig {
                capacity: 4,
                ..Default::default()
            },
        )
        .unwrap();
        assert!(book.len() <= 4, "got {}", book.len());
        assert_eq!(book.distinct_survivors(), 12);
        assert_eq!(book.occupancy.iter().sum::<usize>(), 12);
    }

    #[test]
    fn every_code_decodes_to_a_topology_that_was_executed() {
        let n = 5;
        let mut records = Vec::new();
        let mut executed = Vec::new();
        for seed in 0..12u64 {
            let t = Topology::erdos_renyi(n, 0.5, seed + 1).unwrap();
            executed.push(t.clone());
            records.push(rec("t", t, 1.0, 100));
        }
        let book = Codebook::fit(
            &set(records),
            &CodebookConfig {
                capacity: 4,
                ..Default::default()
            },
        )
        .unwrap();
        for code in book.codes() {
            assert!(
                executed.iter().any(|e| e == code),
                "code {} was never executed",
                code.key()
            );
        }
    }

    #[test]
    fn encode_finds_the_nearest_code() {
        let n = 4;
        let book = Codebook::from_topologies(vec![
            Topology::empty(n).unwrap(),
            Topology::complete(n).unwrap(),
        ])
        .unwrap();
        let mut nearly_complete = Topology::complete(n).unwrap();
        nearly_complete.set_edge(0, 1, false).unwrap();
        assert_eq!(book.encode(&nearly_complete).unwrap(), 1);
        assert_eq!(book.encode(&Topology::chain(n).unwrap()).unwrap(), 0);
    }

    #[test]
    fn encode_rejects_a_mismatched_team_size() {
        let book = Codebook::from_topologies(vec![Topology::complete(4).unwrap()]).unwrap();
        assert!(matches!(
            book.encode(&Topology::complete(5).unwrap()),
            Err(TopologyError::SizeMismatch { .. })
        ));
    }

    #[test]
    fn cold_start_book_indexes_the_coordination_shapes() {
        let n = 4;
        let shapes: Vec<Topology> = CoordinationShape::ALL
            .iter()
            .map(|s| s.topology(n).unwrap())
            .collect();
        let book = Codebook::from_topologies(shapes).unwrap();
        // Solo and Swarm share the empty adjacency, so five shapes give four
        // codes — the dedup is the honest count, not a lost shape.
        assert_eq!(book.len(), 4);
    }

    #[test]
    fn zero_capacity_is_rejected() {
        let records = set(six_family_records(4));
        assert!(matches!(
            Codebook::fit(
                &records,
                &CodebookConfig {
                    capacity: 0,
                    ..Default::default()
                }
            ),
            Err(TopologyError::BadConfig { .. })
        ));
    }
}