ParticleGaussian.h

Go to the documentation of this file.

#pragma once

 
#include <expected>
#include <random>
#include <span>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Print.h>
#include <EmrysFilter/Core/Resampling.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Conditionals.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Hybrid.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/Params.h>
#include <EmrysFilter/Distributions/ParticleSet.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{
 
template <
    Space Space_,
    typename Path_,
    typename Covariance_ = SharedCovariance,
    RealScalar Scalar_ = double,
    typename Allocator_ = DefaultAllocator,
    typename Component_ = Gaussian<typename internal::BlockSplit<Space_, Path_>::RestSpace, Scalar_>>
class ParticleGaussian;
 
template <
    Space Space_,
    typename Path_,
    typename Model_,
    typename Covariance_,
    typename Method_,
    RealScalar Scalar_,
    typename Allocator_,
    typename Upkeep_>
class RbParticleFilter;
 
namespace internal
{

template <typename Space_, typename Split, typename Particles>
class ParticlePoints
{
public:
 
    using Scalar = typename Particles::Scalar;        
    using SpaceType = Space_;                         
    using PathType = typename Split::PathType;        
    using PinnedSpace = typename Split::PinnedSpace;  
    using RestSpace = typename Split::RestSpace;      
    static constexpr bool HasKernels = false;         

    ParticlePoints(Space_ const & space, Split const & split, Particles const & particles)
        : m_space{space}
        , m_split{split}
        , m_particles{particles}
    {}

    [[nodiscard]] Space_ const &
    space() const
    {
        return m_space;
    }

    [[nodiscard]] RestSpace const &
    restSpace() const
    {
        return m_split.restSpace();
    }

    [[nodiscard]] BlockRows
    pinnedRows() const
    {
        return m_split.pinnedRows();
    }

    [[nodiscard]] auto const &
    restIndices() const
    {
        return m_split.restIndices();
    }

    [[nodiscard]] BlockRows
    pinnedDifferenceRows() const
    {
        return m_split.pinnedDifferenceRows();
    }

    [[nodiscard]] auto const &
    restDifferenceIndices() const
    {
        return m_split.restDifferenceIndices();
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return m_particles.size();
    }

    [[nodiscard]] auto
    center(Eigen::Index i) const
    {
        return m_particles.state(i);
    }
 
private:
 
    Space_ const & m_space;         
    Split const & m_split;          
    Particles const & m_particles;  
};
 
}  // namespace internal

template <
    Space Space_,
    typename Path_,
    typename Covariance_ = SharedCovariance,
    RealScalar Scalar_ = double,
    typename Allocator_ = DefaultAllocator,
    typename Component_ = Gaussian<typename internal::BlockSplit<Space_, Path_>::RestSpace, Scalar_>>
class ParticleGaussianParams
{
    static_assert(internal::CovarianceForm<Covariance_>, "the covariance is SharedCovariance or PerPointCovariance");
    static constexpr bool Shared = std::is_same_v<Covariance_, SharedCovariance>;  
    using Split = internal::BlockSplit<Space_, Path_>;  

    static constexpr bool IsHybrid = internal::isGaussianHybrid<Params<Component_>>;
 
    static_assert(
        std::is_same_v<
            Params<Component_>,
            GaussianParams<typename Params<Component_>::SpaceType, typename Params<Component_>::Scalar>>
            || IsHybrid,
        "a particle holds a Gaussian or a hybrid of Gaussians"
    );
    static_assert(
        std::is_same_v<typename Params<Component_>::SpaceType, typename Split::RestSpace>
            && std::is_same_v<typename Params<Component_>::Scalar, Scalar_>,
        "a particle's distribution is over the rest of the state, in the state's scalar"
    );
    static_assert(
        !(IsHybrid && Shared),
        "mixing a particle's modes gives each particle covariances of its own, so a hybrid per particle needs "
        "PerPointCovariance"
    );
 
public:
 
    using Scalar = Scalar_;                           
    using SpaceType = Space_;                         
    using PathType = Path_;                           
    using PinnedSpace = typename Split::PinnedSpace;  
    using RestSpace = typename Split::RestSpace;      
    using CovarianceType = Covariance_;               
    using ConditionalType = Params<Component_>;  
    using allocator_type = Allocator_;           
    using ParticlesType = ParticleSet<PinnedSpace, Scalar_, Allocator_>;  
    using WeightArrayType = typename ParticlesType::WeightArrayType;      
    using LogWeightsType = typename ParticlesType::LogWeightsType;        
    using MeansType = SampleSet<RestSpace, Scalar_>;                      
    using DensityType =
        ParticleGaussian<Space_, Path_, Covariance_, Scalar_, Allocator_, Component_>;  
 
    static constexpr bool IsParams = true;  

    ParticleGaussianParams(Space_ space, ParticlesType particles, ConditionalType const & conditional, Covariance_ = {})
        : m_space{std::move(space)}
        , m_split{m_space}
        , m_particles{std::move(particles)}
        , m_kept(IndexAllocator(m_particles.get_allocator()))
        , m_conditionals{conditional, m_particles.size(), m_particles.get_allocator()}
    {
        check();
        EMRYS_FILTER_ASSERT(conditional.space() == restSpace() && "the Gaussian must be over the rest of the state");
    }

    template <typename MeansDerived>
    ParticleGaussianParams(
        Space_ space,
        ParticlesType particles,
        Eigen::MatrixBase<MeansDerived> const & means,
        typename Moments<RestSpace, Scalar_>::MatrixType const & cov,
        Covariance_ = {}
    )
        requires(!IsHybrid)
        : m_space{std::move(space)}
        , m_split{m_space}
        , m_particles{std::move(particles)}
        , m_kept(IndexAllocator(m_particles.get_allocator()))
        , m_conditionals{
              ConditionalType{m_split.restSpace(), means.col(0), cov},
              m_particles.size(),
              m_particles.get_allocator()
          }
    {
        check();
        EMRYS_FILTER_ASSERT(
            means.cols() == size() && means.rows() == valueDimension(restSpace())
            && "give one mean over the rest of the state per particle"
        );
        if constexpr (Shared)
        {
            m_conditionals.means() = means;
        }
        else
        {
            for (Eigen::Index i = 1; i < size(); ++i)
            {
                m_conditionals.conditional(i).setMean(means.col(i));
            }
        }
    }

    ParticleGaussianParams(Space_ space, ParticlesType particles, std::span<ConditionalType const> conditionals)
        requires(!Shared)
        : ParticleGaussianParams{std::move(space), std::move(particles), conditionals.front()}
    {
        EMRYS_FILTER_ASSERT(Eigen::Index(conditionals.size()) == size() && "give one Gaussian per particle");
        for (Eigen::Index i = 0; i < size(); ++i)
        {
            EMRYS_FILTER_ASSERT(
                conditionals[std::size_t(i)].space() == restSpace()
                && "the Gaussians must be over the rest of the state"
            );
            m_conditionals.conditional(i) = conditionals[std::size_t(i)];
        }
    }

    [[nodiscard]] Space_ const &
    space() const
    {
        return m_space;
    }

    [[nodiscard]] PinnedSpace const &
    pinnedSpace() const
    {
        return m_split.pinnedSpace();
    }

    [[nodiscard]] RestSpace const &
    restSpace() const
    {
        return m_split.restSpace();
    }

    [[nodiscard]] BlockRows
    pinnedRows() const
    {
        return m_split.pinnedRows();
    }

    [[nodiscard]] auto const &
    restIndices() const
    {
        return m_split.restIndices();
    }

    [[nodiscard]] BlockRows
    pinnedDifferenceRows() const
    {
        return m_split.pinnedDifferenceRows();
    }

    [[nodiscard]] auto const &
    restDifferenceIndices() const
    {
        return m_split.restDifferenceIndices();
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return m_particles.size();
    }

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

    void
    setCapacity(Eigen::Index capacity)
    {
        m_stamp.touchShape();
        m_particles.setCapacity(capacity);
        m_conditionals.reserve(capacity);
        m_kept.reserve(std::size_t(capacity));
    }

    [[nodiscard]] ParticlesType const &
    particles() const
    {
        return m_particles;
    }

    [[nodiscard]] auto
    particle(Eigen::Index i) const
    {
        return m_particles.state(i);
    }

    [[nodiscard]] LogWeightsType
    logWeights() const
    {
        return m_particles.logWeights();
    }

    [[nodiscard]] Scalar
    logWeight(Eigen::Index i) const
    {
        return m_particles.logWeight(i);
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    logTotalWeight() const
    {
        return m_particles.logTotalWeight();
    }

    template <ExecutionPolicy P>
    [[nodiscard]] std::expected<Scalar, WeightError>
    logTotalWeight(P && policy) const
    {
        return m_particles.logTotalWeight(std::forward<P>(policy));
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    ess() const
    {
        return m_particles.ess();
    }

    template <ExecutionPolicy P>
    [[nodiscard]] std::expected<Scalar, WeightError>
    ess(P && policy) const
    {
        return m_particles.ess(std::forward<P>(policy));
    }

    [[nodiscard]] Eigen::Index
    mostProbable() const
    {
        Eigen::Index best = 0;
        logWeights().maxCoeff(&best);
        return best;
    }

    [[nodiscard]] decltype(auto)
    conditional(Eigen::Index i) const
    {
        EMRYS_FILTER_ASSERT(i >= 0 && i < size() && "the particle is out of range");
        return m_conditionals.conditional(i);
    }

    [[nodiscard]] auto
    modes() const
        requires IsHybrid
    {
        return internal::summedModes(m_conditionals, logWeights());
    }

    [[nodiscard]] typename Moments<RestSpace, Scalar_>::MatrixType const &
    sharedCovariance() const
        requires Shared
    {
        return m_conditionals.sharedCovariance();
    }

    [[nodiscard]] std::expected<void, NumericError>
    factorize() const
    {
        if (auto const total = m_particles.logTotalWeight(); !total)
        {
            return std::unexpected{total.error()};
        }
        return m_conditionals.factorize();
    }

    [[nodiscard]] std::expected<DensityType, NumericError>
    toDensity() const &
    {
        if (auto result = factorize(); !result)
        {
            return std::unexpected{result.error()};
        }
        return DensityType{*this};
    }

    void
    toDensity() const && = delete;

    template <std::uniform_random_bit_generator G, typename R = resamplers::Systematic>
        requires Resampler<R, Scalar, G>
    [[nodiscard]] std::expected<void, WeightError>
    resample(G & rng, Eigen::Index newSize, R const & resampler = {})
    {
        return resample(execution::sequential, rng, newSize, resampler);
    }

    template <std::uniform_random_bit_generator G, typename R = resamplers::Systematic>
        requires Resampler<R, Scalar, G>
    [[nodiscard]] std::expected<void, WeightError>
    resample(G & rng, R const & resampler = {})
    {
        return resample(execution::sequential, rng, size(), resampler);
    }

    template <ExecutionPolicy P, std::uniform_random_bit_generator G, typename R = resamplers::Systematic>
        requires Resampler<R, Scalar, G>
    [[nodiscard]] std::expected<void, WeightError>
    resample(P && policy, G & rng, Eigen::Index newSize, R const & resampler = {})
    {
        auto const resampled = m_particles.resample(policy, rng, newSize, resampler);
        if (!resampled)
        {
            return resampled;
        }
        m_stamp.touchShape();
        m_conditionals.select(std::forward<P>(policy), m_particles.ancestors());
        return {};
    }

    template <ExecutionPolicy P, std::uniform_random_bit_generator G, typename R = resamplers::Systematic>
        requires Resampler<R, Scalar, G>
    [[nodiscard]] std::expected<void, WeightError>
    resample(P && policy, G & rng, R const & resampler = {})
    {
        return resample(std::forward<P>(policy), rng, size(), resampler);
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneBelow(Scalar logThreshold)
    {
        return pruneWith(internal::PruneBelowRule<Scalar>{logThreshold});
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneTail(Scalar logFraction)
    {
        return pruneWith(internal::PruneTailRule<Scalar>{logFraction});
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneToSize(Eigen::Index maxSize)
    {
        return pruneWith(internal::PruneToSizeRule{maxSize});
    }

    template <typename Derived>
    void
    setParticleStates(Eigen::MatrixBase<Derived> const & states)
    {
        EMRYS_FILTER_ASSERT(states.cols() == size() && "give one value per particle");
        m_stamp.touchValues();
        m_particles.setStates(states);
    }
 
private:

    template <Space, typename, typename, typename, typename, RealScalar, typename, typename>
    friend class RbParticleFilter;
 
    friend DensityType;  

    template <typename>
    friend class internal::ValuesCheck;

    using ConditionalsType = internal::Conditionals<RestSpace, Covariance_, Scalar_, Allocator_, ConditionalType>;

    using IndexAllocator = internal::ReBound<Allocator_, Eigen::Index>;

    template <typename Rule>
    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneWith(Rule const & rule)
    {
        m_kept.clear();
        auto const removed = internal::ParticleSetAccess::prune(m_particles, rule, [&](Eigen::Index i) {
            m_kept.push_back(i);
        });
        if (removed)
        {
            m_stamp.touchShape();
            m_conditionals.keep(std::span<Eigen::Index const>{m_kept.data(), m_kept.size()});
        }
        return removed;
    }

    [[nodiscard]] internal::ParticlePoints<Space_, Split, ParticlesType>
    points() const
    {
        return {m_space, m_split, m_particles};
    }

    [[nodiscard]] GaussianParams<Space_, Scalar>
    collapseParticles() const
    {
        if constexpr (IsHybrid)
        {
            return internal::collapseGivenHybrids(points(), m_conditionals, logWeights(), get_allocator());
        }
        else
        {
            return internal::collapseGivenPoints(points(), m_conditionals, logWeights());
        }
    }

    void
    check() const
    {
        EMRYS_FILTER_ASSERT(!m_particles.empty() && "a Rao-Blackwellized state needs a particle");
        EMRYS_FILTER_ASSERT(m_particles.valueDimension() == pinnedRows().size && "the particles do not fit the block");
    }

    template <ExecutionPolicy P, typename Derived>
    [[nodiscard]] std::pair<Scalar, bool>
    reweight(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods, ConditionalsType & posterior)
    {
        return reweightThen(std::forward<P>(policy), logLikelihoods, [&] {
            m_conditionals.swap(posterior);
        });
    }

    template <ExecutionPolicy P, typename Derived>
    [[nodiscard]] std::pair<Scalar, bool>
    reweight(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods)
    {
        return reweightThen(std::forward<P>(policy), logLikelihoods, [] {});
    }

    template <ExecutionPolicy P, typename Derived, typename Place>
    [[nodiscard]] std::pair<Scalar, bool>
    reweightThen(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods, Place const & place)
    {
        EMRYS_FILTER_ASSERT(logLikelihoods.size() == size() && "give one log-likelihood per particle");
        auto const logTotal = m_particles.logTotalWeight(policy);
        EMRYS_FILTER_ASSERT(logTotal && "the particles' weights must be valid");
        if (!logTotal)
        {
            return {Eigen::NumTraits<Scalar>::quiet_NaN(), false};
        }
        auto const logWeights = m_particles.logWeights() - *logTotal + logLikelihoods.derived();
        // 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");
        if (logLikelihood == -Eigen::NumTraits<Scalar>::infinity())
        {
            return {logLikelihood, false};
        }
        m_stamp.touchShape();
        internal::ParticleSetAccess::setNormalizedLogWeights(m_particles, logWeights - logLikelihood);
        place();
        return {logLikelihood, true};
    }
 
    Space_ m_space;                                     
    Split m_split;                                      
    ParticlesType m_particles;                          
    internal::Buffer<Eigen::Index, Allocator_> m_kept;  
 
    ConditionalsType m_conditionals;  
    internal::Stamp m_stamp;          
};

template <
    Space Space_,
    typename Path_,
    typename Covariance_,
    RealScalar Scalar_,
    typename Allocator_,
    typename Component_>
class ParticleGaussian
{
public:
 
    using ParamsType =
        ParticleGaussianParams<Space_, Path_, Covariance_, Scalar_, Allocator_, Component_>;  
    using Scalar = Scalar_;                                                                   
    using SpaceType = Space_;                                                                 
    using PinnedSpace = typename ParamsType::PinnedSpace;          
    using RestSpace = typename ParamsType::RestSpace;              
    using ConditionalType = typename ParamsType::ConditionalType;  
    using ParticlesType = typename ParamsType::ParticlesType;      
    using MomentsType = Moments<Space_, Scalar_>;                  
    using allocator_type = Allocator_;                             
 
    static constexpr bool IsParams = false;  

    [[nodiscard]] Space_ const &
    space() const
    {
        return source().space();
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return source().size();
    }

    [[nodiscard]] ParticlesType const &
    particles() const
    {
        return source().particles();
    }

    [[nodiscard]] decltype(auto)
    conditional(Eigen::Index i) const
    {
        return source().conditional(i);
    }

    [[nodiscard]] GaussianParams<Space_, Scalar_>
    collapsed() const
    {
        return source().collapseParticles();
    }

    [[nodiscard]] MomentsType
    moments() const
    {
        return collapsed().moments();
    }
 
private:
 
    friend ParamsType;  

    explicit ParticleGaussian(ParamsType const & source)
        : m_check{source}
    {}

    [[nodiscard]] ParamsType const &
    source() const
    {
        return m_check.get();
    }
 
    internal::ValuesCheck<ParamsType> m_check;  
};

template <typename P, Space S, typename Scalar, typename Allocator, typename Covariance = SharedCovariance>
[[nodiscard]] auto
particleGaussian(
    S space,
    ParticleSet<typename internal::BlockSplit<S, P>::PinnedSpace, Scalar, Allocator> particles,
    GaussianParams<typename internal::BlockSplit<S, P>::RestSpace, Scalar> const & conditional,
    Covariance covariance = {}
)
{
    return ParticleGaussianParams<S, P, Covariance, Scalar, Allocator>{
        std::move(space),
        std::move(particles),
        conditional,
        covariance
    };
}

template <
    typename P,
    Space S,
    typename Scalar,
    typename Allocator,
    typename Component,
    typename ModeSpace,
    typename HybridAllocator>
[[nodiscard]] auto
particleGaussian(
    S space,
    ParticleSet<typename internal::BlockSplit<S, P>::PinnedSpace, Scalar, Allocator> particles,
    HybridParams<Component, ModeSpace, HybridAllocator> const & conditional
)
{
    return ParticleGaussianParams<
        S,
        P,
        PerPointCovariance,
        Scalar,
        Allocator,
        Hybrid<Component, ModeSpace, HybridAllocator>>{std::move(space), std::move(particles), conditional};
}

template <
    typename P,
    typename Covariance = SharedCovariance,
    typename S,
    typename Scalar,
    std::uniform_random_bit_generator G,
    typename Allocator = DefaultAllocator>
[[nodiscard]] auto
drawParticleGaussian(
    GaussianParams<S, Scalar> const & joint,
    Eigen::Index count,
    G & rng,
    Covariance = {},
    Allocator const & allocator = {}
)
{
    using Result = ParticleGaussianParams<S, P, Covariance, Scalar, Allocator>;
    using PinnedSamples = SampleSet<typename Result::PinnedSpace, Scalar>;
    using PinnedNormals = SampleSet<DifferenceSpace<typename Result::PinnedSpace>, Scalar>;
    EMRYS_FILTER_ASSERT(count > 0 && "draw at least one particle");
    auto const split = internal::BlockSplit<S, P>{joint.space()};
    auto const conditioning = internal::pinnedConditioning(split, joint.covarianceLower());
    PinnedNormals const z = internal::standardNormal<PinnedNormals>(split.pinnedDifferenceRows().size, count, rng);
    auto pinned = PinnedSamples{split.pinnedRows().size, count};
    auto means = typename Result::MeansType{split.restIndices().size(), count};
    for (Eigen::Index i = 0; i < count; ++i)
    {
        internal::drawPinned(split, conditioning, joint.mean(), z.col(i), pinned.col(i), means.col(i));
    }
    auto particles = typename Result::ParticlesType{split.pinnedSpace(), pinned, allocator};
    return Result{joint.space(), std::move(particles), means, conditioning.cov, Covariance{}};
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <
    typename Space_,
    typename Path_,
    typename Covariance_,
    typename Scalar_,
    typename Allocator_,
    typename Component_>
void
describe(
    std::ostream & stream,
    ParticleGaussianParams<Space_, Path_, Covariance_, Scalar_, Allocator_, Component_> const & state,
    Indent indent
)
{
    Eigen::Index const dimension = valueDimension(state.space());
    stream << indent << "Rao-Blackwellized particle set of ";
    printCount(stream, state.size(), "particle");
    printDimension(stream, dimension, differenceDimension(state.space()));
    auto const ess = state.ess();
    auto const density = state.toDensity();
    if (ess)
    {
        stream << ", effective sample size " << *ess;
    }
    if (density)
    {
        printMoments(stream, indent.deeper(), density->moments());
    }
    else
    {
        stream << '\n';
        describe(stream, density.error(), indent.deeper());
    }
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO