Add save_npy function

CMakeLists.txt

@@ -134,6 +134,7 @@ target_sources(
         src/ddc/discrete_element.cpp
         src/ddc/discrete_space.cpp
         src/ddc/discrete_vector.cpp
+        src/ddc/save_npy.cpp
         src/ddc/print.cpp
     PUBLIC
         FILE_SET core_public
@@ -176,6 +177,7 @@ target_sources(
                 src/ddc/print.hpp
                 src/ddc/real_type.hpp
                 src/ddc/reducer.hpp
+                src/ddc/save_npy.hpp
                 src/ddc/scope_guard.hpp
                 src/ddc/sparse_discrete_domain.hpp
                 src/ddc/strided_discrete_domain.hpp

src/ddc/ddc.hpp

@@ -80,3 +80,4 @@ namespace ddc {
 
 // Output
 #include "print.hpp"
+#include "save_npy.hpp"

src/ddc/save_npy.cpp

@@ -0,0 +1,90 @@
+// Copyright (C) The DDC development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT
+
+#include <bit>
+#include <cstddef>
+#include <cstdint>
+#include <filesystem>
+#include <fstream>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "save_npy.hpp"
+
+namespace ddc::detail {
+
+NpyByteOrder get_byte_order(std::size_t itemsize) noexcept
+{
+    if (itemsize == 1) {
+        return NpyByteOrder::not_applicable;
+    }
+
+    if (std::endian::native == std::endian::little) {
+        return NpyByteOrder::little_endian;
+    }
+
+    return NpyByteOrder::big_endian;
+}
+
+std::string NpyDtype::str() const
+{
+    return std::string(1, static_cast<char>(byte_order)) + static_cast<char>(kind)
+           + std::to_string(itemsize);
+}
+
+// See specification at https://numpy.org/neps/nep-0001-npy-format.html#format-specification-version-1-0
+void save_npy(std::ostream& os, NpyArrayView const& view)
+{
+    // Build shape string: (d0, d1, ..., dN,)
+    std::string shape_str = "(";
+    for (std::size_t ext : view.shape) {
+        shape_str += std::to_string(ext);
+        shape_str += ", ";
+    }
+    shape_str += ")";
+
+    std::string const header_dict
+            = std::string("{'descr': '") + view.dtype.str() + "', 'fortran_order': "
+              + (view.fortran_order ? "True" : "False") + ", 'shape': " + shape_str + ", }";
+
+    // Pad header to a multiple of 16
+    std::size_t const non_padded_header_len = header_dict.size() + 1;
+    // magic(6) + major(1) + minor(1) + header_len(2) + header
+    std::size_t const padding = 16 - (6 + 1 + 1 + 2 + non_padded_header_len) % 16;
+    if (non_padded_header_len + padding > std::numeric_limits<std::uint16_t>::max()) {
+        throw std::runtime_error("save_npy: header too large for npy v1.0.");
+    }
+    auto const header_len = static_cast<std::uint16_t>(non_padded_header_len + padding);
+
+    // magic string
+    os.write("\x93NUMPY", 6);
+    // major version
+    os.put(1);
+    // minor version
+    os.put(0);
+
+    // header length
+    os.write(reinterpret_cast<char const*>(&header_len), sizeof(header_len));
+    // header + padding + newline
+    os.write(header_dict.data(), header_dict.size());
+    os.write("               ", padding);
+    os.put('\n');
+
+    // Raw data
+    std::size_t const n_elems
+            = std::accumulate(view.shape.begin(), view.shape.end(), 1ULL, std::multiplies<> {});
+    os.write(reinterpret_cast<char const*>(view.data), n_elems * view.dtype.itemsize);
+}
+
+void save_npy(std::filesystem::path const& filename, NpyArrayView const& view)
+{
+    std::ofstream file(filename, std::ios::binary);
+    file.exceptions(std::ios::failbit | std::ios::badbit);
+
+    save_npy(file, view);
+}
+
+} // namespace ddc::detail

src/ddc/save_npy.hpp

@@ -0,0 +1,153 @@
+// Copyright (C) The DDC development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <complex>
+#include <cstddef>
+#include <filesystem>
+#include <iosfwd>
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include <Kokkos_Core.hpp>
+
+namespace ddc::detail {
+
+enum class NpyByteOrder : char { little_endian = '<', big_endian = '>', not_applicable = '|' };
+
+NpyByteOrder get_byte_order(std::size_t itemsize) noexcept;
+
+enum class NpyKind : char {
+    boolean = 'b',
+    signed_int = 'i',
+    unsigned_int = 'u',
+    floating_point = 'f',
+    complex = 'c',
+    other = 'V',
+};
+
+struct NpyDtype
+{
+    NpyByteOrder byte_order;
+    NpyKind kind;
+    std::size_t itemsize; // in bytes
+
+    std::string str() const;
+};
+
+template <typename T>
+NpyDtype convert_to_npy_dtype()
+{
+    std::size_t const itemsize = sizeof(T);
+    NpyByteOrder const byte_order = get_byte_order(itemsize);
+    NpyKind kind;
+
+    if constexpr (std::is_same_v<T, bool>) {
+        // ── Single-byte / untyped ─────────────────────────────────────────
+        kind = NpyKind::boolean;
+    } else if constexpr (std::is_same_v<T, std::byte>) {
+        // std::byte → raw byte buffer, no arithmetic meaning
+        kind = NpyKind::other;
+    } else if constexpr (std::is_same_v<T, char>) {
+        // char is a distinct type; its signedness is implementation-defined
+        kind = std::is_signed_v<char> ? NpyKind::signed_int : NpyKind::unsigned_int;
+    } else if constexpr (
+            std::is_same_v<T, std::complex<float>> || std::is_same_v<T, std::complex<double>>) {
+        // ── Complex ───────────────────────────────────────────────────────
+        // NumPy 'c' dtype stores interleaved real+imag, same layout as std::complex
+        kind = NpyKind::complex;
+    } else if constexpr (std::is_floating_point_v<T>) {
+        // ── Floating-point ────────────────────────────────────────────────
+        static_assert(
+                !std::is_same_v<T, long double>,
+                "long double is platform-specific (80/96/128-bit); cast to double first.");
+        kind = NpyKind::floating_point;
+    } else if constexpr (std::is_signed_v<T>) {
+        // ── Integers ──────────────────────────────────────────────────────
+        kind = NpyKind::signed_int;
+    } else if constexpr (std::is_unsigned_v<T>) {
+        kind = NpyKind::unsigned_int;
+    } else {
+        static_assert(sizeof(T) == 0, "Unsupported type for NpyDtype::of<T>()");
+    }
+
+    return {.byte_order = byte_order, .kind = kind, .itemsize = itemsize};
+}
+
+struct NpyArrayView
+{
+    void const* data;
+    NpyDtype dtype;
+    std::vector<std::size_t> shape;
+    bool fortran_order;
+};
+
+void save_npy(std::ostream& os, NpyArrayView const& view);
+
+void save_npy(std::filesystem::path const& filename, NpyArrayView const& view);
+
+} // namespace ddc::detail
+
+namespace ddc::experimental {
+
+template <typename T, typename Extents, typename Layout, typename Accessor>
+void save_npy(std::ostream& os, Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
+{
+    static_assert(
+            std::is_same_v<Layout, Kokkos::layout_left>
+                    || std::is_same_v<Layout, Kokkos::layout_right>,
+            "save_npy: only contiguous layouts supported.");
+    static_assert(
+            std::is_same_v<Accessor, Kokkos::default_accessor<T>>
+                    || std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
+            "save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
+
+    std::vector<std::size_t> shape(Extents::rank());
+    for (std::size_t i = 0; i < Extents::rank(); ++i) {
+        shape[i] = mds.extent(i);
+    }
+
+    ddc::detail::save_npy(
+            os,
+            ddc::detail::NpyArrayView {
+                    .data = mds.data_handle(),
+                    .dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
+                    .shape = std::move(shape),
+                    .fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
+            });
+}
+
+template <typename T, typename Extents, typename Layout, typename Accessor>
+void save_npy(
+        std::filesystem::path const& filename,
+        Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
+{
+    static_assert(
+            std::is_same_v<Layout, Kokkos::layout_left>
+                    || std::is_same_v<Layout, Kokkos::layout_right>,
+            "save_npy: only contiguous layouts supported.");
+    static_assert(
+            std::is_same_v<Accessor, Kokkos::default_accessor<T>>
+                    || std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
+            "save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
+
+    std::vector<std::size_t> shape(Extents::rank());
+    for (std::size_t i = 0; i < Extents::rank(); ++i) {
+        shape[i] = mds.extent(i);
+    }
+
+    ddc::detail::save_npy(
+            filename,
+            ddc::detail::NpyArrayView {
+                    .data = mds.data_handle(),
+                    .dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
+                    .shape = std::move(shape),
+                    .fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
+            });
+}
+
+} // namespace ddc::experimental

tests/CMakeLists.txt

@@ -6,6 +6,8 @@ include(GoogleTest)
 
 add_subdirectory(discrete_space)
 
+add_subdirectory(save_npy)
+
 add_library(ddc_gtest_main)
 target_compile_features(ddc_gtest_main PRIVATE cxx_std_20)
 target_link_libraries(ddc_gtest_main PRIVATE DDC::core GTest::gtest)
@@ -41,6 +43,7 @@ add_executable(
     reducer.cpp
     relocatable_device_code.cpp
     relocatable_device_code_initialization.cpp
+    save_npy.cpp
     sparse_discrete_domain.cpp
     strided_discrete_domain.cpp
     tagged_vector.cpp

tests/save_npy.cpp

@@ -0,0 +1,104 @@
+// Copyright (C) The DDC development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT
+
+// #include <array>
+#include <cmath>
+// #include <complex>
+#include <stdexcept>
+
+#include <ddc/ddc.hpp>
+
+#include <gtest/gtest.h>
+
+#include <Kokkos_Core.hpp>
+
+// namespace {
+
+// std::array constexpr ns {2, 3, 4};
+// int constexpr n = ns[0] * ns[1] * ns[2];
+
+// template <typename T>
+// constexpr T make_value()
+// {
+//     std::complex<double> const base_value(2.3, 0.4);
+//     if constexpr (
+//             std::is_same_v<T, std::complex<float>> || std::is_same_v<T, std::complex<double>>) {
+//         return T(base_value);
+//     } else if constexpr (std::is_floating_point_v<T>) {
+//         return T(std::real(base_value));
+//     } else {
+//         return T(std::llround(std::real(base_value)));
+//     }
+// }
+
+// } // namespace
+
+// template <typename T>
+// struct SaveNpyTest : public ::testing::Test
+// {
+//     using data_type = T;
+// };
+
+// using SaveNpyTypes = ::testing::Types<
+//         float,
+//         double,
+//         std::complex<float>,
+//         std::complex<double>,
+//         char,
+//         signed char,
+//         signed short,
+//         signed int,
+//         signed long,
+//         signed long long,
+//         unsigned char,
+//         unsigned short,
+//         unsigned int,
+//         unsigned long,
+//         unsigned long long>;
+
+// TYPED_TEST_SUITE(SaveNpyTest, SaveNpyTypes);
+
+// TYPED_TEST(SaveNpyTest, SaveNpy0d)
+// {
+//     using data_type = typename TestFixture::data_type;
+//     std::string const label(typeid(data_type).name());
+//     Kokkos::View<data_type, Kokkos::HostSpace> const alloc(label);
+//     Kokkos::deep_copy(alloc, make_value<data_type>());
+//     Kokkos::mdspan const view(alloc.data());
+
+//     ddc::experimental::save_npy("test.npy", view);
+// }
+
+// TYPED_TEST(SaveNpyTest, SaveNpy1d)
+// {
+//     using data_type = typename TestFixture::data_type;
+//     std::string const label(typeid(data_type).name());
+//     Kokkos::View<data_type*, Kokkos::HostSpace> const alloc(label, n);
+//     Kokkos::deep_copy(alloc, make_value<data_type>());
+//     Kokkos::mdspan const view(alloc.data(), n);
+
+//     ddc::experimental::save_npy("test_" + label + "_1d.npy", view);
+// }
+
+// TYPED_TEST(SaveNpyTest, SaveNpy3d)
+// {
+//     using data_type = typename TestFixture::data_type;
+//     std::string const label(typeid(data_type).name());
+//     Kokkos::View<data_type*, Kokkos::HostSpace> const alloc(label, n);
+//     Kokkos::deep_copy(alloc, make_value<data_type>());
+//     Kokkos::mdspan const view(alloc.data(), ns);
+
+//     ddc::experimental::save_npy("test_" + label + "_3d.npy", view);
+// }
+
+TEST(SaveNpy, LargeHeader)
+{
+    ddc::detail::NpyArrayView const view {
+            .data = nullptr,
+            .dtype = ddc::detail::convert_to_npy_dtype<char>(),
+            .shape = std::vector<std::size_t>(25'000, 0),
+            .fortran_order = true,
+    };
+    EXPECT_THROW(ddc::detail::save_npy("test.npy", view), std::runtime_error);
+}