data_beans/alg/

dc_poisson.rs

//! Degree-corrected Poisson refinement core.
//!
//! Shared Poisson-scoring machinery for top-down multi-level refinement.
//! Given per-entity sparse profiles (gene sums or a dense projection) and
//! current group labels, refinement reassigns entities between groups by
//! scoring moves under a Poisson likelihood with size-factor offsets.
//! Candidate-set construction is delegated to [`CandidateProposer`] so
//! different front-ends (cross-batch BBKNN, spatial KNN graph, ...) can
//! plug in their own neighborhood definition.
//!
//! The name "DC-Poisson" is deliberate: we are *not* inferring a stochastic
//! block model — the blocks are given externally (by a prior coarsening
//! step) and we only score membership moves. Compare to the full DC-SBM
//! where blocks themselves are latent variables.
//!
//! Entry points:
//! - [`refine_with_candidates`] — lowest-level sweep driver given pre-built candidates.
//! - [`refine_with_proposer`] — generic driver that first asks a [`CandidateProposer`].
//!
//! See the `BbknnProposer` in `refine_multilevel` for the data_beans::alg
//! front-end, and pinto's `GraphProposer` for the spatial-graph front-end.

use log::info;
use nalgebra::DMatrix;
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::{RngExt, SeedableRng};
use rayon::prelude::*;
use rustc_hash::FxHashMap as HashMap;

/// Additive floor to keep `ln(.)` finite when a block or feature has zero
/// mass.
pub const LOG_EPS: f64 = 1e-9;

/// `ln(x)` in `f32`: libm's single-precision log, near half the cost of
/// the `f64` one. Its error is `f32` rounding (about 1e-7 relative), the
/// precision the cached block logs are stored at anyway. NaN, infinite and
/// non-positive inputs keep their usual `ln` values.
#[inline]
fn fast_ln(x: f64) -> f64 {
    f64::from((x as f32).ln())
}

//////////////////////////
// Public configuration //
//////////////////////////

/// Feature representation used to score entity → block moves.
#[derive(Clone, Debug)]
pub enum ProfileSource {
    /// Sparse entity × feature counts (the existing `gene_sums` format —
    /// the field name is historical; the axis is generic).
    /// [`FeatureWeighting`] is applied per-feature on the sparse profile.
    Raw,
    /// Dense projection: `basis * indicator(entity)` per entity. `basis` is
    /// `proj_dim × num_cells`-shaped; the per-entity profile is the sum over
    /// its member cells' projection columns. [`FeatureWeighting`] is skipped
    /// because the feature axis is no longer count-distributed.
    Projected { basis: DMatrix<f32> },
}

/// Per-feature weighting applied to the sparse profile before DC-Poisson
/// scoring. All variants multiply each nonzero entry `y_{e,f}` by a scalar
/// `w_f` and recompute per-row size factors; only the formula for `w_f`
/// differs.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FeatureWeighting {
    /// No reweighting: `w_f = 1`. Recovers proper DC-Poisson with entity-level
    /// degree correction (Karrer–Newman MAP under `Gamma(α, 0)` on `θ`).
    None,
    /// Fisher-information weight under the fitted NB trend (genes-as-features
    /// case): `w_f = 1 / (1 + π_f · s̄ · φ(μ_f))`. Bounded in `(0, 1]`,
    /// attenuates high-mean / high-dispersion features, recovers
    /// `w_f = 1` in the Poisson limit (`φ → 0`).
    FisherInfoNb,
}

/// Parameters controlling the refinement pass.
#[derive(Clone, Debug)]
pub struct RefineParams {
    /// Gibbs sweeps per level (0 disables Gibbs; greedy still runs).
    pub num_gibbs: usize,
    /// Greedy sweeps per level (early-exits on zero moves).
    pub num_greedy: usize,
    /// Per-feature weighting scheme (only meaningful for `Raw` profile source).
    pub feature_weighting: FeatureWeighting,
    /// Seed for Gibbs RNG.
    pub seed: u64,
    /// Gibbs stagnation threshold (fraction of entities moving per sweep).
    /// Breaks early when three consecutive sweeps are below this bound.
    /// `0.0` disables early exit.
    pub gibbs_stagnation: f64,
    /// Feature source. Defaults to `Raw`.
    pub profile_source: ProfileSource,
    /// Jacobi-style parallel sweeps (default `true`).
    ///
    /// When `true`, every entity proposes its move against a frozen pre-sweep
    /// snapshot of the sufficient stats (rayon `par_iter`); proposals are
    /// then applied sequentially with the move guard. Per-entity RNG seeds
    /// are derived from `(sweep_seed, entity_idx)`, so the result is
    /// deterministic given `seed` and independent of thread count (per
    /// `rand` release — `SmallRng`'s stream is not guaranteed stable across
    /// `rand` versions).
    ///
    /// When `false`, runs the classic Gauss–Seidel loop: each entity scores
    /// against the live stats (updated by every accepted move so far in the
    /// sweep), single-threaded.
    ///
    /// Jacobi is biased — within a sweep, entities can't see each other's
    /// moves — but converges fine for the multilevel DC-Poisson use case
    /// and is roughly P× faster on large pb-sample counts.
    pub parallel: bool,
}

impl Default for RefineParams {
    fn default() -> Self {
        Self {
            num_gibbs: 20,
            num_greedy: 10,
            feature_weighting: FeatureWeighting::FisherInfoNb,
            seed: 42,
            gibbs_stagnation: 0.005,
            profile_source: ProfileSource::Raw,
            parallel: true,
        }
    }
}

////////////////////////////////////
// Sparse row-per-entity profiles //
////////////////////////////////////

/// Sparse row-per-entity profile, owned by the refinement pass.
///
/// Materialized from either a sparse gene-sum format (`from_gene_sums`) or a
/// projected centroid matrix (`from_projection`). Values are stored sorted
/// by feature index within each row so downstream accumulators can rely on
/// the invariant.
pub struct Profiles {
    pub rows: Vec<Vec<(u32, f32)>>,
    pub size_factor: Vec<f32>,
    pub num_entities: usize,
    pub num_features: usize,
}

impl Profiles {
    pub fn from_gene_sums(gene_sums: &[Vec<(usize, f32)>], num_features: usize) -> Self {
        let num_entities = gene_sums.len();
        let (rows, size_factor): (Vec<Vec<(u32, f32)>>, Vec<f32>) = gene_sums
            .par_iter()
            .map(|row| {
                let mut out: Vec<(u32, f32)> = row
                    .iter()
                    .filter(|(_, v)| *v > 0.0)
                    .map(|(g, v)| (*g as u32, *v))
                    .collect();
                out.sort_unstable_by_key(|&(g, _)| g);
                let sf: f32 = out.iter().map(|(_, v)| *v).sum();
                (out, sf)
            })
            .unzip();
        Self {
            rows,
            size_factor,
            num_entities,
            num_features,
        }
    }

    /// Build profiles from a dense projection.
    ///
    /// `basis` is `proj_dim × num_cells`; each entity profile is the sum of
    /// basis columns over the cells constituting it. Gene weighting is not
    /// applied — projection dims aren't gene-aligned.
    ///
    /// The Poisson score needs nonnegative profiles, so the accumulated
    /// per-dimension sums must be `>= 0` — a signed basis (e.g. a raw random
    /// projection) is rejected loudly rather than silently truncated.
    pub fn from_projection(basis: &DMatrix<f32>, entity_to_cells: &[Vec<usize>]) -> Self {
        let num_features = basis.nrows();
        let num_entities = entity_to_cells.len();
        let (rows, size_factor): (Vec<Vec<(u32, f32)>>, Vec<f32>) = entity_to_cells
            .par_iter()
            .map(|cells| {
                let mut acc = vec![0f32; num_features];
                for &c in cells {
                    for d in 0..num_features {
                        acc[d] += basis[(d, c)];
                    }
                }
                let mut out = Vec::with_capacity(num_features);
                let mut sf = 0f32;
                for (d, &v) in acc.iter().enumerate() {
                    assert!(
                        v >= 0.0,
                        "projection profile is negative at dim {d} ({v}); \
                         DC-Poisson requires a nonnegative basis"
                    );
                    if v > 0.0 {
                        out.push((d as u32, v));
                        sf += v;
                    }
                }
                (out, sf)
            })
            .unzip();
        Self {
            rows,
            size_factor,
            num_entities,
            num_features,
        }
    }

    /// In-place reweighting by a caller-supplied per-feature weight vector.
    /// Used by [`Profiles::apply_feature_weighting`] for the NB Fisher-info path.
    pub fn weight_by_vec(&mut self, w: &[f32]) {
        assert_eq!(
            w.len(),
            self.num_features,
            "weight vector length must match num_features"
        );
        self.rows
            .par_iter_mut()
            .zip(self.size_factor.par_iter_mut())
            .for_each(|(row, sf)| {
                let mut new_sf = 0f32;
                for (g, v) in row.iter_mut() {
                    *v *= w[*g as usize];
                    new_sf += *v;
                }
                *sf = new_sf;
            });
    }

    /// Apply the chosen feature weighting in place.
    ///
    /// [`FeatureWeighting::None`] is a no-op; [`FeatureWeighting::FisherInfoNb`]
    /// fits an NB dispersion trend from the current profiles and reweights each
    /// feature by `1 / (1 + π_f · s̄ · φ(μ_f))`.
    pub fn apply_feature_weighting(&mut self, method: FeatureWeighting) {
        match method {
            FeatureWeighting::None => {}
            FeatureWeighting::FisherInfoNb => {
                let w = self.nb_fisher_weights();
                self.weight_by_vec(&w);
            }
        }
    }

    /// Compute per-feature Fisher-info weights under an NB trend fit from
    /// the current profile contents. Returned vector has length `num_features`
    /// and is suitable for [`Profiles::weight_by_vec`].
    pub fn nb_fisher_weights(&self) -> Vec<f32> {
        use crate::alg::nb_dispersion::DispersionTrend;
        use legume_numeric::matrix::sparse_stat::SparseRunningStatistics;
        use legume_numeric::matrix::traits::RunningStatOps;

        // Rows are folded in parallel and merged in a fixed order, so the
        // (non-integer) sums come out the same on any number of threads.
        // Avoids the single-threaded `for row in &self.rows` over ~10⁴–10⁵
        // pb-samples.
        const ROWS_PER_LEAF: usize = 256;
        let num_features = self.num_features;
        let (stats, ..) = crate::qc::fold_in_fixed_order(
            &self.rows,
            ROWS_PER_LEAF,
            &|| {
                (
                    SparseRunningStatistics::<f32>::new(num_features),
                    Vec::<usize>::new(),
                    Vec::<f32>::new(),
                )
            },
            &|(acc, col_rows, col_vals), row| {
                col_rows.clear();
                col_vals.clear();
                for &(g, v) in row {
                    col_rows.push(g as usize);
                    col_vals.push(v);
                }
                acc.add_sparse_column(col_rows, col_vals);
                Ok(())
            },
            &|(a, ..), (b, ..)| a.merge(&b),
        )
        .expect("the row fold cannot fail");

        // `π_g = sum[g] / Σ sum` is derived directly from the stats without
        // a second sparse traversal. `mean_g` uses the entity-count denominator
        // that `SparseRunningStatistics` already provides.
        let trend = DispersionTrend::from_sparse_stats(&stats);
        let means = stats.mean();
        let sums = stats.sum();
        let total_mass: f64 = sums.iter().map(|&s| s as f64).sum();
        let avg_s = if self.num_entities > 0 {
            (total_mass / self.num_entities as f64) as f32
        } else {
            1.0
        };
        let inv_total = if total_mass > 0.0 {
            1.0 / total_mass as f32
        } else {
            0.0
        };
        (0..self.num_features)
            .map(|g| trend.fisher_weight(sums[g] * inv_total, avg_s, means[g]))
            .collect()
    }
}

///////////////////////////////////
// Poisson sufficient statistics //
///////////////////////////////////

/// Sufficient statistics with cached log quantities for O(K · nnz(row)) scoring.
///
/// `gene_sum[g·K + k]` = Σ_{e : z_e = k} y_eg; `size_sum[k]` = Σ_{e : z_e = k} s_e.
/// The per-feature tables are laid out **feature-major** (`g·K + k`): scoring
/// one entity walks its nonzero features once and, for each, reads the `K`
/// group values side by side, instead of striding across a `K × M` table
/// once per candidate group. Use [`DcPoissonStats::idx`] to address them.
/// Log caches score every non-current block with no log; the current block's
/// leave-one-out score takes one `fast_ln` per nonzero feature, and
/// `delta_move` refreshes the logs of the columns a move touches.
///
/// Log caches are `f32` — values live in roughly `[-20, +25]` (dominated by
/// `-ln(LOG_EPS)` at the floor), and scoring is noise-dominated so f32's
/// ~7-digit precision is ample. Halves the largest allocation versus f64.
#[derive(Clone)]
pub struct DcPoissonStats {
    pub k: usize,
    pub num_features: usize,
    pub membership: Vec<usize>,
    /// Feature-major `M × K`: entry `g·K + k`.
    pub gene_sum: Vec<f64>,
    pub size_sum: Vec<f64>,
    /// Feature-major `M × K`: entry `g·K + k`.
    pub log_gene: Vec<f32>,
    pub log_size_offset: Vec<f32>,
}

impl DcPoissonStats {
    /// Index of (group `k`, feature `g`) in `gene_sum` / `log_gene`.
    #[inline(always)]
    pub fn idx(&self, k: usize, g: usize) -> usize {
        g * self.k + k
    }

    pub fn from_profiles(profiles: &Profiles, k: usize, membership: &[usize]) -> Self {
        let m = profiles.num_features;
        let mut gene_sum = vec![0f64; k * m];
        let mut size_sum = vec![0f64; k];
        for (e, row) in profiles.rows.iter().enumerate() {
            let z = membership[e];
            assert!(z < k, "membership[{}]={} out of range 0..{}", e, z, k);
            for &(g, v) in row {
                gene_sum[g as usize * k + z] += v as f64;
            }
            size_sum[z] += profiles.size_factor[e] as f64;
        }
        let log_gene: Vec<f32> = gene_sum
            .iter()
            .map(|&s| fast_ln(s + LOG_EPS) as f32)
            .collect();
        let m_eps = m as f64 * LOG_EPS;
        let log_size_offset: Vec<f32> = size_sum
            .iter()
            .map(|&s| -fast_ln(s + m_eps) as f32)
            .collect();
        Self {
            k,
            num_features: m,
            membership: membership.to_vec(),
            gene_sum,
            size_sum,
            log_gene,
            log_size_offset,
        }
    }

    /// Apply an entity's reassignment and refresh only the affected log entries.
    pub fn delta_move(&mut self, e: usize, k_from: usize, k_to: usize, profiles: &Profiles) {
        if k_from == k_to {
            return;
        }
        let kk = self.k;
        let m = self.num_features;
        let m_eps = m as f64 * LOG_EPS;

        // Subtractions clamp at zero: the running sums accumulate rounding
        // residues, and a block drained back to (true) zero can otherwise be
        // left with a negative residue larger than LOG_EPS — feeding
        // `ln(negative) = NaN` into the log caches.
        for &(g, v) in &profiles.rows[e] {
            let base = g as usize * kk;
            let (i_from, i_to) = (base + k_from, base + k_to);
            self.gene_sum[i_from] = (self.gene_sum[i_from] - v as f64).max(0.0);
            self.gene_sum[i_to] += v as f64;
            self.log_gene[i_from] = fast_ln(self.gene_sum[i_from] + LOG_EPS) as f32;
            self.log_gene[i_to] = fast_ln(self.gene_sum[i_to] + LOG_EPS) as f32;
        }
        let sf = profiles.size_factor[e] as f64;
        self.size_sum[k_from] = (self.size_sum[k_from] - sf).max(0.0);
        self.size_sum[k_to] += sf;
        self.log_size_offset[k_from] = -fast_ln(self.size_sum[k_from] + m_eps) as f32;
        self.log_size_offset[k_to] = -fast_ln(self.size_sum[k_to] + m_eps) as f32;

        self.membership[e] = k_to;
    }

    /// Full recompute from current `membership` (slow path, used for tests).
    #[cfg(test)]
    fn recompute(&mut self, profiles: &Profiles) {
        let k = self.k;
        let m = self.num_features;
        self.gene_sum.iter_mut().for_each(|x| *x = 0.0);
        self.size_sum.iter_mut().for_each(|x| *x = 0.0);
        for (e, row) in profiles.rows.iter().enumerate() {
            let z = self.membership[e];
            for &(g, v) in row {
                self.gene_sum[g as usize * k + z] += v as f64;
            }
            self.size_sum[z] += profiles.size_factor[e] as f64;
        }
        for i in 0..k * m {
            self.log_gene[i] = fast_ln(self.gene_sum[i] + LOG_EPS) as f32;
        }
        let m_eps = m as f64 * LOG_EPS;
        for i in 0..k {
            self.log_size_offset[i] = -fast_ln(self.size_sum[i] + m_eps) as f32;
        }
    }
}

/////////////////////
// Scoring kernels //
/////////////////////

/// Score entity `e` for one destination `k` (Poisson plug-in MAP, up to a
/// common constant):
/// `  s(e, k) = Σ_{g : y_eg > 0} y_eg · ln(gene_sum⁻ᵉ[k, g] + ε) − size_factor[e] · ln(size_sum⁻ᵉ[k] + Mε)`
///
/// where `⁻ᵉ` means entity `e`'s own contribution is excluded. Candidate
/// blocks never contain `e`, so they read straight from the log caches; the
/// entity's *current* block subtracts `e`'s row before taking logs
/// (leave-one-out). Scoring the current block with `e` included would give
/// "stay put" a self-inclusion bonus that grows with the entity's size
/// factor and shrinks with block mass — anchoring large entities in place.
///
/// Reference for [`compute_log_probs_restricted`], which scores all allowed
/// blocks in one pass with the same per-block summation order (so the two
/// agree bit for bit).
#[cfg(test)]
fn score_move(e: usize, k: usize, stats: &DcPoissonStats, profiles: &Profiles) -> f64 {
    let m = stats.num_features;
    let sf = profiles.size_factor[e] as f64;
    let row = &profiles.rows[e];
    if stats.membership[e] == k {
        let m_eps = m as f64 * LOG_EPS;
        let loo_size = (stats.size_sum[k] - sf).max(0.0);
        let mut acc = -sf * fast_ln(loo_size + m_eps);
        for &(g, v) in row {
            let vg = v as f64;
            let loo = (stats.gene_sum[stats.idx(k, g as usize)] - vg).max(0.0);
            acc += vg * fast_ln(loo + LOG_EPS);
        }
        acc
    } else {
        let mut acc = sf * stats.log_size_offset[k] as f64;
        for &(g, v) in row {
            acc += v as f64 * stats.log_gene[stats.idx(k, g as usize)] as f64;
        }
        acc
    }
}

/// Log-probability of placing entity `e` into each of the `allowed` blocks.
///
/// `  s(e, k) = Σ_{g : y_eg > 0} y_eg · ln(gene_sum⁻ᵉ[k, g] + ε) − size_factor[e] · ln(size_sum⁻ᵉ[k] + Mε)`
///
/// `⁻ᵉ` means entity `e`'s own contribution is excluded: candidate blocks
/// never contain `e` and read the log caches, while `e`'s *current* block
/// subtracts `e`'s row before taking logs (leave-one-out) — otherwise
/// "stay put" earns a self-inclusion bonus that anchors large entities.
///
/// All blocks are scored in one pass over `e`'s nonzero features: per
/// feature, the `K` cached values sit side by side (feature-major layout),
/// so the pass touches one short contiguous run per feature. Each block's
/// sum is still accumulated in feature order, the same order as scoring the
/// blocks one at a time, so the scores match that bit for bit. Partial
/// candidate lists accumulate by *position* in `allowed`, so a block listed
/// twice gets the same, correct score in both places.
///
/// Only the `allowed` slots of `log_probs` are written, unless `allowed`
/// is as long as `K`, in which case the first `K` slots are written.
/// The rest keep whatever stale values a previous call left (the caller
/// reuses one buffer across sweeps). Downstream consumers
/// ([`sample_categorical_log_restricted`], [`argmax_log_restricted`]) read
/// the `allowed` slots exclusively.
pub fn compute_log_probs_restricted(
    e: usize,
    stats: &DcPoissonStats,
    profiles: &Profiles,
    allowed: &[usize],
    log_probs: &mut [f64],
) {
    let kk = stats.k;
    let m = stats.num_features;
    let sf = profiles.size_factor[e] as f64;
    let row = &profiles.rows[e];
    let cur = stats.membership[e];
    // Every block 0..K is scored straight into `log_probs`, so repeats in
    // `allowed` (which then misses some block) are harmless.
    let dense = allowed.len() == kk;
    let score_cur = dense || allowed.contains(&cur);

    // Leave-one-out score of the current block.
    let mut cur_acc = if score_cur {
        let loo_size = (stats.size_sum[cur] - sf).max(0.0);
        -sf * fast_ln(loo_size + m as f64 * LOG_EPS)
    } else {
        0.0
    };
    let mut add_cur = |base: usize, vg: f64| {
        if score_cur {
            let loo = (stats.gene_sum[base + cur] - vg).max(0.0);
            cur_acc += vg * fast_ln(loo + LOG_EPS);
        }
    };

    // Cached (non-leave-one-out) scores for every allowed block, including
    // the current one, which is overwritten at the end.
    if dense {
        for (slot, &o) in log_probs[..kk].iter_mut().zip(&stats.log_size_offset) {
            *slot = sf * o as f64;
        }
        for &(g, v) in row {
            let vg = v as f64;
            let base = g as usize * kk;
            let logs = &stats.log_gene[base..base + kk];
            for (slot, &l) in log_probs[..kk].iter_mut().zip(logs) {
                *slot += vg * l as f64;
            }
            add_cur(base, vg);
        }
    } else {
        thread_local! {
            static SCRATCH: std::cell::RefCell<Vec<f64>> = const { std::cell::RefCell::new(Vec::new()) };
        }
        SCRATCH.with(|scratch| {
            let mut scratch = scratch.borrow_mut();
            scratch.clear();
            scratch.extend(
                allowed
                    .iter()
                    .map(|&k| sf * stats.log_size_offset[k] as f64),
            );
            for &(g, v) in row {
                let vg = v as f64;
                let base = g as usize * kk;
                let logs = &stats.log_gene[base..base + kk];
                for (acc, &k) in scratch.iter_mut().zip(allowed) {
                    *acc += vg * logs[k] as f64;
                }
                add_cur(base, vg);
            }
            for (&acc, &k) in scratch.iter().zip(allowed) {
                log_probs[k] = acc;
            }
        });
    }
    if score_cur {
        log_probs[cur] = cur_acc;
    }
}

/// Unrestricted variant used only from tests.
#[cfg(test)]
fn compute_log_probs(e: usize, stats: &DcPoissonStats, profiles: &Profiles, log_probs: &mut [f64]) {
    for (k, slot) in log_probs.iter_mut().enumerate().take(stats.k) {
        *slot = score_move(e, k, stats, profiles);
    }
}

/// Gumbel-max categorical sampling over the `allowed` slots of `log_probs`,
/// skipping non-finite entries. Falls back to `current` (the entity's
/// present label — always a legal "move") when no allowed slot is finite,
/// so a degenerate score vector can never produce an out-of-range label.
pub fn sample_categorical_log_restricted(
    log_probs: &[f64],
    allowed: &[usize],
    current: usize,
    rng: &mut SmallRng,
) -> usize {
    let mut best_key = f64::NEG_INFINITY;
    let mut best_idx = current;
    for &k in allowed {
        let lp = log_probs[k];
        if lp.is_finite() {
            let u: f64 = rng.random_range(1e-12..1.0_f64);
            let g = -(-u.ln()).ln();
            let key = lp + g;
            if key > best_key {
                best_key = key;
                best_idx = k;
            }
        }
    }
    best_idx
}

/// Argmax over `allowed` indices.
pub fn argmax_log_restricted(log_probs: &[f64], allowed: &[usize]) -> usize {
    let mut best = allowed[0];
    let mut best_val = log_probs[best];
    for &k in &allowed[1..] {
        if log_probs[k] > best_val {
            best = k;
            best_val = log_probs[k];
        }
    }
    best
}

/// Argmax over `allowed` indices that leaves `current` only for a strictly
/// better block: a tie with the current label (an entity with no evidence
/// scores every block alike) is not a move. With `current` outside
/// `allowed`, ties go to the first allowed index.
pub fn argmax_log_restricted_from(log_probs: &[f64], allowed: &[usize], current: usize) -> usize {
    let mut best = if allowed.contains(&current) {
        current
    } else {
        allowed[0]
    };
    let mut best_val = log_probs[best];
    for &k in allowed {
        if log_probs[k] > best_val {
            best = k;
            best_val = log_probs[k];
        }
    }
    best
}

////////////////////////////////////
// Hierarchy / relabeling helpers //
////////////////////////////////////

/// Compact a label vector to dense `0..K` ids, preserving relative order of
/// first appearance. Returns the new labels and the new `K`. Generic over
/// the key type so it serves both plain group ids (`usize`) and composite
/// keys (e.g. `(child, parent)` pairs in `project_to_refinement`).
pub fn compact_labels<K>(labels: &[K]) -> (Vec<usize>, usize)
where
    K: Copy + Eq + std::hash::Hash,
{
    let mut map: HashMap<K, usize> = HashMap::default();
    let mut next = 0usize;
    let mut out = Vec::with_capacity(labels.len());
    for &g in labels {
        let new = *map.entry(g).or_insert_with(|| {
            let id = next;
            next += 1;
            id
        });
        out.push(new);
    }
    (out, next)
}

/// For each entity, siblings at `level` = groups at `level` sharing the
/// same parent group at `level + 1`.
///
/// `refined` follows the finest → coarsest convention: `refined[0]` is the
/// finest level, `refined[refined.len() - 1]` is the coarsest. At the
/// coarsest level (index `num_levels - 1`), the virtual root groups every
/// entity together → siblings = all groups.
pub fn compute_sibling_sets(
    refined: &[Vec<usize>],
    level: usize,
    num_groups_at_level: usize,
) -> Vec<Vec<usize>> {
    let num_levels = refined.len();
    let num_entities = refined[level].len();

    if level + 1 >= num_levels {
        let all_groups: Vec<usize> = (0..num_groups_at_level).collect();
        return vec![all_groups; num_entities];
    }

    let mut parent_to_children: HashMap<usize, Vec<usize>> = HashMap::default();
    for (child, parent) in refined[level].iter().zip(refined[level + 1].iter()) {
        let entry = parent_to_children.entry(*parent).or_default();
        if !entry.contains(child) {
            entry.push(*child);
        }
    }
    for v in parent_to_children.values_mut() {
        v.sort_unstable();
    }

    (0..num_entities)
        .map(|e| {
            let parent = refined[level + 1][e];
            parent_to_children.get(&parent).cloned().unwrap_or_default()
        })
        .collect()
}

////////////////////////////////////////////
// Candidate proposer + move guard traits //
////////////////////////////////////////////

/// Propose candidate group labels each entity may move into on one sweep.
///
/// The returned slice `candidates[e]` MUST include `labels[e]` (staying put
/// is always legal) and MUST index into the current 0..k label space.
pub trait CandidateProposer {
    fn propose(&self, labels: &[usize]) -> Vec<Vec<usize>>;
}

/// Veto individual moves after they've been picked by the sweep.
///
/// Called once per *proposed* accepted move (`from != to`) before the move
/// is applied to the sufficient statistics. The destination label is
/// intentionally not passed — guards so far only need the source cluster
/// (e.g. articulation tests ask "does removing this entity disconnect
/// `from`?"). Add a destination-aware variant if a concrete use case
/// surfaces.
pub trait MoveGuard {
    fn accept_move(&self, entity: usize, from: usize, labels: &[usize]) -> bool;
}

/// Trivial guard that accepts every move. Zero-cost default used by the
/// unguarded [`refine_with_candidates`] and [`refine_with_proposer`].
pub struct NoGuard;

impl MoveGuard for NoGuard {
    #[inline(always)]
    fn accept_move(&self, _e: usize, _from: usize, _labels: &[usize]) -> bool {
        true
    }
}

/// Intersect a per-entity sibling set with a per-entity neighbor-group set,
/// with sibling fallback when the intersection is empty and guaranteed
/// inclusion of the entity's current label so staying put is always a
/// legal move.
///
/// Shared between [`crate::alg::refine_multilevel::BbknnProposer`] and pinto's
/// `GraphProposer` — both proposers differ only in how they gather
/// `neighbor_groups`, not in how they combine with siblings.
///
/// * `siblings` — sorted, deduped sibling group list.
/// * `neighbor_groups` — sorted, deduped group list from the proposer's
///   spatial/structural neighborhood.
/// * `current` — the entity's current label; appended if the intersection
///   excluded it.
pub fn intersect_with_siblings_fallback(
    siblings: &[usize],
    neighbor_groups: &[usize],
    current: usize,
) -> Vec<usize> {
    if siblings.is_empty() {
        return Vec::new();
    }
    if siblings.len() == 1 {
        return siblings.to_vec();
    }
    let intersect: Vec<usize> = siblings
        .iter()
        .copied()
        .filter(|g| neighbor_groups.binary_search(g).is_ok())
        .collect();
    if intersect.is_empty() {
        return siblings.to_vec();
    }
    if intersect.contains(&current) {
        intersect
    } else {
        let mut c = intersect;
        c.push(current);
        c.sort_unstable();
        c
    }
}

///////////////////
// Sweep drivers //
///////////////////

/// Refine `labels` in place for one level by running Gibbs + greedy sweeps
/// over `profiles` with the given pre-built `candidates` and a custom
/// [`MoveGuard`]. Returns total accepted moves across all sweeps.
///
/// A proposed move `from → to` is applied only when `guard.accept_move(...)`
/// returns `true`. In Gibbs sweeps, a vetoed move leaves the entity in its
/// current group for this sweep. Greedy sweeps likewise skip the move.
///
/// Shared sweep-driver context bundling the readonly inputs that every
/// `refine_with_*` variant needs. Kept as a struct so the driver fns don't
/// exceed clippy's argument-count threshold.
pub struct RefineContext<'a> {
    pub profiles: &'a Profiles,
    pub k: usize,
    pub params: &'a RefineParams,
    pub level_label: &'a str,
}

/// Outcome of one sweep over all entities.
#[derive(Default, Clone, Copy)]
struct SweepCounts {
    moves: usize,
    vetoed: usize,
}

/// Apply a precomputed `proposals` vector to `stats` sequentially, consulting
/// the move guard. Shared by both Jacobi sweeps after the parallel proposal
/// phase. The guard sees the membership-being-updated, matching the
/// sequential semantics.
fn apply_proposals<G: MoveGuard>(
    proposals: &[usize],
    stats: &mut DcPoissonStats,
    profiles: &Profiles,
    guard: &G,
) -> SweepCounts {
    let mut sc = SweepCounts::default();
    for (e, &new) in proposals.iter().enumerate() {
        let old = stats.membership[e];
        if new == old {
            continue;
        }
        if guard.accept_move(e, old, &stats.membership) {
            stats.delta_move(e, old, new, profiles);
            sc.moves += 1;
        } else {
            sc.vetoed += 1;
        }
    }
    sc
}

/// Gauss–Seidel sequential sweep. For each entity visited in `order`, scores
/// the candidate destinations under the live (mutated) stats, picks one via
/// `pick(log_probs, candidates, current_label)`, and applies the move through
/// the guard. Gibbs callers shuffle `order` and pass a categorical-sampling
/// closure; greedy callers pass `0..n` and an argmax closure.
fn sweep_sequential<G, F, I>(
    candidates: &[Vec<usize>],
    stats: &mut DcPoissonStats,
    profiles: &Profiles,
    guard: &G,
    log_probs: &mut [f64],
    order: I,
    mut pick: F,
) -> SweepCounts
where
    G: MoveGuard,
    F: FnMut(&[f64], &[usize], usize) -> usize,
    I: IntoIterator<Item = usize>,
{
    let mut sc = SweepCounts::default();
    for e in order {
        let cand = &candidates[e];
        if cand.len() < 2 {
            continue;
        }
        compute_log_probs_restricted(e, stats, profiles, cand, log_probs);
        let old = stats.membership[e];
        let new = pick(log_probs, cand, old);
        if new == old {
            continue;
        }
        if guard.accept_move(e, old, &stats.membership) {
            stats.delta_move(e, old, new, profiles);
            sc.moves += 1;
        } else {
            sc.vetoed += 1;
        }
    }
    sc
}

/// Jacobi sweep: every entity computes its proposal against a frozen pre-sweep
/// snapshot of `stats` (rayon `par_iter`); proposals are then applied
/// sequentially with the move guard.
///
/// `pick(log_probs, cand, e, current)` chooses the destination given the
/// restricted log-likelihoods. For Gibbs it derives a per-entity RNG from
/// `(sweep_seed, e)` and samples categorically; for greedy it ignores `e`
/// and argmaxes.
/// `proposals` is a reusable scratch buffer of length `candidates.len()` —
/// hoisted by the driver across sweeps to avoid re-allocating.
fn sweep_jacobi<G: MoveGuard, F>(
    candidates: &[Vec<usize>],
    stats: &mut DcPoissonStats,
    profiles: &Profiles,
    guard: &G,
    k: usize,
    proposals: &mut [usize],
    pick: F,
) -> SweepCounts
where
    F: Fn(&[f64], &[usize], usize, usize) -> usize + Sync,
{
    debug_assert_eq!(proposals.len(), candidates.len());
    {
        // Immutable reborrow so the rayon closure can capture `&stats: Sync`
        // while the outer `&mut` is held across the propose / apply phases.
        let stats: &DcPoissonStats = stats;
        proposals
            .par_iter_mut()
            .enumerate()
            .with_min_len(256)
            .for_each_init(
                || vec![f64::NEG_INFINITY; k],
                |log_probs, (e, prop)| {
                    let cand = &candidates[e];
                    let current = stats.membership[e];
                    if cand.len() < 2 {
                        *prop = current;
                        return;
                    }
                    compute_log_probs_restricted(e, stats, profiles, cand, log_probs);
                    *prop = pick(log_probs, cand, e, current);
                },
            );
    }
    apply_proposals(proposals, stats, profiles, guard)
}

/// This is the lowest-level driver. Pass [`NoGuard`] to get the unguarded
/// behavior; most callers should prefer the convenience wrappers
/// [`refine_with_candidates`] / [`refine_with_proposer`] /
/// [`refine_with_proposer_guarded`].
///
/// Sweep dispatch follows `ctx.params.parallel`: Jacobi (parallel proposal,
/// sequential apply) when true, Gauss–Seidel (single-threaded) when false.
/// See [`RefineParams::parallel`] for trade-offs.
pub fn refine_with_candidates_guarded<G: MoveGuard>(
    labels: &mut [usize],
    candidates: &[Vec<usize>],
    guard: &G,
    rng: &mut SmallRng,
    ctx: &RefineContext,
) -> usize {
    let RefineContext {
        profiles,
        k,
        params,
        level_label,
    } = *ctx;
    let num_entities = labels.len();
    let mut stats = DcPoissonStats::from_profiles(profiles, k, labels);
    let mut log_probs = vec![f64::NEG_INFINITY; k];
    let mut total_moves = 0usize;
    let mut total_vetoed = 0usize;

    // Reusable proposal scratch for Jacobi sweeps — hoisted so the
    // num_sweeps × num_entities allocation isn't repeated each sweep.
    let mut proposals: Vec<usize> = if params.parallel {
        vec![0usize; num_entities]
    } else {
        Vec::new()
    };

    // Reusable shuffled-order buffer for sequential Gibbs sweeps. Greedy
    // sequential iterates `0..num_entities` directly and ignores this.
    let mut order: Vec<usize> = if !params.parallel {
        (0..num_entities).collect()
    } else {
        Vec::new()
    };

    let max_sweeps = (params.num_gibbs + params.num_greedy) as u64;
    let prog_bar = legume_numeric::matrix::progress::new_progress_bar(max_sweeps)
        .with_message(format!("{level_label} sweeps"));
    // Refresh on a timer so the bar visibly animates between sweep ticks
    // (a single sweep at large D can take several seconds — without this
    // the bar is silent until the next inc(1)).
    prog_bar.enable_steady_tick(std::time::Duration::from_millis(100));

    // Jacobi per-sweep seed base: draw a single odd u64 from `rng` so the
    // sweep loop is reproducible given `params.seed` regardless of how many
    // cells `shuffle` would have consumed in the sequential path.
    let jacobi_base_seed = rng.random::<u64>() | 1;

    if params.num_gibbs > 0 {
        let mut low_sweeps = 0usize;
        for sweep in 0..params.num_gibbs {
            let sc = if params.parallel {
                let sweep_seed = jacobi_base_seed.wrapping_mul(sweep as u64 + 1);
                sweep_jacobi(
                    candidates,
                    &mut stats,
                    profiles,
                    guard,
                    k,
                    &mut proposals,
                    |log_probs, cand, e, current| {
                        let vertex_seed = sweep_seed ^ (e as u64).wrapping_mul(2654435761);
                        let mut rng = SmallRng::seed_from_u64(vertex_seed);
                        sample_categorical_log_restricted(log_probs, cand, current, &mut rng)
                    },
                )
            } else {
                order.shuffle(rng);
                sweep_sequential(
                    candidates,
                    &mut stats,
                    profiles,
                    guard,
                    &mut log_probs,
                    order.iter().copied(),
                    |log_probs, cand, current| {
                        sample_categorical_log_restricted(log_probs, cand, current, rng)
                    },
                )
            };
            total_moves += sc.moves;
            total_vetoed += sc.vetoed;
            prog_bar.inc(1);
            if params.gibbs_stagnation > 0.0 {
                if (sc.moves as f64) < params.gibbs_stagnation * (num_entities as f64) {
                    low_sweeps += 1;
                    if low_sweeps >= 3 {
                        break;
                    }
                } else {
                    low_sweeps = 0;
                }
            }
        }
    }

    for _sweep in 0..params.num_greedy {
        let sc = if params.parallel {
            sweep_jacobi(
                candidates,
                &mut stats,
                profiles,
                guard,
                k,
                &mut proposals,
                |log_probs, cand, _e, current| argmax_log_restricted_from(log_probs, cand, current),
            )
        } else {
            sweep_sequential(
                candidates,
                &mut stats,
                profiles,
                guard,
                &mut log_probs,
                0..num_entities,
                argmax_log_restricted_from,
            )
        };
        total_moves += sc.moves;
        total_vetoed += sc.vetoed;
        prog_bar.inc(1);
        if sc.moves == 0 {
            break;
        }
    }
    prog_bar.finish_and_clear();

    if total_vetoed > 0 {
        log::debug!(
            "{}: {} moves, {} vetoed by MoveGuard",
            level_label,
            total_moves,
            total_vetoed
        );
    }

    labels.copy_from_slice(&stats.membership);
    total_moves
}

/// Unguarded variant of [`refine_with_candidates_guarded`].
pub fn refine_with_candidates(
    labels: &mut [usize],
    candidates: &[Vec<usize>],
    rng: &mut SmallRng,
    ctx: &RefineContext,
) -> usize {
    refine_with_candidates_guarded(labels, candidates, &NoGuard, rng, ctx)
}

/// Generic driver: ask `proposer` for candidates, then run the guarded sweep.
pub fn refine_with_proposer_guarded<P: CandidateProposer, G: MoveGuard>(
    labels: &mut [usize],
    proposer: &P,
    guard: &G,
    rng: &mut SmallRng,
    ctx: &RefineContext,
) -> usize {
    let candidates = proposer.propose(labels);
    let moves = refine_with_candidates_guarded(labels, &candidates, guard, rng, ctx);
    info!("  {}: {} DC-Poisson moves", ctx.level_label, moves);
    moves
}

/// Unguarded generic driver: ask `proposer` for candidates, then run the
/// sweep with [`NoGuard`].
pub fn refine_with_proposer<P: CandidateProposer>(
    labels: &mut [usize],
    proposer: &P,
    rng: &mut SmallRng,
    ctx: &RefineContext,
) -> usize {
    let candidates = proposer.propose(labels);
    let moves = refine_with_candidates(labels, &candidates, rng, ctx);
    info!("  {}: {} DC-Poisson moves", ctx.level_label, moves);
    moves
}

////////////////////////////
// Tests (algorithm core) //
////////////////////////////

#[cfg(test)]
#[path = "dc_poisson_tests.rs"]
mod tests;