vectorize implementation

Author: JohanMabille

include/xtensor-python/pyarray.hpp

@@ -114,8 +114,8 @@ namespace xt
                          handle base = handle());
 
         size_type dimension() const;
-        shape_type shape() const;
-        strides_type strides() const;
+        const shape_type& shape() const;
+        const strides_type& strides() const;
         backstrides_type backstrides() const;
 
         void reshape(const shape_type& shape);
@@ -182,6 +182,9 @@ namespace xt
 
         static PyObject *ensure_(PyObject* ptr);
 
+        mutable shape_type m_shape;
+        mutable strides_type m_strides;
+
     };
 
     /**************************************
@@ -198,7 +201,7 @@ namespace xt
     inline auto pyarray_backstrides<A>::operator[](size_type i) const -> value_type
     {
         value_type sh = p_a->shape()[i];
-        value_type res = sh == 1 ? 0 : sh * p_a->strides()[i] / sizeof(typename A::value_type);
+        value_type res = sh == 1 ? 0 : (sh - 1) * p_a->strides()[i] / sizeof(typename A::value_type);
         return  res;
     }
 
@@ -250,21 +253,21 @@ namespace xt
     }
 
     template <class T, int ExtraFlags>
-    inline auto pyarray<T, ExtraFlags>::shape() const -> shape_type
+    inline auto pyarray<T, ExtraFlags>::shape() const -> const shape_type&
     {
         // Until we have the CRTP on shape types, we copy the shape.
-        shape_type shape(dimension());
-        std::copy(pybind_array::shape(), pybind_array::shape() + dimension(), shape.begin());
-        return shape;
+        m_shape.resize(dimension());
+        std::copy(pybind_array::shape(), pybind_array::shape() + dimension(), m_shape.begin());
+        return m_shape;
     }
 
     template <class T, int ExtraFlags>
-    inline auto pyarray<T, ExtraFlags>::strides() const -> strides_type
+    inline auto pyarray<T, ExtraFlags>::strides() const -> const strides_type&
     {
-        strides_type strides(dimension());
-        std::transform(pybind_array::strides(), pybind_array::strides() + dimension(), strides.begin(),
+        m_strides.resize(dimension());
+        std::transform(pybind_array::strides(), pybind_array::strides() + dimension(), m_strides.begin(),
             [](size_type str) { return str / sizeof(value_type); });
-        return strides;
+        return m_strides;
     }
 
     template <class T, int ExtraFlags>

include/xtensor-python/pyvectorize.hpp

@@ -0,0 +1,55 @@
+/***************************************************************************
+* Copyright (c) 2016, Johan Mabille and Sylvain Corlay                     *
+*                                                                          *
+* Distributed under the terms of the BSD 3-Clause License.                 *
+*                                                                          *
+* The full license is in the file LICENSE, distributed with this software. *
+****************************************************************************/
+
+#ifndef PY_VECTORIZE_HPP
+#define PY_VECTORIZE_HPP
+
+#include "pyarray.hpp"
+#include "xtensor/xvectorize.hpp"
+
+namespace xt
+{
+
+    template <class Func, class R, class... Args>
+    struct pyvectorizer
+    {
+        xvectorizer<Func, R> m_vectorizer;
+
+        template <class F>
+        pyvectorizer(F&& func)
+            : m_vectorizer(std::forward<F>(func))
+        {
+        }
+
+        pybind11::object operator()(pyarray<Args, pybind_array::c_style | pybind_array::forcecast>... args)
+        {
+            pyarray<R> res = m_vectorizer(args...);
+            return res;
+        }
+    };
+
+    template <class R, class... Args>
+    inline pyvectorizer<R(*)(Args...), R, Args...> pyvectorize(R(*f) (Args...))
+    {
+        return pyvectorizer<R(*) (Args...), R, Args...>(f);
+    }
+
+    template <class F, class R, class... Args>
+    inline pyvectorizer<F, R, Args...> pyvectorize(F&& f, R(*) (Args...))
+    {
+        return pyvectorizer<F, R, Args...>(std::forward<F>(f));
+    }
+
+    template <class F>
+    inline auto pyvectorize(F&& f) -> decltype(pyvectorize(std::forward<F>(f), (detail::get_function_type<F>*)nullptr))
+    {
+        return pyvectorize(std::forward<F>(f), (detail::get_function_type<F>*)nullptr);
+    }
+}
+
+#endif

test/main.cpp

@@ -1,15 +1,22 @@
 #include <pybind11/pybind11.h>
 #include "xtensor/xarray.hpp"
 #include "xtensor-python/pyarray.hpp"
+#include "xtensor-python/pyvectorize.hpp"
+#include <iostream>
 
 namespace py = pybind11;
 
+int add(int i, int j) {
+    return i + j;
+}
+
 double test0(xt::pyarray<double> &m)
 {
     return m(0);
 }
 
-xt::pyarray<double> test1(xt::pyarray<double> &m) {
+xt::pyarray<double> test1(xt::pyarray<double> &m)
+{
     return m + 2;
 }
 
@@ -19,6 +26,7 @@ PYBIND11_PLUGIN(xtensor_python_test)
 
     m.def("test0", test0, "");
     m.def("test1", test1, "");
+    m.def("vec_add", xt::pyvectorize(add), "");
 
     return m.ptr();
 }

test/test_pyarray.py

@@ -11,8 +11,23 @@
 
 from unittest import TestCase
 import xtensor_python_test as xt
+import numpy as np
 
 class ExampleTest(TestCase):
 
-    def test_example(self):
+    def test_example0(self):
         self.assertEqual(4, xt.test0([4, 5, 6]))
+
+    def test_example1(self):
+        x = np.array([[0., 1.], [2., 3.]])
+        res = np.array([[2., 3.], [4., 5.]])
+        y = xt.test1(x)
+        np.testing.assert_allclose(y, res, 1e-12)
+
+    def test_vectorize(self):
+        x1 = np.array([[0, 1], [2, 3]])
+        x2 = np.array([0, 1])
+        res = np.array([[0, 2], [2, 4]])
+        y = xt.vec_add(x1, x2)
+        np.testing.assert_array_equal(y, res)
+