test_numpy_vectorize.cpp revision 12391
11689SN/A/* 28733Sgeoffrey.blake@arm.com tests/test_numpy_vectorize.cpp -- auto-vectorize functions over NumPy array 37783SGiacomo.Gabrielli@arm.com arguments 47783SGiacomo.Gabrielli@arm.com 57783SGiacomo.Gabrielli@arm.com Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> 67783SGiacomo.Gabrielli@arm.com 77783SGiacomo.Gabrielli@arm.com All rights reserved. Use of this source code is governed by a 87783SGiacomo.Gabrielli@arm.com BSD-style license that can be found in the LICENSE file. 97783SGiacomo.Gabrielli@arm.com*/ 107783SGiacomo.Gabrielli@arm.com 117783SGiacomo.Gabrielli@arm.com#include "pybind11_tests.h" 127783SGiacomo.Gabrielli@arm.com#include <pybind11/numpy.h> 137783SGiacomo.Gabrielli@arm.com 142316SN/Adouble my_func(int x, float y, double z) { 151689SN/A py::print("my_func(x:int={}, y:float={:.0f}, z:float={:.0f})"_s.format(x, y, z)); 161689SN/A return (float) x*y*z; 171689SN/A} 181689SN/A 191689SN/ATEST_SUBMODULE(numpy_vectorize, m) { 201689SN/A try { py::module::import("numpy"); } 211689SN/A catch (...) { return; } 221689SN/A 231689SN/A // test_vectorize, test_docs, test_array_collapse 241689SN/A // Vectorize all arguments of a function (though non-vector arguments are also allowed) 251689SN/A m.def("vectorized_func", py::vectorize(my_func)); 261689SN/A 271689SN/A // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the vectorization) 281689SN/A m.def("vectorized_func2", 291689SN/A [](py::array_t<int> x, py::array_t<float> y, float z) { 301689SN/A return py::vectorize([z](int x, float y) { return my_func(x, y, z); })(x, y); 311689SN/A } 321689SN/A ); 331689SN/A 341689SN/A // Vectorize a complex-valued function 351689SN/A m.def("vectorized_func3", py::vectorize( 361689SN/A [](std::complex<double> c) { return c * std::complex<double>(2.f); } 371689SN/A )); 381689SN/A 392665SN/A // test_type_selection 402665SN/A // Numpy function which only accepts specific data types 411689SN/A m.def("selective_func", [](py::array_t<int, py::array::c_style>) { return "Int branch taken."; }); 421061SN/A m.def("selective_func", [](py::array_t<float, py::array::c_style>) { return "Float branch taken."; }); 435953Ssaidi@eecs.umich.edu m.def("selective_func", [](py::array_t<std::complex<float>, py::array::c_style>) { return "Complex float branch taken."; }); 445596Sgblack@eecs.umich.edu 458779Sgblack@eecs.umich.edu 461061SN/A // test_passthrough_arguments 471061SN/A // Passthrough test: references and non-pod types should be automatically passed through (in the 485596Sgblack@eecs.umich.edu // function definition below, only `b`, `d`, and `g` are vectorized): 498502Sgblack@eecs.umich.edu struct NonPODClass { 507720Sgblack@eecs.umich.edu NonPODClass(int v) : value{v} {} 515596Sgblack@eecs.umich.edu int value; 528502Sgblack@eecs.umich.edu }; 534637SN/A py::class_<NonPODClass>(m, "NonPODClass").def(py::init<int>()); 544637SN/A m.def("vec_passthrough", py::vectorize( 554637SN/A [](double *a, double b, py::array_t<double> c, const int &d, int &e, NonPODClass f, const double g) { 564637SN/A return *a + b + c.at(0) + d + e + f.value + g; 574637SN/A } 588502Sgblack@eecs.umich.edu )); 598502Sgblack@eecs.umich.edu 608502Sgblack@eecs.umich.edu // test_method_vectorization 611061SN/A struct VectorizeTestClass { 622292SN/A VectorizeTestClass(int v) : value{v} {}; 632292SN/A float method(int x, float y) { return y + (float) (x + value); } 642292SN/A int value = 0; 652292SN/A }; 662292SN/A py::class_<VectorizeTestClass> vtc(m, "VectorizeTestClass"); 675596Sgblack@eecs.umich.edu vtc .def(py::init<int>()) 682292SN/A .def_readwrite("value", &VectorizeTestClass::value); 691464SN/A 701464SN/A // Automatic vectorizing of methods 711464SN/A vtc.def("method", py::vectorize(&VectorizeTestClass::method)); 722292SN/A 733782SN/A // test_trivial_broadcasting 741464SN/A // Internal optimization test for whether the input is trivially broadcastable: 751464SN/A py::enum_<py::detail::broadcast_trivial>(m, "trivial") 762292SN/A .value("f_trivial", py::detail::broadcast_trivial::f_trivial) 773782SN/A .value("c_trivial", py::detail::broadcast_trivial::c_trivial) 782292SN/A .value("non_trivial", py::detail::broadcast_trivial::non_trivial); 791464SN/A m.def("vectorized_is_trivial", []( 807783SGiacomo.Gabrielli@arm.com py::array_t<int, py::array::forcecast> arg1, 817783SGiacomo.Gabrielli@arm.com py::array_t<float, py::array::forcecast> arg2, 828471SGiacomo.Gabrielli@arm.com py::array_t<double, py::array::forcecast> arg3 838471SGiacomo.Gabrielli@arm.com ) { 848471SGiacomo.Gabrielli@arm.com ssize_t ndim; 858471SGiacomo.Gabrielli@arm.com std::vector<ssize_t> shape; 868471SGiacomo.Gabrielli@arm.com std::array<py::buffer_info, 3> buffers {{ arg1.request(), arg2.request(), arg3.request() }}; 878471SGiacomo.Gabrielli@arm.com return py::detail::broadcast(buffers, ndim, shape); 888471SGiacomo.Gabrielli@arm.com }); 898471SGiacomo.Gabrielli@arm.com} 908471SGiacomo.Gabrielli@arm.com