From 27fbd8982dc014e058f36dc898b1139036f29efe Mon Sep 17 00:00:00 2001 From: lucaskloss Date: Tue, 28 Jul 2026 11:15:55 +0000 Subject: [PATCH] Serialization for C++ system --- docs/src/core/reference/cxx/system.rst | 17 ++++ metatomic-core/include/metatomic.hpp | 1 + metatomic-core/include/metatomic/io.hpp | 116 ++++++++++++++++++++++++ metatomic-core/tests/cxx/system.cpp | 110 ++++++++++++++++++++++ 4 files changed, 244 insertions(+) create mode 100644 metatomic-core/include/metatomic/io.hpp diff --git a/docs/src/core/reference/cxx/system.rst b/docs/src/core/reference/cxx/system.rst index c6471bfce..54cad2391 100644 --- a/docs/src/core/reference/cxx/system.rst +++ b/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) diff --git a/metatomic-core/include/metatomic.hpp b/metatomic-core/include/metatomic.hpp index a714290e3..e4fa4c930 100644 --- a/metatomic-core/include/metatomic.hpp +++ b/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 diff --git a/metatomic-core/include/metatomic/io.hpp b/metatomic-core/include/metatomic/io.hpp new file mode 100644 index 000000000..a314a8190 --- /dev/null +++ b/metatomic-core/include/metatomic/io.hpp @@ -0,0 +1,116 @@ +#pragma once + +#include +#include +#include + +#include + +#include +#include + +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` and copied into +/// the requested container type. +/// +/// @tparam Buffer byte-container type, such as `std::vector` +/// @param system system to serialize +/// @return serialized system data +template +Buffer save_buffer(const System& system) { + auto buffer = metatomic::io::save_buffer>(system); + return Buffer(buffer.begin(), buffer.end()); +} + +/// Serialize a system into a `std::vector`. +/// +/// 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 save_buffer>(const System& system) { + std::vector buffer; + + auto* ptr = buffer.data(); + auto size = buffer.size(); + + auto realloc = [](void* user_data, uint8_t*, uintptr_t new_size) { + auto* buffer = reinterpret_cast*>(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 +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(buffer.data()), buffer.size(), create_array); +} + +} // namespace io +} // namespace metatomic diff --git a/metatomic-core/tests/cxx/system.cpp b/metatomic-core/tests/cxx/system.cpp index 79532d16a..4ebf5c405 100644 --- a/metatomic-core/tests/cxx/system.cpp +++ b/metatomic-core/tests/cxx/system.cpp @@ -1,4 +1,6 @@ #include +#include +#include #include #include #include @@ -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 +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(expected->dl_tensor) == metatensor::details::vector_from_dlpack(loaded->dl_tensor)); +} + +// Helper function to check that two TensorBlocks have the same metadata and values. +template +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() == expected.values()); + + REQUIRE(loaded.gradients_list() == expected.gradients_list()); + for (const auto& parameter : expected.gradients_list()) { + check_tensor_blocks(expected.gradient(parameter), loaded.gradient(parameter)); + } +} + +// Helper function to check that two TensorMaps have the same metadata and blocks. +template +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(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(system.types().as_dlpack(), loaded.types().as_dlpack()); + check_tensors(system.positions().as_dlpack(), loaded.positions().as_dlpack()); + check_tensors(system.cell().as_dlpack(), loaded.cell().as_dlpack()); + check_tensors(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(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(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(system.types()->dl_tensor); + CHECK((types == std::vector{1, 6, 7, 8})); + + auto positions = metatensor::details::vector_from_dlpack(system.positions()->dl_tensor); + CHECK((positions == std::vector{ + 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>(system); + REQUIRE_FALSE(buffer.empty()); + + auto loaded = metatomic::io::load_buffer(buffer); + check_systems(system, loaded); + } +}