metatensor 0.2.4

Self-describing sparse tensor data format for atomistic machine learning and beyond
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

use std::ffi::CString;

use crate::errors::{check_ptr, check_status};
use crate::{Error, TensorBlock, TensorBlockRef};

use super::{realloc_vec, create_ndarray};

/// Load previously saved `TensorBlock` from the file at the given path.
pub fn load_block(path: impl AsRef<std::path::Path>) -> Result<TensorBlock, Error> {
    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
    let path = CString::new(path).expect("this path contains a NULL byte");

    let ptr = unsafe {
        crate::c_api::mts_block_load(
            path.as_ptr(),
            Some(create_ndarray)
        )
    };

    check_ptr(ptr)?;

    return Ok(unsafe { TensorBlock::from_raw(ptr) });
}

/// Load a serialized `TensorBlock` from a `buffer`.
pub fn load_block_buffer(buffer: &[u8]) -> Result<TensorBlock, Error> {
    let ptr = unsafe {
        crate::c_api::mts_block_load_buffer(
            buffer.as_ptr(),
            buffer.len(),
            Some(create_ndarray)
        )
    };

    check_ptr(ptr)?;

    return Ok(unsafe { TensorBlock::from_raw(ptr) });
}

/// Save the given `block` to a file.
///
/// If the file already exists, it is overwritten. The recomended file extension
/// when saving data is `.mts`, to prevent confusion with generic `.npz`.
pub fn save_block(path: impl AsRef<std::path::Path>, block: TensorBlockRef) -> Result<(), Error> {
    let path = path.as_ref().as_os_str().to_str().expect("this path is not valid UTF8");
    let path = CString::new(path).expect("this path contains a NULL byte");

    unsafe {
        check_status(crate::c_api::mts_block_save(path.as_ptr(), block.as_ptr()))
    }
}


/// Save the given `block` to an in-memory `buffer`.
///
/// This function will grow the buffer as required to fit the data.
pub fn save_block_buffer(block: TensorBlockRef, buffer: &mut Vec<u8>) -> Result<(), Error> {
    let mut buffer_ptr = buffer.as_mut_ptr();
    let mut buffer_count = buffer.len();

    unsafe {
        check_status(crate::c_api::mts_block_save_buffer(
            &mut buffer_ptr,
            &mut buffer_count,
            (buffer as *mut Vec<u8>).cast(),
            Some(realloc_vec),
            block.as_ptr(),
        ))?;
    }

    buffer.resize(buffer_count, 0);

    Ok(())
}