bixverse-rs 0.4.3

//! Helpers around pseudo-bulking and meta-cell aggregations of single cells.

use faer::Mat;
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use std::time::Instant;

use crate::prelude::*;

/////////////
// Helpers //
/////////////

/// Helper function to aggregate the meta cells
///
/// The function will generate the metacells based on the provided indices.
/// Per meta-cell it will aggregate the raw counts and recalculate the norm
/// counts based on the aggregated counts.
///
/// ### Params
///
/// * `reader` - The reader structure to get the cells from disk
/// * `meta_cells` - The indices of the meta cells
/// * `target_size` - Float defining the target size for the normalisation
///   procedure. Usually defaults to `1e4` in single cell.
/// * `n_genes` - Total number of genes in the data
///
/// ### Return
///
/// `CompressedSparseData2` in CSR format with aggregated raw counts and re-
/// normalised counts per meta cell.
pub fn aggregate_meta_cells<S: SingleCellReading>(
    reader: &S,
    metacells: &[&[usize]],
    target_size: f32,
    n_genes: usize,
) -> Result<CompressedSparseData2<u32, f32>, BixverseErrors> {
    let n_metacells = metacells.len();
    let mut all_data: Vec<u32> = Vec::new();
    let mut all_data_norm: Vec<f32> = Vec::new();
    let mut all_indices: Vec<usize> = Vec::new();
    let mut all_indptr: Vec<usize> = vec![0];
    const CHUNK_SIZE: usize = 1000;

    for chunk_start in (0..n_metacells).step_by(CHUNK_SIZE) {
        let chunk_end = (chunk_start + CHUNK_SIZE).min(n_metacells);
        let chunk = &metacells[chunk_start..chunk_end];

        let results: Vec<(Vec<usize>, Vec<u32>, Vec<f32>)> = chunk
            .par_iter()
            .map(|cell_idx| -> Result<_, BixverseErrors> {
                let cells = reader.read_cells_parallel(cell_idx)?;
                let mut gene_counts: FxHashMap<usize, u32> = FxHashMap::default();
                let mut library_size: u32 = 0;
                for cell in &cells {
                    for (idx, count) in cell.indices.iter().zip(cell.data_raw.iter()) {
                        *gene_counts.entry(*idx as usize).or_insert(0) += count;
                        library_size += count;
                    }
                }
                let mut entries: Vec<(usize, u32)> = gene_counts.into_iter().collect();
                entries.sort_by_key(|(idx, _)| *idx);
                let indices: Vec<usize> = entries.iter().map(|(idx, _)| *idx).collect();
                let raw_counts: Vec<u32> = entries.iter().map(|(_, count)| *count).collect();
                let norm_counts: Vec<f32> = entries
                    .iter()
                    .map(|(_, count)| {
                        let norm = (*count as f32 / library_size as f32) * target_size;
                        (norm + 1.0).ln()
                    })
                    .collect();
                Ok((indices, raw_counts, norm_counts))
            })
            .collect::<Result<Vec<_>, _>>()?;

        for (indices, raw_counts, norm_counts) in results {
            all_indices.extend(indices);
            all_data.extend(raw_counts);
            all_data_norm.extend(norm_counts);
            all_indptr.push(all_indices.len());
        }
    }

    Ok(CompressedSparseData2::new_csr(
        &all_data,
        &all_indices.index_cast(),
        &all_indptr.index_cast(),
        Some(&all_data_norm),
        (n_metacells, n_genes),
    ))
}

/// Convert metacell groups to flat assignments, handling unassigned cells
///
/// ### Params
///
/// * `metacells` - Vector of cell groups (metacell → `[cells]`)
/// * `n_cells` - Total number of cells
///
/// ### Returns
///
/// Flat assignment vector where `assignments[cell_id] = Some(metacell_id)`
/// or None if cell is unassigned
pub fn metacells_to_assignments(metacells: &[&[usize]], n_cells: usize) -> Vec<Option<usize>> {
    let mut assignments = vec![None; n_cells];

    for (metacell_id, &cells) in metacells.iter().enumerate() {
        for &cell_id in cells {
            if cell_id < n_cells {
                assignments[cell_id] = Some(metacell_id);
            }
        }
    }

    assignments
}

/// Remap subset assignments back to original index space
///
/// Takes metacell assignments computed on a subset of cells and maps them
/// back to the full original index space. Cells not in the subset will have
/// `None` assignments.
///
/// ### Params
///
/// * `subset_assignments` - Vector of metacell assignments in subset index space
/// * `subset_to_orig` - Mapping from subset indices to original indices
/// * `n_total` - Total number of cells in original space
///
/// ### Return
///
/// Vector of assignments in original index space with `None` for cells not
/// in the subset
pub fn remap_assignments_to_original(
    subset_assignments: &[Option<usize>],
    subset_to_orig: &[usize],
    n_total: usize,
) -> Vec<Option<usize>> {
    let mut full_assignments = vec![None; n_total];
    for (subset_idx, &metacell_id) in subset_assignments.iter().enumerate() {
        if let Some(orig_idx) = subset_to_orig.get(subset_idx) {
            full_assignments[*orig_idx] = metacell_id;
        }
    }
    full_assignments
}

/// Remap metacell indices from subset space to original space
///
/// Takes metacell groups where cell indices are in subset space and transforms
/// all indices back to original space. This is used when metacells are computed
/// on a subset of cells but need to be aggregated from the full dataset.
///
/// ### Params
///
/// * `metacells` - Vector of metacell groups with cell indices in subset space
/// * `subset_to_orig` - Mapping from subset indices to original indices
///
/// ### Return
///
/// Vector of metacell groups with cell indices in original space
pub fn remap_metacells_to_original(
    metacells: &[&[usize]],
    subset_to_orig: &[usize],
) -> Vec<Vec<usize>> {
    metacells
        .iter()
        .map(|&cells| cells.iter().map(|&idx| subset_to_orig[idx]).collect())
        .collect()
}

////////////////////
// Pseudo-bulking //
////////////////////

/// Enum for Pseudo-bulking
#[derive(Debug, Clone, Default)]
pub enum PseudoBulk {
    #[default]
    /// Shall raw counts be pseudo-bulked
    Raw,
    /// Shall normalised counts be pseudo-bulked
    Norm,
}

/// Helper function to parse pseudo-bulk type
///
/// ### Params
///
/// * `s` - Type of pseudo-bulk to perform
///
/// ### Returns
///
/// Option of the PseudoBulk enum
pub fn parse_pseudo_bulk(s: &str) -> Option<PseudoBulk> {
    match s.to_lowercase().as_str() {
        "raw" => Some(PseudoBulk::Raw),
        "norm" | "normalised" | "normalized" => Some(PseudoBulk::Norm),
        _ => None,
    }
}

/// Pseudo-bulk data across cells based on cell indices (dense output)
///
/// ### Params
///
/// * `reader` - Reader for the cell-based store.
/// * `cell_indices` - Slice of indices to pseudo-bulk.
/// * `bulk_type` - Whether to pseudo-bulk raw (sum) or normalised (average)
///   counts.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Dense matrix of samples x genes pseudo-bulked.
pub fn get_pseudo_bulked_counts_dense<S: SingleCellReading>(
    reader: &S,
    cell_indices: &[Vec<usize>],
    bulk_type: PseudoBulk,
    verbose: usize,
) -> Result<Mat<f64>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let n_genes = reader.get_header().total_genes;
    let n_groups = cell_indices.len();
    let mut result = Mat::zeros(n_groups, n_genes);

    for (group_idx, indices) in cell_indices.iter().enumerate() {
        let start_group = Instant::now();
        let chunks = reader.read_cells_parallel(indices)?;
        let n_cells = indices.len() as f64;

        for chunk in chunks {
            match bulk_type {
                PseudoBulk::Raw => {
                    for (value, &gene_idx) in chunk.data_raw.iter().zip(chunk.indices.iter()) {
                        result[(group_idx, gene_idx as usize)] += value as f64;
                    }
                }
                PseudoBulk::Norm => {
                    for (value, &gene_idx) in chunk.data_norm.iter().zip(chunk.indices.iter()) {
                        result[(group_idx, gene_idx as usize)] += value.to_f64();
                    }
                }
            }
        }

        if matches!(bulk_type, PseudoBulk::Norm) {
            for gene_idx in 0..n_genes {
                result[(group_idx, gene_idx)] /= n_cells;
            }
        }

        if verbosity.normal_verbosity() && (group_idx + 1) % 10 == 0 {
            let elapsed = start_group.elapsed();
            let pct_complete = ((group_idx + 1) as f32 / n_groups as f32) * 100.0;
            println!(
                "Processed group {} out of {} (took {:.2?}, completed {:.1}%)",
                group_idx + 1,
                n_groups,
                elapsed,
                pct_complete
            );
        }
    }

    Ok(result)
}

/// Pseudo-bulk data across cells based on cell indices (sparse CSR output)
///
/// ### Params
///
/// * `reader` - Reader for the cell-based store.
/// * `cell_indices` - Slice of indices to pseudo-bulk.
/// * `bulk_type` - Whether to pseudo-bulk raw (sum) or normalised (average)
///   counts.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Sparse CSR matrix of samples x genes pseudo-bulked.
pub fn get_pseudo_bulked_counts_sparse<S: SingleCellReading>(
    reader: &S,
    cell_indices: &[Vec<usize>],
    bulk_type: PseudoBulk,
    verbose: usize,
) -> Result<CompressedSparseData2<f64>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let n_genes = reader.get_header().total_genes;
    let n_groups = cell_indices.len();
    let mut row_data: Vec<FxHashMap<usize, f64>> = vec![FxHashMap::default(); n_groups];

    for (group_idx, indices) in cell_indices.iter().enumerate() {
        let start_group = Instant::now();
        let chunks = reader.read_cells_parallel(indices)?;
        let n_cells = indices.len() as f64;

        for chunk in chunks {
            match bulk_type {
                PseudoBulk::Raw => {
                    for (value, &gene_idx) in chunk.data_raw.iter().zip(chunk.indices.iter()) {
                        *row_data[group_idx].entry(gene_idx as usize).or_insert(0.0) +=
                            value as f64;
                    }
                }
                PseudoBulk::Norm => {
                    for (value, &gene_idx) in chunk.data_norm.iter().zip(chunk.indices.iter()) {
                        *row_data[group_idx].entry(gene_idx as usize).or_insert(0.0) +=
                            value.to_f64();
                    }
                }
            }
        }

        if matches!(bulk_type, PseudoBulk::Norm) {
            for value in row_data[group_idx].values_mut() {
                *value /= n_cells;
            }
        }

        if verbosity.normal_verbosity() && (group_idx + 1) % 10 == 0 {
            let elapsed = start_group.elapsed();
            let pct_complete = ((group_idx + 1) as f32 / n_groups as f32) * 100.0;
            println!(
                "Processed group {} out of {} (took {:.2?}, completed {:.1}%)",
                group_idx + 1,
                n_groups,
                elapsed,
                pct_complete
            );
        }
    }

    let mut data = Vec::new();
    let mut indices = Vec::new();
    let mut indptr = vec![0];

    for row_map in row_data {
        let mut sorted_entries: Vec<_> = row_map.into_iter().collect();
        sorted_entries.sort_by_key(|(idx, _)| *idx);

        for (idx, value) in sorted_entries {
            data.push(value);
            indices.push(idx);
        }
        indptr.push(data.len());
    }

    Ok(CompressedSparseData2 {
        data,
        indices: indices.index_cast(),
        indptr: indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csr,
        data_2: None,
        shape: (n_groups, n_genes),
    })
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;

    //////////////////////
    // Index remapping //
    //////////////////////

    /// Cell ids at or beyond `n_cells` are dropped without a word. Pinned so a
    /// caller passing the wrong `n_cells` cannot start silently losing cells.
    #[test]
    fn metacells_to_assignments_drops_out_of_range_cell_ids() {
        let mc0: &[usize] = &[0, 5];
        let mc1: &[usize] = &[2, 99];
        let assignments = metacells_to_assignments(&[mc0, mc1], 3);

        assert_eq!(assignments, vec![Some(0), None, Some(1)]);
    }

    /// On overlap the last metacell in the slice wins. The function has no
    /// disjointness check, so the ordering is the whole contract.
    #[test]
    fn metacells_to_assignments_lets_later_metacells_overwrite() {
        let mc0: &[usize] = &[0, 1];
        let mc1: &[usize] = &[1, 2];
        let assignments = metacells_to_assignments(&[mc0, mc1], 3);

        assert_eq!(assignments, vec![Some(0), Some(1), Some(1)]);
    }

    /// Cells outside the subset stay `None`, and the mapping is by position in
    /// `subset_to_orig`, not by value.
    #[test]
    fn remap_assignments_to_original_scatters_by_subset_position() {
        let subset = vec![Some(0), None, Some(1)];
        let subset_to_orig = [4, 2, 0];
        let full = remap_assignments_to_original(&subset, &subset_to_orig, 6);

        assert_eq!(full, vec![Some(1), None, None, None, Some(0), None]);
    }

    /// A `subset_assignments` longer than `subset_to_orig` is truncated rather
    /// than panicking, because the lookup goes through `.get`.
    #[test]
    fn remap_assignments_to_original_ignores_trailing_assignments() {
        let subset = vec![Some(0), Some(1), Some(2)];
        let subset_to_orig = [1];
        let full = remap_assignments_to_original(&subset, &subset_to_orig, 3);

        assert_eq!(full, vec![None, Some(0), None]);
    }

    /// Straight index translation, group structure untouched.
    #[test]
    fn remap_metacells_to_original_translates_every_index() {
        let mc0: &[usize] = &[0, 2];
        let mc1: &[usize] = &[1];
        let remapped = remap_metacells_to_original(&[mc0, mc1], &[10, 20, 30]);

        assert_eq!(remapped, vec![vec![10, 30], vec![20]]);
    }

    /// Defect pin: unlike its sibling this one indexes `subset_to_orig`
    /// directly, so a stale subset index aborts the process instead of
    /// returning an error.
    #[test]
    #[should_panic(expected = "index out of bounds")]
    fn remap_metacells_to_original_panics_on_a_stale_index() {
        let mc0: &[usize] = &[0, 5];
        let _ = remap_metacells_to_original(&[mc0], &[10, 20]);
    }

    //////////////////////
    // Pseudo-bulk enum //
    //////////////////////

    /// The string aliases are the R-facing contract and break silently on a
    /// rename, including the British/American spelling pair.
    #[test]
    fn parse_pseudo_bulk_accepts_both_spellings_and_any_case() {
        assert!(matches!(parse_pseudo_bulk("raw"), Some(PseudoBulk::Raw)));
        assert!(matches!(parse_pseudo_bulk("RAW"), Some(PseudoBulk::Raw)));
        assert!(matches!(parse_pseudo_bulk("norm"), Some(PseudoBulk::Norm)));
        assert!(matches!(
            parse_pseudo_bulk("normalised"),
            Some(PseudoBulk::Norm)
        ));
        assert!(matches!(
            parse_pseudo_bulk("Normalized"),
            Some(PseudoBulk::Norm)
        ));
    }

    /// Anything unrecognised falls back to `None` rather than defaulting to
    /// raw counts.
    #[test]
    fn parse_pseudo_bulk_rejects_unknown_labels() {
        assert!(parse_pseudo_bulk("").is_none());
        assert!(parse_pseudo_bulk("counts").is_none());
        assert!(parse_pseudo_bulk("normalise").is_none());
    }

    /// The derived `Default` is what an omitted R argument lands on.
    #[test]
    fn pseudo_bulk_defaults_to_raw() {
        assert!(matches!(PseudoBulk::default(), PseudoBulk::Raw));
    }
}