Condition.h

Go to the documentation of this file.

#pragma once

 
#include <array>
#include <cstddef>
#include <expected>
#include <optional>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <Eigen/Cholesky>
 
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Models/Linearization.h>
#include <EmrysFilter/Spaces/Reset.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{

template <typename Space, typename OutputSpace, typename Scalar = double>
struct Conditioned
{
    GaussianParams<Space, Scalar> posterior;     
    Innovation<OutputSpace, Scalar> innovation;  

    [[nodiscard]] Scalar
    logLikelihood() const
    {
        return innovation.logLikelihood();
    }
};

namespace updates
{

struct Standard
{};

[[nodiscard]] constexpr Standard
standard()
{
    return {};
}

struct Symmetric
{};

[[nodiscard]] constexpr Symmetric
symmetric()
{
    return {};
}

struct Joseph
{};

[[nodiscard]] constexpr Joseph
joseph()
{
    return {};
}
 
}  // namespace updates

template <typename T>
concept CovarianceUpdate =
    std::is_same_v<T, updates::Standard> || std::is_same_v<T, updates::Symmetric> || std::is_same_v<T, updates::Joseph>;

template <typename InputSpace, typename OutputSpace, typename Scalar = double>
struct StatisticalLinearization
{
    using JacobianMatrix = internal::MatrixOf<
        Scalar,
        valueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
        valueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<InputSpace>>>;

    using ErrorMatrix = internal::MatrixOf<
        Scalar,
        valueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
        valueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>>;

    JacobianMatrix jacobian;

    ErrorMatrix errorCov;
};
 
namespace internal
{

template <typename Jacobian, typename Noise>
struct JosephTerms
{
    Jacobian jacobian;  
    Noise noise;        
};

template <typename Root, typename Whitened, typename Noise>
struct JosephRootTerms
{
    Root root;          
    Whitened whitened;  
    Noise noise;        
};

template <typename Matrix, typename Gain, typename Cross, typename OutputCov, typename Linearize>
[[nodiscard]] Matrix
posteriorCov(
    updates::Standard,
    Matrix const & cov,
    Gain const & gain,
    Cross const & cross,
    OutputCov const & /*outputCov*/,
    Linearize const & /*linearize*/
)
{
    return Matrix{cov - multiply(gain, cross.transpose())};
}

template <typename Matrix, typename Gain, typename Cross, typename OutputCov, typename Linearize>
[[nodiscard]] Matrix
posteriorCov(
    updates::Symmetric,
    Matrix const & cov,
    Gain const & gain,
    Cross const & /*cross*/,
    OutputCov const & outputCov,
    Linearize const & /*linearize*/
)
{
    using Scalar = typename Matrix::Scalar;
    OutputCov const symmetric = outputCov.template selfadjointView<Eigen::Lower>();
    Matrix result = cov - multiply(multiply(gain, symmetric), gain.transpose());
    return Matrix{Scalar(0.5) * (result + result.transpose())};
}

template <typename Matrix, typename Gain, typename Cross, typename OutputCov, typename Linearize>
[[nodiscard]] Matrix
posteriorCov(
    updates::Joseph,
    Matrix const & cov,
    Gain const & gain,
    Cross const & /*cross*/,
    OutputCov const & /*outputCov*/,
    Linearize const & linearize
)
{
    using Scalar = typename Matrix::Scalar;
    auto const terms = linearize();
    Matrix result = multiply(multiply(gain, terms.noise), gain.transpose());
    if constexpr (requires { terms.root; })
    {
        Matrix const factor = terms.root - multiply(gain, terms.whitened.transpose());
        result.noalias() += multiply(factor, factor.transpose());
    }
    else
    {
        Matrix const factor = Matrix::Identity(cov.rows(), cov.cols()) - multiply(gain, terms.jacobian);
        result.noalias() += multiply(multiply(factor, cov), factor.transpose());
    }
    return Matrix{Scalar(0.5) * (result + result.transpose())};
}

template <std::size_t K>
using ConsideredRows = std::array<BlockRows, K>;

using NoConsidered = ConsideredRows<0>;

template <typename Gain, std::size_t K>
void
zeroConsideredRows(Gain & gain, ConsideredRows<K> const & considered)
{
    for (BlockRows const rows : considered)
    {
        gain.middleRows(rows.start, rows.size).setZero();
    }
}

template <typename Matrix, std::size_t K>
void
keepConsideredCovariance(Matrix & cov, Matrix const & prior, ConsideredRows<K> const & considered)
{
    for (BlockRows const a : considered)
    {
        for (BlockRows const b : considered)
        {
            cov.block(a.start, b.start, a.size, b.size) = prior.block(a.start, b.start, a.size, b.size);
        }
    }
}

template <typename Space, typename Scalar>
struct KalmanConditioned
{
    GaussianParams<Space, Scalar> posterior;                 
    std::optional<InnovationStatistics<Scalar>> statistics;  
};

template <typename Scalar>
struct ConditionedInPlace
{
    std::optional<InnovationStatistics<Scalar>> statistics;  
    bool accepted = true;                                    
};

template <
    bool Statistics,
    typename Repair,
    typename Space,
    typename Scalar,
    typename Residual,
    typename OutputMatrix,
    typename Cross,
    typename Form,
    typename Linearize,
    std::size_t K,
    typename Accept = AcceptEvery>
[[nodiscard]] std::expected<ConditionedInPlace<Scalar>, GaussianError>
kalmanConditionInPlace(
    GaussianParams<Space, Scalar> & state,
    Eigen::MatrixBase<Residual> const & residual,
    OutputMatrix & outputCov,
    Cross const & cross,
    Form const & form,
    Linearize const & linearize,
    [[maybe_unused]] ConsideredRows<K> const & considered,
    Accept const & accept = {}
)
{
    using Matrix = typename GaussianParams<Space, Scalar>::MatrixType;
    [[maybe_unused]] Eigen::Index const nd = state.differenceDimension();
    [[maybe_unused]] Eigen::Index const md = outputCov.rows();
    EMRYS_FILTER_ASSERT(cross.rows() == nd && cross.cols() == md && "the prediction was not made from this prior");
    EMRYS_FILTER_ASSERT(residual.size() == md);
 
    auto const factored = innovationCholesky<Repair>(outputCov);
    if (!factored)
    {
        return std::unexpected{factored.error()};
    }
    auto const & llt = *factored;
    auto result = ConditionedInPlace<Scalar>{};
    if constexpr (Statistics || acceptsByNis<Accept>)
    {
        result.statistics = innovationStatistics(llt, residual);
    }
    if constexpr (acceptsByNis<Accept>)
    {
        if (!accept(result.statistics->nis, md))
        {
            result.accepted = false;
            return result;
        }
    }
 
    // With nothing to condition, only the likelihood remains; MSVC can't form an empty gain.
    if constexpr (valueDimensionAtCompileTime<DifferenceSpace<Space>> != 0)
    {
        auto gain = transposedSolve<typename Cross::PlainObject>(llt, cross.transpose());
        auto & moments = GaussianAccess::moments(state);
        Matrix & cov = moments.covarianceLower();
        if constexpr (K > 0)
        {
            auto const symmetric = symmetricFromLower<Matrix>(cov);
            if constexpr (std::is_same_v<Form, updates::Joseph>)
            {
                zeroConsideredRows(gain, considered);
                cov = posteriorCov(form, symmetric, gain, cross, outputCov, linearize);
            }
            else
            {
                cov = posteriorCov(form, symmetric, gain, cross, outputCov, linearize);
                zeroConsideredRows(gain, considered);
            }
            keepConsideredCovariance(cov, symmetric, considered);
        }
        else if constexpr (std::is_same_v<Form, updates::Standard>)
        {
            cov.template triangularView<Eigen::StrictlyUpper>() = cov.transpose();
            cov.noalias() -= multiply(gain, cross.transpose());
        }
        else
        {
            auto const symmetric = symmetricFromLower<Matrix>(cov);
            cov = posteriorCov(form, symmetric, gain, cross, outputCov, linearize);
        }
        if constexpr (hasIdentityReset<Space>)
        {
            moments.mean() = Emrys::plus(state.space(), moments.mean(), multiply(gain, residual));
        }
        else
        {
            using Vector = typename GaussianParams<Space, Scalar>::VectorType;
            using Difference = typename GaussianParams<Space, Scalar>::DifferenceVectorType;
            Vector const start = moments.mean();
            auto step = Difference{};
            step.resize(nd);
            step.noalias() = multiply(gain, residual);
            moments.mean() = Emrys::plus(state.space(), start, step);
            resetCovarianceAt(state.space(), start, step, moments.mean(), cov);
        }
    }
    return result;
}

template <
    bool Statistics,
    typename Repair,
    typename Prior,
    typename Residual,
    typename OutputMatrix,
    typename Cross,
    typename Form,
    typename Linearize,
    std::size_t K>
[[nodiscard]] std::expected<KalmanConditioned<typename Prior::SpaceType, typename Prior::Scalar>, GaussianError>
kalmanCondition(
    Prior const & prior,
    Eigen::MatrixBase<Residual> const & residual,
    OutputMatrix & outputCov,
    Cross const & cross,
    Form const & form,
    Linearize const & linearize,
    ConsideredRows<K> const & considered
)
{
    using Posterior = GaussianParams<typename Prior::SpaceType, typename Prior::Scalar>;
    auto result = KalmanConditioned<typename Prior::SpaceType, typename Prior::Scalar>{
        Posterior{
            prior.space(),
            typename Posterior::VectorType{prior.mean()},
            typename Posterior::MatrixType{covarianceLowerOf(prior)}
        },
        std::nullopt
    };
    auto conditioned = kalmanConditionInPlace<Statistics, Repair>(
        result.posterior,
        residual,
        outputCov,
        cross,
        form,
        linearize,
        considered
    );
    if (!conditioned)
    {
        return std::unexpected{conditioned.error()};
    }
    result.statistics = conditioned->statistics;
    return result;
}

template <
    typename Prior,
    typename OutputSpace,
    typename LinearizationForm,
    typename Derived,
    typename Form,
    typename Linearize>
[[nodiscard]] std::expected<Conditioned<typename Prior::SpaceType, OutputSpace, typename Prior::Scalar>, GaussianError>
kalmanUpdate(
    Prior const & prior,
    OutputWithCross<typename Prior::SpaceType, OutputSpace, LinearizationForm, typename Prior::Scalar> const &
        predicted,
    Eigen::MatrixBase<Derived> const & observed,
    Form const & form,
    Linearize const & linearize
)
{
    using Output = GaussianParams<OutputSpace, typename Prior::Scalar>;
    Output const & output = predicted.output();
    EMRYS_FILTER_ASSERT(observed.size() == output.valueDimension());
    typename Output::DifferenceVectorType residual = Emrys::minus(output.space(), observed, output.mean());
    auto const & cross = predicted.crossCovariance();
    auto conditioned = kalmanCondition<false, NoRepair>(
        prior,
        residual,
        output.covarianceLower(),
        cross,
        form,
        linearize,
        NoConsidered{}
    );
    if (!conditioned)
    {
        return std::unexpected{conditioned.error()};
    }
    return Conditioned<typename Prior::SpaceType, OutputSpace, typename Prior::Scalar>{
        std::move(conditioned->posterior),
        {std::move(residual), output.covarianceLower()}
    };
}

template <typename Form, typename InputSpace, typename OutputSpace, typename LinearizationForm, typename Scalar>
[[nodiscard]] auto
linearizeFor(OutputWithCross<InputSpace, OutputSpace, LinearizationForm, Scalar> const & predicted)
{
    if constexpr (std::is_same_v<Form, updates::Joseph>)
    {
        using Linearization = StatisticalLinearization<InputSpace, OutputSpace, Scalar>;
        using Predicted = OutputWithCross<InputSpace, OutputSpace, LinearizationForm, Scalar>;
        typename Linearization::ErrorMatrix noise = predicted.inputsCovariance() + predicted.unexplainedCovariance();
        if constexpr (std::is_same_v<LinearizationForm, RootForm>)
        {
            using Terms = JosephRootTerms<
                typename Predicted::RootMatrix const &,
                typename Predicted::CrossMatrix const &,
                typename Linearization::ErrorMatrix>;
            return Terms{predicted.root(), predicted.whitenedCross(), std::move(noise)};
        }
        else
        {
            return JosephTerms<typename Linearization::JacobianMatrix const &, typename Linearization::ErrorMatrix>{
                predicted.jacobian(),
                std::move(noise)
            };
        }
    }
    else
    {
        return std::tuple<>{};
    }
}
 
}  // namespace internal

template <typename InputSpace, typename OutputSpace, typename Form, typename Scalar>
[[nodiscard]] StatisticalLinearization<InputSpace, OutputSpace, Scalar>
statisticalLinearization(OutputWithCross<InputSpace, OutputSpace, Form, Scalar> const & predicted)
{
    using Result = StatisticalLinearization<InputSpace, OutputSpace, Scalar>;
    typename Result::ErrorMatrix errorCov = predicted.inputsCovariance() + predicted.unexplainedCovariance();
    return Result{typename Result::JacobianMatrix{predicted.jacobian()}, std::move(errorCov)};
}

template <
    UncertainInput Prior,
    typename OutputSpace,
    typename LinearizationForm,
    typename Derived,
    CovarianceUpdate Form = updates::Standard>
[[nodiscard]] auto
condition(
    Prior const & prior,
    OutputWithCross<typename Prior::SpaceType, OutputSpace, LinearizationForm, typename Prior::Scalar> const &
        predicted,
    Eigen::MatrixBase<Derived> const & observed,
    Form const & form = {}
)
{
    return internal::kalmanUpdate(prior, predicted, observed, form, [&] {
        return internal::linearizeFor<Form>(predicted);
    });
}
 
}  // namespace Emrys