Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 1 | /* |
| 2 | tests/test_numpy_array.cpp -- test core array functionality |
| 3 | |
| 4 | Copyright (c) 2016 Ivan Smirnov <i.s.smirnov@gmail.com> |
| 5 | |
| 6 | All rights reserved. Use of this source code is governed by a |
| 7 | BSD-style license that can be found in the LICENSE file. |
| 8 | */ |
| 9 | |
| 10 | #include "pybind11_tests.h" |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 11 | |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 12 | #include <pybind11/numpy.h> |
| 13 | #include <pybind11/stl.h> |
| 14 | |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 15 | #include <cstdint> |
| 16 | #include <vector> |
| 17 | |
| 18 | using arr = py::array; |
| 19 | using arr_t = py::array_t<uint16_t, 0>; |
| 20 | |
| 21 | template<typename... Ix> arr data(const arr& a, Ix&&... index) { |
| 22 | return arr(a.nbytes() - a.offset_at(index...), (const uint8_t *) a.data(index...)); |
| 23 | } |
| 24 | |
| 25 | template<typename... Ix> arr data_t(const arr_t& a, Ix&&... index) { |
| 26 | return arr(a.size() - a.index_at(index...), a.data(index...)); |
| 27 | } |
| 28 | |
| 29 | arr& mutate_data(arr& a) { |
| 30 | auto ptr = (uint8_t *) a.mutable_data(); |
| 31 | for (size_t i = 0; i < a.nbytes(); i++) |
| 32 | ptr[i] = (uint8_t) (ptr[i] * 2); |
| 33 | return a; |
| 34 | } |
| 35 | |
| 36 | arr_t& mutate_data_t(arr_t& a) { |
| 37 | auto ptr = a.mutable_data(); |
| 38 | for (size_t i = 0; i < a.size(); i++) |
| 39 | ptr[i]++; |
| 40 | return a; |
| 41 | } |
| 42 | |
| 43 | template<typename... Ix> arr& mutate_data(arr& a, Ix&&... index) { |
| 44 | auto ptr = (uint8_t *) a.mutable_data(index...); |
| 45 | for (size_t i = 0; i < a.nbytes() - a.offset_at(index...); i++) |
| 46 | ptr[i] = (uint8_t) (ptr[i] * 2); |
| 47 | return a; |
| 48 | } |
| 49 | |
| 50 | template<typename... Ix> arr_t& mutate_data_t(arr_t& a, Ix&&... index) { |
| 51 | auto ptr = a.mutable_data(index...); |
| 52 | for (size_t i = 0; i < a.size() - a.index_at(index...); i++) |
| 53 | ptr[i]++; |
| 54 | return a; |
| 55 | } |
| 56 | |
| 57 | template<typename... Ix> size_t index_at(const arr& a, Ix&&... idx) { return a.index_at(idx...); } |
| 58 | template<typename... Ix> size_t index_at_t(const arr_t& a, Ix&&... idx) { return a.index_at(idx...); } |
| 59 | template<typename... Ix> size_t offset_at(const arr& a, Ix&&... idx) { return a.offset_at(idx...); } |
| 60 | template<typename... Ix> size_t offset_at_t(const arr_t& a, Ix&&... idx) { return a.offset_at(idx...); } |
| 61 | template<typename... Ix> size_t at_t(const arr_t& a, Ix&&... idx) { return a.at(idx...); } |
| 62 | template<typename... Ix> arr_t& mutate_at_t(arr_t& a, Ix&&... idx) { a.mutable_at(idx...)++; return a; } |
| 63 | |
| 64 | #define def_index_fn(name, type) \ |
| 65 | sm.def(#name, [](type a) { return name(a); }); \ |
| 66 | sm.def(#name, [](type a, int i) { return name(a, i); }); \ |
| 67 | sm.def(#name, [](type a, int i, int j) { return name(a, i, j); }); \ |
| 68 | sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); }); |
| 69 | |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 70 | test_initializer numpy_array([](py::module &m) { |
Ivan Smirnov | aca6bca | 2016-09-08 23:03:35 +0100 | [diff] [blame] | 71 | auto sm = m.def_submodule("array"); |
| 72 | |
| 73 | sm.def("ndim", [](const arr& a) { return a.ndim(); }); |
| 74 | sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); }); |
| 75 | sm.def("shape", [](const arr& a, size_t dim) { return a.shape(dim); }); |
| 76 | sm.def("strides", [](const arr& a) { return arr(a.ndim(), a.strides()); }); |
| 77 | sm.def("strides", [](const arr& a, size_t dim) { return a.strides(dim); }); |
| 78 | sm.def("writeable", [](const arr& a) { return a.writeable(); }); |
| 79 | sm.def("size", [](const arr& a) { return a.size(); }); |
| 80 | sm.def("itemsize", [](const arr& a) { return a.itemsize(); }); |
| 81 | sm.def("nbytes", [](const arr& a) { return a.nbytes(); }); |
| 82 | sm.def("owndata", [](const arr& a) { return a.owndata(); }); |
| 83 | |
| 84 | def_index_fn(data, const arr&); |
| 85 | def_index_fn(data_t, const arr_t&); |
| 86 | def_index_fn(index_at, const arr&); |
| 87 | def_index_fn(index_at_t, const arr_t&); |
| 88 | def_index_fn(offset_at, const arr&); |
| 89 | def_index_fn(offset_at_t, const arr_t&); |
| 90 | def_index_fn(mutate_data, arr&); |
| 91 | def_index_fn(mutate_data_t, arr_t&); |
| 92 | def_index_fn(at_t, const arr_t&); |
| 93 | def_index_fn(mutate_at_t, arr_t&); |
Wenzel Jakob | 43f6aa6 | 2016-10-12 23:34:06 +0200 | [diff] [blame] | 94 | |
| 95 | sm.def("make_f_array", [] { |
| 96 | return py::array_t<float>({ 2, 2 }, { 4, 8 }); |
| 97 | }); |
| 98 | |
| 99 | sm.def("make_c_array", [] { |
| 100 | return py::array_t<float>({ 2, 2 }, { 8, 4 }); |
| 101 | }); |
Wenzel Jakob | 369e9b3 | 2016-10-13 00:57:42 +0200 | [diff] [blame] | 102 | |
| 103 | sm.def("wrap", [](py::array a) { |
| 104 | return py::array( |
| 105 | a.dtype(), |
| 106 | std::vector<size_t>(a.shape(), a.shape() + a.ndim()), |
| 107 | std::vector<size_t>(a.strides(), a.strides() + a.ndim()), |
| 108 | a.data(), |
| 109 | a |
| 110 | ); |
| 111 | }); |
Wenzel Jakob | fac7c09 | 2016-10-13 10:37:52 +0200 | [diff] [blame] | 112 | |
| 113 | struct ArrayClass { |
| 114 | int data[2] = { 1, 2 }; |
| 115 | ArrayClass() { py::print("ArrayClass()"); } |
| 116 | ~ArrayClass() { py::print("~ArrayClass()"); } |
| 117 | }; |
| 118 | |
| 119 | py::class_<ArrayClass>(sm, "ArrayClass") |
| 120 | .def(py::init<>()) |
| 121 | .def("numpy_view", [](py::object &obj) { |
| 122 | py::print("ArrayClass::numpy_view()"); |
| 123 | ArrayClass &a = obj.cast<ArrayClass&>(); |
| 124 | return py::array_t<int>({2}, {4}, a.data, obj); |
| 125 | } |
| 126 | ); |
Wenzel Jakob | 496feac | 2016-10-28 00:37:07 +0200 | [diff] [blame^] | 127 | |
| 128 | sm.def("function_taking_uint64", [](uint64_t){ }); |
Ivan Smirnov | 91b3d68 | 2016-08-29 02:41:05 +0100 | [diff] [blame] | 129 | }); |