InfoGaussian.h

Go to the documentation of this file.

#pragma once

 
#include <expected>
#include <type_traits>
#include <utility>
 
#include <Eigen/Cholesky>
 
#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 InfoGaussianParams
{
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<InfoGaussianParams>;  

    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 = false;  

    InfoGaussianParams()
        : InfoGaussianParams{Space_{}}
    {}

    explicit InfoGaussianParams(Space_ space)
        : m_space{std::move(space)}
    {
        m_reference.setZero(Emrys::valueDimension(m_space));
        m_infoMatrix.setZero(differenceDimension(), differenceDimension());
        m_infoVector.setZero(differenceDimension());
    }

    template <MatrixExpression Reference, MatrixExpression Matrix, MatrixExpression Vector>
        requires internal::SizedByLength<Space_>
    InfoGaussianParams(Reference && reference, Matrix && infoMatrix, Vector && infoVector)
        : m_reference{std::forward<Reference>(reference)}
        , m_infoMatrix{std::forward<Matrix>(infoMatrix)}
        , m_infoVector{std::forward<Vector>(infoVector)}
    {
        fitSpace();
    }

    template <MatrixExpression Reference, MatrixExpression Matrix, MatrixExpression Vector>
    InfoGaussianParams(Space_ space, Reference && reference, Matrix && infoMatrix, Vector && infoVector)
        : m_reference{std::forward<Reference>(reference)}
        , m_infoMatrix{std::forward<Matrix>(infoMatrix)}
        , m_infoVector{std::forward<Vector>(infoVector)}
        , m_space{std::move(space)}
    {
        fitSpace();
    }

    InfoGaussianParams(InfoGaussianParams const & other)
        : m_reference{other.m_reference}
        , m_infoMatrix{other.m_infoMatrix}
        , m_infoVector{other.m_infoVector}
        , m_space{other.m_space}
        , m_llt{other.m_factorValid ? other.m_llt : std::nullopt}
        , m_factorValid{other.m_factorValid}
    {
        if (m_factorValid)
        {
            m_logDeterminant = other.m_logDeterminant;
        }
    }

    InfoGaussianParams(InfoGaussianParams && other) noexcept
        : m_reference{std::move(other.m_reference)}
        , m_infoMatrix{std::move(other.m_infoMatrix)}
        , m_infoVector{std::move(other.m_infoVector)}
        , m_space{std::move(other.m_space)}
        , m_llt{other.m_factorValid ? std::move(other.m_llt) : std::nullopt}
        , m_factorValid{other.m_factorValid}
    {
        if (m_factorValid)
        {
            m_logDeterminant = other.m_logDeterminant;
        }
        other.m_stamp.touchShape();
        other.m_factorValid = false;
    }

    InfoGaussianParams &
    operator=(InfoGaussianParams const & other)
    {
        m_stamp.touchShape();
        m_reference = other.m_reference;
        m_infoMatrix = other.m_infoMatrix;
        m_infoVector = other.m_infoVector;
        m_space = other.m_space;
        if (other.m_factorValid)
        {
            m_llt = other.m_llt;
            m_logDeterminant = other.m_logDeterminant;
        }
        m_factorValid = other.m_factorValid;
        return *this;
    }

    InfoGaussianParams &
    operator=(InfoGaussianParams && other) noexcept
    {
        m_stamp.touchShape();
        other.m_stamp.touchShape();
        m_reference = std::move(other.m_reference);
        m_infoMatrix = std::move(other.m_infoMatrix);
        m_infoVector = std::move(other.m_infoVector);
        m_space = std::move(other.m_space);
        if (other.m_factorValid)
        {
            m_llt = std::move(other.m_llt);
            m_logDeterminant = other.m_logDeterminant;
        }
        m_factorValid = other.m_factorValid;
        other.m_factorValid = false;
        return *this;
    }

    [[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 &
    infoMatrix() const
    {
        return m_infoMatrix;
    }

    [[nodiscard]] DifferenceVectorType const &
    infoVector() const
    {
        return m_infoVector;
    }

    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);
        // The product and the reset read the precision whole, so mirror the lower triangle first:
        // a dense product keeps fixed sizes fixed (internal::multiply()), where a self-adjoint view's never does.
        m_infoMatrix.template triangularView<Eigen::StrictlyUpper>() = m_infoMatrix.transpose();
        m_infoVector.noalias() -= internal::multiply(m_infoMatrix, shift);
        if constexpr (!hasIdentityReset<Space_>)
        {
            internal::resetInfoAt(m_space, m_reference, shift, moved, m_infoMatrix, m_infoVector);
            m_factorValid = false;
        }
        m_reference = std::move(moved);
    }

    template <MatrixExpression Matrix, MatrixExpression Vector>
    void
    setInfo(Matrix && infoMatrix, Vector && infoVector)
    {
        m_stamp.touchValues();
        EMRYS_FILTER_ASSERT(
            infoMatrix.rows() == differenceDimension() && infoMatrix.cols() == differenceDimension()
            && infoVector.size() == differenceDimension()
        );
        m_infoMatrix = std::forward<Matrix>(infoMatrix);
        m_infoVector = std::forward<Vector>(infoVector);
        m_factorValid = false;
    }

    template <typename MatrixDerived, typename VectorDerived>
    void
    addInfo(Eigen::MatrixBase<MatrixDerived> const & infoMatrix, Eigen::MatrixBase<VectorDerived> const & infoVector)
    {
        m_stamp.touchValues();
        EMRYS_FILTER_ASSERT(
            infoMatrix.rows() == differenceDimension() && infoMatrix.cols() == differenceDimension()
            && infoVector.size() == differenceDimension()
        );
        m_infoMatrix.template triangularView<Eigen::Lower>() += infoMatrix;
        m_infoVector += infoVector;
        m_factorValid = false;
    }

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

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

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    covariance() const
    {
        if (auto factorized = factorize(); !factorized)
        {
            return std::unexpected{factorized.error()};
        }
        MatrixType result = MatrixType::Identity(differenceDimension(), differenceDimension());
        internal::solveInPlace(*m_llt, result);
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    precision() const
    {
        MatrixType result = m_infoMatrix.template selfadjointView<Eigen::Lower>();
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    covarianceRoot() const
    {
        if (auto factorized = factorize(); !factorized)
        {
            return std::unexpected{factorized.error()};
        }
        MatrixType result = MatrixType::Identity(differenceDimension(), differenceDimension());
        internal::solveInPlace(m_llt->matrixU(), result);
        return result;
    }

    [[nodiscard]] std::expected<MatrixType, GaussianError>
    precisionRoot() const
    {
        if (auto factorized = factorize(); !factorized)
        {
            return std::unexpected{factorized.error()};
        }
        MatrixType result = m_llt->matrixU();
        return result;
    }

    [[nodiscard]] std::expected<void, GaussianError>
    factorize() const
    {
        if (m_factorValid)
        {
            return {};
        }
        if (!m_llt)
        {
            m_llt.emplace(differenceDimension());
        }
        m_llt->compute(m_infoMatrix);
        if (m_llt->info() != Eigen::Success)
        {
            return std::unexpected{GaussianError::Singular};
        }
        // As in GaussianParams::factorize(): NaN passes LLT's test but reaches the factor's
        // diagonal.
        Scalar const logDeterminant = Scalar(-2) * m_llt->matrixLLT().diagonal().array().log().sum();
        if (!isFinite(logDeterminant))
        {
            return std::unexpected{GaussianError::NonFinite};
        }
        m_logDeterminant = logDeterminant;
        m_factorValid = true;
        return {};
    }

    [[nodiscard]] std::expected<DensityType, GaussianError>
    toDensity() const &
    {
        if (auto factorized = factorize(); !factorized)
        {
            return std::unexpected{factorized.error()};
        }
        m_densityMean = solvedMean();
        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 = InfoGaussianParams<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 CrossMatrix = internal::MatrixOf<Scalar, TargetSize, GivenSize, TargetMaxSize, GivenMaxSize>;
                EMRYS_FILTER_ASSERT(values.size() == given.values.size());
 
                typename Result::VectorType reference = m_reference(target.values);
                typename Result::MatrixType matrix =
                    internal::lowerSymmetricBlock(m_infoMatrix, target.differences, target.differences);
                typename Result::DifferenceVectorType vector = m_infoVector(target.differences);
                // 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));
                    CrossMatrix const cross =
                        internal::lowerSymmetricBlock(m_infoMatrix, target.differences, given.differences);
                    vector.noalias() -= internal::multiply(cross, offset);
                }
                return Result{targetSpace, std::move(reference), std::move(matrix), std::move(vector)};
            }
        );
    }
 
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>;

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] std::expected<InfoGaussianParams<KeptSpace, Scalar>, GaussianError>
    marginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        using Result = InfoGaussianParams<KeptSpace, Scalar>;
        using Expected = std::expected<Result, GaussianError>;
        constexpr int D = DifferenceDimensionAtCompileTime;
        constexpr int DMax = MaxDifferenceDimensionAtCompileTime;
        constexpr int KeptSize = Result::DifferenceDimensionAtCompileTime;
        constexpr int KeptMaxSize = Result::MaxDifferenceDimensionAtCompileTime;
        Eigen::Index const n = differenceDimension();
        auto const rest =
            internal::resolveIndices<D, DMax>(n, internal::complementIndices<D, DMax>(n, kept.differences));
        using RestIndices = std::remove_cvref_t<decltype(rest)>;
        constexpr int RestSize = RestIndices::RowsAtCompileTime;
        constexpr int RestMaxSize = RestIndices::MaxRowsAtCompileTime;
        using RestMatrix = internal::MatrixOf<Scalar, RestSize, RestSize, RestMaxSize, RestMaxSize>;
        using RestVector = Eigen::Matrix<Scalar, RestSize, 1, Eigen::ColMajor, RestMaxSize, 1>;
        using CrossMatrix = internal::MatrixOf<Scalar, RestSize, KeptSize, RestMaxSize, KeptMaxSize>;
 
        typename Result::VectorType reference = m_reference(kept.values);
        typename Result::MatrixType matrix =
            internal::lowerSymmetricBlock(m_infoMatrix, kept.differences, kept.differences);
        typename Result::DifferenceVectorType vector = m_infoVector(kept.differences);
        // MSVC can't form the empty products of a marginal over the whole space or over nothing.
        if constexpr (RestSize != 0 && KeptSize != 0)
        {
            if (rest.size() > 0 && kept.differences.size() > 0)
            {
                auto const llt =
                    Eigen::LLT<RestMatrix>{RestMatrix{internal::lowerSymmetricBlock(m_infoMatrix, rest, rest)}};
                if (llt.info() != Eigen::Success)
                {
                    return Expected{std::unexpected{GaussianError::Singular}};
                }
                if (!allFinite(llt.matrixLLT().diagonal()))
                {
                    return Expected{std::unexpected{GaussianError::NonFinite}};
                }
                CrossMatrix const cross = internal::lowerSymmetricBlock(m_infoMatrix, rest, kept.differences);
                CrossMatrix solved = cross;
                internal::solveInPlace(llt, solved);
                RestVector const restVector = m_infoVector(rest);
                matrix.noalias() -= internal::multiply(cross.transpose(), solved);
                vector.noalias() -= internal::multiply(solved.transpose(), restVector);
            }
        }
        return Expected{Result{keptSpace, std::move(reference), std::move(matrix), std::move(vector)}};
    }

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] InfoGaussianParams<KeptSpace, Scalar>
    densityMarginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        using Result = InfoGaussianParams<KeptSpace, Scalar>;
        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();
 
        MatrixType const factor = m_llt->matrixU();
        DifferenceVectorType whitened = m_infoVector;
        internal::solveInPlace(m_llt->matrixL(), whitened);
        AugmentedType stacked(n, n + 1);
        stacked.leftCols(r) = factor(Eigen::placeholders::all, rest);
        stacked.middleCols(r, k) = factor(Eigen::placeholders::all, kept.differences);
        stacked.col(n) = whitened;
        triangularize(stacked);
 
        typename Result::VectorType reference = m_reference(kept.values);
        auto matrix = typename Result::MatrixType{};
        matrix.resize(k, k);
        auto vector = typename Result::DifferenceVectorType{};
        vector.resize(k);
        // MSVC can't form the empty products of a marginal over nothing.
        if constexpr (Result::DifferenceDimensionAtCompileTime != 0)
        {
            typename Result::MatrixType const trailing = stacked.block(r, r, k, k);
            typename Result::DifferenceVectorType const tail = stacked.col(n).tail(k);
            matrix.noalias() = internal::multiply(trailing.transpose(), trailing);
            vector.noalias() = internal::multiply(trailing.transpose(), tail);
        }
        return Result{keptSpace, std::move(reference), std::move(matrix), std::move(vector)};
    }

    [[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_llt->solve(m_infoVector);
        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::Lower>(m_llt->matrixLLT()).transpose(), residual).eval();
    }

    void
    colorInto(auto & colored, auto const & normals) const
    {
        colored = normals;
        internal::solveInPlace(m_llt->matrixU(), colored);
    }

    void
    fitSpace()
    {
        internal::fitValueDimension(m_space, m_reference.size());
        EMRYS_FILTER_ASSERT(
            m_infoMatrix.rows() == differenceDimension() && m_infoMatrix.cols() == differenceDimension()
            && m_infoVector.size() == differenceDimension()
            && "the precision and information vector must have a row per difference coordinate of the space"
        );
        normalizeValues(m_space, m_reference);
    }
 
    VectorType m_reference;                         
    MatrixType m_infoMatrix;                        
    DifferenceVectorType m_infoVector;              
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Space_ m_space;  
 
    // Factorization cache, filled lazily by factorize().
    mutable std::optional<Eigen::LLT<MatrixType>> m_llt;
 
    mutable Scalar m_logDeterminant{};  
    mutable bool m_factorValid{false};  
    mutable VectorType m_densityMean;   
    internal::Stamp m_stamp;            
};

template <Space Space_, RealScalar Scalar_ = double>
using InfoGaussian = FormGaussian<InfoGaussianParams<Space_, Scalar_>>;

template <Space Space_, RealScalar Scalar_>
[[nodiscard]] InfoGaussianParams<Space_, Scalar_>
multiply(InfoGaussianParams<Space_, Scalar_> const & first, InfoGaussianParams<Space_, Scalar_> const & second)
{
    using Result = InfoGaussianParams<Space_, Scalar_>;
    EMRYS_FILTER_ASSERT(first.differenceDimension() == second.differenceDimension());
    Result moved = second;
    moved.setReference(first.reference());
    // Built from parts rather than as a copy of first, whose cached factorization would be copied
    // only to be invalidated. Only the lower triangles are reliable, as in every stored matrix.
    return Result{
        first.space(),
        first.reference(),
        typename Result::MatrixType{first.infoMatrix() + moved.infoMatrix()},
        typename Result::DifferenceVectorType{first.infoVector() + moved.infoVector()}
    };
}

template <GaussianForm G>
[[nodiscard]] std::expected<InfoGaussianParams<typename G::SpaceType, typename G::Scalar>, GaussianError>
toInfo(G const & form)
{
    using Result = InfoGaussianParams<typename G::SpaceType, typename G::Scalar>;
    auto precision = form.precision();
    if (!precision)
    {
        return std::unexpected{precision.error()};
    }
    if constexpr (isInfoForm<G>)
    {
        return Result{form.space(), form.reference(), std::move(*precision), form.infoVector()};
    }
    else
    {
        std::expected<typename G::VectorType, GaussianError> mean = form.mean();
        if (!mean)
        {
            return std::unexpected{mean.error()};
        }
        return Result{
            form.space(),
            std::move(*mean),
            std::move(*precision),
            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, InfoGaussianParams<Space_, Scalar_> const & gaussian, Indent indent)
{
    stream << indent << "information Gaussian";
    printDimension(stream, gaussian.valueDimension(), gaussian.differenceDimension());
    printMoments(stream, indent.deeper(), gaussian.moments());
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO