test_numpy_array.cpp revision 11986
111986Sandreas.sandberg@arm.com/* 211986Sandreas.sandberg@arm.com tests/test_numpy_array.cpp -- test core array functionality 311986Sandreas.sandberg@arm.com 411986Sandreas.sandberg@arm.com Copyright (c) 2016 Ivan Smirnov <i.s.smirnov@gmail.com> 511986Sandreas.sandberg@arm.com 611986Sandreas.sandberg@arm.com All rights reserved. Use of this source code is governed by a 711986Sandreas.sandberg@arm.com BSD-style license that can be found in the LICENSE file. 811986Sandreas.sandberg@arm.com*/ 911986Sandreas.sandberg@arm.com 1011986Sandreas.sandberg@arm.com#include "pybind11_tests.h" 1111986Sandreas.sandberg@arm.com 1211986Sandreas.sandberg@arm.com#include <pybind11/numpy.h> 1311986Sandreas.sandberg@arm.com#include <pybind11/stl.h> 1411986Sandreas.sandberg@arm.com 1511986Sandreas.sandberg@arm.com#include <cstdint> 1611986Sandreas.sandberg@arm.com#include <vector> 1711986Sandreas.sandberg@arm.com 1811986Sandreas.sandberg@arm.comusing arr = py::array; 1911986Sandreas.sandberg@arm.comusing arr_t = py::array_t<uint16_t, 0>; 2011986Sandreas.sandberg@arm.com 2111986Sandreas.sandberg@arm.comtemplate<typename... Ix> arr data(const arr& a, Ix... index) { 2211986Sandreas.sandberg@arm.com return arr(a.nbytes() - a.offset_at(index...), (const uint8_t *) a.data(index...)); 2311986Sandreas.sandberg@arm.com} 2411986Sandreas.sandberg@arm.com 2511986Sandreas.sandberg@arm.comtemplate<typename... Ix> arr data_t(const arr_t& a, Ix... index) { 2611986Sandreas.sandberg@arm.com return arr(a.size() - a.index_at(index...), a.data(index...)); 2711986Sandreas.sandberg@arm.com} 2811986Sandreas.sandberg@arm.com 2911986Sandreas.sandberg@arm.comarr& mutate_data(arr& a) { 3011986Sandreas.sandberg@arm.com auto ptr = (uint8_t *) a.mutable_data(); 3111986Sandreas.sandberg@arm.com for (size_t i = 0; i < a.nbytes(); i++) 3211986Sandreas.sandberg@arm.com ptr[i] = (uint8_t) (ptr[i] * 2); 3311986Sandreas.sandberg@arm.com return a; 3411986Sandreas.sandberg@arm.com} 3511986Sandreas.sandberg@arm.com 3611986Sandreas.sandberg@arm.comarr_t& mutate_data_t(arr_t& a) { 3711986Sandreas.sandberg@arm.com auto ptr = a.mutable_data(); 3811986Sandreas.sandberg@arm.com for (size_t i = 0; i < a.size(); i++) 3911986Sandreas.sandberg@arm.com ptr[i]++; 4011986Sandreas.sandberg@arm.com return a; 4111986Sandreas.sandberg@arm.com} 4211986Sandreas.sandberg@arm.com 4311986Sandreas.sandberg@arm.comtemplate<typename... Ix> arr& mutate_data(arr& a, Ix... index) { 4411986Sandreas.sandberg@arm.com auto ptr = (uint8_t *) a.mutable_data(index...); 4511986Sandreas.sandberg@arm.com for (size_t i = 0; i < a.nbytes() - a.offset_at(index...); i++) 4611986Sandreas.sandberg@arm.com ptr[i] = (uint8_t) (ptr[i] * 2); 4711986Sandreas.sandberg@arm.com return a; 4811986Sandreas.sandberg@arm.com} 4911986Sandreas.sandberg@arm.com 5011986Sandreas.sandberg@arm.comtemplate<typename... Ix> arr_t& mutate_data_t(arr_t& a, Ix... index) { 5111986Sandreas.sandberg@arm.com auto ptr = a.mutable_data(index...); 5211986Sandreas.sandberg@arm.com for (size_t i = 0; i < a.size() - a.index_at(index...); i++) 5311986Sandreas.sandberg@arm.com ptr[i]++; 5411986Sandreas.sandberg@arm.com return a; 5511986Sandreas.sandberg@arm.com} 5611986Sandreas.sandberg@arm.com 5711986Sandreas.sandberg@arm.comtemplate<typename... Ix> size_t index_at(const arr& a, Ix... idx) { return a.index_at(idx...); } 5811986Sandreas.sandberg@arm.comtemplate<typename... Ix> size_t index_at_t(const arr_t& a, Ix... idx) { return a.index_at(idx...); } 5911986Sandreas.sandberg@arm.comtemplate<typename... Ix> size_t offset_at(const arr& a, Ix... idx) { return a.offset_at(idx...); } 6011986Sandreas.sandberg@arm.comtemplate<typename... Ix> size_t offset_at_t(const arr_t& a, Ix... idx) { return a.offset_at(idx...); } 6111986Sandreas.sandberg@arm.comtemplate<typename... Ix> size_t at_t(const arr_t& a, Ix... idx) { return a.at(idx...); } 6211986Sandreas.sandberg@arm.comtemplate<typename... Ix> arr_t& mutate_at_t(arr_t& a, Ix... idx) { a.mutable_at(idx...)++; return a; } 6311986Sandreas.sandberg@arm.com 6411986Sandreas.sandberg@arm.com#define def_index_fn(name, type) \ 6511986Sandreas.sandberg@arm.com sm.def(#name, [](type a) { return name(a); }); \ 6611986Sandreas.sandberg@arm.com sm.def(#name, [](type a, int i) { return name(a, i); }); \ 6711986Sandreas.sandberg@arm.com sm.def(#name, [](type a, int i, int j) { return name(a, i, j); }); \ 6811986Sandreas.sandberg@arm.com sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); }); 6911986Sandreas.sandberg@arm.com 7011986Sandreas.sandberg@arm.comtest_initializer numpy_array([](py::module &m) { 7111986Sandreas.sandberg@arm.com auto sm = m.def_submodule("array"); 7211986Sandreas.sandberg@arm.com 7311986Sandreas.sandberg@arm.com sm.def("ndim", [](const arr& a) { return a.ndim(); }); 7411986Sandreas.sandberg@arm.com sm.def("shape", [](const arr& a) { return arr(a.ndim(), a.shape()); }); 7511986Sandreas.sandberg@arm.com sm.def("shape", [](const arr& a, size_t dim) { return a.shape(dim); }); 7611986Sandreas.sandberg@arm.com sm.def("strides", [](const arr& a) { return arr(a.ndim(), a.strides()); }); 7711986Sandreas.sandberg@arm.com sm.def("strides", [](const arr& a, size_t dim) { return a.strides(dim); }); 7811986Sandreas.sandberg@arm.com sm.def("writeable", [](const arr& a) { return a.writeable(); }); 7911986Sandreas.sandberg@arm.com sm.def("size", [](const arr& a) { return a.size(); }); 8011986Sandreas.sandberg@arm.com sm.def("itemsize", [](const arr& a) { return a.itemsize(); }); 8111986Sandreas.sandberg@arm.com sm.def("nbytes", [](const arr& a) { return a.nbytes(); }); 8211986Sandreas.sandberg@arm.com sm.def("owndata", [](const arr& a) { return a.owndata(); }); 8311986Sandreas.sandberg@arm.com 8411986Sandreas.sandberg@arm.com def_index_fn(data, const arr&); 8511986Sandreas.sandberg@arm.com def_index_fn(data_t, const arr_t&); 8611986Sandreas.sandberg@arm.com def_index_fn(index_at, const arr&); 8711986Sandreas.sandberg@arm.com def_index_fn(index_at_t, const arr_t&); 8811986Sandreas.sandberg@arm.com def_index_fn(offset_at, const arr&); 8911986Sandreas.sandberg@arm.com def_index_fn(offset_at_t, const arr_t&); 9011986Sandreas.sandberg@arm.com def_index_fn(mutate_data, arr&); 9111986Sandreas.sandberg@arm.com def_index_fn(mutate_data_t, arr_t&); 9211986Sandreas.sandberg@arm.com def_index_fn(at_t, const arr_t&); 9311986Sandreas.sandberg@arm.com def_index_fn(mutate_at_t, arr_t&); 9411986Sandreas.sandberg@arm.com 9511986Sandreas.sandberg@arm.com sm.def("make_f_array", [] { 9611986Sandreas.sandberg@arm.com return py::array_t<float>({ 2, 2 }, { 4, 8 }); 9711986Sandreas.sandberg@arm.com }); 9811986Sandreas.sandberg@arm.com 9911986Sandreas.sandberg@arm.com sm.def("make_c_array", [] { 10011986Sandreas.sandberg@arm.com return py::array_t<float>({ 2, 2 }, { 8, 4 }); 10111986Sandreas.sandberg@arm.com }); 10211986Sandreas.sandberg@arm.com 10311986Sandreas.sandberg@arm.com sm.def("wrap", [](py::array a) { 10411986Sandreas.sandberg@arm.com return py::array( 10511986Sandreas.sandberg@arm.com a.dtype(), 10611986Sandreas.sandberg@arm.com std::vector<size_t>(a.shape(), a.shape() + a.ndim()), 10711986Sandreas.sandberg@arm.com std::vector<size_t>(a.strides(), a.strides() + a.ndim()), 10811986Sandreas.sandberg@arm.com a.data(), 10911986Sandreas.sandberg@arm.com a 11011986Sandreas.sandberg@arm.com ); 11111986Sandreas.sandberg@arm.com }); 11211986Sandreas.sandberg@arm.com 11311986Sandreas.sandberg@arm.com struct ArrayClass { 11411986Sandreas.sandberg@arm.com int data[2] = { 1, 2 }; 11511986Sandreas.sandberg@arm.com ArrayClass() { py::print("ArrayClass()"); } 11611986Sandreas.sandberg@arm.com ~ArrayClass() { py::print("~ArrayClass()"); } 11711986Sandreas.sandberg@arm.com }; 11811986Sandreas.sandberg@arm.com 11911986Sandreas.sandberg@arm.com py::class_<ArrayClass>(sm, "ArrayClass") 12011986Sandreas.sandberg@arm.com .def(py::init<>()) 12111986Sandreas.sandberg@arm.com .def("numpy_view", [](py::object &obj) { 12211986Sandreas.sandberg@arm.com py::print("ArrayClass::numpy_view()"); 12311986Sandreas.sandberg@arm.com ArrayClass &a = obj.cast<ArrayClass&>(); 12411986Sandreas.sandberg@arm.com return py::array_t<int>({2}, {4}, a.data, obj); 12511986Sandreas.sandberg@arm.com } 12611986Sandreas.sandberg@arm.com ); 12711986Sandreas.sandberg@arm.com 12811986Sandreas.sandberg@arm.com sm.def("function_taking_uint64", [](uint64_t) { }); 12911986Sandreas.sandberg@arm.com 13011986Sandreas.sandberg@arm.com sm.def("isinstance_untyped", [](py::object yes, py::object no) { 13111986Sandreas.sandberg@arm.com return py::isinstance<py::array>(yes) && !py::isinstance<py::array>(no); 13211986Sandreas.sandberg@arm.com }); 13311986Sandreas.sandberg@arm.com 13411986Sandreas.sandberg@arm.com sm.def("isinstance_typed", [](py::object o) { 13511986Sandreas.sandberg@arm.com return py::isinstance<py::array_t<double>>(o) && !py::isinstance<py::array_t<int>>(o); 13611986Sandreas.sandberg@arm.com }); 13711986Sandreas.sandberg@arm.com 13811986Sandreas.sandberg@arm.com sm.def("default_constructors", []() { 13911986Sandreas.sandberg@arm.com return py::dict( 14011986Sandreas.sandberg@arm.com "array"_a=py::array(), 14111986Sandreas.sandberg@arm.com "array_t<int32>"_a=py::array_t<std::int32_t>(), 14211986Sandreas.sandberg@arm.com "array_t<double>"_a=py::array_t<double>() 14311986Sandreas.sandberg@arm.com ); 14411986Sandreas.sandberg@arm.com }); 14511986Sandreas.sandberg@arm.com 14611986Sandreas.sandberg@arm.com sm.def("converting_constructors", [](py::object o) { 14711986Sandreas.sandberg@arm.com return py::dict( 14811986Sandreas.sandberg@arm.com "array"_a=py::array(o), 14911986Sandreas.sandberg@arm.com "array_t<int32>"_a=py::array_t<std::int32_t>(o), 15011986Sandreas.sandberg@arm.com "array_t<double>"_a=py::array_t<double>(o) 15111986Sandreas.sandberg@arm.com ); 15211986Sandreas.sandberg@arm.com }); 15311986Sandreas.sandberg@arm.com}); 154