metatensor/io/
tensor.rs

1use std::ffi::CString;
2
3use metatensor_sys::mts_create_array_callback_t;
4
5use crate::errors::{check_status, check_ptr};
6use crate::{TensorMap, Error};
7
8use super::{realloc_vec, create_ndarray};
9
10/// Load the serialized tensor map from the given path.
11///
12/// `TensorMap` are serialized using numpy's NPZ format, i.e. a ZIP file
13/// without compression (storage method is STORED), where each file is stored as
14/// a `.npy` array. Both the ZIP and NPY format are well documented:
15///
16/// - ZIP: <https://pkware.cachefly.net/webdocs/casestudies/APPNOTE.TXT>
17/// - NPY: <https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html>
18///
19/// We add other restriction on top of these formats when saving/loading data.
20/// First, `Labels` instances are saved as structured array, see the `labels`
21/// module for more information. Only 32-bit integers are supported for Labels,
22/// and only 64-bit floats are supported for data (values and gradients).
23///
24/// Second, the path of the files in the archive also carry meaning. The keys of
25/// the `TensorMap` are stored in `/keys.npy`, and then different blocks are
26/// stored as
27///
28/// ```bash
29/// /  blocks / <block_id>  / values / samples.npy
30///                         / values / components  / 0.npy
31///                                                / <...>.npy
32///                                                / <n_components>.npy
33///                         / values / properties.npy
34///                         / values / data.npy
35///
36///                         # optional sections for gradients, one by parameter
37///                         /   gradients / <parameter> / samples.npy
38///                                                     /   components  / 0.npy
39///                                                                     / <...>.npy
40///                                                                     / <n_components>.npy
41///                                                     /   data.npy
42/// ```
43///
44/// All arrays are loaded as `ndarray::ArrayD` by default. If you want to load
45/// arrays as a custom type, use [`load_custom_array`] instead.
46pub fn load(path: impl AsRef<std::path::Path>) -> Result<TensorMap, Error> {
47    return load_custom_array(path, Some(create_ndarray));
48}
49
50/// Load the serialized tensor map from the given path.
51///
52/// See the [`load`] function for more information on the data format.
53///
54/// All arrays will be created through the `create_array` callback.
55pub fn load_custom_array(path: impl AsRef<std::path::Path>, create_array: mts_create_array_callback_t) -> Result<TensorMap, Error> {
56    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
57    let path = CString::new(path).expect("this path contains a NULL byte");
58
59    let ptr = unsafe {
60        crate::c_api::mts_tensormap_load(path.as_ptr(), create_array)
61    };
62
63    check_ptr(ptr)?;
64
65    return Ok(unsafe { TensorMap::from_raw(ptr) });
66}
67
68/// Load a serialized `TensorMap` from a `buffer`.
69///
70/// See the [`load`] function for more information on the data format.
71///
72/// All arrays are loaded as `ndarray::ArrayD` by default. If you want to load
73/// arrays as a custom type, use [`load_buffer_custom_array`] instead.
74pub fn load_buffer(buffer: &[u8]) -> Result<TensorMap, Error> {
75    return load_buffer_custom_array(buffer, Some(create_ndarray));
76}
77
78/// Load a serialized `TensorMap` from a `buffer` using a custom array
79/// creation function.
80///
81/// All arrays will be created through the `create_array` callback.
82pub fn load_buffer_custom_array(buffer: &[u8], create_array: mts_create_array_callback_t) -> Result<TensorMap, Error> {
83    let ptr = unsafe {
84        crate::c_api::mts_tensormap_load_buffer(
85            buffer.as_ptr(),
86            buffer.len(),
87            create_array
88        )
89    };
90
91    check_ptr(ptr)?;
92
93    return Ok(unsafe { TensorMap::from_raw(ptr) });
94}
95
96/// Save the given tensor to a file.
97///
98/// If the file already exists, it is overwritten. The recommended file extension
99/// when saving data is `.mts`, to prevent confusion with generic `.npz`.
100///
101/// The format used is documented in the [`load`] function, and consists of a
102/// zip archive containing NPY files.
103pub fn save(path: impl AsRef<std::path::Path>, tensor: &TensorMap) -> Result<(), Error> {
104    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
105    let path = CString::new(path).expect("this path contains a NULL byte");
106
107    unsafe {
108        check_status(crate::c_api::mts_tensormap_save(path.as_ptr(), tensor.ptr))
109    }
110}
111
112
113/// Save the given `tensor` to an in-memory `buffer`.
114///
115/// This function will grow the buffer as required to fit the whole tensor.
116pub fn save_buffer(tensor: &TensorMap, buffer: &mut Vec<u8>) -> Result<(), Error> {
117    let mut buffer_ptr = buffer.as_mut_ptr();
118    let mut buffer_count = buffer.len();
119
120    unsafe {
121        check_status(crate::c_api::mts_tensormap_save_buffer(
122            &mut buffer_ptr,
123            &mut buffer_count,
124            std::ptr::from_mut(buffer).cast(),
125            Some(realloc_vec),
126            tensor.ptr,
127        ))?;
128    }
129
130    buffer.resize(buffer_count, 0);
131
132    Ok(())
133}
metatensor/io/tensor.rs

metatensor/io/
tensor.rs
