data-beans 0.12.1

//! Lock-free parallel streaming pass over a `SparseIoVec` that yields a
//! merged `SparseRunningStatistics<f32>`.
//!
//! Both HVG selection (`hvg::select_hvg_streaming`) and NB-Fisher
//! weighting (`gene_weighting::compute_nb_fisher_weights`) need the same
//! per-feature `(npos, sum, sum_sq, n_cells)` aggregates from a sparse
//! backend. Sharing the skeleton keeps the parallel / lock-free
//! invariants in one place and gives every caller a labelled progress
//! bar with consistent styling.

use crate::qc::{blocks_per_leaf, fold_in_fixed_order};
use crate::sparse_io_vector::SparseIoVec;
use legume_numeric::matrix::sparse_stat::SparseRunningStatistics;
use legume_numeric::matrix::utils::generate_minibatch_intervals;

/// Stream cells through `data_vec` in `block_size` chunks, accumulating
/// per-feature NB sufficient statistics in parallel and merging them in a
/// fixed order ([`fold_in_fixed_order`]), so the result does not depend on
/// the thread count.
///
/// `progress_label` is shown in the indicatif progress bar (e.g.
/// `"HVG"`, `"NB-Fisher"`).
pub fn streaming_sparse_running_stats(
    data_vec: &SparseIoVec,
    block_size: Option<usize>,
    progress_label: &str,
) -> anyhow::Result<SparseRunningStatistics<f32>> {
    let n_features = data_vec.num_rows();
    let n_total = data_vec.num_columns();
    let jobs = generate_minibatch_intervals(n_total, n_features, block_size);

    let prog_bar = legume_numeric::matrix::progress::new_progress_bar(jobs.len() as u64)
        .with_message(format!("{progress_label} blocks"));

    let stats = fold_in_fixed_order(
        &jobs,
        blocks_per_leaf(jobs.len()),
        &|| SparseRunningStatistics::<f32>::new(n_features),
        &|acc, &(lb, ub)| {
            let chunk = data_vec.read_columns_csc(lb..ub)?;
            acc.add_csc(&chunk);
            prog_bar.inc(1);
            Ok(())
        },
        &|a, b| a.merge(&b),
    )?;
    prog_bar.finish_and_clear();

    Ok(stats)
}

/// The same pass, accumulating a second set of statistics on a folded axis.
///
/// Returns `(row_stats, folded_stats)`. `row_to_gene[r]` names the folded
/// bucket row `r` belongs to; buckets are `0..n_folded`.
///
/// The fold happens inside the block read, so it costs `O(nnz)` arithmetic and
/// no extra I/O. It exists because two of the four statistics do NOT survive a
/// fold applied afterwards:
///
/// - `s2` of a sum is not the sum of `s2`. The cross term `2ab` is gone, so a
///   variance rebuilt from per-row squares understates a bucket whose rows
///   covary, which two splice tracks of one gene certainly do.
/// - `npos` of a sum is not the sum of `npos`. A cell detected on both rows of
///   a bucket is one detection, not two, and the difference is exactly the
///   bimodality that breaks a detection-based cutoff on a channelized matrix.
///
/// `s1` and the mean do survive, which is why a caller that needs only those
/// can fold a row-axis result afterwards instead of asking for this.
pub fn streaming_sparse_running_stats_folded(
    data_vec: &SparseIoVec,
    block_size: Option<usize>,
    progress_label: &str,
    row_to_gene: &[u32],
    n_folded: usize,
) -> anyhow::Result<(SparseRunningStatistics<f32>, SparseRunningStatistics<f32>)> {
    let n_features = data_vec.num_rows();
    anyhow::ensure!(
        row_to_gene.len() == n_features,
        "row_to_gene has {} entries for a {}-row matrix",
        row_to_gene.len(),
        n_features
    );
    if let Some(&bad) = row_to_gene.iter().max() {
        anyhow::ensure!(
            (bad as usize) < n_folded,
            "row_to_gene names bucket {} but only {} were declared",
            bad,
            n_folded
        );
    }

    let n_total = data_vec.num_columns();
    let jobs = generate_minibatch_intervals(n_total, n_features, block_size);

    let prog_bar = legume_numeric::matrix::progress::new_progress_bar(jobs.len() as u64)
        .with_message(format!("{progress_label} blocks"));

    // The dense scratch lives in the fold accumulator, so it is allocated once
    // per leaf of the fold rather than once per column, and it is cleared in
    // O(touched) rather than O(n_folded).
    // The accumulator carries a `seen` stamp alongside the value buffer. Using
    // `buf[g] == 0.0` as the "not yet touched" marker would be wrong: a stored
    // explicit zero leaves the buffer at zero, so the bucket is pushed onto
    // `touched` a second time when its other row arrives, and that bucket's
    // npos, s1 and s2 are then all counted twice for that column. Sparse
    // backends do not promise to prune explicit zeros, and a producer that
    // emits both splice tracks whenever a gene is detected writes exactly that
    // pattern.
    type Acc = (
        SparseRunningStatistics<f32>,
        SparseRunningStatistics<f32>,
        Vec<f32>,
        Vec<bool>,
        Vec<usize>,
        Vec<f32>,
    );

    let (row_stats, folded_stats, ..) = fold_in_fixed_order(
        &jobs,
        blocks_per_leaf(jobs.len()),
        &|| -> Acc {
            (
                SparseRunningStatistics::<f32>::new(n_features),
                SparseRunningStatistics::<f32>::new(n_folded),
                vec![0.0f32; n_folded],
                vec![false; n_folded],
                Vec::<usize>::new(),
                Vec::<f32>::new(),
            )
        },
        &|(rows, folded, buf, seen, touched, vals), &(lb, ub)| {
            let chunk = data_vec.read_columns_csc(lb..ub)?;
            rows.add_csc(&chunk);
            for col in chunk.col_iter() {
                touched.clear();
                for (&r, &v) in col.row_indices().iter().zip(col.values().iter()) {
                    let g = row_to_gene[r] as usize;
                    if !seen[g] {
                        seen[g] = true;
                        touched.push(g);
                    }
                    buf[g] += v;
                }
                // `add_sparse_column` accumulates per row and never reads
                // the order, so the touched list is left unsorted.
                vals.clear();
                vals.extend(touched.iter().map(|&g| buf[g]));
                folded.add_sparse_column(touched, vals);
                for &g in touched.iter() {
                    buf[g] = 0.0;
                    seen[g] = false;
                }
            }
            prog_bar.inc(1);
            Ok(())
        },
        &|(ra, fa, ..), (rb, fb, ..)| {
            ra.merge(&rb);
            fa.merge(&fb);
        },
    )?;
    prog_bar.finish_and_clear();

    Ok((row_stats, folded_stats))
}

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