Reductions.hpp

Go to the documentation of this file.

 
// TODO: Avoid the initial owning Tensor copy when reducing strided views.
 
#pragma once
 
#include <stratax/core/dtypes/Concepts.hpp>
#include <stratax/exceptions/Exceptions.hpp>
#include <stratax/core/Shape.hpp>
#include <stratax/containers/Tensor.hpp>
#include <stratax/containers/Matrix.hpp>
#include <stratax/containers/Vector.hpp>
#include <stratax/core/Slice.hpp>
#include <stratax/indexing/Slicing.hpp>
#include <stratax/algorithms/Conversion.hpp>
#include <stratax/core/ReductionTraits.hpp>
 
#include <numeric>
#include <functional>
#include <cmath>
#include <type_traits>
#include <utility>
 
namespace reduction {
 
namespace detail {
 
template<bool FindMaximum, Array A>
requires Ordered<typename A::value_type>
std::pair<typename A::value_type, std::size_t> ordered_extreme(const A& arr)
{
    auto current = arr.begin();
    const auto end = arr.end();
    typename A::value_type selected = *current;
    std::size_t selected_index = 0;
    std::size_t current_index = 1;
 
    // Preserve unordered floating values explicitly: standard extrema algorithms,
    // including MSVC's vectorized path, can assume NaN-free inputs.
    for (++current; current != end; ++current, ++current_index)
    {
        const typename A::value_type value = *current;
        const bool replace = FindMaximum ? selected < value : value < selected;
        if (replace)
        {
            selected = value;
            selected_index = current_index;
        }
    }
 
    return {selected, selected_index};
}
 
inline bool advance(
    const stratax::core::Shape& shape,
    std::vector<std::size_t>& indices)
{
    for (std::size_t d = shape.rank(); d-- > 0;)
    {
        ++indices[d];
 
        if (indices[d] < shape[d])
        {
            return true;
        }
 
        indices[d] = 0;
    }
 
    return false;
}
 
template<Array A>
inline int normalize_axis(const A& arr, int axis)
{
    int init = axis;
    if (axis < 0)
    {
        init += arr.rank();
    }
 
    return init;
}
 
template<Array A>
inline std::vector<std::size_t> result_shape(const A& arr, int axis, bool keepdims)
{
    stratax::core::Shape input_shape = arr.shape();
 
    std::vector<std::size_t> result_dimensions;
    
    if (keepdims)
    {
        for (std::size_t dimension = 0; dimension < input_shape.rank(); ++dimension)
        {
            if (dimension == static_cast<std::size_t>(axis))
            {
                result_dimensions.push_back(1);
            }
 
            else
            {
                result_dimensions.push_back(input_shape[dimension]);
            }
        }
    }
 
    else 
    {
        for (std::size_t dimension = 0; dimension < input_shape.rank(); ++dimension)
        {
            if (dimension != static_cast<std::size_t>(axis))
            {
                result_dimensions.push_back(input_shape[dimension]);
            }
        }
    }
 
    return result_dimensions;
}
 
template<Array A, typename Func>
using axis_reduce_value_t =
    decltype(std::declval<Func>()(
        std::declval<const stratax::container::Tensor<typename A::value_type>&>()));
 
} // namespace detail
 
template<Array A, typename Func>
stratax::container::Tensor<detail::axis_reduce_value_t<A, Func>>
axis_reduce(const A& array, int axis, Func func, bool keepdims = false)
{
    using ResultType = detail::axis_reduce_value_t<A, Func>;
 
    int Axis = detail::normalize_axis(array, axis);
    
    if (Axis < 0 || Axis >= static_cast<int>(array.rank()))
    {
        throw Exceptions::AxisError("Axis is out of range.");
    }
 
    stratax::container::Tensor<typename A::value_type> arr =
        stratax::conversion::to_tensor(array);
 
    const stratax::core::Shape input_shape = arr.shape();
 
    std::vector<std::size_t> result_dims =
        detail::result_shape(array, Axis, keepdims);
 
    // A zero-dimensional tensor cannot store values in the current API.
    // Represent scalar reductions as a single-element tensor.
    if (result_dims.empty())
    {
        ResultType scalar_result = func(arr);
        return stratax::container::Tensor<ResultType>(stratax::core::Shape{1}, scalar_result);
    }
 
    stratax::container::Tensor<ResultType> result(stratax::core::Shape{result_dims});
    if (result.empty())
    {
        return result;
    }
 
    std::vector<std::size_t> output_index(result.rank(), 0);
    std::vector<stratax::core::Slice> slices;
    
    do {
        slices.clear();
        std::size_t output_position = 0;
        for (std::size_t dimension = 0; dimension < input_shape.rank(); ++dimension)
        {
            if (static_cast<int>(dimension) == Axis)
            {
                slices.push_back(stratax::core::Slice{static_cast<std::ptrdiff_t>(0), 
                static_cast<std::ptrdiff_t>(input_shape[dimension])});
            }
 
            else
            {
                const std::size_t index = keepdims
                    ? output_index[dimension]
                    : output_index[output_position++];
 
                slices.push_back(stratax::core::Slice{
                    static_cast<std::ptrdiff_t>(index),
                    static_cast<std::ptrdiff_t>(index + 1)
                });
            }
        }
        auto s = stratax::indexing::slice(arr, slices);
        ResultType value = func(s);
        result(output_index) = value;
    }
        while (detail::advance(result.shape(), output_index));
 
    return result;
}
 
// Global Reductions
 
template<Array A>
requires Numeric<typename A::value_type>
auto sum(const A& arr)
{
    using result_type =
        reduction_sum_t<typename A::value_type>;
 
    return std::accumulate(
        arr.begin(),
        arr.end(),
        result_type{0});
}
 
template<Array A>
requires Numeric<typename A::value_type>
auto prod(const A& arr)
{
    using result_type =
        reduction_prod_t<typename A::value_type>;
 
    return std::accumulate(
        arr.begin(),
        arr.end(),
        result_type{1},
        std::multiplies<result_type>{});
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto max(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::IndexError("Cannot find the maximum of an empty array.");
    }
 
    return detail::ordered_extreme<true>(arr).first;
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto min(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::IndexError("Cannot find the minimum of an empty array.");
    }
 
    return detail::ordered_extreme<false>(arr).first;
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmax(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::IndexError("Cannot find argmax of an empty array.");
    }
 
    return static_cast<stratax::dtype::int64>(detail::ordered_extreme<true>(arr).second);
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmin(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::IndexError("Cannot find argmin of an empty array.");
    }
 
    return static_cast<stratax::dtype::int64>(detail::ordered_extreme<false>(arr).second);
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
double mean(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::ZeroDivisionError("Cannot compute the mean of an empty array.");
    }
 
    long double total = 0.0L;
    for (const auto& value : arr)
    {
        total += static_cast<long double>(value);
    }
 
    return static_cast<double>(total / static_cast<long double>(arr.size()));
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
double var(const A& arr)
{
    if (arr.empty())
    {
        throw Exceptions::ZeroDivisionError("Cannot compute the variance of an empty array.");
    }
 
    double count = 0.0;
    double mean_value = 0.0;
    double m2 = 0.0;
 
    for (const auto& value : arr)
    {
        count += 1.0;
        const double delta = static_cast<double>(value) - mean_value;
        mean_value += delta / count;
        const double delta2 = static_cast<double>(value) - mean_value;
        m2 += delta * delta2;
    }
 
    return m2 / count;
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
double std(const A& arr)
{
    auto vars = var(arr);
    return std::sqrt(vars);
}
 
 
// Axis Reductions
 
template<Array A>
requires Numeric<typename A::value_type>
auto sum(const A& arr, int axis)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) {
            return reduction::sum(s);
        });
}
 
template<Array A>
requires Numeric<typename A::value_type>
auto sum(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) {
            return reduction::sum(s);
        },
        keepdims);
}
 
template<Array A>
requires Numeric<typename A::value_type>
auto prod(const A& arr, int axis)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) {
            return reduction::prod(s);
        });
}
 
template<Array A>
requires Numeric<typename A::value_type>
auto prod(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) {
            return reduction::prod(s);
        },
        keepdims);
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto max(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::max(s); });
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto max(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::max(s); },
        keepdims);
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto min(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::min(s); });
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto min(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::min(s); },
        keepdims);
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmax(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::argmax(s); });
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmax(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::argmax(s); },
        keepdims);
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmin(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::argmin(s); });
}
 
template<Array A>
requires Ordered<typename A::value_type>
auto argmin(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::argmin(s); },
        keepdims);
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
mean(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::mean(s); },
        keepdims);
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
mean(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::mean(s); });
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
var(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::var(s); },
        keepdims);
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
var(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::var(s); });
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
std(const A& arr, int axis, bool keepdims)
{
    return axis_reduce(
        arr,
        axis,
        [](const auto& s) { return reduction::std(s); },
        keepdims);
}
 
template<Array A>
requires (
    Numeric<typename A::value_type> &&
    Ordered<typename A::value_type>
)
stratax::container::Tensor<double>
std(const A& arr, int axis)
{
    return axis_reduce(
        arr, axis, [](const auto& s) { return reduction::std(s); });
}
 
} // namespace reduction