Skip to main content

data_beans/export/
mtx.rs

1use crate::sparse_io::*;
2use crate::utilities::name_matching::split_id_name;
3
4use legume_numeric::matrix::common_io::{mkdir, open_buf_writer, write_lines};
5use log::{info, warn};
6use std::io::Write;
7
8/// Cells per block when streaming the matrix out (bounds peak memory to one
9/// block of triplets regardless of matrix size).
10const COL_BLOCK: usize = 512;
11
12/// How [`to_mtx`] writes the triplet.
13#[derive(Clone, Debug)]
14pub struct MtxExport {
15    /// Gzip the three files (`matrix.mtx.gz`, ...), the Cell Ranger v3
16    /// layout; else plain `matrix.mtx`, `features.tsv`, `barcodes.tsv`.
17    pub gzip: bool,
18    /// Written as every feature's `feature_type`, the third column of
19    /// `features.tsv`: 10x uses `Gene Expression`, `Antibody Capture` or
20    /// `Peaks`.
21    pub feature_type: Box<str>,
22}
23
24impl Default for MtxExport {
25    /// Gzipped, every feature `Gene Expression`: what `to-mtx` writes.
26    fn default() -> Self {
27        Self {
28            gzip: true,
29            feature_type: "Gene Expression".into(),
30        }
31    }
32}
33
34/// Export `data` to a 10x Genomics MEX (MatrixMarket) triplet in the
35/// directory `out_dir`, created if missing, readable by scanpy / Seurat.
36/// This is the inverse of [`crate::import::from_mtx`]; the directory is
37/// returned.
38///
39/// data-beans already stores the matrix as (features x cells), which is 10x's
40/// orientation (features are rows, barcodes are columns), so no transpose is
41/// needed — unlike `to_h5ad`. A feature named `id_name` is split back into
42/// its id and name columns.
43pub fn to_mtx(
44    data: &dyn SparseIo<IndexIter = Vec<usize>>,
45    out_dir: &str,
46    opts: &MtxExport,
47) -> anyhow::Result<Box<str>> {
48    let nrow = data
49        .num_rows()
50        .ok_or_else(|| anyhow::anyhow!("backend has no `nrow`"))?;
51    let ncol = data
52        .num_columns()
53        .ok_or_else(|| anyhow::anyhow!("backend has no `ncol`"))?;
54    let nnz = data
55        .num_non_zeros()
56        .ok_or_else(|| anyhow::anyhow!("backend has no `nnz`"))?;
57
58    let row_names = data.row_names()?; // features
59    let col_names = data.column_names()?; // barcodes
60    if row_names.len() != nrow {
61        anyhow::bail!("row name count {} != nrow {}", row_names.len(), nrow);
62    }
63    if col_names.len() != ncol {
64        anyhow::bail!("column name count {} != ncol {}", col_names.len(), ncol);
65    }
66
67    let dir = out_dir.trim_end_matches('/');
68    mkdir(dir)?;
69    let suffix = if opts.gzip { ".gz" } else { "" };
70    let mtx_path = format!("{}/matrix.mtx{}", dir, suffix);
71    let features_path = format!("{}/features.tsv{}", dir, suffix);
72    let barcodes_path = format!("{}/barcodes.tsv{}", dir, suffix);
73
74    info!(
75        "Exporting to {}/ ({} features x {} barcodes, {} non-zeros)",
76        dir, nrow, ncol, nnz
77    );
78
79    // features.tsv: 10x 3-column layout (id, name, feature_type). Recover the
80    // id/name split from the backend's composite `id_name` display name.
81    let feature_lines: Vec<Box<str>> = row_names
82        .iter()
83        .map(|composite| {
84            let (id, name) = split_id_name(composite);
85            format!("{}\t{}\t{}", id, name, opts.feature_type).into_boxed_str()
86        })
87        .collect();
88    write_lines(&feature_lines, &features_path)?;
89    info!("wrote {}", features_path);
90
91    // barcodes.tsv: one barcode per line.
92    write_lines(&col_names, &barcodes_path)?;
93    info!("wrote {}", barcodes_path);
94
95    // matrix.mtx: stream the CSC columns as 1-based (feature, barcode, count).
96    write_mtx_10x(data, &mtx_path, nrow, ncol, nnz)?;
97    info!("wrote {}", mtx_path);
98
99    info!("done: {}/", dir);
100    Ok(dir.into())
101}
102
103/// Write a Cell Ranger v3 style `matrix.mtx[.gz]`: a MatrixMarket coordinate
104/// file with an `integer` field, a `%metadata_json` comment, 1-based indices
105/// and features as rows / barcodes as columns. Values are streamed a column
106/// block at a time. Non-integer values (a non-count backend) are rounded, with
107/// a single warning — 10x MEX is defined over integer counts.
108fn write_mtx_10x(
109    data: &dyn SparseIo<IndexIter = Vec<usize>>,
110    path: &str,
111    nrow: usize,
112    ncol: usize,
113    nnz: usize,
114) -> anyhow::Result<()> {
115    let mut buf = open_buf_writer(path)?;
116    writeln!(buf, "%%MatrixMarket matrix coordinate integer general")?;
117    writeln!(
118        buf,
119        "%metadata_json: {{\"software_version\": \"data-beans-{}\", \"format_version\": 2}}",
120        env!("CARGO_PKG_VERSION")
121    )?;
122    writeln!(buf, "{} {} {}", nrow, ncol, nnz)?;
123
124    let mut warned_fractional = false;
125    let mut lo = 0usize;
126    while lo < ncol {
127        let hi = (lo + COL_BLOCK).min(ncol);
128        let (_, _, triplets) = data.read_triplets_by_columns((lo..hi).collect())?;
129        for (row_i, col_local, val) in triplets {
130            let count = val.round();
131            if !warned_fractional && (val - count).abs() > 1e-6 {
132                warn!(
133                    "matrix holds non-integer values (e.g. {}); rounding to integers \
134                     for the 10x `integer` MatrixMarket field",
135                    val
136                );
137                warned_fractional = true;
138            }
139            // 1-based; global column = block offset + local index.
140            writeln!(
141                buf,
142                "{} {} {}",
143                row_i + 1,
144                lo as u64 + col_local + 1,
145                count as i64
146            )?;
147        }
148        lo = hi;
149    }
150    buf.flush()?;
151    Ok(())
152}