Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 1 | /* |
Dean Moldovan | a0c1ccf | 2016-08-12 13:50:00 +0200 | [diff] [blame] | 2 | tests/test_numpy_vectorize.cpp -- auto-vectorize functions over NumPy array |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 3 | arguments |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 4 | |
Wenzel Jakob | 8cb6cb3 | 2016-04-17 20:21:41 +0200 | [diff] [blame] | 5 | Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 6 | |
| 7 | All rights reserved. Use of this source code is governed by a |
| 8 | BSD-style license that can be found in the LICENSE file. |
| 9 | */ |
| 10 | |
Dean Moldovan | a0c1ccf | 2016-08-12 13:50:00 +0200 | [diff] [blame] | 11 | #include "pybind11_tests.h" |
Wenzel Jakob | 8f4eb00 | 2015-10-15 18:13:33 +0200 | [diff] [blame] | 12 | #include <pybind11/numpy.h> |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 13 | |
| 14 | double my_func(int x, float y, double z) { |
Dean Moldovan | 81511be | 2016-09-07 00:50:10 +0200 | [diff] [blame] | 15 | py::print("my_func(x:int={}, y:float={:.0f}, z:float={:.0f})"_s.format(x, y, z)); |
Boris Schäling | 20ee935 | 2016-05-28 12:26:18 +0200 | [diff] [blame] | 16 | return (float) x*y*z; |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 17 | } |
| 18 | |
Wenzel Jakob | 43398a8 | 2015-07-28 16:12:20 +0200 | [diff] [blame] | 19 | std::complex<double> my_func3(std::complex<double> c) { |
| 20 | return c * std::complex<double>(2.f); |
| 21 | } |
| 22 | |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 23 | test_initializer numpy_vectorize([](py::module &m) { |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 24 | // Vectorize all arguments of a function (though non-vector arguments are also allowed) |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 25 | m.def("vectorized_func", py::vectorize(my_func)); |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 26 | |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 27 | // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the vectorization) |
| 28 | m.def("vectorized_func2", |
Wenzel Jakob | b50872a | 2015-10-13 17:38:22 +0200 | [diff] [blame] | 29 | [](py::array_t<int> x, py::array_t<float> y, float z) { |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 30 | return py::vectorize([z](int x, float y) { return my_func(x, y, z); })(x, y); |
| 31 | } |
| 32 | ); |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 33 | |
| 34 | // Vectorize a complex-valued function |
Wenzel Jakob | 43398a8 | 2015-07-28 16:12:20 +0200 | [diff] [blame] | 35 | m.def("vectorized_func3", py::vectorize(my_func3)); |
Wenzel Jakob | b47a9de | 2016-05-19 16:02:09 +0200 | [diff] [blame] | 36 | |
| 37 | /// Numpy function which only accepts specific data types |
Dean Moldovan | 665e880 | 2016-08-12 22:28:31 +0200 | [diff] [blame] | 38 | m.def("selective_func", [](py::array_t<int, py::array::c_style>) { return "Int branch taken."; }); |
| 39 | m.def("selective_func", [](py::array_t<float, py::array::c_style>) { return "Float branch taken."; }); |
| 40 | m.def("selective_func", [](py::array_t<std::complex<float>, py::array::c_style>) { return "Complex float branch taken."; }); |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 41 | |
| 42 | |
| 43 | // Internal optimization test for whether the input is trivially broadcastable: |
Jason Rhinelander | b0292c1 | 2017-03-18 21:11:59 -0300 | [diff] [blame] | 44 | py::enum_<py::detail::broadcast_trivial>(m, "trivial") |
| 45 | .value("f_trivial", py::detail::broadcast_trivial::f_trivial) |
| 46 | .value("c_trivial", py::detail::broadcast_trivial::c_trivial) |
| 47 | .value("non_trivial", py::detail::broadcast_trivial::non_trivial); |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 48 | m.def("vectorized_is_trivial", []( |
| 49 | py::array_t<int, py::array::forcecast> arg1, |
| 50 | py::array_t<float, py::array::forcecast> arg2, |
| 51 | py::array_t<double, py::array::forcecast> arg3 |
| 52 | ) { |
Cris Luengo | 30d43c4 | 2017-04-14 14:33:44 -0600 | [diff] [blame^] | 53 | ssize_t ndim; |
| 54 | std::vector<ssize_t> shape; |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 55 | std::array<py::buffer_info, 3> buffers {{ arg1.request(), arg2.request(), arg3.request() }}; |
| 56 | return py::detail::broadcast(buffers, ndim, shape); |
| 57 | }); |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 58 | }); |