Support streaming lazy masked expressions

Author: Alex-PLACET

include/xtensor/core/xmath.hpp

@@ -21,6 +21,7 @@
 #include <type_traits>
 
 #include <xtl/xcomplex.hpp>
+#include <xtl/xmasked_value.hpp>
 #include <xtl/xsequence.hpp>
 #include <xtl/xtype_traits.hpp>
 
@@ -569,13 +570,64 @@ namespace xt
 
     namespace math
     {
+        namespace detail
+        {
+            template <typename T>
+            constexpr decltype(auto) masked_data(const T& value) noexcept
+            {
+                return value;
+            }
+
+            template <typename T, typename B>
+            constexpr decltype(auto) masked_data(const xtl::xmasked_value<T, B>& value) noexcept
+            {
+                return value.value();
+            }
+
+            template <typename T>
+            constexpr bool masked_visible(const T&) noexcept
+            {
+                return true;
+            }
+
+            template <typename T, typename B>
+            constexpr bool masked_visible(const xtl::xmasked_value<T, B>& value) noexcept
+            {
+                return static_cast<bool>(value.visible());
+            }
+        }
+
         template <class T = void>
         struct minimum
         {
             template <class A1, class A2>
             constexpr auto operator()(const A1& t1, const A2& t2) const noexcept
             {
-                return xtl::select(t1 < t2, t1, t2);
+                if constexpr (xtl::is_xmasked_value<std::decay_t<A1>>::value || xtl::is_xmasked_value<std::decay_t<A2>>::value)
+                {
+                    using value_type = xtl::promote_type_t<
+                        std::decay_t<decltype(detail::masked_data(t1))>,
+                        std::decay_t<decltype(detail::masked_data(t2))>>;
+                    using return_type = xtl::xmasked_value<value_type, bool>;
+
+                    if (detail::masked_visible(t1) && detail::masked_visible(t2))
+                    {
+                        return return_type(
+                            static_cast<value_type>(
+                                detail::masked_data(t1) < detail::masked_data(t2)
+                                    ? detail::masked_data(t1)
+                                    : detail::masked_data(t2)
+                            ),
+                            true
+                        );
+                    }
+
+                    return return_type(value_type(0), false);
+                }
+                else
+                {
+                    return xtl::select(t1 < t2, t1, t2);
+                }
             }
 
             template <class A1, class A2>
@@ -591,7 +643,31 @@ namespace xt
             template <class A1, class A2>
             constexpr auto operator()(const A1& t1, const A2& t2) const noexcept
             {
-                return xtl::select(t1 > t2, t1, t2);
+                if constexpr (xtl::is_xmasked_value<std::decay_t<A1>>::value || xtl::is_xmasked_value<std::decay_t<A2>>::value)
+                {
+                    using value_type = xtl::promote_type_t<
+                        std::decay_t<decltype(detail::masked_data(t1))>,
+                        std::decay_t<decltype(detail::masked_data(t2))>>;
+                    using return_type = xtl::xmasked_value<value_type, bool>;
+
+                    if (detail::masked_visible(t1) && detail::masked_visible(t2))
+                    {
+                        return return_type(
+                            static_cast<value_type>(
+                                detail::masked_data(t1) > detail::masked_data(t2)
+                                    ? detail::masked_data(t1)
+                                    : detail::masked_data(t2)
+                            ),
+                            true
+                        );
+                    }
+
+                    return return_type(value_type(0), false);
+                }
+                else
+                {
+                    return xtl::select(t1 > t2, t1, t2);
+                }
             }
 
             template <class A1, class A2>

include/xtensor/views/xmasked_view.hpp

@@ -118,14 +118,15 @@ namespace xt
         using bool_load_type = xtl::xmasked_value<typename data_type::bool_load_type, mask_type>;
 
         using shape_type = typename data_type::shape_type;
-        using strides_type = typename data_type::strides_type;
+        using strides_type = xtl::mpl::eval_if_t<has_strides<data_type>, detail::expr_strides_type<data_type>, get_strides_type<shape_type>>;
+        using backstrides_type = xtl::mpl::eval_if_t<has_strides<data_type>, detail::expr_backstrides_type<data_type>, get_strides_type<shape_type>>;
 
         static constexpr layout_type static_layout = data_type::static_layout;
         static constexpr bool contiguous_layout = false;
 
         using inner_shape_type = typename data_type::inner_shape_type;
-        using inner_strides_type = typename data_type::inner_strides_type;
-        using inner_backstrides_type = typename data_type::inner_backstrides_type;
+        using inner_strides_type = xtl::mpl::eval_if_t<has_strides<data_type>, detail::expr_inner_strides_type<data_type>, get_strides_type<shape_type>>;
+        using inner_backstrides_type = xtl::mpl::eval_if_t<has_strides<data_type>, detail::expr_inner_backstrides_type<data_type>, get_strides_type<shape_type>>;
 
         using expression_tag = xtensor_expression_tag;
 
@@ -163,7 +164,12 @@ namespace xt
 
         size_type size() const noexcept;
         const inner_shape_type& shape() const noexcept;
+        template <typename DT = data_type>
+            requires has_strides<DT>::value
         const inner_strides_type& strides() const noexcept;
+
+        template <typename DT = data_type>
+            requires has_strides<DT>::value
         const inner_backstrides_type& backstrides() const noexcept;
         using accessible_base::dimension;
         using accessible_base::shape;
@@ -202,6 +208,9 @@ namespace xt
         template <class S>
         bool has_linear_assign(const S& strides) const noexcept;
 
+        template <typename S>
+        bool broadcast_shape(S& shape, bool reuse_cache = false) const;
+
         data_type& value() noexcept;
         const data_type& value() const noexcept;
 
@@ -338,6 +347,8 @@ namespace xt
      * Returns the strides of the xmasked_view.
      */
     template <class CTD, class CTM>
+    template <typename DT>
+        requires has_strides<DT>::value
     inline auto xmasked_view<CTD, CTM>::strides() const noexcept -> const inner_strides_type&
     {
         return m_data.strides();
@@ -347,6 +358,8 @@ namespace xt
      * Returns the backstrides of the xmasked_view.
      */
     template <class CTD, class CTM>
+    template <typename DT>
+        requires has_strides<DT>::value
     inline auto xmasked_view<CTD, CTM>::backstrides() const noexcept -> const inner_backstrides_type&
     {
         return m_data.backstrides();
@@ -370,6 +383,13 @@ namespace xt
         return false;
     }
 
+    template <typename CTD, typename CTM>
+    template <typename S>
+    inline bool xmasked_view<CTD, CTM>::broadcast_shape(S& shape, bool) const
+    {
+        return xt::broadcast_shape(m_data.shape(), shape);
+    }
+
     /**
      * Fills the data with the given value.
      * @param value the value to fill the data with.

test/test_xmasked_view.cpp

@@ -9,6 +9,7 @@
 
 #include <sstream>
 
+#include "xtensor/core/xmath.hpp"
 #include "xtensor/io/xio.hpp"
 #include "xtensor/optional/xoptional_assembly.hpp"
 #include "xtensor/views/xmasked_view.hpp"
@@ -40,6 +41,42 @@ namespace xt
         return masked_view(data, std::move(mask));
     }
 
+    template <class A, class B, class M, class = void>
+    struct is_masked_minimum_streamable : std::false_type
+    {
+    };
+
+    template <class A, class B, class M>
+    struct is_masked_minimum_streamable<
+        A,
+        B,
+        M,
+        std::void_t<decltype(
+            std::declval<std::ostream&>()
+            << minimum(
+                masked_view(std::declval<const A&>(), std::declval<const M&>()),
+                masked_view(std::declval<const B&>(), std::declval<const M&>())
+            ))>> : std::true_type
+    {
+    };
+
+    template <class A, class B, class M, class = void>
+    struct is_masked_view_of_minimum_streamable : std::false_type
+    {
+    };
+
+    template <class A, class B, class M>
+    struct is_masked_view_of_minimum_streamable<
+        A,
+        B,
+        M,
+        std::void_t<decltype(
+            std::declval<std::ostream&>()
+            << masked_view(minimum(std::declval<const A&>(), std::declval<const B&>()), std::declval<const M&>()))>>
+        : std::true_type
+    {
+    };
+
     TEST(xmasked_view, dimension)
     {
         auto data = make_test_data();
@@ -226,6 +263,27 @@ namespace xt
                                      " {masked,      5, masked}}";
         EXPECT_EQ(out.str(), expected);
     }
+    
+    TEST(xmasked_view, lazy_expression_stream)
+    {
+        using array_type = xarray<double>;
+        const array_type a = {1., 1., 1., 1.};
+        const array_type b = {0.1, 0.7, 0.3, 0.9};
+
+        const auto mask = b < 0.5;
+        const auto expected = eval(masked_view(minimum(a, b), mask));
+
+        std::stringstream expected_out;
+        expected_out << expected;
+
+        std::stringstream masked_min_out;
+        masked_min_out << minimum(masked_view(a, mask), masked_view(b, mask));
+        EXPECT_EQ(masked_min_out.str(), expected_out.str());
+
+        std::stringstream masked_expr_out;
+        masked_expr_out << masked_view(minimum(a, b), mask);
+        EXPECT_EQ(masked_expr_out.str(), expected_out.str());
+    }
 
     TEST(xmasked_view, assign)
     {
@@ -240,6 +298,21 @@ namespace xt
         EXPECT_EQ(data, expected1);
     }
 
+    TEST(xmasked_view, assign_const_masked_view_rhs)
+    {
+        xarray<double> data = {{1., -2., 3.}, {4., 5., -6.}, {7., 8., -9.}};
+        const xarray<double> data2 = {{0.1, 0.2, 0.3}, {0.4, 0.5, 0.6}, {0.7, 0.8, 0.9}};
+        xarray<bool> mask = {{true, true, true}, {true, false, false}, {true, false, true}};
+
+        auto masked_data = masked_view(data, mask);
+        const auto masked_data2 = masked_view(data2, mask);
+
+        masked_data = masked_data2;
+
+        xarray<double> expected = {{0.1, 0.2, 0.3}, {0.4, 5., -6.}, {0.7, 8., 0.9}};
+        EXPECT_EQ(data, expected);
+    }
+
     TEST(xmasked_view, view)
     {
         xt::xarray<size_t> data = {{0, 1}, {2, 3}, {4, 5}};