ParticleFilter.h

Go to the documentation of this file.

#pragma once

 
#include <expected>
#include <random>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Models/Measurement.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/ParticleSet.h>
#include <EmrysFilter/Models/ProcessModel.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Filters/FilterUpkeep.h>
#include <EmrysFilter/Filters/Recursion.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Models/Transform.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Core/Print.h>
 
namespace Emrys
{

template <typename Scalar = double>
class ParticleFilterUpdate
{
public:

    [[nodiscard]] Scalar
    logLikelihood() const
    {
        return m_logLikelihood;
    }
 
    bool accepted = false;  
 
private:
 
    friend struct internal::UpdateAccess;
 
    Scalar m_logLikelihood = 0;  
};
 
namespace internal
{

template <typename Input>
struct DrawnInput
{
    DrawnInput(Input const & input, Eigen::Index & /*row*/)
        : input{input}
    {}
 
    static constexpr int RowsAtCompileTime = 0;     
    static constexpr int MaxRowsAtCompileTime = 0;  
 
    Input const & input;  

    template <typename Derived, typename G>
    void
    fill(Derived && /*draws*/, G & /*rng*/) const
    {}

    template <typename Derived>
    [[nodiscard]] Input const &
    at(Eigen::MatrixBase<Derived> const & /*draws*/) const
    {
        return input;
    }
};

template <UncertainInput Input>
struct DrawnInput<Input>
{
    using Space = typename Input::SpaceType;  
    using Scalar = typename Input::Scalar;    

    DrawnInput(Input const & input, Eigen::Index & row)
        : input{input}
        , factor{defaultRoot(covarianceLowerOf(input))}
        , firstRow{row}
    {
        row += input.differenceDimension();
    }

    static constexpr int RowsAtCompileTime = rowsAtCompileTime<true, Space>;

    static constexpr int MaxRowsAtCompileTime = maxValueDimensionAtCompileTime<DifferenceSpace<Space>>;
 
    Input const & input;  

    typename Moments<Space, Scalar>::MatrixType factor;
 
    Eigen::Index firstRow;  

    template <typename Derived, typename G>
    void
    fill(Derived && draws, G & rng, std::normal_distribution<Scalar> & normal) const
    {
        fillStandardNormal(middleRowsOf<RowsAtCompileTime>(draws, firstRow, input.differenceDimension()), rng, normal);
    }

    template <typename Derived>
    [[nodiscard]] typename Moments<Space, Scalar>::VectorType
    at(Eigen::MatrixBase<Derived> const & draws) const
    {
        typename Moments<Space, Scalar>::DifferenceVectorType const z =
            middleRowsOf<RowsAtCompileTime>(draws.derived(), firstRow, input.differenceDimension());
        return Emrys::plus(input.space(), input.mean(), multiply(factor, z));
    }
};

template <typename Input>
[[nodiscard]] Eigen::Index
drawnRows([[maybe_unused]] Input const & input)
{
    if constexpr (UncertainInput<Input>)
    {
        return input.differenceDimension();
    }
    else
    {
        return 0;
    }
}
 
}  // namespace internal

template <
    Space Space_,
    typename Model_,
    RealScalar Scalar_ = double,
    typename Allocator_ = DefaultAllocator,
    typename Upkeep_ = upkeep::None>
class ParticleFilter
    : public internal::HeldUpkeep<ParticleFilter<Space_, Model_, Scalar_, Allocator_, Upkeep_>, Upkeep_>
{
    using UpkeepBase = internal::HeldUpkeep<ParticleFilter, Upkeep_>;  
 
public:
 
    using Scalar = Scalar_;                                      
    using SpaceType = Space_;                                    
    using ModelType = Model_;                                    
    using StateType = ParticleSet<Space_, Scalar_, Allocator_>;  
    using allocator_type = Allocator_;                           

    ParticleFilter(std::type_identity_t<Model_> model, StateType state, std::type_identity_t<Upkeep_> upkeep = {})
        : UpkeepBase{std::move(upkeep)}
        , m_model{std::move(model)}
        , m_state{std::move(state)}
        , m_logLikelihoods(ScalarAllocator(m_state.get_allocator()))
        , m_draws(ScalarAllocator(m_state.get_allocator()))
        , m_batch(ScalarAllocator(m_state.get_allocator()))
    {
        EMRYS_FILTER_ASSERT(!m_state.empty() && "the initial particle set must have a particle");
        internal::assertPartitions(m_model, m_state.space());
        reserveScratch();
    }

    [[nodiscard]] allocator_type
    get_allocator() const
    {
        return m_state.get_allocator();
    }

    [[nodiscard]] StateType const &
    state() const
    {
        return m_state;
    }

    [[nodiscard]] StateType &
    state()
    {
        return m_state;
    }

    void
    setState(StateType state)
    {
        EMRYS_FILTER_ASSERT(!state.empty() && "the particle set must have a particle");
        internal::assertPartitions(m_model, state.space());
        m_state = std::move(state);
        if (m_logLikelihoods.get_allocator() != ScalarAllocator(m_state.get_allocator()))
        {
            // The assignment carried the new set's allocator in, so the scratch follows it.
            m_logLikelihoods = internal::Buffer<Scalar_, Allocator_>(ScalarAllocator(m_state.get_allocator()));
            m_draws = internal::Buffer<Scalar_, Allocator_>(ScalarAllocator(m_state.get_allocator()));
            m_batch = internal::Buffer<Scalar_, Allocator_>(ScalarAllocator(m_state.get_allocator()));
        }
        reserveScratch();
    }

    [[nodiscard]] Model_ const &
    model() const
    {
        return m_model;
    }

    [[nodiscard]] Model_ &
    model()
    {
        return m_model;
    }

    template <std::uniform_random_bit_generator G, typename... Inputs>
        requires(!isFixedStep<Model_>)
    void
    predict(G & rng, TimeStep<Scalar_> auto const & dt, Inputs const &... inputs)
    {
        predict(execution::sequential, rng, dt, inputs...);
    }

    template <ExecutionPolicy P, std::uniform_random_bit_generator G, typename... Inputs>
        requires(!isFixedStep<Model_>)
    void
    predict(P && policy, G & rng, TimeStep<Scalar_> auto const & dt, Inputs const &... inputs)
    {
        auto const & space = m_state.space();
        Scalar_ const step = internal::toScalarTime<Scalar_>(dt);
        auto const call = internal::modelWith(std::cref(m_model), inputs...);
        auto const resolved = internal::resolveProcess(call, space);
        std::apply(
            [&](auto const &... uncertain) {
                if constexpr (std::remove_cvref_t<decltype(resolved)>::BatchedFunction)
                {
                    moveInBatches(
                        std::forward<P>(policy),
                        rng,
                        internal::defaultRoot(resolved.template noise<Scalar_>(step)),
                        resolved,
                        step,
                        internal::inputOf(uncertain)...
                    );
                }
                else
                {
                    moveParticles(
                        std::forward<P>(policy),
                        rng,
                        internal::defaultRoot(resolved.template noise<Scalar_>(step)),
                        [&](auto const & x, auto const &... in) {
                            return internal::valueOf(resolved.propagate(x, step, in...));
                        },
                        internal::inputOf(uncertain)...
                    );
                }
            },
            resolved.inputs()
        );
    }

    template <std::uniform_random_bit_generator G, typename... Inputs>
        requires(isFixedStep<Model_> && !internal::leadsWithSteps<Inputs...>)
    void
    predict(G & rng, Inputs const &... inputs)
    {
        internal::rejectTimeStep<Inputs...>();
        predict(execution::sequential, rng, inputs...);
    }

    template <ExecutionPolicy P, std::uniform_random_bit_generator G, typename... Inputs>
        requires(isFixedStep<Model_> && !internal::leadsWithSteps<Inputs...>)
    void
    predict(P && policy, G & rng, Inputs const &... inputs)
    {
        internal::rejectTimeStep<Inputs...>();
        auto const & space = m_state.space();
        auto const call = internal::modelWith(std::cref(m_model.model()), inputs...);
        auto const resolved = internal::resolveProcess(call, space);
        std::apply(
            [&](auto const &... uncertain) {
                if constexpr (std::remove_cvref_t<decltype(resolved)>::BatchedFunction)
                {
                    moveInBatches(
                        std::forward<P>(policy),
                        rng,
                        m_model.noiseRoot(),
                        resolved,
                        m_model.timeStep(),
                        internal::inputOf(uncertain)...
                    );
                }
                else
                {
                    moveParticles(
                        std::forward<P>(policy),
                        rng,
                        m_model.noiseRoot(),
                        [&](auto const & x, auto const &... in) {
                            return internal::valueOf(resolved.propagate(x, m_model.timeStep(), in...));
                        },
                        internal::inputOf(uncertain)...
                    );
                }
            },
            resolved.inputs()
        );
    }

    template <std::uniform_random_bit_generator G, typename... Inputs>
        requires isFixedStep<Model_>
    void
    predict(G & rng, Steps steps, Inputs const &... inputs)
    {
        predict(execution::sequential, rng, steps, inputs...);
    }

    template <ExecutionPolicy P, std::uniform_random_bit_generator G, typename... Inputs>
        requires isFixedStep<Model_>
    void
    predict(P && policy, G & rng, Steps steps, Inputs const &... inputs)
    {
        EMRYS_FILTER_ASSERT(steps.count >= 0 && "the number of steps must be nonnegative");
        for (Eigen::Index i = 0; i < steps.count; ++i)
        {
            predict(policy, rng, inputs...);
        }
    }

    template <internal::WeighsPoints M>
    [[nodiscard]] std::expected<ParticleFilterUpdate<Scalar_>, WeightError>
    update(M const & measurement)
    {
        return update(execution::sequential, measurement);
    }

    template <ExecutionPolicy P, internal::WeighsPoints M>
    [[nodiscard]] std::expected<ParticleFilterUpdate<Scalar_>, WeightError>
    update(P && policy, M const & measurement)
    {
        auto const logTotal = m_state.logTotalWeight(policy);
        if (!logTotal)
        {
            return std::unexpected{logTotal.error()};
        }
        m_logLikelihoods.resize(std::size_t(m_state.size()));
        auto logWeights = Eigen::Map<typename StateType::WeightArrayType>{m_logLikelihoods.data(), m_state.size()};
        internal::fillLogLikelihoods(policy, measurement, m_state.space(), m_state.states(), logWeights, m_batch);
        logWeights += m_state.logWeights() - *logTotal;
        // A NaN log-likelihood at any particle makes the evidence NaN.
        Scalar_ const logLikelihood = logSumExp(std::forward<P>(policy), logWeights);
        EMRYS_FILTER_ASSERT(!isNaN(logLikelihood) && "the measurement's log-likelihood must not be NaN");
        auto result = ParticleFilterUpdate<Scalar_>{};
        internal::UpdateAccess::logLikelihood(result) = logLikelihood;
        result.accepted = logLikelihood > -Eigen::NumTraits<Scalar_>::infinity();
        if (result.accepted)
        {
            internal::ParticleSetAccess::setNormalizedLogWeights(m_state, logWeights - logLikelihood);
        }
        return result;
    }
 
private:
 
    using ScalarAllocator = internal::ReBound<Allocator_, Scalar_>;  

    void
    reserveScratch()
    {
        auto const count = std::size_t(m_state.capacity().value_or(m_state.size()));
        m_logLikelihoods.reserve(count);
        m_draws.reserve(std::size_t(m_state.differenceDimension()) * count);
    }

    template <typename P, typename G, typename Derived, typename Propagate, typename... Inputs>
    void
    moveParticles(
        P && policy,
        G & rng,
        Eigen::MatrixBase<Derived> const & noiseFactor,
        Propagate const & propagate,
        Inputs const &... inputs
    )
    {
        using Vector = ValueVector<Space_, Scalar_>;
        using Difference = DifferenceVector<Space_, Scalar_>;
        constexpr int NoiseRows = internal::rowsAtCompileTime<true, Space_>;
        constexpr int Rows = internal::sizeSum<NoiseRows, internal::DrawnInput<Inputs>::RowsAtCompileTime...>;
        constexpr int MaxRows = internal::sizeSum<
            maxValueDimensionAtCompileTime<DifferenceSpace<Space_>>,
            internal::DrawnInput<Inputs>::MaxRowsAtCompileTime...>;
        auto const & space = m_state.space();
        Eigen::Index const count = m_state.size();
        EMRYS_FILTER_ASSERT(count > 0 && "the particle set must have a particle");
        Eigen::Index const noiseRows = differenceDimension(space);
        // Initialized in order, each input taking its rows after those before it.
        Eigen::Index row = noiseRows;
        auto const drawn = std::tuple{internal::DrawnInput<Inputs>{inputs, row}...};
        Eigen::Index const rows = noiseRows + (Eigen::Index{0} + ... + internal::drawnRows(inputs));
        EMRYS_FILTER_ASSERT(row == rows);
        // One distribution for every draw, which may keep a value between them, so that the draws
        // are those of filling every particle's column at once.
        auto normal = std::normal_distribution<Scalar_>{};
        auto const draw = [&](auto && column) {
            internal::fillStandardNormal(internal::middleRowsOf<NoiseRows>(column, 0, noiseRows), rng, normal);
            std::apply(
                [&](auto const &... inputDraws) {
                    (inputDraws.fill(column, rng, normal), ...);
                },
                drawn
            );
        };
        auto const move = [&](Eigen::Index i, auto const & column) {
            Difference const z = internal::middleRowsOf<NoiseRows>(column, 0, noiseRows);
            Vector const moved = std::apply(
                [&](auto const &... inputDraws) {
                    return Vector{propagate(m_state.state(i), inputDraws.at(column)...)};
                },
                drawn
            );
            m_state.setState(i, Emrys::plus(space, moved, Difference{internal::multiply(noiseFactor, z)}));
        };
        if constexpr (internal::isSequential<P>)
        {
            auto column = internal::MatrixOf<Scalar_, Rows, 1, MaxRows, 1>{};
            column.resize(rows);
            for (Eigen::Index i = 0; i < count; ++i)
            {
                draw(column);
                move(i, column);
            }
        }
        else
        {
            using Draws = internal::MatrixOf<Scalar_, Rows, Eigen::Dynamic, MaxRows, Eigen::Dynamic>;
            m_draws.resize(std::size_t(rows * count));
            auto draws = Eigen::Map<Draws>{m_draws.data(), rows, count};
            for (Eigen::Index i = 0; i < count; ++i)
            {
                draw(draws.col(i));
            }
            internal::forEachIndex(std::forward<P>(policy), count, [&](Eigen::Index i) {
                move(i, draws.col(i));
            });
        }
    }

    template <typename P, typename G, typename Derived, typename R, typename TimeScalar, typename... Inputs>
    void
    moveInBatches(
        P && policy,
        G & rng,
        Eigen::MatrixBase<Derived> const & noiseFactor,
        R const & resolved,
        TimeScalar dt,
        Inputs const &... inputs
    )
    {
        using Difference = DifferenceVector<Space_, Scalar_>;
        constexpr int NoiseRows = internal::rowsAtCompileTime<true, Space_>;
        constexpr int Rows = internal::sizeSum<NoiseRows, internal::DrawnInput<Inputs>::RowsAtCompileTime...>;
        constexpr int MaxRows = internal::sizeSum<
            maxValueDimensionAtCompileTime<DifferenceSpace<Space_>>,
            internal::DrawnInput<Inputs>::MaxRowsAtCompileTime...>;
        using Draws = internal::MatrixOf<Scalar_, Rows, Eigen::Dynamic, MaxRows, Eigen::Dynamic>;
        using X = typename R::template Columns<typename R::In, Scalar_>;
        using Y = typename R::template Columns<typename R::Out, Scalar_>;
        auto const & space = m_state.space();
        Eigen::Index const count = m_state.size();
        EMRYS_FILTER_ASSERT(count > 0 && "the particle set must have a particle");
        Eigen::Index const noiseRows = differenceDimension(space);
        Eigen::Index row = noiseRows;
        auto const drawn = std::tuple{internal::DrawnInput<Inputs>{inputs, row}...};
        Eigen::Index const rows = noiseRows + (Eigen::Index{0} + ... + internal::drawnRows(inputs));
        EMRYS_FILTER_ASSERT(row == rows);
        auto normal = std::normal_distribution<Scalar_>{};
        m_draws.resize(std::size_t(rows * count));
        auto draws = Eigen::Map<Draws>{m_draws.data(), rows, count};
        for (Eigen::Index i = 0; i < count; ++i)
        {
            auto column = draws.col(i);
            internal::fillStandardNormal(internal::middleRowsOf<NoiseRows>(column, 0, noiseRows), rng, normal);
            std::apply(
                [&](auto const &... inputDraws) {
                    (inputDraws.fill(column, rng, normal), ...);
                },
                drawn
            );
        }
        Eigen::Index const inRows = valueDimension(resolved.in());
        Eigen::Index const outRows = valueDimension(resolved.out());
        m_batch.resize(std::size_t((inRows + outRows) * count));
        internal::forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
            auto x = Eigen::Map<X>{m_batch.data() + inRows * begin, inRows, length};
            auto y = Eigen::Map<Y>{m_batch.data() + inRows * count + outRows * begin, outRows, length};
            for (Eigen::Index j = 0; j < length; ++j)
            {
                std::apply(
                    [&](auto const &... inputDraws) {
                        resolved.template gatherColumn<Scalar_>(
                            x.col(j),
                            m_state.state(begin + j),
                            inputDraws.at(draws.col(begin + j))...
                        );
                    },
                    drawn
                );
            }
            resolved.stepBatch(Eigen::Map<X const>{x.data(), inRows, length}, dt, y);
            for (Eigen::Index j = 0; j < length; ++j)
            {
                Eigen::Index const i = begin + j;
                auto const column = draws.col(i);
                Difference const z = internal::middleRowsOf<NoiseRows>(column, 0, noiseRows);
                auto const moved = resolved.scattered(m_state.state(i), y.col(j));
                m_state.setState(i, Emrys::plus(space, moved, Difference{internal::multiply(noiseFactor, z)}));
            }
        });
    }
 
    Model_ m_model;                                          
    StateType m_state;                                       
    internal::Buffer<Scalar_, Allocator_> m_logLikelihoods;  
    internal::Buffer<Scalar_, Allocator_> m_draws;  
    internal::Buffer<Scalar_, Allocator_> m_batch;  
};

template <typename Space, typename Model, typename Scalar, typename Allocator>
ParticleFilter(Model, ParticleSet<Space, Scalar, Allocator>)
    -> ParticleFilter<Space, internal::BoundStep<Model, Space, Scalar>, Scalar, Allocator>;
 
template <typename Space, typename Model, typename Scalar, typename Allocator, typename U>
ParticleFilter(Model, ParticleSet<Space, Scalar, Allocator>, U)
    -> ParticleFilter<Space, internal::BoundStep<Model, Space, Scalar>, Scalar, Allocator, U>;
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Space_, typename Model_, typename Scalar_, typename Allocator_, typename Upkeep_>
void
describe(
    std::ostream & stream,
    ParticleFilter<Space_, Model_, Scalar_, Allocator_, Upkeep_> const & filter,
    Indent indent
)
{
    describeFilter(stream, "particle filter", filter.state(), indent);
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO