test_numpy_vectorize.cpp
Go to the documentation of this file.
1 /*
2  tests/test_numpy_vectorize.cpp -- auto-vectorize functions over NumPy array
3  arguments
4 
5  Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.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 <pybind11/numpy.h>
12 
13 #include "pybind11_tests.h"
14 
15 #include <utility>
16 
17 double my_func(int x, float y, double z) {
18  py::print("my_func(x:int={}, y:float={:.0f}, z:float={:.0f})"_s.format(x, y, z));
19  return (float) x * y * z;
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} {};
80  float method(int x, float y) const { return y + (float) (x + value); }
81  int value = 0;
82  };
83  py::class_<VectorizeTestClass> vtc(m, "VectorizeTestClass");
84  vtc.def(py::init<int>()).def_readwrite("value", &VectorizeTestClass::value);
85 
86  // Automatic vectorizing of methods
87  vtc.def("method", py::vectorize(&VectorizeTestClass::method));
88 
89  // test_trivial_broadcasting
90  // Internal optimization test for whether the input is trivially broadcastable:
91  py::enum_<py::detail::broadcast_trivial>(m, "trivial")
92  .value("f_trivial", py::detail::broadcast_trivial::f_trivial)
93  .value("c_trivial", py::detail::broadcast_trivial::c_trivial)
94  .value("non_trivial", py::detail::broadcast_trivial::non_trivial);
95  m.def("vectorized_is_trivial",
96  [](const py::array_t<int, py::array::forcecast> &arg1,
97  const py::array_t<float, py::array::forcecast> &arg2,
98  const py::array_t<double, py::array::forcecast> &arg3) {
99  py::ssize_t ndim = 0;
100  std::vector<py::ssize_t> shape;
101  std::array<py::buffer_info, 3> buffers{
102  {arg1.request(), arg2.request(), arg3.request()}};
103  return py::detail::broadcast(buffers, ndim, shape);
104  });
105 
106  m.def("add_to", py::vectorize([](NonPODClass &x, int a) { x.value += a; }));
107 }
Matrix3f m
double my_func(int x, float y, double z)
Scalar * y
Scalar * b
Definition: benchVecAdd.cpp:17
Scalar Scalar * c
Definition: benchVecAdd.cpp:17
EIGEN_STRONG_INLINE Packet4f print(const Packet4f &a)
void g(const string &key, int i)
Definition: testBTree.cpp:41
detail::vectorize_helper< Return(*)(Args...), Return, Args... > vectorize(Return(*f)(Args...))
Definition: numpy.h:1945
Array< int, Dynamic, 1 > v
Point2(* f)(const Point3 &, OptionalJacobian< 2, 3 >)
Array< double, 1, 3 > e(1./3., 0.5, 2.)
TEST_SUBMODULE(numpy_vectorize, m)
set noclip points set clip one set noclip two set bar set border lt lw set xdata set ydata set zdata set x2data set y2data set boxwidth set dummy x
broadcast_trivial broadcast(const std::array< buffer_info, N > &buffers, ssize_t &ndim, std::vector< ssize_t > &shape)
Definition: numpy.h:1647


gtsam
Author(s):
autogenerated on Tue Jul 4 2023 02:37:46