Eye function

include/xtensor/xbuilder.hpp

@@ -14,6 +14,7 @@
 #define XBUILDER_HPP
 
 #include <utility>
+#include <functional>
 #include <cmath>
 
 #include "xfunction.hpp"
@@ -133,6 +134,103 @@ namespace xt
                 return m_start;
             }
         };
+
+        template <class F>
+        struct fn_impl
+        {
+            using value_type = typename F::value_type;
+            using size_type = std::size_t;
+
+            fn_impl(F&& f) : m_ft(f)
+            {
+            }
+
+            inline value_type operator()() const
+            {
+                // special case when called without args (happens when printing)
+                return value_type();
+            }
+
+            template <class... Args>
+            inline value_type operator()(Args... args) const
+            {
+                size_type idx [sizeof...(Args)] = {static_cast<size_type>(args)...};
+                return access_impl(std::begin(idx), std::end(idx));
+            }
+
+            inline value_type operator[](const xindex& idx) const
+            {
+                return access_impl(idx.begin(), idx.end());
+            }
+
+            template <class It>
+            inline value_type element(It first, It last) const
+            {
+                return access_impl(first, last);
+            }
+
+        private:
+            F m_ft;
+            template <class It>
+            inline value_type access_impl(const It& begin, const It& end) const
+            {
+                return m_ft(begin, end);
+            }
+        };
+
+        template <class T>
+        struct eye_fn
+        {
+            using value_type = T;
+
+            eye_fn(int k) : m_k(k)
+            {
+            }
+
+            template <class It>
+            inline T operator()(const It& /*begin*/, const It& end) const
+            {
+                // workaround windows compile error by using temporary 
+                // iterators and operator-=
+                auto end_1 = end;
+                auto end_2 = end;
+                end_1 -= 1;
+                end_2 -= 2;
+                return *(end_1) == *(end_2) + m_k ? T(1) : T(0);
+            }
+
+        private:
+            int m_k;
+        };
+    }
+
+    /**
+     * @function eye(const std::vector<std::size_t>& shape, int k = 0)
+     * @brief generate array with ones on the diagonal
+     *
+     * @param shape shape of the resulting expression
+     * @param k index of the diagonal. 0 (default) refers to the main diagonal,
+     *          a positive value refers to an upper diagonal, and a negative
+     *          value to a lower diagonal.
+     *
+     * @tparam T value_type of xexpression
+     *
+     * @return xgenerator that generates the values on access
+     */
+    template <class T = bool>
+    inline auto eye(const std::vector<std::size_t>& shape, int k = 0)
+    {
+        return detail::make_xgenerator(detail::fn_impl<detail::eye_fn<T>>(detail::eye_fn<T>(k)), shape);
+    }
+
+    /**
+     * @function eye(std::size_t n, int k = 0)
+     * @brief like eye with a shape of n x n
+     */
+    template <class T = bool>
+    inline auto eye(std::size_t n, int k = 0)
+    {
+        return eye<T>({n, n}, k);
     }
 
     /**

test/test_xbuilder.cpp

@@ -144,4 +144,22 @@ namespace xt
         ASSERT_EQ(m_assigned(3), 1000);
     }
 
+    TEST(xbuilder, eye)
+    {
+        auto e = eye(5);
+        ASSERT_EQ(2, e.dimension());
+        shape_t expected_shape = {5, 5};
+        ASSERT_EQ(expected_shape, e.shape());
+
+        ASSERT_EQ(1, e(1, 1));
+        xindex idx({1, 0});
+        ASSERT_EQ(0, e[idx]);
+
+        xarray<bool> m_assigned = e;
+        ASSERT_EQ(1, m_assigned(2, 2));
+        ASSERT_EQ(0, m_assigned(4, 2));
+
+        xt::xindex idx2({2, 2});
+        ASSERT_EQ(1, e.element(idx2.begin(), idx2.end()));
+    }
 }