KalmanFilter.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <expected>
#include <optional>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Distributions/Component.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/GaussianForm.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Filters/FilterUpkeep.h>
#include <EmrysFilter/Filters/Recursion.h>
#include <EmrysFilter/Models/FormSteps.h>
#include <EmrysFilter/Models/Measurement.h>
#include <EmrysFilter/Models/ProcessModel.h>
#include <EmrysFilter/Models/Reduction.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Models/Transform.h>
#include <EmrysFilter/Core/Print.h>
 
namespace Emrys
{
 
namespace internal
{

template <typename C>
inline constexpr bool isKalmanState = Component<C> || (GaussianForm<Params<C>> && std::same_as<Density<Params<C>>, C>);
 
}  // namespace internal

template <typename C>
concept KalmanState = requires { requires internal::isKalmanState<C>; };
 
namespace internal
{

template <typename C, typename M>
[[nodiscard]] consteval auto
kalmanUpdateType()
{
    if constexpr (isOtherGaussianForm<C>)
    {
        return std::type_identity<decltype(updateInPlace<false>(std::declval<C &>(), std::declval<M const &>()))>{};
    }
    else
    {
        using Measured =
            typename decltype(Emrys::update(std::declval<C const &>(), std::declval<M const &>()))::value_type;
        return std::type_identity<std::expected<
            FilterUpdateWithScratch<
                typename C::SpaceType,
                typename Measured::MeasurementSpaceType,
                Measured::BranchCount,
                typename C::Scalar,
                DefaultAllocator>,
            GaussianError>>{};
    }
}
 
}  // namespace internal

template <KalmanState Component_, typename Model_, typename Method_ = DefaultMethodFor<Model_>>
class KalmanRecursion
{
    static_assert(
        internal::predictsBy<Model_, Method_>,
        "a filter with a fixed-step model predicts by the model's method"
    );
    static_assert(internal::checkPredictionMethod<Method_>());

    static constexpr bool IsForm = internal::isOtherGaussianForm<Params<Component_>>;
 
public:
 
    using Scalar = typename Params<Component_>::Scalar;        
    using SpaceType = typename Params<Component_>::SpaceType;  
    using ModelType = Model_;                                  
    using MethodType = Method_;                                
    using StateType = Params<Component_>;                      
    static constexpr bool Reentrant = true;                    

    template <Measurement M>
    using UpdateType = typename decltype(internal::kalmanUpdateType<Params<Component_>, M>())::type;

    template <Measurement M>
    using ResultType = typename UpdateType<M>::value_type;

    explicit KalmanRecursion(Model_ model, MethodFor<Model_, Method_> method = {})
        : m_model{std::move(model)}
        , m_method{internal::predictionMethod(m_model, std::move(method))}
    {}

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

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

    [[nodiscard]] Method_ const &
    method() const
    {
        return m_method;
    }

    [[nodiscard]] bool
    partitions(SpaceType const & space) const
    {
        return m_model.partitions(space);
    }

    template <typename... Inputs>
        requires(!isFixedStep<Model_>)
    [[nodiscard]] auto
    predict(Params<Component_> & state, TimeStep<Scalar> auto const & dt, Inputs const &... inputs) const
    {
        return internal::predictInPlace(m_method, state, m_model, internal::toScalarTime<Scalar>(dt), inputs...);
    }

    template <typename... Inputs>
        requires isFixedStep<Model_>
    [[nodiscard]] auto
    predict(Params<Component_> & state, Inputs const &... inputs) const
    {
        return internal::predictInPlace(state, m_model, inputs...);
    }

    template <typename... Inputs>
        requires(!isFixedStep<Model_> && !IsForm)
    [[nodiscard]] PredictedJoint<SpaceType, Scalar>
    predictJoint(Params<Component_> & state, TimeStep<Scalar> auto const & dt, Inputs const &... inputs) const
    {
        auto joint = Emrys::predictJoint(m_method, state, m_model, dt, inputs...);
        state = joint.predicted;
        return joint;
    }

    template <typename... Inputs>
        requires(isFixedStep<Model_> && !IsForm)
    [[nodiscard]] PredictedJoint<SpaceType, Scalar>
    predictJoint(Params<Component_> & state, Inputs const &... inputs) const
    {
        auto joint = Emrys::predictJoint(state, m_model, inputs...);
        state = joint.predicted;
        return joint;
    }

    template <Measurement M, typename Policy = reductions::Collapse>
        requires(
            (IsForm || ReductionPolicy<Policy, internal::MeasuredUpdate<Params<Component_>, M>, Params<Component_>>)
            && !internal::isLinearizedAbout<Policy>
        )
    [[nodiscard]] UpdateType<M>
    update(Params<Component_> & state, M const & measurement, Policy const & policy = {}) const
    {
        return updateAbout(state, measurement, internal::NoReference{}, policy);
    }

    template <Measurement M, typename Reference, typename Policy = reductions::Collapse>
        requires(IsForm || ReductionPolicy<Policy, internal::MeasuredUpdate<Params<Component_>, M>, Params<Component_>>)
    [[nodiscard]] UpdateType<M>
    update(
        Params<Component_> & state,
        M const & measurement,
        LinearizedAbout<Reference> const & about,
        Policy const & policy = {}
    ) const
    {
        return updateAbout(state, measurement, about, policy);
    }

    template <Measurement M, typename Policy, typename Result>
        requires(
            (IsForm || ReductionPolicy<Policy, internal::MeasuredUpdate<Params<Component_>, M>, Params<Component_>>)
            && !internal::isLinearizedAbout<Policy>
        )
    [[nodiscard]] std::expected<void, GaussianError>
    update(Params<Component_> & state, M const & measurement, Policy const & policy, Result & result) const
    {
        return updateAboutInto(state, measurement, internal::NoReference{}, policy, result);
    }

    template <Measurement M, typename Reference, typename Policy, typename Result>
        requires(IsForm || ReductionPolicy<Policy, internal::MeasuredUpdate<Params<Component_>, M>, Params<Component_>>)
    [[nodiscard]] std::expected<void, GaussianError>
    update(
        Params<Component_> & state,
        M const & measurement,
        LinearizedAbout<Reference> const & about,
        Policy const & policy,
        Result & result
    ) const
    {
        return updateAboutInto(state, measurement, about, policy, result);
    }
 
private:

    template <Measurement M, typename Reference, typename Policy>
    [[nodiscard]] UpdateType<M>
    updateAbout(
        Params<Component_> & state,
        M const & measurement,
        Reference const & reference,
        Policy const & policy
    ) const
    {
        constexpr bool Statistics = internal::policyReadsStatistics<Policy>;
        if constexpr (IsForm)
        {
            static_assert(
                internal::isOneBranchInPlacePolicy<Policy>,
                "a square-root or information form has one branch and takes Collapse, Select, or a Gate over them; "
                "a policy of the user's own is the covariance form's (decisions.md [gaussian-forms])"
            );
            return internal::updateInPlace<Statistics>(
                state,
                measurement,
                reference,
                internal::oneBranchAcceptor(policy)
            );
        }
        else
        {
            using Measured = typename decltype(internal::updateBranches<Statistics>(state, measurement))::value_type;
            if constexpr (internal::isOneBranchInPlacePolicy<Policy> && Measured::BranchCount == 1)
            {
                // One branch is its own collapse, and is the new state, conditioned where it is
                // (decisions.md [borrowed-intermediates]); a gate decides on its NIS before the state
                // is written, so a rejection copies nothing [update-outcomes].
                return internal::updateInPlace<Statistics>(
                    state,
                    measurement,
                    reference,
                    internal::oneBranchAcceptor(policy)
                );
            }
            else if constexpr (Measured::BranchCount == std::dynamic_extent)
            {
                auto result = UpdateType<M>{};
                auto const reduced = reduceBranches(state, measurement, reference, policy, *result);
                if (!reduced)
                {
                    return std::unexpected{reduced.error()};
                }
                return result;
            }
            else
            {
                auto measured = internal::updateBranches<Statistics>(state, measurement, reference);
                if (!measured)
                {
                    return std::unexpected{measured.error()};
                }
                // The policy writes what it keeps to the state, which it never reads.
                auto out = internal::OneComponentOutput<Params<Component_>, Scalar>{state};
                internal::reduceInto(policy, std::move(measured->branches), out);
                auto result = internal::filterUpdate(std::move(*measured), out.added);
                internal::MixtureEvidenceAccess::setLogKeptWeight(result, out.logWeight);
                return result;
            }
        }
    }

    template <Measurement M, typename Reference, typename Policy, typename Result>
    [[nodiscard]] std::expected<void, GaussianError>
    updateAboutInto(
        Params<Component_> & state,
        M const & measurement,
        Reference const & reference,
        Policy const & policy,
        Result & result
    ) const
    {
        if constexpr (!IsForm && Result::BranchCount == std::dynamic_extent)
        {
            return reduceBranches(state, measurement, reference, policy, result);
        }
        else
        {
            auto updated = updateAbout(state, measurement, reference, policy);
            if (!updated)
            {
                return std::unexpected{updated.error()};
            }
            result = std::move(*updated);
            return {};
        }
    }

    template <Measurement M, typename Reference, typename Policy, typename Result>
    [[nodiscard]] std::expected<void, GaussianError>
    reduceBranches(
        Params<Component_> & state,
        M const & measurement,
        Reference const & reference,
        Policy const & policy,
        Result & result
    ) const
    {
        constexpr bool Statistics = internal::policyReadsStatistics<Policy>;
        auto const measured = internal::updateBranchesInto<Statistics>(state, measurement, result.measured, reference);
        if (!measured)
        {
            return measured;
        }
        // The policy writes what it keeps to the state, which it never reads.
        auto out = internal::OneComponentOutput<Params<Component_>, Scalar>{state};
        internal::reduceInto(policy, std::move(result.measured.branches), out);
        internal::filterUpdateInto(result.measured, out.added, result);
        internal::MixtureEvidenceAccess::setLogKeptWeight(result, out.logWeight);
        return {};
    }
 
    Model_ m_model;                                   
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Method_ m_method;  
};
 
namespace internal
{

template <KalmanState Component_, typename Model, typename Method, typename Allocator>
struct RecursionFor
{
    using type = KalmanRecursion<Component_, Model, Method>;  
};

template <KalmanState Component_, typename Model, typename Method, typename Allocator>
using RecursionForType = typename RecursionFor<Component_, Model, Method, Allocator>::type;
 
}  // namespace internal

template <
    KalmanState Component_,
    typename Model_,
    typename Method_ = DefaultMethodFor<Model_>,
    RecursionOver<Component_> Recursion_ = KalmanRecursion<Component_, Model_, Method_>,
    typename Upkeep_ = upkeep::None>
    requires std::same_as<typename Recursion_::ModelType, Model_>
          && std::same_as<typename Recursion_::MethodType, Method_>
class KalmanFilter
    : public internal::HeldUpkeep<KalmanFilter<Component_, Model_, Method_, Recursion_, Upkeep_>, Upkeep_>
{
    using UpkeepBase = internal::HeldUpkeep<KalmanFilter, Upkeep_>;  

    static constexpr bool IsForm = internal::isOtherGaussianForm<Params<Component_>>;
 
public:
 
    using Scalar = typename Params<Component_>::Scalar;        
    using SpaceType = typename Params<Component_>::SpaceType;  
    using ModelType = Model_;                                  
    using MethodType = Method_;                                
    using StateType = Params<Component_>;                      
    using RecursionType = Recursion_;                          

    KalmanFilter(
        Model_ model,
        Params<Component_> state,
        MethodFor<Model_, Method_> method = {},
        std::type_identity_t<Upkeep_> upkeep = {}
    )
        : UpkeepBase{std::move(upkeep)}
        , m_recursion{std::move(model), std::move(method)}
        , m_state{std::move(state)}
    {
        internal::assertPartitions(m_recursion.model(), m_state.space());
    }

    KalmanFilter(Model_ model, Params<Component_> state, std::type_identity_t<Upkeep_> upkeep)
        requires(!std::same_as<Upkeep_, Emrys::upkeep::None>)
        : KalmanFilter{std::move(model), std::move(state), MethodFor<Model_, Method_>{}, std::move(upkeep)}
    {}

    KalmanFilter(Recursion_ recursion, StateType state, std::type_identity_t<Upkeep_> upkeep = {})
        : UpkeepBase{std::move(upkeep)}
        , m_recursion{std::move(recursion)}
        , m_state{std::move(state)}
    {
        internal::assertPartitions(m_recursion.model(), m_state.space());
    }

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

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

    void
    setState(StateType state)
    {
        internal::assertPartitions(m_recursion.model(), state.space());
        m_state = std::move(state);
    }

    [[nodiscard]] RecursionType const &
    recursion() const
    {
        return m_recursion;
    }

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

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

    [[nodiscard]] Method_ const &
    method() const
    {
        return m_recursion.method();
    }

    template <typename... Inputs>
        requires(!isFixedStep<Model_>)
    [[nodiscard]] auto
    predict(TimeStep<Scalar> auto const & dt, Inputs const &... inputs)
    {
        return m_recursion.predict(m_state, internal::toScalarTime<Scalar>(dt), inputs...);
    }

    template <typename... Inputs>
        requires(isFixedStep<Model_> && !internal::leadsWithSteps<Inputs...>)
    [[nodiscard]] auto
    predict(Inputs const &... inputs)
    {
        internal::rejectTimeStep<Inputs...>();
        return m_recursion.predict(m_state, inputs...);
    }

    template <typename... Inputs>
        requires isFixedStep<Model_>
    [[nodiscard]] auto
    predict(Steps steps, Inputs const &... inputs)
    {
        EMRYS_FILTER_ASSERT(steps.count >= 0 && "the number of steps must be nonnegative");
        if constexpr (IsForm)
        {
            using Result = decltype(m_recursion.predict(m_state, inputs...));
            for (Eigen::Index i = 0; i < steps.count; ++i)
            {
                if (Result result = m_recursion.predict(m_state, inputs...); !result)
                {
                    return result;
                }
            }
            return Result{};
        }
        else
        {
            for (Eigen::Index i = 0; i < steps.count; ++i)
            {
                m_recursion.predict(m_state, inputs...);
            }
        }
    }

    template <typename... Inputs>
        requires(!isFixedStep<Model_> && !IsForm)
    [[nodiscard]] PredictedJoint<SpaceType, Scalar>
    predictJoint(TimeStep<Scalar> auto const & dt, Inputs const &... inputs)
    {
        return m_recursion.predictJoint(m_state, dt, inputs...);
    }

    template <typename... Inputs>
        requires(isFixedStep<Model_> && !IsForm && !internal::leadsWithSteps<Inputs...>)
    [[nodiscard]] PredictedJoint<SpaceType, Scalar>
    predictJoint(Inputs const &... inputs)
    {
        internal::rejectTimeStep<Inputs...>();
        return m_recursion.predictJoint(m_state, inputs...);
    }

    template <Measurement M, typename Policy = reductions::Collapse>
        requires(!internal::isLinearizedAbout<Policy>)
    [[nodiscard]] auto
    update(M const & measurement, Policy const & policy = {})
    {
        return m_recursion.update(m_state, measurement, policy);
    }

    template <Measurement M, typename Reference, typename Policy = reductions::Collapse>
    [[nodiscard]] auto
    update(M const & measurement, LinearizedAbout<Reference> const & about, Policy const & policy = {})
    {
        return m_recursion.update(m_state, measurement, about, policy);
    }

    template <Measurement M, typename Policy, typename Result>
        requires(!internal::isLinearizedAbout<Policy>)
    [[nodiscard]] std::expected<void, GaussianError>
    update(M const & measurement, Policy const & policy, Result & result)
    {
        return m_recursion.update(m_state, measurement, policy, result);
    }

    template <Measurement M, typename Reference, typename Policy, typename Result>
    [[nodiscard]] std::expected<void, GaussianError>
    update(M const & measurement, LinearizedAbout<Reference> const & about, Policy const & policy, Result & result)
    {
        return m_recursion.update(m_state, measurement, about, policy, result);
    }
 
private:
 
    RecursionType m_recursion;  
    StateType m_state;          
};

template <typename P, typename Model>
    requires KalmanState<Density<P>> && (!Recursion<Model>)
KalmanFilter(Model, P) -> KalmanFilter<
    Density<P>,
    internal::BoundStep<Model, typename P::SpaceType, typename P::Scalar>,
    DefaultMethodFor<Model>>;
 
template <typename P, Recursion R>
    requires KalmanState<Density<P>> && std::same_as<typename R::StateType, P>
KalmanFilter(R, P) -> KalmanFilter<Density<P>, typename R::ModelType, typename R::MethodType, R>;
 
// A third argument is a method or an upkeep, told apart by its type.
template <typename P, typename Model, typename X>
    requires KalmanState<Density<P>> && (!Recursion<Model>)
KalmanFilter(Model, P, X) -> KalmanFilter<
    Density<P>,
    internal::BoundStep<Model, typename P::SpaceType, typename P::Scalar>,
    internal::MethodArgumentFor<Model, X>,
    KalmanRecursion<
        Density<P>,
        internal::BoundStep<Model, typename P::SpaceType, typename P::Scalar>,
        internal::MethodArgumentFor<Model, X>>,
    internal::UpkeepArgumentFor<X>>;
 
template <typename P, typename Model, typename Method, typename U>
    requires KalmanState<Density<P>>
KalmanFilter(Model, P, Method, U) -> KalmanFilter<
    Density<P>,
    internal::BoundStep<Model, typename P::SpaceType, typename P::Scalar>,
    Method,
    KalmanRecursion<Density<P>, internal::BoundStep<Model, typename P::SpaceType, typename P::Scalar>, Method>,
    U>;
 
template <typename P, Recursion R, typename U>
    requires KalmanState<Density<P>> && std::same_as<typename R::StateType, P> && internal::UpkeepArgument<U>
KalmanFilter(R, P, U) -> KalmanFilter<Density<P>, typename R::ModelType, typename R::MethodType, R, U>;
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Component_, typename Model_, typename Method_, typename Recursion_, typename Upkeep_>
void
describe(
    std::ostream & stream,
    KalmanFilter<Component_, Model_, Method_, Recursion_, Upkeep_> const & filter,
    Indent indent
)
{
    describeFilter(stream, "Kalman filter", filter.state(), indent);
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO