complex benchmark

Author: JohanMabille

benchmark/CMakeLists.txt

@@ -78,6 +78,8 @@ endif()
 configure_file(benchmark_pyarray.py benchmark_pyarray.py COPYONLY)
 configure_file(benchmark_pytensor.py benchmark_pytensor.py COPYONLY)
 configure_file(benchmark_pybind_array.py benchmark_pybind_array.py COPYONLY)
+configure_file(benchmark_pyvectorize.py benchmark_pyvectorize.py COPYONLY)
+configure_file(benchmark_pybind_vectorize.py benchmark_pybind_vectorize.py COPYONLY)
 
 add_custom_target(xbenchmark DEPENDS ${XTENSOR_PYTHON_BENCHMARK_TARGET})
 

benchmark/benchmark_pybind_vectorize.py

@@ -0,0 +1,6 @@
+from benchmark_xtensor_python import pybind_rect_to_polar
+import numpy as np
+
+from timeit import timeit
+w = np.ones(100000, dtype=complex)
+print (timeit('pybind_rect_to_polar(w[::2])', 'from __main__ import w, pybind_rect_to_polar', number=1000))

benchmark/benchmark_pyvectorize.py

@@ -0,0 +1,6 @@
+from benchmark_xtensor_python import rect_to_polar
+import numpy as np
+
+from timeit import timeit
+w = np.ones(100000, dtype=complex)
+print (timeit('rect_to_polar(w[::2])', 'from __main__ import w, rect_to_polar', number=1000))

benchmark/main.cpp

@@ -6,8 +6,10 @@
 #include "xtensor/xarray.hpp"
 #include "xtensor-python/pyarray.hpp"
 #include "xtensor-python/pytensor.hpp"
+#include "xtensor-python/pyvectorize.hpp"
 
 #include <complex>
+using complex_t = std::complex<double>;
 
 namespace py = pybind11;
 
@@ -47,5 +49,18 @@ PYBIND11_PLUGIN(benchmark_xtensor_python)
         }
     );
 
+    m.def("rect_to_polar", [](xt::pyarray<complex_t> const& a) {
+            return py::make_tuple(xt::pyvectorize([](complex_t x) { return std::abs(x); })(a),
+                                  xt::pyvectorize([](complex_t x) { return std::arg(x); })(a));
+    });
+
+    m.def("pybind_rect_to_polar", [](py::array a) {
+            if (py::isinstance<py::array_t<complex_t>>(a))
+                return py::make_tuple(py::vectorize([](complex_t x) { return std::abs(x); })(a),
+                                      py::vectorize([](complex_t x) { return std::arg(x); })(a));
+            else
+                throw py::type_error("rect_to_polar unhandled type");
+    });
+
     return m.ptr();
 }

include/xtensor-python/pycontainer.hpp

@@ -14,6 +14,7 @@
 #include <cmath>
 #include "pybind11/pybind11.h"
 #include "pybind11/common.h"
+#include "pybind11/complex.h"
 // Because of layout, else xiterator and xtensor_forward are sufficient
 #include "xtensor/xcontainer.hpp"
 
@@ -188,6 +189,13 @@ namespace xt
                     std::is_same<T, double>::value ? 1 : std::is_same<T, long double>::value ? 2 : 0));
         };
 
+        template <class T>
+        struct is_fmt_numeric<std::complex<T>>
+        {
+            static constexpr bool value = true;
+            static constexpr int index = is_fmt_numeric<T>::index + 3;
+        };
+
         template <class T>
         struct numpy_traits
         {