SqrtGaussian.h

Go to the documentation of this file.

#pragma once

 
#include <expected>
#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 <typename C>
struct ComponentTraits;
 
namespace internal
{
 
template <typename Source, typename Selection, typename SubSpace>
class SqrtComponentMarginal;
 
}  // namespace internal

template <Space Space_, RealScalar Scalar_ = double>
class SqrtGaussianParams
{
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<SqrtGaussianParams>;                     

    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 = false;  
    static constexpr bool IsSqrtForm = true;   

    SqrtGaussianParams()
        : SqrtGaussianParams{Space_{}}
    {}

    explicit SqrtGaussianParams(Space_ space)
        : m_space{std::move(space)}
    {
        m_mean.setZero(Emrys::valueDimension(m_space));
        m_root.setZero(differenceDimension(), differenceDimension());
    }

    template <MatrixExpression Mean, MatrixExpression Root>
        requires internal::SizedByLength<Space_>
    SqrtGaussianParams(Mean && mean, Root && root)
        : m_mean{std::forward<Mean>(mean)}
        , m_root{std::forward<Root>(root)}
    {
        fitSpace();
    }

    template <MatrixExpression Mean, MatrixExpression Root>
    SqrtGaussianParams(Space_ space, Mean && mean, Root && root)
        : m_mean{std::forward<Mean>(mean)}
        , m_root{std::forward<Root>(root)}
        , 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 &
    mean() const
    {
        return m_mean;
    }

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

    [[nodiscard]] MomentsType
    moments() const
    {
        return MomentsType{m_mean, *(covariance())};
    }

    template <MatrixExpression Mean>
    void
    setMean(Mean && mean)
    {
        m_stamp.touchValues();
        EMRYS_FILTER_ASSERT(mean.size() == valueDimension());
        m_mean = std::forward<Mean>(mean);
        normalizeValues(m_space, m_mean);
    }

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

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

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

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

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

    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);
        });
    }

    [[nodiscard]] std::expected<void, GaussianError>
    factorize() const
    {
        if (auto checked = checkRoot(); !checked)
        {
            return std::unexpected{checked.error()};
        }
        // Unchanged, it writes nothing, so factorizing what is factorized only reads.
        Scalar const logDeterminant = Scalar(2) * m_root.diagonal().array().abs().log().sum();
        if (m_logDeterminant != logDeterminant)
        {
            m_logDeterminant = logDeterminant;
        }
        return {};
    }

    template <typename R>
    [[nodiscard]] std::expected<void, GaussianError>
    factorize(R const & /*repair*/) const
    {
        static_assert(
            sizeof(R) == 0,
            "a square-root form's root is positive semidefinite by construction and has no repair; repairing a "
            "covariance is the covariance form's (decisions.md [gaussian-forms])"
        );
        return factorize();
    }

    [[nodiscard]] std::expected<DensityType, GaussianError>
    toDensity() const &
    {
        if (auto factorized = factorize(); !factorized)
        {
            return std::unexpected{factorized.error()};
        }
        return DensityType{*this};
    }

    void
    toDensity() const && = delete;
 
private:
 
    friend DensityType;  

    friend struct ComponentTraits<SqrtGaussianParams>;

    template <typename, typename, typename>
    friend class internal::SqrtComponentMarginal;

    template <typename>
    friend class internal::ValuesCheck;

    [[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::NotPositiveDefinite};
        }
        return {};
    }

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] SqrtGaussianParams<KeptSpace, Scalar>
    marginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        using Result = SqrtGaussianParams<KeptSpace, Scalar>;
        using Rows = internal::MatrixOf<
            Scalar,
            Result::DifferenceDimensionAtCompileTime,
            DifferenceDimensionAtCompileTime,
            Result::MaxDifferenceDimensionAtCompileTime,
            MaxDifferenceDimensionAtCompileTime>;
        typename Result::VectorType mean = m_mean(kept.values);
        Rows const rows = m_root(kept.differences, Eigen::placeholders::all);
        return Result{keptSpace, std::move(mean), internal::lowerTriangularRoot(rows)};
    }

    template <typename Kept, typename KeptSpace>
    [[nodiscard]] SqrtGaussianParams<KeptSpace, Scalar>
    densityMarginalAt(Kept const & kept, KeptSpace const & keptSpace) const
    {
        return marginalAt(kept, keptSpace);
    }

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

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

    void
    whitenInPlace(DifferenceVectorType & residual) const
    {
        m_root.template triangularView<Eigen::Lower>().solveInPlace(residual);
    }

    void
    colorInto(auto & colored, auto const & normals) const
    {
        colored.noalias() = internal::multiply(internal::asTriangular<Eigen::Lower>(m_root), normals);
    }

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

template <typename Source, typename Selection, typename SubSpace>
class SqrtComponentMarginal : public SqrtGaussianParams<SubSpace, typename Source::Scalar>
{
public:
 
    using ParamsType = SqrtGaussianParams<SubSpace, typename Source::Scalar>;  

    SqrtComponentMarginal(Source const & source, Selection shared)
        : ParamsType{source.marginalAt(
              Coordinates<
                  decltype(shared.resolved->indices) const &,
                  decltype(shared.resolved->differenceIndices) const &>{
                  shared.resolved->indices,
                  shared.resolved->differenceIndices
              },
              shared.resolved->space
          )}
    {}

    [[nodiscard]] ParamsType
    copy() const
    {
        return *this;
    }
};
 
}  // namespace internal

template <Space Space_, RealScalar Scalar_ = double>
using SqrtGaussian = FormGaussian<SqrtGaussianParams<Space_, Scalar_>>;

template <GaussianForm G>
[[nodiscard]] std::expected<SqrtGaussianParams<typename G::SpaceType, typename G::Scalar>, GaussianError>
toSqrt(G const & form)
{
    using Result = SqrtGaussianParams<typename G::SpaceType, typename G::Scalar>;
    std::expected<typename G::VectorType, GaussianError> mean = form.mean();
    if (!mean)
    {
        return std::unexpected{mean.error()};
    }
    auto root = form.covarianceRoot();
    if (!root)
    {
        return std::unexpected{root.error()};
    }
    if constexpr (isInfoForm<G>)
    {
        typename Result::MatrixType lower = internal::lowerTriangularRoot(*root);
        if constexpr (!hasIdentityReset<typename G::SpaceType>)
        {
            auto const step = Emrys::minus(form.space(), *mean, form.reference());
            auto const g = resetJacobian(form.space(), form.reference(), step);
            lower = internal::lowerTriangularRoot(typename Result::MatrixType{internal::multiply(g, lower)});
        }
        return Result{form.space(), std::move(*mean), std::move(lower)};
    }
    else
    {
        return Result{form.space(), std::move(*mean), std::move(*root)};
    }
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

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