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

use std::ffi::CString;

use crate::c_api::mts_labels_t;
use crate::errors::check_status;
use crate::{Labels, Error};

use super::realloc_vec;

/// Load previously saved `Labels` from the file at the given path.
pub fn load_labels(path: impl AsRef<std::path::Path>) -> Result<Labels, 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 mut labels = mts_labels_t::null();

    unsafe {
        check_status(crate::c_api::mts_labels_load(path.as_ptr(), &mut labels))?;
    }

    return Ok(unsafe { Labels::from_raw(labels) });
}

/// Load previously saved `Labels` from an in-memory `buffer`.
pub fn load_labels_buffer(buffer: &[u8]) -> Result<Labels, Error> {
    let mut labels = mts_labels_t::null();

    unsafe {
        check_status(crate::c_api::mts_labels_load_buffer(
            buffer.as_ptr(),
            buffer.len(),
            &mut labels
        ))?;
    }

    return Ok(unsafe { Labels::from_raw(labels) });
}

/// Save the given `Labels` 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_labels(path: impl AsRef<std::path::Path>, labels: &Labels) -> 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_labels_save(path.as_ptr(), labels.raw))
    }
}


/// Save the given `labels` to an in-memory `buffer`.
///
/// This function will grow the buffer as required to fit the labels.
pub fn save_labels_buffer(labels: &Labels, 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_labels_save_buffer(
            &mut buffer_ptr,
            &mut buffer_count,
            (buffer as *mut Vec<u8>).cast(),
            Some(realloc_vec),
            labels.raw,
        ))?;
    }

    buffer.resize(buffer_count, 0);

    Ok(())
}