Weights.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <concepts>
#include <expected>
#include <limits>
#include <random>
#include <ranges>
#include <span>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/NumericError.h>
 
namespace Emrys
{

template <typename Scalar>
class LogSumExp
{
public:

    void
    add(Scalar x)
    {
        if (x == -Eigen::NumTraits<Scalar>::infinity())
        {
            return;
        }
        if (x <= m_max)
        {
            m_sum += Eigen::numext::exp(x - m_max);
        }
        else
        {
            m_sum = m_sum * Eigen::numext::exp(m_max - x) + Scalar(1);
            m_max = x;
        }
    }

    template <typename Derived>
    void
    add(Eigen::ArrayBase<Derived> const & xs)
    {
        if constexpr (std::floating_point<Scalar>)
        {
            if (xs.size() == 0)
            {
                return;
            }
            Scalar const max = xs.template maxCoeff<Eigen::PropagateNaN>();
            if (max == -Eigen::NumTraits<Scalar>::infinity())
            {
                return;
            }
            auto range = LogSumExp{};
            range.m_max = max;
            // Each +inf input is the new maximum, so each counts one relative to it.
            range.m_sum = max == Eigen::NumTraits<Scalar>::infinity() ? Scalar((xs == max).count())
                                                                      : Scalar((xs - max).exp().sum());
            add(range);
        }
        else
        {
            for (Eigen::Index i = 0; i < xs.size(); ++i)
            {
                add(Scalar(xs(i)));
            }
        }
    }

    void
    add(LogSumExp const & other)
    {
        if (other.m_sum == Scalar(0))
        {
            return;
        }
        if (m_sum == Scalar(0))
        {
            *this = other;
            return;
        }
        if (other.m_max <= m_max)
        {
            m_sum += other.m_sum * Eigen::numext::exp(other.m_max - m_max);
        }
        else
        {
            m_sum = m_sum * Eigen::numext::exp(m_max - other.m_max) + other.m_sum;
            m_max = other.m_max;
        }
    }

    [[nodiscard]] Scalar
    value() const
    {
        return m_sum == Scalar(0) ? -Eigen::NumTraits<Scalar>::infinity() : m_max + Eigen::numext::log(m_sum);
    }
 
private:
 
    Scalar m_max = -Eigen::NumTraits<Scalar>::infinity();  
    Scalar m_sum = 0;                                      
};

template <typename Scalar, ExecutionPolicy P, typename LogAt>
[[nodiscard]] Scalar
logSumExp(P && policy, Eigen::Index count, LogAt const & logAt)
{
    return internal::reduceRanges(
               std::forward<P>(policy),
               count,
               [&](Eigen::Index begin, Eigen::Index length) {
                   auto sum = LogSumExp<Scalar>{};
                   for (Eigen::Index i = begin; i < begin + length; ++i)
                   {
                       sum.add(Scalar(logAt(i)));
                   }
                   return sum;
               },
               [](LogSumExp<Scalar> a, LogSumExp<Scalar> const & b) {
                   a.add(b);
                   return a;
               }
    ).value();
}

template <typename Scalar, typename LogAt>
[[nodiscard]] Scalar
logSumExp(Eigen::Index count, LogAt const & logAt)
{
    return logSumExp<Scalar>(execution::sequential, count, logAt);
}

template <ExecutionPolicy P, typename Derived>
[[nodiscard]] typename Derived::Scalar
logSumExp(P && policy, Eigen::ArrayBase<Derived> const & logs)
{
    using Scalar = typename Derived::Scalar;
    return internal::reduceRanges(
               std::forward<P>(policy),
               logs.size(),
               [&](Eigen::Index begin, Eigen::Index length) {
                   auto sum = LogSumExp<Scalar>{};
                   sum.add(logs.derived().segment(begin, length));
                   return sum;
               },
               [](LogSumExp<Scalar> a, LogSumExp<Scalar> const & b) {
                   a.add(b);
                   return a;
               }
    ).value();
}

template <typename Derived>
[[nodiscard]] typename Derived::Scalar
logSumExp(Eigen::ArrayBase<Derived> const & logs)
{
    return logSumExp(execution::sequential, logs);
}
 
namespace internal
{
 
template <typename Scalar, ExecutionPolicy P>
[[nodiscard]] std::expected<Scalar, WeightError>
logTotalWeight(P && policy, std::span<Scalar const> logWeights);

template <typename Scalar, std::ranges::input_range R>
[[nodiscard]] std::expected<Scalar, WeightError>
logTotalWeight(R && logWeights)
{
    if constexpr (std::ranges::contiguous_range<R>
                  && std::same_as<std::remove_cv_t<std::ranges::range_value_t<R>>, Scalar>)
    {
        return logTotalWeight<Scalar>(execution::sequential, std::span<Scalar const>{logWeights});
    }
    else
    {
        auto total = LogSumExp<Scalar>{};
        bool empty = true;
        for (Scalar const w : logWeights)
        {
            if (isNaN(w) || w == Eigen::NumTraits<Scalar>::infinity())
            {
                return std::unexpected{WeightError::NonFinite};
            }
            total.add(w);
            empty = false;
        }
        if (empty)
        {
            return std::unexpected{WeightError::Empty};
        }
        Scalar const result = total.value();
        if (result == -Eigen::NumTraits<Scalar>::infinity())
        {
            return std::unexpected{WeightError::ZeroSum};
        }
        return result;
    }
}

template <typename Scalar, ExecutionPolicy P>
[[nodiscard]] std::expected<Scalar, WeightError>
logTotalWeight(P && policy, std::span<Scalar const> logWeights)
{
    // A range's sum, and whether a weight in it is NaN or +inf.
    struct Total
    {
        LogSumExp<Scalar> sum;   
        bool nonFinite = false;  
    };
 
    if (logWeights.empty())
    {
        return std::unexpected{WeightError::Empty};
    }
    Total const total = reduceRanges(
        std::forward<P>(policy),
        Eigen::Index(logWeights.size()),
        [&](Eigen::Index begin, Eigen::Index length) {
            auto range = Total{};
            auto const weights = logWeights.subspan(std::size_t(begin), std::size_t(length));
            if constexpr (std::floating_point<Scalar>)
            {
                auto const map = Eigen::Map<Eigen::Array<Scalar, Eigen::Dynamic, 1> const>{weights.data(), length};
                // NaN and +inf are the weights that are not below +inf.
                range.nonFinite = !(map < Eigen::NumTraits<Scalar>::infinity()).all();
                if (!range.nonFinite)
                {
                    range.sum.add(map);
                }
            }
            else
            {
                for (Scalar const w : weights)
                {
                    if (isNaN(w) || w == Eigen::NumTraits<Scalar>::infinity())
                    {
                        range.nonFinite = true;
                        break;
                    }
                    range.sum.add(w);
                }
            }
            return range;
        },
        [](Total a, Total const & b) {
            a.sum.add(b.sum);
            a.nonFinite = a.nonFinite || b.nonFinite;
            return a;
        }
    );
    if (total.nonFinite)
    {
        return std::unexpected{WeightError::NonFinite};
    }
    Scalar const result = total.sum.value();
    if (result == -Eigen::NumTraits<Scalar>::infinity())
    {
        return std::unexpected{WeightError::ZeroSum};
    }
    return result;
}

template <typename Scalar>
class PruneCut
{
public:

    PruneCut() = default;

    PruneCut(Scalar cutoff, Eigen::Index keepTies)
        : m_cutoff{cutoff}
        , m_keepTies{keepTies}
    {}

    [[nodiscard]] bool
    removes(Scalar logWeight)
    {
        if (logWeight < m_cutoff)
        {
            return true;
        }
        if (logWeight == m_cutoff)
        {
            return m_tiesSeen++ >= m_keepTies;
        }
        return false;
    }
 
private:
 
    Scalar m_cutoff = -Eigen::NumTraits<Scalar>::infinity();             
    Eigen::Index m_keepTies = std::numeric_limits<Eigen::Index>::max();  
    Eigen::Index m_tiesSeen = 0;                                         
};

template <typename Scalar>
[[nodiscard]] PruneCut<Scalar>
lightestCut(std::span<Scalar const> sorted, Eigen::Index count)
{
    if (count == 0)
    {
        return {};
    }
    Scalar const cutoff = sorted[std::size_t(count - 1)];
    auto const [first, last] = std::equal_range(sorted.begin(), sorted.end(), cutoff);
    auto const lighter = Eigen::Index(first - sorted.begin());
    auto const ties = Eigen::Index(last - first);
    return {cutoff, ties - (count - lighter)};
}

template <typename Scalar>
struct PruneBelowRule
{
    Scalar logThreshold;  

    template <typename LogAt, typename Scratch>
    [[nodiscard]] PruneCut<Scalar>
    operator()(Eigen::Index count, LogAt const & logAt, Scalar logTotal, Scratch & /*scratch*/) const
    {
        Scalar const cut = logThreshold + logTotal;
        Scalar heaviest = -Eigen::NumTraits<Scalar>::infinity();
        for (Eigen::Index i = 0; i < count; ++i)
        {
            heaviest = std::max(heaviest, Scalar(logAt(i)));
        }
        if (heaviest < cut)
        {
            // Every element is below the threshold: the first of the heaviest stays.
            return {heaviest, 1};
        }
        return {cut, count};
    }
};

template <typename Scalar>
struct PruneTailRule
{
    Scalar logFraction;  

    template <typename LogAt, typename Scratch>
    [[nodiscard]] PruneCut<Scalar>
    operator()(Eigen::Index count, LogAt const & logAt, Scalar logTotal, Scratch & scratch) const
    {
        EMRYS_FILTER_ASSERT(logFraction < Scalar(0) && "the fraction pruned is below one");
        scratch.resize(std::size_t(count));
        for (Eigen::Index i = 0; i < count; ++i)
        {
            scratch[std::size_t(i)] = Scalar(logAt(i));
        }
        std::sort(scratch.begin(), scratch.end());
        Scalar const budget = logFraction + logTotal;
        auto removed = LogSumExp<Scalar>{};
        Eigen::Index taken = 0;
        while (taken < count - 1)
        {
            auto next = removed;
            next.add(scratch[std::size_t(taken)]);
            if (!(next.value() < budget))
            {
                break;
            }
            removed = next;
            ++taken;
        }
        return lightestCut(std::span<Scalar const>{scratch.data(), scratch.size()}, taken);
    }
};

struct PruneToSizeRule
{
    Eigen::Index maxSize;  

    template <typename Scalar, typename LogAt, typename Scratch>
    [[nodiscard]] PruneCut<Scalar>
    operator()(Eigen::Index count, LogAt const & logAt, Scalar /*logTotal*/, Scratch & scratch) const
    {
        EMRYS_FILTER_ASSERT(maxSize >= 1 && "a prune leaves at least one element");
        if (count <= maxSize)
        {
            return {};
        }
        scratch.resize(std::size_t(count));
        for (Eigen::Index i = 0; i < count; ++i)
        {
            scratch[std::size_t(i)] = Scalar(logAt(i));
        }
        // The lightest elements to remove come first; the ties among them are found in a sort of
        // those up to the one at the cut.
        Eigen::Index const removed = count - maxSize;
        auto const at = scratch.begin() + std::ptrdiff_t(removed - 1);
        std::nth_element(scratch.begin(), at, scratch.end());
        Scalar const cutoff = *at;
        auto const lighter = Eigen::Index(std::count_if(scratch.begin(), scratch.end(), [&](Scalar w) {
            return w < cutoff;
        }));
        auto const ties = Eigen::Index(std::count(scratch.begin(), scratch.end(), cutoff));
        return {cutoff, ties - (removed - lighter)};
    }
};

template <typename Scalar, typename Weight, std::uniform_random_bit_generator G>
[[nodiscard]] Eigen::Index
drawIndex(Eigen::Index count, Weight const & weight, G & rng)
{
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    EMRYS_FILTER_ASSERT(count > 0 && "draw from at least one element");
    Scalar const u = std::uniform_real_distribution<Scalar>{}(rng);
    Scalar sum = 0;
    Eigen::Index last = 0;
    for (Eigen::Index i = 0; i < count; ++i)
    {
        Scalar const w = weight(i);
        if (w > 0)
        {
            sum += w;
            last = i;
            if (u < sum)
            {
                return i;
            }
        }
    }
    return last;
}
 
}  // namespace internal
 
}  // namespace Emrys