GaussianForm.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <expected>
#include <random>
#include <utility>
 
#include <EmrysFilter/Core/ChiSquared.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Reset.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

template <typename G>
concept GaussianForm = requires(G const & g) {
    typename G::SpaceType;
    typename G::Scalar;
    typename G::VectorType;
    typename G::MatrixType;
    { G::IsInfoForm } -> std::convertible_to<bool>;
    { G::IsSqrtForm } -> std::convertible_to<bool>;
    { g.space() } -> std::same_as<typename G::SpaceType const &>;
    { g.valueDimension() } -> std::same_as<Eigen::Index>;
    { g.differenceDimension() } -> std::same_as<Eigen::Index>;
    { g.mean() } -> std::convertible_to<std::expected<typename G::VectorType, GaussianError>>;
    { g.covariance() } -> std::same_as<std::expected<typename G::MatrixType, GaussianError>>;
    { g.precision() } -> std::same_as<std::expected<typename G::MatrixType, GaussianError>>;
    { g.covarianceRoot() } -> std::same_as<std::expected<typename G::MatrixType, GaussianError>>;
    { g.precisionRoot() } -> std::same_as<std::expected<typename G::MatrixType, GaussianError>>;
};

template <GaussianForm G>
inline constexpr bool isInfoForm = G::IsInfoForm;

template <GaussianForm G>
inline constexpr bool isSqrtForm = G::IsSqrtForm;

template <GaussianForm G>
[[nodiscard]] std::expected<GaussianParams<typename G::SpaceType, typename G::Scalar>, GaussianError>
toGaussian(G const & form)
{
    std::expected<typename G::VectorType, GaussianError> mean = form.mean();
    if (!mean)
    {
        return std::unexpected{mean.error()};
    }
    auto covariance = form.covariance();
    if (!covariance)
    {
        return std::unexpected{covariance.error()};
    }
    if constexpr (isInfoForm<G> && !hasIdentityReset<typename G::SpaceType>)
    {
        auto const step = Emrys::minus(form.space(), *mean, form.reference());
        resetCovariance(form.space(), form.reference(), step, *covariance);
    }
    return GaussianParams<typename G::SpaceType, typename G::Scalar>{
        form.space(),
        std::move(*mean),
        std::move(*covariance)
    };
}

template <typename Params_>
class FormGaussian
{
public:
 
    using ParamsType = Params_;                                           
    using SpaceType = typename Params_::SpaceType;                        
    using Scalar = typename Params_::Scalar;                              
    using Index = Eigen::Index;                                           
    using VectorType = typename Params_::VectorType;                      
    using DifferenceVectorType = typename Params_::DifferenceVectorType;  
    using MatrixType = typename Params_::MatrixType;                      
    using MomentsType = typename Params_::MomentsType;                    

    static constexpr int ValueDimensionAtCompileTime = Params_::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = Params_::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = Params_::DifferenceDimensionAtCompileTime;

    static constexpr int MaxDifferenceDimensionAtCompileTime = Params_::MaxDifferenceDimensionAtCompileTime;
 
    static constexpr bool IsParams = false;  

    [[nodiscard]] Index
    valueDimension() const
    {
        return params().valueDimension();
    }

    [[nodiscard]] Index
    differenceDimension() const
    {
        return params().differenceDimension();
    }

    [[nodiscard]] SpaceType const &
    space() const
    {
        return params().space();
    }

    [[nodiscard]] VectorType const &
    mean() const
    {
        return params().densityMean();
    }

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

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

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

    [[nodiscard]] Scalar
    logDeterminant() const
    {
        return params().m_logDeterminant;
    }

    [[nodiscard]] MatrixType const &
    root() const
        requires(Params_::IsSqrtForm && !Params_::IsInfoForm)
    {
        return params().root();
    }

    [[nodiscard]] Scalar
    logPeak() const
    {
        Scalar const log2Pi = Emrys::log2Pi<Scalar>();
        return Scalar(-0.5) * (internal::scalarFromCount<Scalar>(differenceDimension()) * log2Pi + logDeterminant());
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    squaredMahalanobisDistance(Eigen::MatrixBase<Derived> const & x) const
    {
        static_assert(Derived::IsVectorAtCompileTime, "squaredMahalanobisDistance() expects a vector");
        EMRYS_FILTER_ASSERT(x.size() == valueDimension());
        DifferenceVectorType residual = Emrys::minus(space(), x, mean());
        params().whitenInPlace(residual);
        return residual.squaredNorm();
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    confidence(Eigen::MatrixBase<Derived> const & x) const
    {
        if (differenceDimension() == 0)
        {
            return Scalar(0);
        }
        return chiSquaredCdf<Scalar>(
            squaredMahalanobisDistance(x),
            internal::scalarFromCount<Scalar>(differenceDimension())
        );
    }

    template <typename Derived>
    [[nodiscard]] bool
    contains(Eigen::MatrixBase<Derived> const & x, Scalar probability) const
    {
        EMRYS_FILTER_ASSERT(probability > 0 && probability < 1 && "the probability must be in (0, 1)");
        if (differenceDimension() == 0)
        {
            return true;
        }
        return squaredMahalanobisDistance(x)
            <= chiSquaredQuantile<Scalar>(probability, internal::scalarFromCount<Scalar>(differenceDimension()));
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        return logPeak() - Scalar(0.5) * squaredMahalanobisDistance(x);
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    density(Eigen::MatrixBase<Derived> const & x) const
    {
        return Eigen::numext::exp(logDensity(x));
    }

    template <typename Derived>
    [[nodiscard]] typename Derived::PlainObject
    colorNormals(Eigen::MatrixBase<Derived> const & z) const
    {
        EMRYS_FILTER_ASSERT(z.rows() == differenceDimension());
        auto result = typename Derived::PlainObject{};
        result.resize(z.rows(), z.cols());
        params().colorInto(result, z);
        return result;
    }
 
private:
 
    friend Params_;  

    explicit FormGaussian(Params_ const & params)
        : m_check{params}
    {}

    [[nodiscard]] Params_ const &
    params() const
    {
        return m_check.get();
    }
 
    internal::ValuesCheck<Params_> m_check;  
};

template <typename Params_, std::uniform_random_bit_generator G>
[[nodiscard]] SampleSet<typename Params_::SpaceType, typename Params_::Scalar>
drawSamples(FormGaussian<Params_> const & gaussian, Eigen::Index n, G & rng)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using Space = typename Params_::SpaceType;
    using Scalar = typename Params_::Scalar;
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const z = internal::standardNormal<Normals>(gaussian.differenceDimension(), n, rng);
    return SampleSet<Space, Scalar>{Emrys::plus(gaussian.space(), gaussian.mean(), gaussian.colorNormals(z))};
}

template <typename Params_, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<typename Params_::SpaceType, typename Params_::Scalar>::VectorType
drawSample(FormGaussian<Params_> const & gaussian, G & rng)
{
    using MomentsType = Moments<typename Params_::SpaceType, typename Params_::Scalar>;
    using Normals = DifferenceVector<typename Params_::SpaceType, typename Params_::Scalar>;
    Normals const z = internal::standardNormalVector<Normals>(gaussian.differenceDimension(), rng);
    return typename MomentsType::VectorType{Emrys::plus(gaussian.space(), gaussian.mean(), gaussian.colorNormals(z))};
}

template <typename Params_, std::uniform_random_bit_generator G>
[[nodiscard]] SamplesWithLogDensity<typename Params_::SpaceType, typename Params_::Scalar>
drawSamplesWithLogDensity(FormGaussian<Params_> const & gaussian, Eigen::Index n, G & rng)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using Space = typename Params_::SpaceType;
    using Scalar = typename Params_::Scalar;
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const z = internal::standardNormal<Normals>(gaussian.differenceDimension(), n, rng);
    auto result = SamplesWithLogDensity<Space, Scalar>{
        Emrys::plus(gaussian.space(), gaussian.mean(), gaussian.colorNormals(z)),
        {}
    };
    if constexpr (isLinear<Space>)
    {
        result.logDensity = gaussian.logPeak() - Scalar(0.5) * z.colwise().squaredNorm().transpose().array();
    }
    else
    {
        result.logDensity.resize(n);
        for (Eigen::Index j = 0; j < n; ++j)
        {
            result.logDensity(j) = gaussian.logDensity(result.samples.col(j));
        }
    }
    return result;
}

template <typename Params_, std::uniform_random_bit_generator G>
[[nodiscard]] SampleWithLogDensity<typename Params_::SpaceType, typename Params_::Scalar>
drawSampleWithLogDensity(FormGaussian<Params_> const & gaussian, G & rng)
{
    using Space = typename Params_::SpaceType;
    using Scalar = typename Params_::Scalar;
    using MomentsType = Moments<Space, Scalar>;
    using Normals = DifferenceVector<Space, Scalar>;
    Normals const z = internal::standardNormalVector<Normals>(gaussian.differenceDimension(), rng);
    auto result = SampleWithLogDensity<Space, Scalar>{
        typename MomentsType::VectorType{Emrys::plus(gaussian.space(), gaussian.mean(), gaussian.colorNormals(z))},
        0
    };
    if constexpr (isLinear<Space>)
    {
        result.logDensity = gaussian.logPeak() - Scalar(0.5) * z.squaredNorm();
    }
    else
    {
        result.logDensity = gaussian.logDensity(result.sample);
    }
    return result;
}
 
}  // namespace Emrys