Component.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <type_traits>
 
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Triangularization.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/GaussianForm.h>
#include <EmrysFilter/Distributions/Params.h>
#include <EmrysFilter/Distributions/SqrtGaussian.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{
 
namespace internal
{

template <typename C, typename F>
decltype(auto)
keepingSpace(C & component, F & f)
{
    [[maybe_unused]] auto const before = typename C::SpaceType{component.space()};
    if constexpr (std::is_void_v<std::invoke_result_t<F &, C &>>)
    {
        f(component);
        EMRYS_FILTER_ASSERT(component.space() == before && "the component must keep its space");
    }
    else
    {
        std::invoke_result_t<F &, C &> result = f(component);
        EMRYS_FILTER_ASSERT(component.space() == before && "the component must keep its space");
        return result;
    }
}

template <typename Scalar>
struct CollapseWeights
{
    Scalar logTotal;    
    Eigen::Index kept;  
    Eigen::Index only;  
};

template <typename Scalar, typename LogWeightAt>
[[nodiscard]] CollapseWeights<Scalar>
collapseWeights(Eigen::Index count, LogWeightAt const & logWeightAt)
{
    EMRYS_FILTER_ASSERT(count > 0 && "a collapse needs a component");
    auto total = LogSumExp<Scalar>{};
    auto result = CollapseWeights<Scalar>{Scalar(0), 0, 0};
    for (Eigen::Index i = 0; i < count; ++i)
    {
        Scalar const logWeight = logWeightAt(i);
        total.add(logWeight);
        if (logWeight > -Eigen::NumTraits<Scalar>::infinity())
        {
            ++result.kept;
            result.only = i;
        }
    }
    result.logTotal = total.value();
    EMRYS_FILTER_ASSERT(isFinite(result.logTotal) && "the weights must be finite and not all zero");
    return result;
}
 
}  // namespace internal

template <typename C>
struct ComponentTraits;

template <typename C>
inline constexpr bool isGaussian = false;

template <Space Space_, RealScalar Scalar_>
inline constexpr bool isGaussian<GaussianParams<Space_, Scalar_>>{true};


template <Space Space_, RealScalar Scalar_>
struct ComponentTraits<GaussianParams<Space_, Scalar_>>
{
    using Type = GaussianParams<Space_, Scalar_>;  

    template <typename Selection, typename SubSpace>
    using MarginalType = GaussianMarginal<Type, Selection, SubSpace>;

    template <typename Derived>
    [[nodiscard]] static Scalar_
    logDensity(Type const & gaussian, Eigen::MatrixBase<Derived> const & x)
    {
        return internal::FactorizedGaussian::logDensity(gaussian, x);
    }

    template <typename ComponentAt, typename LogWeightAt>
    [[nodiscard]] static Type
    collapse(Eigen::Index count, ComponentAt componentAt, LogWeightAt logWeightAt)
    {
        auto const weights = internal::collapseWeights<Scalar_>(count, logWeightAt);
        if (weights.kept == 1)
        {
            return Type{componentAt(weights.only)};
        }
        Type const & first = componentAt(0);
        return Type{
            first.space(),
            mixtureMomentsIn<typename Type::MomentsType>(
                first.space(),
                count,
                [&](Eigen::Index i) -> decltype(auto) {
                    return componentAt(i).mean();
                },
                [&](Eigen::Index i) -> decltype(auto) {
                    return componentAt(i).covarianceLower();
                },
                logWeightAt,
                weights.logTotal
            )
        };
    }

    [[nodiscard]] static typename Type::MomentsType const &
    moments(Type const & gaussian)
    {
        return gaussian.moments();
    }
};

template <Space Space_, RealScalar Scalar_>
struct ComponentTraits<SqrtGaussianParams<Space_, Scalar_>>
{
    using Type = SqrtGaussianParams<Space_, Scalar_>;  

    template <typename Selection, typename SubSpace>
    using MarginalType = internal::SqrtComponentMarginal<Type, Selection, SubSpace>;

    template <typename Derived>
    [[nodiscard]] static Scalar_
    logDensity(Type const & gaussian, Eigen::MatrixBase<Derived> const & x)
    {
        typename Type::DifferenceVectorType residual = Emrys::minus(gaussian.space(), x, gaussian.mean());
        gaussian.whitenInPlace(residual);
        Scalar_ const dimension = internal::scalarFromCount<Scalar_>(gaussian.differenceDimension());
        return Scalar_(-0.5) * (dimension * log2Pi<Scalar_>() + gaussian.m_logDeterminant + residual.squaredNorm());
    }

    template <typename ComponentAt, typename LogWeightAt>
    [[nodiscard]] static Type
    collapse(Eigen::Index count, ComponentAt componentAt, LogWeightAt logWeightAt)
    {
        auto const weights = internal::collapseWeights<Scalar_>(count, logWeightAt);
        if (weights.kept == 1)
        {
            return Type{componentAt(weights.only)};
        }
        Type const & first = componentAt(0);
        auto const & space = first.space();
        typename Type::VectorType mean = pointMomentsIn<typename Type::MomentsType>(
                                             space,
                                             count,
                                             [&](Eigen::Index i) -> decltype(auto) {
                                                 return componentAt(i).mean();
                                             },
                                             logWeightAt,
                                             weights.logTotal
        )
                                             .mean();
        constexpr int Rows = Type::DifferenceDimensionAtCompileTime;
        constexpr int MaxRows = Type::MaxDifferenceDimensionAtCompileTime;
        using Stack = internal::MatrixOf<
            Scalar_,
            Rows,
            Rows == Eigen::Dynamic ? Eigen::Dynamic : 2 * Rows + 1,
            MaxRows,
            MaxRows == Eigen::Dynamic ? Eigen::Dynamic : 2 * MaxRows + 1>;
        Eigen::Index const n = first.differenceDimension();
        typename Type::MatrixType factor = Type::MatrixType::Zero(n, n);
        auto stack = Stack{n, 2 * n + 1};
        for (Eigen::Index i = 0; i < count; ++i)
        {
            Scalar_ const logWeight = logWeightAt(i);
            if (logWeight == -Eigen::NumTraits<Scalar_>::infinity())
            {
                continue;
            }
            Type const & component = componentAt(i);
            Scalar_ const root = Eigen::numext::exp(Scalar_(0.5) * (logWeight - weights.logTotal));
            stack.leftCols(n) = factor;
            stack.middleCols(n, n) = root * component.root();
            stack.col(2 * n) = root * Emrys::minus(space, component.mean(), mean);
            factor = internal::lowerTriangularRoot(stack);
        }
        return Type{space, std::move(mean), std::move(factor)};
    }

    [[nodiscard]] static typename Type::MomentsType
    moments(Type const & gaussian)
    {
        return gaussian.moments();
    }
};

template <typename C>
concept Component = std::same_as<Density<Params<C>>, C> && requires(Params<C> const & c) {
    typename Params<C>::SpaceType;
    typename Params<C>::Scalar;
    typename Params<C>::MomentsType;
    { Params<C>::ValueDimensionAtCompileTime } -> std::convertible_to<int>;
    { Params<C>::MaxValueDimensionAtCompileTime } -> std::convertible_to<int>;
    { Params<C>::DifferenceDimensionAtCompileTime } -> std::convertible_to<int>;
    { Params<C>::MaxDifferenceDimensionAtCompileTime } -> std::convertible_to<int>;
    { c.space() } -> std::same_as<typename Params<C>::SpaceType const &>;
    { c.valueDimension() } -> std::same_as<Eigen::Index>;
    { c.differenceDimension() } -> std::same_as<Eigen::Index>;
    { ComponentTraits<Params<C>>::moments(c) } -> std::convertible_to<typename Params<C>::MomentsType>;
    { c.factorize() };
    typename ComponentTraits<Params<C>>::Type;
};
 
namespace internal
{

template <typename C>
concept InfoFormDensity = GaussianForm<Params<C>> && std::same_as<Density<Params<C>>, C> && isInfoForm<Params<C>>;

template <typename C>
concept ComponentOrInfoForm = Component<C> || InfoFormDensity<C>;
 
}  // namespace internal
 
}  // namespace Emrys