data_beans/sparse_io_vector/
scales.rs

1//! Read-time column scales, and the depth factors they come from.
2//!
3//! A column's depth is read off its total count, but a few rows detected in
4//! nearly every column can take a large share of the counts, and that share
5//! can differ from one group of columns to the next. Divided by a total that
6//! holds them, every other row of a group whose broad rows take more would
7//! read lower: a group effect the normalization itself makes. So
8//! [`SparseIoVec::depth_scales`] reads each column's count over the rows no
9//! [`RowFilter`] leaves out, over the median of that count across columns.
10//!
11//! The filters only choose which rows count toward the depth. With
12//! [`SparseIoVec::set_column_scales`], every read divides each column by its
13//! factor, the left-out rows included; without it, reads return the counts
14//! as stored.
15
16use super::*;
17
18/// Rows a column's depth leaves out. Several filters combine by union.
19#[derive(Clone, Debug)]
20pub enum RowFilter {
21    /// One flag per row, `true` to leave the row out.
22    Mask(Vec<bool>),
23    /// Rows whose names match.
24    Names(regex::Regex),
25    /// Rows detected in more than `share` of the columns of every group,
26    /// `groups` giving each column's group. A row common in one group alone
27    /// is kept, since that may be what tells the groups apart.
28    Broad { share: f64, groups: Vec<Box<str>> },
29}
30
31/// Each column's depth factor, and the rows its depth left out.
32#[derive(Clone, Debug)]
33pub struct ColumnScales {
34    /// Per column: its depth over the median of the positive depths. A
35    /// column with no count off the left-out rows gets 1, so its reads stay
36    /// as stored rather than blown up.
37    pub factor: Vec<f32>,
38    /// Per row: `true` when some filter left it out of the depth.
39    pub ignored_rows: Vec<bool>,
40}
41
42impl SparseIoVec {
43    /// The depth factors of the columns, read off the counts as stored, in
44    /// blocks of `block_size` columns; refused once scales are set, since
45    /// the reads would then be scaled. A [`RowFilter::Broad`] filter takes
46    /// one pass over the columns to find its rows before the pass that reads
47    /// the depths.
48    pub fn depth_scales(
49        &self,
50        ignore: &[RowFilter],
51        block_size: Option<usize>,
52    ) -> anyhow::Result<ColumnScales> {
53        anyhow::ensure!(
54            self.column_scales.is_none(),
55            "depth_scales on a scaled vector: read the depths before set_column_scales"
56        );
57        let (n_rows, n_cols) = (self.num_rows(), self.num_columns());
58        let blocks = generate_minibatch_intervals(n_cols, n_rows, block_size);
59
60        let mut ignored = vec![false; n_rows];
61        for filter in ignore {
62            match filter {
63                RowFilter::Mask(mask) => {
64                    anyhow::ensure!(
65                        mask.len() == n_rows,
66                        "row mask has {} flags for {n_rows} rows",
67                        mask.len()
68                    );
69                    ignored.iter_mut().zip(mask).for_each(|(i, &o)| *i |= o);
70                }
71                RowFilter::Names(re) => {
72                    let names = self.row_names()?;
73                    ignored
74                        .iter_mut()
75                        .zip(&names)
76                        .for_each(|(i, n)| *i |= re.is_match(n));
77                }
78                RowFilter::Broad { share, groups } => {
79                    let broad = self.broad_rows(*share, groups, &blocks)?;
80                    ignored.iter_mut().zip(broad).for_each(|(i, o)| *i |= o);
81                }
82            }
83        }
84
85        let kept: Vec<usize> = (0..n_rows).filter(|&r| !ignored[r]).collect();
86        let (_, depth, _, _) =
87            crate::qc::collect_column_stat_across_vec(self, Some(&kept), block_size)?.to_f32_vecs();
88
89        Ok(ColumnScales {
90            factor: factors(&depth),
91            ignored_rows: ignored,
92        })
93    }
94
95    /// Rows detected in more than `share` of the columns of every group.
96    fn broad_rows(
97        &self,
98        share: f64,
99        groups: &[Box<str>],
100        blocks: &[(usize, usize)],
101    ) -> anyhow::Result<Vec<bool>> {
102        let (n_rows, n_cols) = (self.num_rows(), self.num_columns());
103        anyhow::ensure!(
104            share.is_finite() && (0.0..1.0).contains(&share),
105            "broad share {share} is not in [0, 1)"
106        );
107        anyhow::ensure!(
108            groups.len() == n_cols,
109            "{} group labels for {n_cols} columns",
110            groups.len()
111        );
112        let mut index: HashMap<&str, usize> = HashMap::default();
113        let g: Vec<usize> = groups
114            .iter()
115            .map(|l| {
116                let next = index.len();
117                *index.entry(l).or_insert(next)
118            })
119            .collect();
120        let n_groups = index.len();
121        let mut size = vec![0usize; n_groups];
122        g.iter().for_each(|&k| size[k] += 1);
123
124        // `detected[k * n_rows + r]`: the columns of group `k` holding row `r`.
125        let zero = || vec![0u32; n_groups * n_rows];
126        let detected = blocks
127            .par_iter()
128            .try_fold(zero, |mut d, &(lb, ub)| -> anyhow::Result<Vec<u32>> {
129                let x = self.read_columns_csc(lb..ub)?;
130                for (k, col) in x.col_iter().enumerate() {
131                    let dk = &mut d[g[lb + k] * n_rows..][..n_rows];
132                    for (&r, &v) in col.row_indices().iter().zip(col.values()) {
133                        if v > 0.0 {
134                            dk[r] += 1;
135                        }
136                    }
137                }
138                Ok(d)
139            })
140            .try_reduce(zero, |mut a, b| {
141                a.iter_mut().zip(b).for_each(|(a, b)| *a += b);
142                Ok(a)
143            })?;
144
145        Ok((0..n_rows)
146            .map(|r| {
147                (0..n_groups).all(|k| detected[k * n_rows + r] as f64 > share * size[k] as f64)
148            })
149            .collect())
150    }
151
152    /// Divide every read of column `c` by `factor[c]` from now on; `None`
153    /// reads the counts as stored again. Each factor must be positive and
154    /// finite, one per column. Pushing another backend afterwards is
155    /// refused, since its columns would have none.
156    ///
157    /// Every read through this vector is scaled, collapse and QC included:
158    /// keep a clone without scales for whatever must see the stored counts.
159    pub fn set_column_scales(&mut self, factor: Option<Vec<f32>>) -> anyhow::Result<()> {
160        if let Some(f) = &factor {
161            anyhow::ensure!(
162                f.len() == self.num_columns(),
163                "{} column scales for {} columns",
164                f.len(),
165                self.num_columns()
166            );
167            if let Some((c, v)) = f
168                .iter()
169                .enumerate()
170                .find(|(_, v)| !(v.is_finite() && **v > 0.0))
171            {
172                anyhow::bail!("column {c} has scale {v}; scales must be positive and finite");
173            }
174        }
175        self.column_scales = factor;
176        Ok(())
177    }
178
179    /// The factors every read divides by, one per column, when set.
180    pub fn column_scales(&self) -> Option<&[f32]> {
181        self.column_scales.as_deref()
182    }
183}
184
185/// Each depth over the median of the positive ones; a column with none gets
186/// 1. All ones when no depth is positive.
187fn factors(depth: &[f32]) -> Vec<f32> {
188    let positive: Vec<f32> = depth.iter().copied().filter(|&d| d > 0.0).collect();
189    if positive.is_empty() {
190        return vec![1.0; depth.len()];
191    }
192    let median = median(&positive);
193    depth
194        .iter()
195        .map(|&d| if d > 0.0 { d / median } else { 1.0 })
196        .collect()
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::sparse_io::{create_sparse_from_triplets, SparseIoBackend};
203
204    /// The matched and neighbour reads go through `read_matched`.
205    #[test]
206    fn a_matched_column_read_divides_by_its_factor() {
207        let dir = tempfile::tempdir().unwrap();
208        let path = dir.path().join("a.zarr");
209        let t = [(0, 0, 4.0), (1, 1, 6.0), (0, 2, 9.0)];
210        let mut b = create_sparse_from_triplets(
211            &t,
212            (2, 3, t.len()),
213            Some(path.to_str().unwrap()),
214            Some(&SparseIoBackend::Zarr),
215        )
216        .unwrap();
217        b.register_row_names_vec(&["R0".into(), "R1".into()]);
218        b.register_column_names_vec(&["c0".into(), "c1".into(), "c2".into()]);
219        let mut data = SparseIoVec::new();
220        data.push(Arc::from(b), None).unwrap();
221        data.set_column_scales(Some(vec![2.0, 3.0, 9.0])).unwrap();
222        let matched = data
223            .read_matched(vec![0, 0], vec![2, 1], vec![0.0, 0.0])
224            .unwrap();
225        assert_eq!(matched.shape, (2, 2));
226        let mut out = matched.triplets;
227        out.sort_by_key(|&(r, c, _)| (c, r));
228        assert_eq!(out, [(0, 0, 1.0), (1, 1, 2.0)]);
229    }
230}
data_beans/sparse_io_vector/scales.rs

data_beans/sparse_io_vector/
scales.rs
