single_statistics/testing/
utils.rs

1use nalgebra_sparse::CsrMatrix;
2use single_utilities::traits::FloatOpsTS;
3
4pub fn extract_unique_groups(group_ids: &[usize]) -> Vec<usize> {
5    let mut unique_groups = group_ids.to_vec();
6    unique_groups.sort();
7    unique_groups.dedup();
8    unique_groups
9}
10
11/// Get indices for each group
12pub fn get_group_indices(group_ids: &[usize], unique_groups: &[usize]) -> (Vec<usize>, Vec<usize>) {
13    let group1 = unique_groups[0];
14    let group2 = unique_groups[1];
15
16    let group1_indices = group_ids.iter()
17        .enumerate()
18        .filter_map(|(i, &g)| if g == group1 { Some(i) } else { None })
19        .collect();
20
21    let group2_indices = group_ids.iter()
22        .enumerate()
23        .filter_map(|(i, &g)| if g == group2 { Some(i) } else { None })
24        .collect();
25
26    (group1_indices, group2_indices)
27}
28
29pub(crate) fn accumulate_gene_statistics_two_groups<T>(
30    matrix: &CsrMatrix<T>,
31    group1_indices: &[usize],
32    group2_indices: &[usize],
33) -> anyhow::Result<(Vec<T>, Vec<T>, Vec<T>, Vec<T>)>
34where
35    T: FloatOpsTS,
36{
37    let n_genes = matrix.ncols();
38
39    // Pre-allocate all accumulation vectors
40    let mut group1_sums = vec![T::zero(); n_genes];
41    let mut group1_sum_squares = vec![T::zero(); n_genes];
42    let mut group2_sums = vec![T::zero(); n_genes];
43    let mut group2_sum_squares = vec![T::zero(); n_genes];
44    
45    for &row_idx in group1_indices {
46        let row = matrix.row(row_idx);
47        for (col_idx, &value) in row.col_indices().iter().zip(row.values().iter()) {
48            group1_sums[*col_idx] += value;
49            group1_sum_squares[*col_idx] += value * value;
50        }
51    }
52    
53    for &row_idx in group2_indices {
54        let row = matrix.row(row_idx);
55        for (col_idx, &value) in row.col_indices().iter().zip(row.values().iter()) {
56            group2_sums[*col_idx] += value;
57            group2_sum_squares[*col_idx] += value * value;
58        }
59    }
60
61    Ok((group1_sums, group1_sum_squares, group2_sums, group2_sum_squares))
62}
63
64pub(crate) fn extract_gene_values_optimized<T>(
65    matrix: &CsrMatrix<T>,
66    gene_idx: usize,
67    group1_indices: &[usize],
68    group2_indices: &[usize],
69) -> (Vec<T>, Vec<T>)
70where
71    T: FloatOpsTS,
72{
73    // Pre-allocate with known capacity
74    let mut group1_values = Vec::with_capacity(group1_indices.len());
75    let mut group2_values = Vec::with_capacity(group2_indices.len());
76
77    // Extract values efficiently
78    for &row_idx in group1_indices {
79        let value = matrix.get_entry(row_idx, gene_idx)
80            .map(|entry| entry.into_value())
81            .unwrap_or(T::zero());
82        group1_values.push(value);
83    }
84
85    for &row_idx in group2_indices {
86        let value = matrix.get_entry(row_idx, gene_idx)
87            .map(|entry| entry.into_value())
88            .unwrap_or(T::zero());
89        group2_values.push(value);
90    }
91
92    (group1_values, group2_values)
93}