data-beans 0.12.0

//! Pseudobulk-sample (pb-sample) construction.
//!
//! `build_pb_samples` aggregates per-(batch, group) intersections of the
//! finest hash partition into a layout + per-pb-sample gene sums used
//! downstream by the refinement step (`refine_multilevel` BBKNN +
//! DC-SBM) and by the cross-batch matched-stat path.

use super::*;

pub struct PbSampleLayout {
    /// Centroid matrix: proj_dim x num_pb_samples
    pub centroids: DMatrix<f32>,
    /// Number of cells in each pb-sample
    pub cell_counts: Vec<f32>,
    /// Batch index for each pb-sample
    pub pb_sample_to_batch: Vec<usize>,
    /// Sample/group index for each pb-sample
    pub pb_sample_to_group: Vec<usize>,
    /// Maps (batch, group) → pb-sample index
    pub bg_to_pbsamp: HashMap<(usize, usize), usize>,
    /// Global cell index → owning pb-sample index.
    pub cell_to_pbsamp: Vec<usize>,
    /// For a singleton pb-sample (an anchored reference column or a bulk
    /// sample), the one column it stands for; `None` for ordinary
    /// (batch, group) pb-samples. The layout is the single owner of "which
    /// pb-samples are already summaries" — downstream steps
    /// (`split_anchored_finest_groups`) read it here instead of re-deriving
    /// membership from batch labels.
    pub singleton_col: Vec<Option<usize>>,
    /// Batches whose columns are bulk samples: singletons like anchors, but
    /// excluded from cross-batch matching in BOTH directions (see
    /// [`bbknn_match_one_pbsamp`]).
    pub bulk_batches: Vec<usize>,
    /// Per-pb-sample CNV stratum when collapse is stratified. Same-stratum
    /// BBKNN filter reads this; `None` disables the filter.
    pub pb_sample_to_stratum: Option<Vec<usize>>,
}

impl PbSampleLayout {
    /// Whether this pb-sample is a bulk sample (a singleton that takes no
    /// part in matching).
    #[must_use]
    pub fn is_bulk(&self, pbsamp: usize) -> bool {
        self.is_bulk_batch(self.pb_sample_to_batch[pbsamp])
    }

    /// Whether `batch` is a bulk batch — the one spelling of the predicate,
    /// so per-pb-sample and per-batch callers cannot drift apart.
    #[must_use]
    pub fn is_bulk_batch(&self, batch: usize) -> bool {
        self.bulk_batches.contains(&batch)
    }
}

/// Pre-aggregated pb-sample data for fast cross-batch matching.
/// Each pb-sample is the intersection of a (batch, sample) pair.
pub struct PbSampleCollection {
    pub layout: PbSampleLayout,
    /// Sparse gene sums per pb-sample: Vec of (gene_index, sum)
    pub gene_sums: Vec<Vec<(usize, f32)>>,
    /// Number of genes
    pub num_genes: usize,
}

/// Intermediate per-batch accumulator used during pb-sample construction.
pub(super) struct BatchAccumulator {
    centroid_sum: Vec<f32>,
    gene_sum: HashMap<usize, f32>,
    count: usize,
}

/// A single pb-sample produced from a (batch, group) intersection.
pub(super) struct PbSampleData {
    centroid: Vec<f32>,
    gene_sums: Vec<(usize, f32)>,
    cell_count: f32,
    batch: usize,
    group: usize,
}

/// Build the shared pb-sample layout from (batch, group) intersections.
///
/// For each non-empty (batch, group) block:
/// - Centroid = mean of projection vectors
/// - Cell count = number of cells in block
///
/// `anchor_batches` are batches whose columns are ALREADY pseudobulks — a
/// prior run's carried reference. Re-averaging those through this run's
/// partition would flatten resolution the parent already paid for (averages
/// of averages), so each anchored column becomes its own singleton pb-sample
/// instead, keeping the reference at its stored granularity for matching.
/// `bulk_batches` get the same singleton treatment (a bulk sample is already
/// a summary) but are additionally barred from matching — see
/// [`bbknn_match_one_pbsamp`].
///
/// This only uses `proj_kn` (no CSC reads), so it can be shared across layers.
pub(super) fn build_pb_sample_layout(
    group_to_cols: &[Vec<usize>],
    col_to_batch: &[usize],
    proj_kn: &DMatrix<f32>,
    col_weight: Option<&[f32]>,
    anchor_batches: &[usize],
    bulk_batches: &[usize],
    cell_to_stratum: Option<&[usize]>,
) -> anyhow::Result<PbSampleLayout> {
    let proj_dim = proj_kn.nrows();

    /// Intermediate per-batch accumulator for centroid computation.
    struct CentroidAccum {
        centroid_sum: Vec<f32>,
        /// Summed multiplicity, not a raw column count: a column standing for
        /// `m` cells must weigh `m` in both the centroid and `cell_counts`,
        /// or a carried pseudobulk counts the same as one cell.
        count: f32,
    }
    let weight_of = |c: usize| col_weight.map_or(1.0, |w| w[c]);

    // Collect centroid data per group in parallel. The last tuple field is
    // `Some(column)` for a singleton (anchored) pb-sample, `None` for an
    // ordinary (batch, group) block.
    type CentroidTuple = (usize, usize, Vec<f32>, f32, Option<usize>);
    let per_group_results: Vec<Vec<CentroidTuple>> = group_to_cols
        .par_iter()
        .enumerate()
        .map(|(group, cells)| {
            let mut batch_data: HashMap<usize, CentroidAccum> = HashMap::default();
            let mut singletons: Vec<CentroidTuple> = Vec::new();

            for &glob_idx in cells {
                let batch = col_to_batch[glob_idx];
                let w = weight_of(glob_idx);
                if anchor_batches.contains(&batch) || bulk_batches.contains(&batch) {
                    let centroid: Vec<f32> =
                        (0..proj_dim).map(|d| proj_kn[(d, glob_idx)]).collect();
                    singletons.push((batch, group, centroid, w, Some(glob_idx)));
                    continue;
                }
                let acc = batch_data.entry(batch).or_insert_with(|| CentroidAccum {
                    centroid_sum: vec![0f32; proj_dim],
                    count: 0.0,
                });
                for d in 0..proj_dim {
                    acc.centroid_sum[d] += proj_kn[(d, glob_idx)] * w;
                }
                acc.count += w;
            }

            let mut out: Vec<CentroidTuple> = batch_data
                .into_iter()
                .filter(|(_, acc)| acc.count > 0.0)
                .map(|(batch, acc)| {
                    let inv_count = 1.0 / acc.count;
                    let centroid: Vec<f32> =
                        acc.centroid_sum.iter().map(|v| v * inv_count).collect();
                    (batch, group, centroid, acc.count, None)
                })
                .collect::<Vec<_>>();
            out.extend(singletons);
            out
        })
        .collect();

    // Flatten into a single list
    let all_pbsamp: Vec<_> = per_group_results.into_iter().flatten().collect();
    let num_pb = all_pbsamp.len();

    if num_pb == 0 {
        return Err(anyhow::anyhow!("no pb-samples built"));
    }

    // Build centroid matrix and metadata
    let mut centroids = DMatrix::<f32>::zeros(proj_dim, num_pb);
    let mut cell_counts = Vec::with_capacity(num_pb);
    let mut pbsamp_to_batch = Vec::with_capacity(num_pb);
    let mut pbsamp_to_group = Vec::with_capacity(num_pb);
    let mut singleton_col = Vec::with_capacity(num_pb);
    let mut bg_to_pbsamp = HashMap::default();

    let ncols = col_to_batch.len();
    let mut cell_to_pbsamp = vec![usize::MAX; ncols];
    for (i, (batch, group, centroid, count, sc)) in all_pbsamp.into_iter().enumerate() {
        for (d, &v) in centroid.iter().enumerate() {
            centroids[(d, i)] = v;
        }
        cell_counts.push(count);
        pbsamp_to_batch.push(batch);
        pbsamp_to_group.push(group);
        singleton_col.push(sc);
        match sc {
            // A singleton column IS its pb-sample; `(batch, group)` is not a
            // unique key for singletons, so they are mapped directly and stay
            // out of `bg_to_pbsamp`.
            Some(col) => cell_to_pbsamp[col] = i,
            None => {
                bg_to_pbsamp.insert((batch, group), i);
            }
        }
    }

    // Cell → pb-sample inversion for the ordinary (batch, group) blocks;
    // singleton columns were mapped directly above.
    for (group, cells) in group_to_cols.iter().enumerate() {
        for &c in cells {
            let b = col_to_batch[c];
            if cell_to_pbsamp[c] == usize::MAX {
                if let Some(&pbsamp) = bg_to_pbsamp.get(&(b, group)) {
                    cell_to_pbsamp[c] = pbsamp;
                }
            }
        }
    }

    let pb_sample_to_stratum = if let Some(strata) = cell_to_stratum {
        anyhow::ensure!(
            strata.len() == ncols,
            "cell_to_stratum has {} entries, layout has {} columns",
            strata.len(),
            ncols
        );
        let mut out = vec![0usize; num_pb];
        let mut seen = vec![false; num_pb];
        for (c, &pbsamp) in cell_to_pbsamp.iter().enumerate() {
            if pbsamp == usize::MAX {
                continue;
            }
            if !seen[pbsamp] {
                out[pbsamp] = strata[c];
                seen[pbsamp] = true;
            }
        }
        Some(out)
    } else {
        None
    };

    Ok(PbSampleLayout {
        centroids,
        cell_counts,
        pb_sample_to_batch: pbsamp_to_batch,
        pb_sample_to_group: pbsamp_to_group,
        bg_to_pbsamp,
        cell_to_pbsamp,
        singleton_col,
        bulk_batches: bulk_batches.to_vec(),
        pb_sample_to_stratum,
    })
}

/// Collect gene sums for each pb-sample from a single `SparseIoVec` layer.
///
/// Keyed on the layout's `cell_to_pbsamp` so ordinary (batch, group) blocks
/// and singleton (anchored) pb-samples accumulate through one path.
pub(super) fn collect_pb_sample_gene_sums(
    data_vec: &SparseIoVec,
    group_to_cols: &[Vec<usize>],
    cell_to_pbsamp: &[usize],
    num_pb: usize,
) -> anyhow::Result<Vec<Vec<(usize, f32)>>> {
    use indicatif::ParallelProgressIterator;
    let num_genes = data_vec.num_rows();
    let prog_bar = styled_progress_bar(group_to_cols.len() as u64, "groups (pb-sample gene sums)");
    // One dense gene accumulator per worker, reset through the list of genes
    // a pb-sample touched, so each nonzero costs an indexed add.
    let per_group: Vec<Vec<_>> = group_to_cols
        .par_iter()
        .progress_with(prog_bar.clone())
        .map_init(
            || (vec![0f32; num_genes], vec![false; num_genes], Vec::new()),
            |(sums, seen, touched), cells| -> anyhow::Result<_> {
                let yy = data_vec.read_columns_csc(cells.iter().cloned())?;

                // A pb-sample lies within one group; the stable sort lines its
                // cells up as one run, in group order.
                let mut by_pbsamp: Vec<(usize, usize)> = cells
                    .iter()
                    .enumerate()
                    .map(|(local_idx, &col)| (cell_to_pbsamp[col], local_idx))
                    .filter(|&(pbsamp, _)| pbsamp != usize::MAX)
                    .collect();
                by_pbsamp.sort_by_key(|&(pbsamp, _)| pbsamp);

                let mut out: Vec<(usize, Vec<(usize, f32)>)> = Vec::new();
                for run in by_pbsamp.chunk_by(|a, b| a.0 == b.0) {
                    for &(_, local_idx) in run {
                        // Weighted to stay consistent with the weighted
                        // `cell_counts` in the layout:
                        // `collect_matched_stat_coarse` divides one by the
                        // other, and the quotient has to remain the per-cell
                        // rate.
                        let w = data_vec.column_multiplicity(cells[local_idx]);
                        let y_j = yy.col(local_idx);
                        for (&gene, &val) in y_j.row_indices().iter().zip(y_j.values().iter()) {
                            if !seen[gene] {
                                seen[gene] = true;
                                touched.push(gene);
                            }
                            sums[gene] += val * w;
                        }
                    }
                    touched.sort_unstable();
                    let sorted: Vec<(usize, f32)> = touched
                        .drain(..)
                        .map(|gene| {
                            seen[gene] = false;
                            (gene, std::mem::take(&mut sums[gene]))
                        })
                        .collect();
                    out.push((run[0].0, sorted));
                }
                Ok(out)
            },
        )
        .collect::<anyhow::Result<_>>()?;
    prog_bar.finish_and_clear();

    let mut gene_sums: Vec<Vec<(usize, f32)>> = vec![vec![]; num_pb];
    for (pbsamp_idx, sorted) in per_group.into_iter().flatten() {
        gene_sums[pbsamp_idx] = sorted;
    }

    Ok(gene_sums)
}

/// Build pb-samples (layout + gene sums) from a single `SparseIoVec`.
pub(super) fn build_pb_samples(
    data_vec: &SparseIoVec,
    proj_kn: &DMatrix<f32>,
    num_genes: usize,
    anchor_batches: &[usize],
    bulk_batches: &[usize],
    cell_to_stratum: Option<&[usize]>,
) -> anyhow::Result<PbSampleCollection> {
    let group_to_cols = data_vec
        .take_grouped_columns()
        .ok_or(anyhow::anyhow!("columns not assigned to groups"))?;
    let col_to_batch: Vec<usize> = (0..proj_kn.ncols())
        .map(|c| data_vec.get_batch_membership(std::iter::once(c))[0])
        .collect();

    let weights: Option<Vec<f32>> = data_vec.has_column_multiplicity().then(|| {
        (0..proj_kn.ncols())
            .map(|c| data_vec.column_multiplicity(c))
            .collect()
    });
    let layout = build_pb_sample_layout(
        group_to_cols,
        &col_to_batch,
        proj_kn,
        weights.as_deref(),
        anchor_batches,
        bulk_batches,
        cell_to_stratum,
    )?;
    let num_pb = layout.cell_counts.len();
    let gene_sums =
        collect_pb_sample_gene_sums(data_vec, group_to_cols, &layout.cell_to_pbsamp, num_pb)?;

    Ok(PbSampleCollection {
        layout,
        gene_sums,
        num_genes,
    })
}

/// How many distinct pb-samples each batch's cell index can return, in total
/// and (when the layout is stratified) per stratum. A match search for `knn`
/// pb-samples is complete once it holds `min(knn, count)` of them, so it need
/// not exhaust a batch that has fewer than `knn` to offer.
pub(crate) struct MatchablePbSamples {
    per_batch: Vec<usize>,
    per_batch_stratum: Option<Vec<HashMap<usize, usize>>>,
}

impl MatchablePbSamples {
    pub(crate) fn new(layout: &PbSampleLayout, batch_knn_lookup: &[ColumnDict<usize>]) -> Self {
        let strata = layout.pb_sample_to_stratum.as_deref();
        let (per_batch, by_stratum): (Vec<usize>, Vec<HashMap<usize, usize>>) = batch_knn_lookup
            .par_iter()
            .map(|bknn| {
                let mut pbsamps: Vec<usize> = bknn
                    .names()
                    .iter()
                    .map(|&c| layout.cell_to_pbsamp[c])
                    .filter(|&p| p != usize::MAX)
                    .collect();
                pbsamps.sort_unstable();
                pbsamps.dedup();
                let mut by_stratum: HashMap<usize, usize> = HashMap::default();
                if let Some(st) = strata {
                    for &p in &pbsamps {
                        *by_stratum.entry(st[p]).or_default() += 1;
                    }
                }
                (pbsamps.len(), by_stratum)
            })
            .unzip();
        Self {
            per_batch,
            per_batch_stratum: strata.map(|_| by_stratum),
        }
    }

    /// Distinct pb-samples batch `batch` holds for a query of `stratum`
    /// (`None` = unstratified: all of them).
    fn count(&self, batch: usize, stratum: Option<usize>) -> usize {
        match (self.per_batch_stratum.as_ref(), stratum) {
            (Some(by_stratum), Some(s)) => by_stratum[batch].get(&s).copied().unwrap_or(0),
            _ => self.per_batch[batch],
        }
    }
}

/// Recover up to `knn` *distinct* pb-samples (closest, by min cell distance)
/// from a single batch's cell-level HNSW for the given `query` centroid.
///
/// Grows the HNSW query adaptively until `knn` distinct pb-samples are found or
/// the batch's cells are exhausted (bounded by `bknn.len()`). This replaces the
/// old fixed `4·knn+1` oversample, which under-recovered whenever a batch's
/// nearest cells collapsed into fewer than `knn` pb-samples (small / sorting-
/// biased batches) — starving the cross-batch δ estimate of matches.
///
/// `available` is the number of distinct pb-samples in `bknn` the stratum
/// filter lets through ([`MatchablePbSamples`]): the search is complete once
/// it holds `min(knn, available)` of them, and has nothing to do when there
/// are none. A count above the true number only defers the stop to exhaustion.
///
/// `own_pbsamp` and `usize::MAX` (unassigned) cells are skipped. `stratum` is
/// the per-pb-sample stratum table with the query's own stratum; only
/// pb-samples of that stratum are kept. Returned distance per pb-sample is the
/// minimum over its collapsed cells.
pub(crate) fn knn_distinct_pbsamples_in_batch(
    bknn: &ColumnDict<usize>,
    query: &[f32],
    knn: usize,
    available: usize,
    cell_to_pbsamp: &[usize],
    own_pbsamp: usize,
    stratum: Option<(&[usize], usize)>,
) -> anyhow::Result<Vec<(usize, f32)>> {
    let n = bknn.num_points();
    let wanted = knn.min(available);
    if n == 0 || wanted == 0 {
        return Ok(Vec::new());
    }
    // Start at the historical oversample, then grow ×4 per round. `search_by_
    // query_data` self-clamps the request to the batch size, so query_k == n
    // returns every cell — our exhaustion signal.
    let mut query_k = (knn * 4 + 1).min(n);
    let mut best: HashMap<usize, f32> = HashMap::default();
    loop {
        let (cell_ids, dists) = bknn.search_by_query_data(query, query_k)?;
        best.clear();
        for (&c, &d) in cell_ids.iter().zip(dists.iter()) {
            let other_pbsamp = cell_to_pbsamp[c];
            if other_pbsamp == usize::MAX || other_pbsamp == own_pbsamp {
                continue;
            }
            if stratum.is_some_and(|(st, own_s)| st[other_pbsamp] != own_s) {
                continue;
            }
            best.entry(other_pbsamp)
                .and_modify(|old| {
                    if d < *old {
                        *old = d;
                    }
                })
                .or_insert(d);
        }
        if best.len() >= wanted || query_k >= n {
            break;
        }
        query_k = query_k.saturating_mul(4).min(n);
    }
    let mut per_batch: Vec<(usize, f32)> = best.drain().collect();
    per_batch.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
    per_batch.truncate(knn);
    Ok(per_batch)
}

/// For a single `pbsamp`, gather its cross-batch match set: up to `knn`
/// distinct pb-samples from **each** non-own batch (closest to `pbsamp`'s
/// centroid), flattened across batches via [`knn_distinct_pbsamples_in_batch`].
///
/// The per-pb-sample body of [`per_batch_sc_neighbors`]; `matchable` is the
/// per-batch census of `batch_knn_lookup` that bounds each search.
fn bbknn_match_one_pbsamp(
    layout: &PbSampleLayout,
    batch_knn_lookup: &[ColumnDict<usize>],
    matchable: &MatchablePbSamples,
    knn: usize,
    pbsamp: usize,
    anchor_batches: Option<&[usize]>,
) -> anyhow::Result<Vec<(usize, f32)>> {
    // Bulk is a RECEIVER but never a SOURCE — greedy correction, the same
    // discipline the carried pb_reference uses. A bulk column draws its
    // counterfactual from the cell frame and is corrected toward it; the
    // cells self-match through the anchor path, so their frame never moves.
    // Letting bulk serve as a source is what would drag the cells, and that
    // is the only direction excluded here (below, and by keeping bulk out of
    // every anchor set).
    let pbsamp_batch = layout.pb_sample_to_batch[pbsamp];
    let centroid: Vec<f32> = layout.centroids.column(pbsamp).iter().copied().collect();
    let stratum = layout
        .pb_sample_to_stratum
        .as_deref()
        .map(|st| (st, st[pbsamp]));
    let own_stratum = stratum.map(|(_, own_s)| own_s);
    let mut all_hits: Vec<(usize, f32)> = Vec::new();
    match anchor_batches {
        // Pooled matching: every non-own, non-bulk batch contributes
        // counterfactual candidates, symmetrically. A bulk batch cannot be a
        // source: matching a cell pb-sample to a mixture is the same
        // composition-into-δ leak in the other direction.
        None => {
            for (b, bknn) in batch_knn_lookup.iter().enumerate() {
                if b == pbsamp_batch || layout.is_bulk_batch(b) {
                    continue;
                }
                let per_batch = knn_distinct_pbsamples_in_batch(
                    bknn,
                    &centroid,
                    knn,
                    matchable.count(b, own_stratum),
                    &layout.cell_to_pbsamp,
                    pbsamp,
                    stratum,
                )?;
                all_hits.extend(per_batch);
            }
        }
        // Greedy (anchored) matching: counterfactuals come from the anchor
        // batches ONLY, for everyone. New batches are corrected toward the
        // anchor frame; an anchor pb-sample's nearest anchor is itself
        // (distance 0 → softmax self-match), so the frame stays fixed and its
        // δ settles at the prior — the reference is never re-adjusted.
        Some(anchors) => {
            for &b in anchors {
                let Some(bknn) = batch_knn_lookup.get(b) else {
                    continue;
                };
                let per_batch = knn_distinct_pbsamples_in_batch(
                    bknn,
                    &centroid,
                    knn,
                    matchable.count(b, own_stratum),
                    &layout.cell_to_pbsamp,
                    usize::MAX, // self-match allowed
                    stratum,
                )?;
                all_hits.extend(per_batch);
            }
        }
    }
    Ok(all_hits)
}

/// Per-pb-sample, for each non-own batch return up to `knn` distinct
/// pb-samples whose member cells are closest to `pbsamp`'s centroid.
///
/// Queries `SparseIoVec::batch_knn_lookup` (per-batch HNSW over cells) via
/// [`knn_distinct_pbsamples_in_batch`], which grows each per-batch search until
/// `knn` distinct pb-samples are recovered or the batch is exhausted.
///
/// With `anchor_batches`, the sources are the anchor batches only, for every
/// pb-sample (see [`bbknn_match_one_pbsamp`]).
///
/// Returns: `result[pbsamp] = Vec<(other_pbsamp, distance)>` flattened across all
/// non-own batches.
pub(crate) fn per_batch_sc_neighbors(
    layout: &PbSampleLayout,
    batch_knn_lookup: &[ColumnDict<usize>],
    knn: usize,
    anchor_batches: Option<&[usize]>,
) -> anyhow::Result<Vec<Vec<(usize, f32)>>> {
    use indicatif::ParallelProgressIterator;
    let num_pb = layout.cell_counts.len();
    let matchable = MatchablePbSamples::new(layout, batch_knn_lookup);

    let prog_bar = styled_progress_bar(num_pb as u64, "pb-samples (BBKNN match)");
    let result = (0..num_pb)
        .into_par_iter()
        .progress_with(prog_bar.clone())
        .map(|pbsamp| {
            bbknn_match_one_pbsamp(
                layout,
                batch_knn_lookup,
                &matchable,
                knn,
                pbsamp,
                anchor_batches,
            )
        })
        .collect();
    prog_bar.finish_and_clear();
    result
}

/// Match pb-samples across batches and accumulate counterfactual
/// statistics into `stat.imputed_sum_ds` and `stat.matched_bs`.
///
/// `pbsamp_to_group` is the per-pb-sample group assignment to use when writing
/// into stat columns; callers pass `&layout.pb_sample_to_group` for the
/// hash-partition mapping, or a refined mapping from
/// `refine_multilevel::refine_assignments`.
///
/// `knn` is now the per-other-batch neighbour count: each pb-sample draws
/// up to `knn` distinct foreign pb-samples from **each** non-own batch, so
/// the total match set is up to `knn · (num_batches − 1)`.
pub(super) fn build_pb_sample_to_cells(layout: &PbSampleLayout) -> Vec<Vec<usize>> {
    let num_pb = layout.cell_counts.len();
    let mut out: Vec<Vec<usize>> = vec![vec![]; num_pb];
    for (c, &pbsamp) in layout.cell_to_pbsamp.iter().enumerate() {
        if pbsamp != usize::MAX {
            out[pbsamp].push(c);
        }
    }
    out
}

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