gwseq-io 0.2.0

//! `save_to`: writing a read straight to a `.npz` instead of through Python.
//!
//! The point of it is the read that does not fit in
//! memory twice — a whole-chromosome contact matrix is hundreds of megabytes,
//! and handing it to numpy only for `np.savez` to copy it out again doubles
//! that. `ndarray-npy` does the encoding; what this file owns is *which* arrays
//! go in under *which* keys, since that is the part callers depend on.

use ndarray::{Array1, Array2, ArrayView1};
use ndarray_npy::NpzWriter;

use crate::arrays::CooMatrix;
use crate::error::{Error, Result};

/// Map an npz failure onto an I/O error naming the file, which is what the
/// caller has to act on either way.
fn wrap(path: &str, err: impl std::fmt::Display) -> Error {
    Error::io(path, std::io::Error::other(err.to_string()))
}

fn writer(path: &str) -> Result<NpzWriter<std::fs::File>> {
    let file = std::fs::File::create(path).map_err(|e| Error::io(path, e))?;
    Ok(NpzWriter::new(file))
}

/// Write a dense matrix under the `values` key.
pub fn write_dense(path: &str, values: &Array2<f32>) -> Result<()> {
    let mut npz = writer(path)?;
    npz.add_array("values", values).map_err(|e| wrap(path, e))?;
    npz.finish().map_err(|e| wrap(path, e))?;
    Ok(())
}

/// Write a COO matrix under the four keys `read_sparse_values` documents.
///
/// `shape` goes in as a two-element `uint32` array rather than a tuple: an npz
/// holds arrays, and this is the spelling `numpy.load(...)["shape"]` reads back
/// as the pair the dict form hands out.
pub fn write_sparse(path: &str, coo: &CooMatrix) -> Result<()> {
    let mut npz = writer(path)?;
    // Views, not clones. This module exists so that a HiC read reaches a file
    // without being held twice, and three `Vec::clone`s of a sparse matrix that
    // can run to millions of entries was the thing it was there to avoid.
    npz.add_array("values", &ArrayView1::from(&coo.values[..]))
        .map_err(|e| wrap(path, e))?;
    npz.add_array("row", &ArrayView1::from(&coo.row[..]))
        .map_err(|e| wrap(path, e))?;
    npz.add_array("col", &ArrayView1::from(&coo.col[..]))
        .map_err(|e| wrap(path, e))?;
    npz.add_array(
        "shape",
        &Array1::from(vec![coo.shape.0 as u32, coo.shape.1 as u32]),
    )
    .map_err(|e| wrap(path, e))?;
    npz.finish().map_err(|e| wrap(path, e))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The local-name field of every entry of a stored zip, in order.
    fn entry_names(bytes: &[u8]) -> Vec<String> {
        let mut names = Vec::new();
        let mut i = 0;
        while i + 30 <= bytes.len() {
            if bytes[i..i + 4] != [0x50, 0x4b, 0x03, 0x04] {
                i += 1;
                continue;
            }
            let name_len = u16::from_le_bytes([bytes[i + 26], bytes[i + 27]]) as usize;
            names.push(String::from_utf8_lossy(&bytes[i + 30..i + 30 + name_len]).into_owned());
            i += 30 + name_len;
        }
        names
    }

    #[test]
    fn a_dense_save_writes_one_entry_named_values() {
        let dir = std::env::temp_dir().join("gwseq_npz_dense");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("out.npz");
        let values = Array2::from_shape_vec((2, 3), (0..6).map(|v| v as f32).collect()).unwrap();
        write_dense(path.to_str().unwrap(), &values).unwrap();

        let bytes = std::fs::read(&path).unwrap();
        assert_eq!(entry_names(&bytes), ["values.npy"]);
        // The header is the one numpy reads: the magic (0x93 "NUMPY", which is
        // why this looks for the bytes rather than the text), and the shape it
        // was given rather than a flattened one.
        assert!(
            bytes.windows(6).any(|w| w == b"\x93NUMPY"),
            "no .npy magic in the entry"
        );
        let text = String::from_utf8_lossy(&bytes[..200]);
        assert!(text.contains("'shape': (2, 3)"), "{text:?}");
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn a_sparse_save_writes_the_four_documented_keys() {
        let dir = std::env::temp_dir().join("gwseq_npz_sparse");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("out.npz");
        let coo = CooMatrix {
            values: vec![1.0, 2.0],
            row: vec![0, 1],
            col: vec![1, 0],
            shape: (2, 2),
        };
        write_sparse(path.to_str().unwrap(), &coo).unwrap();

        let bytes = std::fs::read(&path).unwrap();
        assert_eq!(
            entry_names(&bytes),
            ["values.npy", "row.npy", "col.npy", "shape.npy"]
        );
        std::fs::remove_dir_all(&dir).ok();
    }
}