Mixture.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <concepts>
#include <expected>
#include <functional>
#include <iterator>
#include <random>
#include <span>
#include <numeric>
#include <optional>
#include <ranges>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Execution.h>
#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/Resampling.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Component.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Composite.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
 
namespace Emrys
{

struct ComponentError
{
    Eigen::Index index;  
    NumericError error;  

    friend bool
    operator==(ComponentError const &, ComponentError const &) = default;
};

using MixtureError = std::variant<WeightError, ComponentError>;

using GaussianMixtureError = MixtureError;
 
template <internal::ComponentOrInfoForm Component_, typename Allocator_>
class Mixture;
 
template <internal::ComponentOrInfoForm Component_, typename Allocator_>
class MixtureParams;
 
template <typename Source, typename Selection, typename SubSpace>
class MixtureMarginal;
 
namespace internal
{
struct MixtureLineageAccess;
}  // namespace internal

template <RealScalar Scalar_, typename Allocator_ = DefaultAllocator>
class MixtureLineage
{
public:
 
    using Scalar = Scalar_;             
    using allocator_type = Allocator_;  
    using Index = Eigen::Index;         
 
    MixtureLineage() = default;

    explicit MixtureLineage(Allocator_ const & allocator)
        : m_logShares(ScalarAllocator(allocator))
    {}

    MixtureLineage(MixtureLineage const & other)
        : MixtureLineage{other, other.get_allocator()}
    {}

    MixtureLineage(MixtureLineage const & other, Allocator_ const & allocator)
        : m_logShares(other.m_logShares, ScalarAllocator(allocator))
        , m_parentCount{other.m_parentCount}
        , m_count{other.m_count}
    {}

    MixtureLineage(MixtureLineage && other) = default;

    MixtureLineage(MixtureLineage && other, Allocator_ const & allocator)
        : m_logShares(std::move(other.m_logShares), ScalarAllocator(allocator))
        , m_parentCount{other.m_parentCount}
        , m_count{other.m_count}
    {}

    MixtureLineage &
    operator=(MixtureLineage const & other) = default;

    MixtureLineage &
    operator=(MixtureLineage && other) = default;
 
    ~MixtureLineage() = default;

    [[nodiscard]] allocator_type
    get_allocator() const
    {
        return allocator_type(m_logShares.get_allocator());
    }

    [[nodiscard]] Index
    count() const
    {
        return m_count;
    }

    [[nodiscard]] Index
    parentCount() const
    {
        return m_parentCount;
    }

    [[nodiscard]] Index
    size() const
    {
        return m_parentCount == 0 ? Index(0) : Index(m_logShares.size()) / m_parentCount;
    }

    [[nodiscard]] std::span<Scalar const>
    logShares(Index i) const
    {
        EMRYS_FILTER_ASSERT(i >= 0 && i < size() && "the component has no row");
        return {m_logShares.data() + i * m_parentCount, std::size_t(m_parentCount)};
    }

    [[nodiscard]] Scalar
    logShare(Index i, Index parent) const
    {
        EMRYS_FILTER_ASSERT(parent >= 0 && parent < m_parentCount && "no such parent");
        return logShares(i)[std::size_t(parent)];
    }

    template <typename Parents, typename Children>
    void
    carry(Eigen::MatrixBase<Parents> const & parents, Eigen::MatrixBase<Children> & children, Scalar orphan = 0) const
    {
        EMRYS_FILTER_ASSERT(
            parents.size() == parentCount() && children.size() == size() && "one value per parent and per component"
        );
        Scalar const none = -Eigen::NumTraits<Scalar>::infinity();
        for (Index i = 0; i < size(); ++i)
        {
            auto const shares = logShares(i);
            Scalar const top = *std::ranges::max_element(shares);
            if (top == none)
            {
                children(i) = orphan;
                continue;
            }
            Scalar sum = 0;
            Scalar total = 0;
            for (Index parent = 0; parent < m_parentCount; ++parent)
            {
                Scalar const share = Eigen::numext::exp(Scalar(shares[std::size_t(parent)] - top));
                sum += share * parents(parent);
                total += share;
            }
            children(i) = sum / total;
        }
    }
 
private:

    template <internal::ComponentOrInfoForm, typename>
    friend class MixtureParams;
 
    friend struct internal::MixtureLineageAccess;  
 
    using ScalarAllocator = internal::ReBound<Allocator_, Scalar>;  

    [[nodiscard]] bool
    tracking() const
    {
        return m_count > 0;
    }

    void
    start(Index components, std::optional<Index> capacity)
    {
        ++m_count;
        m_parentCount = components;
        if (capacity)
        {
            m_logShares.reserve(std::size_t(*capacity * components));
        }
        m_logShares.assign(std::size_t(components * components), -Eigen::NumTraits<Scalar>::infinity());
        for (Index i = 0; i < components; ++i)
        {
            m_logShares[std::size_t(i * components + i)] = Scalar(0);
        }
    }

    void
    reset()
    {
        m_logShares.clear();
        m_parentCount = 0;
        m_count = 0;
    }

    void
    clearRows()
    {
        m_logShares.clear();
    }

    void
    pushOrphan()
    {
        if (tracking())
        {
            m_logShares.resize(m_logShares.size() + std::size_t(m_parentCount), -Eigen::NumTraits<Scalar>::infinity());
        }
    }

    void
    push(std::span<Scalar const> row)
    {
        if (tracking())
        {
            EMRYS_FILTER_ASSERT(Index(row.size()) == m_parentCount && "the row must have a share per parent");
            m_logShares.insert(m_logShares.end(), row.begin(), row.end());
        }
    }

    void
    pushCopyOf(Index i)
    {
        if (tracking())
        {
            std::size_t const end = m_logShares.size();
            m_logShares.resize(end + std::size_t(m_parentCount));
            std::copy_n(
                m_logShares.begin() + i * m_parentCount,
                m_parentCount,
                m_logShares.begin() + std::ptrdiff_t(end)
            );
        }
    }

    void
    moveRow(Index from, Index to)
    {
        if (tracking() && from != to)
        {
            std::copy_n(
                m_logShares.begin() + from * m_parentCount,
                m_parentCount,
                m_logShares.begin() + to * m_parentCount
            );
        }
    }

    void
    truncate(Index rows)
    {
        if (tracking())
        {
            m_logShares.resize(std::size_t(rows * m_parentCount));
        }
    }

    void
    removeAt(Index i)
    {
        if (tracking())
        {
            moveRow(size() - 1, i);
            truncate(size() - 1);
        }
    }

    void
    merge(Index keep, Index absorb, Scalar logKeep, Scalar logAbsorb)
    {
        if (!tracking())
        {
            return;
        }
        auto total = LogSumExp<Scalar>{};
        total.add(logKeep);
        total.add(logAbsorb);
        Scalar const logTotal = total.value();
        if (!isFinite(logTotal))
        {
            return;
        }
        Scalar * const kept = m_logShares.data() + keep * m_parentCount;
        Scalar const * const absorbed = m_logShares.data() + absorb * m_parentCount;
        for (Index n = 0; n < m_parentCount; ++n)
        {
            auto share = LogSumExp<Scalar>{};
            share.add(kept[n] + logKeep - logTotal);
            share.add(absorbed[n] + logAbsorb - logTotal);
            kept[n] = share.value();
        }
    }

    void
    swap(MixtureLineage & other)
    {
        m_logShares.swap(other.m_logShares);
        std::swap(m_parentCount, other.m_parentCount);
        std::swap(m_count, other.m_count);
    }
 
    internal::Buffer<Scalar, Allocator_> m_logShares;  
    Index m_parentCount = 0;                           
    Index m_count = 0;                                 
};
 
namespace internal
{

struct MixtureLineageAccess
{
    template <typename Scalar, typename ToAllocator, typename FromAllocator>
    static void
    assign(MixtureLineage<Scalar, ToAllocator> & to, MixtureLineage<Scalar, FromAllocator> const & from)
    {
        to.m_logShares.assign(from.m_logShares.begin(), from.m_logShares.end());
        to.m_parentCount = from.m_parentCount;
        to.m_count = from.m_count;
    }

    template <typename MixtureType>
    static void
    copy(MixtureType & to, MixtureType const & from)
    {
        to.m_lineage = from.m_lineage;
    }

    template <typename MixtureType>
    static void
    startFrom(MixtureType & to, MixtureType const & from)
    {
        EMRYS_FILTER_ASSERT(to.empty() && "the mixture must have no components");
        to.m_lineage.m_logShares.clear();
        to.m_lineage.m_parentCount = from.m_lineage.m_parentCount;
        to.m_lineage.m_count = from.m_lineage.m_count;
    }

    template <typename MixtureType>
    static void
    inherit(MixtureType & to, MixtureType const & from, Eigen::Index i)
    {
        if (from.m_lineage.tracking())
        {
            to.m_lineage.truncate(to.size() - 1);
            to.m_lineage.push(from.m_lineage.logShares(i));
        }
    }
};

template <typename Source, typename Target, typename Given, typename Derived, typename TargetSpace, typename GivenSpace>
[[nodiscard]] auto
conditionalMixtureIn(
    Source const & source,
    Target const & target,
    Given const & given,
    Eigen::MatrixBase<Derived> const & values,
    TargetSpace const & targetSpace,
    GivenSpace const & givenSpace
)
{
    EMRYS_FILTER_ASSERT(distinctIndices(target.values, given.values) && "target and given indices must be distinct");
    using Params = MixtureParams<Gaussian<TargetSpace, typename Source::Scalar>, typename Source::allocator_type>;
    using Result = std::expected<Params, GaussianError>;
    auto result = Params{targetSpace, source.get_allocator()};
    result.reserve(source.size());
    for (auto const & entry : source)
    {
        auto conditioned = conditionMomentsWithEvidence(
            entry.component.mean(),
            entry.component.covarianceLower(),
            target,
            given,
            values,
            givenSpace,
            targetSpace
        );
        if (!conditioned)
        {
            return Result{std::unexpect, GaussianError::NotPositiveDefinite};
        }
        result.add(
            entry.logWeight + conditioned->logLikelihood,
            std::move(conditioned->moments.mean),
            std::move(conditioned->moments.cov)
        );
    }
    return Result{std::move(result)};
}

template <typename GivenPath, typename Source, typename Derived>
[[nodiscard]] auto
tagConditionalMixture(Source const & source, Eigen::MatrixBase<Derived> const & values)
{
    return tagConditional<GivenPath>(
        source.space(),
        [&](auto const & target, auto const & given, auto const & targetSpace, auto const & givenSpace) {
            return conditionalMixtureIn(source, target, given, values, targetSpace, givenSpace);
        }
    );
}
 
}  // namespace internal

template <internal::ComponentOrInfoForm Component_, typename Allocator_ = DefaultAllocator>
class MixtureParams
{
    static_assert(
        Component<Component_>,
        "a Gaussian in an information form is not a mixture component, having no collapse; convert each "
        "with toGaussian() or toSqrt() and hold a Gaussian or a SqrtGaussian (decisions.md [gaussian-forms])"
    );
 
public:
 
    using ComponentType = Params<Component_>;                                 
    using SpaceType = typename ComponentType::SpaceType;                      
    using allocator_type = Allocator_;                                        
    using MomentsType = typename ComponentType::MomentsType;                  
    using Scalar = typename ComponentType::Scalar;                            
    using Index = Eigen::Index;                                               
    using VectorType = typename MomentsType::VectorType;                      
    using DifferenceVectorType = typename MomentsType::DifferenceVectorType;  
    using MatrixType = typename MomentsType::MatrixType;                      
    using DensityType = Mixture<Component_, Allocator_>;                      
    using LineageType = MixtureLineage<Scalar, Allocator_>;                   

    static constexpr int ValueDimensionAtCompileTime = ComponentType::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = ComponentType::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = ComponentType::DifferenceDimensionAtCompileTime;

    static constexpr int MaxDifferenceDimensionAtCompileTime = ComponentType::MaxDifferenceDimensionAtCompileTime;
 
    static constexpr bool IsParams = true;  

    struct Entry : internal::AllocatorTypeOf<ComponentType>
    {
        Entry(Scalar weight, ComponentType held)
            : logWeight{weight}
            , component{std::move(held)}
        {}

        template <typename A>
            requires(std::is_constructible_v<ComponentType, ComponentType const &, A const &>)
        Entry(Entry const & other, A const & allocator)
            : logWeight{other.logWeight}
            , component{other.component, allocator}
        {}

        template <typename A>
            requires(std::is_constructible_v<ComponentType, ComponentType &&, A const &>)
        Entry(Entry && other, A const & allocator)
            : logWeight{other.logWeight}
            , component{std::move(other.component), allocator}
        {}
 
        Entry(Entry const &) = default;  
        Entry(Entry &&) = default;       

        Entry &
        operator=(Entry const &) = default;

        Entry &
        operator=(Entry &&) = default;
 
        ~Entry() = default;
 
        Scalar logWeight;         
        ComponentType component;  
    };
 
    using const_iterator = typename internal::Buffer<Entry, Allocator_>::const_iterator;  

    template <bool ReturnRemovedWeight>
    using RemovalResult = std::conditional_t<ReturnRemovedWeight, std::expected<Scalar, WeightError>, void>;

    class MutableEntry
    {
    public:
 
        MutableEntry(MutableEntry const &) = delete;  

        MutableEntry &
        operator=(MutableEntry const &) = delete;

        [[nodiscard]] Scalar
        logWeight() const
        {
            return m_entry.logWeight;
        }

        [[nodiscard]] Index
        index() const
        {
            return m_index;
        }

        [[nodiscard]] ComponentType const &
        component() const
        {
            return m_entry.component;
        }

        template <typename C>
            requires std::same_as<std::remove_cvref_t<C>, ComponentType>
        void
        setComponent(C && component)
        {
            EMRYS_FILTER_ASSERT(
                component.space() == m_entry.component.space() && "the component must be in the mixture's space"
            );
            MixtureParams::assignComponent(m_entry, std::forward<C>(component));
        }

        template <typename F>
            requires std::invocable<F &, ComponentType &>
        decltype(auto)
        updateComponent(F && f)
        {
            return internal::keepingSpace(m_entry.component, f);
        }

        [[nodiscard]] MomentsType const &
        moments() const
            requires(isGaussian<ComponentType>)
        {
            return component().moments();
        }

        [[nodiscard]] VectorType const &
        mean() const
            requires(isGaussian<ComponentType>)
        {
            return component().mean();
        }

        [[nodiscard]] MatrixType const &
        covarianceLower() const
            requires(isGaussian<ComponentType>)
        {
            return component().covarianceLower();
        }

        template <MatrixExpression Mean>
            requires(isGaussian<ComponentType>)
        void
        setMean(Mean && mean)
        {
            m_entry.component.setMean(std::forward<Mean>(mean));
        }

        template <MatrixExpression Cov>
            requires(isGaussian<ComponentType>)
        void
        setCovariance(Cov && cov)
        {
            m_entry.component.setCovariance(std::forward<Cov>(cov));
        }

        void
        setMoments(MomentsType moments)
            requires(isGaussian<ComponentType>)
        {
            MixtureParams::assignComponent(m_entry, std::move(moments));
        }
 
    private:
 
        friend MixtureParams;  

        MutableEntry(Entry & entry, Index index)
            : m_entry{entry}
            , m_index{index}
        {}
 
        Entry & m_entry;  
        Index m_index;    
    };
 
    MixtureParams() = default;

    explicit MixtureParams(Allocator_ const & allocator)
        : m_entries(EntryAllocator(allocator))
        , m_weightScratch(ScalarAllocator(allocator))
        , m_lineage(allocator)
    {}

    explicit MixtureParams(Capacity capacity, Allocator_ const & allocator = {})
        requires FixedSizeSpace<SpaceType>
        : MixtureParams{allocator}
    {
        setCapacity(capacity.count);
    }

    explicit MixtureParams(Index valueDimension, Allocator_ const & allocator = {})
        requires internal::SizedByLength<SpaceType>
        : MixtureParams{allocator}
    {
        fitSpace(valueDimension);
    }

    explicit MixtureParams(SpaceType space, Allocator_ const & allocator = {})
        : MixtureParams{allocator}
    {
        m_space = std::move(space);
    }

    MixtureParams(SpaceType space, Capacity capacity, Allocator_ const & allocator = {})
        : MixtureParams{std::move(space), allocator}
    {
        setCapacity(capacity.count);
    }

    MixtureParams(SpaceType space, Index valueDimension, Allocator_ const & allocator = {})
        : MixtureParams{std::move(space), allocator}
    {
        fitSpace(valueDimension);
    }

    MixtureParams(MixtureParams const & other)
        : MixtureParams{other, other.get_allocator()}
    {}

    MixtureParams(MixtureParams const & other, Allocator_ const & allocator)
        : MixtureParams{allocator}
    {
        m_space = other.m_space;
        if (other.m_capacity)
        {
            setCapacity(*(other.m_capacity));
        }
        m_entries.assign(other.m_entries.begin(), other.m_entries.end());
    }

    MixtureParams(MixtureParams && other) noexcept
        : m_entries{std::move(other.m_entries)}
        , m_space{std::move(other.m_space)}
        , m_capacity{other.m_capacity}
        , m_weightScratch{std::move(other.m_weightScratch)}
        , m_lineage{std::move(other.m_lineage)}
    {
        other.m_stamp.touchShape();
    }

    MixtureParams(MixtureParams && other, Allocator_ const & allocator)
        : m_entries{std::move(other.m_entries), EntryAllocator(allocator)}
        , m_space{std::move(other.m_space)}
        , m_capacity{other.m_capacity}
        , m_weightScratch(ScalarAllocator(allocator))
        , m_lineage{std::move(other.m_lineage), allocator}
    {
        other.m_stamp.touchShape();
    }

    MixtureParams &
    operator=(MixtureParams const & other)
    {
        m_stamp.touchShape();
        if (this == &other)
        {
            return *this;
        }
        assertCapacity(other.size());
        m_entries.assign(other.m_entries.begin(), other.m_entries.end());
        m_space = other.m_space;
        m_lineage.reset();
        return *this;
    }

    MixtureParams &
    operator=(MixtureParams && other) noexcept
    {
        m_stamp.touchShape();
        other.m_stamp.touchShape();
        if (this == &other)
        {
            return *this;
        }
        if (std::allocator_traits<Allocator_>::propagate_on_container_move_assignment::value
            || get_allocator() == other.get_allocator())
        {
            if (get_allocator() != other.get_allocator())
            {
                // The entries take other's allocator, so the scratch does too.
                m_weightScratch = std::move(other.m_weightScratch);
            }
            m_entries = std::move(other.m_entries);
            m_capacity = other.m_capacity;
            m_lineage = std::move(other.m_lineage);
        }
        else
        {
            // Copied into this mixture's storage, which keeps its capacity.
            assertCapacity(other.size());
            m_entries.assign(
                std::make_move_iterator(other.m_entries.begin()),
                std::make_move_iterator(other.m_entries.end())
            );
            m_lineage = other.m_lineage;
        }
        m_space = std::move(other.m_space);
        return *this;
    }

    [[nodiscard]] allocator_type
    get_allocator() const
    {
        return allocator_type(m_entries.get_allocator());
    }

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

    [[nodiscard]] bool
    empty() const
    {
        return m_entries.empty();
    }

    [[nodiscard]] std::optional<Index>
    capacity() const
    {
        return m_capacity;
    }

    void
    setCapacity(Index capacity)
    {
        EMRYS_FILTER_ASSERT(capacity >= size() && "the capacity is smaller than the mixture");
        m_capacity = capacity;
        reserveScratch(capacity);
    }

    [[nodiscard]] Index
    valueDimension() const
    {
        return Emrys::valueDimension(m_space);
    }

    [[nodiscard]] Index
    differenceDimension() const
    {
        return Emrys::differenceDimension(m_space);
    }

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

    void
    resizeComponents(Index n)
        requires(isGaussian<ComponentType> && internal::SizedByLength<SpaceType>)
    {
        m_stamp.touchShape();
        internal::fitValueDimension(m_space, n);
        for (auto & entry : m_entries)
        {
            entry.component.resize(n);
        }
    }

    template <typename Path, typename MeanDerived, typename CovDerived, typename CrossDerived>
        requires(isGaussian<ComponentType>)
    void
    insert(
        Index element,
        Eigen::MatrixBase<MeanDerived> const & mean,
        Eigen::MatrixBase<CovDerived> const & cov,
        Eigen::MatrixBase<CrossDerived> const & cross
    )
    {
        m_stamp.touchShape();
        static_cast<void>(internal::growBlock<Path>(m_space, element, mean.size()));
        for (auto & entry : m_entries)
        {
            entry.component.template insert<Path>(element, mean, cov, cross);
        }
    }

    template <typename Path, typename MeanDerived, typename CovDerived>
        requires(isGaussian<ComponentType>)
    void
    insert(Index element, Eigen::MatrixBase<MeanDerived> const & mean, Eigen::MatrixBase<CovDerived> const & cov)
    {
        using Cross = internal::MatrixOf<
            Scalar,
            Eigen::Dynamic,
            Eigen::Dynamic,
            CovDerived::MaxRowsAtCompileTime,
            MomentsType::MatrixType::MaxColsAtCompileTime>;
        insert<Path>(element, mean, cov, Cross::Zero(cov.rows(), differenceDimension()));
    }

    template <typename Path>
        requires(isGaussian<ComponentType>)
    void
    erase(Index element, Index count = 1)
    {
        m_stamp.touchShape();
        static_cast<void>(internal::shrinkBlock<Path>(m_space, element, count));
        for (auto & entry : m_entries)
        {
            entry.component.template erase<Path>(element, count);
        }
    }

    void
    reserve(Index n)
    {
        m_stamp.touchValues();
        assertCapacity(n);
        m_entries.reserve(std::size_t(n));
        reserveScratch(n);
    }

    void
    clear()
    {
        m_stamp.touchShape();
        m_entries.clear();
        m_lineage.clearRows();
    }

    void
    swap(MixtureParams & other)
    {
        m_stamp.touchShape();
        other.m_stamp.touchShape();
        EMRYS_FILTER_ASSERT(get_allocator() == other.get_allocator() && "the mixtures must have equal allocators");
        assertCapacity(other.size());
        other.assertCapacity(size());
        m_entries.swap(other.m_entries);
        std::swap(m_space, other.m_space);
        m_lineage.swap(other.m_lineage);
    }

    [[nodiscard]] LineageType const &
    lineage() const
    {
        return m_lineage;
    }

    void
    startLineage()
    {
        m_stamp.touchValues();
        m_lineage.start(size(), m_capacity);
    }

    [[nodiscard]] const_iterator
    begin() const
    {
        return m_entries.cbegin();
    }

    [[nodiscard]] const_iterator
    end() const
    {
        return m_entries.cend();
    }

    [[nodiscard]] Entry const &
    operator[](Index i) const
    {
        assertIndex(i);
        return at(i);
    }

    Index
    add(Scalar logWeight, ComponentType component)
    {
        m_stamp.touchShape();
        EMRYS_FILTER_ASSERT(component.space() == m_space && "the component must be in the mixture's space");
        assertCapacity(size() + 1);
        m_entries.push_back(Entry{logWeight, std::move(component)});
        m_lineage.pushOrphan();
        return size() - 1;
    }

    Index
    add(Scalar logWeight, MomentsType moments)
        requires(isGaussian<ComponentType>)
    {
        EMRYS_FILTER_ASSERT(
            moments.valueDimension() == valueDimension() && moments.differenceDimension() == differenceDimension()
        );
        return add(logWeight, ComponentType{m_space, std::move(moments)});
    }

    template <MatrixExpression Mean, MatrixExpression Cov>
        requires(isGaussian<ComponentType>)
    Index
    add(Scalar logWeight, Mean && mean, Cov && cov)
    {
        return add(logWeight, MomentsType{std::forward<Mean>(mean), std::forward<Cov>(cov)});
    }

    void
    remove(Index i)
    {
        assertIndex(i);
        m_stamp.touchShape();
        removeAt(i);
    }

    void
    setComponent(Index i, ComponentType component)
    {
        EMRYS_FILTER_ASSERT(component.space() == m_space && "the component must be in the mixture's space");
        assignComponent(mutableAt(i), std::move(component));
    }

    void
    setMoments(Index i, MomentsType moments)
        requires(isGaussian<ComponentType>)
    {
        assignComponent(mutableAt(i), std::move(moments));
    }

    template <MatrixExpression Mean>
        requires(isGaussian<ComponentType>)
    void
    setMean(Index i, Mean && mean)
    {
        mutableAt(i).component.setMean(std::forward<Mean>(mean));
    }

    template <MatrixExpression Cov>
        requires(isGaussian<ComponentType>)
    void
    setCovariance(Index i, Cov && cov)
    {
        mutableAt(i).component.setCovariance(std::forward<Cov>(cov));
    }

    void
    setLogWeight(Index i, Scalar logWeight)
    {
        mutableAt(i).logWeight = logWeight;
    }

    void
    scaleWeight(Index i, Scalar logScale)
    {
        mutableAt(i).logWeight += logScale;
    }

    void
    scaleWeights(Scalar logScale)
    {
        m_stamp.touchValues();
        for (auto & entry : m_entries)
        {
            entry.logWeight += logScale;
        }
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    logTotalWeight() const
    {
        return internal::logTotalWeight<Scalar>(m_entries | std::views::transform(&Entry::logWeight));
    }

    [[nodiscard]] std::expected<void, WeightError>
    normalize()
    {
        m_stamp.touchValues();
        auto const total = logTotalWeight();
        if (!total)
        {
            return std::unexpected{total.error()};
        }
        scaleWeights(-*total);
        return {};
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    ess() const
    {
        return logTotalWeight().transform([this](Scalar total) {
            return essGiven(total);
        });
    }

    [[nodiscard]] Index
    mostProbable() const
    {
        EMRYS_FILTER_ASSERT(!empty() && "the most probable component of an empty mixture");
        return Index(std::ranges::max_element(m_entries, std::ranges::less{}, &Entry::logWeight) - m_entries.begin());
    }

    template <std::ranges::input_range R>
    void
    assign(R && components)
    {
        clear();
        if constexpr (std::ranges::sized_range<R>)
        {
            reserve(Index(std::ranges::size(components)));
        }
        for (auto && [logWeight, component] : components)
        {
            add(logWeight, component);
        }
    }

    void
    append(MixtureParams const & other)
    {
        m_stamp.touchShape();
        EMRYS_FILTER_ASSERT(other.m_space == m_space && "cannot append a mixture in a different space");
        std::size_t const n = other.m_entries.size();
        assertCapacity(size() + Index(n));
        m_entries.reserve(m_entries.size() + n);
        for (std::size_t i = 0; i < n; ++i)
        {
            m_entries.push_back(other.m_entries[i]);
            m_lineage.pushOrphan();
        }
    }

    template <typename F>
    void
    forEach(F && f)
    {
        forEach(execution::sequential, std::forward<F>(f));
    }

    template <ExecutionPolicy P, typename F>
    void
    forEach(P && policy, F && f)
    {
        static_assert(
            std::is_void_v<std::invoke_result_t<F &, MutableEntry &>>,
            "forEach() takes a callback that returns nothing; one that returns a log scale factor goes to reweight()"
        );
        visitEntries(std::forward<P>(policy), [&](MutableEntry & handle, Entry &) {
            std::invoke(f, handle);
        });
    }

    template <typename F>
    void
    reweight(F && f)
    {
        reweight(execution::sequential, std::forward<F>(f));
    }

    template <ExecutionPolicy P, typename F>
    void
    reweight(P && policy, F && f)
    {
        static_assert(
            std::floating_point<std::remove_cvref_t<std::invoke_result_t<F &, MutableEntry &>>>,
            "reweight() takes a callback that returns a log scale factor, a floating-point number"
        );
        visitEntries(std::forward<P>(policy), [&](MutableEntry & handle, Entry & entry) {
            entry.logWeight += Scalar(std::invoke(f, handle));
        });
    }

    Index
    merge(Index keep, Index absorb)
    {
        return merge(keep, absorb, merged(keep, absorb));
    }

    [[nodiscard]] ComponentType
    merged(Index first, Index second) const
    {
        assertIndex(first);
        assertIndex(second);
        Index const pair[2] = {first, second};
        EMRYS_FILTER_ASSERT(pair[0] != pair[1] && "cannot merge a component with itself");
        return ComponentTraits<ComponentType>::collapse(
            Index(2),
            [&](Index i) -> ComponentType const & {
                return at(pair[i]).component;
            },
            [&](Index i) {
                return at(pair[i]).logWeight;
            }
        );
    }

    Index
    merge(Index keep, Index absorb, ComponentType component)
    {
        return mergeInto(keep, absorb, std::move(component));
    }

    Index
    merge(Index keep, Index absorb, MomentsType moments)
        requires(isGaussian<ComponentType>)
    {
        return mergeInto(keep, absorb, std::move(moments));
    }

    template <std::ranges::forward_range R>
    void
    split(Index i, R && children)
    {
        assertIndex(i);
        m_stamp.touchShape();
        EMRYS_FILTER_ASSERT(!std::ranges::empty(children));
        auto relativeTotal = LogSumExp<Scalar>{};
        for ([[maybe_unused]] auto && [logRelativeWeight, child] : children)
        {
            relativeTotal.add(logRelativeWeight);
        }
        Scalar const logTotal = relativeTotal.value();
        EMRYS_FILTER_ASSERT(isFinite(logTotal) && "relative weights must have a finite, nonzero sum");
        Scalar const logBase = at(i).logWeight - logTotal;
        bool first = true;
        for (auto && [logRelativeWeight, child] : children)
        {
            if (first)
            {
                at(i).logWeight = logBase + logRelativeWeight;
                assignComponent(at(i), child);
                first = false;
            }
            else
            {
                add(logBase + logRelativeWeight, child);
                m_lineage.truncate(size() - 1);
                m_lineage.pushCopyOf(i);
            }
        }
    }

    template <bool ReturnRemovedWeight = false, typename Pred>
        requires(!ReturnRemovedWeight)
    void
    removeIf(Pred pred)
    {
        removeEntriesIf<false>(std::move(pred));
    }

    template <bool ReturnRemovedWeight, typename Pred>
        requires ReturnRemovedWeight
    [[nodiscard]] std::expected<Scalar, WeightError>
    removeIf(Pred pred)
    {
        return removeEntriesIf<true>(std::move(pred));
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneBelow(Scalar logThreshold)
    {
        return pruneWith(internal::PruneBelowRule<Scalar>{logThreshold});
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneTail(Scalar logFraction)
    {
        return pruneWith(internal::PruneTailRule<Scalar>{logFraction});
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneToSize(Index maxSize)
    {
        return pruneWith(internal::PruneToSizeRule{maxSize});
    }

    template <typename Path>
    [[nodiscard]] auto
    marginal() const &
    {
        return internal::tagMarginal<MixtureMarginal, Path>(*this);
    }

    template <typename Path>
    void
    marginal() const && = delete;

    template <typename Path>
    [[nodiscard]] auto
    marginal(Index element) const &
    {
        return internal::elementMarginal<MixtureMarginal, Path>(*this, element);
    }

    template <typename Path>
    void
    marginal(Index element) const && = delete;

    template <typename GivenPath, typename Derived>
        requires(isGaussian<ComponentType>)
    [[nodiscard]] auto
    conditional(Eigen::MatrixBase<Derived> const & values) const
    {
        return internal::tagConditionalMixture<GivenPath>(*this, values);
    }

    [[nodiscard]] std::expected<void, MixtureError>
    factorize() const
    {
        return factorizeComponents(*this, internal::NoRepair{});
    }

    template <typename R>
    [[nodiscard]] std::expected<void, MixtureError>
    factorize(R const & repair)
    {
        m_stamp.touchValues();
        return factorizeComponents(*this, repair);
    }

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

    void
    toDensity() const && = delete;

    template <typename R>
    [[nodiscard]] std::expected<DensityType, MixtureError>
    toDensity(R const & repair) &
    {
        if (auto result = factorize(repair); !result)
        {
            return std::unexpected{result.error()};
        }
        return DensityType{*this};
    }

    template <typename R>
    void
    toDensity(R const &) && = delete;
 
private:
 
    friend DensityType;  

    [[nodiscard]] Scalar
    essGiven(Scalar total) const
    {
        Scalar sumOfSquares = 0;
        for (auto const & entry : m_entries)
        {
            sumOfSquares += Eigen::numext::abs2(Eigen::numext::exp(entry.logWeight - total));
        }
        return Scalar(1) / sumOfSquares;
    }

    template <typename>
    friend class internal::ValuesCheck;

    template <typename>
    friend class internal::ShapeCheck;
 
    friend struct internal::MixtureLineageAccess;  

    template <typename Self, typename R>
    [[nodiscard]] static std::expected<void, MixtureError>
    factorizeComponents(Self & self, R const & repair)
    {
        constexpr bool Repairs = !std::is_same_v<R, internal::NoRepair>;
        auto const factorizeOne = [&](auto & component) {
            if constexpr (Repairs)
            {
                return component.factorize(repair);
            }
            else
            {
                return component.factorize();
            }
        };
        auto const total = self.logTotalWeight();
        if (!total)
        {
            return std::unexpected{MixtureError{total.error()}};
        }
        if constexpr (Repairs)
        {
            // Every repair is tried before any is kept, on a copy of a component that needs one,
            // so that a component that can't be repaired leaves the mixture as it was. Repairs are
            // rare, so the copies are too.
            for (Index i = 0; i < self.size(); ++i)
            {
                auto const & component = std::as_const(self.at(i).component);
                if (!component.factorize())
                {
                    auto trial = ComponentType{component};
                    if (auto result = factorizeOne(trial); !result)
                    {
                        return std::unexpected{MixtureError{ComponentError{i, result.error()}}};
                    }
                }
            }
        }
        for (Index i = 0; i < self.size(); ++i)
        {
            if (auto result = factorizeOne(self.at(i).component); !result)
            {
                return std::unexpected{MixtureError{ComponentError{i, result.error()}}};
            }
        }
        // Unchanged weights write nothing, so factorizing parameters that are already factorized
        // only reads, as toDensity() from several threads needs (decisions.md [execution-policies]).
        if (self.m_logTotalWeight != *total)
        {
            self.m_logTotalWeight = *total;
        }
        return {};
    }
 
    using EntryAllocator = internal::ReBound<Allocator_, Entry>;    
    using ScalarAllocator = internal::ReBound<Allocator_, Scalar>;  

    void
    fitSpace(Index n)
    {
        internal::fitValueDimension(m_space, n);
    }

    static void
    assignComponent(Entry & entry, MomentsType const & moments)
        requires(isGaussian<ComponentType>)
    {
        entry.component.setMean(moments.mean());
        entry.component.setCovariance(moments.covarianceLower());
    }

    static void
    assignComponent(Entry & entry, MomentsType && moments)
        requires(isGaussian<ComponentType>)
    {
        entry.component.setMean(std::move(moments.mean()));
        entry.component.setCovariance(std::move(moments.covarianceLower()));
    }

    template <bool ReturnRemovedWeight, typename Pred>
    RemovalResult<ReturnRemovedWeight>
    removeEntriesIf(Pred pred)
    {
        m_stamp.touchShape();
        auto const guard = internal::CallbackGuard{m_stamp};
        auto const removeAtIndex = [&](Index i) {
            bool remove = false;
            if constexpr (std::is_invocable_v<Pred &, Index, Entry const &>)
            {
                remove = static_cast<bool>(std::invoke(pred, i, std::as_const(at(i))));
            }
            else
            {
                remove = static_cast<bool>(std::invoke(pred, std::as_const(at(i))));
            }
            guard.check();
            return remove;
        };
        if constexpr (ReturnRemovedWeight)
        {
            auto const total = logTotalWeight();
            if (!total)
            {
                return std::unexpected{total.error()};
            }
            return compact(removeAtIndex) - *total;
        }
        else
        {
            compact(removeAtIndex);
        }
    }

    template <ExecutionPolicy P, typename Visit>
    void
    visitEntries(P && policy, Visit const & visit)
    {
        m_stamp.touchValues();
        auto const guard = internal::CallbackGuard{m_stamp};
        internal::forEachIndex(std::forward<P>(policy), size(), [&](Index i) {
            auto & entry = m_entries[std::size_t(i)];
            auto handle = MutableEntry{entry, i};
            visit(handle, entry);
            guard.check();
        });
    }

    static void
    assignComponent(Entry & entry, ComponentType const & component)
    {
        entry.component = component;
    }

    static void
    assignComponent(Entry & entry, ComponentType && component)
    {
        entry.component = std::move(component);
    }

    template <typename Replacement>
    Index
    mergeInto(Index k, Index a, Replacement && replacement)
    {
        assertIndex(k);
        assertIndex(a);
        m_stamp.touchShape();
        EMRYS_FILTER_ASSERT(k != a && "cannot merge a component with itself");
        auto weight = LogSumExp<Scalar>{};
        weight.add(at(k).logWeight);
        weight.add(at(a).logWeight);
        m_lineage.merge(k, a, at(k).logWeight, at(a).logWeight);
        at(k).logWeight = weight.value();
        assignComponent(at(k), std::forward<Replacement>(replacement));
        bool const keepWasLast = k == size() - 1;
        removeAt(a);
        return keepWasLast ? a : k;
    }

    void
    assertCapacity([[maybe_unused]] Index count) const
    {
        EMRYS_FILTER_ASSERT(
            (!m_capacity || count <= *m_capacity) && "the mixture would hold more components than its capacity"
        );
    }

    void
    reserveScratch(Index count)
    {
        auto const n = std::size_t(count);
        m_entries.reserve(n);
        m_weightScratch.reserve(n);
    }

    void
    assertIndex([[maybe_unused]] Index i) const
    {
        EMRYS_FILTER_ASSERT(i >= 0 && i < size() && "component index out of range");
    }

    [[nodiscard]] Entry &
    at(Index i)
    {
        return m_entries[std::size_t(i)];
    }

    [[nodiscard]] Entry const &
    at(Index i) const
    {
        return m_entries[std::size_t(i)];
    }

    [[nodiscard]] Entry &
    mutableAt(Index i)
    {
        assertIndex(i);
        m_stamp.touchValues();
        return at(i);
    }

    void
    removeAt(Index i)
    {
        if (i != size() - 1)
        {
            at(i) = std::move(m_entries.back());
        }
        m_entries.pop_back();
        m_lineage.removeAt(i);
    }

    template <typename Rule>
    [[nodiscard]] std::expected<Scalar, WeightError>
    pruneWith(Rule const & rule)
    {
        auto const total = logTotalWeight();
        if (!total)
        {
            return std::unexpected{total.error()};
        }
        m_stamp.touchShape();
        auto cut = rule(
            size(),
            [&](Index i) {
                return at(i).logWeight;
            },
            *total,
            m_weightScratch
        );
        Scalar const removed = compact([&](Index i) {
            return cut.removes(at(i).logWeight);
        });
        // The kept components include one of positive weight, so their total is finite.
        scaleWeights(-*(logTotalWeight()));
        return removed - *total;
    }

    template <typename RemoveAtIndex>
    Scalar
    compact(RemoveAtIndex removeAtIndex)
    {
        auto removed = LogSumExp<Scalar>{};
        Index out = 0;
        for (Index i = 0; i < size(); ++i)
        {
            if (removeAtIndex(i))
            {
                removed.add(at(i).logWeight);
                continue;
            }
            if (out != i)
            {
                at(out) = std::move(at(i));
                m_lineage.moveRow(i, out);
            }
            ++out;
        }
        m_entries.erase(m_entries.begin() + out, m_entries.end());
        m_lineage.truncate(out);
        return removed.value();
    }
 
    internal::Buffer<Entry, Allocator_> m_entries;     
    EMRYS_FILTER_NO_UNIQUE_ADDRESS SpaceType m_space;  
    std::optional<Index> m_capacity;                   
 
    // Scratch, kept between calls: the weights pruneTail() and pruneToSize() order.
    internal::Buffer<Scalar, Allocator_> m_weightScratch;  
 
    LineageType m_lineage;  
 
    // Filled by the const factorize() for use by views.
    mutable Scalar m_logTotalWeight{};  
 
    internal::Stamp m_stamp;  
};

template <Space Space_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
using GaussianMixtureParams = MixtureParams<Gaussian<Space_, Scalar_>, Allocator_>;

template <internal::ComponentOrInfoForm Component_, typename Allocator_ = DefaultAllocator>
class Mixture
{
public:
 
    using ParamsType = MixtureParams<Component_, Allocator_>;                
    using ComponentType = Params<Component_>;                                
    using SpaceType = typename ParamsType::SpaceType;                        
    using allocator_type = Allocator_;                                       
    using Entry = typename ParamsType::Entry;                                
    using const_iterator = typename ParamsType::const_iterator;              
    using MomentsType = typename ParamsType::MomentsType;                    
    using Scalar = typename ParamsType::Scalar;                              
    using Index = typename ParamsType::Index;                                
    using VectorType = typename ParamsType::VectorType;                      
    using DifferenceVectorType = typename ParamsType::DifferenceVectorType;  
    using MatrixType = typename ParamsType::MatrixType;                      

    static constexpr int ValueDimensionAtCompileTime = ParamsType::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = ParamsType::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = ParamsType::DifferenceDimensionAtCompileTime;

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

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

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

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

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

    [[nodiscard]] allocator_type
    get_allocator() const
    {
        return source().get_allocator();
    }

    [[nodiscard]] const_iterator
    begin() const
    {
        return source().begin();
    }

    [[nodiscard]] const_iterator
    end() const
    {
        return source().end();
    }

    [[nodiscard]] typename ParamsType::Entry const &
    operator[](Index i) const
    {
        return source()[i];
    }

    [[nodiscard]] Scalar
    logTotalWeight() const
    {
        return source().m_logTotalWeight;
    }

    [[nodiscard]] Scalar
    ess() const
    {
        return source().essGiven(logTotalWeight());
    }

    [[nodiscard]] Index
    mostProbable() const
    {
        return source().mostProbable();
    }

    [[nodiscard]] ComponentType
    collapsed() const
    {
        return ComponentTraits<ComponentType>::collapse(
            size(),
            [&](Index i) -> ComponentType const & {
                return (*this)[i].component;
            },
            [&](Index i) {
                return (*this)[i].logWeight;
            }
        );
    }

    [[nodiscard]] MomentsType
    moments() const
    {
        return ComponentTraits<ComponentType>::moments(collapsed());
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        static_assert(Derived::IsVectorAtCompileTime, "logDensity() expects a vector");
        EMRYS_FILTER_ASSERT(x.size() == valueDimension());
        auto sum = LogSumExp<Scalar>{};
        for (auto const & entry : *this)
        {
            if (entry.logWeight == -Eigen::NumTraits<Scalar>::infinity())
            {
                continue;
            }
            sum.add(entry.logWeight + ComponentTraits<ComponentType>::logDensity(entry.component, x));
        }
        return sum.value() - logTotalWeight();
    }

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

    template <typename Path>
    [[nodiscard]] auto
    marginal() const
    {
        return internal::tagMarginal<MixtureMarginal, Path>(*this);
    }

    template <typename Path>
    [[nodiscard]] auto
    marginal(Index element) const
    {
        return internal::elementMarginal<MixtureMarginal, Path>(*this, element);
    }

    template <typename GivenPath, typename Derived>
        requires(isGaussian<ComponentType>)
    [[nodiscard]] auto
    conditional(Eigen::MatrixBase<Derived> const & values) const
    {
        return internal::tagConditionalMixture<GivenPath>(*this, values);
    }
 
private:
 
    friend ParamsType;  

    explicit Mixture(ParamsType const & source)
        : m_check{source}
    {}

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

template <Space Space_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
using GaussianMixture = Mixture<Gaussian<Space_, Scalar_>, Allocator_>;
 
namespace internal
{

template <typename Source>
struct ElementsOf;

template <Component Component_, typename Allocator_>
struct ElementsOf<MixtureParams<Component_, Allocator_>>
{
    template <typename Source>
    [[nodiscard]] static auto
    logWeight(Source const & source, Eigen::Index i)
    {
        return source[i].logWeight;
    }

    template <typename Source>
    [[nodiscard]] static auto const &
    component(Source const & source, Eigen::Index i)
    {
        return source[i].component;
    }

    template <typename ComponentParams>
    using CopyType = MixtureParams<Density<ComponentParams>, Allocator_>;

    template <typename Source, typename Marginal>
    [[nodiscard]] static auto
    copy(Source const &, Marginal const & marginal, Allocator_ const & allocator)
    {
        auto result = CopyType<typename Marginal::ComponentMarginalType::ParamsType>{marginal.space(), allocator};
        result.reserve(marginal.size());
        for (auto const & element : marginal)
        {
            result.add(element.logWeight, element.component.copy());
        }
        return result;
    }
};

template <Component Component_, typename Allocator_>
struct ElementsOf<Mixture<Component_, Allocator_>> : ElementsOf<MixtureParams<Component_, Allocator_>>
{};
 
}  // namespace internal

template <typename Source, typename Selection, typename SubSpace>
class MixtureMarginal
{
public:
 
    using Scalar = typename Source::Scalar;  
    using Index = Eigen::Index;              

    static constexpr int ValueDimensionAtCompileTime = valueDimensionAtCompileTime<SubSpace>;

    static constexpr int MaxValueDimensionAtCompileTime = maxValueDimensionAtCompileTime<SubSpace>;

    static constexpr int DifferenceDimensionAtCompileTime = valueDimensionAtCompileTime<DifferenceSpace<SubSpace>>;

    static constexpr int MaxDifferenceDimensionAtCompileTime =
        maxValueDimensionAtCompileTime<DifferenceSpace<SubSpace>>;
 
    using SpaceType = SubSpace;                              
    using allocator_type = typename Source::allocator_type;  
    using ComponentType = typename Source::ComponentType;    

    using ResolvedType = typename internal::SelectionHolder<Source, Selection, SubSpace>::ResolvedType;

    using IndexVectorType = decltype(ResolvedType::indices);

    using DifferenceIndexVectorType = decltype(ResolvedType::differenceIndices);

    using ComponentMarginalType = typename ComponentTraits<
        ComponentType>::template MarginalType<internal::SharedSelection<ResolvedType>, SubSpace>;

    struct Element
    {
        Scalar logWeight;                 
        ComponentMarginalType component;  
    };
 
    using value_type = Element;                                         
    using const_iterator = internal::ElementIterator<MixtureMarginal>;  
    using iterator = const_iterator;                                    

    using ParamsType =
        typename internal::ElementsOf<Source>::template CopyType<typename ComponentMarginalType::ParamsType>;

    MixtureMarginal(Source const & source, Selection const & selection, SubSpace space)
        requires(!internal::isSharedSelection<Selection>)
        : m_holder{source}
        , m_selection{source, selection, std::move(space)}
    {}

    MixtureMarginal(Source const & source, Selection shared)
        requires(internal::isSharedSelection<Selection>)
        : m_holder{source}
        , m_selection{shared}
    {}

    [[nodiscard]] SubSpace const &
    space() const
    {
        return m_selection.get().space;
    }

    [[nodiscard]] allocator_type
    get_allocator() const
    {
        return source().get_allocator();
    }

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

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

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

    [[nodiscard]] IndexVectorType const &
    indices() const
    {
        return m_selection.get().indices;
    }

    [[nodiscard]] DifferenceIndexVectorType const &
    differenceIndices() const
    {
        return m_selection.get().differenceIndices;
    }

    [[nodiscard]] Element
    operator[](Index i) const
    {
        EMRYS_FILTER_ASSERT(i >= 0 && i < size() && "component index out of range");
        Source const & from = source();
        return Element{
            Elements::logWeight(from, i),
            ComponentMarginalType{
                Elements::component(from, i),
                internal::SharedSelection<ResolvedType>{&m_selection.get()}
            }
        };
    }

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

    [[nodiscard]] const_iterator
    end() const
    {
        return const_iterator{*this, size()};
    }

    [[nodiscard]] ParamsType
    copy() const
    {
        return copy(get_allocator());
    }

    [[nodiscard]] ParamsType
    copy(allocator_type const & allocator) const
    {
        return Elements::copy(source(), *this, allocator);
    }
 
private:
 
    using Elements = internal::ElementsOf<Source>;  

    [[nodiscard]] Source const &
    source() const
    {
        return m_holder.get();
    }
 
    internal::HolderFor<Source> m_holder;  
    internal::SelectionHolder<Source, Selection, SubSpace> m_selection;  
};

template <
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SampleSet<Space, Scalar>
drawSamples(
    Mixture<Gaussian<Space, Scalar>, Allocator> const & mixture,
    Eigen::Index n,
    G & rng,
    R const & resampler = {}
)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using SampleSetType = SampleSet<Space, Scalar>;
    auto weights = internal::bufferOf<Scalar>(mixture.get_allocator());
    weights.reserve(std::size_t(mixture.size()));
    for (auto const & entry : mixture)
    {
        weights.push_back(Eigen::numext::exp(entry.logWeight - mixture.logTotalWeight()));
    }
    auto ancestors = internal::bufferOf<Eigen::Index>(mixture.get_allocator());
    ancestors.resize(std::size_t(n));
    resampler(std::span<Scalar const>{weights}, std::span<Eigen::Index>{ancestors}, rng);
 
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const noise = internal::standardNormal<Normals>(mixture.differenceDimension(), n, rng);
    auto samples = SampleSetType{mixture.valueDimension(), n};
    for (Eigen::Index j = 0; j < n; ++j)
    {
        auto const gaussian = mixture[ancestors[std::size_t(j)]].component.toDensity();
        samples.col(j) = Emrys::plus(
            mixture.space(),
            gaussian->mean(),
            internal::multiply(internal::asTriangular<Eigen::Lower>(gaussian->llt().matrixLLT()), noise.col(j))
        );
    }
    return samples;
}

template <typename Scalar, typename Space, typename Allocator, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<Space, Scalar>::VectorType
drawSample(Mixture<Gaussian<Space, Scalar>, Allocator> const & mixture, G & rng)
{
    Eigen::Index const picked = internal::drawIndex<Scalar>(
        mixture.size(),
        [&](Eigen::Index k) {
            return Eigen::numext::exp(Scalar(mixture[k].logWeight - mixture.logTotalWeight()));
        },
        rng
    );
    auto const gaussian = mixture[picked].component.toDensity();
    return drawSample(*gaussian, rng);
}

template <
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(
    Mixture<Gaussian<Space, Scalar>, Allocator> const & mixture,
    Eigen::Index n,
    G & rng,
    R const & resampler = {}
)
{
    return drawSamplesWithLogDensity(execution::sequential, mixture, n, rng, resampler);
}

template <
    ExecutionPolicy P,
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(
    P && policy,
    Mixture<Gaussian<Space, Scalar>, Allocator> const & mixture,
    Eigen::Index n,
    G & rng,
    R const & resampler = {}
)
{
    auto result = SamplesWithLogDensity<Space, Scalar>{drawSamples(mixture, n, rng, resampler), {}};
    result.logDensity.resize(n);
    internal::forEachIndex(std::forward<P>(policy), n, [&](Eigen::Index j) {
        result.logDensity(j) = mixture.logDensity(result.samples.col(j));
    });
    return result;
}

template <typename Scalar, typename Space, typename Allocator, std::uniform_random_bit_generator G>
[[nodiscard]] SampleWithLogDensity<Space, Scalar>
drawSampleWithLogDensity(Mixture<Gaussian<Space, Scalar>, Allocator> const & mixture, G & rng)
{
    auto result = SampleWithLogDensity<Space, Scalar>{drawSample(mixture, rng), 0};
    result.logDensity = mixture.logDensity(result.sample);
    return result;
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Component_, typename Allocator_>
void
describe(std::ostream & stream, MixtureParams<Component_, Allocator_> const & mixture, Indent indent)
{
    using Scalar = typename MixtureParams<Component_, Allocator_>::Scalar;
    stream << indent << "mixture of ";
    printCount(stream, mixture.size(), "component");
    printDimension(stream, mixture.valueDimension(), mixture.differenceDimension());
    auto const logTotal = mixture.logTotalWeight();
    if (logTotal)
    {
        Params<Component_> const collapsed = ComponentTraits<Params<Component_>>::collapse(
            mixture.size(),
            [&](Eigen::Index i) -> Params<Component_> const & {
                return mixture[i].component;
            },
            [&](Eigen::Index i) {
                return mixture[i].logWeight;
            }
        );
        printMoments(stream, indent.deeper(), ComponentTraits<Params<Component_>>::moments(collapsed));
    }
    else
    {
        stream << '\n';
        describe(stream, logTotal.error(), indent.deeper());
    }
    Eigen::Index i = 0;
    for (auto const & entry : mixture)
    {
        stream << '\n' << indent.deeper() << "component " << i << ", ";
        if (logTotal)
        {
            stream << "weight " << Eigen::numext::exp(Scalar(entry.logWeight - *logTotal));
        }
        else
        {
            stream << "log-weight " << entry.logWeight;
        }
        stream << ":\n";
        describe(stream, entry.component, indent.deeper().deeper());
        ++i;
    }
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO