forked from xtensor-stack/xtensor-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.cpp
More file actions
60 lines (50 loc) · 1.66 KB
/
Copy pathmain.cpp
File metadata and controls
60 lines (50 loc) · 1.66 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
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
#include "pybind11/pybind11.h"
#include "pybind11/numpy.h"
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include "numpy/arrayobject.h"
#include "xtensor/xtensor.hpp"
#include "xtensor/xarray.hpp"
#include "xtensor-python/pyarray.hpp"
#include "xtensor-python/pytensor.hpp"
#include "xtensor-python/pyvectorize.hpp"
using complex_t = std::complex<double>;
namespace py = pybind11;
PYBIND11_PLUGIN(benchmark_xtensor_python)
{
if(_import_array() < 0)
{
PyErr_SetString(PyExc_ImportError, "numpy.core.multiarray failed to import");
return nullptr;
}
py::module m("benchmark_xtensor_python", "Benchmark module for xtensor python bindings");
m.def("sum_array", [](xt::pyarray<double> const& x) {
double sum = 0;
for(auto e : x)
sum += e;
return sum;
});
m.def("sum_tensor", [](xt::pytensor<double, 1> const& x) {
double sum = 0;
for(auto e : x)
sum += e;
return sum;
});
m.def("pybind_sum_array", [](py::array_t<double> const& x) {
double sum = 0;
size_t size = x.size();
const double* data = x.data(0);
for(size_t i = 0; i < size; ++i)
sum += data[i];
return sum;
});
m.def("rect_to_polar", [](xt::pyarray<complex_t> const& a) {
return py::vectorize([](complex_t x) { return std::abs(x); })(a);
});
m.def("pybind_rect_to_polar", [](py::array a) {
if (py::isinstance<py::array_t<complex_t>>(a))
return py::vectorize([](complex_t x) { return std::abs(x); })(a);
else
throw py::type_error("rect_to_polar unhandled type");
});
return m.ptr();
}