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
|
#pragma once
#include <array>
#include <vector>
#include "common.h"
#include <nanobind/ndarray.h>
template<typename T>
using cpu_array = nb::ndarray<T, nb::shape<-1>, nb::device::cpu>;
template<typename T>
using cpu_c_array = nb::ndarray<T, nb::shape<-1>, nb::device::cpu, nb::c_contig>;
using cpu_miller_array = nb::ndarray<int, nb::shape<-1,3>, nb::device::cpu>;
using cpu_c_miller_array = nb::ndarray<int, nb::shape<-1,3>, nb::device::cpu, nb::c_contig>;
template<typename T>
auto numpy_array_from_vector(std::vector<T>&& original_vec) {
using V = std::vector<T>;
auto v = new V(std::move(original_vec));
nb::capsule owner(v, [](void* p) noexcept { delete static_cast<V*>(p); });
return nb::ndarray<nb::numpy, T, nb::shape<-1>>(v->data(), {v->size()}, owner);
}
template<typename T, size_t N>
auto py_array2d_from_vector(std::vector<std::array<T,N>>&& original_vec) {
using V = std::vector<std::array<T,N>>;
auto v = new V(std::move(original_vec));
nb::capsule owner(v, [](void* p) noexcept { delete static_cast<V*>(p); });
return nb::ndarray<nb::numpy, T, nb::shape<-1, N>>(v->data(), {v->size(), N}, owner);
}
// to be used with rv_policy::reference_internal
template<typename T, typename S>
auto vector_member_array(std::vector<S>& vec, T S::*ptr) {
constexpr int64_t stride = static_cast<int64_t>(sizeof(S) / sizeof(T));
static_assert(stride * sizeof(T) == sizeof(S),
"vector_member_array(): problem with stride");
return nb::ndarray<nb::numpy, T, nb::shape<-1>>(
&(vec.data()->*ptr),
{vec.size()},
nb::handle(),
{stride});
}
template<typename T>
auto make_numpy_array(std::initializer_list<size_t> size,
std::initializer_list<int64_t> strides={}) {
size_t total_size = 1;
for (size_t i : size)
total_size *= i;
T* c_array = new T[total_size];
nb::capsule owner(c_array, [](void* p) noexcept { delete [] static_cast<T*>(p); });
return nb::ndarray<nb::numpy, T>(c_array, size, owner, strides);
}
template<typename Ret, typename Obj, typename Func>
auto miller_function(const Obj& obj, Func func, const cpu_miller_array& hkl) {
auto h = hkl.view();
size_t n = h.shape(0);
auto result = make_numpy_array<Ret>({n});
Ret* rptr = result.data();
for (size_t i = 0; i < n; ++i)
rptr[i] = (obj.*func)({h(i, 0), h(i, 1), h(i, 2)});
return result;
}
|
