GaussianMixtureHistory.h

Go to the documentation of this file.

#pragma once

 
#include <cstddef>
#include <list>
#include <optional>
#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/Weights.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Filters/MixtureFilter.h>
#include <EmrysFilter/Smoothers/GaussianHistory.h>
#include <EmrysFilter/Smoothers/Lag.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{
 
namespace internal
{

template <typename GaussianType, typename Allocator>
struct MixtureSmoothingScratch
{
    using Scalar = typename GaussianType::Scalar;  
 
    Buffer<GaussianType, Allocator> pairs;     
    Buffer<Scalar, Allocator> pairLogWeights;  
    Buffer<Eigen::Index, Allocator> offsets;   
    Buffer<Scalar, Allocator> logWeights;      

    explicit MixtureSmoothingScratch(Allocator const & allocator = {})
        : pairs(ReBound<Allocator, GaussianType>(allocator))
        , pairLogWeights(ReBound<Allocator, Scalar>(allocator))
        , offsets(ReBound<Allocator, Eigen::Index>(allocator))
        , logWeights(ReBound<Allocator, Scalar>(allocator))
    {}
};

template <
    ExecutionPolicy P,
    typename Joint,
    typename Gains,
    typename Next,
    typename NextLineage,
    typename Smoothed,
    typename Scratch>
void
smoothMixtureStep(
    P && policy,
    Joint const & joint,
    Gains const & gains,
    Next const & next,
    NextLineage const & nextLineage,
    Smoothed & smoothed,
    Scratch & scratch
)
{
    using Scalar = typename Joint::Scalar;
    using GaussianType = typename Joint::JointType::GaussianType;
    using Index = Eigen::Index;
    Scalar const minusInfinity = -Eigen::NumTraits<Scalar>::infinity();
    Index const parents = Index(joint.joints.size());
    Index const children = next.size();
    EMRYS_FILTER_ASSERT(smoothed.size() == parents && "the smoothed mixture must have a component per parent");
    EMRYS_FILTER_ASSERT(Index(gains.size()) == parents && "the gains must have one per parent");
    EMRYS_FILTER_ASSERT(
        nextLineage.parentCount() == parents && nextLineage.size() == children
        && "the lineage must have a row per component at the step's end and a share per component at its start"
    );
    auto const logTotal = next.logTotalWeight();
    EMRYS_FILTER_ASSERT(logTotal && "the smoothed mixture at the step's end must have valid weights");
    auto const logChildWeight = [&](Index p) {
        return next[p].logWeight - *logTotal;
    };
    auto const paired = [&](Index p, Index n) {
        return logChildWeight(p) > minusInfinity && nextLineage.logShare(p, n) > minusInfinity;
    };
 
    // The pairs (p, n) of child p at the end and its parent n at the start, grouped by parent.
    scratch.offsets.assign(std::size_t(parents + 1), 0);
    for (Index n = 0; n < parents; ++n)
    {
        Index count = 0;
        for (Index p = 0; p < children; ++p)
        {
            count += paired(p, n) ? 1 : 0;
        }
        scratch.offsets[std::size_t(n + 1)] = scratch.offsets[std::size_t(n)] + count;
    }
    std::size_t const total = std::size_t(scratch.offsets.back());
    scratch.pairs.resize(total);
    scratch.pairLogWeights.resize(total);
    scratch.logWeights.resize(std::size_t(parents));
 
    internal::forEachIndex(policy, parents, [&](Index n) {
        auto const & parentJoint = joint.joints[std::size_t(n)];
        Index const first = scratch.offsets[std::size_t(n)];
        Index slot = first;
        for (Index p = 0; p < children; ++p)
        {
            if (paired(p, n))
            {
                // An RTS step from parent n to child p, weighted as (11) with the normalized share of (15).
                scratch.pairs[std::size_t(slot)] = smoothStep(parentJoint, gains[std::size_t(n)], next[p].component);
                scratch.pairLogWeights[std::size_t(slot)] = logChildWeight(p) + nextLineage.logShare(p, n);
                ++slot;
            }
        }
        auto weight = LogSumExp<Scalar>{};
        for (Index i = first; i < slot; ++i)
        {
            weight.add(scratch.pairLogWeights[std::size_t(i)]);
        }
        scratch.logWeights[std::size_t(n)] = weight.value();
    });
    smoothed.forEach(policy, [&](auto & entry) {
        Index const n = entry.index();
        Index const first = scratch.offsets[std::size_t(n)];
        Index const count = scratch.offsets[std::size_t(n + 1)] - first;
        if (count == 0)
        {
            // No child at the end descends from it: it keeps its filtered Gaussian, with no weight.
            entry.setComponent(joint.joints[std::size_t(n)].filtered);
            return;
        }
        // The pairs that share parent n are merged into one, the group merge of Algorithm 2.
        entry.setComponent(
            ComponentTraits<GaussianType>::collapse(
                count,
                [&](Index i) -> GaussianType const & {
                    return scratch.pairs[std::size_t(first + i)];
                },
                [&](Index i) {
                    return scratch.pairLogWeights[std::size_t(first + i)];
                }
            )
        );
    });
    for (Index n = 0; n < parents; ++n)
    {
        smoothed.setLogWeight(n, scratch.logWeights[std::size_t(n)]);
    }
    // A child with weight descends from some parent with a share, so some parent has weight.
    [[maybe_unused]] auto const normalized = smoothed.normalize();
    EMRYS_FILTER_ASSERT(normalized && "nothing at the step's end descends from a component at its start");
}
 
}  // namespace internal

template <Space Space_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
class GaussianMixtureHistory
{
public:
 
    using Scalar = Scalar_;                                                        
    using SpaceType = Space_;                                                      
    using JointType = GaussianMixturePredictedJoint<Space_, Scalar_, Allocator_>;  
    using SmoothedType = GaussianMixtureParams<Space_, Scalar_, Allocator_>;       
    using allocator_type = Allocator_;  
    using Index = Eigen::Index;         

    class Step
    {
    public:
 
        using allocator_type = Allocator_;  

        using GainType = typename GaussianParams<Space_, Scalar_>::MatrixType;

        using GainAllocator = internal::ReBound<Allocator_, GainType>;

        explicit Step(JointType const & recorded, Allocator_ const & allocator = {})
            : joint{copyJoint(recorded, allocator)}
            , smoothed{joint.joints.front().filtered.space(), allocator}
            , gains(GainAllocator(allocator))
        {
            smoothed.reserve(Index(joint.joints.size()));
            resetSmoothed();
            computeGains();
        }

        explicit Step(JointType && recorded, Allocator_ const & allocator = {})
            : joint{moveJoint(std::move(recorded), allocator)}
            , smoothed{joint.joints.front().filtered.space(), allocator}
            , gains(GainAllocator(allocator))
        {
            smoothed.reserve(Index(joint.joints.size()));
            resetSmoothed();
            computeGains();
        }

        Step(Step const & other, Allocator_ const & allocator)
            : joint{copyJoint(other.joint, allocator)}
            , smoothed{other.smoothed, allocator}
            , gains(other.gains, GainAllocator(allocator))
        {}

        Step(Step && other, Allocator_ const & allocator)
            : joint{moveJoint(std::move(other.joint), allocator)}
            , smoothed{std::move(other.smoothed), allocator}
            , gains(std::move(other.gains), GainAllocator(allocator))
        {}
 
        Step(Step const &) = default;  
        Step(Step &&) = default;       

        Step &
        operator=(Step const &) = default;

        Step &
        operator=(Step &&) = default;
 
        ~Step() = default;

        void
        resetSmoothed()
        {
            smoothed.clear();
            for (std::size_t n = 0; n < joint.joints.size(); ++n)
            {
                static_cast<void>(smoothed.add(joint.logWeights[n], joint.joints[n].filtered));
            }
        }

        void
        computeGains()
        {
            gains.resize(joint.joints.size());
            for (std::size_t n = 0; n < gains.size(); ++n)
            {
                gains[n] = internal::rtsGain(joint.joints[n]);
            }
        }
 
        JointType joint;  

        SmoothedType smoothed;

        internal::Buffer<GainType, Allocator_> gains;
 
    private:

        [[nodiscard]] static JointType
        copyJoint(JointType const & recorded, Allocator_ const & allocator)
        {
            return {
                typename JointType::LogWeightsType(recorded.logWeights, allocator),
                typename JointType::JointsType(recorded.joints, allocator),
                typename JointType::LineageType(recorded.lineage, allocator)
            };
        }

        [[nodiscard]] static JointType
        moveJoint(JointType && recorded, Allocator_ const & allocator)
        {
            return {
                typename JointType::LogWeightsType(std::move(recorded.logWeights), allocator),
                typename JointType::JointsType(std::move(recorded.joints), allocator),
                typename JointType::LineageType(std::move(recorded.lineage), allocator)
            };
        }
    };
 
private:
 
    using StepAllocator = internal::ReBound<Allocator_, Step>;  
    using ListType = std::list<Step, StepAllocator>;            

    using ScratchType = internal::MixtureSmoothingScratch<GaussianParams<Space_, Scalar_>, Allocator_>;
 
public:
 
    using const_iterator = typename ListType::const_iterator;  

    GaussianMixtureHistory() = default;

    explicit GaussianMixtureHistory(Allocator_ const & allocator)
        : m_steps(StepAllocator(allocator))
        , m_scratch{allocator}
    {}

    explicit GaussianMixtureHistory(Lag lag, Allocator_ const & allocator = {})
        : m_steps(StepAllocator(allocator))
        , m_scratch{allocator}
        , m_lag{lag.steps}
    {
        EMRYS_FILTER_ASSERT(lag.steps > 0 && "a fixed-lag history keeps at least one step");
    }

    GaussianMixtureHistory(GaussianMixtureHistory const & other)
        : GaussianMixtureHistory{other, other.get_allocator()}
    {}

    GaussianMixtureHistory(GaussianMixtureHistory const & other, Allocator_ const & allocator)
        : m_steps(other.m_steps, StepAllocator(allocator))
        , m_scratch{allocator}
        , m_lag{other.m_lag}
    {}

    GaussianMixtureHistory(GaussianMixtureHistory && other) = default;

    GaussianMixtureHistory(GaussianMixtureHistory && other, Allocator_ const & allocator)
        : m_steps(std::move(other.m_steps), StepAllocator(allocator))
        , m_scratch{allocator}
        , m_lag{other.m_lag}
    {}

    GaussianMixtureHistory &
    operator=(GaussianMixtureHistory const & other) = default;

    GaussianMixtureHistory &
    operator=(GaussianMixtureHistory && other) = default;
 
    ~GaussianMixtureHistory() = default;

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

    [[nodiscard]] std::optional<Index>
    lag() const
    {
        return m_lag > 0 ? std::optional<Index>{m_lag} : std::nullopt;
    }

    [[nodiscard]] Index
    size() const
    {
        return Index(m_steps.size());
    }

    [[nodiscard]] bool
    empty() const
    {
        return m_steps.empty();
    }

    [[nodiscard]] Step const &
    front() const
    {
        EMRYS_FILTER_ASSERT(!empty() && "the history is empty");
        return m_steps.front();
    }

    [[nodiscard]] Step const &
    back() const
    {
        EMRYS_FILTER_ASSERT(!empty() && "the history is empty");
        return m_steps.back();
    }

    [[nodiscard]] const_iterator
    begin() const
    {
        return m_steps.begin();
    }

    [[nodiscard]] const_iterator
    end() const
    {
        return m_steps.end();
    }

    void
    push(JointType const & joint)
    {
        record(joint);
    }

    void
    push(JointType && joint)
    {
        record(std::move(joint));
    }

    template <typename LatestAllocator>
    void
    smooth(GaussianMixtureParams<Space_, Scalar_, LatestAllocator> const & latest)
    {
        smooth(execution::sequential, latest);
    }

    template <ExecutionPolicy P, typename LatestAllocator>
    void
    smooth(P && policy, GaussianMixtureParams<Space_, Scalar_, LatestAllocator> const & latest)
    {
        auto step = m_steps.rbegin();
        if (step == m_steps.rend())
        {
            return;
        }
        assertFollows(latest.lineage(), *step);
        internal::smoothMixtureStep(
            policy,
            step->joint,
            step->gains,
            latest,
            latest.lineage(),
            step->smoothed,
            m_scratch
        );
        for (auto next = step++; step != m_steps.rend(); next = step++)
        {
            internal::smoothMixtureStep(
                policy,
                step->joint,
                step->gains,
                next->smoothed,
                next->joint.lineage,
                step->smoothed,
                m_scratch
            );
        }
    }

    void
    clear()
    {
        m_steps.clear();
    }
 
private:

    template <typename Joint>
    void
    record(Joint && joint)
    {
        EMRYS_FILTER_ASSERT(
            !joint.joints.empty() && joint.logWeights.size() == joint.joints.size() && "the joint has no components"
        );
        if (!empty())
        {
            assertFollows(joint.lineage, back());
        }
        if (m_lag > 0 && size() == m_lag)
        {
            m_steps.splice(m_steps.end(), m_steps, m_steps.begin());
            Step & step = m_steps.back();
            step.joint.logWeights.assign(joint.logWeights.begin(), joint.logWeights.end());
            step.joint.joints.assign(joint.joints.begin(), joint.joints.end());
            step.joint.lineage = joint.lineage;
            step.resetSmoothed();
            step.computeGains();
        }
        else
        {
            m_steps.emplace_back(std::forward<Joint>(joint));
        }
    }

    template <typename Lineage>
    static void
    assertFollows([[maybe_unused]] Lineage const & lineage, [[maybe_unused]] Step const & step)
    {
        EMRYS_FILTER_ASSERT(
            lineage.count() == step.joint.lineage.count() + 1
            && lineage.parentCount() == Index(step.joint.joints.size())
            && "the lineage must start at the newest record's predictJoint(): one predictJoint() apart, and no copy of "
               "the state between"
        );
    }
 
    ListType m_steps;       
    ScratchType m_scratch;  
    Index m_lag = 0;        
};
 
}  // namespace Emrys