SqrtInfoGaussian.h

Go to the documentation of this file.

#pragma once

 
#include <expected>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Core/Forwarding.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Print.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Core/Triangularization.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/GaussianForm.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Spaces/Composite.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Reset.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

template <Space Space_, RealScalar Scalar_ = double>
class SqrtInfoGaussianParams
{
public:
 
    using MomentsType = Moments<Space_, Scalar_>;                             
    using Scalar = typename MomentsType::Scalar;                              
    using Index = typename MomentsType::Index;                                
    using VectorType = typename MomentsType::VectorType;                      
    using DifferenceVectorType = typename MomentsType::DifferenceVectorType;  
    using MatrixType = typename MomentsType::MatrixType;                      
    using SpaceType = Space_;                                                 
    using DensityType = FormGaussian<SqrtInfoGaussianParams>;                 

    static constexpr int ValueDimensionAtCompileTime = MomentsType::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = MomentsType::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = MomentsType::DifferenceDimensionAtCompileTime;

    static constexpr int MaxDifferenceDimensionAtCompileTime = MomentsType::MaxDifferenceDimensionAtCompileTime;
 
    static constexpr bool IsParams = true;    
    static constexpr bool IsInfoForm = true;  
    static constexpr bool IsSqrtForm = true;  

    SqrtInfoGaussianParams()
        : SqrtInfoGaussianParams{Space_{}}
    {}

    explicit SqrtInfoGaussianParams(Space_ space)
        : m_space{std::move(space)}
    {
        m_reference.setZero(Emrys::valueDimension(m_space));
        m_root.setZero(differenceDimension(), differenceDimension());
        m_whitenedOffset.setZero(differenceDimension());
        m_densityMean.setZero(valueDimension());
    }

    template <MatrixExpression Reference, MatrixExpression Root, MatrixExpression Offset>
        requires internal::SizedByLength<Space_>
    SqrtInfoGaussianParams(Reference && reference, Root && root, Offset && whitenedOffset)
        : m_reference{std::forward<Reference>(reference)}
        , m_root{std::forward<Root>(root)}
        , m_whitenedOffset{std::forward<Offset>(whitenedOffset)}
    {
        fitSpace();
    }

    template <MatrixExpression Reference, MatrixExpression Root, MatrixExpression Offset>
    SqrtInfoGaussianParams(Space_ space, Reference && reference, Root && root, Offset && whitenedOffset)
        : m_reference{std::forward<Reference>(reference)}
        , m_root{std::forward<Root>(root)}
        , m_whitenedOffset{std::forward<Offset>(whitenedOffset)}
        , m_space{std::move(space)}
    {
        fitSpace();
    }

    [[nodiscard]] Index
    valueDimension() const
    {
        return Emrys::valueDimension(m_space);
    }

    [[nodiscard]] Index
    differenceDimension() const
    {
        return Emrys::differenceDimension(m_space);
    }

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

    [[nodiscard]] VectorType const &
    reference() const
    {
        return m_reference;
    }

    [[nodiscard]] MatrixType const &
    infoRoot() const
    {
        return m_root;
    }

    [[nodiscard]] DifferenceVectorType const &
    whitenedOffset() const
    {
        return m_whitenedOffset;
    }

    [[nodiscard]] DifferenceVectorType
    infoVector() const
    {
        DifferenceVectorType result =
            internal::multiply(internal::asTriangular<Eigen::Upper>(m_root).transpose(), m_whitenedOffset);
        return result;
    }

    template <typename Derived>
    void
    setReference(Eigen::MatrixBase<Derived> const & reference)
    {
        m_stamp.touchValues();
        EMRYS_FILTER_ASSERT(reference.size() == valueDimension());
        VectorType moved = reference;
        normalizeValues(m_space, moved);
        DifferenceVectorType const shift = Emrys::minus(m_space, moved, m_reference);
        m_whitenedOffset.noalias() -= internal::multiply(internal::asTriangular<Eigen::Upper>(m_root), shift);
        if constexpr (!hasIdentityReset<Space_>)
        {
            Eigen::Index const n = differenceDimension();
            AugmentedType stacked(n, n + 1);
            stacked.leftCols(n) = m_root;
            internal::forEachInverseResetJacobian(
                m_space,
                m_reference,
                shift,
                moved,
                [&](Eigen::Index differences, auto const & gInverse) {
                    Eigen::Index const k = gInverse.rows();
                    stacked.middleCols(differences, k) = m_root.middleCols(differences, k) * gInverse;
                }
            );
            stacked.col(n) = m_whitenedOffset;
            triangularize(stacked);
            m_root = stacked.leftCols(n);
            m_whitenedOffset = stacked.col(n);
        }
        m_reference = std::move(moved);
    }

    template <MatrixExpression Root, MatrixExpression Offset>
    void
    setRoot(Root && root, Offset && whitenedOffset)
    {
        m_stamp.touchValues();
        EMRYS_FILTER_ASSERT(
            root.rows() == differenceDimension() && root.cols() == differenceDimension()
            && whitenedOffset.size() == differenceDimension()
        );
        m_root = std::forward<Root>(root);
        m_root.template triangularView<Eigen::StrictlyLower>().setZero();
        m_whitenedOffset = std::forward<Offset>(whitenedOffset);
    }

    [[nodiscard]] std::expected<VectorType, GaussianError>
    mean() const
    {
        if (auto checked = checkRoot(); !checked)
        {
            return std::unexpected{checked.error()};
        }
        return solvedMean();
    }

    [[nodiscard]] std::expected<MomentsType, GaussianError>
    moments() const
    {
        if (auto checked = checkRoot(); !checked)
        {
            return std::unexpected{checked.error()};
        }
        return momentsWithMean(solvedMean());
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    covariance() const
    {
        auto inverse = covarianceRoot();
        if (!inverse)
        {
            return std::unexpected{inverse.error()};
        }
        MatrixType result = internal::multiply(*inverse, inverse->transpose());
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    precision() const
    {
        MatrixType result = internal::multiply(m_root.transpose(), internal::asTriangular<Eigen::Upper>(m_root));
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    covarianceRoot() const
    {
        if (auto checked = checkRoot(); !checked)
        {
            return std::unexpected{checked.error()};
        }
        MatrixType result = MatrixType::Identity(differenceDimension(), differenceDimension());
        internal::solveInPlace(m_root.template triangularView<Eigen::Upper>(), result);
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    precisionRoot() const
    {
        return m_root;
    }

    [[nodiscard]] std::expected<DensityType, GaussianError>
    toDensity() const &
    {
        if (auto checked = checkRoot(); !checked)
        {
            return std::unexpected{checked.error()};
        }
        m_densityMean = solvedMean();
        m_logDeterminant = Scalar(-2) * m_root.diagonal().array().abs().log().sum();
        return DensityType{*this};
    }

    void
    toDensity() const && = delete;

    template <typename Path>
    [[nodiscard]] auto
    marginal() const
    {
        return internal::marginalOverPath<Path>(m_space, [&](auto const & kept, auto const & keptSpace) {
            return marginalAt(kept, keptSpace);
        });
    }

    template <typename Path>
    [[nodiscard]] auto
    marginal(Index element) const
    {
        return internal::marginalOverElement<Path>(m_space, element, [&](auto const & kept, auto const & keptSpace) {
            return marginalAt(kept, keptSpace);
        });
    }

    template <typename GivenPath, typename Derived>
    [[nodiscard]] auto
    conditional(Eigen::MatrixBase<Derived> const & values) const
    {
        return internal::tagConditional<GivenPath>(
            m_space,
            [&](auto const & target, auto const & given, auto const & targetSpace, auto const & givenSpace) {
                using GivenSpace = std::remove_cvref_t<decltype(givenSpace)>;
                using Result = SqrtInfoGaussianParams<std::remove_cvref_t<decltype(targetSpace)>, Scalar>;
                constexpr int GivenSize = valueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>;
                constexpr int GivenMaxSize = maxValueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>;
                constexpr int TargetSize = Result::DifferenceDimensionAtCompileTime;
                constexpr int TargetMaxSize = Result::MaxDifferenceDimensionAtCompileTime;
                using GivenVector = Eigen::Matrix<Scalar, GivenSize, 1, Eigen::ColMajor, GivenMaxSize, 1>;
                using Stacked = internal::MatrixOf<
                    Scalar,
                    DifferenceDimensionAtCompileTime,
                    TargetSize == Eigen::Dynamic ? Eigen::Dynamic : TargetSize + 1,
                    MaxDifferenceDimensionAtCompileTime,
                    TargetMaxSize == Eigen::Dynamic ? Eigen::Dynamic : TargetMaxSize + 1>;
                EMRYS_FILTER_ASSERT(values.size() == given.values.size());
 
                Eigen::Index const n = differenceDimension();
                Eigen::Index const t = target.differences.size();
                typename Result::VectorType reference = m_reference(target.values);
                auto root = typename Result::MatrixType{};
                root.resize(t, t);
                auto whitened = typename Result::DifferenceVectorType{};
                whitened.resize(t);
                // MSVC can't form the empty products of a conditional on everything or on nothing.
                if constexpr (TargetSize != 0 && GivenSize != 0)
                {
                    GivenVector const offset = Emrys::minus(givenSpace, values, m_reference(given.values));
                    Stacked stacked(n, t + 1);
                    stacked.leftCols(t) = m_root(Eigen::placeholders::all, target.differences);
                    stacked.col(t) = m_whitenedOffset
                                   - internal::multiply(m_root(Eigen::placeholders::all, given.differences), offset);
                    triangularize(stacked);
                    root = stacked.topLeftCorner(t, t);
                    whitened = stacked.col(t).head(t);
                }
                else
                {
                    root = m_root(target.differences, target.differences);
                    whitened = m_whitenedOffset(target.differences);
                }
                return Result{targetSpace, std::move(reference), std::move(root), std::move(whitened)};
            }
        );
    }
 
private:
 
    friend DensityType;  

    template <typename>
    friend class internal::ValuesCheck;

    using AugmentedType = internal::MatrixOf<
        Scalar,
        DifferenceDimensionAtCompileTime,
        DifferenceDimensionAtCompileTime == Eigen::Dynamic ? Eigen::Dynamic : DifferenceDimensionAtCompileTime + 1,
        MaxDifferenceDimensionAtCompileTime,
        MaxDifferenceDimensionAtCompileTime == Eigen::Dynamic ? Eigen::Dynamic
                                                              : MaxDifferenceDimensionAtCompileTime + 1>;

    [[nodiscard]] std::expected<void, GaussianError>
    checkRoot() const
    {
        if (!allFinite(m_root))
        {
            return std::unexpected{GaussianError::NonFinite};
        }
        if ((m_root.diagonal().array() == Scalar(0)).any())
        {
            return std::unexpected{GaussianError::Singular};
        }
        return {};
    }

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] std::expected<SqrtInfoGaussianParams<KeptSpace, Scalar>, GaussianError>
    marginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        using Result = SqrtInfoGaussianParams<KeptSpace, Scalar>;
        using Expected = std::expected<Result, GaussianError>;
        constexpr int D = DifferenceDimensionAtCompileTime;
        constexpr int DMax = MaxDifferenceDimensionAtCompileTime;
        Eigen::Index const n = differenceDimension();
        Eigen::Index const k = kept.differences.size();
        auto const rest =
            internal::resolveIndices<D, DMax>(n, internal::complementIndices<D, DMax>(n, kept.differences));
        Eigen::Index const r = rest.size();
 
        typename Result::VectorType reference = m_reference(kept.values);
        AugmentedType stacked(n, n + 1);
        stacked.leftCols(r) = m_root(Eigen::placeholders::all, rest);
        stacked.middleCols(r, k) = m_root(Eigen::placeholders::all, kept.differences);
        stacked.col(n) = m_whitenedOffset;
        triangularize(stacked);
        if ((stacked.diagonal().head(r).array() == Scalar(0)).any())
        {
            return Expected{std::unexpected{GaussianError::Singular}};
        }
        typename Result::MatrixType root = stacked.block(r, r, k, k);
        typename Result::DifferenceVectorType whitened = stacked.col(n).tail(k);
        return Expected{Result{keptSpace, std::move(reference), std::move(root), std::move(whitened)}};
    }

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] SqrtInfoGaussianParams<KeptSpace, Scalar>
    densityMarginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        auto result = marginalAt(kept, keptSpace);
        EMRYS_FILTER_ASSERT(result.has_value() && "a checked root's marginal cannot be singular");
        return std::move(*result);
    }

    [[nodiscard]] MomentsType
    momentsWithMean(VectorType mean) const
    {
        auto result = MomentsType{std::move(mean), *(covariance())};
        if constexpr (!hasIdentityReset<Space_>)
        {
            DifferenceVectorType const step = Emrys::minus(m_space, result.mean(), m_reference);
            resetCovariance(m_space, m_reference, step, result.covarianceLower());
        }
        return result;
    }

    [[nodiscard]] MomentsType
    densityMoments() const
    {
        return momentsWithMean(m_densityMean);
    }

    [[nodiscard]] VectorType
    solvedMean() const
    {
        DifferenceVectorType const offset = m_root.template triangularView<Eigen::Upper>().solve(m_whitenedOffset);
        VectorType result = Emrys::plus(m_space, m_reference, offset);
        return result;
    }

    [[nodiscard]] VectorType const &
    densityMean() const
    {
        return m_densityMean;
    }

    void
    whitenInPlace(DifferenceVectorType & residual) const
    {
        residual = internal::multiply(internal::asTriangular<Eigen::Upper>(m_root), residual).eval();
    }

    void
    colorInto(auto & colored, auto const & normals) const
    {
        colored = normals;
        internal::solveInPlace(m_root.template triangularView<Eigen::Upper>(), colored);
    }

    void
    fitSpace()
    {
        internal::fitValueDimension(m_space, m_reference.size());
        EMRYS_FILTER_ASSERT(
            m_root.rows() == differenceDimension() && m_root.cols() == differenceDimension()
            && m_whitenedOffset.size() == differenceDimension()
            && "the root and whitened offset must have a row per difference coordinate of the space"
        );
        normalizeValues(m_space, m_reference);
        m_root.template triangularView<Eigen::StrictlyLower>().setZero();
        m_densityMean.setZero(valueDimension());
    }
 
    VectorType m_reference;                         
    MatrixType m_root;                              
    DifferenceVectorType m_whitenedOffset;          
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Space_ m_space;  
    mutable Scalar m_logDeterminant{};              
    mutable VectorType m_densityMean;               
    internal::Stamp m_stamp;                        
};

template <Space Space_, RealScalar Scalar_ = double>
using SqrtInfoGaussian = FormGaussian<SqrtInfoGaussianParams<Space_, Scalar_>>;

template <Space Space_, RealScalar Scalar_>
[[nodiscard]] SqrtInfoGaussianParams<Space_, Scalar_>
multiply(SqrtInfoGaussianParams<Space_, Scalar_> const & first, SqrtInfoGaussianParams<Space_, Scalar_> const & second)
{
    using Result = SqrtInfoGaussianParams<Space_, Scalar_>;
    constexpr int D = Result::DifferenceDimensionAtCompileTime;
    constexpr int DMax = Result::MaxDifferenceDimensionAtCompileTime;
    using Stacked = internal::MatrixOf<
        Scalar_,
        D == Eigen::Dynamic ? Eigen::Dynamic : 2 * D,
        D == Eigen::Dynamic ? Eigen::Dynamic : D + 1,
        DMax == Eigen::Dynamic ? Eigen::Dynamic : 2 * DMax,
        DMax == Eigen::Dynamic ? Eigen::Dynamic : DMax + 1>;
    EMRYS_FILTER_ASSERT(first.differenceDimension() == second.differenceDimension());
    Eigen::Index const n = first.differenceDimension();
 
    Result moved = second;
    moved.setReference(first.reference());
    Stacked stacked(2 * n, n + 1);
    stacked.topLeftCorner(n, n) = first.infoRoot();
    stacked.topRightCorner(n, 1) = first.whitenedOffset();
    stacked.bottomLeftCorner(n, n) = moved.infoRoot();
    stacked.bottomRightCorner(n, 1) = moved.whitenedOffset();
    triangularize(stacked);
 
    // Built from parts rather than as a copy of first, whose root and offset would be overwritten.
    auto root = typename Result::MatrixType{stacked.topLeftCorner(n, n)};
    root.template triangularView<Eigen::StrictlyLower>().setZero();
    return Result{
        first.space(),
        first.reference(),
        std::move(root),
        typename Result::DifferenceVectorType{stacked.topRightCorner(n, 1)}
    };
}

template <GaussianForm G>
[[nodiscard]] std::expected<SqrtInfoGaussianParams<typename G::SpaceType, typename G::Scalar>, GaussianError>
toSqrtInfo(G const & form)
{
    using Result = SqrtInfoGaussianParams<typename G::SpaceType, typename G::Scalar>;
    if constexpr (std::is_same_v<G, Result>)
    {
        return form;
    }
    else
    {
        auto root = form.precisionRoot();
        if (!root)
        {
            return std::unexpected{root.error()};
        }
        if constexpr (isInfoForm<G>)
        {
            typename Result::DifferenceVectorType whitened =
                root->transpose().template triangularView<Eigen::Lower>().solve(form.infoVector());
            return Result{form.space(), form.reference(), std::move(*root), std::move(whitened)};
        }
        else
        {
            std::expected<typename G::VectorType, GaussianError> mean = form.mean();
            if (!mean)
            {
                return std::unexpected{mean.error()};
            }
            triangularize(*root);
            return Result{
                form.space(),
                std::move(*mean),
                std::move(*root),
                Result::DifferenceVectorType::Zero(form.differenceDimension())
            };
        }
    }
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Space_, typename Scalar_>
void
describe(std::ostream & stream, SqrtInfoGaussianParams<Space_, Scalar_> const & gaussian, Indent indent)
{
    stream << indent << "square-root information Gaussian";
    printDimension(stream, gaussian.valueDimension(), gaussian.differenceDimension());
    printMoments(stream, indent.deeper(), gaussian.moments());
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO