bixverse-rs 0.6.4

//! Implementation of the VISION framework to score spatially correlated
//! gene sets. See DeTomaso, et al., Nat. Commun., 2019

use extendr_api::{Conversions, List};
use rayon::prelude::*;
use std::time::Instant;

use crate::core::math::vector_helpers::rank_vector;
use crate::prelude::*;
use crate::single_cell::sc_processing::knn::knn_distance_weights;

////////////
// VISION //
////////////

/// Structure to store the indices of the SignatureGenes
#[derive(Clone, Debug)]
pub struct SignatureGenes {
    /// The gene indices of the positive genes
    pub positive: Vec<usize>,
    /// The gene indices of the negative genes
    pub negative: Vec<usize>,
}

/// Helper function to transform an R gene set list to `Vec<SignatureGenes>`
///
/// ### Params
///
/// * `gs_list` - Initial R list with the gene sets for VISION
///
/// ### Returns
///
/// The vector of SignatureGenes
pub fn r_list_to_sig_genes(gs_list: List) -> extendr_api::Result<Vec<SignatureGenes>> {
    let mut gene_signatures: Vec<SignatureGenes> = Vec::with_capacity(gs_list.len());

    for i in 0..gs_list.len() {
        let r_obj = gs_list.elt(i)?;
        let gs_list_i = r_obj.as_list().ok_or_else(|| {
            extendr_api::Error::from(
                "The lists in the gs_list could not be converted. Please check!",
            )
        })?;
        gene_signatures.push(SignatureGenes::from_r_list(gs_list_i)?);
    }

    Ok(gene_signatures)
}

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

/// Calculate VISION signature scores for a single cell
///
/// The cell-level mean and variance run over all `total_genes`, including the
/// implicit zeros. Zeros contribute nothing to either sum, so dividing the sums
/// over the stored values by `total_genes` is exactly that.
///
/// Both accumulate in one pass. Splitting them costs a second `f16` decode of
/// every stored value, which is what the loop is bound by, so fusing them
/// halves the per-cell statistics. The accumulators are `f64` to match
/// [`crate::single_cell::mc_analysis::vision_mc`]; on this path that buys no
/// accuracy, since the `f16` storage of `data_norm` already sets the floor an
/// order of magnitude above the accumulation error, and it costs about one per
/// cent.
///
/// ### Params
///
/// * `cell` - The CsrCellChunk
/// * `signatures` - Slice of `SignatureGenes` to calculate the scores for
/// * `total_genes` - Total number of represented genes
/// * `scratch` - Zeroed dense buffer of length `total_genes`; the stored
///   values are scattered into it and cleared again before returning
///
/// ### Returns
///
/// A Vec<f32> with a score for each of the `SignatureGenes` that were supplied
fn calculate_vision_scores_for_cell(
    cell: &CsrCellChunk,
    signatures: &[SignatureGenes],
    total_genes: usize,
    scratch: &mut [f32],
) -> Vec<f32> {
    for (&idx, value) in cell.indices.iter().zip(cell.data_norm.iter()) {
        scratch[idx as usize] = value.to_f32();
    }
    let get_expr = |gene_idx: usize| -> f64 { scratch[gene_idx] as f64 };

    // Cell-level statistics (ALL genes including zeros)
    let mut sum = 0.0_f64;
    let mut sum_sq = 0.0_f64;
    for value in &cell.data_norm {
        let value = value.to_f32() as f64;
        sum += value;
        sum_sq += value * value;
    }

    let mu_cell = sum / total_genes as f64;
    // clamped so a residual negative variance yields zero rather than a NaN
    // score, which would poison the ranks in `calc_autocorr_with_clusters`
    let sigma_cell = ((sum_sq / total_genes as f64) - (mu_cell * mu_cell))
        .max(0.0)
        .sqrt();

    // Score signatures
    let scores = signatures
        .iter()
        .map(|sig| {
            let signature_size = (sig.positive.len() + sig.negative.len()) as f64;

            if signature_size == 0.0 || sigma_cell == 0.0 {
                return 0.0;
            }

            let sum_pos: f64 = sig.positive.iter().map(|&idx| get_expr(idx)).sum();
            let sum_neg: f64 = sig.negative.iter().map(|&idx| get_expr(idx)).sum();

            let mean_sig = (sum_pos - sum_neg) / signature_size;

            // R "znorm_columns" formula
            ((mean_sig - mu_cell) / sigma_cell) as f32
        })
        .collect();

    for &idx in &cell.indices {
        scratch[idx as usize] = 0.0;
    }

    scores
}

/// Signatures transposed per tile, so that each cell is read one cache line
/// at a time rather than once per signature
const CONSISTENCY_TILE: usize = 64;

/// Geary's C based consistency (`1 - C`) of every signature column
///
/// The scores are stored cell-major, so a signature is a strided gather over
/// all cells. Signatures are processed in tiles: each tile is transposed in
/// one sequential pass over the cells and the signatures of the tile are then
/// ranked and scored in parallel.
///
/// ### Params
///
/// * `scores` - Cells x signatures scores, one `Vec` per cell
/// * `n_signatures` - Number of signature columns
/// * `knn_indices` - Neighbour indices per cell
/// * `knn_weights` - Neighbour weights matching `knn_indices`
///
/// ### Returns
///
/// `1 - C` per signature, in column order
fn consistency_per_signature(
    scores: &[Vec<f32>],
    n_signatures: usize,
    knn_indices: &[Vec<usize>],
    knn_weights: &[Vec<f32>],
) -> Vec<f64> {
    let n_cells = scores.len();
    let mut out = Vec::with_capacity(n_signatures);

    for start in (0..n_signatures).step_by(CONSISTENCY_TILE) {
        let width = CONSISTENCY_TILE.min(n_signatures - start);
        let mut tile: Vec<Vec<f32>> = (0..width).map(|_| Vec::with_capacity(n_cells)).collect();
        for cell in scores {
            for (col, &v) in tile.iter_mut().zip(&cell[start..start + width]) {
                col.push(v);
            }
        }

        let tile_res: Vec<f64> = tile
            .par_iter()
            .map(|col| {
                let ranks = rank_vector(col);
                1.0 - geary_c(&ranks.r_float_convert(), knn_indices, knn_weights)
            })
            .collect();
        out.extend(tile_res);
    }

    out
}

/// Calculate Geary's C for a single signature (pathway)
///
/// This implements a modified approach akin to VISION.
///
/// ### Params
///
/// * `scores` - Ranked signature scores for all cells
/// * `knn_indices` - KNN indices matrix (cells x k)
/// * `knn_weights` - KNN weights matrix (cells x k)
///
/// ### Returns
///
/// Geary's C statistic
fn geary_c(scores: &[f64], knn_indices: &[Vec<usize>], knn_weights: &[Vec<f32>]) -> f64 {
    let n = scores.len();

    let mean: f64 = scores.iter().sum::<f64>() / n as f64;
    let variance: f64 = scores.iter().map(|x| (x - mean).powi(2)).sum::<f64>();

    if variance == 0.0 {
        return 0.0;
    }

    let mut numerator = 0.0;
    let mut total_weight = 0.0;

    // unsafe unchecked access in hot loop
    for (i, (indices, weights)) in knn_indices.iter().zip(knn_weights.iter()).enumerate() {
        let xi = unsafe { *scores.get_unchecked(i) };
        for (&j, &w) in indices.iter().zip(weights.iter()) {
            let xj = unsafe { *scores.get_unchecked(j) };
            numerator += w as f64 * (xi - xj).powi(2);
            total_weight += w as f64;
        }
    }

    let norm = 2.0 * total_weight * variance / (n as f64 - 1.0);

    numerator / norm
}

/// Calculate KNN weights using the VISION exponential kernel
///
/// `exp(-d^2 / sigma^2)` with `sigma` the distance to the furthest of the `k`
/// neighbours, row-normalised so every cell contributes the same total weight.
/// That is `computeKNNWeights` in `VISION/R/Projections.R`, and it is
/// [`knn_distance_weights`] at `neighborhood_factor = 1`: the width neighbour
/// is `ceil(k / 1) = k`, i.e. the furthest one.
///
/// The row normalisation matters. [`geary_c`] divides by the total weight, so
/// a global rescaling would cancel, but normalising per row does not: it is
/// what stops cells whose neighbours sit at similar distances from carrying
/// more of the statistic than cells with one close neighbour and a long tail.
///
/// ### Params
///
/// * `knn_distances` - KNN distances (cells x k), ascending
///
/// ### Returns
///
/// KNN weights matrix (cells x k), each row summing to one
///
/// ### References
///
/// DeTomaso, et al., Nat. Commun., 2019
fn calc_knn_weights(knn_distances: &[Vec<f32>]) -> Vec<Vec<f32>> {
    knn_distance_weights(knn_distances, 1.0)
}

//////////
// Main //
//////////

/// Calculate VISION signature scores across a set of cells
///
/// ### Params
///
/// * `reader` - Reader over the cell-based count store.
/// * `gene_signs` - Slice of `SignatureGenes` to calculate the scores for
/// * `cells_to_keep` - Vector of indices with the cells to keep.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A `Vec<Vec<f32>>` with cells x scores per gene set (pair)
pub fn calculate_vision<S: SingleCellReading>(
    reader: &S,
    gene_signs: &[SignatureGenes],
    cells_to_keep: &[usize],
    verbose: usize,
) -> Result<Vec<Vec<f32>>, BixverseErrors> {
    let verbosity = parse_verbosity_level(verbose);

    let start_read = Instant::now();
    let no_genes = reader.get_header().total_genes;
    let cell_chunks: Vec<CsrCellChunk> = reader.read_cells_parallel(cells_to_keep)?;
    let end_read = start_read.elapsed();

    if verbosity.normal_verbosity() {
        println!("Loaded in data: {:.2?}", end_read);
    }

    let start_signatures = Instant::now();
    let signature_scores: Vec<Vec<f32>> = cell_chunks
        .par_iter()
        .map_init(
            || vec![0.0_f32; no_genes],
            |scratch, chunk| calculate_vision_scores_for_cell(chunk, gene_signs, no_genes, scratch),
        )
        .collect();
    let end_signatures = start_signatures.elapsed();

    if verbosity.normal_verbosity() {
        println!("Calculated VISION scores: {:.2?}", end_signatures);
    }

    // cells x signatures
    Ok(signature_scores)
}

/// Calculate VISION signature scores across a set of cells (streaming)
///
/// The streaming version of the function.
///
/// ### Params
///
/// * `reader` - Reader over the cell-based count store.
/// * `gene_signs` - Slice of `SignatureGenes` to calculate the scores for
/// * `cells_to_keep` - Vector of indices with the cells to keep.
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// A `Vec<Vec<f32>>` with cells x scores per gene set (pair)
pub fn calculate_vision_streaming<S: SingleCellReading>(
    reader: &S,
    gene_signs: &[SignatureGenes],
    cells_to_keep: &[usize],
    verbose: usize,
) -> Result<Vec<Vec<f32>>, BixverseErrors> {
    const CHUNK_SIZE: usize = 50000;

    let verbosity = parse_verbosity_level(verbose);

    let total_chunks = cells_to_keep.len().div_ceil(CHUNK_SIZE);
    let no_genes = reader.get_header().total_genes;

    let mut all_results: Vec<Vec<f32>> = Vec::with_capacity(cells_to_keep.len());

    for (chunk_idx, cell_indices_chunk) in cells_to_keep.chunks(CHUNK_SIZE).enumerate() {
        let start_chunk = Instant::now();

        let cell_chunks = reader.read_cells_parallel(cell_indices_chunk)?;

        let chunk_scores: Vec<Vec<f32>> = cell_chunks
            .par_iter()
            .map_init(
                || vec![0.0_f32; no_genes],
                |scratch, chunk| {
                    calculate_vision_scores_for_cell(chunk, gene_signs, no_genes, scratch)
                },
            )
            .collect();

        all_results.extend(chunk_scores);

        if verbosity.detailed_verbosity() {
            let elapsed = start_chunk.elapsed();
            let pct_complete = ((chunk_idx + 1) as f32 / total_chunks as f32) * 100.0;
            println!(
                "Processing chunk {} out of {} (took {:.2?}, completed {:.1}%)",
                chunk_idx + 1,
                total_chunks,
                elapsed,
                pct_complete
            );
        }
    }

    Ok(all_results)
}

/// Calculate VISION local autocorrelation scores
///
/// ### Params
///
/// * `pathway_scores` - Vector representing the actual vision scores:
///   `cells x pathways`
/// * `random_scores_by_cluster` - Vector representing the random scores by
///   cluster: `clusters -> (cells x sigs)`
/// * `cluster_membership` - Vector representing to which cluster a given
///   gene set belongs.
/// * `knn_indices` - KNN indices from embedding (cells x k)
/// * `knn_distances` - KNN distances (cells x k), ascending
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Tuple of (consistency_scores, p_values) for each pathway
pub fn calc_autocorr_with_clusters(
    pathway_scores: &[Vec<f32>],
    random_scores_by_cluster: &[Vec<Vec<f32>>],
    cluster_membership: &[usize],
    knn_indices: Vec<Vec<usize>>,
    knn_distances: Vec<Vec<f32>>,
    verbose: usize,
) -> (Vec<f64>, Vec<f64>) {
    let verbosity = parse_verbosity_level(verbose);

    let start = Instant::now();

    let knn_weights = calc_knn_weights(&knn_distances);

    if verbosity.normal_verbosity() {
        println!("Computed KNN weights: {:.2?}", start.elapsed());
    }

    let n_pathways = pathway_scores[0].len();

    // calculate Geary's C for actual pathways
    let pathway_consistency =
        consistency_per_signature(pathway_scores, n_pathways, &knn_indices, &knn_weights);

    if verbosity.normal_verbosity() {
        println!("Calculated pathway consistency: {:.2?}", start.elapsed());
    }

    let cluster_bg_consistency: Vec<Vec<f64>> = random_scores_by_cluster
        .iter()
        .map(|cluster_scores| {
            let n_random = cluster_scores[0].len();

            consistency_per_signature(cluster_scores, n_random, &knn_indices, &knn_weights)
        })
        .collect();

    // calculate p vals
    let p_vals: Vec<f64> = pathway_consistency
        .iter()
        .enumerate()
        .map(|(i, &fg_c)| {
            let cluster_idx = cluster_membership[i];
            let bg_dist = &cluster_bg_consistency[cluster_idx];

            let n = bg_dist.len();
            let num_greater_equal = bg_dist.iter().filter(|&&bg_c| bg_c >= fg_c).count();

            (num_greater_equal + 1) as f64 / (n + 1) as f64
        })
        .collect();

    if verbosity.normal_verbosity() {
        println!("Calculated p-values: {:.2?}", start.elapsed());
    }

    (pathway_consistency, p_vals)
}

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

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

    /// Values from `VISION/R/Projections.R::computeKNNWeights` run on the same
    /// input: `sigma = rowMaxs(d)`, `exp(-1 * (d * d) / sigma^2)`, then divided
    /// by `rowSums`. The row normalisation is the part that was missing.
    #[test]
    fn test_vision_kernel_matches_the_r_reference() {
        let distances = vec![
            vec![0.5, 1.0, 2.0, 4.0],
            vec![1.0, 1.0, 1.0, 1.0],
            vec![0.0, 0.0, 0.0, 0.0],
        ];

        let expected = [
            vec![0.320_621_29, 0.305_938_97, 0.253_632_3, 0.119_807_42],
            vec![0.25, 0.25, 0.25, 0.25],
            vec![0.25, 0.25, 0.25, 0.25],
        ];

        let weights = calc_knn_weights(&distances);

        for (got, want) in weights.iter().zip(expected.iter()) {
            for (a, b) in got.iter().zip(want.iter()) {
                assert_relative_eq!(a, b, epsilon = 1e-6);
            }
            assert_relative_eq!(got.iter().sum::<f32>(), 1.0, epsilon = 1e-6);
        }
    }
}