Skip to content

Commit 75c6dbe

Browse files
tpadioleautpadioleau
committed
Add save_npy function
1 parent e712093 commit 75c6dbe

10 files changed

Lines changed: 541 additions & 3 deletions

File tree

‎.github/workflows/python-checks.yaml‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,10 @@ jobs:
4040
runs-on: ubuntu-latest
4141
steps:
4242
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
43-
- name: Install dependencies
44-
run: |
45-
pipx install pylint~=4.0
43+
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
44+
with:
45+
python-version: '3.14'
46+
- run: pip install numpy pylint~=4.0
4647
- name: Analysing the code with pylint
4748
run: |
4849
pylint $(git ls-files '*.py' ':!benchmarks/splines_plot.py')

‎CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ target_sources(
138138
src/ddc/non_uniform_point_sampling.cpp
139139
src/ddc/periodic_sampling.cpp
140140
src/ddc/print.cpp
141+
src/ddc/save_npy.cpp
141142
src/ddc/scope_guard.cpp
142143
src/ddc/uniform_point_sampling.cpp
143144
PUBLIC
@@ -182,6 +183,7 @@ target_sources(
182183
src/ddc/print.hpp
183184
src/ddc/real_type.hpp
184185
src/ddc/reducer.hpp
186+
src/ddc/save_npy.hpp
185187
src/ddc/scope_guard.hpp
186188
src/ddc/sparse_discrete_domain.hpp
187189
src/ddc/strided_discrete_domain.hpp

‎src/ddc/ddc.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,3 +82,4 @@ namespace ddc {
8282

8383
// Output
8484
#include "print.hpp"
85+
#include "save_npy.hpp"

‎src/ddc/save_npy.cpp‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// Copyright (C) The DDC development team, see COPYRIGHT.md file
2+
//
3+
// SPDX-License-Identifier: MIT
4+
5+
#include <bit>
6+
#include <cstddef>
7+
#include <cstdint>
8+
#include <filesystem>
9+
#include <fstream>
10+
#include <numeric>
11+
#include <stdexcept>
12+
#include <string>
13+
#include <vector>
14+
15+
#include "save_npy.hpp"
16+
17+
namespace ddc::detail {
18+
19+
NpyByteOrder get_byte_order(std::size_t itemsize) noexcept
20+
{
21+
if (itemsize == 1) {
22+
return NpyByteOrder::not_applicable;
23+
}
24+
25+
if (std::endian::native == std::endian::little) {
26+
return NpyByteOrder::little_endian;
27+
}
28+
29+
return NpyByteOrder::big_endian;
30+
}
31+
32+
std::string NpyDtype::str() const
33+
{
34+
return std::string(1, static_cast<char>(byte_order)) + static_cast<char>(kind)
35+
+ std::to_string(itemsize);
36+
}
37+
38+
// See specification at https://numpy.org/neps/nep-0001-npy-format.html#format-specification-version-1-0
39+
void save_npy(std::ostream& os, NpyArrayView const& view)
40+
{
41+
// Build shape string: (d0, d1, ..., dN,)
42+
std::string shape_str = "(";
43+
for (std::size_t ext : view.shape) {
44+
shape_str += std::to_string(ext);
45+
shape_str += ", ";
46+
}
47+
shape_str += ")";
48+
49+
std::string const header_dict
50+
= std::string("{'descr': '") + view.dtype.str() + "', 'fortran_order': "
51+
+ (view.fortran_order ? "True" : "False") + ", 'shape': " + shape_str + ", }";
52+
53+
// Pad header to a multiple of 16
54+
std::size_t const non_padded_header_len = header_dict.size() + 1;
55+
// magic(6) + major(1) + minor(1) + header_len(2) + header
56+
std::size_t const padding = 16 - (6 + 1 + 1 + 2 + non_padded_header_len) % 16;
57+
if (non_padded_header_len + padding > std::numeric_limits<std::uint16_t>::max()) {
58+
throw std::runtime_error("save_npy: header too large for npy v1.0.");
59+
}
60+
auto const header_len = static_cast<std::uint16_t>(non_padded_header_len + padding);
61+
62+
// magic string
63+
os.write("\x93NUMPY", 6);
64+
// major version
65+
os.put(1);
66+
// minor version
67+
os.put(0);
68+
69+
// header length
70+
os.write(reinterpret_cast<char const*>(&header_len), sizeof(header_len));
71+
// header + padding + newline
72+
os.write(header_dict.data(), header_dict.size());
73+
os.write(" ", padding);
74+
os.put('\n');
75+
76+
// Raw data
77+
std::size_t const n_elems
78+
= std::accumulate(view.shape.begin(), view.shape.end(), 1ULL, std::multiplies<> {});
79+
os.write(reinterpret_cast<char const*>(view.data), n_elems * view.dtype.itemsize);
80+
}
81+
82+
void save_npy(std::filesystem::path const& filename, NpyArrayView const& view)
83+
{
84+
std::ofstream file(filename, std::ios::binary);
85+
file.exceptions(std::ios::failbit | std::ios::badbit);
86+
87+
save_npy(file, view);
88+
}
89+
90+
} // namespace ddc::detail

‎src/ddc/save_npy.hpp‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
// Copyright (C) The DDC development team, see COPYRIGHT.md file
2+
//
3+
// SPDX-License-Identifier: MIT
4+
5+
#pragma once
6+
7+
#include <complex>
8+
#include <cstddef>
9+
#include <filesystem>
10+
#include <iosfwd>
11+
#include <string>
12+
#include <type_traits>
13+
#include <utility>
14+
#include <vector>
15+
16+
#include <Kokkos_Core.hpp>
17+
18+
namespace ddc::detail {
19+
20+
enum class NpyByteOrder : char { little_endian = '<', big_endian = '>', not_applicable = '|' };
21+
22+
NpyByteOrder get_byte_order(std::size_t itemsize) noexcept;
23+
24+
enum class NpyKind : char {
25+
boolean = 'b',
26+
signed_int = 'i',
27+
unsigned_int = 'u',
28+
floating_point = 'f',
29+
complex = 'c',
30+
other = 'V',
31+
};
32+
33+
struct NpyDtype
34+
{
35+
NpyByteOrder byte_order;
36+
NpyKind kind;
37+
std::size_t itemsize; // in bytes
38+
39+
std::string str() const;
40+
};
41+
42+
template <typename T>
43+
NpyDtype convert_to_npy_dtype()
44+
{
45+
std::size_t const itemsize = sizeof(T);
46+
NpyByteOrder const byte_order = get_byte_order(itemsize);
47+
NpyKind kind;
48+
49+
if constexpr (std::is_same_v<T, bool>) {
50+
// ── Single-byte / untyped ─────────────────────────────────────────
51+
kind = NpyKind::boolean;
52+
} else if constexpr (std::is_same_v<T, std::byte>) {
53+
// std::byte → raw byte buffer, no arithmetic meaning
54+
kind = NpyKind::other;
55+
} else if constexpr (std::is_same_v<T, char>) {
56+
// char is a distinct type; its signedness is implementation-defined
57+
kind = std::is_signed_v<char> ? NpyKind::signed_int : NpyKind::unsigned_int;
58+
} else if constexpr (
59+
std::is_same_v<T, std::complex<float>> || std::is_same_v<T, std::complex<double>>) {
60+
// ── Complex ───────────────────────────────────────────────────────
61+
// NumPy 'c' dtype stores interleaved real+imag, same layout as std::complex
62+
kind = NpyKind::complex;
63+
} else if constexpr (std::is_floating_point_v<T>) {
64+
// ── Floating-point ────────────────────────────────────────────────
65+
static_assert(
66+
!std::is_same_v<T, long double>,
67+
"long double is platform-specific (80/96/128-bit); cast to double first.");
68+
kind = NpyKind::floating_point;
69+
} else if constexpr (std::is_signed_v<T>) {
70+
// ── Integers ──────────────────────────────────────────────────────
71+
kind = NpyKind::signed_int;
72+
} else if constexpr (std::is_unsigned_v<T>) {
73+
kind = NpyKind::unsigned_int;
74+
} else {
75+
static_assert(sizeof(T) == 0, "Unsupported type for NpyDtype::of<T>()");
76+
}
77+
78+
return {.byte_order = byte_order, .kind = kind, .itemsize = itemsize};
79+
}
80+
81+
struct NpyArrayView
82+
{
83+
void const* data;
84+
NpyDtype dtype;
85+
std::vector<std::size_t> shape;
86+
bool fortran_order;
87+
};
88+
89+
void save_npy(std::ostream& os, NpyArrayView const& view);
90+
91+
void save_npy(std::filesystem::path const& filename, NpyArrayView const& view);
92+
93+
} // namespace ddc::detail
94+
95+
namespace ddc::experimental {
96+
97+
template <typename T, typename Extents, typename Layout, typename Accessor>
98+
void save_npy(std::ostream& os, Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
99+
{
100+
static_assert(
101+
std::is_same_v<Layout, Kokkos::layout_left>
102+
|| std::is_same_v<Layout, Kokkos::layout_right>,
103+
"save_npy: only contiguous layouts supported.");
104+
static_assert(
105+
std::is_same_v<Accessor, Kokkos::default_accessor<T>>
106+
|| std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
107+
"save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
108+
109+
std::vector<std::size_t> shape(Extents::rank());
110+
for (std::size_t i = 0; i < Extents::rank(); ++i) {
111+
shape[i] = mds.extent(i);
112+
}
113+
114+
ddc::detail::save_npy(
115+
os,
116+
ddc::detail::NpyArrayView {
117+
.data = mds.data_handle(),
118+
.dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
119+
.shape = std::move(shape),
120+
.fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
121+
});
122+
}
123+
124+
template <typename T, typename Extents, typename Layout, typename Accessor>
125+
void save_npy(
126+
std::filesystem::path const& filename,
127+
Kokkos::mdspan<T, Extents, Layout, Accessor> const& mds)
128+
{
129+
static_assert(
130+
std::is_same_v<Layout, Kokkos::layout_left>
131+
|| std::is_same_v<Layout, Kokkos::layout_right>,
132+
"save_npy: only contiguous layouts supported.");
133+
static_assert(
134+
std::is_same_v<Accessor, Kokkos::default_accessor<T>>
135+
|| std::is_same_v<Accessor, Kokkos::default_accessor<T const>>,
136+
"save_npy: non-host-accessible accessor. Use create_mirror_view + deep_copy first.");
137+
138+
std::vector<std::size_t> shape(Extents::rank());
139+
if constexpr (Extents::rank() > 0) {
140+
for (std::size_t i = 0; i < Extents::rank(); ++i) {
141+
shape[i] = mds.extent(i);
142+
}
143+
}
144+
145+
ddc::detail::save_npy(
146+
filename,
147+
ddc::detail::NpyArrayView {
148+
.data = mds.data_handle(),
149+
.dtype = ddc::detail::convert_to_npy_dtype<std::remove_const_t<T>>(),
150+
.shape = std::move(shape),
151+
.fortran_order = std::is_same_v<Layout, Kokkos::layout_left>,
152+
});
153+
}
154+
155+
} // namespace ddc::experimental

‎tests/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ include(GoogleTest)
66

77
add_subdirectory(discrete_space)
88

9+
add_subdirectory(save_npy)
10+
911
add_library(ddc_gtest_main)
1012
target_compile_features(ddc_gtest_main PRIVATE cxx_std_20)
1113
target_link_libraries(ddc_gtest_main PRIVATE DDC::core GTest::gtest)
@@ -42,6 +44,7 @@ add_executable(
4244
reducer.cpp
4345
relocatable_device_code.cpp
4446
relocatable_device_code_initialization.cpp
47+
save_npy.cpp
4548
sparse_discrete_domain.cpp
4649
strided_discrete_domain.cpp
4750
tagged_vector.cpp

0 commit comments

Comments
 (0)