sparse_npz 0.1.0

//! Reader and writer for **SciPy sparse matrices** stored in the NumPy `.npz`
//! format.
//!
//! `scipy.sparse.save_npz` writes a zip bundle of `.npy` arrays (`format`,
//! `shape`, `data`, `indices`, `indptr`). This crate reads those bundles, in
//! either CSC or CSR form, and writes them back in CSC form, exchanging the
//! matrix as a [`CscMatrix`]. The element dtype is preserved end to end: the
//! reader decodes `data` into a typed [`Values`] (boolean, any signed or
//! unsigned 8/16/32/64-bit integer, or 32/64-bit float) and the writer emits
//! that same dtype.
//!
//! Members are read whether they are stored, DEFLATE-compressed, or
//! Zstandard-compressed; they are written DEFLATE-compressed by default (as
//! SciPy does), or with Zstandard via [`Compression::Zstd`].
//!
//! Only little-endian, C-order `.npy` data is handled, which is what NumPy and
//! SciPy produce on the platforms targeted here.

#![forbid(unsafe_code)]

mod npy;

use std::io::{Cursor, Read, Seek, Write};
use std::path::Path;

use npy::NpyArray;
use zip::write::SimpleFileOptions;
use zip::CompressionMethod;

/// Anything that can go wrong while reading or writing a sparse `.npz`.
#[derive(Debug)]
pub enum Error {
    /// An underlying I/O or zip-container failure.
    Io(String),
    /// A malformed or unsupported `.npy` member.
    Npy(String),
    /// A structurally invalid or unsupported sparse layout.
    Format(String),
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::Io(m) => write!(f, "{m}"),
            Error::Npy(m) => write!(f, "invalid .npy member: {m}"),
            Error::Format(m) => write!(f, "unsupported sparse .npz: {m}"),
        }
    }
}

impl std::error::Error for Error {}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Self {
        Error::Io(e.to_string())
    }
}

macro_rules! numeric_values {
    ($($variant:ident => ($ty:ty, $descr:literal)),+ $(,)?) => {
        /// The stored (nonzero) values of a sparse matrix, tagged with their
        /// NumPy element dtype so it survives a read/write round trip. The
        /// vector length is the number of stored entries.
        #[derive(Clone, Debug, PartialEq)]
        pub enum Values {
            /// NumPy `bool_` (`|b1`).
            Bool(Vec<bool>),
            $( $variant(Vec<$ty>), )+
        }

        impl Values {
            /// The number of stored values.
            pub fn len(&self) -> usize {
                match self {
                    Values::Bool(v) => v.len(),
                    $( Values::$variant(v) => v.len(), )+
                }
            }

            /// Whether there are no stored values.
            pub fn is_empty(&self) -> bool {
                self.len() == 0
            }

            /// Every value widened to `f64` (booleans map to `1.0`/`0.0`).
            pub fn to_f64(&self) -> Vec<f64> {
                match self {
                    Values::Bool(v) => v.iter().map(|&b| if b { 1.0 } else { 0.0 }).collect(),
                    $( Values::$variant(v) => v.iter().map(|&x| x as f64).collect(), )+
                }
            }

            fn at_f64(&self, k: usize) -> f64 {
                match self {
                    Values::Bool(v) => if v[k] { 1.0 } else { 0.0 },
                    $( Values::$variant(v) => v[k] as f64, )+
                }
            }

            /// The NumPy dtype string this variant serializes as.
            fn descr(&self) -> &'static str {
                match self {
                    Values::Bool(_) => "|b1",
                    $( Values::$variant(_) => $descr, )+
                }
            }

            fn to_le_bytes(&self) -> Vec<u8> {
                match self {
                    Values::Bool(v) => v.iter().map(|&b| b as u8).collect(),
                    $( Values::$variant(v) => {
                        let mut out = Vec::with_capacity(v.len() * std::mem::size_of::<$ty>());
                        for &x in v {
                            out.extend_from_slice(&x.to_le_bytes());
                        }
                        out
                    } )+
                }
            }

            /// Builds a new `Values` whose element `i` is `self[order[i]]`.
            fn gather(&self, order: &[usize]) -> Values {
                match self {
                    Values::Bool(v) => Values::Bool(order.iter().map(|&k| v[k]).collect()),
                    $( Values::$variant(v) => Values::$variant(order.iter().map(|&k| v[k]).collect()), )+
                }
            }

            fn from_npy(a: &NpyArray) -> Result<Values, Error> {
                let d = &a.data;
                Ok(match a.descr.as_str() {
                    "|b1" | "<b1" => Values::Bool(d.iter().map(|&b| b != 0).collect()),
                    // Endianness-agnostic spellings of the 1-byte integer dtypes;
                    // the `|i1`/`|u1` forms NumPy emits are handled by the macro.
                    "<i1" => Values::I8(decode(d, i8::from_le_bytes)?),
                    "<u1" => Values::U8(decode(d, u8::from_le_bytes)?),
                    $( $descr => Values::$variant(decode(d, <$ty>::from_le_bytes)?), )+
                    other => return Err(Error::Npy(format!("unsupported data dtype {other}"))),
                })
            }
        }
    };
}

numeric_values! {
    I8 => (i8, "|i1"),
    I16 => (i16, "<i2"),
    I32 => (i32, "<i4"),
    I64 => (i64, "<i8"),
    U8 => (u8, "|u1"),
    U16 => (u16, "<u2"),
    U32 => (u32, "<u4"),
    U64 => (u64, "<u8"),
    F32 => (f32, "<f4"),
    F64 => (f64, "<f8"),
}

/// The on-disk sparse orientation of a SciPy `.npz`: compressed by column
/// (`csc`) or by row (`csr`).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Format {
    Csc,
    Csr,
}

/// The ZIP compression method used for the `.npy` members of a written `.npz`.
/// `Deflate` is what NumPy and SciPy emit and is universally readable; `Zstd`
/// (ZIP method 93) is smaller and faster but is only readable by newer readers
/// (numpy/scipy on Python 3.14 or later). Reading auto-detects either method.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Compression {
    /// DEFLATE, with an optional compression `level`. The usable range is
    /// `1..=264` (1..=9 via flate2, 10..=264 via Zopfli); `None` selects the
    /// default level of 6.
    Deflate { level: Option<i32> },
    /// Zstandard, with an optional compression `level`. The usable range is
    /// roughly `-7..=22`; `None` (or `0`) selects Zstandard's default level of 3.
    Zstd { level: Option<i32> },
}

impl Compression {
    fn options(self) -> SimpleFileOptions {
        let base = SimpleFileOptions::default();
        let (method, level) = match self {
            Compression::Deflate { level } => (CompressionMethod::Deflated, level),
            Compression::Zstd { level } => (CompressionMethod::Zstd, level),
        };
        let opts = base.compression_method(method);
        match level {
            Some(l) => opts.compression_level(Some(l as i64)),
            None => opts,
        }
    }
}

/// A sparse matrix in compressed-sparse-column (CSC) form, matching SciPy's
/// `csc_matrix` attributes. `col_ptr` has length `cols + 1`; the entries of
/// column `j` are `row_indices[col_ptr[j]..col_ptr[j + 1]]` with the parallel
/// values in `values`. This is the crate's internal orientation; CSR inputs are
/// transposed into it on read and can be written back out as CSR.
#[derive(Clone, Debug, PartialEq)]
pub struct CscMatrix {
    pub rows: usize,
    pub cols: usize,
    pub col_ptr: Vec<usize>,
    pub row_indices: Vec<usize>,
    pub values: Values,
}

impl CscMatrix {
    /// The number of stored (nonzero) entries.
    pub fn nnz(&self) -> usize {
        self.values.len()
    }

    /// Iterates the stored entries as `(row, column, value)` triples in
    /// column-major order, with each value widened to `f64`.
    pub fn entries(&self) -> impl Iterator<Item = (usize, usize, f64)> + '_ {
        (0..self.cols).flat_map(move |j| {
            (self.col_ptr[j]..self.col_ptr[j + 1]).map(move |k| (self.row_indices[k], j, self.values.at_f64(k)))
        })
    }

    /// Reads a SciPy sparse matrix from a `.npz` file, discarding the on-disk
    /// orientation. CSR inputs are transposed into CSC; CSC inputs are taken as
    /// is. The element dtype is preserved.
    pub fn read_npz<P: AsRef<Path>>(path: P) -> Result<CscMatrix, Error> {
        Ok(Self::read_npz_with_format(path)?.0)
    }

    /// Like [`read_npz`](CscMatrix::read_npz), but also returns the source
    /// [`Format`], so a caller can write the result back in the same orientation.
    pub fn read_npz_with_format<P: AsRef<Path>>(path: P) -> Result<(CscMatrix, Format), Error> {
        Self::read_npz_reader(std::fs::File::open(path)?)
    }

    /// Like [`read_npz_with_format`](CscMatrix::read_npz_with_format), but reads
    /// the `.npz` bundle from an in-memory byte slice instead of a file, so it can
    /// consume a matrix piped in on a standard input stream.
    pub fn read_npz_bytes(bytes: &[u8]) -> Result<(CscMatrix, Format), Error> {
        Self::read_npz_reader(Cursor::new(bytes))
    }

    /// Reads a `.npz` bundle from any seekable byte source.
    fn read_npz_reader<R: Read + Seek>(reader: R) -> Result<(CscMatrix, Format), Error> {
        let mut zip = zip::ZipArchive::new(reader).map_err(|e| Error::Io(e.to_string()))?;

        let member = |zip: &mut zip::ZipArchive<R>, name: &str| -> Result<NpyArray, Error> {
            let mut entry = zip.by_name(name).map_err(|_| Error::Format(format!("missing {name}")))?;
            let mut bytes = Vec::new();
            entry.read_to_end(&mut bytes)?;
            NpyArray::parse(&bytes).map_err(Error::Npy)
        };

        let format = member(&mut zip, "format.npy")?.as_ascii();
        let shape = member(&mut zip, "shape.npy")?.as_i64().map_err(Error::Npy)?;
        if shape.len() != 2 {
            return Err(Error::Format(format!("shape is not 2-D (got {} dims)", shape.len())));
        }
        let (rows, cols) = (shape[0] as usize, shape[1] as usize);
        let indptr = to_usize(member(&mut zip, "indptr.npy")?.as_i64().map_err(Error::Npy)?);
        let indices = to_usize(member(&mut zip, "indices.npy")?.as_i64().map_err(Error::Npy)?);
        let values = Values::from_npy(&member(&mut zip, "data.npy")?)?;

        match format.as_str() {
            "csc" => Ok((CscMatrix { rows, cols, col_ptr: indptr, row_indices: indices, values }, Format::Csc)),
            "csr" => Ok((csr_to_csc(rows, cols, &indptr, &indices, values), Format::Csr)),
            other => Err(Error::Format(format!("format '{other}' is not csc or csr"))),
        }
    }

    /// Writes this matrix to a `.npz` file in SciPy's CSC layout, DEFLATE-compressed.
    pub fn write_npz<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
        self.write_npz_with(path, Format::Csc, Compression::Deflate { level: None })
    }

    /// Writes this matrix to a `.npz` file in the given orientation, DEFLATE-compressed.
    pub fn write_npz_as<P: AsRef<Path>>(&self, path: P, format: Format) -> Result<(), Error> {
        self.write_npz_with(path, format, Compression::Deflate { level: None })
    }

    /// Writes this matrix to a `.npz` file in the given orientation and with the
    /// given ZIP compression, emitting the values in their own dtype. Index
    /// arrays use the narrowest of `int32`/`int64` that fits, matching SciPy's
    /// own choice.
    pub fn write_npz_with<P: AsRef<Path>>(&self, path: P, format: Format, compression: Compression) -> Result<(), Error> {
        self.write_npz_writer(std::fs::File::create(path)?, format, compression)?;
        Ok(())
    }

    /// Like [`write_npz_with`](CscMatrix::write_npz_with), but serializes the
    /// `.npz` bundle to a byte vector instead of a file, so it can be piped out on
    /// a standard output stream.
    pub fn write_npz_bytes(&self, format: Format, compression: Compression) -> Result<Vec<u8>, Error> {
        let cursor = self.write_npz_writer(Cursor::new(Vec::new()), format, compression)?;
        Ok(cursor.into_inner())
    }

    /// Writes this matrix's `.npz` bundle to any seekable byte sink, returning it.
    fn write_npz_writer<W: Write + Seek>(&self, writer: W, format: Format, compression: Compression) -> Result<W, Error> {
        let mut shape_bytes = Vec::with_capacity(16);
        shape_bytes.extend_from_slice(&(self.rows as i64).to_le_bytes());
        shape_bytes.extend_from_slice(&(self.cols as i64).to_le_bytes());

        // CSC stores its own arrays; CSR is the transpose, computed on demand.
        let transposed = match format {
            Format::Csc => None,
            Format::Csr => Some(self.to_csr()),
        };
        let (tag, indptr, indices, values) = match &transposed {
            None => ("csc", &self.col_ptr, &self.row_indices, &self.values),
            Some((row_ptr, col_indices, values)) => ("csr", row_ptr, col_indices, values),
        };

        let (indices_descr, indices_bytes) = encode_indices(indices);
        let (indptr_descr, indptr_bytes) = encode_indices(indptr);
        let members = [
            ("indices.npy", npy::write(indices_descr, &[indices.len()], &indices_bytes)),
            ("indptr.npy", npy::write(indptr_descr, &[indptr.len()], &indptr_bytes)),
            ("format.npy", npy::write("|S3", &[], tag.as_bytes())),
            ("shape.npy", npy::write("<i8", &[2], &shape_bytes)),
            ("data.npy", npy::write(values.descr(), &[values.len()], &values.to_le_bytes())),
        ];

        let mut zip = zip::ZipWriter::new(writer);
        let options = compression.options();
        for (name, bytes) in members {
            zip.start_file(name, options).map_err(|e| Error::Io(e.to_string()))?;
            zip.write_all(&bytes)?;
        }
        zip.finish().map_err(|e| Error::Io(e.to_string()))
    }

    /// Transposes this CSC matrix into CSR arrays `(row_ptr, col_indices,
    /// values)`, carrying the values (and their dtype) through the reordering.
    fn to_csr(&self) -> (Vec<usize>, Vec<usize>, Values) {
        let nnz = self.nnz();
        let mut row_ptr = vec![0usize; self.rows + 1];
        for &r in &self.row_indices {
            row_ptr[r + 1] += 1;
        }
        for i in 0..self.rows {
            row_ptr[i + 1] += row_ptr[i];
        }

        let mut next = row_ptr.clone();
        let mut col_indices = vec![0usize; nnz];
        let mut order = vec![0usize; nnz];
        for j in 0..self.cols {
            for k in self.col_ptr[j]..self.col_ptr[j + 1] {
                let r = self.row_indices[k];
                let dst = next[r];
                next[r] += 1;
                col_indices[dst] = j;
                order[dst] = k;
            }
        }

        (row_ptr, col_indices, self.values.gather(&order))
    }
}

fn decode<const N: usize, T>(data: &[u8], f: impl Fn([u8; N]) -> T) -> Result<Vec<T>, Error> {
    if data.len() % N != 0 {
        return Err(Error::Npy("data length not a multiple of element size".into()));
    }
    Ok(data
        .chunks_exact(N)
        .map(|c| {
            let mut a = [0u8; N];
            a.copy_from_slice(c);
            f(a)
        })
        .collect())
}

fn to_usize(v: Vec<i64>) -> Vec<usize> {
    v.into_iter().map(|x| x as usize).collect()
}

/// Transposes a CSR triple into the equivalent CSC matrix, carrying the values
/// (and their dtype) through the reordering.
fn csr_to_csc(rows: usize, cols: usize, indptr: &[usize], indices: &[usize], values: Values) -> CscMatrix {
    let nnz = values.len();
    let mut col_ptr = vec![0usize; cols + 1];
    for &c in indices {
        col_ptr[c + 1] += 1;
    }
    for j in 0..cols {
        col_ptr[j + 1] += col_ptr[j];
    }

    let mut next = col_ptr.clone();
    let mut row_indices = vec![0usize; nnz];
    let mut order = vec![0usize; nnz];
    for i in 0..rows {
        for k in indptr[i]..indptr[i + 1] {
            let c = indices[k];
            let dst = next[c];
            next[c] += 1;
            row_indices[dst] = i;
            order[dst] = k;
        }
    }

    CscMatrix { rows, cols, col_ptr, row_indices, values: values.gather(&order) }
}

/// Encodes an index array as the narrowest little-endian NumPy int dtype that
/// holds it, matching SciPy's own `int32`/`int64` choice.
fn encode_indices(values: &[usize]) -> (&'static str, Vec<u8>) {
    let fits_i32 = values.iter().all(|&v| v <= i32::MAX as usize);
    if fits_i32 {
        let mut bytes = Vec::with_capacity(values.len() * 4);
        for &v in values {
            bytes.extend_from_slice(&(v as i32).to_le_bytes());
        }
        ("<i4", bytes)
    } else {
        let mut bytes = Vec::with_capacity(values.len() * 8);
        for &v in values {
            bytes.extend_from_slice(&(v as i64).to_le_bytes());
        }
        ("<i8", bytes)
    }
}