sfs_core/array/
npy.rs

1//! Reading and writing in the numpy npy format.
2//!
3//! The npy format is described [here][spec]. Only a subset required to read/write SFS
4//! is supported. Only simple type descriptors for the basic integer and float types are
5//! supported. In addition, only reading/writing C-order is supported; trying to read a
6//! Fortran-order npy file will result in a run-time error.
7//!
8//! [spec]: https://numpy.org/neps/nep-0001-npy-format.html
9
10use std::io;
11
12use super::{Array, Shape};
13
14mod header;
15use header::{Endian, Header, HeaderDict, Type, TypeDescriptor, Version};
16
17/// The npy magic number.
18pub(crate) const MAGIC: [u8; 6] = *b"\x93NUMPY";
19
20/// Reads an array in npy format from a reader.
21///
22/// The stream is assumed to be positioned at the start.
23pub fn read_array<R>(reader: &mut R) -> io::Result<Array<f64>>
24where
25    R: io::BufRead,
26{
27    let header = Header::read(reader)?;
28    let dict = header.dict;
29
30    match (dict.type_descriptor, dict.fortran_order) {
31        (_, true) => Err(io::Error::new(
32            io::ErrorKind::InvalidData,
33            "Fortran order not supported when reading npy",
34        )),
35        (descr, false) => {
36            let values = descr.read(reader)?;
37
38            Array::new(values, Shape(dict.shape)).map_err(|_| {
39                io::Error::new(io::ErrorKind::InvalidData, "npy shape does not fit values")
40            })
41        }
42    }
43}
44
45/// Writes an array in npy format to a writer.
46pub fn write_array<W>(writer: &mut W, array: &Array<f64>) -> io::Result<()>
47where
48    W: io::Write,
49{
50    let header = Header::new(
51        Version::V1,
52        HeaderDict::new(
53            TypeDescriptor::new(Endian::Little, Type::F8),
54            false,
55            array.shape().as_ref().to_vec(),
56        ),
57    );
58
59    header.write(writer)?;
60
61    for v in array.iter() {
62        writer.write_all(&v.to_le_bytes())?;
63    }
64
65    Ok(())
66}
sfs_core/array/npy.rs

sfs_core/array/
npy.rs
