From d8dad512d9bf7afbb69846879ee1e70e0c096bc5 Mon Sep 17 00:00:00 2001 From: Ahmed Sheikh Date: Thu, 22 Oct 2026 13:49:34 -0400 Subject: [PATCH] perf: parse computeMatrix matrices in Rust in read_matrix_file heatmapper.read_matrix_file, used by plotHeatmap, plotProfile and computeMatrixOperations, splits every line in Python and converts each value with float(). For an 82k x 3,600 matrix this takes about 20 s, half of plotHeatmap's runtime. hp.r_read_matrix (matrixreader.rs) decompresses the file on one thread and parses chunks of lines on a rayon pool with the GIL released. Values are parsed with str::parse::, which is correctly rounded like Python's float() and accepts nan/inf. It returns the parameter lines, the six region fields per row and the values as a float64 bytearray. read_matrix_file wraps the buffer with np.frombuffer as a MaskedArray without a mask (the same type the per-line reader produced) and builds regions and groups as before. The reader uses up to 16 threads from the CPU affinity mask. The per-line reader is kept as _read_matrix_file_legacy and is used for anything the fast path does not handle: file objects, rows of unequal length, tokens str::parse rejects (e.g. "1_0"), parameter lines after data rows, or an hp module without r_read_matrix. test_readMatrixFile.py reads every matrix in the test data with both readers and checks that matrix type, dtype, shape, mask, value bits (NaN and -0.0 included), regions, parameters, labels, boundaries and sorting settings are identical, and that the fallback is taken for a file object and for an unusual token. --- pydeeptools/deeptools/heatmapper.py | 41 +++++ .../deeptools/test/test_readMatrixFile.py | 94 +++++++++++ src/lib.rs | 2 + src/matrixreader.rs | 152 ++++++++++++++++++ 4 files changed, 289 insertions(+) create mode 100644 pydeeptools/deeptools/test/test_readMatrixFile.py create mode 100644 src/matrixreader.rs diff --git a/pydeeptools/deeptools/heatmapper.py b/pydeeptools/deeptools/heatmapper.py index 942c0cb887..c30aeef55a 100644 --- a/pydeeptools/deeptools/heatmapper.py +++ b/pydeeptools/deeptools/heatmapper.py @@ -709,6 +709,44 @@ def read_matrix_file(self, matrix_file): # file, this is considered as a delimiter # to split the heatmap into groups + import json + import os + + # Decompress and parse the matrix in Rust (in parallel); the values come back as float64 bytes. + # Anything the fast reader does not handle (file objects, ragged rows, unexpected tokens) + # is read with the per-line reader below. + try: + from deeptools.hp import r_read_matrix + try: + nthreads = len(os.sched_getaffinity(0)) + except AttributeError: + nthreads = os.cpu_count() or 1 + path = os.fsdecode(os.fspath(matrix_file)) + params, metas, buf, nrows, ncols = r_read_matrix(path, min(nthreads, 16)) + except (ImportError, TypeError, ValueError): + return self._read_matrix_file_legacy(matrix_file) + if not params or nrows == 0: + return self._read_matrix_file_legacy(matrix_file) + + for line in params: + self.parameters = json.loads(line.strip()[1:].strip()) + regions = [] + current_group_index = 0 + max_group_bound = self.parameters['group_boundaries'][1] + for chrom, start, end, name, score, strand in metas: + regs = [(int(x), int(y)) for x, y in zip(start.split(","), end.split(","))] + # get the group index + if len(regions) >= max_group_bound: + current_group_index += 1 + max_group_bound = self.parameters['group_boundaries'][current_group_index + 1] + regions.append([chrom, regs, name, max_group_bound, strand, score]) + # same array type as the per-line reader produces (MaskedArray without a mask) + matrix = np.ma.MaskedArray(np.frombuffer(buf, dtype=np.float64).reshape(nrows, ncols)) + del buf + self._set_matrix(regions, matrix) + + def _read_matrix_file_legacy(self, matrix_file): + # per-line reader, used for inputs the fast reader does not handle import json regions = [] matrix_rows = [] @@ -742,6 +780,9 @@ def read_matrix_file(self, matrix_file): regions.append([chrom, regs, name, max_group_bound, strand, score]) matrix = np.vstack(matrix_rows) + self._set_matrix(regions, matrix) + + def _set_matrix(self, regions, matrix): self.matrix = _matrix(regions, matrix, self.parameters['group_boundaries'], self.parameters['sample_boundaries'], group_labels=self.parameters['group_labels'], diff --git a/pydeeptools/deeptools/test/test_readMatrixFile.py b/pydeeptools/deeptools/test/test_readMatrixFile.py new file mode 100644 index 0000000000..63bf67497c --- /dev/null +++ b/pydeeptools/deeptools/test/test_readMatrixFile.py @@ -0,0 +1,94 @@ +import glob +import gzip +import os + +import numpy as np +import pytest + +from deeptools import heatmapper + +ROOT = os.path.dirname(os.path.abspath(__file__)) + + +def _matrix_files(): + files = [] + for f in sorted(glob.glob(os.path.join(ROOT, "**", "*.gz"), recursive=True)): + with gzip.open(f) as fh: + if fh.read(1) == b"@": + files.append(f) + return files + + +MATRIX_FILES = _matrix_files() + + +def _read_fast(path, monkeypatch): + # make sure the Rust reader is used and does not fall back to the per-line reader + def no_fallback(self, matrix_file): + raise AssertionError(f"fell back to the per-line reader for {matrix_file}") + + hm = heatmapper.heatmapper() + with monkeypatch.context() as m: + m.setattr(heatmapper.heatmapper, "_read_matrix_file_legacy", no_fallback) + hm.read_matrix_file(path) + return hm + + +def _read_legacy(path): + hm = heatmapper.heatmapper() + hm._read_matrix_file_legacy(path) + return hm + + +def assert_same_heatmapper(a, b): + ma, mb = a.matrix, b.matrix + assert type(ma.matrix) is type(mb.matrix) + assert ma.matrix.dtype == mb.matrix.dtype + assert ma.matrix.shape == mb.matrix.shape + assert np.array_equal(np.ma.getmaskarray(ma.matrix), np.ma.getmaskarray(mb.matrix)) + # compare the bits, so that NaN and -0.0 are checked as well + da = np.ascontiguousarray(np.ma.getdata(ma.matrix)) + db = np.ascontiguousarray(np.ma.getdata(mb.matrix)) + assert np.array_equal(da.view(np.int64), db.view(np.int64)) + assert ma.regions == mb.regions + assert ma.group_boundaries == mb.group_boundaries + assert ma.sample_boundaries == mb.sample_boundaries + assert ma.group_labels == mb.group_labels + assert ma.sample_labels == mb.sample_labels + assert ma.sort_method == mb.sort_method + assert ma.sort_using == mb.sort_using + assert a.parameters == b.parameters + + +def test_found_matrix_files(): + assert len(MATRIX_FILES) > 50 + + +@pytest.mark.parametrize("path", MATRIX_FILES, ids=lambda p: os.path.relpath(p, ROOT)) +def test_read_matrix_file_matches_legacy(path, monkeypatch): + assert_same_heatmapper(_read_fast(path, monkeypatch), _read_legacy(path)) + + +def test_read_matrix_file_fallback(tmp_path): + src = os.path.join(ROOT, "test_data", "computeMatrixOperations.mat.gz") + expected = _read_legacy(src) + + # file objects are read by the per-line reader + hm = heatmapper.heatmapper() + with open(src, "rb") as fh: + hm.read_matrix_file(fh) + assert_same_heatmapper(hm, expected) + + # so are values that only Python's float() accepts ("1_0" == 10.0) + with gzip.open(src, "rt") as fh: + lines = fh.read().split("\n") + fields = lines[1].split("\t") + fields[6] = "1_0" + lines[1] = "\t".join(fields) + odd = tmp_path / "odd.mat.gz" + with gzip.open(odd, "wt") as fh: + fh.write("\n".join(lines)) + hm = heatmapper.heatmapper() + hm.read_matrix_file(str(odd)) + assert hm.matrix.matrix[0, 0] == 10.0 + assert_same_heatmapper(hm, _read_legacy(str(odd))) diff --git a/src/lib.rs b/src/lib.rs index 9a37b0398e..b89b316eff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ mod computematrix; mod covcalc; mod filehandler; mod filtering; +mod matrixreader; mod multibamsummary; mod normalization; #[cfg(test)] @@ -19,5 +20,6 @@ fn hp(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(computematrix::r_computematrix, m)?)?; m.add_function(wrap_pyfunction!(alignmentsieve::r_alignmentsieve, m)?)?; m.add_function(wrap_pyfunction!(multibamsummary::r_mbams, m)?)?; + m.add_function(wrap_pyfunction!(matrixreader::r_read_matrix, m)?)?; Ok(()) } diff --git a/src/matrixreader.rs b/src/matrixreader.rs new file mode 100644 index 0000000000..7f73695722 --- /dev/null +++ b/src/matrixreader.rs @@ -0,0 +1,152 @@ +//! Parallel reader for computeMatrix's gzipped text matrix (used by +//! `heatmapper.read_matrix_file`, i.e. plotHeatmap / plotProfile / +//! computeMatrixOperations). +//! +//! Semantics follow the per-line Python reader: each line is stripped, lines +//! starting with '@' carry the JSON parameters, other lines are split on tabs +//! into 6 region fields and the values. Values are parsed with +//! `str::parse::`, which is correctly rounded like Python's `float()` +//! (and accepts nan/inf). Anything the fast path does not handle (a parameter +//! line after data rows, ragged rows, unparsable tokens) is reported as a +//! ValueError so the caller can fall back to the Python reader. + +use flate2::read::MultiGzDecoder; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyByteArray; +use rayon::prelude::*; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::sync::mpsc::sync_channel; + +/// Decompressed bytes per parse chunk. +const CHUNK_BYTES: usize = 32 << 20; + +type Meta = (String, String, String, String, String, String); + +fn parse_line(line: &[u8]) -> Result<(Meta, Vec), String> { + let line = std::str::from_utf8(line) + .map_err(|_| "non UTF-8 line".to_string())? + .trim_matches(|c: char| c.is_ascii_whitespace() || c == '\x0b'); + let mut it = line.split('\t'); + let mut f = || { + it.next() + .map(|s| s.to_string()) + .ok_or_else(|| "line with fewer than 6 fields".to_string()) + }; + let meta = (f()?, f()?, f()?, f()?, f()?, f()?); + let vals = it + .map(|t| { + t.trim_ascii() + .parse::() + .map_err(|_| format!("cannot parse value '{}'", t)) + }) + .collect::, String>>()?; + Ok((meta, vals)) +} + +/// Returns (parameter_lines, region_fields, values_as_f64_bytes, nrows, ncols). +#[pyfunction] +pub fn r_read_matrix<'py>( + py: Python<'py>, + path: &str, + nproc: usize, +) -> PyResult<( + Vec, + Vec, + Bound<'py, PyByteArray>, + usize, + usize, +)> { + let path = path.to_string(); + let res: Result<(Vec, Vec, Vec, usize), String> = py.detach(move || { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(nproc.max(1)) + .build() + .map_err(|e| e.to_string())?; + let file = File::open(&path).map_err(|e| format!("{}: {}", path, e))?; + // Decompress on a separate thread, handing over chunks of whole lines. + let (tx, rx) = sync_channel::>, String>>(2); + let reader = std::thread::spawn(move || { + let mut rd = BufReader::with_capacity(1 << 20, MultiGzDecoder::new(file)); + loop { + let mut lines: Vec> = Vec::new(); + let mut bytes = 0usize; + while bytes < CHUNK_BYTES { + let mut l = Vec::new(); + match rd.read_until(b'\n', &mut l) { + Ok(0) => break, + Ok(n) => { + bytes += n; + lines.push(l); + } + Err(e) => { + let _ = tx.send(Err(e.to_string())); + return; + } + } + } + let done = lines.is_empty(); + if !done && tx.send(Ok(lines)).is_err() { + return; + } + if done || bytes < CHUNK_BYTES { + return; + } + } + }); + let mut params: Vec = Vec::new(); + let mut metas: Vec = Vec::new(); + let mut values: Vec = Vec::new(); + let mut ncols: Option = None; + for chunk in rx { + let mut lines = chunk?; + // Parameter lines ('@') are only allowed before the first data row. + let mut skip = 0; + while skip < lines.len() && lines[skip].trim_ascii_start().first() == Some(&b'@') { + if !metas.is_empty() { + return Err("parameter line after data rows".to_string()); + } + let l = String::from_utf8(std::mem::take(&mut lines[skip])) + .map_err(|_| "non UTF-8 parameter line".to_string())?; + params.push(l); + skip += 1; + } + let parsed: Vec<(Meta, Vec)> = pool.install(|| { + lines[skip..] + .par_iter() + .map(|l| { + if l.trim_ascii_start().first() == Some(&b'@') { + Err("parameter line after data rows".to_string()) + } else { + parse_line(l) + } + }) + .collect::, String>>() + })?; + for (meta, row) in parsed { + match ncols { + None => ncols = Some(row.len()), + Some(n) if n != row.len() => return Err("rows of unequal length".to_string()), + _ => {} + } + values.extend_from_slice(&row); + metas.push(meta); + } + } + reader + .join() + .map_err(|_| "decompression thread panicked".to_string())?; + Ok((params, metas, values, ncols.unwrap_or(0))) + }); + let (params, metas, values, ncols) = res.map_err(PyValueError::new_err)?; + let nrows = metas.len(); + let nbytes = values.len() * std::mem::size_of::(); + let buf = PyByteArray::new_with(py, nbytes, |b: &mut [u8]| { + for (dst, v) in b.chunks_exact_mut(8).zip(values.iter()) { + dst.copy_from_slice(&v.to_ne_bytes()); + } + Ok(()) + })?; + Ok((params, metas, buf, nrows, ncols)) +}