data_beans/alg/

refine_multilevel.rs

//! BBKNN-based Poisson refinement over multi-level pb-sample partitions.
//!
//! Given a hash-initialized hierarchy of pb-sample → group mappings and a
//! `PbSampleLayout` (HNSW over centroids, batch assignments), refine each
//! level from coarsest to finest by proposing moves that keep each pb-sample
//! under the same parent group (sibling-constrained) and are drawn from the
//! batch-balanced KNN neighborhood. Moves are scored by a DC-Poisson
//! log-likelihood with NB Fisher-info feature weighting (see
//! [`crate::alg::dc_poisson::FeatureWeighting`]).
//!
//! Candidate-set construction is BBKNN-specific and lives here ([`BbknnProposer`]).
//! The generic scoring core lives in [`crate::alg::dc_poisson`] and is shared with
//! other front-ends (e.g. pinto's spatial-graph proposer).
//!
//! Entry point: [`refine_assignments`].

#![allow(dead_code)]

use crate::alg::collapse_data::{per_batch_sc_neighbors, PbSampleLayout};
use crate::alg::dc_poisson::{
    compute_sibling_sets, intersect_with_siblings_fallback, refine_with_proposer,
    CandidateProposer, ProfileSource, Profiles, RefineContext,
};
use log::{debug, info};
use rand::rngs::SmallRng;
use rand::SeedableRng;
use rustc_hash::FxHashMap;

pub use crate::alg::collapse_data::GeneSums;
// Re-export the DC-Poisson core items so existing callers / integration
// tests that expect `refine_multilevel::{RefineParams, compact_labels, ...}`
// keep working after the core moved.
pub use crate::alg::dc_poisson::{compact_labels, RefineParams};

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

/// Per-level refined pb-sample → group mapping.
///
/// `pbsamp_to_group[level][pbsamp_idx] = refined group id at that level`. Level order
/// matches the existing collapse convention (`level_dims[0]` = finest).
#[derive(Debug, Clone)]
pub struct RefinedAssignment {
    pub pbsamp_to_group: Vec<Vec<usize>>,
    pub num_groups_per_level: Vec<usize>,
}

/////////////////////////////////////////////////////////////////////////
// Candidate set construction (siblings ∩ BBKNN with sibling fallback) //
/////////////////////////////////////////////////////////////////////////

/// The candidate pb-sample ids of a match table — distances are unused for
/// refinement.
fn candidate_ids(matches: &[Vec<(usize, f32)>]) -> Vec<Vec<usize>> {
    matches
        .iter()
        .map(|hits| hits.iter().map(|&(p, _)| p).collect())
        .collect()
}

/// Candidate set per pb-sample = siblings ∩ (groups of BBKNN neighbors),
/// via the shared [`intersect_with_siblings_fallback`] helper.
///
/// With `pb_sample_to_stratum`, a pb-sample's siblings are first cut down to
/// its own group and the groups made of its stratum only, so neither the
/// intersection nor the sibling fallback (a pb-sample without a same-stratum
/// match in another batch has no BBKNN neighbors) can move it across strata.
fn build_candidate_sets(
    siblings: &[Vec<usize>],
    bbknn: &[Vec<usize>],
    pbsamp_to_group_at_level: &[usize],
    pb_sample_to_stratum: Option<&[usize]>,
) -> Vec<Vec<usize>> {
    let strata = pb_sample_to_stratum
        .map(|stratum| (stratum, group_strata(pbsamp_to_group_at_level, stratum)));
    siblings
        .iter()
        .enumerate()
        .map(|(pbsamp, sib)| {
            let own = pbsamp_to_group_at_level[pbsamp];
            let mut neighbor_groups: Vec<usize> = bbknn[pbsamp]
                .iter()
                .map(|&j| pbsamp_to_group_at_level[j])
                .collect();
            neighbor_groups.sort_unstable();
            neighbor_groups.dedup();
            match strata.as_ref() {
                Some((stratum, of_group)) => {
                    let same_stratum: Vec<usize> = sib
                        .iter()
                        .copied()
                        .filter(|&g| g == own || of_group.get(g) == Some(&Some(stratum[pbsamp])))
                        .collect();
                    intersect_with_siblings_fallback(&same_stratum, &neighbor_groups, own)
                }
                None => intersect_with_siblings_fallback(sib, &neighbor_groups, own),
            }
        })
        .collect()
}

/// The one stratum of each group's members, for the groups up to the largest
/// label in use; `None` for an empty group and for a group that mixes strata.
fn group_strata(pbsamp_to_group: &[usize], stratum: &[usize]) -> Vec<Option<usize>> {
    let k = pbsamp_to_group.iter().max().map_or(0, |&m| m + 1);
    let mut of_group: Vec<Option<usize>> = vec![None; k];
    let mut mixed = vec![false; k];
    for (&g, &s) in pbsamp_to_group.iter().zip(stratum) {
        match of_group[g] {
            None => of_group[g] = Some(s),
            Some(seen) if seen != s => mixed[g] = true,
            Some(_) => {}
        }
    }
    for (slot, &is_mixed) in of_group.iter_mut().zip(&mixed) {
        if is_mixed {
            *slot = None;
        }
    }
    of_group
}

///////////////////////////////////
// BBKNN-based CandidateProposer //
///////////////////////////////////

/// Proposes move candidates for pb-samples via cross-batch BBKNN.
///
/// For each pb-sample, the candidate set is
/// `siblings ∩ (groups of BBKNN neighbors)`, with sibling fallback when
/// the intersection is empty. Siblings are the other pb-samples sharing
/// the same parent group at the next-coarser level (passed in at
/// construction). With strata ([`Self::with_strata`]) the candidates are
/// further limited to groups of the pb-sample's own stratum.
pub struct BbknnProposer<'a> {
    siblings: Vec<Vec<usize>>,
    bbknn: &'a [Vec<usize>],
    pb_sample_to_stratum: Option<&'a [usize]>,
}

impl<'a> BbknnProposer<'a> {
    pub fn new(siblings: Vec<Vec<usize>>, bbknn: &'a [Vec<usize>]) -> Self {
        Self {
            siblings,
            bbknn,
            pb_sample_to_stratum: None,
        }
    }

    /// Keep every pb-sample within groups of its own stratum
    /// (`PbSampleLayout::pb_sample_to_stratum`); `None` leaves the proposer
    /// unstratified.
    #[must_use]
    pub fn with_strata(mut self, pb_sample_to_stratum: Option<&'a [usize]>) -> Self {
        self.pb_sample_to_stratum = pb_sample_to_stratum;
        self
    }
}

impl<'a> CandidateProposer for BbknnProposer<'a> {
    fn propose(&self, labels: &[usize]) -> Vec<Vec<usize>> {
        build_candidate_sets(
            &self.siblings,
            self.bbknn,
            labels,
            self.pb_sample_to_stratum,
        )
    }
}

///////////////////////
// Refinement driver //
///////////////////////

/// Fixed inputs describing *what* to refine. `batch_knn_lookup` is the
/// per-batch HNSW over cells, typically obtained from
/// `SparseIoVec::batch_knn_lookup()` after `build_hnsw_per_batch`.
/// `k_per_batch` is the BBKNN fan-out — up to this many distinct
/// pb-samples are drawn from **each** non-own batch as move candidates.
#[derive(Clone, Copy)]
pub struct RefineInputs<'a> {
    pub layout: &'a PbSampleLayout,
    pub gene_sums: &'a GeneSums,
    pub num_genes: usize,
    pub pb_sample_to_cells: &'a [Vec<usize>],
    pub batch_knn_lookup: &'a [legume_numeric::matrix::knn_match::ColumnDict<usize>],
    pub k_per_batch: usize,
    pub initial_sc_to_group_per_level: &'a [Vec<usize>],
    /// Per-level reprojection offsets used to re-anchor each finer level in its
    /// refined parent (the child hash *relative to its parent* — the "extra
    /// bits"). `reproject_offsets_per_level[level]` is the per-pb-sample offset
    /// for subdividing `level` against `level+1`; the coarsest level's entry is
    /// unused (no parent). An empty outer slice (or empty per-level entry)
    /// falls back to the positional [`child_offset_within_parent`], which is
    /// also bounded but not hash-semantic. See `collapse_data::refine`.
    pub reproject_offsets_per_level: &'a [Vec<usize>],
}

/// Top-down BBKNN + Poisson refinement.
///
/// Levels in `inputs.initial_sc_to_group_per_level` follow the existing
/// finest → coarsest convention. The returned `RefinedAssignment` uses the
/// same ordering.
pub fn refine_assignments(
    inputs: &RefineInputs<'_>,
    params: &RefineParams,
) -> anyhow::Result<RefinedAssignment> {
    refine_assignments_with_matches(inputs, params, None)
}

/// [`refine_assignments`] with the pooled cross-batch match table already at
/// hand: `pooled_matches` is `per_batch_sc_neighbors(layout, batch_knn_lookup,
/// k_per_batch, None)`, which the caller also needs for the matched
/// statistics. `None` computes it here.
pub(crate) fn refine_assignments_with_matches(
    inputs: &RefineInputs<'_>,
    params: &RefineParams,
    pooled_matches: Option<&[Vec<(usize, f32)>]>,
) -> anyhow::Result<RefinedAssignment> {
    let RefineInputs {
        layout,
        gene_sums,
        num_genes,
        pb_sample_to_cells,
        batch_knn_lookup,
        k_per_batch,
        initial_sc_to_group_per_level,
        reproject_offsets_per_level,
    } = *inputs;
    let num_levels = initial_sc_to_group_per_level.len();
    if num_levels == 0 {
        return Err(anyhow::anyhow!("no levels"));
    }
    let num_pb = layout.cell_counts.len();
    for (i, lvl) in initial_sc_to_group_per_level.iter().enumerate() {
        if lvl.len() != num_pb {
            return Err(anyhow::anyhow!(
                "level {} has {} entries, expected {}",
                i,
                lvl.len(),
                num_pb
            ));
        }
    }

    // Compact each level's initial labels to a dense 0..K range.
    let mut refined: Vec<Vec<usize>> = Vec::with_capacity(num_levels);
    let mut ks: Vec<usize> = Vec::with_capacity(num_levels);
    for lvl in initial_sc_to_group_per_level {
        let (compact, k) = compact_labels(lvl);
        refined.push(compact);
        ks.push(k);
    }

    // Both sweep counts zero ⇒ no DC-Poisson moves can happen; skip the
    // (expensive) profile build, NB-Fisher weighting, and BBKNN construction
    // and return the compacted initial labels.
    if params.num_gibbs == 0 && params.num_greedy == 0 {
        info!("Skipping DC-Poisson refinement: --pb-refine-gibbs=0 and --pb-refine-greedy=0");
        return Ok(RefinedAssignment {
            pbsamp_to_group: refined,
            num_groups_per_level: ks,
        });
    }

    // Build profiles once.
    info!(
        "Building DC-Poisson profiles ({} pb-samples × {} genes) ...",
        num_pb, num_genes
    );
    let mut profiles = match &params.profile_source {
        ProfileSource::Raw => Profiles::from_gene_sums(gene_sums, num_genes),
        ProfileSource::Projected { basis } => Profiles::from_projection(basis, pb_sample_to_cells),
    };
    if matches!(params.profile_source, ProfileSource::Raw)
        && !matches!(
            params.feature_weighting,
            crate::alg::dc_poisson::FeatureWeighting::None
        )
    {
        info!("Computing NB Fisher-info feature weights ...");
        profiles.apply_feature_weighting(params.feature_weighting);
    }

    // BBKNN proposals via the shared per-batch cell HNSW: for every non-own
    // batch, that batch's cell-level index is queried with the pb-sample's
    // centroid and the cells returned are deduplicated to their owning
    // pb-samples, up to `k_per_batch` per other batch.
    info!(
        "Building BBKNN candidate sets (knn={} per non-own batch) ...",
        k_per_batch
    );
    let bbknn = match pooled_matches {
        Some(matches) => candidate_ids(matches),
        // Pooled (None): the DC-SBM refinement graph should connect across
        // every batch pair regardless of anchoring — anchoring shapes the
        // counterfactual for δ, not the clustering topology.
        None => candidate_ids(&per_batch_sc_neighbors(
            layout,
            batch_knn_lookup,
            k_per_batch,
            None,
        )?),
    };

    let mut rng = SmallRng::seed_from_u64(params.seed);

    // Walk coarsest → finest (highest level index down to 0).
    for level in (0..num_levels).rev() {
        // Refining the coarser level moves pb-samples across parents, so this
        // finer level must be re-anchored to nest strictly in its *refined*
        // parent before sibling sets are computed (else a straddling group
        // lands in two parents' sibling sets and the hierarchy breaks,
        // violating `fine_to_coarse_from_refined`'s invariant).
        //
        // We subdivide the refined parent by the child hash *relative to its
        // parent* (`child_offset_within_parent`), not by the full child code.
        // Both give a strict refinement, but the relative form is bounded by
        // `2^child_dim` — crossing the scrambled parent with the full child
        // code instead inflates the count ~20× (e.g. 1024 leaf codes →
        // 19098 finest groups), which is the level-graduation pathology.
        if level + 1 < num_levels {
            // Prefer the hash-semantic extra-bit offsets the caller precomputed;
            // fall back to the positional offset (still bounded) when absent.
            let offset: Vec<usize> = match reproject_offsets_per_level.get(level) {
                Some(o) if !o.is_empty() => o.clone(),
                _ => child_offset_within_parent(
                    &initial_sc_to_group_per_level[level],
                    &initial_sc_to_group_per_level[level + 1],
                ),
            };
            let (reprojected, new_k) = project_to_refinement(&offset, &refined[level + 1]);
            refined[level] = reprojected;
            ks[level] = new_k;
        }
        let k = ks[level];
        debug!("refining level {} (k={}, num_pb={})", level, k, num_pb);
        let siblings = compute_sibling_sets(&refined, level, k);

        let proposer = BbknnProposer::new(siblings, &bbknn)
            .with_strata(layout.pb_sample_to_stratum.as_deref());

        let pbsamp_to_group = &mut refined[level];
        let label = format!("Refine L{}/{}", num_levels - level, num_levels);
        let moves = refine_with_proposer(
            pbsamp_to_group,
            &proposer,
            &mut rng,
            &RefineContext {
                profiles: &profiles,
                k,
                params,
                level_label: &label,
            },
        );
        // Compact in case greedy emptied groups (keeps K monotone without gaps).
        let (compact, new_k) = compact_labels(pbsamp_to_group);
        *pbsamp_to_group = compact;
        ks[level] = new_k;
        info!(
            "  level {} refined: {} moves; {} of {} groups remain",
            level, moves, new_k, k
        );
    }

    Ok(RefinedAssignment {
        pbsamp_to_group: refined,
        num_groups_per_level: ks,
    })
}

/// Re-label `child` so it is a strict refinement of `parent`: two entities
/// share an output label iff they agree on **both** their current child
/// label and their parent label. Any child group that straddles multiple
/// parents is split. Returns the relabeled vector and its group count
/// (already compact `0..k` — `compact_labels` over the `(child, parent)`
/// key does the first-appearance dense relabel).
fn project_to_refinement(child: &[usize], parent: &[usize]) -> (Vec<usize>, usize) {
    debug_assert_eq!(child.len(), parent.len());
    let pairs: Vec<(usize, usize)> = child.iter().copied().zip(parent.iter().copied()).collect();
    compact_labels(&pairs)
}

/// Local index of each entity's `child` label within its `parent` label, in
/// first-seen order.
///
/// Used to subdivide a refined parent by the child hash *relative to its
/// parent* (the "extra bits") rather than by the full child code. Because the
/// child hash nests in the parent hash, each parent holds at most
/// `2^(child_dim − parent_dim)` distinct child codes, so the offset — and any
/// `(refined_parent, offset)` subdivision built from it — stays bounded by
/// `2^child_dim` even after DC-SBM has scrambled the parent. Crossing the
/// refined parent with the *full* child code instead (the old
/// `project_to_refinement` against `refined[level]`) cross-products the
/// scrambled parent with every child code and inflates the count ~20×.
fn child_offset_within_parent(child: &[usize], parent: &[usize]) -> Vec<usize> {
    debug_assert_eq!(child.len(), parent.len());
    let mut per_parent: FxHashMap<usize, FxHashMap<usize, usize>> = FxHashMap::default();
    let mut offsets = vec![0usize; child.len()];
    for (i, (&c, &p)) in child.iter().zip(parent.iter()).enumerate() {
        let local = per_parent.entry(p).or_default();
        let next = local.len();
        offsets[i] = *local.entry(c).or_insert(next);
    }
    offsets
}

////////////////////////////
// Tests (BBKNN-specific) //
////////////////////////////

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