Resampling.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <random>
#include <span>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
 
namespace Emrys
{


template <typename R, typename Scalar, typename G>
concept Resampler = requires(R const & r, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) {
    r(weights, ancestors, rng);
};

template <typename R, typename Scalar, typename G>
concept ResamplerWithScratch = Resampler<R, Scalar, G>
                            && requires(
                                   R const & r,
                                   std::span<Scalar const> weights,
                                   std::span<Eigen::Index> ancestors,
                                   G & rng,
                                   std::span<Scalar> scratch
                            ) { r(weights, ancestors, rng, scratch); };

template <typename R, typename P, typename Scalar, typename G>
concept ResamplerUnder =
    ExecutionPolicy<P>
    && requires(R const & r, P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) {
           r(std::forward<P>(policy), weights, ancestors, rng);
       };
 
namespace internal
{

template <std::floating_point Scalar, std::uniform_random_bit_generator G>
[[nodiscard]] Scalar
uniform01(G & rng)
{
    auto uniform = std::uniform_real_distribution<Scalar>{Scalar(0), Scalar(1)};
    return std::min(uniform(rng), std::nextafter(Scalar(1), Scalar(0)));
}

template <typename Scalar, typename Positions>
void
invertCumulative(std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, Positions u)
{
    EMRYS_FILTER_ASSERT(!weights.empty());
    Scalar total = 0;
    std::size_t lastPositive = 0;
    for (std::size_t i = 0; i < weights.size(); ++i)
    {
        EMRYS_FILTER_ASSERT(weights[i] >= Scalar(0) && "weights must be nonnegative");
        total += weights[i];
        if (weights[i] > Scalar(0))
        {
            lastPositive = i;
        }
    }
    EMRYS_FILTER_ASSERT(total > Scalar(0) && "weights must have a positive sum");
 
    std::size_t i = 0;
    Scalar cumulative = weights[0];
    for (std::size_t j = 0; j < ancestors.size(); ++j)
    {
        Scalar const position = u(j) * total;
        while (i < lastPositive && position >= cumulative)
        {
            cumulative += weights[++i];
        }
        ancestors[j] = Eigen::Index(i);
    }
}

template <ExecutionPolicy P, typename Scalar, typename Positions>
void
invertCumulative(P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, Positions const & u)
{
    if constexpr (isSequential<P>)
    {
        invertCumulative(weights, ancestors, u);
    }
    else
    {
        EMRYS_FILTER_ASSERT(!weights.empty());
        Eigen::Index const count = Eigen::Index(weights.size());
        auto cumulative = scratchVector<Scalar>(policy, weights.size());
        cumulativeSum(policy, weights, std::span<Scalar>{cumulative});
        // The last positive weight, or -1 for none, as the largest index of each range's.
        Eigen::Index const lastPositive = reduceRanges(
            policy,
            count,
            [&](Eigen::Index begin, Eigen::Index length) {
                Eigen::Index last = -1;
                for (Eigen::Index i = begin; i < begin + length; ++i)
                {
                    EMRYS_FILTER_ASSERT(weights[std::size_t(i)] >= Scalar(0) && "weights must be nonnegative");
                    if (weights[std::size_t(i)] > Scalar(0))
                    {
                        last = i;
                    }
                }
                return last;
            },
            [](Eigen::Index a, Eigen::Index b) {
                return b >= 0 ? b : a;
            }
        );
        Scalar const total = cumulative.back();
        EMRYS_FILTER_ASSERT(lastPositive >= 0 && total > Scalar(0) && "weights must have a positive sum");
        forEachIndex(std::forward<P>(policy), Eigen::Index(ancestors.size()), [&](Eigen::Index j) {
            Scalar const position = u(std::size_t(j)) * total;
            auto const above = std::upper_bound(cumulative.begin(), cumulative.end(), position);
            ancestors[std::size_t(j)] = std::min(Eigen::Index(above - cumulative.begin()), lastPositive);
        });
    }
}

template <std::floating_point Scalar, typename P, std::uniform_random_bit_generator G>
[[nodiscard]] auto
uniforms01(P const & policy, std::size_t count, G & rng)
{
    auto draws = scratchVector<Scalar>(policy, count);
    for (Scalar & draw : draws)
    {
        draw = uniform01<Scalar>(rng);
    }
    return draws;
}
 
}  // namespace internal

namespace resamplers
{

struct Multinomial
{
    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(
        std::span<Scalar const> weights,
        std::span<Eigen::Index> ancestors,
        G & rng,
        std::span<Scalar> scratch
    ) const
    {
        EMRYS_FILTER_ASSERT(scratch.size() >= ancestors.size() && "the scratch must hold a draw per ancestor");
        auto const positions = scratch.first(ancestors.size());
        for (auto & p : positions)
        {
            p = internal::uniform01<Scalar>(rng);
        }
        std::ranges::sort(positions);
        internal::invertCumulative(weights, ancestors, [&](std::size_t j) {
            return positions[j];
        });
    }

    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        auto positions = std::vector<Scalar>(ancestors.size());
        (*this)(weights, ancestors, rng, std::span<Scalar>{positions});
    }

    template <ExecutionPolicy P, std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        if constexpr (internal::isSequential<P>)
        {
            (*this)(weights, ancestors, rng);
        }
        else
        {
            auto const positions = internal::uniforms01<Scalar>(policy, ancestors.size(), rng);
            internal::invertCumulative(std::forward<P>(policy), weights, ancestors, [&](std::size_t j) {
                return positions[j];
            });
        }
    }
};

[[nodiscard]] constexpr Multinomial
multinomial()
{
    return {};
}

struct Stratified
{
    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        Scalar const m = Scalar(ancestors.size());
        internal::invertCumulative(weights, ancestors, [&](std::size_t j) {
            return (Scalar(j) + internal::uniform01<Scalar>(rng)) / m;
        });
    }

    template <ExecutionPolicy P, std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        if constexpr (internal::isSequential<P>)
        {
            (*this)(weights, ancestors, rng);
        }
        else
        {
            Scalar const m = Scalar(ancestors.size());
            auto const offsets = internal::uniforms01<Scalar>(policy, ancestors.size(), rng);
            internal::invertCumulative(std::forward<P>(policy), weights, ancestors, [&](std::size_t j) {
                return (Scalar(j) + offsets[j]) / m;
            });
        }
    }
};

[[nodiscard]] constexpr Stratified
stratified()
{
    return {};
}

struct Systematic
{
    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        Scalar const m = Scalar(ancestors.size());
        Scalar const offset = internal::uniform01<Scalar>(rng);
        internal::invertCumulative(weights, ancestors, [&](std::size_t j) {
            return (Scalar(j) + offset) / m;
        });
    }

    template <ExecutionPolicy P, std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        Scalar const m = Scalar(ancestors.size());
        Scalar const offset = internal::uniform01<Scalar>(rng);
        internal::invertCumulative(std::forward<P>(policy), weights, ancestors, [&](std::size_t j) {
            return (Scalar(j) + offset) / m;
        });
    }
};

[[nodiscard]] constexpr Systematic
systematic()
{
    return {};
}

template <typename Inner = Multinomial>
struct Residual
{
    Inner fractional{};  

    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(
        std::span<Scalar const> weights,
        std::span<Eigen::Index> ancestors,
        G & rng,
        std::span<Scalar> scratch
    ) const
    {
        EMRYS_FILTER_ASSERT(
            scratch.size() >= weights.size() + ancestors.size()
            && "the scratch must hold a fraction per weight and a draw per ancestor"
        );
        Scalar total = 0;
        for (Scalar const w : weights)
        {
            EMRYS_FILTER_ASSERT(w >= Scalar(0) && "weights must be nonnegative");
            total += w;
        }
        EMRYS_FILTER_ASSERT(total > Scalar(0) && "weights must have a positive sum");
 
        Scalar const m = Scalar(ancestors.size());
        auto const fractions = scratch.first(weights.size());
        std::size_t out = 0;
        for (std::size_t i = 0; i < weights.size(); ++i)
        {
            Scalar const expected = m * weights[i] / total;
            Scalar const copies = std::floor(expected);
            fractions[i] = expected - copies;
            // Converting a negative or NaN count to std::size_t is undefined, so it takes no copies.
            std::size_t const count = copies > Scalar(0) ? std::size_t(copies) : 0;
            for (std::size_t k = 0; k < count && out < ancestors.size(); ++k)
            {
                ancestors[out++] = Eigen::Index(i);
            }
        }
        if (out < ancestors.size())
        {
            if constexpr (ResamplerWithScratch<Inner, Scalar, G>)
            {
                fractional(
                    std::span<Scalar const>{fractions},
                    ancestors.subspan(out),
                    rng,
                    scratch.subspan(weights.size())
                );
            }
            else
            {
                fractional(std::span<Scalar const>{fractions}, ancestors.subspan(out), rng);
            }
        }
    }

    template <std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        auto scratch = std::vector<Scalar>(weights.size() + ancestors.size());
        (*this)(weights, ancestors, rng, std::span<Scalar>{scratch});
    }

    template <ExecutionPolicy P, std::floating_point Scalar, std::uniform_random_bit_generator G>
    void
    operator()(P && policy, std::span<Scalar const> weights, std::span<Eigen::Index> ancestors, G & rng) const
    {
        if constexpr (internal::isSequential<P>)
        {
            (*this)(weights, ancestors, rng);
        }
        else
        {
            Eigen::Index const count = Eigen::Index(weights.size());
            Scalar const total = internal::reduceRanges(
                policy,
                count,
                [&](Eigen::Index begin, Eigen::Index length) {
                    Scalar sum = 0;
                    for (Scalar const w : weights.subspan(std::size_t(begin), std::size_t(length)))
                    {
                        sum += w;
                    }
                    return sum;
                },
                [](Scalar a, Scalar b) {
                    return a + b;
                }
            );
            EMRYS_FILTER_ASSERT(total > Scalar(0) && "weights must have a positive sum");
 
            Scalar const m = Scalar(ancestors.size());
            auto fractions = internal::scratchVector<Scalar>(policy, weights.size());
            auto ends = internal::scratchVector<Eigen::Index>(policy, weights.size());
            internal::forEachIndex(policy, count, [&](Eigen::Index i) {
                Scalar const expected = m * weights[std::size_t(i)] / total;
                Scalar const copies = std::floor(expected);
                fractions[std::size_t(i)] = expected - copies;
                ends[std::size_t(i)] = Eigen::Index(copies);
            });
            internal::cumulativeSum(policy, std::span<Eigen::Index const>{ends}, std::span<Eigen::Index>{ends});
            Eigen::Index const size = Eigen::Index(ancestors.size());
            internal::forEachIndex(policy, count, [&](Eigen::Index i) {
                Eigen::Index const begin = i == 0 ? 0 : ends[std::size_t(i - 1)];
                for (Eigen::Index k = begin; k < std::min(ends[std::size_t(i)], size); ++k)
                {
                    ancestors[std::size_t(k)] = i;
                }
            });
            std::size_t const out = std::size_t(std::min(ends.back(), size));
            if (out < ancestors.size())
            {
                auto const remainder = std::span<Scalar const>{fractions};
                if constexpr (ResamplerUnder<Inner, P, Scalar, G>)
                {
                    fractional(std::forward<P>(policy), remainder, ancestors.subspan(out), rng);
                }
                else
                {
                    fractional(remainder, ancestors.subspan(out), rng);
                }
            }
        }
    }
};

template <typename Inner = Multinomial>
[[nodiscard]] constexpr Residual<Inner>
residual(Inner inner = {})
{
    return {std::move(inner)};
}
 
}  // namespace resamplers
 
}  // namespace Emrys