bixverse_sc_data/
plotting.rs

1//! Helpers to extract genes for plotting. Can densify raw and normalised
2//! counts, do Z-score normalisation for groups and caclulate proportion
3//! of cells.
4
5use indexmap::IndexSet;
6use rayon::prelude::*;
7
8use crate::prelude::*;
9use bixverse_core::utils::simd::*;
10
11/////////////
12// Helpers //
13/////////////
14
15/// Result for grouped gene statistics (dot plot / heatmap data).
16pub struct GroupedGeneStats {
17    /// Gene indices (original), length = n_genes
18    pub gene_indices: Vec<usize>,
19    /// Group labels in order, length = n_groups
20    pub group_labels: Vec<String>,
21    /// Mean normalised expression per gene per group,
22    /// row-major (genes x groups)
23    pub mean_expression: Vec<f32>,
24    /// Fraction of cells with count > 0 per gene per group, same layout
25    pub pct_expressed: Vec<f32>,
26}
27
28/// Z-score in place using SIMD helpers, optionally clamp.
29///
30/// ### Params
31///
32/// * `data` - Mutable reference to the data
33/// * `clip` - Optional clipping values
34fn scale_and_clip(data: &mut [f32], clip: Option<f32>) {
35    let n = data.len() as f32;
36    let mean = sum_simd_f32(data) / n;
37    let var = sum_squared_dev_simd_f32(data, mean) / n;
38    let sd = var.sqrt();
39    if sd > 1e-8 {
40        for x in data.iter_mut() {
41            *x = (*x - mean) / sd;
42        }
43    }
44    if let Some(c) = clip {
45        for x in data.iter_mut() {
46            *x = x.clamp(-c, c);
47        }
48    }
49}
50
51////////////////////
52// Main functions //
53////////////////////
54
55/// Extract raw counts for a single gene as a dense vector over the given cells.
56///
57/// ### Params
58///
59/// * `reader` - Reader over the gene-based count store.
60/// * `cell_indices` - Slice of indices for the cells.
61/// * `gene_index` - The gene index for which to get the values.
62///
63/// ### Returns
64///
65/// Returns a vector of length `cell_indices.len()` in the same order as
66/// `cell_indices`.
67pub fn extract_raw_counts<S: SingleCellReading>(
68    reader: &S,
69    cell_indices: &[usize],
70    gene_index: usize,
71) -> Result<Vec<u32>, ScDataErrors> {
72    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
73
74    let mut chunk = reader
75        .read_gene_parallel(&[gene_index])?
76        .into_iter()
77        .next()
78        .unwrap();
79    chunk.filter_selected_cells(&cell_set);
80
81    let n = cell_indices.len();
82    let mut dense = vec![0u32; n];
83    for (i, &cell_idx) in chunk.indices.iter().enumerate() {
84        dense[cell_idx as usize] = chunk.data_raw.get(i);
85    }
86    Ok(dense)
87}
88
89/// Extract normalised (log1p) counts for a single gene, dense, over given cells.
90///
91/// If `scale` is true, z-scores across all selected cells (including zeros).
92/// If `clip` is provided, scaled values are clamped to `[-clip, clip]`.
93///
94/// ### Params
95///
96/// * `reader` - Reader over the gene-based count store.
97/// * `cell_indices` - Slice of indices for the cells.
98/// * `gene_index` - The gene index for which to get the values.
99/// * `scale` - Shall the data be scaled.
100/// * `clip` - Shall scaled data be clipped between specific values.
101///
102/// ### Returns
103///
104/// Returns a vector of length `cell_indices.len()` in the same order as
105/// `cell_indices`.
106pub fn extract_norm_counts<S: SingleCellReading>(
107    reader: &S,
108    cell_indices: &[usize],
109    gene_index: usize,
110    scale: bool,
111    clip: Option<f32>,
112) -> Result<Vec<f32>, ScDataErrors> {
113    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
114
115    let mut chunk = reader
116        .read_gene_parallel(&[gene_index])?
117        .into_iter()
118        .next()
119        .unwrap();
120    chunk.filter_selected_cells(&cell_set);
121
122    let n = cell_indices.len();
123    let mut dense = vec![0.0f32; n];
124    for (i, &cell_idx) in chunk.indices.iter().enumerate() {
125        dense[cell_idx as usize] = chunk.data_norm[i].to_f32();
126    }
127
128    if scale {
129        scale_and_clip(&mut dense, clip);
130    }
131
132    Ok(dense)
133}
134
135/// Extract normalised (log1p) counts for a multiple genes, dense, over given
136/// cells.
137///
138/// If `scale` is true, z-scores across all selected cells (including zeros).
139/// If `clip` is provided, scaled values are clamped to `[-clip, clip]`.
140///
141/// ### Params
142///
143/// * `reader` - Reader over the gene-based count store.
144/// * `cell_indices` - Slice of indices for the cells.
145/// * `gene_index` - The gene index for which to get the values.
146/// * `scale` - Shall the data be scaled.
147/// * `clip` - Shall scaled data be clipped between specific values.
148///
149/// ### Returns
150///
151/// Returns a vector of vectors of length `cell_indices.len()` in the same order
152/// as `cell_indices` with the densified gene expression values.
153pub fn extract_norm_counts_multi<S: SingleCellReading>(
154    reader: &S,
155    cell_indices: &[usize],
156    gene_indices: &[usize],
157    scale: bool,
158    clip: Option<f32>,
159) -> Result<Vec<Vec<f32>>, ScDataErrors> {
160    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
161    let n = cell_indices.len();
162
163    let mut chunks = reader.read_gene_parallel(gene_indices)?;
164
165    chunks.par_iter_mut().for_each(|chunk| {
166        chunk.filter_selected_cells(&cell_set);
167    });
168
169    let res = chunks
170        .par_iter()
171        .map(|chunk| {
172            let mut dense = vec![0.0f32; n];
173            for (i, &cell_idx) in chunk.indices.iter().enumerate() {
174                dense[cell_idx as usize] = chunk.data_norm[i].to_f32();
175            }
176            if scale {
177                scale_and_clip(&mut dense, clip);
178            }
179            dense
180        })
181        .collect();
182
183    Ok(res)
184}
185
186/// Compute per-group mean expression and percent expressed for a set of genes.
187///
188/// `group_levels` defines the factor levels (labels). `group_ids` is parallel
189/// to `cell_indices` and contains indices into `group_levels`, i.e.
190/// `group_levels[group_ids[i]]` is the label for `cell_indices[i]`.
191///
192/// This mirrors R's factor representation: an integer vector plus a levels
193/// attribute.
194///
195/// ### Params
196///
197/// * `reader` - Reader over the gene-based count store.
198/// * `cell_indices` - Slice of indices for the cells.
199/// * `gene_indices` - Slice of indices for the genes.
200/// * `group_ids` - Integer group assignments parallel to `cell_indices`,
201///   indexing into `group_levels`.
202/// * `group_levels` - Factor level labels. `group_ids` values must be valid
203///   indices into this slice.
204///
205/// ### Returns
206///
207/// A `GroupedGeneStats` containing per-gene, per-group mean normalised
208/// expression and fraction of expressing cells. Output vectors are laid out
209/// row-major (genes x n_levels).
210pub fn extract_grouped_gene_stats<S: SingleCellReading>(
211    reader: &S,
212    cell_indices: &[usize],
213    gene_indices: &[usize],
214    group_ids: &[usize],
215    group_levels: &[String],
216) -> Result<GroupedGeneStats, ScDataErrors> {
217    assert_eq!(
218        cell_indices.len(),
219        group_ids.len(),
220        "cell_indices and group_ids must have equal length"
221    );
222
223    let cell_set: IndexSet<u32> = cell_indices.iter().map(|&x| x as u32).collect();
224    let n_groups = group_levels.len();
225
226    let mut group_sizes = vec![0usize; n_groups];
227    for &g in group_ids {
228        debug_assert!(
229            g < n_groups,
230            "group_id {} out of bounds for {} levels",
231            g,
232            n_groups
233        );
234        group_sizes[g] += 1;
235    }
236
237    let mut chunks = reader.read_gene_parallel(gene_indices)?;
238
239    chunks.par_iter_mut().for_each(|chunk| {
240        chunk.filter_selected_cells(&cell_set);
241    });
242
243    let n_genes = gene_indices.len();
244    let mut mean_expression = vec![0.0f32; n_genes * n_groups];
245    let mut pct_expressed = vec![0.0f32; n_genes * n_groups];
246
247    for (g, chunk) in chunks.iter().enumerate() {
248        let mut sums = vec![0.0f32; n_groups];
249        let mut counts = vec![0usize; n_groups];
250
251        for (i, &new_cell_idx) in chunk.indices.iter().enumerate() {
252            let grp = group_ids[new_cell_idx as usize];
253            sums[grp] += chunk.data_norm[i].to_f32();
254            counts[grp] += 1;
255        }
256
257        let offset = g * n_groups;
258        for grp in 0..n_groups {
259            if group_sizes[grp] > 0 {
260                let size = group_sizes[grp] as f32;
261                mean_expression[offset + grp] = sums[grp] / size;
262                pct_expressed[offset + grp] = counts[grp] as f32 / size;
263            }
264        }
265    }
266
267    Ok(GroupedGeneStats {
268        gene_indices: chunks.iter().map(|c| c.original_index).collect(),
269        group_labels: group_levels.to_vec(),
270        mean_expression,
271        pct_expressed,
272    })
273}
bixverse_sc_data/plotting.rs

bixverse_sc_data/
plotting.rs
