diff --git a/docs/source/api/xsort.rst b/docs/source/api/xsort.rst index 082532d52..f88a57e01 100644 --- a/docs/source/api/xsort.rst +++ b/docs/source/api/xsort.rst @@ -9,6 +9,33 @@ xsort Defined in ``xtensor/misc/xsort.hpp`` +Medians over multiple axes +-------------------------- + +Pass a sequence of axes to ``xt::median`` to group their elements before taking +one median per output element. The remaining axes retain their original order. +For example:: + + xt::xarray a = {{{0., 0., 0., 100.}}, + {{10., 10., 10., 10.}}}; + auto result = xt::median(a, {0, 2}); // shape {1}, value {10.} + +Taking successive medians instead gives ``{5.}``. Materialize the intermediate +result before calling the single-axis overload again:: + + xt::xarray intermediate = xt::median(a, 2); + auto successive = xt::median(intermediate, 0); // shape {1}, value {5.} + +The grouped result is independent of the order of the selected axes. Negative +axes count from the end; repeated or invalid axes are rejected. +An empty sequence keeps the input shape, while selecting every axis returns a +scalar expression. Selecting an axis of length zero is an error. + +The input is not modified. This overload follows the result-type rules of the +single-axis overload, including its promotion of integral input. This promotion +also applies to an empty axis sequence and can round large integers. Calling +``xt::median(a)`` without axes retains the existing scalar API and result type. + .. cpp:namespace-push:: xt .. doxygengroup:: xt_xsort diff --git a/include/xtensor/misc/xsort.hpp b/include/xtensor/misc/xsort.hpp index 619bbb6de..5d2251b6b 100644 --- a/include/xtensor/misc/xsort.hpp +++ b/include/xtensor/misc/xsort.hpp @@ -12,7 +12,11 @@ #include #include +#include #include +#include +#include +#include #include #include @@ -1141,6 +1145,120 @@ namespace xt } } + /** + * Find the median over a group of axes. + * + * All elements along the selected axes form one group for each index of the + * remaining axes. The median is computed once per group, not successively + * along each axis. The order of the axes does not affect the result. + * The input is not modified. The result has the same value type as the + * single-axis overload, including its promotion of integral input. + * + * @ingroup xt_xsort + * @param e input xexpression + * @param axes sequence of distinct axes, with negative axes counted from the end. + * An empty sequence retains the shape, with the same value-type + * promotion as a single-axis median (which can round large integers). + * @return expression with the selected axes removed, in the original order + * of the remaining axes. Selecting all axes gives a scalar expression. + * @throw std::runtime_error if an axis is invalid or repeated, or a selected + * axis has length zero. + */ + template + requires std::ranges::range + inline auto median(E&& e, const X& axes) + { + dynamic_shape reduced_axes; + for (auto axis : axes) + { + using axis_type = std::decay_t; + static_assert(xtl::is_integral::value, "Median axes must be integral."); + // Keep the full axis width until after validating its range. The + // axis type can be wider than size_t, or narrower than the rank. + using unsigned_axis_type = std::make_unsigned_t< + std::conditional_t::value, unsigned char, axis_type>>; + using comparison_type = std::common_type_t; + auto dimension = static_cast(e.dimension()); + auto ax = static_cast(axis); + if constexpr (std::is_signed::value) + { + if (axis < 0) + { + // Avoid negating the minimum signed value. + auto from_end = static_cast(-(axis + 1)); + if (from_end >= dimension) + { + XTENSOR_THROW(std::runtime_error, "Median axis is out of bounds."); + } + ax = dimension - from_end - 1; + } + } + if (ax >= dimension) + { + XTENSOR_THROW(std::runtime_error, "Median axis is out of bounds."); + } + reduced_axes.push_back(static_cast(ax)); + } + std::sort(reduced_axes.begin(), reduced_axes.end()); + if (std::adjacent_find(reduced_axes.begin(), reduced_axes.end()) != reduced_axes.end()) + { + XTENSOR_THROW(std::runtime_error, "Median axes should not contain duplicates."); + } + + dynamic_shape permutation; + dynamic_shape grouped_shape; + std::size_t group_size = 1; + for (std::size_t ax = 0; ax < e.dimension(); ++ax) + { + if (std::binary_search(reduced_axes.begin(), reduced_axes.end(), ax)) + { + auto extent = e.shape()[ax]; + if (extent == 0) + { + XTENSOR_THROW(std::runtime_error, "Cannot take a median over an empty axis."); + } + if (extent > (std::numeric_limits::max)() / group_size) + { + XTENSOR_THROW(std::runtime_error, "Median group is too large."); + } + group_size *= extent; + } + else + { + permutation.push_back(ax); + grouped_shape.push_back(e.shape()[ax]); + } + } + permutation.insert(permutation.end(), reduced_axes.begin(), reduced_axes.end()); + grouped_shape.push_back(group_size); + + // Materialize logical groups in row-major order, even for strided input + // or an arbitrary input layout. The existing single-axis median owns its + // partitioned copy and preserves its usual result-type rules. + using buffer_type = xarray::value_type, layout_type::row_major>; + buffer_type grouped; + if constexpr (strided_expression>) + { + grouped = transpose(std::forward(e), permutation); + } + else + { + // Lazy functions and broadcasts do not expose strides to transpose. + // Evaluate them before permuting, without adding this copy for views + // and containers that already have a usable stride interface. + buffer_type evaluated = std::forward(e); + grouped = transpose(evaluated, permutation); + } + grouped.reshape(grouped_shape); + return median(std::move(grouped), static_cast(grouped_shape.size() - 1)); + } + + template + inline auto median(E&& e, std::initializer_list axes) + { + return median(std::forward(e), dynamic_shape(axes)); + } + namespace detail { template diff --git a/test/test_xsort.cpp b/test/test_xsort.cpp index c8f9dc143..6ea1ba97a 100644 --- a/test/test_xsort.cpp +++ b/test/test_xsort.cpp @@ -7,6 +7,13 @@ * The full license is in the file LICENSE, distributed with this software. * ****************************************************************************/ +#include +#include +#include +#include +#include +#include + #include "xtensor/containers/xadapt.hpp" #include "xtensor/containers/xarray.hpp" #include "xtensor/containers/xfixed.hpp" @@ -23,6 +30,72 @@ namespace xt { + namespace + { + struct median_axes_range + { + std::array axes; + }; + + const int* begin(const median_axes_range& range) + { + return range.axes.data(); + } + + const int* end(const median_axes_range& range) + { + return range.axes.data() + range.axes.size(); + } + + // Enumerate logical coordinates and fully sort each output lane. This + // oracle does not transpose, reshape, partition, or call xt::median. + template + xarray + grouped_median_reference(const E& e, const std::vector& axes) + { + std::vector shape; + std::size_t lane_count = 1; + for (std::size_t ax = 0; ax < e.dimension(); ++ax) + { + if (std::find(axes.begin(), axes.end(), ax) == axes.end()) + { + shape.push_back(e.shape()[ax]); + lane_count *= e.shape()[ax]; + } + } + std::vector> lanes(lane_count); + std::vector index(e.dimension()); + for (std::size_t linear = 0; linear < e.size(); ++linear) + { + auto remainder = linear; + for (std::size_t ax = e.dimension(); ax != 0; --ax) + { + index[ax - 1] = remainder % e.shape()[ax - 1]; + remainder /= e.shape()[ax - 1]; + } + std::size_t lane = 0; + for (std::size_t ax = 0; ax < e.dimension(); ++ax) + { + if (std::find(axes.begin(), axes.end(), ax) == axes.end()) + { + lane = lane * e.shape()[ax] + index[ax]; + } + } + lanes[lane].push_back(static_cast(e.element(index.begin(), index.end()))); + } + auto result = xarray::from_shape(shape); + for (std::size_t lane = 0; lane < lane_count; ++lane) + { + auto& values = lanes[lane]; + std::sort(values.begin(), values.end()); + auto middle = values.size() / 2; + result.flat(lane) = values.size() % 2 == 0 ? (values[middle - 1] + values[middle]) / 2. + : values[middle]; + } + return result; + } + } + TEST(xsort, argsort) { xarray a1 = {2, 3, 1}; @@ -564,4 +637,155 @@ namespace xt EXPECT_EQ(ma0, ma0_exp); EXPECT_EQ(ma1, ma1_exp); } + + TEST(xsort, median_grouped_axes) + { + const xtensor a = { + {{0., 0., 0., 100.}, {1., 2., 99., 100.}, {-9., -8., -7., -6.}}, + {{10., 10., 10., 10.}, {3., 4., 5., 6.}, {-5., 30., 40., 50.}} + }; + const auto original = a; + const xarray expected = {10., 4.5, -5.5}; + EXPECT_EQ(grouped_median_reference(a, {0, 2}), expected); + EXPECT_EQ(median(a, {0, 2}), expected); + EXPECT_EQ(median(a, {2, 0}), expected); + EXPECT_EQ(median(a, {-1, -3}), expected); + const int array_axes[] = {0, 2}; + EXPECT_EQ(median(a, array_axes), expected); + const std::initializer_list list_axes = {0, 2}; + EXPECT_EQ(median(a, list_axes), expected); + xarray last_axis = median(a, 2); + xarray first_axis = median(a, 0); + EXPECT_NE(median(last_axis, 0), expected); + EXPECT_NE(median(first_axis, 1), expected); + + const std::vector axes = {0, 2}; + const std::array reversed = {2, 0}; + const int raw_axes[] = {0, 2}; + EXPECT_EQ(median(a, axes), expected); + EXPECT_EQ(median(a, reversed), expected); + EXPECT_EQ(median(a, raw_axes), expected); + EXPECT_EQ(median(a, {0, 1}), grouped_median_reference(a, {0, 1})); + EXPECT_EQ(median(a, {1, 2}), grouped_median_reference(a, {1, 2})); + auto all_axes = median(a, {2, 0, 1}); + EXPECT_EQ(all_axes.dimension(), 0); + EXPECT_EQ(all_axes(), median(a)); + EXPECT_EQ(a, original); + } + + TEST(xsort, median_grouped_axes_layouts_and_views) + { + xtensor a = xtensor::from_shape({3, 2, 3}); + for (std::size_t i = 0; i < a.size(); ++i) + { + a.flat(i) = static_cast((i * 11) % 29) - 14; + } + const auto original = a; + const auto expected = grouped_median_reference(a, {0, 2}); + xarray column_major = a; + EXPECT_EQ(median(a, {0, 2}), expected); + EXPECT_EQ(median(column_major, {2, 0}), expected); + auto lazy = a + 2; + static_assert(!strided_expression); + EXPECT_EQ(median(lazy, {0, 2}), expected + 2); + EXPECT_EQ(median(a + 2, {0, 2}), expected + 2); + + auto transposed = transpose(a, {2, 0, 1}); + EXPECT_EQ(median(transposed, {0, 1}), expected); + auto strided = view(a, range(0, 3, 2), all(), range(0, 3, 2)); + EXPECT_EQ(median(strided, {0, 2}), grouped_median_reference(strided, {0, 2})); + auto reversed = view(a, all(), all(), range(placeholders::_, placeholders::_, -1)); + EXPECT_EQ(median(reversed, {0, 2}), expected); + auto broadcasted = broadcast(a, {2, 3, 2, 3}); + static_assert(!strided_expression); + EXPECT_EQ(median(broadcasted, {0, 1, 3}), grouped_median_reference(broadcasted, {0, 1, 3})); + EXPECT_EQ(median(broadcast(a, {2, 3, 2, 3}), {3, 1, 0}), grouped_median_reference(broadcasted, {0, 1, 3})); + EXPECT_EQ(a, original); + EXPECT_EQ(column_major, original); + } + + TEST(xsort, median_grouped_axes_types_and_legacy) + { + xarray integers = {{1, 2}, {5, 8}}; + auto grouped = median(integers, {0, 1}); + auto single = median(integers, 0); + static_assert(std::is_same::value, ""); + EXPECT_EQ(grouped(), 3.5); + EXPECT_EQ(median(integers), 3); // The existing scalar API keeps its input type. + EXPECT_EQ(median(integers, {0}), single); + EXPECT_EQ(median(integers, {1}), median(integers, 1)); + EXPECT_EQ(median(integers, {-1}), median(integers, -1)); + + enum axis + { + rows = 0 + }; + + EXPECT_EQ(median(integers, rows), single); + EXPECT_EQ(median(integers, std::integral_constant{}), median(integers, 1)); + static_assert(std::ranges::range); + EXPECT_EQ(median(integers, median_axes_range{{0, 1}}), grouped); + EXPECT_EQ(median(integers, std::vector{}), integers); + + xarray floats = {{1.f, 2.f}, {5.f, 8.f}}; + auto grouped_float = median(floats, {0, 1}); + auto single_float = median(floats, 0); + static_assert( + std::is_same::value, + "" + ); + EXPECT_EQ(grouped_float(), 3.5f); + EXPECT_EQ(median(floats, std::array{}), floats); + xarray large_integers = {9007199254740993ULL}; + auto empty_axes = median(large_integers, std::vector{}); + auto one_axis = median(large_integers, 0); + static_assert(std::is_same::value, ""); + static_assert(std::is_same::value, ""); + EXPECT_EQ(empty_axes.dimension(), 1); + EXPECT_EQ(empty_axes.shape()[0], 1); + EXPECT_EQ(empty_axes(0), 9007199254740992.); + EXPECT_EQ(one_axis(), empty_axes(0)); + EXPECT_EQ(large_integers(0), 9007199254740993ULL); + xarray scalar = 7.; + auto scalar_result = median(scalar, std::vector{}); + EXPECT_EQ(scalar_result.dimension(), 0); + EXPECT_EQ(scalar_result(), 7.); + + auto empty_lanes = xarray::from_shape({0, 3, 2}); + auto empty_result = median(empty_lanes, {1, 2}); + EXPECT_EQ(empty_result.dimension(), 1); + EXPECT_EQ(empty_result.shape()[0], 0); + } + + TEST(xsort, median_grouped_axes_invalid) + { + xarray a = {{1., 2.}, {3., 4.}}; + EXPECT_THROW(median(a, {0, 0}), std::runtime_error); + EXPECT_THROW(median(a, {0, -2}), std::runtime_error); + EXPECT_THROW(median(a, {2}), std::runtime_error); + EXPECT_THROW(median(a, {-3}), std::runtime_error); + EXPECT_THROW(median(a, {std::numeric_limits::min()}), std::runtime_error); + EXPECT_THROW(median(a, {std::numeric_limits::max()}), std::runtime_error); + EXPECT_THROW(median(a, {std::numeric_limits::min()}), std::runtime_error); + if constexpr (sizeof(std::size_t) < sizeof(long long)) + { + EXPECT_THROW(median(a, {4294967296LL}), std::runtime_error); + EXPECT_THROW(median(a, {-4294967297LL}), std::runtime_error); + EXPECT_THROW(median(a, {4294967296ULL}), std::runtime_error); + } + EXPECT_EQ(median(a, std::array{-1}), (xarray{1.5, 3.5})); + EXPECT_EQ(median(a, std::array{0}), (xarray{2., 3.})); + EXPECT_EQ(median(a, std::array{true}), (xarray{1.5, 3.5})); + EXPECT_THROW(median(a, {std::numeric_limits::min()}), std::runtime_error); + auto high_rank = xarray::from_shape(std::vector(130, 1)); + high_rank.fill(7.); + xarray narrow_axis = median(high_rank, {std::numeric_limits::min()}); + EXPECT_EQ(narrow_axis.dimension(), 129); + EXPECT_EQ(narrow_axis.size(), 1); + EXPECT_EQ(narrow_axis.flat(0), 7.); + auto empty = xarray::from_shape({2, 0, 3}); + EXPECT_THROW(median(empty, {0, 1}), std::runtime_error); + xarray scalar = 1.; + EXPECT_THROW(median(scalar, {0}), std::runtime_error); + } }