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.
135fn median(totals: &[f64]) -> f64 {
136    let mut sorted: Vec<f64> = totals.to_vec();
137    sorted.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
138    let n = sorted.len();
139    if n == 0 {
140        return 0.0;
141    }
142    if n % 2 == 1 {
143        sorted[n / 2]
144    } else {
145        0.5 * (sorted[n / 2 - 1] + sorted[n / 2])
146    }
147}
148
149/// Per-cell total counts, indexed by cell.
150fn cell_totals(m: &CountMatrix) -> Vec<f64> {
151    let mut totals = vec![0.0_f64; m.n_cells];
152    for e in &m.entries {
153        totals[e.cell as usize] += e.value;
154    }
155    totals
156}
157
158/// Normalize each cell to `target_sum` (or the median of totals) then optionally
159/// apply `ln(1+x)`. Mirrors scanpy: a zero-count cell's scaling factor collapses
160/// to 1, leaving its (empty) row untouched, and the sparsity pattern is exact
161/// because `log1p(0) = 0`.
162pub fn normalize(m: &mut CountMatrix, params: &NormalizeParams) {
163    let totals = cell_totals(m);
164    let target = params.target_sum.unwrap_or_else(|| median(&totals));
165
166    let scale: Vec<f64> = totals
167        .iter()
168        .map(|&t| {
169            let s = t / target;
170            if s == 0.0 { 1.0 } else { s }
171        })
172        .collect();
173
174    let log1p = params.log1p;
175    m.entries.par_iter_mut().for_each(|e| {
176        let v = e.value / scale[e.cell as usize];
177        e.value = if log1p { v.ln_1p() } else { v };
178    });
179}
180
181/// Write the matrix back in genes×cells MatrixMarket real coordinate layout,
182/// preserving the input entry order. A big buffer plus ryu float formatting
183/// keeps the matrix-sized write I/O-bound rather than format-bound.
184pub fn write_mtx(m: &CountMatrix, out: impl Write) -> Result<()> {
185    let mut w = BufWriter::with_capacity(1 << 20, out);
186    w.write_all(b"%%MatrixMarket matrix coordinate real general\n")
187        .map_err(RsomicsError::Io)?;
188    let mut header = itoa_line(m.n_genes, m.n_cells, m.entries.len());
189    header.push('\n');
190    w.write_all(header.as_bytes()).map_err(RsomicsError::Io)?;
191
192    let mut fmt = ryu::Buffer::new();
193    let mut buf: Vec<u8> = Vec::with_capacity(64);
194    for e in &m.entries {
195        buf.clear();
196        write_uint(&mut buf, e.gene as u64 + 1);
197        buf.push(b' ');
198        write_uint(&mut buf, e.cell as u64 + 1);
199        buf.push(b' ');
200        buf.extend_from_slice(fmt.format(e.value).as_bytes());
201        buf.push(b'\n');
202        w.write_all(&buf).map_err(RsomicsError::Io)?;
203    }
204    w.flush().map_err(RsomicsError::Io)?;
205    Ok(())
206}
207
208fn itoa_line(a: usize, b: usize, c: usize) -> String {
209    format!("{a} {b} {c}")
210}
211
212fn write_uint(buf: &mut Vec<u8>, mut n: u64) {
213    if n == 0 {
214        buf.push(b'0');
215        return;
216    }
217    let start = buf.len();
218    while n > 0 {
219        buf.push(b'0' + (n % 10) as u8);
220        n /= 10;
221    }
222    buf[start..].reverse();
223}
224
225fn parse_usize(tok: Option<&str>) -> Result<usize> {
226    tok.ok_or_else(|| RsomicsError::InvalidInput("MTX header missing a dimension".into()))?
227        .parse::<usize>()
228        .map_err(Into::into)
229}
230
231/// End-to-end: read the 10x matrix from `dir`, normalize, write to `out`.
232pub fn run(dir: &Path, params: &NormalizeParams, out: impl Write) -> Result<(usize, usize)> {
233    let mut m = parse_mtx(open_mtx(dir)?)?;
234    let shape = (m.n_genes, m.n_cells);
235    normalize(&mut m, params);
236    write_mtx(&m, out)?;
237    Ok(shape)
238}
239
240/// `--target-sum` accepts a positive float or the literal `median`.
241pub fn parse_target_sum(s: &str) -> Result<Option<f64>> {
242    if s.eq_ignore_ascii_case("median") {
243        return Ok(None);
244    }
245    let v = s
246        .parse::<f64>()
247        .map_err(|_| RsomicsError::InvalidInput(format!("invalid --target-sum '{s}'")))?;
248    if v <= 0.0 || !v.is_finite() {
249        return Err(RsomicsError::InvalidInput(
250            "--target-sum must be a positive finite number or 'median'".into(),
251        ));
252    }
253    Ok(Some(v))
254}
255
256/// Output destination — stdout for `-`, otherwise a file.
257pub fn open_output(path: &str) -> Result<Box<dyn Write>> {
258    if path == "-" {
259        Ok(Box::new(std::io::stdout().lock()))
260    } else {
261        Ok(Box::new(
262            File::create(PathBuf::from(path)).map_err(RsomicsError::Io)?,
263        ))
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    fn tiny() -> CountMatrix {
272        let mut entries = Vec::new();
273        let push = |v: &mut Vec<Entry>, g: u32, c: u32, val: f64| {
274            v.push(Entry {
275                gene: g,
276                cell: c,
277                value: val,
278            })
279        };
280        push(&mut entries, 0, 0, 3.0);
281        push(&mut entries, 2, 0, 1.0);
282        push(&mut entries, 4, 0, 2.0);
283        push(&mut entries, 1, 1, 5.0);
284        push(&mut entries, 0, 2, 1.0);
285        push(&mut entries, 1, 2, 1.0);
286        push(&mut entries, 2, 2, 1.0);
287        push(&mut entries, 3, 2, 1.0);
288        CountMatrix {
289            n_genes: 5,
290            n_cells: 4,
291            entries,
292        }
293    }
294
295    #[test]
296    fn median_includes_zero_cells() {
297        assert_eq!(median(&[6.0, 5.0, 4.0, 0.0]), 4.5);
298    }
299
300    #[test]
301    fn matches_scanpy_tiny() {
302        let mut m = tiny();
303        normalize(
304            &mut m,
305            &NormalizeParams {
306                target_sum: None,
307                log1p: true,
308            },
309        );
310        let want = [
311            (0u32, 0u32, 1.178655_f64),
312            (2, 0, 0.5596158),
313            (4, 0, 0.91629076),
314            (1, 1, 1.7047482),
315            (0, 2, 0.7537718),
316            (1, 2, 0.7537718),
317            (2, 2, 0.7537718),
318            (3, 2, 0.7537718),
319        ];
320        for (e, (_, _, exp)) in m.entries.iter().zip(want.iter()) {
321            assert!((e.value - exp).abs() < 1e-5, "{} vs {}", e.value, exp);
322        }
323    }
324
325    #[test]
326    fn target_sum_parsing() {
327        assert_eq!(parse_target_sum("median").unwrap(), None);
328        assert_eq!(parse_target_sum("1e4").unwrap(), Some(10000.0));
329        assert!(parse_target_sum("-1").is_err());
330        assert!(parse_target_sum("abc").is_err());
331    }
332
333    #[test]
334    fn roundtrip_mtx() {
335        let m = tiny();
336        let mut buf = Vec::new();
337        write_mtx(&m, &mut buf).unwrap();
338        let parsed = parse_mtx(&buf[..]).unwrap();
339        assert_eq!(parsed.n_genes, 5);
340        assert_eq!(parsed.n_cells, 4);
341        assert_eq!(parsed.entries.len(), 8);
342    }
343}
rsomics_sc_normalize/lib.rs

rsomics_sc_normalize/
lib.rs
