Metatomic-core - io.hpp - serialization functions

docs/src/core/reference/cxx/system.rst

@@ -6,3 +6,20 @@ System
 
 .. doxygenclass:: metatomic::PairListOptions
    :members:
+
+Serialization
+-------------
+
+Systems can be saved to a file or serialized into an in-memory byte buffer.
+Loading a system requires an array-creation callback, which allocates the
+arrays of the reconstructed system.
+
+.. doxygenfunction:: metatomic::io::save
+
+.. doxygenfunction:: metatomic::io::save_buffer
+
+.. doxygenfunction:: metatomic::io::load
+
+.. doxygenfunction:: metatomic::io::load_buffer(const uint8_t* buffer, uintptr_t buffer_count, mts_create_array_callback_t create_array)
+
+.. doxygenfunction:: metatomic::io::load_buffer(const Buffer& buffer, mts_create_array_callback_t create_array)

metatomic-core/include/metatomic.hpp

@@ -4,3 +4,4 @@
 #include "metatomic/plugin.hpp"   // IWYU pragma: export
 #include "metatomic/errors.hpp"   // IWYU pragma: export
 #include "metatomic/metadata.hpp" // IWYU pragma: export
+#include "metatomic/io.hpp"       // IWYU pragma: export

metatomic-core/include/metatomic/io.hpp

@@ -0,0 +1,116 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <metatomic.h>
+
+#include <metatomic/errors.hpp>
+#include <metatomic/system.hpp>
+
+namespace metatomic {
+namespace io {
+
+/// Save a system to a file.
+///
+/// @param path path of the file to create or overwrite
+/// @param system system to serialize
+inline void save(const std::string& path, const System& system) {
+    details::check_status(mta_save(path.c_str(), system.as_mta_system_t()));
+}
+
+/// Serialize a system into a byte container.
+///
+/// `Buffer` must be constructible from a pair of iterators over bytes. The
+/// serialization is performed using a `std::vector<uint8_t>` and copied into
+/// the requested container type.
+///
+/// @tparam Buffer byte-container type, such as `std::vector<uint8_t>`
+/// @param system system to serialize
+/// @return serialized system data
+template <typename Buffer>
+Buffer save_buffer(const System& system) {
+    auto buffer = metatomic::io::save_buffer<std::vector<uint8_t>>(system);
+    return Buffer(buffer.begin(), buffer.end());
+}
+
+/// Serialize a system into a `std::vector<uint8_t>`.
+///
+/// The C API grows the vector through a reallocation callback. The returned
+/// vector contains exactly the number of bytes produced by the serializer.
+///
+/// @param system system to serialize
+/// @return serialized system data
+template <>
+inline std::vector<uint8_t> save_buffer<std::vector<uint8_t>>(const System& system) {
+    std::vector<uint8_t> buffer;
+
+    auto* ptr = buffer.data();
+    auto size = buffer.size();
+
+    auto realloc = [](void* user_data, uint8_t*, uintptr_t new_size) {
+        auto* buffer = reinterpret_cast<std::vector<uint8_t>*>(user_data);
+        buffer->resize(new_size, '\0');
+        return buffer->data();
+    };
+
+    details::check_status(mta_save_buffer(&ptr, &size, &buffer, realloc, system.as_mta_system_t()));
+
+    buffer.resize(size, '\0');
+
+    return buffer;
+}
+
+/// Load a system from a file.
+///
+/// @param path path of the serialized system file
+/// @param create_array callback used to create arrays during deserialization
+/// @return reconstructed system
+inline System load(
+    const std::string& path,
+    mts_create_array_callback_t create_array = metatensor::details::default_create_array
+) {
+    mta_system_t* ptr = nullptr;
+    details::check_status(mta_load(path.c_str(), create_array, &ptr));
+    details::check_pointer(ptr);
+    return System::unsafe_from_ptr(ptr);
+}
+
+/// Load a system from a contiguous byte buffer.
+///
+/// @param buffer serialized system data
+/// @param buffer_count number of bytes available at `buffer`
+/// @param create_array callback used to create arrays during deserialization
+/// @return reconstructed system
+inline System load_buffer(
+    const uint8_t* buffer,
+    uintptr_t buffer_count,
+    mts_create_array_callback_t create_array = metatensor::details::default_create_array
+) {
+    mta_system_t* ptr = nullptr;
+    details::check_status(mta_load_buffer(buffer, buffer_count, create_array, &ptr));
+    details::check_pointer(ptr);
+    return System::unsafe_from_ptr(ptr);
+}
+
+/// Load a system from a byte container.
+///
+/// The container must provide contiguous storage through `data()` and report
+/// its size in bytes through `size()`.
+///
+/// @tparam Buffer contiguous byte-container type
+/// @param buffer serialized system data
+/// @param create_array callback used to create arrays during deserialization
+/// @return reconstructed system
+template <typename Buffer>
+System load_buffer(
+    const Buffer& buffer,
+    mts_create_array_callback_t create_array = metatensor::details::default_create_array
+) {
+    static_assert(sizeof(typename Buffer::value_type) == sizeof(uint8_t), "`Buffer` must be a container of uint8_t or equivalent");
+    return metatomic::io::load_buffer(reinterpret_cast<const uint8_t*>(buffer.data()), buffer.size(), create_array);
+}
+
+} // namespace io
+} // namespace metatomic

metatomic-core/tests/cxx/system.cpp

@@ -1,4 +1,6 @@
 #include <cstdint>
+#include <cstdio>
+#include <filesystem>
 #include <memory>
 #include <string>
 #include <utility>
@@ -123,6 +125,70 @@ static metatensor::TensorMap custom_data() {
     return metatensor::TensorMap(keys, std::move(blocks));
 }
 
+// Helper function to check that two DLPack tensors have the same shape, strides, dtype, and data.
+template <typename T>
+static void check_tensors(const DLManagedTensorVersioned* expected, const DLManagedTensorVersioned* loaded) {
+    REQUIRE((expected != nullptr && loaded != nullptr));
+    CHECK(loaded->dl_tensor.ndim == expected->dl_tensor.ndim);
+    CHECK(loaded->dl_tensor.dtype.code == expected->dl_tensor.dtype.code);
+    CHECK(loaded->dl_tensor.dtype.bits == expected->dl_tensor.dtype.bits);
+    CHECK(loaded->dl_tensor.dtype.lanes == expected->dl_tensor.dtype.lanes);
+
+    for (int64_t i = 0; i < expected->dl_tensor.ndim; i++) {
+        CHECK(loaded->dl_tensor.shape[i] == expected->dl_tensor.shape[i]);
+        CHECK(loaded->dl_tensor.strides[i] == expected->dl_tensor.strides[i]);
+    }
+
+    CHECK(metatensor::details::vector_from_dlpack<T>(expected->dl_tensor) == metatensor::details::vector_from_dlpack<T>(loaded->dl_tensor));
+}
+
+// Helper function to check that two TensorBlocks have the same metadata and values.
+template <typename T>
+static void check_tensor_blocks(metatensor::TensorBlock&& expected, metatensor::TensorBlock&& loaded) {
+    CHECK(loaded.samples() == expected.samples());
+    CHECK(loaded.components() == expected.components());
+    CHECK(loaded.properties() == expected.properties());
+    CHECK(loaded.values_shape() == expected.values_shape());
+    CHECK(loaded.values<T>() == expected.values<T>());
+
+    REQUIRE(loaded.gradients_list() == expected.gradients_list());
+    for (const auto& parameter : expected.gradients_list()) {
+        check_tensor_blocks<T>(expected.gradient(parameter), loaded.gradient(parameter));
+    }
+}
+
+// Helper function to check that two TensorMaps have the same metadata and blocks.
+template <typename T>
+static void check_tensor_maps(metatensor::TensorMap&& expected, metatensor::TensorMap&& loaded) {
+    REQUIRE(loaded.keys() == expected.keys());
+    for (uintptr_t i = 0; i < expected.keys().count(); i++) {
+        check_tensor_blocks<T>(expected.block_by_id(i), loaded.block_by_id(i));
+    }
+}
+
+// Helper function to check that two Systems are equal
+static void check_systems(const metatomic::System& system, const metatomic::System& loaded) {
+    CHECK(loaded.size() == system.size());
+    CHECK(loaded.length_unit() == system.length_unit());
+
+    check_tensors<int32_t>(system.types().as_dlpack(), loaded.types().as_dlpack());
+    check_tensors<float>(system.positions().as_dlpack(), loaded.positions().as_dlpack());
+    check_tensors<float>(system.cell().as_dlpack(), loaded.cell().as_dlpack());
+    check_tensors<bool>(system.pbc().as_dlpack(), loaded.pbc().as_dlpack());
+
+    REQUIRE(loaded.known_pairs().size() == system.known_pairs().size());
+    for (size_t i = 0; i < system.known_pairs().size(); i++) {
+        CHECK(loaded.known_pairs()[i] == system.known_pairs()[i]);
+        CHECK(loaded.known_pairs()[i].requestors() == system.known_pairs()[i].requestors());
+        check_tensor_blocks<float>(system.pairs(system.known_pairs()[i]), loaded.pairs(loaded.known_pairs()[i]));
+    }
+
+    REQUIRE(loaded.known_custom_data() == system.known_custom_data());
+    for (const auto& name : system.known_custom_data()) {
+        check_tensor_maps<float>(system.custom_data(name), loaded.custom_data(name));
+    }
+}
+
 
 TEST_CASE("System basics") {
     auto system = test_system(4);
@@ -291,3 +357,47 @@ TEST_CASE("System ownership") {
         CHECK(system.size() == 4);
     }
 }
+
+TEST_CASE("System serialization") {
+    SECTION("save and load to a file") {
+        auto system = test_system(4);
+        const std::string path = "metatomic-test-system.mta";
+
+        metatomic::io::save(path, system);
+        auto loaded = metatomic::io::load(path);
+
+        check_systems(system, loaded);
+
+        std::remove(path.c_str());
+    }
+
+    SECTION("load a legacy file") {
+        auto path = std::filesystem::path(__FILE__).parent_path().parent_path() / "data" / "legacy.mta";
+        auto system = metatomic::io::load(path.string());
+
+        CHECK(system.as_mta_system_t() != nullptr);
+        CHECK(system.size() == 4);
+        CHECK(system.length_unit().empty());
+
+        auto types = metatensor::details::vector_from_dlpack<int32_t>(system.types()->dl_tensor);
+        CHECK((types == std::vector<int32_t>{1, 6, 7, 8}));
+
+        auto positions = metatensor::details::vector_from_dlpack<double>(system.positions()->dl_tensor);
+        CHECK((positions == std::vector<double>{
+            0.0, 0.0, 0.0,
+            1.0, 2.0, 3.0,
+            4.0, 5.0, 6.0,
+            7.0, 8.0, 9.0,
+        }));
+    }
+
+    SECTION("save and load to an in-memory buffer") {
+        auto system = test_system(4);
+
+        auto buffer = metatomic::io::save_buffer<std::vector<uint8_t>>(system);
+        REQUIRE_FALSE(buffer.empty());
+
+        auto loaded = metatomic::io::load_buffer(buffer);
+        check_systems(system, loaded);
+    }
+}