bixverse-sc 0.9.2

//! Implements the cxds method from Bais & Kostka, Bioinformatics,
//! 2020. The idea: heterotypic doublets co-express marker genes
//! that normally belong to different cell types. Gene pairs whose
//! co-expression is rarer than expected under independence receive
//! high scores. Cells that co-express many such unusual pairs are
//! likely doublets.

use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};

use crate::prelude::*;

////////////
// Consts //
////////////

/// Default number of genes to select for cxds scoring.
pub const CXDS_NTOP: usize = 500;

/// Cells decoded per block when fitting
const CXDS_READ_BLOCK: usize = 50_000;

/// Sentinel in the dense gene lookup for genes outside the selected set
const CXDS_NO_GENE: u16 = u16::MAX;

///////////
// Types //
///////////

/// Per-cell expressed gene set, indexed into the cxds gene set.
///
/// Stored as a sorted `Vec<u16>` since the gene set is at most
/// `CXDS_NTOP` genes wide.
type CellGeneSet = Vec<u16>;

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

/// Sum of `pair_scores[i, j]` over all pairs `i < j` of a sorted gene set.
///
/// Four independent accumulators break the add dependency chain of the
/// gathers.
///
/// ### Params
///
/// * `pair_scores` - Symmetric `[ntop, ntop]` score matrix, row-major
/// * `ntop` - Row length
/// * `expressed` - Sorted gene positions expressed in the cell
///
/// ### Returns
///
/// The summed pair score
#[inline]
fn pair_score_sum(pair_scores: &[f32], ntop: usize, expressed: &[u16]) -> f32 {
    let mut total = 0.0f32;
    for (a, &gi) in expressed.iter().enumerate() {
        let i = gi as usize;
        let row = &pair_scores[i * ntop..(i + 1) * ntop];
        let rest = &expressed[a + 1..];
        let mut acc = [0.0f32; 4];
        let mut quads = rest.chunks_exact(4);
        for q in &mut quads {
            for (slot, &j) in acc.iter_mut().zip(q) {
                *slot += row[j as usize];
            }
        }
        let mut tail = 0.0f32;
        for &j in quads.remainder() {
            tail += row[j as usize];
        }
        total += (acc[0] + acc[1]) + (acc[2] + acc[3]) + tail;
    }
    total
}

/// Approximate the standard normal CDF.
///
/// Uses the Abramowitz & Stegun rational approximation (formula
/// 26.2.17) with accuracy ~1.5e-7.
///
/// ### Params
///
/// * `x` - Quantile.
///
/// ### Returns
///
/// P(Z <= x) for Z ~ N(0, 1).
fn norm_cdf(x: f64) -> f64 {
    if x < -8.0 {
        return 0.0;
    }
    if x > 8.0 {
        return 1.0;
    }
    let t = 1.0 / (1.0 + 0.231_641_9 * x.abs());
    let poly = t
        * (0.319_381_530
            + t * (-0.356_563_782
                + t * (1.781_477_937 + t * (-1.821_255_978 + t * 1.330_274_429))));
    let pdf = (-x * x / 2.0).exp() * std::f64::consts::FRAC_2_SQRT_PI * 0.5;
    let tail = pdf * poly;
    if x >= 0.0 { 1.0 - tail } else { tail }
}

/// Compute `-log(P(X >= k))` where `X ~ Binomial(n, p)`.
///
/// Uses the normal approximation with continuity correction, which is accurate
/// for `n * p * (1 - p) >= 5`. For degenerate cases (zero variance), returns
/// 0.0.
///
/// ### Params
///
/// * `k` - Observed count.
/// * `n` - Number of trials.
/// * `p` - Success probability under the null.
///
/// ### Returns
///
/// Negative log upper-tail p-value, clamped to `[0, 700]`.
fn neg_log_binom_upper(k: u32, n: u32, p: f64) -> f32 {
    let np = n as f64 * p;
    let var = np * (1.0 - p);
    if var < 1e-10 {
        return 0.0;
    }
    // Continuity correction: P(X >= k) ~= P(Z >= (k - 0.5 - np) / sd)
    let z = (k as f64 - 0.5 - np) / var.sqrt();
    let tail_p = 1.0 - norm_cdf(z);
    if tail_p < 1e-300 {
        return 700.0; // cap
    }
    (-tail_p.ln()) as f32
}

//////////
// cxds //
//////////

/// Precomputed cxds gene-pair score model.
///
/// Built from observed cells, then applied to both observed and
/// simulated cells to produce per-cell cxds doublet scores.
#[allow(dead_code)]
pub struct CxdsModel {
    /// Gene-pair score matrix, `ntop x ntop`, stored row-major.
    ///
    /// `pair_scores[i * ntop + j]` holds `-log(P)` for the gene
    /// pair `(i, j)`. Symmetric; only `i < j` entries are nonzero.
    pair_scores: Vec<f32>,
    /// Original gene indices of the selected cxds genes (into the
    /// full gene space).
    selected_genes: Vec<usize>,
    /// Number of selected genes.
    ntop: usize,
    /// Reverse lookup: original gene index -> position in
    /// `selected_genes`.
    gene_map: FxHashMap<usize, u16>,
}

impl CxdsModel {
    /// Build the cxds model from observed cells.
    ///
    /// Single-pass approach: reads each cell from disk once, buffering
    /// the expressed HVG indices in memory. After gene selection, the
    /// buffered indices are filtered to the selected set and used for
    /// parallel co-expression counting.
    ///
    /// ### Params
    ///
    /// * `reader` - Reader over the cell-based count store.
    /// * `cells_to_keep` - Observed cell indices.
    /// * `hvg_genes` - Highly variable gene indices (the candidate
    ///   pool for cxds gene selection).
    /// * `ntop` - Number of genes to select.
    ///
    /// ### Returns
    ///
    /// `(model, obs_gene_sets)` where `obs_gene_sets[i]` is the
    /// expressed gene set for observed cell `i`.
    pub fn fit<S: SingleCellReading>(
        reader: &S,
        cells_to_keep: &[usize],
        hvg_genes: &[usize],
        ntop: usize,
    ) -> Result<(Self, Vec<CellGeneSet>), ScErrors> {
        let n_cells = cells_to_keep.len();
        let hvg_set: FxHashSet<usize> = hvg_genes.iter().copied().collect();

        let mut gene_counts: FxHashMap<usize, u32> = FxHashMap::default();
        let mut cell_hvg_indices: Vec<Vec<usize>> = Vec::with_capacity(n_cells);

        for block in cells_to_keep.chunks(CXDS_READ_BLOCK) {
            let chunks = reader.read_cells_parallel(block)?;
            let expressed: Vec<Vec<usize>> = chunks
                .par_iter()
                .map(|chunk| {
                    chunk
                        .indices
                        .iter()
                        .enumerate()
                        .filter(|&(i, &g)| {
                            hvg_set.contains(&(g as usize)) && chunk.data_raw.get(i) > 0
                        })
                        .map(|(_, &g)| g as usize)
                        .collect()
                })
                .collect();

            for hvgs in expressed {
                for &gi in &hvgs {
                    *gene_counts.entry(gi).or_insert(0) += 1;
                }
                cell_hvg_indices.push(hvgs);
            }
        }

        // gene selection: top ntop by binomial variance p*(1-p)
        let nf = n_cells as f64;
        let mut gene_vars: Vec<(usize, f64)> = gene_counts
            .iter()
            .map(|(&gene, &count)| {
                let p = count as f64 / nf;
                (gene, p * (1.0 - p))
            })
            .collect();
        gene_vars.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        let actual_ntop = ntop.min(gene_vars.len());
        let selected_genes: Vec<usize> = gene_vars[..actual_ntop].iter().map(|&(g, _)| g).collect();

        let gene_map: FxHashMap<usize, u16> = selected_genes
            .iter()
            .enumerate()
            .map(|(pos, &g)| (g, pos as u16))
            .collect();

        let marginals: Vec<f64> = selected_genes
            .iter()
            .map(|g| *gene_counts.get(g).unwrap_or(&0) as f64 / nf)
            .collect();

        // build per-cell gene sets from buffered HVG indices
        let obs_gene_sets: Vec<CellGeneSet> = cell_hvg_indices
            .par_iter()
            .map(|hvgs| {
                let mut expressed: CellGeneSet = hvgs
                    .iter()
                    .filter_map(|&gi| gene_map.get(&gi).copied())
                    .collect();
                expressed.sort_unstable();
                expressed
            })
            .collect();

        // drop the buffered HVG indices; no longer needed -> force compiler to
        // liberate memory
        drop(cell_hvg_indices);

        // parallel co-expression counting
        let n_coexpress = obs_gene_sets
            .par_iter()
            .fold(
                || vec![0u32; actual_ntop * actual_ntop],
                |mut local_coex, expressed| {
                    let ne = expressed.len();
                    for a in 0..ne {
                        let i = expressed[a] as usize;
                        let row = i * actual_ntop;
                        for b in (a + 1)..ne {
                            let j = expressed[b] as usize;
                            local_coex[row + j] += 1;
                        }
                    }
                    local_coex
                },
            )
            .reduce(
                || vec![0u32; actual_ntop * actual_ntop],
                |mut a, b| {
                    for i in 0..a.len() {
                        a[i] += b[i];
                    }
                    a
                },
            );

        // build gene-pair score matrix
        //
        // n_exclusive[i][j] = n_express[i] + n_express[j] - 2 * n_coexpress[i][j]
        let n_express: Vec<u32> = selected_genes
            .iter()
            .map(|g| *gene_counts.get(g).unwrap_or(&0))
            .collect();

        let n = n_cells as u32;
        let mut pair_scores = vec![0.0f32; actual_ntop * actual_ntop];

        for i in 0..actual_ntop {
            for j in (i + 1)..actual_ntop {
                let pi = marginals[i];
                let pj = marginals[j];
                let p_excl = pi * (1.0 - pj) + (1.0 - pi) * pj;
                let coex = n_coexpress[i * actual_ntop + j];
                let k = n_express[i] + n_express[j] - 2 * coex;
                let score = neg_log_binom_upper(k, n, p_excl);
                pair_scores[i * actual_ntop + j] = score;
                pair_scores[j * actual_ntop + i] = score;
            }
        }

        let model = CxdsModel {
            pair_scores,
            selected_genes,
            ntop: actual_ntop,
            gene_map,
        };

        Ok((model, obs_gene_sets))
    }

    /// Compute cxds scores for cells given their expressed gene sets.
    ///
    /// For each cell, sums the pair scores `S[i][j]` over all pairs
    /// `(i, j)` where both genes `i` and `j` are expressed in the
    /// cell. Cells that co-express unusual gene combinations receive
    /// high scores.
    ///
    /// ### Params
    ///
    /// * `gene_sets` - Per-cell expressed gene sets (indices into
    ///   the cxds gene set).
    ///
    /// ### Returns
    ///
    /// Per-cell cxds scores.
    pub fn score(&self, gene_sets: &[CellGeneSet]) -> Vec<f32> {
        gene_sets
            .par_iter()
            .map(|expressed| pair_score_sum(&self.pair_scores, self.ntop, expressed))
            .collect()
    }

    /// Compute cxds scores for simulated doublets from their (noise-injected)
    /// raw count chunks.
    ///
    /// ### Params
    ///
    /// * `sim_chunks` - Simulated doublet sparse chunks with
    ///   HVG-remapped gene indices (i.e. chunk.indices are positions
    ///   into the HVG array). Raw counts should be post-noise.
    /// * `hvg_to_cxds_pos` - Map from HVG position to cxds gene
    ///   position. Entries missing mean the gene isn't in the cxds
    ///   gene set.
    /// * `bin_thresh` - Binarisation threshold (raw count >= thresh
    ///   counts as expressed). Match what was used in `fit`.
    ///
    /// ### Returns
    ///
    /// Per-doublet cxds scores.
    pub fn score_simulated_from_chunks(
        &self,
        sim_chunks: &[CsrCellChunk],
        selected_genes_to_cxds_pos: &FxHashMap<usize, u16>,
        bin_thresh: u32,
    ) -> Vec<f32> {
        let table_len = selected_genes_to_cxds_pos
            .keys()
            .max()
            .map_or(0, |&m| m + 1);
        let mut pos_table = vec![CXDS_NO_GENE; table_len];
        for (&gene, &pos) in selected_genes_to_cxds_pos {
            pos_table[gene] = pos;
        }

        sim_chunks
            .par_iter()
            .map(|chunk| {
                // Build the cell's cxds-indexed expressed gene set
                // from the raw counts, applying the binarisation
                // threshold.
                let mut expressed: Vec<u16> = Vec::new();
                for i in 0..chunk.indices.len() {
                    let raw = chunk.data_raw.get(i);
                    if raw < bin_thresh {
                        continue;
                    }
                    let cxds_pos = pos_table
                        .get(chunk.indices[i] as usize)
                        .copied()
                        .unwrap_or(CXDS_NO_GENE);
                    if cxds_pos != CXDS_NO_GENE {
                        expressed.push(cxds_pos);
                    }
                }
                expressed.sort_unstable();

                // score through the pair matrix (same as score())
                pair_score_sum(&self.pair_scores, self.ntop, &expressed)
            })
            .collect()
    }

    /// Build an selected gene-position to cxds-position lookup.
    ///
    /// ### Params
    ///
    /// * `selected_genes` - Selected gene indices (into original gene space),
    ///   parallel to selected genes positions.
    ///
    /// ### Returns
    ///
    /// Map from selected gene position to cxds gene-set position for genes
    /// that appear in both.
    pub fn hvg_position_map(&self, selected_genes: &[usize]) -> FxHashMap<usize, u16> {
        selected_genes
            .iter()
            .enumerate()
            .filter_map(|(hvg_pos, orig_gene)| {
                self.gene_map
                    .get(orig_gene)
                    .map(|&cxds_pos| (hvg_pos, cxds_pos))
            })
            .collect()
    }
}