test_numpy_vectorize.cpp revision 12391:ceeca8b41e4b
111723Sar4jc@virginia.edu/* 211723Sar4jc@virginia.edu tests/test_numpy_vectorize.cpp -- auto-vectorize functions over NumPy array 311723Sar4jc@virginia.edu arguments 411723Sar4jc@virginia.edu 511723Sar4jc@virginia.edu Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch> 611723Sar4jc@virginia.edu 711723Sar4jc@virginia.edu All rights reserved. Use of this source code is governed by a 811723Sar4jc@virginia.edu BSD-style license that can be found in the LICENSE file. 911723Sar4jc@virginia.edu*/ 1011723Sar4jc@virginia.edu 1111723Sar4jc@virginia.edu#include "pybind11_tests.h" 1211723Sar4jc@virginia.edu#include <pybind11/numpy.h> 1311723Sar4jc@virginia.edu 1411723Sar4jc@virginia.edudouble my_func(int x, float y, double z) { 1511723Sar4jc@virginia.edu py::print("my_func(x:int={}, y:float={:.0f}, z:float={:.0f})"_s.format(x, y, z)); 1611723Sar4jc@virginia.edu return (float) x*y*z; 1711723Sar4jc@virginia.edu} 1811723Sar4jc@virginia.edu 1911723Sar4jc@virginia.eduTEST_SUBMODULE(numpy_vectorize, m) { 2011723Sar4jc@virginia.edu try { py::module::import("numpy"); } 2111723Sar4jc@virginia.edu catch (...) { return; } 2211723Sar4jc@virginia.edu 2311723Sar4jc@virginia.edu // test_vectorize, test_docs, test_array_collapse 2411723Sar4jc@virginia.edu // Vectorize all arguments of a function (though non-vector arguments are also allowed) 2511723Sar4jc@virginia.edu m.def("vectorized_func", py::vectorize(my_func)); 2611723Sar4jc@virginia.edu 2711723Sar4jc@virginia.edu // Vectorize a lambda function with a capture object (e.g. to exclude some arguments from the vectorization) 2811723Sar4jc@virginia.edu m.def("vectorized_func2", 2911723Sar4jc@virginia.edu [](py::array_t<int> x, py::array_t<float> y, float z) { 3011723Sar4jc@virginia.edu return py::vectorize([z](int x, float y) { return my_func(x, y, z); })(x, y); 3111723Sar4jc@virginia.edu } 3211723Sar4jc@virginia.edu ); 3311723Sar4jc@virginia.edu 3411723Sar4jc@virginia.edu // Vectorize a complex-valued function 3511723Sar4jc@virginia.edu m.def("vectorized_func3", py::vectorize( 3611723Sar4jc@virginia.edu [](std::complex<double> c) { return c * std::complex<double>(2.f); } 3711723Sar4jc@virginia.edu )); 3811723Sar4jc@virginia.edu 3911723Sar4jc@virginia.edu // test_type_selection 4011723Sar4jc@virginia.edu // Numpy function which only accepts specific data types 4111723Sar4jc@virginia.edu m.def("selective_func", [](py::array_t<int, py::array::c_style>) { return "Int branch taken."; }); 4211723Sar4jc@virginia.edu m.def("selective_func", [](py::array_t<float, py::array::c_style>) { return "Float branch taken."; }); 4311723Sar4jc@virginia.edu m.def("selective_func", [](py::array_t<std::complex<float>, py::array::c_style>) { return "Complex float branch taken."; }); 4411723Sar4jc@virginia.edu 4511723Sar4jc@virginia.edu 4611723Sar4jc@virginia.edu // test_passthrough_arguments 4711723Sar4jc@virginia.edu // Passthrough test: references and non-pod types should be automatically passed through (in the 4811723Sar4jc@virginia.edu // function definition below, only `b`, `d`, and `g` are vectorized): 4911723Sar4jc@virginia.edu struct NonPODClass { 5011723Sar4jc@virginia.edu NonPODClass(int v) : value{v} {} 5111723Sar4jc@virginia.edu int value; 5211723Sar4jc@virginia.edu }; 5311723Sar4jc@virginia.edu py::class_<NonPODClass>(m, "NonPODClass").def(py::init<int>()); 5411723Sar4jc@virginia.edu m.def("vec_passthrough", py::vectorize( 5511723Sar4jc@virginia.edu [](double *a, double b, py::array_t<double> c, const int &d, int &e, NonPODClass f, const double g) { 5611723Sar4jc@virginia.edu return *a + b + c.at(0) + d + e + f.value + g; 5711723Sar4jc@virginia.edu } 5811723Sar4jc@virginia.edu )); 5911723Sar4jc@virginia.edu 6011723Sar4jc@virginia.edu // test_method_vectorization 6111723Sar4jc@virginia.edu struct VectorizeTestClass { 6211723Sar4jc@virginia.edu VectorizeTestClass(int v) : value{v} {}; 6311723Sar4jc@virginia.edu float method(int x, float y) { return y + (float) (x + value); } 6411723Sar4jc@virginia.edu int value = 0; 6511723Sar4jc@virginia.edu }; 6611723Sar4jc@virginia.edu py::class_<VectorizeTestClass> vtc(m, "VectorizeTestClass"); 6711723Sar4jc@virginia.edu vtc .def(py::init<int>()) 6811723Sar4jc@virginia.edu .def_readwrite("value", &VectorizeTestClass::value); 6911723Sar4jc@virginia.edu 7011723Sar4jc@virginia.edu // Automatic vectorizing of methods 7111723Sar4jc@virginia.edu vtc.def("method", py::vectorize(&VectorizeTestClass::method)); 7211723Sar4jc@virginia.edu 7311723Sar4jc@virginia.edu // test_trivial_broadcasting 7411723Sar4jc@virginia.edu // Internal optimization test for whether the input is trivially broadcastable: 7511723Sar4jc@virginia.edu py::enum_<py::detail::broadcast_trivial>(m, "trivial") 7611723Sar4jc@virginia.edu .value("f_trivial", py::detail::broadcast_trivial::f_trivial) 7711723Sar4jc@virginia.edu .value("c_trivial", py::detail::broadcast_trivial::c_trivial) 7811723Sar4jc@virginia.edu .value("non_trivial", py::detail::broadcast_trivial::non_trivial); 7911723Sar4jc@virginia.edu m.def("vectorized_is_trivial", []( 8011723Sar4jc@virginia.edu py::array_t<int, py::array::forcecast> arg1, 8111723Sar4jc@virginia.edu py::array_t<float, py::array::forcecast> arg2, 8211723Sar4jc@virginia.edu py::array_t<double, py::array::forcecast> arg3 8311723Sar4jc@virginia.edu ) { 8411723Sar4jc@virginia.edu ssize_t ndim; 8511723Sar4jc@virginia.edu std::vector<ssize_t> shape; 8611723Sar4jc@virginia.edu std::array<py::buffer_info, 3> buffers {{ arg1.request(), arg2.request(), arg3.request() }}; 8711723Sar4jc@virginia.edu return py::detail::broadcast(buffers, ndim, shape); 8811723Sar4jc@virginia.edu }); 8911723Sar4jc@virginia.edu} 9011723Sar4jc@virginia.edu