metatensor/io/
block.rs

1use std::ffi::CString;
2
3use metatensor_sys::mts_create_array_callback_t;
4
5use crate::errors::{check_ptr, check_status};
6use crate::{Error, TensorBlock, TensorBlockRef};
7
8use super::{realloc_vec, create_ndarray};
9
10/// Load previously saved `TensorBlock` from the file at the given path.
11///
12/// All arrays are loaded as `ndarray::ArrayD` by default. If you want to load
13/// arrays as a custom type, use [`load_block_custom_array`] instead.
14pub fn load_block(path: impl AsRef<std::path::Path>) -> Result<TensorBlock, Error> {
15    return load_block_custom_array(path, Some(create_ndarray));
16}
17
18/// Load previously saved `TensorBlock` from the file at the given path.
19///
20/// All arrays will be created through the `create_array` callback.
21pub fn load_block_custom_array(path: impl AsRef<std::path::Path>, create_array: mts_create_array_callback_t) -> Result<TensorBlock, Error> {
22    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
23    let path = CString::new(path).expect("this path contains a NULL byte");
24
25    let ptr = unsafe {
26        crate::c_api::mts_block_load(
27            path.as_ptr(),
28            create_array
29        )
30    };
31
32    check_ptr(ptr)?;
33
34    return Ok(unsafe { TensorBlock::from_raw(ptr) });
35}
36
37/// Load a serialized `TensorBlock` from a `buffer`.
38///
39/// All arrays are loaded as `ndarray::ArrayD` by default. If you want to load
40/// arrays as a custom type, use [`load_block_buffer_custom_array`] instead.
41pub fn load_block_buffer(buffer: &[u8]) -> Result<TensorBlock, Error> {
42    return load_block_buffer_custom_array(buffer, Some(create_ndarray));
43}
44
45/// Load a serialized `TensorBlock` from a `buffer`.
46///
47/// All arrays will be created through the `create_array` callback.
48pub fn load_block_buffer_custom_array(buffer: &[u8], create_array: mts_create_array_callback_t) -> Result<TensorBlock, Error> {
49    let ptr = unsafe {
50        crate::c_api::mts_block_load_buffer(
51            buffer.as_ptr(),
52            buffer.len(),
53            create_array
54        )
55    };
56
57    check_ptr(ptr)?;
58
59    return Ok(unsafe { TensorBlock::from_raw(ptr) });
60}
61
62/// Save the given `block` to a file.
63///
64/// If the file already exists, it is overwritten. The recommended file extension
65/// when saving data is `.mts`, to prevent confusion with generic `.npz`.
66pub fn save_block(path: impl AsRef<std::path::Path>, block: TensorBlockRef) -> Result<(), Error> {
67    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
68    let path = CString::new(path).expect("this path contains a NULL byte");
69
70    unsafe {
71        check_status(crate::c_api::mts_block_save(path.as_ptr(), block.as_ptr()))
72    }
73}
74
75
76/// Save the given `block` to an in-memory `buffer`.
77///
78/// This function will grow the buffer as required to fit the data.
79pub fn save_block_buffer(block: TensorBlockRef, buffer: &mut Vec<u8>) -> Result<(), Error> {
80    let mut buffer_ptr = buffer.as_mut_ptr();
81    let mut buffer_count = buffer.len();
82
83    unsafe {
84        check_status(crate::c_api::mts_block_save_buffer(
85            &mut buffer_ptr,
86            &mut buffer_count,
87            std::ptr::from_mut(buffer).cast(),
88            Some(realloc_vec),
89            block.as_ptr(),
90        ))?;
91    }
92
93    buffer.resize(buffer_count, 0);
94
95    Ok(())
96}
metatensor/io/block.rs

metatensor/io/
block.rs
