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 | f3ce00e | 2017-03-26 00:51:40 -0300 | [diff] [blame] | 23 | struct VectorizeTestClass { |
| 24 | VectorizeTestClass(int v) : value{v} {}; |
| 25 | float method(int x, float y) { return y + (float) (x + value); } |
| 26 | int value = 0; |
| 27 | }; |
| 28 | |
| 29 | struct NonPODClass { |
| 30 | NonPODClass(int v) : value{v} {} |
| 31 | int value; |
| 32 | }; |
| 33 | |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 34 | test_initializer numpy_vectorize([](py::module &m) { |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 35 | // 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] | 36 | m.def("vectorized_func", py::vectorize(my_func)); |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 37 | |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 38 | // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the vectorization) |
| 39 | m.def("vectorized_func2", |
Wenzel Jakob | b50872a | 2015-10-13 17:38:22 +0200 | [diff] [blame] | 40 | [](py::array_t<int> x, py::array_t<float> y, float z) { |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 41 | return py::vectorize([z](int x, float y) { return my_func(x, y, z); })(x, y); |
| 42 | } |
| 43 | ); |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 44 | |
| 45 | // Vectorize a complex-valued function |
Wenzel Jakob | 43398a8 | 2015-07-28 16:12:20 +0200 | [diff] [blame] | 46 | m.def("vectorized_func3", py::vectorize(my_func3)); |
Wenzel Jakob | b47a9de | 2016-05-19 16:02:09 +0200 | [diff] [blame] | 47 | |
| 48 | /// Numpy function which only accepts specific data types |
Dean Moldovan | 665e880 | 2016-08-12 22:28:31 +0200 | [diff] [blame] | 49 | m.def("selective_func", [](py::array_t<int, py::array::c_style>) { return "Int branch taken."; }); |
| 50 | m.def("selective_func", [](py::array_t<float, py::array::c_style>) { return "Float branch taken."; }); |
| 51 | 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] | 52 | |
| 53 | |
Jason Rhinelander | f3ce00e | 2017-03-26 00:51:40 -0300 | [diff] [blame] | 54 | // Passthrough test: references and non-pod types should be automatically passed through (in the |
| 55 | // function definition below, only `b`, `d`, and `g` are vectorized): |
| 56 | py::class_<NonPODClass>(m, "NonPODClass").def(py::init<int>()); |
| 57 | m.def("vec_passthrough", py::vectorize( |
| 58 | [](double *a, double b, py::array_t<double> c, const int &d, int &e, NonPODClass f, const double g) { |
| 59 | return *a + b + c.at(0) + d + e + f.value + g; |
| 60 | } |
| 61 | )); |
| 62 | |
| 63 | py::class_<VectorizeTestClass> vtc(m, "VectorizeTestClass"); |
| 64 | vtc .def(py::init<int>()) |
| 65 | .def_readwrite("value", &VectorizeTestClass::value); |
| 66 | |
| 67 | // Automatic vectorizing of methods |
| 68 | vtc.def("method", py::vectorize(&VectorizeTestClass::method)); |
| 69 | |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 70 | // Internal optimization test for whether the input is trivially broadcastable: |
Jason Rhinelander | b0292c1 | 2017-03-18 21:11:59 -0300 | [diff] [blame] | 71 | py::enum_<py::detail::broadcast_trivial>(m, "trivial") |
| 72 | .value("f_trivial", py::detail::broadcast_trivial::f_trivial) |
| 73 | .value("c_trivial", py::detail::broadcast_trivial::c_trivial) |
| 74 | .value("non_trivial", py::detail::broadcast_trivial::non_trivial); |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 75 | m.def("vectorized_is_trivial", []( |
| 76 | py::array_t<int, py::array::forcecast> arg1, |
| 77 | py::array_t<float, py::array::forcecast> arg2, |
| 78 | py::array_t<double, py::array::forcecast> arg3 |
| 79 | ) { |
Cris Luengo | 30d43c4 | 2017-04-14 14:33:44 -0600 | [diff] [blame] | 80 | ssize_t ndim; |
| 81 | std::vector<ssize_t> shape; |
Jason Rhinelander | ae5a8f7 | 2017-03-15 00:57:56 -0300 | [diff] [blame] | 82 | std::array<py::buffer_info, 3> buffers {{ arg1.request(), arg2.request(), arg3.request() }}; |
| 83 | return py::detail::broadcast(buffers, ndim, shape); |
| 84 | }); |
Jason Rhinelander | 52f4be8 | 2016-09-03 14:54:22 -0400 | [diff] [blame] | 85 | }); |