UniformMixture.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Distributions/Component.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/Product.h>
#include <EmrysFilter/Distributions/Uniform.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

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

    using CountsType =
        internal::IndexVector<MomentsType::ValueDimensionAtCompileTime, MomentsType::MaxValueDimensionAtCompileTime>;

    struct Element
    {
        Scalar logWeight;     
        MomentsType moments;  
    };
 
    using value_type = Element;                                     
    using const_iterator = internal::ElementIterator<UniformGrid>;  
    using iterator = const_iterator;                                

    UniformGrid(
        SpaceType space,
        VectorType corner,
        VectorType width,
        CountsType counts,
        MatrixType covariance,
        Scalar relativeWidth
    )
        : m_space{std::move(space)}
        , m_corner{std::move(corner)}
        , m_width{std::move(width)}
        , m_counts{std::move(counts)}
        , m_covariance{std::move(covariance)}
    {
        EMRYS_FILTER_ASSERT(
            isFinite(relativeWidth) && relativeWidth > Scalar(0) && "the relative width must be positive and finite"
        );
        EMRYS_FILTER_ASSERT(
            m_corner.size() == m_width.size() && m_corner.size() == m_counts.size()
            && m_covariance.rows() == m_corner.size() && m_covariance.cols() == m_corner.size()
            && "one corner, width, count, and covariance row per coordinate"
        );
        Index gridded = 0;
        for (Index i = 0; i < m_counts.size(); ++i)
        {
            EMRYS_FILTER_ASSERT(m_counts(i) >= 0 && "a count must not be negative");
            if (m_counts(i) > 0)
            {
                ++gridded;
                m_size *= m_counts(i);
                Scalar const deviation = relativeWidth * m_width(i) / Scalar(m_counts(i));
                m_covariance(i, i) = deviation * deviation;
            }
        }
        EMRYS_FILTER_ASSERT(gridded > 0 && "the grid needs a gridded coordinate");
        m_logWeight = -Eigen::numext::log(Scalar(m_size));
    }

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

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

    [[nodiscard]] Index
    size() const
    {
        return m_size;
    }

    [[nodiscard]] CountsType const &
    counts() const
    {
        return m_counts;
    }

    [[nodiscard]] Element
    operator[](Index k) const
    {
        EMRYS_FILTER_ASSERT(k >= 0 && k < m_size && "component index out of range");
        VectorType mean = m_corner;
        Index rest = k;
        for (Index i = 0; i < m_counts.size(); ++i)
        {
            Index const count = m_counts(i);
            if (count > 0)
            {
                Index const j = rest % count;
                rest /= count;
                mean(i) = m_corner(i) + (Scalar(j) + Scalar(0.5)) * m_width(i) / Scalar(count);
            }
        }
        return Element{m_logWeight, MomentsType{std::move(mean), m_covariance}};
    }

    [[nodiscard]] const_iterator
    begin() const
    {
        return const_iterator{*this, 0};
    }

    [[nodiscard]] const_iterator
    end() const
    {
        return const_iterator{*this, m_size};
    }
 
private:
 
    SpaceType m_space;        
    VectorType m_corner;      
    VectorType m_width;       
    CountsType m_counts;      
    MatrixType m_covariance;  
    Index m_size = 1;         
    Scalar m_logWeight{};     
};
 
namespace internal
{

template <Space S, RealScalar Scalar, typename CountOf>
[[nodiscard]] UniformGrid<S, Scalar>
uniformGrid(Uniform<S, Scalar> const & uniform, CountOf const & countOf, Scalar relativeWidth)
{
    using Grid = UniformGrid<S, Scalar>;
    Eigen::Index const n = uniform.valueDimension();
    typename Grid::CountsType counts(n);
    for (Eigen::Index i = 0; i < n; ++i)
    {
        counts(i) = countOf(i);
        EMRYS_FILTER_ASSERT(counts(i) >= 1 && "each coordinate needs at least one component");
    }
    return Grid{
        uniform.space(),
        uniform.lower(),
        uniform.width(),
        std::move(counts),
        Grid::MatrixType::Zero(n, n),
        relativeWidth
    };
}

template <typename F>
inline constexpr bool isUniform = false;

template <Space S, RealScalar Scalar>
inline constexpr bool isUniform<Uniform<S, Scalar>>{true};

template <typename F>
inline constexpr bool isGaussianFactor = isGaussian<F>;

template <typename F, typename BlockSpace>
inline constexpr bool isGaussianFactor<EachRepetition<F, BlockSpace>>{isGaussian<F>};


template <Space S, RealScalar Scalar, typename... Factors, typename CountOf>
[[nodiscard]] UniformGrid<S, Scalar>
productGrid(
    Product<S, Scalar, Factors...> const & product,
    CountOf const & countOf,
    Eigen::Index givenCounts,
    Scalar relativeWidth
)
{
    static_assert(
        ((isUniform<Factors> || isGaussianFactor<Factors>) && ...),
        "uniformComponents() of a Product takes Uniform factors placed with on<Path>() and GaussianParams factors"
    );
    static_assert((isUniform<Factors> || ...), "uniformComponents() of a Product needs a Uniform factor to grid");
    using Grid = UniformGrid<S, Scalar>;
    using Index = Eigen::Index;
    Index const n = product.valueDimension();
    auto corner = typename Grid::VectorType(n);
    typename Grid::MatrixType covariance = Grid::MatrixType::Zero(n, n);
    typename Grid::VectorType width = Grid::VectorType::Zero(n);
    typename Grid::CountsType counts = Grid::CountsType::Zero(n);
    [&]<std::size_t... K>(std::index_sequence<K...>) {
        (
            [&] {
                auto const & factor = product.template factor<K>();
                auto const & indices = product.template indices<K>();
                if constexpr (isUniform<std::tuple_element_t<K, std::tuple<Factors...>>>)
                {
                    corner(indices) = factor.lower();
                    width(indices) = factor.width();
                    counts(indices).setOnes();
                }
                else
                {
                    auto const & moments = factor.moments();
                    using FactorMatrix = std::remove_cvref_t<decltype(moments.covarianceLower())>;
                    corner(indices) = moments.mean();
                    covariance(indices, indices) =
                        FactorMatrix{moments.covarianceLower().template selfadjointView<Eigen::Lower>()};
                }
            }(),
            ...
        );
    }(std::index_sequence_for<Factors...>{});
 
    Index gridded = 0;
    for (Index i = 0; i < n; ++i)
    {
        if (counts(i) > 0)
        {
            counts(i) = countOf(gridded++);
            EMRYS_FILTER_ASSERT(counts(i) >= 1 && "each gridded coordinate needs at least one component");
        }
    }
    EMRYS_FILTER_ASSERT(
        (givenCounts < 0 || givenCounts == gridded) && "give one count per coordinate of the Uniform factors"
    );
    return Grid{
        product.space(),
        std::move(corner),
        std::move(width),
        std::move(counts),
        std::move(covariance),
        relativeWidth
    };
}
 
}  // namespace internal

template <Space S, RealScalar Scalar, typename CountsDerived>
    requires std::integral<typename CountsDerived::Scalar>
[[nodiscard]] UniformGrid<S, Scalar>
uniformComponents(
    Uniform<S, Scalar> const & uniform,
    Eigen::DenseBase<CountsDerived> const & counts,
    std::type_identity_t<Scalar> relativeWidth = Scalar(0.6)
)
{
    EMRYS_FILTER_ASSERT(counts.size() == uniform.valueDimension() && "give one count per coordinate");
    return internal::uniformGrid(
        uniform,
        [&](Eigen::Index i) {
            return Eigen::Index(counts(i));
        },
        relativeWidth
    );
}

template <Space S, RealScalar Scalar>
[[nodiscard]] UniformGrid<S, Scalar>
uniformComponents(
    Uniform<S, Scalar> const & uniform,
    Eigen::Index count,
    std::type_identity_t<Scalar> relativeWidth = Scalar(0.6)
)
{
    return internal::uniformGrid(
        uniform,
        [&](Eigen::Index) {
            return count;
        },
        relativeWidth
    );
}

template <Space S, RealScalar Scalar, typename... Factors, typename CountsDerived>
    requires std::integral<typename CountsDerived::Scalar>
[[nodiscard]] UniformGrid<S, Scalar>
uniformComponents(
    Product<S, Scalar, Factors...> const & product,
    Eigen::DenseBase<CountsDerived> const & counts,
    std::type_identity_t<Scalar> relativeWidth = Scalar(0.6)
)
{
    return internal::productGrid(
        product,
        [&](Eigen::Index j) {
            EMRYS_FILTER_ASSERT(j < counts.size() && "give one count per coordinate of the Uniform factors");
            return Eigen::Index(counts(j));
        },
        counts.size(),
        relativeWidth
    );
}

template <Space S, RealScalar Scalar, typename... Factors>
[[nodiscard]] UniformGrid<S, Scalar>
uniformComponents(
    Product<S, Scalar, Factors...> const & product,
    Eigen::Index count,
    std::type_identity_t<Scalar> relativeWidth = Scalar(0.6)
)
{
    return internal::productGrid(
        product,
        [&](Eigen::Index) {
            return count;
        },
        -1,
        relativeWidth
    );
}
 
}  // namespace Emrys