rsomics_sc_normalize/
lib.rs

1use std::fs::File;
2use std::io::{BufRead, BufReader, BufWriter, Read, Write};
3use std::path::{Path, PathBuf};
4
5use flate2::read::MultiGzDecoder;
6use rayon::prelude::*;
7use rsomics_common::{Result, RsomicsError};
8
9/// A single-cell count matrix in 10x MatrixMarket layout: rows are genes,
10/// columns are cells, stored as coordinate triplets. Counts are held as f64
11/// because scanpy normalizes after promoting the integer matrix to float.
12pub struct CountMatrix {
13    pub n_genes: usize,
14    pub n_cells: usize,
15    /// One entry per stored nonzero, in the file's row-major (gene-major) order.
16    pub entries: Vec<Entry>,
17}
18
19#[derive(Clone, Copy)]
20pub struct Entry {
21    pub gene: u32,
22    pub cell: u32,
23    pub value: f64,
24}
25
26pub struct NormalizeParams {
27    /// `None` reproduces scanpy's default: the median of every cell's total
28    /// count, computed over all cells including those with zero counts.
29    pub target_sum: Option<f64>,
30    pub log1p: bool,
31}
32
33/// Resolve the 10x triple inside `dir`, accepting both gzipped (v3) and plain
34/// (v2) layouts. Only the matrix is required for normalization.
35pub fn open_mtx(dir: &Path) -> Result<Box<dyn Read>> {
36    for name in ["matrix.mtx.gz", "matrix.mtx"] {
37        let path = dir.join(name);
38        if path.exists() {
39            return open_maybe_gz(&path);
40        }
41    }
42    Err(RsomicsError::InvalidInput(format!(
43        "no matrix.mtx or matrix.mtx.gz in {}",
44        dir.display()
45    )))
46}
47
48fn open_maybe_gz(path: &Path) -> Result<Box<dyn Read>> {
49    let file = File::open(path)
50        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
51    if path.extension().is_some_and(|e| e == "gz") {
52        Ok(Box::new(MultiGzDecoder::new(file)))
53    } else {
54        Ok(Box::new(file))
55    }
56}
57
58/// Parse a MatrixMarket coordinate file (real, integer, or pattern; general).
59/// 10x stores genes on rows, cells on columns.
60pub fn parse_mtx(reader: impl Read) -> Result<CountMatrix> {
61    let mut reader = BufReader::new(reader);
62    let mut line = String::new();
63
64    reader.read_line(&mut line).map_err(RsomicsError::Io)?;
65    let banner = line.trim();
66    if !banner.starts_with("%%MatrixMarket") {
67        return Err(RsomicsError::InvalidInput(
68            "missing %%MatrixMarket banner".into(),
69        ));
70    }
71    let pattern = banner.contains("pattern");
72
73    let (n_genes, n_cells, nnz) = loop {
74        line.clear();
75        let n = reader.read_line(&mut line).map_err(RsomicsError::Io)?;
76        if n == 0 {
77            return Err(RsomicsError::InvalidInput("truncated MTX header".into()));
78        }
79        let t = line.trim();
80        if t.is_empty() || t.starts_with('%') {
81            continue;
82        }
83        let mut it = t.split_whitespace();
84        let rows = parse_usize(it.next())?;
85        let cols = parse_usize(it.next())?;
86        let nnz = parse_usize(it.next())?;
87        break (rows, cols, nnz);
88    };
89
90    let mut entries = Vec::with_capacity(nnz);
91    for raw in reader.lines() {
92        let raw = raw.map_err(RsomicsError::Io)?;
93        let t = raw.trim();
94        if t.is_empty() {
95            continue;
96        }
97        let mut it = t.split_whitespace();
98        let gene = parse_usize(it.next())?;
99        let cell = parse_usize(it.next())?;
100        let value = if pattern {
101            1.0
102        } else {
103            it.next()
104                .ok_or_else(|| RsomicsError::InvalidInput("MTX entry missing value".into()))?
105                .parse::<f64>()?
106        };
107        if gene == 0 || gene > n_genes || cell == 0 || cell > n_cells {
108            return Err(RsomicsError::InvalidInput(format!(
109                "MTX index out of bounds: ({gene}, {cell})"
110            )));
111        }
112        entries.push(Entry {
113            gene: (gene - 1) as u32,
114            cell: (cell - 1) as u32,
115            value,
116        });
117    }
118    if entries.len() != nnz {
119        return Err(RsomicsError::InvalidInput(format!(
120            "MTX declared {nnz} entries, found {}",
121            entries.len()
122        )));
123    }
124
125    Ok(CountMatrix {
126        n_genes,
127        n_cells,
128        entries,
129    })
130}
131
132/// scanpy's `np.median` over the per-cell totals (linear interpolation: for an
133/// even count the two central order statistics are averaged). Zero-count cells
134/// are part of the population. A single NaN total poisons the whole median
135/// exactly as `np.median` propagates it, which cascades to an all-NaN matrix.
136fn median(totals: &[f64]) -> f64 {
137    let n = totals.len();
138    if n == 0 {
139        return 0.0;
140    }
141    if totals.iter().any(|t| t.is_nan()) {
142        return f64::NAN;
143    }
144    let mut sorted: Vec<f64> = totals.to_vec();
145    sorted.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
146    if n % 2 == 1 {
147        sorted[n / 2]
148    } else {
149        0.5 * (sorted[n / 2 - 1] + sorted[n / 2])
150    }
151}
152
153/// Per-cell total counts, indexed by cell.
154fn cell_totals(m: &CountMatrix) -> Vec<f64> {
155    let mut totals = vec![0.0_f64; m.n_cells];
156    for e in &m.entries {
157        totals[e.cell as usize] += e.value;
158    }
159    totals
160}
161
162/// Normalize each cell to `target_sum` (or the median of totals) then optionally
163/// apply `ln(1+x)`. Mirrors scanpy: a zero-count cell's scaling factor collapses
164/// to 1, leaving its (empty) row untouched, and the sparsity pattern is exact
165/// because `log1p(0) = 0`. A NaN count flows through untouched — its cell's
166/// scale is NaN (never mistaken for the zero case, since `NaN != 0`), so the
167/// affected values normalize to NaN just as scanpy propagates them.
168pub fn normalize(m: &mut CountMatrix, params: &NormalizeParams) {
169    let totals = cell_totals(m);
170    let target = params.target_sum.unwrap_or_else(|| median(&totals));
171
172    let scale: Vec<f64> = totals
173        .iter()
174        .map(|&t| {
175            let s = t / target;
176            if s == 0.0 { 1.0 } else { s }
177        })
178        .collect();
179
180    let log1p = params.log1p;
181    m.entries.par_iter_mut().for_each(|e| {
182        let v = e.value / scale[e.cell as usize];
183        e.value = if log1p { v.ln_1p() } else { v };
184    });
185}
186
187/// Write the matrix back in genes×cells MatrixMarket real coordinate layout,
188/// preserving the input entry order. A big buffer plus ryu float formatting
189/// keeps the matrix-sized write I/O-bound rather than format-bound.
190pub fn write_mtx(m: &CountMatrix, out: impl Write) -> Result<()> {
191    let mut w = BufWriter::with_capacity(1 << 20, out);
192    w.write_all(b"%%MatrixMarket matrix coordinate real general\n")
193        .map_err(RsomicsError::Io)?;
194    let mut header = itoa_line(m.n_genes, m.n_cells, m.entries.len());
195    header.push('\n');
196    w.write_all(header.as_bytes()).map_err(RsomicsError::Io)?;
197
198    let mut fmt = ryu::Buffer::new();
199    let mut buf: Vec<u8> = Vec::with_capacity(64);
200    for e in &m.entries {
201        buf.clear();
202        write_uint(&mut buf, e.gene as u64 + 1);
203        buf.push(b' ');
204        write_uint(&mut buf, e.cell as u64 + 1);
205        buf.push(b' ');
206        buf.extend_from_slice(fmt.format(e.value).as_bytes());
207        buf.push(b'\n');
208        w.write_all(&buf).map_err(RsomicsError::Io)?;
209    }
210    w.flush().map_err(RsomicsError::Io)?;
211    Ok(())
212}
213
214fn itoa_line(a: usize, b: usize, c: usize) -> String {
215    format!("{a} {b} {c}")
216}
217
218fn write_uint(buf: &mut Vec<u8>, mut n: u64) {
219    if n == 0 {
220        buf.push(b'0');
221        return;
222    }
223    let start = buf.len();
224    while n > 0 {
225        buf.push(b'0' + (n % 10) as u8);
226        n /= 10;
227    }
228    buf[start..].reverse();
229}
230
231fn parse_usize(tok: Option<&str>) -> Result<usize> {
232    tok.ok_or_else(|| RsomicsError::InvalidInput("MTX header missing a dimension".into()))?
233        .parse::<usize>()
234        .map_err(Into::into)
235}
236
237/// End-to-end: read the 10x matrix from `dir`, normalize, write to `out`.
238///
239/// A NaN count survives normalization as NaN. Under `log1p` (the default)
240/// scanpy's pipeline then rejects it — `sc.pp.log1p` runs sklearn `check_array`,
241/// which raises on NaN — so we fail loud rather than emit an all-NaN matrix with
242/// a success exit. With `--no-log` the raw `normalize_total` result (all-NaN) is
243/// returned, mirroring scanpy, which does not validate at that step.
244pub fn run(dir: &Path, params: &NormalizeParams, out: impl Write) -> Result<(usize, usize)> {
245    let mut m = parse_mtx(open_mtx(dir)?)?;
246    let shape = (m.n_genes, m.n_cells);
247    normalize(&mut m, params);
248    if params.log1p && m.entries.iter().any(|e| e.value.is_nan()) {
249        return Err(RsomicsError::InvalidInput(
250            "input count matrix contains NaN; log-normalization is undefined".into(),
251        ));
252    }
253    write_mtx(&m, out)?;
254    Ok(shape)
255}
256
257/// `--target-sum` accepts a positive float or the literal `median`.
258pub fn parse_target_sum(s: &str) -> Result<Option<f64>> {
259    if s.eq_ignore_ascii_case("median") {
260        return Ok(None);
261    }
262    let v = s
263        .parse::<f64>()
264        .map_err(|_| RsomicsError::InvalidInput(format!("invalid --target-sum '{s}'")))?;
265    if v <= 0.0 || !v.is_finite() {
266        return Err(RsomicsError::InvalidInput(
267            "--target-sum must be a positive finite number or 'median'".into(),
268        ));
269    }
270    Ok(Some(v))
271}
272
273/// Output destination — stdout for `-`, otherwise a file.
274pub fn open_output(path: &str) -> Result<Box<dyn Write>> {
275    if path == "-" {
276        Ok(Box::new(std::io::stdout().lock()))
277    } else {
278        Ok(Box::new(
279            File::create(PathBuf::from(path)).map_err(RsomicsError::Io)?,
280        ))
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    fn tiny() -> CountMatrix {
289        let mut entries = Vec::new();
290        let push = |v: &mut Vec<Entry>, g: u32, c: u32, val: f64| {
291            v.push(Entry {
292                gene: g,
293                cell: c,
294                value: val,
295            })
296        };
297        push(&mut entries, 0, 0, 3.0);
298        push(&mut entries, 2, 0, 1.0);
299        push(&mut entries, 4, 0, 2.0);
300        push(&mut entries, 1, 1, 5.0);
301        push(&mut entries, 0, 2, 1.0);
302        push(&mut entries, 1, 2, 1.0);
303        push(&mut entries, 2, 2, 1.0);
304        push(&mut entries, 3, 2, 1.0);
305        CountMatrix {
306            n_genes: 5,
307            n_cells: 4,
308            entries,
309        }
310    }
311
312    #[test]
313    fn median_includes_zero_cells() {
314        assert_eq!(median(&[6.0, 5.0, 4.0, 0.0]), 4.5);
315    }
316
317    #[test]
318    fn median_propagates_nan() {
319        assert!(median(&[f64::NAN, 7.0, 4.0]).is_nan());
320        assert!(median(&[1.0, 2.0, f64::NAN, 3.0]).is_nan());
321    }
322
323    #[test]
324    fn normalize_nan_cell_poisons_matrix() {
325        let mut m = CountMatrix {
326            n_genes: 2,
327            n_cells: 2,
328            entries: vec![
329                Entry {
330                    gene: 0,
331                    cell: 0,
332                    value: 3.0,
333                },
334                Entry {
335                    gene: 1,
336                    cell: 0,
337                    value: f64::NAN,
338                },
339                Entry {
340                    gene: 0,
341                    cell: 1,
342                    value: 5.0,
343                },
344            ],
345        };
346        normalize(
347            &mut m,
348            &NormalizeParams {
349                target_sum: None,
350                log1p: false,
351            },
352        );
353        assert!(m.entries.iter().all(|e| e.value.is_nan()));
354    }
355
356    #[test]
357    fn matches_scanpy_tiny() {
358        let mut m = tiny();
359        normalize(
360            &mut m,
361            &NormalizeParams {
362                target_sum: None,
363                log1p: true,
364            },
365        );
366        let want = [
367            (0u32, 0u32, 1.178655_f64),
368            (2, 0, 0.5596158),
369            (4, 0, 0.91629076),
370            (1, 1, 1.7047482),
371            (0, 2, 0.7537718),
372            (1, 2, 0.7537718),
373            (2, 2, 0.7537718),
374            (3, 2, 0.7537718),
375        ];
376        for (e, (_, _, exp)) in m.entries.iter().zip(want.iter()) {
377            assert!((e.value - exp).abs() < 1e-5, "{} vs {}", e.value, exp);
378        }
379    }
380
381    #[test]
382    fn target_sum_parsing() {
383        assert_eq!(parse_target_sum("median").unwrap(), None);
384        assert_eq!(parse_target_sum("1e4").unwrap(), Some(10000.0));
385        assert!(parse_target_sum("-1").is_err());
386        assert!(parse_target_sum("abc").is_err());
387    }
388
389    #[test]
390    fn roundtrip_mtx() {
391        let m = tiny();
392        let mut buf = Vec::new();
393        write_mtx(&m, &mut buf).unwrap();
394        let parsed = parse_mtx(&buf[..]).unwrap();
395        assert_eq!(parsed.n_genes, 5);
396        assert_eq!(parsed.n_cells, 4);
397        assert_eq!(parsed.entries.len(), 8);
398    }
399}
rsomics_sc_normalize/lib.rs

rsomics_sc_normalize/
lib.rs
