bixverse_sc_data/
mtx_multifile_io.rs

1//! Multi-mtx loading: scan multiple mtx files, apply global gene QC,
2//! and write all cells into a single binary file. Mirrors the h5ad
3//! multi-loader design.
4
5use rayon::prelude::*;
6use std::path::Path;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicUsize, Ordering};
9use std::time::Instant;
10use thousands::Separable;
11
12use crate::data_io::*;
13use crate::mtx_io::{MtxSource, parse_block};
14use crate::prelude::*;
15
16////////////////
17// File tasks //
18////////////////
19
20/// Per-file task descriptor for multi-mtx loading
21pub struct MtxFileTask {
22    /// Experimental identifier
23    pub exp_id: String,
24    /// Path to the .mtx file, plain or gzipped
25    pub mtx_path: String,
26    /// file-local gene idx -> universe gene idx; None if gene not in
27    /// the (intersected) universe
28    pub gene_local_to_universe: Vec<Option<usize>>,
29    /// Are cells the rows in this file
30    pub cells_as_rows: bool,
31}
32
33/// Per-file QC output returned to the caller
34pub struct MtxFileQcResult {
35    /// Experimental identifier
36    pub exp_id: String,
37    /// Cells to keep (file-local 0-indexed indices)
38    pub cells_to_keep: Vec<usize>,
39    /// Library sizes of the kept cells
40    pub lib_size: Vec<usize>,
41    /// Number of features per kept cell
42    pub nnz: Vec<usize>,
43}
44
45/// Final result from multi-mtx loading
46pub struct MultiMtxResult {
47    /// Universe indices of the genes that survived the global QC
48    pub global_gene_indices: Vec<usize>,
49    /// Total cells ingested
50    pub total_cells: usize,
51    /// Total genes/features ingested
52    pub total_genes: usize,
53    /// Per file QC information
54    pub per_file: Vec<MtxFileQcResult>,
55}
56
57/////////////
58// Helpers //
59/////////////
60
61/// Open one task's mtx file and read its shape in cell and gene terms.
62///
63/// ### Params
64///
65/// * `task` - The file task descriptor
66///
67/// ### Returns
68///
69/// `(source, no_cells, no_genes)`.
70fn open_task(task: &MtxFileTask) -> Result<(MtxSource, usize, usize), ScDataErrors> {
71    let src = MtxSource::open(Path::new(&task.mtx_path), task.cells_as_rows)?;
72    let [rows, cols, _] = src.shape;
73    let (no_cells, no_genes) = if task.cells_as_rows {
74        (rows, cols)
75    } else {
76        (cols, rows)
77    };
78    Ok((src, no_cells, no_genes))
79}
80
81/// Split a 1-based mtx coordinate into 0-based `(cell, gene)`.
82///
83/// ### Params
84///
85/// * `row` - 1-based row as stored
86/// * `col` - 1-based column as stored
87/// * `cells_as_rows` - Whether rows are cells
88///
89/// ### Returns
90///
91/// `None` for a 0 coordinate, which the format does not allow.
92#[inline]
93fn cell_gene(row: u32, col: u32, cells_as_rows: bool) -> Option<(usize, usize)> {
94    let (cell, gene) = if cells_as_rows {
95        (row, col)
96    } else {
97        (col, row)
98    };
99    Some((
100        (cell as usize).checked_sub(1)?,
101        (gene as usize).checked_sub(1)?,
102    ))
103}
104
105/// Count, per universe gene, the number of cells expressing it in one mtx file.
106///
107/// The blocks of the file are parsed in parallel, one tally per worker.
108///
109/// ### Params
110///
111/// * `task` - The mtx file to scan plus its local-to-universe gene mapping
112/// * `universe_size` - Number of genes in the intersection universe
113///
114/// ### Returns
115///
116/// Per-gene non-zero counts, indexed by universe gene index.
117fn scan_gene_nnz(task: &MtxFileTask, universe_size: usize) -> Result<Vec<usize>, ScDataErrors> {
118    let (src, _, _) = open_task(task)?;
119    let field = src.field;
120    let per_worker = src.par_fold_blocks(
121        || vec![0u32; universe_size],
122        |acc, _, bytes| {
123            parse_block(bytes, field, |row, col, _| {
124                if let Some((_, gene)) = cell_gene(row, col, task.cells_as_rows)
125                    && let Some(&Some(u)) = task.gene_local_to_universe.get(gene)
126                {
127                    acc[u] += 1;
128                }
129            })
130        },
131    )?;
132
133    let mut total = vec![0usize; universe_size];
134    for acc in per_worker {
135        total
136            .iter_mut()
137            .zip(acc)
138            .for_each(|(t, n)| *t += n as usize);
139    }
140    Ok(total)
141}
142
143/// Scan a single mtx file for per-cell NNZ and library size over the final gene
144/// set.
145///
146/// The blocks of the file are parsed in parallel, ignoring genes not present
147/// in `gene_local_to_final`.
148///
149/// ### Params
150///
151/// * `task` - The file task descriptor
152/// * `gene_local_to_final` - Mapping from file-local gene index to final gene
153///   index
154///
155/// ### Returns
156///
157/// Per-cell `(nnz, lib_size)` pairs indexed by file-local cell index
158fn scan_cell_stats(
159    task: &MtxFileTask,
160    gene_local_to_final: &[Option<usize>],
161) -> Result<Vec<(usize, u64)>, ScDataErrors> {
162    let (src, no_cells, _) = open_task(task)?;
163    let field = src.field;
164    let per_worker = src.par_fold_blocks(
165        || vec![(0u32, 0u64); no_cells],
166        |acc, _, bytes| {
167            parse_block(bytes, field, |row, col, value| {
168                if let Some((cell, gene)) = cell_gene(row, col, task.cells_as_rows)
169                    && cell < no_cells
170                    && let Some(&Some(_)) = gene_local_to_final.get(gene)
171                {
172                    acc[cell].0 += 1;
173                    acc[cell].1 += value as u64;
174                }
175            })
176        },
177    )?;
178
179    let mut stats = vec![(0usize, 0u64); no_cells];
180    for acc in per_worker {
181        stats.iter_mut().zip(acc).for_each(|(t, (u, l))| {
182            t.0 += u as usize;
183            t.1 += l;
184        });
185    }
186    Ok(stats)
187}
188
189/// Write the kept cells of a single mtx file into the unified sparse writer.
190///
191/// Reads the file block by block, filters to kept cells and final genes, and
192/// writes each cell's entries sorted by gene via [`write_cell_rows`].
193///
194/// ### Params
195///
196/// * `task` - The file task descriptor
197/// * `cells_to_keep` - File-local 0-indexed cell indices to include
198/// * `gene_local_to_final` - Mapping from file-local gene index to final gene
199///   index
200/// * `target_size` - Target library size for normalisation
201/// * `cell_offset` - Global cell offset for this file's cells in the unified
202///   output
203/// * `writer` - Mutable reference to the unified sparse writer
204/// * `verbose` - Controls verbosity
205///
206/// ### Returns
207///
208/// Per-file QC statistics for the written cells
209fn write_mtx_file_cells(
210    task: &MtxFileTask,
211    cells_to_keep: &[usize],
212    gene_local_to_final: &[Option<usize>],
213    target_size: f32,
214    cell_offset: usize,
215    writer: &mut CellGeneSparseWriter,
216    verbose: bool,
217) -> Result<MtxFileQcResult, ScDataErrors> {
218    if verbose {
219        println!(
220            "  Writing {} ({} cells)...",
221            task.exp_id,
222            cells_to_keep.len().separate_with_underscores()
223        );
224    }
225
226    let (src, _, _) = open_task(task)?;
227    let field = src.field;
228    let cell_map = dense_index_map(cells_to_keep);
229    let gene_map = dense_gene_map_from_options(gene_local_to_final);
230
231    // (gene_final_idx, raw_count) per kept cell
232    let mut cell_data: Vec<Vec<(u32, u32)>> = vec![Vec::new(); cells_to_keep.len()];
233
234    src.for_each_block_seq(|_, block| {
235        parse_block(&block, field, |row, col, value| {
236            let Some((cell, gene)) = cell_gene(row, col, task.cells_as_rows) else {
237                return;
238            };
239            let new_cell = cell_map.get(cell).copied().unwrap_or(INDEX_DROPPED);
240            let final_gene = gene_map.get(gene).copied().unwrap_or(INDEX_DROPPED);
241            if new_cell != INDEX_DROPPED && final_gene != INDEX_DROPPED {
242                cell_data[new_cell as usize].push((final_gene, value));
243            }
244        })
245    })?;
246
247    let (nnz, lib_size) = write_cell_rows(&mut cell_data, cell_offset, target_size, writer)?;
248
249    Ok(MtxFileQcResult {
250        exp_id: task.exp_id.clone(),
251        cells_to_keep: cells_to_keep.to_vec(),
252        lib_size,
253        nnz,
254    })
255}
256
257//////////
258// Main //
259//////////
260
261/// Load multiple mtx files into a single binary.
262///
263/// 1. Parallel per-file scan of gene NNZ against the intersected universe
264/// 2. Apply global `min_cells` to determine the final gene set
265/// 3. Parallel per-file scan of cell stats against the final gene set
266/// 4. Apply per-cell `min_unique_genes` / `min_lib_size`
267/// 5. Stream kept cells into the unified binary
268///
269/// ### Params
270///
271/// * `tasks` - Slice of per-file task descriptors
272/// * `bin_path` - Output path for the unified binary
273/// * `universe_size` - Total number of genes in the shared universe
274/// * `cell_qc` - Cell and gene quality thresholds and target library size
275/// * `verbose` - Controls verbosity
276///
277/// ### Returns
278///
279/// Summary of the global gene set, total cell and gene counts, and per-file QC
280/// results
281pub fn multi_mtx_to_file<P: AsRef<Path>>(
282    tasks: &[MtxFileTask],
283    bin_path: P,
284    universe_size: usize,
285    cell_qc: &MinCellQuality,
286    verbose: bool,
287) -> Result<MultiMtxResult, ScDataErrors> {
288    let total_start = Instant::now();
289
290    if verbose {
291        println!("Scan 1/2: gene NNZ across {} mtx files...", tasks.len());
292    }
293    let completed = Arc::new(AtomicUsize::new(0));
294    let report_interval = (tasks.len() / 10).max(1);
295
296    let per_file_nnz: Vec<Vec<usize>> = tasks
297        .par_iter()
298        .map(|task| {
299            let res = scan_gene_nnz(task, universe_size);
300            if verbose {
301                let done = completed.fetch_add(1, Ordering::Relaxed) + 1;
302                if done.is_multiple_of(report_interval) || done == tasks.len() {
303                    println!("  Scanned {}/{} files", done, tasks.len());
304                }
305            }
306            res
307        })
308        .collect::<Result<Vec<_>, _>>()?;
309
310    let mut global_gene_nnz = vec![0usize; universe_size];
311    for f in &per_file_nnz {
312        for (i, &c) in f.iter().enumerate() {
313            global_gene_nnz[i] += c;
314        }
315    }
316
317    let mut universe_to_final = vec![None; universe_size];
318    let mut final_idx = 0usize;
319    for (u, &nnz) in global_gene_nnz.iter().enumerate() {
320        if nnz >= cell_qc.min_cells {
321            universe_to_final[u] = Some(final_idx);
322            final_idx += 1;
323        }
324    }
325    let global_gene_indices: Vec<usize> = (0..universe_size)
326        .filter(|&i| universe_to_final[i].is_some())
327        .collect();
328    let total_genes = global_gene_indices.len();
329
330    if verbose {
331        println!(
332            "  Genes passing global QC: {} / {}",
333            total_genes.separate_with_underscores(),
334            universe_size.separate_with_underscores()
335        );
336    }
337
338    let composed: Vec<Vec<Option<usize>>> = tasks
339        .iter()
340        .map(|t| {
341            t.gene_local_to_universe
342                .iter()
343                .map(|opt| opt.and_then(|u| universe_to_final[u]))
344                .collect()
345        })
346        .collect();
347
348    if verbose {
349        println!("Scan 2/2: cell stats with final gene set...");
350    }
351    let completed = Arc::new(AtomicUsize::new(0));
352
353    let per_file_cell_stats: Vec<Vec<(usize, u64)>> = tasks
354        .par_iter()
355        .zip(composed.par_iter())
356        .map(|(task, mapping)| {
357            let res = scan_cell_stats(task, mapping);
358            if verbose {
359                let done = completed.fetch_add(1, Ordering::Relaxed) + 1;
360                if done.is_multiple_of(report_interval) || done == tasks.len() {
361                    println!("  Scanned {}/{} files", done, tasks.len());
362                }
363            }
364            res
365        })
366        .collect::<Result<Vec<_>, _>>()?;
367
368    let per_file_cells: Vec<Vec<usize>> = per_file_cell_stats
369        .iter()
370        .map(|stats| {
371            stats
372                .iter()
373                .enumerate()
374                .filter(|(_, (u, l))| {
375                    *u >= cell_qc.min_unique_genes && *l as f32 >= cell_qc.min_lib_size as f32
376                })
377                .map(|(i, _)| i)
378                .collect()
379        })
380        .collect();
381
382    let total_cells: usize = per_file_cells.iter().map(|v| v.len()).sum();
383
384    if verbose {
385        for (i, t) in tasks.iter().enumerate() {
386            println!(
387                "  {}: {} cells passing QC",
388                t.exp_id,
389                per_file_cells[i].len().separate_with_underscores()
390            );
391        }
392        println!("  Total cells: {}", total_cells.separate_with_underscores());
393    }
394
395    if verbose {
396        println!("Writing cells to binary...");
397    }
398    let mut writer = CellGeneSparseWriter::new(
399        &bin_path,
400        true,
401        total_cells,
402        total_genes,
403        cell_qc.target_size,
404    )?;
405    let mut cell_offset = 0usize;
406    let mut per_file_results = Vec::with_capacity(tasks.len());
407
408    for (idx, task) in tasks.iter().enumerate() {
409        let res = write_mtx_file_cells(
410            task,
411            &per_file_cells[idx],
412            &composed[idx],
413            cell_qc.target_size,
414            cell_offset,
415            &mut writer,
416            verbose,
417        )?;
418        cell_offset += per_file_cells[idx].len();
419        per_file_results.push(res);
420    }
421
422    writer.finalise()?;
423
424    if verbose {
425        println!("Multi-mtx loading complete: {:.2?}", total_start.elapsed());
426    }
427
428    Ok(MultiMtxResult {
429        global_gene_indices,
430        total_cells,
431        total_genes,
432        per_file: per_file_results,
433    })
434}
435
436///////////
437// Tests //
438///////////
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use std::io::Write;
444    use std::path::PathBuf;
445
446    /// Removes a test's scratch file even if an assert fails.
447    struct TempPath(PathBuf);
448
449    /// Drop implementation for [`TempPath`]. Errors are ignored: the file may
450    /// already be gone, and this runs during unwind.
451    impl Drop for TempPath {
452        fn drop(&mut self) {
453            let _ = std::fs::remove_file(&self.0);
454        }
455    }
456
457    impl TempPath {
458        /// Scratch file named after the calling test.
459        fn new(name: &str, ext: &str) -> Self {
460            Self(std::env::temp_dir().join(format!("bixverse_mtx_multi_{name}.{ext}")))
461        }
462
463        /// Path as a `String`.
464        fn string(&self) -> String {
465            self.0.to_str().expect("utf-8 temp path").to_string()
466        }
467    }
468
469    /// Three cells x four genes as integer counts, cells as rows.
470    const INTEGER_CELLS_AS_ROWS: &str = "%%MatrixMarket matrix coordinate integer general\n\
4713 4 5\n1 1 5\n1 3 7\n2 2 3\n3 1 2\n3 4 9\n";
472
473    /// Two cells x four genes as integral reals, genes as rows.
474    const REAL_GENES_AS_ROWS: &str = "%%MatrixMarket matrix coordinate real general\n\
4754 2 3\n1 1 4.0\n2 1 1.1e+01\n4 2 6\n";
476
477    /// One cell x four genes as a pattern, cells as rows.
478    const PATTERN_CELLS_AS_ROWS: &str = "%%MatrixMarket matrix coordinate pattern general\n\
4791 4 2\n1 2\n1 4\n";
480
481    /// Task over an identity gene universe of four genes.
482    fn task(name: &str, path: &TempPath, cells_as_rows: bool) -> MtxFileTask {
483        MtxFileTask {
484            exp_id: name.to_string(),
485            mtx_path: path.string(),
486            gene_local_to_universe: (0..4).map(Some).collect(),
487            cells_as_rows,
488        }
489    }
490
491    /// QC that keeps everything.
492    fn keep_all() -> MinCellQuality {
493        MinCellQuality {
494            min_unique_genes: 1,
495            min_lib_size: 1,
496            min_cells: 1,
497            target_size: 1e4,
498        }
499    }
500
501    /// Integer, gzipped real and pattern files load into one store with the
502    /// values each field type implies.
503    #[test]
504    fn multi_mtx_reads_every_field_type_and_gzip() {
505        let integer = TempPath::new("integer", "mtx");
506        std::fs::write(&integer.0, INTEGER_CELLS_AS_ROWS).expect("written");
507        let real = TempPath::new("real", "mtx.gz");
508        let mut enc = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
509        enc.write_all(REAL_GENES_AS_ROWS.as_bytes())
510            .expect("gzipped");
511        std::fs::write(&real.0, enc.finish().expect("gzipped")).expect("written");
512        let pattern = TempPath::new("pattern", "mtx");
513        std::fs::write(&pattern.0, PATTERN_CELLS_AS_ROWS).expect("written");
514        let bin = TempPath::new("store", "bin");
515
516        let tasks = [
517            task("integer", &integer, true),
518            task("real", &real, false),
519            task("pattern", &pattern, true),
520        ];
521        let res = multi_mtx_to_file(&tasks, &bin.0, 4, &keep_all(), false).expect("loaded");
522        assert_eq!(res.total_cells, 6);
523        assert_eq!(res.per_file[1].lib_size, vec![15, 6]);
524        assert_eq!(res.per_file[2].lib_size, vec![2]);
525
526        let reader = ParallelSparseReader::new(&bin.string()).expect("reader opens");
527        let cells: Vec<(Vec<u32>, Vec<u32>)> = reader
528            .read_cells_parallel(&[3, 4, 5])
529            .expect("cells read back")
530            .into_iter()
531            .map(|c| (c.indices.clone(), c.data_raw.iter().collect()))
532            .collect();
533        assert_eq!(
534            cells,
535            vec![
536                (vec![0, 1], vec![4, 11]),
537                (vec![3], vec![6]),
538                (vec![1, 3], vec![1, 1]),
539            ]
540        );
541    }
542
543    /// A fractional value fails the load rather than truncating.
544    #[test]
545    fn multi_mtx_rejects_a_fractional_real() {
546        let real = TempPath::new("fraction", "mtx");
547        std::fs::write(
548            &real.0,
549            "%%MatrixMarket matrix coordinate real general\n1 4 1\n1 1 2.5\n",
550        )
551        .expect("written");
552        let bin = TempPath::new("fraction_store", "bin");
553        assert!(matches!(
554            multi_mtx_to_file(
555                &[task("fraction", &real, true)],
556                &bin.0,
557                4,
558                &keep_all(),
559                false
560            ),
561            Err(ScDataErrors::MtxNotACount { .. })
562        ));
563    }
564}
bixverse_sc_data/mtx_multifile_io.rs

bixverse_sc_data/
mtx_multifile_io.rs
