Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 1 | /* |
Wenzel Jakob | a576e6a | 2015-07-29 17:51:54 +0200 | [diff] [blame] | 2 | example/example10.cpp -- auto-vectorize functions over NumPy array |
| 3 | arguments |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 4 | |
| 5 | Copyright (c) 2015 Wenzel Jakob <wenzel@inf.ethz.ch> |
| 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 | |
| 11 | #include "example.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) { |
| 15 | std::cout << "my_func(x:int=" << x << ", y:float=" << y << ", z:float=" << z << ")" << std::endl; |
| 16 | return x*y*z; |
| 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 | |
Wenzel Jakob | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 23 | void init_ex10(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 | d4258ba | 2015-07-26 16:33:49 +0200 | [diff] [blame] | 36 | } |