perf: parse computeMatrix matrices in Rust in read_matrix_file

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'],

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)))

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(())
 }

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::<f64>`, 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<f64>), 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::<f64>()
+                .map_err(|_| format!("cannot parse value '{}'", t))
+        })
+        .collect::<Result<Vec<f64>, 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<String>,
+    Vec<Meta>,
+    Bound<'py, PyByteArray>,
+    usize,
+    usize,
+)> {
+    let path = path.to_string();
+    let res: Result<(Vec<String>, Vec<Meta>, Vec<f64>, 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::<Result<Vec<Vec<u8>>, 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<u8>> = 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<String> = Vec::new();
+        let mut metas: Vec<Meta> = Vec::new();
+        let mut values: Vec<f64> = Vec::new();
+        let mut ncols: Option<usize> = 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<f64>)> = 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::<Result<Vec<_>, 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::<f64>();
+    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))
+}