-
Notifications
You must be signed in to change notification settings - Fork 31
Expand file tree
/
Copy patharray_like.hpp
More file actions
35 lines (31 loc) · 1.27 KB
/
array_like.hpp
File metadata and controls
35 lines (31 loc) · 1.27 KB
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
// Copyright 2018-2019 Hans Dembinski and Henry Schreiner
//
// Distributed under the 3-Clause BSD License. See accompanying
// file LICENSE or https://github.com/scikit-hep/boost-histogram for details.
#pragma once
#include <bh_python/pybind11.hpp>
#include <boost/core/span.hpp>
#include <vector>
/// Generate empty array with same shape and strides as argument
template <class T>
py::array_t<T> array_like(const py::object& obj) {
if(!py::isinstance<py::array>(obj)) {
py::ssize_t shape[1] = {0}; // if scalar
if(py::isinstance<py::sequence>(obj) && !py::isinstance<py::str>(obj)) {
// if sequence
auto seq = py::cast<py::sequence>(obj);
shape[0] = static_cast<py::ssize_t>(seq.size());
}
return py::array_t<T>(shape);
}
auto arr = py::cast<py::array>(obj);
std::vector<py::ssize_t> strides;
strides.reserve(static_cast<std::size_t>(arr.ndim()));
for(int i = 0; i < arr.ndim(); ++i) {
strides.emplace_back(arr.strides()[i] / arr.itemsize()
* static_cast<py::ssize_t>(sizeof(T)));
}
return py::array_t<T>{boost::span<const py::ssize_t>{
arr.shape(), static_cast<std::size_t>(arr.ndim())},
strides};
}