xtensor-stack · JohanMabille · Feb 28, 2017 · Feb 24, 2017 · Feb 27, 2017
diff --git a/benchmark/CMakeLists.txt b/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})
 
diff --git a/benchmark/benchmark_pybind_vectorize.py b/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))
diff --git a/benchmark/benchmark_pyvectorize.py b/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))
diff --git a/benchmark/main.cpp b/benchmark/main.cpp
@@ -6,8 +6,9 @@
 #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 +48,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();
 }
diff --git a/include/xtensor-python/pybuffer_adaptor.hpp b/include/xtensor-python/pybuffer_adaptor.hpp
@@ -280,6 +280,12 @@ namespace xt
 
         inline self_type operator+(difference_type n) const { return self_type(p_current + n); }
         inline self_type operator-(difference_type n) const { return self_type(p_current - n); }
+        inline self_type operator-(const self_type& rhs) const
+        {
+            self_type tmp(*this);
+            tmp -= (p_current - rhs.p_current);
+            return tmp;
+        }
 
         pointer get_pointer() const { return p_current; }
 

diff --git a/include/xtensor-python/pycontainer.hpp b/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
         {

diff --git a/include/xtensor-python/pyvectorize.hpp b/include/xtensor-python/pyvectorize.hpp
@@ -26,7 +26,7 @@ namespace xt
         {
         }
 
-        pybind11::object operator()(const pyarray<Args>&... args)
+        inline pyarray<R> operator()(const pyarray<Args>&... args)
         {
             pyarray<R> res = m_vectorizer(args...);
             return res;

diff --git a/test/main.cpp b/test/main.cpp
@@ -7,6 +7,7 @@
 #include <numeric>
 
 namespace py = pybind11;
+using complex_t = std::complex<double>;
 
 // Examples
 
@@ -58,5 +59,10 @@ PYBIND11_PLUGIN(xtensor_python_test)
 
     m.def("vectorize_example1", xt::pyvectorize(add), "");
 
+    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));
+    });
+
     return m.ptr();
 }
diff --git a/test/setup.py b/test/setup.py
@@ -107,7 +107,7 @@ def build_extensions(self):
     description='An example project using xtensor-python',
     long_description='',
     ext_modules=ext_modules,
-    install_requires=['pybind11==2.0.1'],
+    install_requires=['pybind11>=2.0.1'],
     cmdclass={'build_ext': BuildExt},
     zip_safe=False,
 )
diff --git a/test/test_pyarray.py b/test/test_pyarray.py
@@ -50,3 +50,9 @@ def test_readme_example2(self):
              [-1.499227,  0.136731,  1.646979,  1.643002,  0.128456],
              [-1.084323, -0.583843,  0.45342 ,  1.073811,  0.706945]], 1e-5)
 
+    def test_rect_to_polar(self):
+        print("test6")
+        x = np.ones(10, dtype=complex)
+        z = xt.rect_to_polar(x[::2]);
+        np.testing.assert_allclose(z, (np.ones(5, dtype=float), np.zeros(5, dtype=float)), 1e-5)
+