lumen_io/
npy.rs

1use lumen_core::Shape;
2use lumen_core::Tensor;
3use lumen_core::WithDType;
4use thiserrorctx::Context;
5use zip::result::ZipError;
6use zip::write::SimpleFileOptions;
7use zip::ZipArchive;
8use zip::ZipWriter;
9use std::collections::HashMap;
10use std::io::Cursor;
11use std::io::Seek;
12use std::str::Utf8Error;
13use std::{fs::File, io::{BufReader, BufWriter, Read, Write}, path::Path};
14use lumen_core::{DynTensor, DType};
15use crate::utils;
16
17#[thiserrorctx::context_error]
18pub enum NpyError {
19    #[error(transparent)]
20    Core(#[from] lumen_core::Error),
21
22    #[error(transparent)]
23    Zip(#[from] ZipError),
24
25    #[error(transparent)]
26    Io(#[from] std::io::Error),
27
28    #[error(transparent)]
29    Utf8(#[from] Utf8Error),
30
31    // === Npy Error ===
32    #[error("not a npy file for can't read correct magic")]
33    NotNpyFile,
34
35    #[error("unsupported NPY version {major}.{minor}")]
36    UnsupportVersion {
37        major: u8,
38        minor: u8,
39    },
40
41    #[error("{error} in {field} field")]
42    Header {
43        error: String,
44        field: &'static str 
45    },
46
47    #[error("no exit {field} field")]
48    LackField {
49        field: &'static str 
50    },
51
52    #[error("unsupport descr {0}")]
53    UnsupportedDescr(String),
54
55    #[error("unsupport dtype {0}")]
56    UnsupportedDType(DType),
57
58    #[error("Numpy does't support usize dtype")]
59    UnsupportUSize,
60}
61
62const NPY_MAGIC: &[u8] = b"\x93NUMPY";
63
64
65/// Load a single `.npy` file into a [`DynTensor`].
66///
67/// # Arguments
68/// * `path` - Path to the `.npy` file.
69///
70/// # Errors
71/// Returns an error if the file cannot be opened, read, or parsed.
72pub fn load_npy_file<P: AsRef<Path>>(path: P) -> NpyResult<DynTensor> {
73    let file = File::open(path).map_err(NpyError::Io).context("open file")?;
74    let mut reader = BufReader::new(file);
75    load_npy(&mut reader)
76}
77
78pub fn load_npy<R: Read + Seek>(mut reader: &mut R) -> NpyResult<DynTensor> {
79    // Read magic
80    let mut magic = [0u8; 6];
81    reader.read_exact(&mut magic)?;
82    if &magic != NPY_MAGIC {
83        Err( NpyError::NotNpyFile )?;
84    }
85
86    // Read version
87    let mut version = [0u8; 2];
88    reader.read_exact(&mut version)?;
89    let version = (version[0], version[1]);
90
91    // Read header
92    let header_len = match version {
93        (1, 0) => {
94            let mut buf = [0u8; 2];
95            reader.read_exact(&mut buf)?;
96            u16::from_le_bytes(buf) as usize
97        }
98        (2, 0) | (3, 0) => {
99            let mut buf = [0u8; 4];
100            reader.read_exact(&mut buf)?;
101            u32::from_le_bytes(buf) as usize
102        }
103        _ => Err( NpyError::UnsupportVersion { major: version.0, minor: version.1 })?,
104    };
105
106    // Read Header
107    let mut header_bytes = vec![0u8; header_len];   
108    reader.read_exact(&mut header_bytes)?;
109    let header_str = str::from_utf8(&header_bytes)?.trim();
110
111    // Parse header
112    // "{'descr': '<f8', 'fortran_order': False, 'shape': (3, 4), }"
113    let header_str = header_str.trim_matches(|c| c == '{' || c == '}');
114
115    // 'descr': '<f8', 'fortran_order': False, 'shape': (3, 4), 
116    let mut descr_opt: Option<&str> = None;
117    let mut fortran_order_opt: Option<bool> = None;
118    let mut shape_opt: Option<Vec<usize>> = None;
119
120    let mut header_str = header_str;
121    while !header_str.is_empty() {
122        if header_str.starts_with("'descr'") {
123            let colon_index = header_str.find(",")
124                .ok_or_else(|| NpyError::Header { 
125                    field: "descr", 
126                    error: "No colon".into() 
127                })?;
128            // ": '<f8'"
129            let descr = &header_str[7..colon_index];
130            let ref1_index = descr.find("'")
131                .ok_or_else(|| NpyError::Header { 
132                    field: "descr", 
133                    error: "No start ' in 'descr' field value".into() 
134                })?;
135            // "<f8'"
136            let descr = &descr[ref1_index+1..];
137            let ref2_index = descr.find("'")
138                .ok_or_else(|| NpyError::Header { 
139                    field: "descr", 
140                    error: "No end ' in 'descr' field value".into() 
141                })?;
142            // "<f8"
143            let descr = &descr[..ref2_index];
144            descr_opt = Some(descr);
145
146            header_str = header_str[colon_index+1..].trim_start();
147        } else if header_str.starts_with("'fortran_order'") {
148            let colon_index = header_str.find(",")
149                .ok_or_else(|| NpyError::Header { 
150                    field: "fortran_order", 
151                    error: "No colon in 'fortran_order' field".into() 
152                })?;
153            // ": False"
154            let fortran_order = &header_str[15..colon_index];
155            let index = fortran_order.find(":")
156                .ok_or_else(|| NpyError::Header { 
157                    field: "fortran_order", 
158                    error: "No : in 'fortran_order' field".into() 
159                })?;
160            // "False"
161            let fortran_order = fortran_order[index+1..].trim();
162            match fortran_order {
163                "False" => fortran_order_opt = Some(false),
164                "True" => fortran_order_opt = Some(true),
165                _ => Err(NpyError::Header { 
166                    field: "fortran_order", 
167                    error: format!("Unsupported value '{}'", fortran_order) 
168                })?,
169            };
170
171            header_str = header_str[colon_index+1..].trim_start();
172        } else if header_str.starts_with("'shape'") {
173            // "'shape': (3, 4), "
174            let left_brace_index = header_str.find("(")
175                .ok_or_else(|| NpyError::Header { 
176                    field: "shape", 
177                    error: "No ( in 'shape' field".into() 
178                })?;
179            let right_brace_index = header_str.find(")")
180                .ok_or_else(|| NpyError::Header { 
181                    field: "shape", 
182                    error: "No ) in 'shape' field".into() 
183                })?;
184            // 3, 4
185            let shape = &header_str[left_brace_index + 1..right_brace_index];
186            let shape: Vec<usize> = shape
187                .split(',')
188                .filter_map(|s| s.trim().parse::<usize>().ok())
189                .collect();
190            shape_opt = Some(shape);
191
192            header_str = header_str[right_brace_index+2..].trim_start();                
193        }
194    }
195
196    // Check header
197    let descr = descr_opt.ok_or_else(|| NpyError::LackField { field: "descr" } )?;
198    let _ = fortran_order_opt.ok_or_else(|| NpyError::LackField { field: "fortran_order" } )?;
199    let shape = shape_opt.ok_or_else(|| NpyError::LackField { field: "shape" } )?;
200    let shape: Shape = shape.into();
201
202    match descr {
203        "|b1" => Ok(utils::load_tensor_reader(DType::Bool, shape, &mut reader).map_err(NpyError::Core)?),
204        "<u1" => Ok(utils::load_tensor_reader(DType::U8, shape, &mut reader).map_err(NpyError::Core)?),
205        "<i4" => Ok(utils::load_tensor_reader(DType::I32, shape, &mut reader).map_err(NpyError::Core)?),
206        "<u4" => Ok(utils::load_tensor_reader(DType::U32, shape, &mut reader).map_err(NpyError::Core)?),
207        "<f4" => Ok(utils::load_tensor_reader(DType::F32, shape, &mut reader).map_err(NpyError::Core)?),
208        "<f8" => Ok(utils::load_tensor_reader(DType::F64, shape, &mut reader).map_err(NpyError::Core)?),
209        _ => Err(NpyError::UnsupportedDescr(descr.to_string()))?
210    }
211}
212
213/// Load multiple arrays from a `.npz` archive into a `HashMap`.
214///
215/// Each entry in the returned map corresponds to one `.npy` file inside the archive,
216/// where the key is the file stem (without extension) and the value is the array.
217///
218/// # Arguments
219/// * `path` - Path to the `.npz` archive.
220///
221/// # Errors
222/// Returns an error if the archive cannot be opened, read, or parsed.
223pub fn load_npz_file<P: AsRef<Path>>(path: P) -> NpyResult<HashMap<String, DynTensor>> {
224    let file = std::fs::File::open(path)?;
225    let mut archive = ZipArchive::new(file)?;
226    let mut arrays = HashMap::new();
227
228    for i in 0..archive.len() {
229        let mut zip_file = archive.by_index(i)?;
230        let mut buffer = Vec::new();
231        use std::io::Read;
232        zip_file.read_to_end(&mut buffer)?;
233        
234        let mut cursor = Cursor::new(buffer);
235        let array = load_npy(&mut cursor)?;
236
237        let name = Path::new(zip_file.name())
238            .file_stem()
239            .and_then(|s| s.to_str())
240            .unwrap_or(zip_file.name())
241            .to_string();
242        arrays.insert(name, array);
243    }
244
245    Ok(arrays)
246}
247
248/// Save a [`DynTensor`] into a `.npy` file.
249///
250/// # Arguments
251/// * `tensor` - The array to save.
252/// * `path` - Path to the output `.npy` file.
253///
254/// # Errors
255/// Returns an error if the file cannot be created or written.
256pub fn save_npy_file<P: AsRef<Path>>(tensor: impl Into<DynTensor>, path: P) -> NpyResult<()> {
257    let file = File::create(path).map_err(NpyError::Io).context("create file")?;
258    let mut writer = BufWriter::new(file);
259    save_npy(tensor, &mut writer)
260}
261
262pub fn save_npy<W: Write>(tensor: impl Into<DynTensor>, writer: &mut W) -> NpyResult<()> {
263    let tensor = tensor.into();
264    match tensor {
265        DynTensor::Bool(t) => _save_npy(t, writer),
266        DynTensor::U8(t) => _save_npy(t, writer),
267        DynTensor::U32(t) => _save_npy(t, writer),
268        DynTensor::I32(t) => _save_npy(t, writer),
269        DynTensor::F32(t) => _save_npy(t, writer),
270        DynTensor::F64(t) => _save_npy(t, writer),
271        DynTensor::Bf16(_) =>
272            Err(NpyError::UnsupportedDType(DType::Bf16))?
273    }
274}
275
276pub fn save_npz_file<P: AsRef<Path>>(tensors: &HashMap<String, DynTensor>, path: P) -> NpyResult<()> {
277    let file = File::create(path)?;
278    let mut zip = ZipWriter::new(file);
279
280    let options = SimpleFileOptions::default()
281        .compression_method(zip::CompressionMethod::Stored)
282        .unix_permissions(0o755);
283    
284    for (name, tensor) in tensors {
285        zip.start_file(format!("{}.npy", name), options)?;
286        save_npy(tensor.clone(), &mut zip)?
287    }
288
289    zip.finish()?;
290    Ok(())
291}
292
293fn _save_npy<D: WithDType, W: Write>(tensor: Tensor<D>, writer: &mut W) -> NpyResult<()> {
294    let descr = dtype_to_descr(D::DTYPE)?;
295    let version = (1, 0);
296    let fortran_order = false;
297    let shape = tensor.dims();
298
299    let storage = tensor.storage_read()?;
300    let data_vec = storage.data();
301    let bytes: Vec<u8> = bytemuck::cast_slice(data_vec).to_vec();
302
303    writer.write_all(NPY_MAGIC)?;
304
305    writer.write_all(&[version.0, version.1])?;
306
307    let dict = format!(
308        "{{'descr': '{}', 'fortran_order': {}, 'shape': ({}), }}",
309        descr,
310        if fortran_order { "True" } else { "False" },
311        shape.iter().map(|s| s.to_string()).collect::<Vec<_>>().join(", ")
312    );
313
314    // header padding 
315    let header_len = dict.len() + 1; // \n
316    if header_len > u16::MAX as usize {
317        panic!("Header too long for version 1.0 NPY");
318    }
319    let padding = (64 - ((10 + 2 + header_len) % 64)) % 64;
320    let mut header_bytes = dict.into_bytes();
321    header_bytes.push(b'\n');
322    header_bytes.extend(vec![b' '; padding]);
323
324    // header_len
325    let header_len_u16: u16 = header_bytes.len().try_into().unwrap();
326    writer.write_all(&header_len_u16.to_le_bytes())?;
327
328    // header
329    writer.write_all(&header_bytes)?;
330
331    // data
332    writer.write_all(&bytes)?;
333
334    writer.flush()?;
335    Ok(())
336}
337
338fn dtype_to_descr(dtype: DType) -> NpyResult<&'static str> {
339    match dtype {
340        DType::Bool => Ok("|b1"),
341        DType::U8 => Ok("<u1"),
342        DType::U32 => Ok("<u4"),
343        DType::I32 => Ok("<i4"),
344        DType::F32 => Ok("<f4"),
345        DType::F64 => Ok("<f8"),
346        DType::Bf16 => Err(NpyError::UnsupportedDType(DType::Bf16))?,
347    }
348}
349
350#[cfg(test)]
351mod test {
352    use std::collections::HashMap;
353    use lumen_core::{DynTensor, Tensor};
354    use tempfile::NamedTempFile;
355
356    use crate::npy::{load_npy_file, save_npy_file};
357
358    use super::{load_npz_file, save_npz_file};
359
360    #[test]
361    fn test_to_tensor() {
362        let tensor = load_npy_file("./bench/test1.npy").unwrap().as_f32().unwrap();
363        println!("{}", tensor);
364
365        let tensor = load_npy_file("./bench/test3.npy").unwrap().as_bool().unwrap();
366        println!("{}", tensor);
367
368        let tensor = load_npy_file("./bench/test4.npy").unwrap().as_i32().unwrap();
369        println!("{}", tensor);
370
371        let _tensor = load_npy_file("./bench/test5.npy").unwrap().as_f64().unwrap();
372    }
373
374    #[test]
375    fn test_write_npy() {
376        let tmpfile = NamedTempFile::new().unwrap();
377        let tensor = Tensor::<f32>::randn(0., 1., (4, 5)).unwrap();
378        save_npy_file(&tensor, tmpfile.path()).unwrap();
379
380        let loaded_tensor = load_npy_file(tmpfile.path()).unwrap().as_f32().unwrap();
381        assert!(loaded_tensor.allclose(&tensor, 1e-6, 1e-6).unwrap());
382    }
383
384    #[test]
385    fn test_load_npz() {
386        let tensors = load_npz_file("./bench/test1.npz").unwrap();
387        for (name, t) in tensors {
388            println!("{}: {}", name, t.shape());
389        }
390    }
391
392    #[test]
393    fn test_save_npz() {
394        let tmpfile = NamedTempFile::new().unwrap();
395
396        let scalar = Tensor::new(1).unwrap();
397        let vector_f32 = Tensor::new(&[1.0f32, 2., 3.]).unwrap();
398        let matrix_f32 = Tensor::new(&[[1, 2, 3], [3, 4, 5]]).unwrap();
399        let ones_f32 = Tensor::<f32>::ones((2, 9)).unwrap();
400        let randn_f64 = Tensor::randn(0.0f64, 1., (1, 2, 3)).unwrap();
401        let fill_f64 = Tensor::full((2, 3, 4), 1.2).unwrap();
402        let arange_f64 = Tensor::arange(0., 10.).unwrap();
403        let trues = Tensor::trues((3, 4)).unwrap();
404        let booleans = Tensor::new(&[[true, false], [false, true]]).unwrap();
405
406        let mut tensors = HashMap::new();
407        tensors.insert("scalar".to_string(), DynTensor::I32(scalar));
408        tensors.insert("vector_f32".to_string(), DynTensor::F32(vector_f32));
409        tensors.insert("matrix_f32".to_string(), DynTensor::I32(matrix_f32));
410        tensors.insert("ones_f32".to_string(), DynTensor::F32(ones_f32));
411        tensors.insert("randn_f64".to_string(), DynTensor::F64(randn_f64));
412        tensors.insert("fill_f64".to_string(), DynTensor::F64(fill_f64));
413        tensors.insert("arange_f64".to_string(), DynTensor::F64(arange_f64));
414        tensors.insert("trues".to_string(), DynTensor::Bool(trues));
415        tensors.insert("booleans".to_string(), DynTensor::Bool(booleans));
416
417        save_npz_file(&tensors, tmpfile.path()).unwrap();
418
419        let tensors = load_npz_file(tmpfile.path()).unwrap();
420        for (name, t) in tensors {
421            println!("{}: {}", name, t.shape());
422            match name.as_str() {
423                "scalar" => assert!(t.as_i32().unwrap().is_scalar()),
424                "vector_f32" => assert_eq!(t.as_f32().unwrap().rank(), 1),
425                "matrix_f32" => assert_eq!(t.as_i32().unwrap().dims(), [2, 3]),
426                _ => {}
427            }
428        }
429    }
430}
lumen_io/npy.rs

lumen_io/
npy.rs
