gwseq_io/
npz.rs

1//! `save_to`: writing a read straight to a `.npz` instead of through Python.
2//!
3//! The point of it is the read that does not fit in
4//! memory twice — a whole-chromosome contact matrix is hundreds of megabytes,
5//! and handing it to numpy only for `np.savez` to copy it out again doubles
6//! that. `ndarray-npy` does the encoding; what this file owns is *which* arrays
7//! go in under *which* keys, since that is the part callers depend on.
8
9use ndarray::{Array1, Array2, ArrayView1};
10use ndarray_npy::NpzWriter;
11
12use crate::arrays::CooMatrix;
13use crate::error::{Error, Result};
14
15/// Map an npz failure onto an I/O error naming the file, which is what the
16/// caller has to act on either way.
17fn wrap(path: &str, err: impl std::fmt::Display) -> Error {
18    Error::io(path, std::io::Error::other(err.to_string()))
19}
20
21fn writer(path: &str) -> Result<NpzWriter<std::fs::File>> {
22    let file = std::fs::File::create(path).map_err(|e| Error::io(path, e))?;
23    Ok(NpzWriter::new(file))
24}
25
26/// Write a dense matrix under the `values` key.
27pub fn write_dense(path: &str, values: &Array2<f32>) -> Result<()> {
28    let mut npz = writer(path)?;
29    npz.add_array("values", values).map_err(|e| wrap(path, e))?;
30    npz.finish().map_err(|e| wrap(path, e))?;
31    Ok(())
32}
33
34/// Write a COO matrix under the four keys `read_sparse_values` documents.
35///
36/// `shape` goes in as a two-element `uint32` array rather than a tuple: an npz
37/// holds arrays, and this is the spelling `numpy.load(...)["shape"]` reads back
38/// as the pair the dict form hands out.
39pub fn write_sparse(path: &str, coo: &CooMatrix) -> Result<()> {
40    let mut npz = writer(path)?;
41    // Views, not clones. This module exists so that a HiC read reaches a file
42    // without being held twice, and three `Vec::clone`s of a sparse matrix that
43    // can run to millions of entries was the thing it was there to avoid.
44    npz.add_array("values", &ArrayView1::from(&coo.values[..]))
45        .map_err(|e| wrap(path, e))?;
46    npz.add_array("row", &ArrayView1::from(&coo.row[..]))
47        .map_err(|e| wrap(path, e))?;
48    npz.add_array("col", &ArrayView1::from(&coo.col[..]))
49        .map_err(|e| wrap(path, e))?;
50    npz.add_array(
51        "shape",
52        &Array1::from(vec![coo.shape.0 as u32, coo.shape.1 as u32]),
53    )
54    .map_err(|e| wrap(path, e))?;
55    npz.finish().map_err(|e| wrap(path, e))?;
56    Ok(())
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62
63    /// The local-name field of every entry of a stored zip, in order.
64    fn entry_names(bytes: &[u8]) -> Vec<String> {
65        let mut names = Vec::new();
66        let mut i = 0;
67        while i + 30 <= bytes.len() {
68            if bytes[i..i + 4] != [0x50, 0x4b, 0x03, 0x04] {
69                i += 1;
70                continue;
71            }
72            let name_len = u16::from_le_bytes([bytes[i + 26], bytes[i + 27]]) as usize;
73            names.push(String::from_utf8_lossy(&bytes[i + 30..i + 30 + name_len]).into_owned());
74            i += 30 + name_len;
75        }
76        names
77    }
78
79    #[test]
80    fn a_dense_save_writes_one_entry_named_values() {
81        let dir = std::env::temp_dir().join("gwseq_npz_dense");
82        std::fs::create_dir_all(&dir).unwrap();
83        let path = dir.join("out.npz");
84        let values = Array2::from_shape_vec((2, 3), (0..6).map(|v| v as f32).collect()).unwrap();
85        write_dense(path.to_str().unwrap(), &values).unwrap();
86
87        let bytes = std::fs::read(&path).unwrap();
88        assert_eq!(entry_names(&bytes), ["values.npy"]);
89        // The header is the one numpy reads: the magic (0x93 "NUMPY", which is
90        // why this looks for the bytes rather than the text), and the shape it
91        // was given rather than a flattened one.
92        assert!(
93            bytes.windows(6).any(|w| w == b"\x93NUMPY"),
94            "no .npy magic in the entry"
95        );
96        let text = String::from_utf8_lossy(&bytes[..200]);
97        assert!(text.contains("'shape': (2, 3)"), "{text:?}");
98        std::fs::remove_dir_all(&dir).ok();
99    }
100
101    #[test]
102    fn a_sparse_save_writes_the_four_documented_keys() {
103        let dir = std::env::temp_dir().join("gwseq_npz_sparse");
104        std::fs::create_dir_all(&dir).unwrap();
105        let path = dir.join("out.npz");
106        let coo = CooMatrix {
107            values: vec![1.0, 2.0],
108            row: vec![0, 1],
109            col: vec![1, 0],
110            shape: (2, 2),
111        };
112        write_sparse(path.to_str().unwrap(), &coo).unwrap();
113
114        let bytes = std::fs::read(&path).unwrap();
115        assert_eq!(
116            entry_names(&bytes),
117            ["values.npy", "row.npy", "col.npy", "shape.npy"]
118        );
119        std::fs::remove_dir_all(&dir).ok();
120    }
121}
gwseq_io/npz.rs

gwseq_io/
npz.rs
