bixverse-methods 0.9.1

//! Implementation of the RCisTarget algorithm from Aibar, et al., Nat Methods,
//! 2018.

use faer::MatRef;
use rayon::prelude::*;

use crate::prelude::*;
use bixverse_core::core::math::vector_helpers::standard_deviation;

////////////////
// Structures //
////////////////

/// Results structure for a single enriched motif
#[derive(Debug, Clone)]
pub struct MotifEnrichment<T> {
    /// Index of the motif
    pub motif_idx: usize,
    /// Normalised enrichment score
    pub nes: T,
    /// Area under the curve
    pub auc: T,
    /// Rank position at maximum enrichment
    pub rank_at_max: u32,
    /// Number of enriched genes
    pub n_enriched: usize,
    /// Indices of enriched genes in the leading edge
    pub enriched_gene_indices: Vec<usize>,
}

/// Enum for the RCC calculation types
#[derive(Clone, Debug, Default)]
pub enum RccType {
    /// Use the `ICisTarget` type for the RCC calculation
    #[default]
    ICisTarget,
    /// Use the `Approx` type for the RCC calculation
    Approx,
}

/// Parse the RCC calculation type
///
/// ### Params
///
/// * `s` - The string to parse
///
/// ### Returns
///
/// Option of the `RccType`
pub fn parse_rcc_type(s: &str) -> Option<RccType> {
    match s.to_lowercase().as_str() {
        "icistarget" => Some(RccType::ICisTarget),
        "approx" => Some(RccType::Approx),
        _ => None,
    }
}

/////////////////////////
// Helpers - AUC & NES //
/////////////////////////

/// Calculate AUC for a single motif's gene ranks
///
/// ### Params
///
/// * `gene_ranks` - Slice of gene ranks for this motif
/// * `auc_threshold` - Maximum rank to consider for AUC calculation
///
/// ### Returns
///
/// Normalised AUC score
fn calculate_auc_single<T>(gene_ranks: &[i32], auc_threshold: i32) -> T
where
    T: BixverseFloat,
{
    let mut filtered = gene_ranks
        .iter()
        .copied()
        .filter(|&r| r < auc_threshold)
        .collect::<Vec<i32>>();

    if filtered.is_empty() {
        return T::zero();
    }

    filtered.sort_unstable();
    filtered.push(auc_threshold);

    let n_genes = T::from_usize(gene_ranks.len()).unwrap();
    let max_auc = T::from_i32(auc_threshold).unwrap() * n_genes;

    let auc = filtered
        .windows(2)
        .enumerate()
        .map(|(i, w)| T::from_i32(w[1] - w[0]).unwrap() * T::from_usize(i + 1).unwrap())
        .fold(T::zero(), |acc, x| acc + x);

    auc / max_auc
}

/// Calculate AUC scores for all motifs given a gene set
///
/// ### Params
///
/// * `matrix` - Matrix reference containing gene-motif ranks
/// * `gene_indices` - Indices of genes in the gene set
/// * `auc_threshold` - Maximum rank to consider for AUC calculation
///
/// ### Returns
///
/// Vector of AUC scores for each motif
pub fn calculate_aucs<T>(
    matrix: MatRef<'_, i32>,
    gene_indices: &[usize],
    auc_threshold: i32,
) -> Vec<T>
where
    T: BixverseFloat,
{
    let n_motifs = matrix.ncols();

    (0..n_motifs)
        .into_par_iter()
        .map(|motif_idx| {
            let gene_ranks: Vec<i32> = gene_indices
                .iter()
                .map(|&row_idx| *matrix.get(row_idx, motif_idx))
                .collect();

            calculate_auc_single(&gene_ranks, auc_threshold)
        })
        .collect()
}

/// Calculate Normalised Enrichment Scores from AUC values
///
/// ### Params
///
/// * `auc_scores` - Slice of AUC scores
///
/// ### Returns
///
/// Vector of z-score normalised enrichment scores
pub fn calculate_nes<T>(auc_scores: &[T]) -> Vec<T>
where
    T: BixverseFloat + std::iter::Sum,
{
    let n = T::from_usize(auc_scores.len()).unwrap();
    let mean = auc_scores.iter().copied().fold(T::zero(), |acc, x| acc + x) / n;

    let std = standard_deviation(auc_scores);

    if std == T::zero() {
        return vec![T::zero(); auc_scores.len()];
    }

    auc_scores
        .par_iter()
        .map(|&auc| (auc - mean) / std)
        .collect()
}

/////////
// RCC //
/////////

/// Calculate rolling mean with centre alignment and extend fill
///
/// ### Params
///
/// * `values` - Slice of values to smooth
/// * `window` - Window size for rolling mean
///
/// ### Returns
///
/// Vector of smoothed values
fn rolling_mean<T: BixverseFloat>(values: &[T], window: usize) -> Vec<T> {
    let n = values.len();
    let mut result = Vec::with_capacity(n);
    let half_window = window / 2;
    let mut sum = T::zero();
    let mut count = 0;

    let start = 0;
    let end = (half_window + 1).min(n);
    for &v in &values[start..end] {
        sum += v;
        count += 1;
    }
    result.push(sum / T::from_usize(count).unwrap());

    for i in 1..n {
        let new_start = i.saturating_sub(half_window);
        let new_end = (i + half_window + 1).min(n);
        let old_start = (i - 1).saturating_sub(half_window);
        let old_end = (i - 1 + half_window + 1).min(n);

        for &v in &values[old_start..new_start] {
            sum -= v;
            count -= 1;
        }

        for &v in &values[old_end..new_end] {
            sum += v;
            count += 1;
        }

        result.push(sum / T::from_usize(count).unwrap());
    }

    result
}

/// Calculate Recovery Curve for a single motif
///
/// ### Params
///
/// * `gene_ranks` - Slice of gene ranks for this motif
/// * `max_rank` - Maximum rank to consider
///
/// ### Returns
///
/// Vector where position i equals the number of genes recovered by rank i
fn calculate_rcc(gene_ranks: &[i32], max_rank: i32) -> Vec<u32> {
    let mut rcc = vec![0u32; max_rank as usize];
    let mut filtered: Vec<i32> = gene_ranks
        .iter()
        .copied()
        .filter(|&r| r > 0 && r < max_rank)
        .collect();

    if filtered.is_empty() {
        return rcc;
    }

    filtered.sort_unstable();
    filtered.push(max_rank);

    let mut prev_rank = 0;
    for (gene_count, &rank) in filtered.iter().enumerate() {
        #[allow(clippy::needless_range_loop)]
        for pos in prev_rank..(rank as usize).min(max_rank as usize) {
            rcc[pos] = gene_count as u32;
        }
        prev_rank = rank as usize;
    }

    rcc
}

/// Calculate mean+2*sd thresholds using approximate method
///
/// ### Params
///
/// * `matrix` - Matrix reference containing gene-motif ranks
/// * `gene_indices` - Indices of genes in the gene set
/// * `max_rank` - Maximum rank to consider
/// * `n_mean` - Window size for smoothing
///
/// ### Returns
///
/// Vector of smoothed threshold values at each rank position
fn calculate_rcc_thresholds_aprox<T>(
    matrix: MatRef<'_, i32>,
    gene_indices: &[usize],
    max_rank: i32,
    n_mean: usize,
) -> Vec<T>
where
    T: BixverseFloat,
{
    let n_motifs = matrix.ncols();
    let n_genes = gene_indices.len();
    let max_rank_extra = (max_rank as usize) + n_mean;

    let mut global_mat = vec![vec![0usize; max_rank_extra]; n_genes + 1];

    for m_idx in 0..n_motifs {
        let mut ranks: Vec<i32> = gene_indices
            .iter()
            .map(|&idx| *matrix.get(idx, m_idx))
            .filter(|&r| r > 0 && r <= max_rank_extra as i32)
            .collect();

        if ranks.is_empty() {
            continue;
        }

        ranks.sort_unstable();

        for (gene_count, &rank) in ranks.iter().enumerate() {
            if (rank as usize) < max_rank_extra {
                global_mat[gene_count + 1][rank as usize] += 1;
            }
        }
    }

    let mut rcc_mean_raw = vec![T::zero(); max_rank_extra];
    let mut rcc_sd_raw = vec![T::zero(); max_rank_extra];

    for rank_pos in 0..max_rank_extra {
        let total_count: usize = global_mat.iter().map(|row| row[rank_pos]).sum();

        if total_count == 0 {
            if rank_pos > 0 {
                rcc_mean_raw[rank_pos] = rcc_mean_raw[rank_pos - 1];
                rcc_sd_raw[rank_pos] = rcc_sd_raw[rank_pos - 1];
            }
            continue;
        }

        let total_count_t = T::from_usize(total_count).unwrap();

        let mut sum = T::zero();
        let mut sum_sq = T::zero();

        for (num_genes, count) in global_mat.iter().enumerate() {
            let c = T::from_usize(count[rank_pos]).unwrap();
            let ng = T::from_usize(num_genes).unwrap();
            sum += ng * c;
            sum_sq += ng * ng * c;
        }

        let mean = sum / total_count_t;
        let variance = (sum_sq / total_count_t) - (mean * mean);

        rcc_mean_raw[rank_pos] = mean;
        rcc_sd_raw[rank_pos] = variance.sqrt();
    }

    let mut rcc_mean = vec![T::zero(); max_rank as usize];
    let mut rcc_sd = vec![T::zero(); max_rank as usize];

    let smoothed_mean = rolling_mean(&rcc_mean_raw, n_mean);
    let smoothed_sd = rolling_mean(&rcc_sd_raw, n_mean);

    for i in 0..(max_rank as usize) {
        if i < 5 {
            rcc_mean[i] = rcc_mean_raw[i];
            rcc_sd[i] = rcc_sd_raw[i];
        } else if i < smoothed_mean.len() {
            rcc_mean[i] = smoothed_mean[i];
            rcc_sd[i] = smoothed_sd[i];
        }
    }

    rcc_mean
        .iter()
        .zip(rcc_sd.iter())
        .map(|(&m, &s)| m + T::from_f64(2.0).unwrap() * s)
        .collect()
}

/// Find the leading edge for a motif
///
/// ### Params
///
/// * `matrix` - Matrix reference containing gene-motif ranks
/// * `gene_indices` - Indices of genes in the gene set
/// * `motif_idx` - Index of the motif to analyse
/// * `rcc_m2sd` - Pre-calculated threshold values
/// * `all_rccs` - Optional pre-calculated RCC curves (for exact method)
///
/// ### Returns
///
/// Tuple of (rank at maximum enrichment, number of enriched genes, enriched gene indices)
fn find_leading_edge<T>(
    matrix: MatRef<'_, i32>,
    gene_indices: &[usize],
    motif_idx: usize,
    rcc_m2sd: &[T],
    all_rccs: Option<&Vec<Vec<u32>>>,
) -> (u32, usize, Vec<usize>)
where
    T: BixverseFloat,
{
    let current_rcc = match all_rccs {
        Some(rccs) => &rccs[motif_idx],
        None => &calculate_rcc(
            &gene_indices
                .iter()
                .map(|&idx| *matrix.get(idx, motif_idx))
                .collect::<Vec<_>>(),
            rcc_m2sd.len() as i32,
        ),
    };

    let mut max_enrichment = T::neg_infinity();
    let mut best_rank = 0;

    for (rank_pos, &threshold) in rcc_m2sd.iter().enumerate() {
        if rank_pos < current_rcc.len() {
            let enrichment = T::from_u32(current_rcc[rank_pos]).unwrap() - threshold;
            if enrichment > max_enrichment {
                max_enrichment = enrichment;
                best_rank = rank_pos;
            }
        }
    }

    let rank_at_max = best_rank as u32;

    let enriched: Vec<usize> = gene_indices
        .iter()
        .filter(|&&idx| {
            let rank = *matrix.get(idx, motif_idx);
            rank > 0 && rank <= rank_at_max as i32
        })
        .copied()
        .collect();

    let n_enriched = enriched.len();

    (rank_at_max, n_enriched, enriched)
}

/// Calculate RCC for all motifs
///
/// ### Params
///
/// * `matrix` - Matrix reference containing gene-motif ranks
/// * `gene_indices` - Indices of genes in the gene set
/// * `max_rank` - Maximum rank to consider
///
/// ### Returns
///
/// Vector of RCC curves for each motif
fn calculate_all_rccs(
    matrix: MatRef<'_, i32>,
    gene_indices: &[usize],
    max_rank: i32,
) -> Vec<Vec<u32>> {
    let n_motifs = matrix.ncols();

    (0..n_motifs)
        .into_par_iter()
        .map(|motif_idx| {
            let gene_ranks: Vec<i32> = gene_indices
                .iter()
                .map(|&idx| *matrix.get(idx, motif_idx))
                .collect();

            calculate_rcc(&gene_ranks, max_rank)
        })
        .collect()
}

/// Calculate mean+2*sd thresholds from pre-calculated RCCs (exact method)
///
/// ### Params
///
/// * `all_rccs` - Pre-calculated RCC curves for all motifs
/// * `max_rank` - Maximum rank to consider
///
/// ### Returns
///
/// Vector of threshold values at each rank position
fn calculate_rcc_thresholds<T>(all_rccs: &[Vec<u32>], max_rank: i32) -> Vec<T>
where
    T: BixverseFloat,
{
    let max_rank_usize = max_rank as usize;
    let n_motifs = T::from_usize(all_rccs.len()).unwrap();

    // one pass per curve, so the reads are contiguous
    let mut sums = vec![T::zero(); max_rank_usize];
    let mut sums_sq = vec![T::zero(); max_rank_usize];
    for rcc in all_rccs {
        for ((sum, sum_sq), &count) in sums.iter_mut().zip(sums_sq.iter_mut()).zip(rcc) {
            let val = T::from_u32(count).unwrap();
            *sum += val;
            *sum_sq += val * val;
        }
    }

    sums.into_iter()
        .zip(sums_sq)
        .map(|(sum, sum_sq)| {
            let mean = sum / n_motifs;
            let variance = (sum_sq / n_motifs) - (mean * mean);
            mean + T::from_f64(2.0).unwrap() * variance.sqrt()
        })
        .collect()
}

/// Process a single gene set - calculate AUC, NES, and find significant motifs
///
/// ### Params
///
/// * `matrix` - Matrix reference containing gene-motif ranks
/// * `gene_indices` - Indices of genes in the gene set
/// * `auc_threshold` - Maximum rank to consider for AUC calculation
/// * `nes_threshold` - Minimum NES value for significance
/// * `max_rank` - Maximum rank to consider for RCC
/// * `method` - Method to use ("icistarget" for exact, other for approximate)
/// * `n_mean` - Window size for smoothing (used in approximate method)
///
/// ### Returns
///
/// Vector of significant motif enrichments
pub fn process_gene_set<T>(
    matrix: MatRef<'_, i32>,
    gene_indices: &[usize],
    auc_threshold: i32,
    nes_threshold: T,
    max_rank: i32,
    method: &RccType,
    n_mean: usize,
) -> Vec<MotifEnrichment<T>>
where
    T: BixverseFloat + std::iter::Sum,
{
    let auc_scores = calculate_aucs(matrix, gene_indices, auc_threshold);

    let nes_scores = calculate_nes(&auc_scores);

    let (rcc_m2sd, all_rccs) = match method {
        RccType::ICisTarget => {
            let all_rccs = calculate_all_rccs(matrix, gene_indices, max_rank);
            let rcc_m2sd: Vec<T> = calculate_rcc_thresholds(&all_rccs, max_rank);
            (rcc_m2sd, Some(all_rccs))
        }
        RccType::Approx => {
            let rcc_m2sd = calculate_rcc_thresholds_aprox(matrix, gene_indices, max_rank, n_mean);
            (rcc_m2sd, None)
        }
    };

    nes_scores
        .iter()
        .enumerate()
        .filter(|&(_, &nes)| nes >= nes_threshold)
        .map(|(motif_idx, &nes)| {
            let (rank_at_max, n_enriched, enriched_gene_indices) = find_leading_edge(
                matrix,
                gene_indices,
                motif_idx,
                &rcc_m2sd,
                all_rccs.as_ref(),
            );

            MotifEnrichment {
                motif_idx,
                nes,
                auc: auc_scores[motif_idx],
                rank_at_max,
                n_enriched,
                enriched_gene_indices,
            }
        })
        .collect()
}

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

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

    /// Regression: the priming loop primed `half_window` elements instead of
    /// `half_window + 1`, so `values[2]` here was dropped and never recovered.
    /// Every mean except the middle one was wrong.
    #[test]
    fn rolling_mean_is_centred_and_clipped_at_the_edges() {
        let got: Vec<f64> = rolling_mean(&[1.0, 2.0, 3.0, 4.0, 5.0], 5);

        assert_eq!(got.len(), 5);
        for (g, w) in got.iter().zip([2.0, 2.5, 3.0, 3.5, 4.0]) {
            assert_relative_eq!(*g, w, epsilon = 1e-12);
        }
    }

    /// A window of one is the identity, and must not divide by a zero count.
    #[test]
    fn rolling_mean_window_of_one_is_the_identity() {
        let values = [1.0_f64, 7.0, 3.0];
        let got: Vec<f64> = rolling_mean(&values, 1);

        for (g, w) in got.iter().zip(values) {
            assert_relative_eq!(*g, w, epsilon = 1e-12);
        }
    }

    /// AUC normalises by the full gene count, not the below-threshold subset:
    /// ranks [1, 3] contribute 2*1 + 2*2 = 6 against a maximum of 5 * 3.
    #[test]
    fn auc_single_normalises_by_the_unfiltered_gene_count() {
        assert_relative_eq!(
            calculate_auc_single::<f64>(&[1, 3, 7], 5),
            0.4,
            epsilon = 1e-12
        );
    }

    /// Every rank at or above the threshold leaves nothing to integrate.
    #[test]
    fn auc_single_is_zero_when_no_gene_beats_the_threshold() {
        assert_relative_eq!(
            calculate_auc_single::<f64>(&[7, 9], 5),
            0.0,
            epsilon = 1e-12
        );
    }

    /// The recovery curve steps at each retained rank and holds until the next.
    #[test]
    fn rcc_steps_at_each_retained_rank() {
        assert_eq!(calculate_rcc(&[2, 4, 9], 6), vec![0, 0, 1, 1, 2, 2]);
    }

    /// Rank 0 is filtered out by the `r > 0` guard, so it counts as absent
    /// rather than as the best possible rank.
    #[test]
    fn rcc_treats_rank_zero_as_absent() {
        assert_eq!(calculate_rcc(&[0, 3], 5), vec![0, 0, 0, 1, 1]);
    }

    /// NES z-scores against the sample standard deviation.
    #[test]
    fn nes_z_scores_against_the_sample_sd() {
        let got: Vec<f64> = calculate_nes(&[1.0, 2.0, 3.0]);

        for (g, w) in got.iter().zip([-1.0, 0.0, 1.0]) {
            assert_relative_eq!(*g, w, epsilon = 1e-12);
        }
    }

    /// A constant input has no spread to normalise by, so every score is zero
    /// rather than NaN.
    #[test]
    fn nes_of_a_constant_input_is_all_zero() {
        assert_eq!(calculate_nes::<f64>(&[2.0; 4]), vec![0.0; 4]);
    }
}