xhistogram.hpp

/***************************************************************************
 * Copyright (c) Johan Mabille, Sylvain Corlay and Wolf Vollprecht          *
 * Copyright (c) QuantStack                                                 *
 *                                                                          *
 * Distributed under the terms of the BSD 3-Clause License.                 *
 *                                                                          *
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/
 
#ifndef XTENSOR_HISTOGRAM_HPP
#define XTENSOR_HISTOGRAM_HPP
 
#include "xset_operation.hpp"
#include "xsort.hpp"
#include "xtensor.hpp"
#include "xview.hpp"
 
using namespace xt::placeholders;
 
namespace xt
{
    template <class E1, class E2>
    inline auto digitize(E1&& data, E2&& bin_edges, bool right = false)
    {
        XTENSOR_ASSERT(bin_edges.dimension() == 1);
        XTENSOR_ASSERT(bin_edges.size() >= 2);
        XTENSOR_ASSERT(std::is_sorted(bin_edges.cbegin(), bin_edges.cend()));
        XTENSOR_ASSERT(xt::amin(data)[0] >= bin_edges[0]);
        XTENSOR_ASSERT(xt::amax(data)[0] <= bin_edges[bin_edges.size() - 1]);
 
        return xt::searchsorted(std::forward<E2>(bin_edges), std::forward<E1>(data), right);
    }
 
    namespace detail
    {
        template <class R = double, class E1, class E2, class E3>
        inline auto histogram_imp(E1&& data, E2&& bin_edges, E3&& weights, bool density, bool equal_bins)
        {
            using size_type = common_size_type_t<std::decay_t<E1>, std::decay_t<E2>, std::decay_t<E3>>;
            using value_type = typename std::decay_t<E3>::value_type;
 
            XTENSOR_ASSERT(data.dimension() == 1);
            XTENSOR_ASSERT(weights.dimension() == 1);
            XTENSOR_ASSERT(bin_edges.dimension() == 1);
            XTENSOR_ASSERT(weights.size() == data.size());
            XTENSOR_ASSERT(bin_edges.size() >= 2);
            XTENSOR_ASSERT(std::is_sorted(bin_edges.cbegin(), bin_edges.cend()));
 
            size_t n_bins = bin_edges.size() - 1;
            xt::xtensor<value_type, 1> count = xt::zeros<value_type>({n_bins});
 
            if (equal_bins)
            {
                std::array<typename std::decay_t<E2>::value_type, 2> bounds = xt::minmax(bin_edges)();
                auto left = static_cast<double>(bounds[0]);
                auto right = static_cast<double>(bounds[1]);
                double norm = 1. / (right - left);
                for (size_t i = 0; i < data.size(); ++i)
                {
                    auto v = static_cast<double>(data(i));
                    // left and right are not bounds of data
                    if (v >= left && v < right)
                    {
                        auto i_bin = static_cast<size_t>(static_cast<double>(n_bins) * (v - left) * norm);
                        count(i_bin) += weights(i);
                    }
                    else if (v == right)
                    {
                        count(n_bins - 1) += weights(i);
                    }
                }
            }
            else
            {
                auto sorter = xt::argsort(data);
 
                size_type ibin = 0;
 
                for (auto& idx : sorter)
                {
                    auto item = data[idx];
 
                    if (item < bin_edges[0])
                    {
                        continue;
                    }
 
                    if (item > bin_edges[n_bins])
                    {
                        break;
                    }
 
                    while (item >= bin_edges[ibin + 1] && ibin < n_bins - 1)
                    {
                        ++ibin;
                    }
 
                    count[ibin] += weights[idx];
                }
            }
 
            xt::xtensor<R, 1> prob = xt::cast<R>(count);
 
            if (density)
            {
                R n = static_cast<R>(data.size());
                for (size_type i = 0; i < bin_edges.size() - 1; ++i)
                {
                    prob[i] /= (static_cast<R>(bin_edges[i + 1] - bin_edges[i]) * n);
                }
            }
 
            return prob;
        }
 
    }  // detail
 
    template <class R = double, class E1, class E2, class E3>
    inline auto histogram(E1&& data, E2&& bin_edges, E3&& weights, bool density = false)
    {
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            std::forward<E2>(bin_edges),
            std::forward<E3>(weights),
            density,
            false
        );
    }
 
    template <class R = double, class E1, class E2>
    inline auto histogram(E1&& data, E2&& bin_edges, bool density = false)
    {
        using value_type = typename std::decay_t<E1>::value_type;
 
        auto n = data.size();
 
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            std::forward<E2>(bin_edges),
            xt::ones<value_type>({n}),
            density,
            false
        );
    }
 
    template <class R = double, class E1>
    inline auto histogram(E1&& data, std::size_t bins = 10, bool density = false)
    {
        using value_type = typename std::decay_t<E1>::value_type;
 
        auto n = data.size();
 
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            histogram_bin_edges(data, xt::ones<value_type>({n}), bins),
            xt::ones<value_type>({n}),
            density,
            true
        );
    }
 
    template <class R = double, class E1, class E2>
    inline auto histogram(E1&& data, std::size_t bins, E2 left, E2 right, bool density = false)
    {
        using value_type = typename std::decay_t<E1>::value_type;
 
        auto n = data.size();
 
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            histogram_bin_edges(data, left, right, bins),
            xt::ones<value_type>({n}),
            density,
            true
        );
    }
 
    template <class R = double, class E1, class E2>
    inline auto histogram(E1&& data, std::size_t bins, E2&& weights, bool density = false)
    {
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            histogram_bin_edges(data, weights, bins),
            std::forward<E2>(weights),
            density,
            true
        );
    }
 
    template <class R = double, class E1, class E2, class E3>
    inline auto histogram(E1&& data, std::size_t bins, E2&& weights, E3 left, E3 right, bool density = false)
    {
        return detail::histogram_imp<R>(
            std::forward<E1>(data),
            histogram_bin_edges(data, weights, left, right, bins),
            std::forward<E2>(weights),
            density,
            true
        );
    }
 
    enum class histogram_algorithm
    {
        automatic,
        linspace,
        logspace,
        uniform
    };
 
    template <class E1, class E2, class E3>
    inline auto histogram_bin_edges(
        E1&& data,
        E2&& weights,
        E3 left,
        E3 right,
        std::size_t bins = 10,
        histogram_algorithm mode = histogram_algorithm::automatic
    )
    {
        // counter and return type
        using size_type = common_size_type_t<std::decay_t<E1>, std::decay_t<E2>>;
        using value_type = typename std::decay_t<E1>::value_type;
        using weights_type = typename std::decay_t<E2>::value_type;
 
        // basic checks
        // - rank
        XTENSOR_ASSERT(data.dimension() == 1);
        XTENSOR_ASSERT(weights.dimension() == 1);
        // - size
        XTENSOR_ASSERT(weights.size() == data.size());
        // - bounds
        XTENSOR_ASSERT(left <= right);
        // - non-empty
        XTENSOR_ASSERT(bins > std::size_t(0));
 
        // act on different modes
        switch (mode)
        {
            // bins of equal width
            case histogram_algorithm::automatic:
            {
                xt::xtensor<value_type, 1> bin_edges = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }
 
            // bins of equal width
            case histogram_algorithm::linspace:
            {
                xt::xtensor<value_type, 1> bin_edges = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }
 
            // bins of logarithmically increasing size
            case histogram_algorithm::logspace:
            {
                using rhs_value_type = std::conditional_t<xtl::is_integral<value_type>::value, double, value_type>;
 
                xtensor<value_type, 1> bin_edges = xt::cast<value_type>(
                    xt::logspace<rhs_value_type>(std::log10(left), std::log10(right), bins + 1)
                );
 
                // TODO: replace above with below after 'xsimd' fix
                // xt::xtensor<value_type,1> bin_edges = xt::logspace<value_type>(
                //     std::log10(left), std::log10(right), bins+1);
 
                return bin_edges;
            }
 
            // same amount of data in each bin
            case histogram_algorithm::uniform:
            {
                // indices that sort "data"
                auto sorter = xt::argsort(data);
 
                // histogram: all of equal 'height'
                // - height
                weights_type w = xt::sum<weights_type>(weights)[0] / static_cast<weights_type>(bins);
                // - apply to all bins
                xt::xtensor<weights_type, 1> count = w * xt::ones<weights_type>({bins});
 
                // take cumulative sum, to act as easy look-up
                count = xt::cumsum(count);
 
                // edges
                // - allocate
                std::vector<size_t> shape = {bins + 1};
                xt::xtensor<value_type, 1> bin_edges = xtensor<value_type, 1>::from_shape(shape);
                // - first/last edge
                bin_edges[0] = left;
                bin_edges[bins] = right;
                // - cumulative weight
                weights_type cum_weight = static_cast<weights_type>(0);
                // - current bin
                size_type ibin = 0;
                // - loop to find interior bin-edges
                for (size_type i = 0; i < weights.size(); ++i)
                {
                    if (cum_weight >= count[ibin])
                    {
                        bin_edges[ibin + 1] = data[sorter[i]];
                        ++ibin;
                    }
                    cum_weight += weights[sorter[i]];
                }
                return bin_edges;
            }
 
            // bins of equal width
            default:
            {
                xt::xtensor<value_type, 1> bin_edges = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }
        }
    }
 
    template <class E1, class E2>
    inline auto histogram_bin_edges(
        E1&& data,
        E2&& weights,
        std::size_t bins = 10,
        histogram_algorithm mode = histogram_algorithm::automatic
    )
    {
        using value_type = typename std::decay_t<E1>::value_type;
        std::array<value_type, 2> left_right;
        left_right = xt::minmax(data)();
 
        return histogram_bin_edges(
            std::forward<E1>(data),
            std::forward<E2>(weights),
            left_right[0],
            left_right[1],
            bins,
            mode
        );
    }
 
    template <class E1>
    inline auto
    histogram_bin_edges(E1&& data, std::size_t bins = 10, histogram_algorithm mode = histogram_algorithm::automatic)
    {
        using value_type = typename std::decay_t<E1>::value_type;
 
        auto n = data.size();
        std::array<value_type, 2> left_right;
        left_right = xt::minmax(data)();
 
        return histogram_bin_edges(
            std::forward<E1>(data),
            xt::ones<value_type>({n}),
            left_right[0],
            left_right[1],
            bins,
            mode
        );
    }
 
    template <class E1, class E2>
    inline auto histogram_bin_edges(
        E1&& data,
        E2 left,
        E2 right,
        std::size_t bins = 10,
        histogram_algorithm mode = histogram_algorithm::automatic
    )
    {
        using value_type = typename std::decay_t<E1>::value_type;
 
        auto n = data.size();
 
        return histogram_bin_edges(std::forward<E1>(data), xt::ones<value_type>({n}), left, right, bins, mode);
    }
 
    template <class E1, class E2, XTL_REQUIRES(is_xexpression<std::decay_t<E2>>)>
    inline auto bincount(E1&& data, E2&& weights, std::size_t minlength = 0)
    {
        using result_value_type = typename std::decay_t<E2>::value_type;
        using input_value_type = typename std::decay_t<E1>::value_type;
        using size_type = typename std::decay_t<E1>::size_type;
 
        static_assert(
            xtl::is_integral<typename std::decay_t<E1>::value_type>::value,
            "Bincount data has to be integral type."
        );
        XTENSOR_ASSERT(data.dimension() == 1);
        XTENSOR_ASSERT(weights.dimension() == 1);
 
        std::array<input_value_type, 2> left_right;
        left_right = xt::minmax(data)();
 
        if (left_right[0] < input_value_type(0))
        {
            XTENSOR_THROW(std::runtime_error, "Data argument for bincount can only contain positive integers!");
        }
 
        xt::xtensor<result_value_type, 1> res = xt::zeros<result_value_type>(
            {(std::max)(minlength, std::size_t(left_right[1] + 1))}
        );
 
        for (size_type i = 0; i < data.size(); ++i)
        {
            res(data(i)) += weights(i);
        }
 
        return res;
    }
 
    template <class E1>
    inline auto bincount(E1&& data, std::size_t minlength = 0)
    {
        return bincount(
            std::forward<E1>(data),
            xt::ones<typename std::decay_t<E1>::value_type>(data.shape()),
            minlength
        );
    }
 
    template <class E>
    inline xt::xtensor<size_t, 1> bin_items(size_t N, E&& weights)
    {
        if (weights.size() <= std::size_t(1))
        {
            xt::xtensor<size_t, 1> n = N * xt::ones<size_t>({1});
            return n;
        }
 
#ifdef XTENSOR_ENABLE_ASSERT
        using value_type = typename std::decay_t<E>::value_type;
 
        XTENSOR_ASSERT(xt::all(weights >= static_cast<value_type>(0)));
        XTENSOR_ASSERT(xt::sum(weights)() > static_cast<value_type>(0));
#endif
 
        xt::xtensor<double, 1> P = xt::cast<double>(weights) / static_cast<double>(xt::sum(weights)());
        xt::xtensor<size_t, 1> n = xt::ceil(static_cast<double>(N) * P);
 
        if (xt::sum(n)() == N)
        {
            return n;
        }
 
        xt::xtensor<size_t, 1> d = xt::zeros<size_t>(P.shape());
        xt::xtensor<size_t, 1> sorter = xt::argsort(P);
        sorter = xt::view(sorter, xt::range(P.size(), _, -1));
        sorter = xt::view(sorter, xt::range(0, xt::sum(n)(0) - N));
        xt::view(d, xt::keep(sorter)) = 1;
        n -= d;
 
        return n;
    }
 
    inline xt::xtensor<size_t, 1> bin_items(size_t N, size_t bins)
    {
        return bin_items(N, xt::ones<double>({bins}));
    }
}
 
#endif