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)
+