| 20 | } |
| 21 | |
| 22 | TEST_SUBMODULE(numpy_vectorize, m) { |
| 23 | try { |
| 24 | py::module_::import("numpy"); |
| 25 | } catch (const py::error_already_set &) { |
| 26 | return; |
| 27 | } |
| 28 | |
| 29 | // test_vectorize, test_docs, test_array_collapse |
| 30 | // Vectorize all arguments of a function (though non-vector arguments are also allowed) |
| 31 | m.def("vectorized_func", py::vectorize(my_func)); |
| 32 | |
| 33 | // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the |
| 34 | // vectorization) |
| 35 | m.def("vectorized_func2", [](py::array_t<int> x, py::array_t<float> y, float z) { |
| 36 | return py::vectorize([z](int x, float y) { return my_func(x, y, z); })(std::move(x), |
| 37 | std::move(y)); |
| 38 | }); |
| 39 | |
| 40 | // Vectorize a complex-valued function |
| 41 | m.def("vectorized_func3", |
| 42 | py::vectorize([](std::complex<double> c) { return c * std::complex<double>(2.f); })); |
| 43 | |
| 44 | // test_type_selection |
| 45 | // NumPy function which only accepts specific data types |
| 46 | // A lot of these no lints could be replaced with const refs, and probably should at some |
| 47 | // point. |
| 48 | m.def("selective_func", |
| 49 | [](const py::array_t<int, py::array::c_style> &) { return "Int branch taken."; }); |
| 50 | m.def("selective_func", |
| 51 | [](const py::array_t<float, py::array::c_style> &) { return "Float branch taken."; }); |
| 52 | m.def("selective_func", [](const py::array_t<std::complex<float>, py::array::c_style> &) { |
| 53 | return "Complex float branch taken."; |
| 54 | }); |
| 55 | |
| 56 | // test_passthrough_arguments |
| 57 | // Passthrough test: references and non-pod types should be automatically passed through (in |
| 58 | // the function definition below, only `b`, `d`, and `g` are vectorized): |
| 59 | struct NonPODClass { |
| 60 | explicit NonPODClass(int v) : value{v} {} |
| 61 | int value; |
| 62 | }; |
| 63 | py::class_<NonPODClass>(m, "NonPODClass") |
| 64 | .def(py::init<int>()) |
| 65 | .def_readwrite("value", &NonPODClass::value); |
| 66 | m.def("vec_passthrough", |
| 67 | py::vectorize([](const double *a, |
| 68 | double b, |
| 69 | // Changing this broke things |
| 70 | // NOLINTNEXTLINE(performance-unnecessary-value-param) |
| 71 | py::array_t<double> c, |
| 72 | const int &d, |
| 73 | int &e, |
| 74 | NonPODClass f, |
| 75 | const double g) { return *a + b + c.at(0) + d + e + f.value + g; })); |
| 76 | |
| 77 | // test_method_vectorization |
| 78 | struct VectorizeTestClass { |
| 79 | explicit VectorizeTestClass(int v) : value{v} {}; |