Merge pull request #30 from JohanMabille/complex

Complex

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,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();
 }

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; }
 

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
         {

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;

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();
 }

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

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