Linearization.h

Go to the documentation of this file.

#pragma once

 
#include <cmath>
#include <concepts>
#include <expected>
#include <limits>
#include <optional>
#include <type_traits>
#include <utility>
 
#include <Eigen/Cholesky>
 
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

template <typename T>
concept UncertainInput = requires(T const & input) {
    typename T::Scalar;
    typename T::SpaceType;
    input.mean();
    input.space();
    input.valueDimension();
    input.differenceDimension();
} && (requires(T const & input) { input.covarianceLower(); } || requires(T const & input) { input.root(); });
 
namespace internal
{

template <UncertainInput Input>
[[nodiscard]] decltype(auto)
covarianceLowerOf(Input const & input)
{
    if constexpr (requires { input.covarianceLower(); })
    {
        return input.covarianceLower();
    }
    else
    {
        using Matrix = std::remove_cvref_t<decltype(input.root())>;
        Matrix result = multiply(asTriangular<Eigen::Lower>(input.root()), input.root().transpose());
        return result;
    }
}

template <UncertainInput Prior>
[[nodiscard]] decltype(auto)
withCovariance(Prior const & prior)
{
    if constexpr (requires { prior.covarianceLower(); })
    {
        return (prior);
    }
    else
    {
        using Result = GaussianParams<typename Prior::SpaceType, typename Prior::Scalar>;
        return Result{
            prior.space(),
            typename Result::VectorType{prior.mean()},
            typename Result::MatrixType{covarianceLowerOf(prior)}
        };
    }
}
 
}  // namespace internal

struct JacobianForm
{};

struct RootForm
{};
 
namespace internal
{

template <typename T>
inline constexpr bool isLinearizationForm = std::is_same_v<T, JacobianForm> || std::is_same_v<T, RootForm>;

template <typename InputSpace, typename OutputSpace, typename Scalar>
using CrossMatrixOf = MatrixOf<
    Scalar,
    valueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
    valueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>,
    maxValueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
    maxValueDimensionAtCompileTime<DifferenceSpace<OutputSpace>>>;

template <typename Derived>
[[nodiscard]] bool
isInvertibleLowerTriangular(Eigen::MatrixBase<Derived> const & root)
{
    using Scalar = typename Derived::Scalar;
    for (Eigen::Index j = 0; j < root.cols(); ++j)
    {
        if (!(root(j, j) > Scalar(0)))
        {
            return false;
        }
        for (Eigen::Index i = 0; i < j; ++i)
        {
            if (root(i, j) != Scalar(0))
            {
                return false;
            }
        }
    }
    return true;
}

template <typename Result, typename Root, typename Whitened>
[[nodiscard]] Result
solveWithRoot(Eigen::MatrixBase<Root> const & root, Eigen::MatrixBase<Whitened> const & whitened)
{
    using RootMatrix = typename Root::PlainObject;
    Result result = whitened;
    // MSVC can't form an empty solve.
    if constexpr (Root::RowsAtCompileTime != 0)
    {
        if (isInvertibleLowerTriangular(root))
        {
            root.template triangularView<Eigen::Lower>().transpose().solveInPlace(result);
        }
        else
        {
            auto const ldlt = Eigen::LDLT<RootMatrix>{RootMatrix{multiply(root, root.transpose())}};
            result.noalias() = multiply(root, whitened);
            solveInPlace(ldlt, result);
        }
    }
    return result;
}
 
}  // namespace internal

template <typename InputSpace, typename OutputSpace, typename Form, typename Scalar = double>
class OutputWithCross
{
    static_assert(internal::isLinearizationForm<Form>);

    static constexpr bool IsRoot = std::is_same_v<Form, RootForm>;
 
public:

    using FormType = Form;

    using OutputType = GaussianParams<OutputSpace, Scalar>;

    using OutputVector = typename OutputType::VectorType;

    using OutputMatrix = typename OutputType::MatrixType;

    using CrossMatrix = internal::CrossMatrixOf<InputSpace, OutputSpace, Scalar>;

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

    using RootMatrix = internal::MatrixOf<
        Scalar,
        valueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
        valueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<InputSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<InputSpace>>>;

    using FactorMatrix = std::conditional_t<IsRoot, RootMatrix, JacobianMatrix>;

    template <
        MatrixExpression Mean,
        MatrixExpression Factor,
        MatrixExpression Cross,
        MatrixExpression Explained,
        MatrixExpression Inputs,
        MatrixExpression Unexplained>
    OutputWithCross(
        OutputSpace space,
        Mean && mean,
        Factor && factor,
        Cross && cross,
        Explained && explained,
        Inputs && inputs,
        Unexplained && unexplained
    )
        : m_explained{std::forward<Explained>(explained)}
        , m_inputs{std::forward<Inputs>(inputs)}
        , m_unexplained{std::forward<Unexplained>(unexplained)}
        , m_output{std::move(space), std::forward<Mean>(mean), OutputMatrix{m_explained + m_inputs + m_unexplained}}
        , m_factor{std::forward<Factor>(factor)}
        , m_cross{std::forward<Cross>(cross)}
    {
        EMRYS_FILTER_ASSERT(
            m_cross.cols() == m_output.differenceDimension() && m_factor.cols() == m_cross.rows()
            && "the linearization must have a row per input coordinate and a column per output coordinate"
        );
    }

    [[nodiscard]] OutputType const &
    output() const &
    {
        return m_output;
    }

    [[nodiscard]] OutputType
    output() &&
    {
        return std::move(m_output);
    }

    [[nodiscard]] OutputMatrix const &
    explainedCovariance() const
    {
        return m_explained;
    }

    [[nodiscard]] OutputMatrix const &
    inputsCovariance() const
    {
        return m_inputs;
    }

    [[nodiscard]] OutputMatrix const &
    unexplainedCovariance() const
    {
        return m_unexplained;
    }

    [[nodiscard]] CrossMatrix const &
    crossCovariance() const
        requires(!IsRoot)
    {
        return m_cross;
    }

    [[nodiscard]] CrossMatrix
    crossCovariance() const
        requires IsRoot
    {
        return internal::multiply(m_factor, m_cross);
    }

    [[nodiscard]] JacobianMatrix const &
    jacobian() const
        requires(!IsRoot)
    {
        return m_factor;
    }

    [[nodiscard]] JacobianMatrix
    jacobian() const
        requires IsRoot
    {
        return internal::solveWithRoot<CrossMatrix>(m_factor, m_cross).transpose();
    }

    [[nodiscard]] RootMatrix const &
    root() const
        requires IsRoot
    {
        return m_factor;
    }

    [[nodiscard]] CrossMatrix const &
    whitenedCross() const
        requires IsRoot
    {
        return m_cross;
    }
 
private:
 
    OutputMatrix m_explained;    
    OutputMatrix m_inputs;       
    OutputMatrix m_unexplained;  
    OutputType m_output;         
    FactorMatrix m_factor;       
    CrossMatrix m_cross;         
};
 
namespace internal
{

template <typename Scalar>
struct InnovationStatistics
{
    Scalar nis;             
    Scalar logDeterminant;  
};

template <typename R, typename Matrix>
[[nodiscard]] std::expected<Eigen::LLT<std::remove_const_t<Matrix>>, GaussianError>
innovationCholesky(Matrix & cov)
{
    auto llt = Eigen::LLT<std::remove_const_t<Matrix>>{cov};
    if (llt.info() == Eigen::Success)
    {
        return llt;
    }
    if constexpr (!std::is_same_v<R, NoRepair>)
    {
        static_assert(!std::is_const_v<Matrix>, "a repaired innovation covariance replaces the one given");
        auto repaired = R{}.repaired(cov);
        if (!repaired)
        {
            return std::unexpected{repaired.error()};
        }
        llt.compute(*repaired);
        if (llt.info() != Eigen::Success)
        {
            return std::unexpected{GaussianError::NotPositiveDefinite};
        }
        cov = std::move(*repaired);
        return llt;
    }
    else
    {
        return std::unexpected{GaussianError::NotPositiveDefinite};
    }
}

template <typename Matrix, typename Residual>
[[nodiscard]] InnovationStatistics<typename Matrix::Scalar>
innovationStatistics(Eigen::LLT<Matrix> const & llt, Eigen::MatrixBase<Residual> const & residual)
{
    using Scalar = typename Matrix::Scalar;
    typename Residual::PlainObject whitened = residual;
    llt.matrixL().solveInPlace(whitened);
    return {whitened.squaredNorm(), Scalar(2) * llt.matrixLLT().diagonal().array().log().sum()};
}

struct InnovationInPlace
{};

struct AcceptEvery
{
    template <typename Scalar>
    [[nodiscard]] constexpr bool
    operator()(Scalar /*nis*/, Eigen::Index /*dimension*/) const
    {
        return true;
    }
};

template <typename Accept>
inline constexpr bool acceptsByNis = !std::is_same_v<Accept, AcceptEvery>;
 
}  // namespace internal

template <typename Space, typename Scalar = double>
class Innovation
{
public:

    using VectorType = typename GaussianParams<Space, Scalar>::DifferenceVectorType;

    using MatrixType = typename GaussianParams<Space, Scalar>::MatrixType;

    template <MatrixExpression Residual, MatrixExpression Cov>
    Innovation(
        Residual && residual,
        Cov && cov,
        std::optional<internal::InnovationStatistics<Scalar>> statistics = std::nullopt
    )
        : m_residual{std::forward<Residual>(residual)}
        , m_cov{std::forward<Cov>(cov)}
        , m_statistics{statistics}
    {
        EMRYS_FILTER_ASSERT(
            m_cov.rows() == m_residual.size() && m_cov.cols() == m_residual.size()
            && "the covariance must have a row per residual coordinate"
        );
    }

    template <MatrixExpression Residual, typename MakeCov, typename Condition>
    Innovation(internal::InnovationInPlace, Residual && residual, MakeCov const & makeCov, Condition const & condition)
        : m_residual{std::forward<Residual>(residual)}
        , m_cov{makeCov()}
        , m_statistics{condition(std::as_const(m_residual), m_cov)}
    {
        EMRYS_FILTER_ASSERT(
            m_cov.rows() == m_residual.size() && m_cov.cols() == m_residual.size()
            && "the covariance must have a row per residual coordinate"
        );
    }

    [[nodiscard]] VectorType const &
    residual() const
    {
        return m_residual;
    }

    [[nodiscard]] MatrixType const &
    covarianceLower() const
    {
        return m_cov;
    }

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

    [[nodiscard]] Scalar
    nis() const
    {
        return statistics().nis;
    }

    [[nodiscard]] Scalar
    logLikelihood() const
    {
        internal::InnovationStatistics<Scalar> const & known = statistics();
        return Scalar(-0.5) * (Scalar(differenceDimension()) * log2Pi<Scalar>() + known.logDeterminant + known.nis);
    }
 
private:

    [[nodiscard]] internal::InnovationStatistics<Scalar> const &
    statistics() const
    {
        if (!m_statistics)
        {
            auto const llt = Eigen::LLT<MatrixType>{m_cov};
            EMRYS_FILTER_ASSERT(llt.info() == Eigen::Success && "the innovation covariance must be positive definite");
            m_statistics = internal::innovationStatistics(llt, m_residual);
        }
        return *m_statistics;
    }
 
    VectorType m_residual;  
    MatrixType m_cov;       

    mutable std::optional<internal::InnovationStatistics<Scalar>> m_statistics;
};
 
}  // namespace Emrys