comnum_spec/io/
npy.rs

1//! NPY arrays of NumPy (complex, real or integer).
2
3use super::Datatype;
4use super::raw::{decode_all, write_creating_dirs};
5use comnum_core::{Error, Result};
6use comnum_math::C64;
7use std::path::Path;
8
9/// The content of an NPY file.
10#[derive(Clone, Debug, PartialEq)]
11pub enum NpyData {
12    /// A complex array (`<c8`, `<c16`).
13    Complex(Vec<C64>),
14    /// A real array (`<f4`, `<f8`).
15    Real(Vec<f64>),
16    /// An integer or boolean array.
17    Int(Vec<i64>),
18}
19
20/// An NPY array: its shape and its data in C order.
21#[derive(Clone, Debug, PartialEq)]
22pub struct Npy {
23    /// The shape.
24    pub shape: Vec<usize>,
25    /// The data.
26    pub data: NpyData,
27}
28
29impl Npy {
30    /// The data as complex samples (real and integer arrays are converted).
31    pub fn to_complex(&self) -> Vec<C64> {
32        match &self.data {
33            NpyData::Complex(v) => v.clone(),
34            NpyData::Real(v) => v.iter().map(|&x| C64::new(x, 0.0)).collect(),
35            NpyData::Int(v) => v.iter().map(|&x| C64::new(x as f64, 0.0)).collect(),
36        }
37    }
38}
39
40fn header_value<'a>(header: &'a str, key: &str) -> Option<&'a str> {
41    let i = header.find(&format!("'{key}'"))? + key.len() + 2;
42    let rest = header[i..].trim_start().strip_prefix(':')?.trim_start();
43    let end = if rest.starts_with('(') {
44        rest.find(')')? + 1
45    } else if let Some(r) = rest.strip_prefix('\'') {
46        r.find('\'')? + 2
47    } else {
48        rest.find([',', '}']).unwrap_or(rest.len())
49    };
50    Some(rest[..end].trim())
51}
52
53/// Read an NPY array (little-endian numbers or booleans, C order or one dimension).
54///
55/// # Errors
56/// [`Error::Io`] for a missing file, [`Error::Config`] for an unsupported array.
57pub fn read_npy(path: impl AsRef<Path>) -> Result<Npy> {
58    let path = path.as_ref();
59    let b = std::fs::read(path)?;
60    let bad = |m: &str| Error::Config(format!("{}: {m}", path.display()));
61    if b.len() < 10 || &b[..6] != b"\x93NUMPY" {
62        return Err(bad("not an NPY file"));
63    }
64    let (hlen, start) = match b[6] {
65        1 => (u16::from_le_bytes([b[8], b[9]]) as usize, 10),
66        2 | 3 if b.len() >= 12 => (u32::from_le_bytes([b[8], b[9], b[10], b[11]]) as usize, 12),
67        _ => return Err(bad("unsupported NPY version")),
68    };
69    let header = std::str::from_utf8(b.get(start..start + hlen).ok_or_else(|| bad("truncated header"))?)
70        .map_err(|_| bad("header is not text"))?;
71    let descr = header_value(header, "descr").ok_or_else(|| bad("no descr"))?.trim_matches('\'');
72    let fortran = header_value(header, "fortran_order").ok_or_else(|| bad("no fortran_order"))? == "True";
73    let shape: Vec<usize> = header_value(header, "shape")
74        .ok_or_else(|| bad("no shape"))?
75        .trim_matches(|c| c == '(' || c == ')')
76        .split(',')
77        .filter(|t| !t.trim().is_empty())
78        .map(|t| t.trim().parse::<usize>().map_err(|_| bad("invalid shape")))
79        .collect::<Result<_>>()?;
80    if fortran && shape.iter().filter(|&&d| d > 1).count() > 1 {
81        return Err(bad("Fortran-ordered arrays of more than one dimension are not supported"));
82    }
83    let n = shape.iter().try_fold(1usize, |n, &d| n.checked_mul(d)).ok_or_else(|| bad("invalid shape"))?;
84    let data = &b[start + hlen..];
85    let need = |k: usize| -> Result<&[u8]> {
86        n.checked_mul(k).and_then(|len| data.get(..len)).ok_or_else(|| bad("truncated data"))
87    };
88    let le = |d: &str| d.starts_with('<') || d.starts_with('|') || d.starts_with('=');
89    if !le(descr) {
90        return Err(bad(&format!("big-endian or unknown dtype {descr:?}")));
91    }
92    let data = match &descr[1..] {
93        "c16" => NpyData::Complex(decode_all(need(16)?, Datatype::CF64)?),
94        "c8" => NpyData::Complex(decode_all(need(8)?, Datatype::CF32)?),
95        "f8" => NpyData::Real(
96            need(8)?.chunks_exact(8).map(|c| f64::from_le_bytes(c.try_into().unwrap_or([0; 8]))).collect(),
97        ),
98        "f4" => NpyData::Real(
99            need(4)?.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f64).collect(),
100        ),
101        "i8" => {
102            NpyData::Int(need(8)?.chunks_exact(8).map(|c| i64::from_le_bytes(c.try_into().unwrap_or([0; 8]))).collect())
103        }
104        "i4" => NpyData::Int(
105            need(4)?.chunks_exact(4).map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64).collect(),
106        ),
107        "i2" => NpyData::Int(need(2)?.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]]) as i64).collect()),
108        "i1" => NpyData::Int(need(1)?.iter().map(|&c| c as i8 as i64).collect()),
109        "u1" | "b1" => NpyData::Int(need(1)?.iter().map(|&c| c as i64).collect()),
110        "u2" => NpyData::Int(need(2)?.chunks_exact(2).map(|c| u16::from_le_bytes([c[0], c[1]]) as i64).collect()),
111        "u4" => NpyData::Int(
112            need(4)?.chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64).collect(),
113        ),
114        other => return Err(bad(&format!("unsupported dtype {other:?}"))),
115    };
116    Ok(Npy { shape, data })
117}
118
119fn write_npy_bytes(path: &Path, descr: &str, n: usize, payload: &[u8]) -> Result<()> {
120    let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': ({n},), }}");
121    // magic (6) + version (2) + length (2) + header + newline, padded to 64 bytes
122    let total = 10 + header.len() + 1;
123    header.push_str(&" ".repeat((64 - total % 64) % 64));
124    header.push('\n');
125    let mut out = b"\x93NUMPY\x01\x00".to_vec();
126    out.extend((header.len() as u16).to_le_bytes());
127    out.extend(header.as_bytes());
128    out.extend(payload);
129    write_creating_dirs(path, &out)
130}
131
132/// Write a one-dimensional complex NPY array (`<c16`).
133///
134/// # Errors
135/// [`Error::Io`].
136pub fn write_npy(path: impl AsRef<Path>, x: &[C64]) -> Result<()> {
137    let mut p = Vec::with_capacity(16 * x.len());
138    for &v in x {
139        Datatype::CF64.encode(v, &mut p);
140    }
141    write_npy_bytes(path.as_ref(), "<c16", x.len(), &p)
142}
143
144/// Write a one-dimensional integer NPY array (`<i8`), e.g. symbol indices.
145///
146/// # Errors
147/// [`Error::Io`].
148pub fn write_npy_i64(path: impl AsRef<Path>, x: &[i64]) -> Result<()> {
149    let p: Vec<u8> = x.iter().flat_map(|v| v.to_le_bytes()).collect();
150    write_npy_bytes(path.as_ref(), "<i8", x.len(), &p)
151}
comnum_spec/io/npy.rs

comnum_spec/io/
npy.rs
