winsfs_core/sfs/io/
npy.rs

1//! Reading and writing for SFS in the numpy npy format.
2//!
3//! The npy format is described [here][spec]. Only a subset required to read/write an 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::{fs::File, io, path::Path};
11
12use crate::sfs::{
13    generics::{Normalisation, Shape},
14    DynUSfs, SfsBase,
15};
16
17mod header;
18use header::{Endian, Header, HeaderDict, Type, TypeDescriptor, Version};
19
20/// Reads an SFS in npy format from a reader.
21///
22/// The stream is assumed to be positioned at the start.
23pub fn read_sfs<R>(reader: &mut R) -> io::Result<DynUSfs>
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            DynUSfs::from_vec_shape(values, dict.shape.into_boxed_slice()).map_err(|_| {
39                io::Error::new(io::ErrorKind::InvalidData, "npy shape does not fit values")
40            })
41        }
42    }
43}
44
45/// Reads an SFS in npy format from a file path.
46pub fn read_sfs_from_path<P>(path: P) -> io::Result<DynUSfs>
47where
48    P: AsRef<Path>,
49{
50    let mut reader = File::open(path).map(io::BufReader::new)?;
51    read_sfs(&mut reader)
52}
53
54/// Writes an SFS in npy format to a writer.
55pub fn write_sfs<W, S, N>(writer: &mut W, sfs: &SfsBase<S, N>) -> io::Result<()>
56where
57    W: io::Write,
58    S: Shape,
59    N: Normalisation,
60{
61    let header = Header::new(
62        Version::V1,
63        HeaderDict::new(
64            TypeDescriptor::new(Endian::Little, Type::F8),
65            false,
66            sfs.shape().as_ref().to_vec(),
67        ),
68    );
69
70    header.write(writer)?;
71
72    for v in sfs.iter() {
73        writer.write_all(&v.to_le_bytes())?;
74    }
75
76    Ok(())
77}
78
79/// Writes an SFS in npy format to a file path.
80///
81/// If the file already exists, it will be overwritten.
82pub fn write_sfs_to_path<P, S, N>(path: P, sfs: &SfsBase<S, N>) -> io::Result<()>
83where
84    P: AsRef<Path>,
85    S: Shape,
86    N: Normalisation,
87{
88    let mut writer = File::create(path)?;
89    write_sfs(&mut writer, sfs)
90}
winsfs_core/sfs/io/npy.rs

winsfs_core/sfs/io/
npy.rs
