Reduction.h

Go to the documentation of this file.

#pragma once

 
#include <array>
#include <concepts>
#include <cstddef>
#include <expected>
#include <iterator>
#include <optional>
#include <ranges>
#include <span>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Models/Consistency.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Forwarding.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Models/Measurement.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Core/Weights.h>
 
namespace Emrys
{
 
namespace internal
{

template <typename T>
inline constexpr bool isUpdateBranch = false;

template <typename Space, typename MeasurementSpace, typename Scalar>
inline constexpr bool isUpdateBranch<UpdateBranch<Space, MeasurementSpace, Scalar>> = true;


template <typename R>
concept BranchRange = std::ranges::random_access_range<R> && std::ranges::sized_range<R>
                   && isUpdateBranch<std::remove_cvref_t<std::ranges::range_reference_t<R>>>;

template <typename R>
using BranchOf = std::remove_cvref_t<std::ranges::range_reference_t<R>>;

template <typename R>
using BranchScalarOf = decltype(BranchOf<R>::logWeight);

template <typename R>
using PosteriorOf = decltype(BranchOf<R>::posterior);

template <BranchRange R>
[[nodiscard]] BranchScalarOf<R>
logTotalWeight(R const & branches)
{
    auto total = LogSumExp<BranchScalarOf<R>>{};
    for (auto const & branch : branches)
    {
        total.add(branch.logWeight);
    }
    return total.value();
}

template <typename Out>
inline constexpr bool holdsSeveral = false;

template <typename Out>
    requires requires { Out::HoldsSeveral; }
inline constexpr bool holdsSeveral<Out> = Out::HoldsSeveral;


template <typename Component, typename Scalar, typename Target = Component>
struct OneComponentOutput
{
    using ComponentType = Component;
 
    Target & target;       
    Scalar logWeight = 0;  
    bool added = false;    

    void
    add(Scalar weight, Component && component)
    {
        EMRYS_FILTER_ASSERT(!added && "this filter keeps one component per update, so a policy adds at most one");
        target = std::move(component);
        logWeight = weight;
        added = true;
    }

    void
    add(Scalar weight, Component const & component)
    {
        add(weight, Component(component));
    }
};

template <typename Component, typename Scalar>
struct SeveralComponentOutput
{
    using ComponentType = Component;
    static constexpr bool HoldsSeveral = true;

    void
    add(Scalar, Component &&)
    {}

    void
    add(Scalar, Component const &)
    {}
};

template <typename P, typename Branches, typename Component>
concept ReturnsComponent = requires(P const & policy, Branches && branches) {
    { policy(static_cast<Branches &&>(branches)) } -> std::same_as<std::optional<Component>>;
};

template <typename P, typename Branches, typename Out>
concept ReducesInto =
    std::invocable<P const &, Branches, Out &> || ReturnsComponent<P, Branches, typename Out::ComponentType>;

template <typename P, BranchRange Branches, typename Out>
    requires ReducesInto<P, Branches, Out>
void
reduceInto(P const & policy, Branches && branches, Out & out)
{
    if constexpr (std::invocable<P const &, Branches, Out &>)
    {
        policy(std::forward<Branches>(branches), out);
    }
    else
    {
        auto const logTotal = logTotalWeight(branches);
        if (auto reduced = policy(std::forward<Branches>(branches)))
        {
            out.add(logTotal, std::move(*reduced));
        }
    }
}

template <typename Iterator>
struct BranchAt
{
    Iterator first;  

    [[nodiscard]] decltype(auto)
    operator()(std::size_t i) const
    {
        return first[std::iter_difference_t<Iterator>(i)];
    }
};

template <typename R>
using PassingBranches =
    ForwardedLike<R, std::ranges::transform_view<std::span<std::size_t const>, BranchAt<std::ranges::iterator_t<R>>>>;
 
}  // namespace internal

namespace reductions
{

struct Collapse
{
    template <typename Branches>
        requires internal::BranchRange<Branches>
    [[nodiscard]] std::optional<internal::PosteriorOf<Branches>>
    operator()(Branches && branches) const
    {
        using Posterior = internal::PosteriorOf<Branches>;
        using Scalar = internal::BranchScalarOf<Branches>;
        using MomentsType = typename Posterior::MomentsType;
        EMRYS_FILTER_ASSERT(!std::ranges::empty(branches) && "an update has at least one branch");
        auto const at = std::ranges::begin(branches);
        auto && first = at[0].posterior;
        Eigen::Index const count = Eigen::Index(std::ranges::size(branches));
        if (count == 1)
        {
            return internal::forwardLike<Branches>(first);
        }
 
        Scalar const logTotal = internal::logTotalWeight(branches);
        EMRYS_FILTER_ASSERT(isFinite(logTotal) && "the branches' weights must be valid");
        auto const branch = [&](Eigen::Index i) -> auto const & {
            return at[i];
        };
        return Posterior{
            first.space(),
            internal::mixtureMomentsIn<MomentsType>(
                first.space(),
                count,
                [&](Eigen::Index i) -> decltype(auto) {
                    return branch(i).posterior.mean();
                },
                [&](Eigen::Index i) -> decltype(auto) {
                    return branch(i).posterior.covarianceLower();
                },
                [&](Eigen::Index i) {
                    return branch(i).logWeight;
                },
                logTotal
            )
        };
    }
};

[[nodiscard]] constexpr Collapse
collapse()
{
    return {};
}

struct Select
{
    template <typename Branches>
        requires internal::BranchRange<Branches>
    [[nodiscard]] std::optional<internal::PosteriorOf<Branches>>
    operator()(Branches && branches) const
    {
        EMRYS_FILTER_ASSERT(!std::ranges::empty(branches) && "an update has at least one branch");
        auto const at = std::ranges::begin(branches);
        std::ptrdiff_t const count = std::ptrdiff_t(std::ranges::size(branches));
        std::ptrdiff_t best = 0;
        for (std::ptrdiff_t i = 1; i < count; ++i)
        {
            if (at[i].logWeight > at[best].logWeight)
            {
                best = i;
            }
        }
        return internal::forwardLike<Branches>(at[best].posterior);
    }
};

[[nodiscard]] constexpr Select
select()
{
    return {};
}

struct Keep
{
    template <typename Branches, typename Out>
        requires internal::BranchRange<Branches> && internal::holdsSeveral<Out>
    void
    operator()(Branches && branches, Out & out) const
    {
        for (auto && branch : branches)
        {
            out.add(branch.logWeight, internal::forwardLike<Branches>(branch.posterior));
        }
    }
};

[[nodiscard]] constexpr Keep
keep()
{
    return {};
}
 
}  // namespace reductions

template <RealScalar Scalar = double>
struct NisThreshold
{
    Scalar value;  
};

template <typename Value>
    requires std::is_arithmetic_v<Value>
NisThreshold(Value) -> NisThreshold<internal::DeducedScalar<Value>>;

 
namespace reductions
{

template <typename Inner = Collapse, RealScalar Scalar_ = double>
class Gate
{
public:

    using Scalar = Scalar_;

    explicit Gate(Scalar probability, Inner inner = {})
        : m_probability{probability}
        , m_inner{std::move(inner)}
    {
        EMRYS_FILTER_ASSERT(probability > Scalar(0) && probability < Scalar(1) && "the probability must be in (0, 1)");
    }

    explicit Gate(NisThreshold<Scalar> threshold, Inner inner = {})
        : m_threshold{threshold.value}
        , m_inner{std::move(inner)}
    {
        EMRYS_FILTER_ASSERT(threshold.value >= Scalar(0) && "the threshold must be nonnegative");
    }

    template <RealScalar Result = Scalar>
    [[nodiscard]] Result
    threshold(Eigen::Index dimension) const
    {
        if (m_probability)
        {
            return chiSquaredQuantile<Result>(Result(*m_probability), internal::scalarFromCount<Result>(dimension));
        }
        return Result(m_threshold);
    }

    [[nodiscard]] Inner const &
    inner() const
    {
        return m_inner;
    }

    template <typename Branches, typename Out>
        requires internal::BranchRange<Branches> && internal::ReducesInto<Inner, Branches, Out>
              && internal::ReducesInto<Inner, internal::PassingBranches<Branches>, Out>
    void
    operator()(Branches && branches, Out & out) const
    {
        using UpdateScalar = internal::BranchScalarOf<Branches>;
        EMRYS_FILTER_ASSERT(!std::ranges::empty(branches) && "an update has at least one branch");
        auto const at = std::ranges::begin(branches);
        std::size_t const count = std::size_t(std::ranges::size(branches));
        auto const limit = threshold<UpdateScalar>(at[0].innovation.differenceDimension());
        std::size_t passed = 0;
        for (auto const & branch : branches)
        {
            passed += branch.nis() <= limit ? 1 : 0;
        }
        if (passed == 0)
        {
            return;
        }
        if (passed == count)
        {
            internal::reduceInto(m_inner, std::forward<Branches>(branches), out);
            return;
        }
 
        auto const reducePassing = [&](auto & indices) {
            std::size_t n = 0;
            for (std::size_t k = 0; k < count; ++k)
            {
                if (at[std::ptrdiff_t(k)].nis() <= limit)
                {
                    indices[n++] = k;
                }
            }
            auto passing = std::ranges::transform_view{
                std::span<std::size_t const>{indices.data(), passed},
                internal::BranchAt<std::ranges::iterator_t<Branches>>{at}
            };
            internal::reduceInto(m_inner, internal::forwardLike<Branches>(passing), out);
        };
        using B = std::remove_cvref_t<Branches>;
        if constexpr (requires { std::tuple_size<B>::value; })
        {
            auto indices = std::array<std::size_t, std::tuple_size_v<B>>{};
            reducePassing(indices);
        }
        else if (count <= InlinePassing)
        {
            auto indices = std::array<std::size_t, InlinePassing>{};
            reducePassing(indices);
        }
        else
        {
            auto indices = internal::bufferOf<std::size_t>(allocatorOf(branches));
            indices.resize(passed);
            reducePassing(indices);
        }
    }

    static constexpr std::size_t InlinePassing = 32;
 
private:

    template <typename Branches>
    [[nodiscard]] static auto
    allocatorOf(Branches const & branches)
    {
        if constexpr (requires { branches.get_allocator(); })
        {
            return branches.get_allocator();
        }
        else
        {
            return DefaultAllocator{};
        }
    }
 
    std::optional<Scalar> m_probability;           
    Scalar m_threshold = 0;                        
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Inner m_inner;  
};

template <typename Value, typename Inner = Collapse>
    requires std::is_arithmetic_v<Value>
[[nodiscard]] Gate<Inner, internal::DeducedScalar<Value>>
gate(Value probability, Inner inner = {})
{
    using Scalar = internal::DeducedScalar<Value>;
    return Gate<Inner, Scalar>{Scalar(probability), std::move(inner)};
}

template <RealScalar Scalar, typename Inner = Collapse>
[[nodiscard]] Gate<Inner, Scalar>
gate(NisThreshold<Scalar> threshold, Inner inner = {})
{
    return Gate<Inner, Scalar>{threshold, std::move(inner)};
}
 
}  // namespace reductions

template <typename P, typename Update, typename State>
concept ReductionPolicy = internal::ReducesInto<
    P,
    typename std::remove_cvref_t<Update>::Container,
    internal::OneComponentOutput<State, typename std::remove_cvref_t<Update>::Scalar>>;

template <typename P, typename Update, typename Component_>
concept MixtureReductionPolicy = internal::ReducesInto<
    P,
    typename std::remove_cvref_t<Update>::Container,
    internal::SeveralComponentOutput<Params<Component_>, typename std::remove_cvref_t<Update>::Scalar>>;

template <typename State, typename Policy, typename Branches>
    requires internal::BranchRange<Branches>
          && internal::ReducesInto<
                 Policy,
                 Branches,
                 internal::OneComponentOutput<State, internal::BranchScalarOf<Branches>, std::optional<State>>>
[[nodiscard]] std::optional<State>
reduceToOne(Policy const & policy, Branches && branches)
{
    auto kept = std::optional<State>{};
    auto out = internal::OneComponentOutput<State, internal::BranchScalarOf<Branches>, std::optional<State>>{kept};
    internal::reduceInto(policy, std::forward<Branches>(branches), out);
    return kept;
}
 
namespace internal
{

template <typename Policy, typename Branches, typename Component>
inline constexpr bool keepsOne = ReturnsComponent<Policy, Branches, Component>;

template <typename Inner, typename Scalar, typename Branches, typename Component>
inline constexpr bool keepsOne<reductions::Gate<Inner, Scalar>, Branches, Component> =
    keepsOne<Inner, Branches, Component>;


template <typename Policy>
inline constexpr bool policyReadsStatistics = true;

template <>
inline constexpr bool policyReadsStatistics<reductions::Collapse> = false;
 
template <>
inline constexpr bool policyReadsStatistics<reductions::Select> = false;
 
template <>
inline constexpr bool policyReadsStatistics<reductions::Keep> = false;


template <typename Policy>
struct WithStatistics
{
    Policy const & policy;  

    template <typename Branches, typename Out>
        requires BranchRange<Branches> && ReducesInto<Policy, Branches, Out>
    void
    operator()(Branches && branches, Out & out) const
    {
        reduceInto(policy, std::forward<Branches>(branches), out);
    }
};

template <typename Policy>
inline constexpr bool policyReadsStatistics<WithStatistics<Policy>> = true;


template <typename Policy>
inline constexpr bool isCollapsePolicy = false;

template <>
inline constexpr bool isCollapsePolicy<reductions::Collapse> = true;
 
template <typename Policy>
inline constexpr bool isCollapsePolicy<WithStatistics<Policy>> = isCollapsePolicy<Policy>;


template <typename Policy>
inline constexpr bool isOneBranchInPlacePolicy = false;

template <>
inline constexpr bool isOneBranchInPlacePolicy<reductions::Collapse> = true;
 
template <>
inline constexpr bool isOneBranchInPlacePolicy<reductions::Select> = true;
 
template <typename Scalar>
inline constexpr bool isOneBranchInPlacePolicy<reductions::Gate<reductions::Collapse, Scalar>> = true;
 
template <typename Scalar>
inline constexpr bool isOneBranchInPlacePolicy<reductions::Gate<reductions::Select, Scalar>> = true;
 
template <typename Policy>
inline constexpr bool isOneBranchInPlacePolicy<WithStatistics<Policy>> = isOneBranchInPlacePolicy<Policy>;


template <typename Gate>
struct AcceptByGate
{
    Gate const & gate;  

    template <typename Scalar>
    [[nodiscard]] bool
    operator()(Scalar nis, Eigen::Index dimension) const
    {
        return nis <= gate.template threshold<Scalar>(dimension);
    }
};

template <typename Policy>
[[nodiscard]] auto
oneBranchAcceptor(Policy const & policy)
{
    static_assert(isOneBranchInPlacePolicy<Policy>);
    if constexpr (requires { policy.policy; })
    {
        return oneBranchAcceptor(policy.policy);
    }
    else if constexpr (requires { policy.inner(); })
    {
        return AcceptByGate<Policy>{policy};
    }
    else
    {
        return AcceptEvery{};
    }
}

template <typename Scalar, typename Space, typename MeasurementSpace, std::size_t Count, typename Allocator>
[[nodiscard]] FilterUpdate<MeasurementSpace, Count, Scalar, Allocator>
filterUpdate(MeasurementUpdate<Space, MeasurementSpace, Count, Scalar, Allocator> && measured, bool accepted)
{
    using Result = FilterUpdate<MeasurementSpace, Count, Scalar, Allocator>;
    using Branch = typename Result::BranchType;
    auto const strip = [](auto & branch) {
        return Branch{
            branch.logWeight,
            std::move(branch.innovation),
            branch.mode,
            branch.component,
            branch.iterations,
            branch.converged
        };
    };
    if constexpr (Count == std::dynamic_extent)
    {
        auto result = Result{
            typename Result::Container(measured.get_allocator()),
            MixtureEvidenceAccess::mixtureLogLikelihood(measured),
            accepted
        };
        result.branches.reserve(measured.branches.size());
        for (auto & branch : measured.branches)
        {
            result.branches.push_back(strip(branch));
        }
        return result;
    }
    else
    {
        return [&]<std::size_t... I>(std::index_sequence<I...>) {
            return Result{
                {strip(measured.branches[I])...},
                MixtureEvidenceAccess::mixtureLogLikelihood(measured),
                accepted
            };
        }(std::make_index_sequence<Count>{});
    }
}

template <
    typename Scalar,
    typename Space,
    typename MeasurementSpace,
    std::size_t Count,
    typename MeasuredAllocator,
    typename ResultAllocator>
void
filterUpdateInto(
    MeasurementUpdate<Space, MeasurementSpace, Count, Scalar, MeasuredAllocator> & measured,
    bool accepted,
    FilterUpdate<MeasurementSpace, Count, Scalar, ResultAllocator> & result
)
{
    if constexpr (Count == std::dynamic_extent)
    {
        using Branch = typename FilterUpdate<MeasurementSpace, Count, Scalar, ResultAllocator>::BranchType;
        result.branches.clear();
        result.branches.reserve(measured.branches.size());
        for (auto & branch : measured.branches)
        {
            result.branches.push_back(
                Branch{
                    branch.logWeight,
                    std::move(branch.innovation),
                    branch.mode,
                    branch.component,
                    branch.iterations,
                    branch.converged
                }
            );
        }
        MixtureEvidenceAccess::setMixtureLogLikelihood(result, MixtureEvidenceAccess::mixtureLogLikelihood(measured));
        MixtureEvidenceAccess::setLogKeptWeight(result, Scalar(0));
        result.accepted = accepted;
    }
    else
    {
        result = filterUpdate(std::move(measured), accepted);
    }
}

template <bool Statistics, typename Slots, typename Allocator, UncertainInput Prior, Measurement M, typename Reduce>
[[nodiscard]] std::expected<void, GaussianError>
updateSlot(
    Slots & slots,
    std::size_t i,
    Allocator const & allocator,
    Prior const & prior,
    M const & measurement,
    Reduce const & reduce
)
{
    using Update = typename Slots::value_type;
    EMRYS_FILTER_ASSERT(i <= slots.size() && "the slots before this one must exist");
    if constexpr (requires(Update & update) { update.measured; })
    {
        if (i == slots.size())
        {
            slots.push_back(Update(allocator));
        }
        auto & slot = slots[i];
        auto const measured = updateBranchesInto<Statistics>(prior, measurement, slot.measured);
        if (!measured)
        {
            return measured;
        }
        bool const accepted = reduce(slot.measured);
        filterUpdateInto(slot.measured, accepted, slot);
    }
    else
    {
        auto measured = updateBranches<Statistics>(prior, measurement);
        if (!measured)
        {
            return std::unexpected{measured.error()};
        }
        bool const accepted = reduce(*measured);
        auto update = filterUpdate(std::move(*measured), accepted);
        if (i == slots.size())
        {
            slots.push_back(std::move(update));
        }
        else
        {
            slots[i] = std::move(update);
        }
    }
    return {};
}

template <
    bool Statistics,
    typename Slots,
    typename Allocator,
    UncertainInput Prior,
    GivenMeasurementLike G,
    typename Reduce>
[[nodiscard]] std::expected<void, GaussianError>
updateSlot(
    Slots & slots,
    std::size_t i,
    Allocator const & allocator,
    Prior const & prior,
    G const & given,
    Reduce const & reduce
)
{
    using Update = typename Slots::value_type;
    EMRYS_FILTER_ASSERT(i <= slots.size() && "the slots before this one must exist");
    auto measured = updateBranchesWith<Statistics>(prior, given.pinned, given.value);
    if (!measured)
    {
        return std::unexpected{measured.error()};
    }
    bool const accepted = reduce(*measured);
    if constexpr (requires(Update & update) { update.measured; })
    {
        if (i == slots.size())
        {
            slots.push_back(Update(allocator));
        }
        filterUpdateInto(*measured, accepted, slots[i]);
    }
    else
    {
        auto update = filterUpdate(std::move(*measured), accepted);
        if (i == slots.size())
        {
            slots.push_back(std::move(update));
        }
        else
        {
            slots[i] = std::move(update);
        }
    }
    return {};
}
 
}  // namespace internal
 
}  // namespace Emrys