GridGaussian.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <array>
#include <expected>
#include <cmath>
#include <cstddef>
#include <optional>
#include <expected>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Print.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Categorical.h>
#include <EmrysFilter/Distributions/Conditionals.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Hybrid.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/Params.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Discrete.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{
 
template <typename Grid_, typename Model_, typename Covariance_, typename Method_, typename Upkeep_>
class RbPointMassFilter;
 
template <
    typename Grid_,
    typename Covariance_ = SharedCovariance,
    typename Component_ = Gaussian<typename Grid_::RestSpace, typename Grid_::Scalar>>
class GridGaussian;

template <typename Block, int I = 0>
struct GridAxis : Finite<Eigen::Dynamic>
{
    using Finite::Finite;
};
 
namespace internal
{

template <typename Block, std::size_t... I>
std::tuple<GridAxis<Block, int(I)>...> leafGridAxes(std::index_sequence<I...>);

template <typename Block>
struct GridAxesOf
{
    using type = decltype(leafGridAxes<Block>(std::make_index_sequence<std::size_t(Block::FixedValueDimension)>{}));
};

template <typename Block>
    requires IsComposite<Block>
struct GridAxesOf<Block>
{
    template <typename... Blocks>
    static auto
    of(std::tuple<Blocks...> *) -> decltype(std::tuple_cat(std::declval<typename GridAxesOf<Blocks>::type>()...));

    using type = decltype(of(static_cast<CompositeBlocks<Block> *>(nullptr)));
};

template <typename... Axes>
DiscreteProduct<Axes...>
gridProduct(std::tuple<Axes...> *);

template <typename Axis, typename... Axes>
inline constexpr int axisOccurrences = (0 + ... + int(std::is_same_v<Axis, Axes>));

template <typename... Axes>
consteval bool
distinctAxes(std::tuple<Axes...> *)
{
    return ((axisOccurrences<Axes, Axes...> == 1) && ...);
}

template <typename Path>
using GridAxes = typename GridAxesOf<PathTag<Path>>::type;
 
}  // namespace internal

template <typename Path>
struct GridCells : decltype(internal::gridProduct(static_cast<internal::GridAxes<Path> *>(nullptr)))
{
    static_assert(
        internal::distinctAxes(static_cast<internal::GridAxes<Path> *>(nullptr)),
        "two coordinates of the block lie in blocks with the same tag, so their grid axes would share a name"
    );

    using Product = decltype(internal::gridProduct(static_cast<internal::GridAxes<Path> *>(nullptr)));
 
    using Product::Product;
};

template <typename Path>
struct PointCells : Finite<Eigen::Dynamic>
{
    using Finite::Finite;
};
 
namespace internal
{

template <typename Cells>
struct IsGridCells : std::false_type
{};

template <typename Path>
struct IsGridCells<GridCells<Path>> : std::true_type
{};
 
}  // namespace internal

template <
    Space Space_,
    typename Path_,
    DiscreteSpace Cells_,
    typename CellCovariance_ = SharedCovariance,
    RealScalar Scalar_ = double,
    typename Allocator_ = DefaultAllocator>
class PointGrid
{
    static_assert(
        internal::CovarianceForm<CellCovariance_>,
        "the cell covariance is SharedCovariance or PerPointCovariance"
    );
 
    static constexpr bool Shared = std::is_same_v<CellCovariance_, SharedCovariance>;  
 
public:

    using Scalar = Scalar_;

    using SpaceType = Space_;

    using CellsType = Cells_;

    using PathType = Path_;

    using CellCovarianceType = CellCovariance_;

    using allocator_type = Allocator_;

    using PinnedSpace = typename internal::BlockSplit<Space_, Path_>::PinnedSpace;

    using RestSpace = typename internal::BlockSplit<Space_, Path_>::RestSpace;

    using CentersType = SampleSet<PinnedSpace, Scalar_>;

    using PinnedVector = typename Moments<PinnedSpace, Scalar_>::VectorType;

    using PinnedMatrix = typename Moments<PinnedSpace, Scalar_>::MatrixType;

    static constexpr bool Regular = internal::IsGridCells<Cells_>::value;
 
    static_assert(
        !Regular || std::is_same_v<Cells_, GridCells<Path_>>,
        "a regular grid's cells are the GridCells of its own block"
    );

    template <typename CentersDerived>
    PointGrid(
        Space_ space,
        Cells_ cells,
        Eigen::MatrixBase<CentersDerived> const & centers,
        std::optional<PinnedMatrix> cellCovariance = {},
        Allocator_ const & allocator = {}
    )
        requires(Shared && !Regular)
        : PointGrid{Parts{}, std::move(space), std::move(cells), listOf(cellCovariance), allocator}
    {
        setCenters(centers);
    }

    template <typename CentersDerived>
    PointGrid(
        Space_ space,
        Cells_ cells,
        Eigen::MatrixBase<CentersDerived> const & centers,
        std::span<PinnedMatrix const> cellCovariances,
        Allocator_ const & allocator = {}
    )
        requires(!Shared && !Regular)
        : PointGrid{Parts{}, std::move(space), std::move(cells), cellCovariances, allocator}
    {
        setCenters(centers);
    }

    PointGrid(
        Space_ space,
        Cells_ cells,
        PinnedVector first,
        PinnedVector spacing,
        std::optional<PinnedMatrix> cellCovariance = {},
        Allocator_ const & allocator = {}
    )
        requires(Shared && Regular)
        : PointGrid{Parts{}, std::move(space), std::move(cells), listOf(cellCovariance), allocator}
    {
        setGeometry(std::move(first), std::move(spacing));
    }

    PointGrid(
        Space_ space,
        Cells_ cells,
        PinnedVector first,
        PinnedVector spacing,
        std::span<PinnedMatrix const> cellCovariances,
        Allocator_ const & allocator = {}
    )
        requires(!Shared && Regular)
        : PointGrid{Parts{}, std::move(space), std::move(cells), cellCovariances, allocator}
    {
        setGeometry(std::move(first), std::move(spacing));
    }

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

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

    [[nodiscard]] Cells_ const &
    cells() const
    {
        return m_cells;
    }

    [[nodiscard]] PinnedSpace const &
    pinnedSpace() const
    {
        return m_split.pinnedSpace();
    }

    [[nodiscard]] RestSpace const &
    restSpace() const
    {
        return m_split.restSpace();
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return m_cells.cardinality();
    }

    [[nodiscard]] PinnedVector
    center(Eigen::Index cell) const
    {
        EMRYS_FILTER_ASSERT(cell >= 0 && cell < size() && "the cell is out of range");
        if constexpr (Regular)
        {
            auto const coordinates = m_cells.coordinates(cell);
            PinnedVector result = m_first;
            for (std::size_t i = 0; i < coordinates.size(); ++i)
            {
                Eigen::Index const row = Eigen::Index(i);
                result(row) += Scalar(coordinates[i]) * m_spacing(row);
            }
            normalizeValues(pinnedSpace(), result);
            return result;
        }
        else
        {
            return PinnedVector{centers().col(cell)};
        }
    }

    [[nodiscard]] ConstSampleSetMap<PinnedSpace, Scalar_>
    centers() const
        requires(!Regular)
    {
        return {m_centers.data(), pinnedRows().size, size()};
    }

    [[nodiscard]] PinnedVector const &
    first() const
        requires Regular
    {
        return m_first;
    }

    [[nodiscard]] PinnedVector const &
    spacing() const
        requires Regular
    {
        return m_spacing;
    }

    [[nodiscard]] bool
    hasCellCovariance() const
    {
        return !m_cellCovariances.empty();
    }

    [[nodiscard]] PinnedMatrix const &
    cellCovariance(Eigen::Index cell) const
    {
        EMRYS_FILTER_ASSERT(hasCellCovariance() && "the cells are point masses");
        EMRYS_FILTER_ASSERT(cell >= 0 && cell < size() && "the cell is out of range");
        return m_cellCovariances[Shared ? 0 : std::size_t(cell)];
    }

    [[nodiscard]] BlockRows
    pinnedRows() const
    {
        return m_split.pinnedRows();
    }

    [[nodiscard]] auto const &
    restIndices() const
    {
        return m_split.restIndices();
    }

    [[nodiscard]] BlockRows
    pinnedDifferenceRows() const
    {
        return m_split.pinnedDifferenceRows();
    }

    [[nodiscard]] auto const &
    restDifferenceIndices() const
    {
        return m_split.restDifferenceIndices();
    }

    [[nodiscard]] PointGrid
    subdivided(Eigen::Index factor) const
        requires(Shared && Regular)
    {
        EMRYS_FILTER_ASSERT(factor >= 1 && "subdivide each cell into at least one");
        Cells_ cells = [&]<std::size_t... I>(std::index_sequence<I...>) {
            return Cells_{std::tuple_element_t<I, internal::DiscreteProductAxes<Cells_>>{
                m_cells.template axis<I>().cardinality() * factor
            }...};
        }(std::make_index_sequence<Cells_::AxisCount>{});
        PinnedVector spacing = m_spacing / Scalar(factor);
        PinnedVector first = m_first - Scalar(0.5) * Scalar(factor - 1) * spacing;
        auto cellCovariance = std::optional<PinnedMatrix>{};
        if (hasCellCovariance())
        {
            cellCovariance = m_cellCovariances.front() / Scalar(factor * factor);
        }
        return PointGrid{
            m_space,
            std::move(cells),
            std::move(first),
            std::move(spacing),
            std::move(cellCovariance),
            get_allocator()
        };
    }
 
private:

    struct Parts
    {};

    [[nodiscard]] static std::span<PinnedMatrix const>
    listOf(std::optional<PinnedMatrix> const & covariance)
    {
        return covariance ? std::span<PinnedMatrix const>{&*covariance, 1} : std::span<PinnedMatrix const>{};
    }

    PointGrid(
        Parts,
        Space_ space,
        Cells_ cells,
        std::span<PinnedMatrix const> cellCovariances,
        Allocator_ const & allocator
    )
        : m_space{std::move(space)}
        , m_cells{std::move(cells)}
        , m_split{m_space}
        , m_cellCovariances(cellCovariances.begin(), cellCovariances.end(), MatrixAllocator(allocator))
        , m_centers(ScalarAllocator(allocator))
    {
        EMRYS_FILTER_ASSERT(
            (Shared || Eigen::Index(m_cellCovariances.size()) == size()) && "give one covariance per cell"
        );
        for (PinnedMatrix const & covariance : m_cellCovariances)
        {
            checkCovariance(covariance);
        }
    }

    template <typename CentersDerived>
    void
    setCenters(Eigen::MatrixBase<CentersDerived> const & centers)
    {
        EMRYS_FILTER_ASSERT(centers.cols() == m_cells.cardinality() && "give one center per cell");
        EMRYS_FILTER_ASSERT(centers.rows() == pinnedRows().size && "the centers do not fit the block");
        m_centers.resize(std::size_t(centers.size()));
        auto stored = SampleSetMap<PinnedSpace, Scalar_>{m_centers.data(), centers.rows(), centers.cols()};
        stored = centers;
        normalizeValues(pinnedSpace(), stored);
    }

    void
    setGeometry(PinnedVector first, PinnedVector spacing)
    {
        EMRYS_FILTER_ASSERT(
            first.size() == pinnedRows().size && spacing.size() == pinnedRows().size
            && pinnedRows().size == Eigen::Index(Cells_::AxisCount)
            && "give the first point and spacing per axis, one axis per coordinate of the block"
        );
        m_first = std::move(first);
        m_spacing = std::move(spacing);
    }

    void
    checkCovariance(PinnedMatrix const & covariance) const
    {
        EMRYS_FILTER_ASSERT(
            covariance.rows() == pinnedDifferenceRows().size && covariance.cols() == pinnedDifferenceRows().size
            && "a cell covariance does not fit the block"
        );
    }

    using MatrixAllocator = internal::ReBound<Allocator_, PinnedMatrix>;

    using ScalarAllocator = internal::ReBound<Allocator_, Scalar_>;
 
    Space_ m_space;                                                
    Cells_ m_cells;                                                
    internal::BlockSplit<Space_, Path_> m_split;                   
    internal::Buffer<PinnedMatrix, Allocator_> m_cellCovariances;  
    internal::Buffer<Scalar_, Allocator_> m_centers;               
    PinnedVector m_first;                                          
    PinnedVector m_spacing;                                        
};

template <typename P, Space S, typename Derived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
pointGrid(
    S space,
    Eigen::MatrixBase<Derived> const & centers,
    std::optional<typename Moments<internal::PathTag<P>, typename Derived::Scalar>::MatrixType> cellCovariance = {},
    Allocator const & allocator = {}
)
{
    using Grid = PointGrid<S, P, PointCells<P>, SharedCovariance, typename Derived::Scalar, Allocator>;
    return Grid{std::move(space), PointCells<P>{centers.cols()}, centers, std::move(cellCovariance), allocator};
}

template <typename P, Space S, typename Derived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
pointGrid(
    S space,
    Eigen::MatrixBase<Derived> const & centers,
    std::span<typename Moments<internal::PathTag<P>, typename Derived::Scalar>::MatrixType const> cellCovariances,
    Allocator const & allocator = {}
)
{
    using Grid = PointGrid<S, P, PointCells<P>, PerPointCovariance, typename Derived::Scalar, Allocator>;
    return Grid{std::move(space), PointCells<P>{centers.cols()}, centers, cellCovariances, allocator};
}

template <FixedSizeSpace S, typename P, typename Derived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
pointGrid(
    Eigen::MatrixBase<Derived> const & centers,
    std::optional<typename Moments<internal::PathTag<P>, typename Derived::Scalar>::MatrixType> cellCovariance = {},
    Allocator const & allocator = {}
)
{
    return pointGrid<P>(S{}, centers, std::move(cellCovariance), allocator);
}

template <FixedSizeSpace S, typename P, typename Derived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
pointGrid(
    Eigen::MatrixBase<Derived> const & centers,
    std::span<typename Moments<internal::PathTag<P>, typename Derived::Scalar>::MatrixType const> cellCovariances,
    Allocator const & allocator = {}
)
{
    return pointGrid<P>(S{}, centers, cellCovariances, allocator);
}
 
namespace internal
{

template <typename Cells, std::size_t D>
[[nodiscard]] Cells
regularCells(std::array<Eigen::Index, D> const & counts)
{
    return [&]<std::size_t... I>(std::index_sequence<I...>) {
        return Cells{std::tuple_element_t<I, DiscreteProductAxes<Cells>>{counts[I]}...};
    }(std::make_index_sequence<D>{});
}
 
}  // namespace internal

template <typename P, Space S, typename FirstDerived, typename SpacingDerived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
regularGrid(
    S space,
    Eigen::MatrixBase<FirstDerived> const & first,
    Eigen::MatrixBase<SpacingDerived> const & spacing,
    std::array<Eigen::Index, std::size_t(internal::PathTag<P>::FixedValueDimension)> const & counts,
    std::optional<typename Moments<internal::PathTag<P>, typename FirstDerived::Scalar>::MatrixType> cellCovariance =
        {},
    Allocator const & allocator = {}
)
{
    static_assert(internal::PathTag<P>::ValueDimensionStep == 0, "a regular grid needs a fixed-size block");
    using Grid = PointGrid<S, P, GridCells<P>, SharedCovariance, typename FirstDerived::Scalar, Allocator>;
    return Grid{
        std::move(space),
        internal::regularCells<typename Grid::CellsType>(counts),
        typename Grid::PinnedVector{first},
        typename Grid::PinnedVector{spacing},
        std::move(cellCovariance),
        allocator
    };
}

template <typename P, Space S, typename FirstDerived, typename SpacingDerived, typename Allocator = DefaultAllocator>
[[nodiscard]] auto
regularGrid(
    S space,
    Eigen::MatrixBase<FirstDerived> const & first,
    Eigen::MatrixBase<SpacingDerived> const & spacing,
    std::array<Eigen::Index, std::size_t(internal::PathTag<P>::FixedValueDimension)> const & counts,
    std::span<typename Moments<internal::PathTag<P>, typename FirstDerived::Scalar>::MatrixType const> cellCovariances,
    Allocator const & allocator = {}
)
{
    static_assert(internal::PathTag<P>::ValueDimensionStep == 0, "a regular grid needs a fixed-size block");
    using Grid = PointGrid<S, P, GridCells<P>, PerPointCovariance, typename FirstDerived::Scalar, Allocator>;
    return Grid{
        std::move(space),
        internal::regularCells<typename Grid::CellsType>(counts),
        typename Grid::PinnedVector{first},
        typename Grid::PinnedVector{spacing},
        cellCovariances,
        allocator
    };
}

template <
    FixedSizeSpace S,
    typename P,
    typename FirstDerived,
    typename SpacingDerived,
    typename Allocator = DefaultAllocator>
[[nodiscard]] auto
regularGrid(
    Eigen::MatrixBase<FirstDerived> const & first,
    Eigen::MatrixBase<SpacingDerived> const & spacing,
    std::array<Eigen::Index, std::size_t(internal::PathTag<P>::FixedValueDimension)> const & counts,
    std::optional<typename Moments<internal::PathTag<P>, typename FirstDerived::Scalar>::MatrixType> cellCovariance =
        {},
    Allocator const & allocator = {}
)
{
    return regularGrid<P>(S{}, first, spacing, counts, std::move(cellCovariance), allocator);
}

template <
    FixedSizeSpace S,
    typename P,
    typename FirstDerived,
    typename SpacingDerived,
    typename Allocator = DefaultAllocator>
[[nodiscard]] auto
regularGrid(
    Eigen::MatrixBase<FirstDerived> const & first,
    Eigen::MatrixBase<SpacingDerived> const & spacing,
    std::array<Eigen::Index, std::size_t(internal::PathTag<P>::FixedValueDimension)> const & counts,
    std::span<typename Moments<internal::PathTag<P>, typename FirstDerived::Scalar>::MatrixType const> cellCovariances,
    Allocator const & allocator = {}
)
{
    return regularGrid<P>(S{}, first, spacing, counts, cellCovariances, allocator);
}
 
namespace internal
{

template <typename Grid>
class GridPoints
{
public:

    using Scalar = typename Grid::Scalar;

    using SpaceType = typename Grid::SpaceType;

    using PathType = typename Grid::PathType;

    using PinnedSpace = typename Grid::PinnedSpace;

    using RestSpace = typename Grid::RestSpace;

    static constexpr bool HasKernels = true;

    GridPoints(Grid const & grid, std::span<Eigen::Index const> cells)
        : m_grid{grid}
        , m_cells{cells}
    {}

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

    [[nodiscard]] PinnedSpace const &
    pinnedSpace() const
    {
        return m_grid.pinnedSpace();
    }

    [[nodiscard]] RestSpace const &
    restSpace() const
    {
        return m_grid.restSpace();
    }

    [[nodiscard]] BlockRows
    pinnedRows() const
    {
        return m_grid.pinnedRows();
    }

    [[nodiscard]] auto const &
    restIndices() const
    {
        return m_grid.restIndices();
    }

    [[nodiscard]] BlockRows
    pinnedDifferenceRows() const
    {
        return m_grid.pinnedDifferenceRows();
    }

    [[nodiscard]] auto const &
    restDifferenceIndices() const
    {
        return m_grid.restDifferenceIndices();
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return Eigen::Index(m_cells.size());
    }

    [[nodiscard]] auto
    center(Eigen::Index i) const
    {
        return m_grid.center(m_cells[std::size_t(i)]);
    }

    [[nodiscard]] typename Grid::PinnedMatrix const *
    kernel(Eigen::Index i) const
    {
        return m_grid.hasCellCovariance() ? &m_grid.cellCovariance(m_cells[std::size_t(i)]) : nullptr;
    }
 
private:
 
    Grid const & m_grid;                    
    std::span<Eigen::Index const> m_cells;  
};
 
}  // namespace internal

template <
    typename Grid_,
    typename Covariance_ = SharedCovariance,
    typename Component_ = Gaussian<typename Grid_::RestSpace, typename Grid_::Scalar>>
class GridGaussianParams
{
    static_assert(internal::CovarianceForm<Covariance_>, "the covariance is SharedCovariance or PerPointCovariance");
    static constexpr bool Shared = std::is_same_v<Covariance_, SharedCovariance>;  
 
    static_assert(
        !Shared || std::is_same_v<typename Grid_::CellCovarianceType, SharedCovariance>,
        "cells with covariances of their own give each cell its own innovation covariance, so the "
        "filter needs PerPointCovariance"
    );

    static constexpr bool IsHybrid = internal::isGaussianHybrid<Params<Component_>>;
 
    static_assert(
        std::is_same_v<
            Params<Component_>,
            GaussianParams<typename Params<Component_>::SpaceType, typename Params<Component_>::Scalar>>
            || IsHybrid,
        "a cell holds a Gaussian or a hybrid of Gaussians"
    );
    static_assert(
        std::is_same_v<typename Params<Component_>::SpaceType, typename Grid_::RestSpace>
            && std::is_same_v<typename Params<Component_>::Scalar, typename Grid_::Scalar>,
        "a cell's distribution is over the grid's RestSpace, in the grid's scalar"
    );
    static_assert(
        !(IsHybrid && Shared),
        "mixing a cell's modes gives each cell covariances of its own, so a hybrid per cell needs PerPointCovariance"
    );
 
public:

    using GridType = Grid_;

    using Scalar = typename Grid_::Scalar;

    using SpaceType = typename Grid_::SpaceType;

    using CellsType = typename Grid_::CellsType;

    using RestSpace = typename Grid_::RestSpace;

    using PinnedSpace = typename Grid_::PinnedSpace;

    using CovarianceType = Covariance_;

    using ConditionalType = Params<Component_>;

    using WeightArrayType = Eigen::Array<Scalar, Eigen::Dynamic, 1>;

    using LogWeightsType = Eigen::Map<WeightArrayType const>;

    using allocator_type = typename Grid_::allocator_type;

    using DensityType = GridGaussian<Grid_, Covariance_, Component_>;
 
    static constexpr bool IsParams = true;  

    GridGaussianParams(Grid_ grid, ConditionalType const & conditional, Covariance_ = {})
        : m_grid{std::move(grid)}
        , m_cells(IndexAllocator(m_grid.get_allocator()))
        , m_logWeights(ScalarAllocator(m_grid.get_allocator()))
        , m_conditionals{conditional, m_grid.size(), m_grid.get_allocator()}
        , m_kept(IndexAllocator(m_grid.get_allocator()))
        , m_weightScratch(ScalarAllocator(m_grid.get_allocator()))
        , m_children(ChildAllocator(m_grid.get_allocator()))
    {
        EMRYS_FILTER_ASSERT(
            conditional.space() == m_grid.restSpace() && "the Gaussian must be over the rest of the state"
        );
        m_cells.resize(std::size_t(m_grid.size()));
        for (std::size_t c = 0; c < m_cells.size(); ++c)
        {
            m_cells[c] = Eigen::Index(c);
        }
        m_logWeights.assign(std::size_t(m_grid.size()), -Eigen::numext::log(Scalar(m_grid.size())));
        m_kept.reserve(m_cells.size());
        m_weightScratch.reserve(m_cells.size());
    }

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

    [[nodiscard]] Grid_ const &
    grid() const
    {
        return m_grid;
    }

    [[nodiscard]] Eigen::Index
    size() const
    {
        return Eigen::Index(m_cells.size());
    }

    [[nodiscard]] Eigen::Index
    cell(Eigen::Index i) const
    {
        assertPosition(i);
        return m_cells[std::size_t(i)];
    }

    [[nodiscard]] std::span<Eigen::Index const>
    cells() const
    {
        return {m_cells.data(), m_cells.size()};
    }

    [[nodiscard]] std::optional<Eigen::Index>
    find(Eigen::Index cell) const
    {
        auto const found = std::lower_bound(m_cells.begin(), m_cells.end(), cell);
        if (found == m_cells.end() || *found != cell)
        {
            return std::nullopt;
        }
        return Eigen::Index(found - m_cells.begin());
    }

    [[nodiscard]] auto
    center(Eigen::Index i) const
    {
        return m_grid.center(cell(i));
    }

    [[nodiscard]] LogWeightsType
    logWeights() const
    {
        return {m_logWeights.data(), size()};
    }

    [[nodiscard]] Scalar
    logWeight(Eigen::Index i) const
    {
        assertPosition(i);
        return m_logWeights[std::size_t(i)];
    }

    [[nodiscard]] Eigen::Index
    mostProbable() const
    {
        Eigen::Index best = 0;
        logWeights().maxCoeff(&best);
        return best;
    }

    [[nodiscard]] CategoricalParams<CellsType, Scalar>
    toCategorical() const
    {
        using ArrayType = typename CategoricalParams<CellsType, Scalar>::ArrayType;
        ArrayType logP = ArrayType::Constant(m_grid.size(), -Eigen::NumTraits<Scalar>::infinity());
        for (Eigen::Index i = 0; i < size(); ++i)
        {
            logP(m_cells[std::size_t(i)]) = logWeight(i);
        }
        return CategoricalParams<CellsType, Scalar>{m_grid.cells(), logP};
    }

    template <typename P>
    [[nodiscard]] auto
    marginal() const
    {
        auto const [stride, axis] = internal::locateAxisByPath<P>(m_grid.cells());
        using Axis = std::remove_cvref_t<decltype(*axis)>;
        using ArrayType = typename CategoricalParams<Axis, Scalar>::ArrayType;
        Eigen::Index const count = axis->cardinality();
        // One pass over the cells per value, so the totals take only the axis's own storage.
        auto logP = ArrayType{};
        logP.resize(count);
        for (Eigen::Index v = 0; v < count; ++v)
        {
            auto total = LogSumExp<Scalar>{};
            for (Eigen::Index i = 0; i < size(); ++i)
            {
                if ((m_cells[std::size_t(i)] / stride) % count == v)
                {
                    total.add(logWeight(i));
                }
            }
            logP(v) = total.value();
        }
        return CategoricalParams<Axis, Scalar>{*axis, logP};
    }

    [[nodiscard]] std::expected<void, WeightError>
    setWeights(CategoricalParams<CellsType, Scalar> const & weights)
    {
        EMRYS_FILTER_ASSERT(weights.space() == m_grid.cells() && "the weights are over different cells");
        m_stamp.touchShape();
        m_weightScratch.resize(m_logWeights.size());
        for (Eigen::Index i = 0; i < size(); ++i)
        {
            m_weightScratch[std::size_t(i)] = weights.logWeight(m_cells[std::size_t(i)]);
        }
        auto const total =
            internal::logTotalWeight<Scalar>(std::span<Scalar const>{m_weightScratch.data(), m_weightScratch.size()});
        if (!total)
        {
            return std::unexpected{total.error()};
        }
        m_logWeights.swap(m_weightScratch);
        normalizeAndRemoveDead();
        return {};
    }

    [[nodiscard]] decltype(auto)
    conditional(Eigen::Index i) const
    {
        assertPosition(i);
        return m_conditionals.conditional(i);
    }

    [[nodiscard]] auto
    modes() const
        requires IsHybrid
    {
        return internal::summedModes(m_conditionals, logWeights());
    }

    [[nodiscard]] typename Moments<RestSpace, Scalar>::MatrixType const &
    sharedCovariance() const
        requires Shared
    {
        return m_conditionals.sharedCovariance();
    }

    [[nodiscard]] std::expected<void, NumericError>
    factorize() const
    {
        return m_conditionals.factorize();
    }

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

    void
    toDensity() const && = delete;

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

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

    [[nodiscard]] Scalar
    pruneToSize(Eigen::Index maxSize)
    {
        return pruneWith(internal::PruneToSizeRule{maxSize});
    }

    void
    subdivide(Eigen::Index factor)
        requires(Grid_::Regular && std::is_same_v<typename Grid_::CellCovarianceType, SharedCovariance>)
    {
        constexpr std::size_t D = CellsType::AxisCount;
        EMRYS_FILTER_ASSERT(factor >= 1 && "subdivide each cell into at least one");
        m_stamp.touchShape();
        Grid_ refined = m_grid.subdivided(factor);
        Eigen::Index offsets = 1;
        for (std::size_t d = 0; d < D; ++d)
        {
            offsets *= factor;
        }
        auto & children = m_children;
        children.clear();
        children.reserve(m_cells.size() * std::size_t(offsets));
        for (Eigen::Index i = 0; i < size(); ++i)
        {
            auto const parent = m_grid.cells().coordinates(m_cells[std::size_t(i)]);
            for (Eigen::Index o = 0; o < offsets; ++o)
            {
                auto child = std::array<Eigen::Index, D>{};
                Eigen::Index rest = o;
                for (std::size_t d = D; d-- > 0;)
                {
                    child[d] = factor * parent[d] + rest % factor;
                    rest /= factor;
                }
                children.emplace_back(refined.cells().value(child), i);
            }
        }
        std::sort(children.begin(), children.end());
        m_kept.clear();
        m_weightScratch.clear();
        Scalar const share = Eigen::numext::log(Scalar(offsets));
        for (std::size_t k = 0; k < children.size(); ++k)
        {
            m_kept.push_back(children[k].second);
            m_weightScratch.push_back(m_logWeights[std::size_t(children[k].second)] - share);
        }
        m_conditionals.select(m_kept);
        m_grid = std::move(refined);
        m_cells.resize(children.size());
        for (std::size_t k = 0; k < children.size(); ++k)
        {
            m_cells[k] = children[k].first;
        }
        m_logWeights.swap(m_weightScratch);
    }
 
private:

    template <typename, typename, typename, typename, typename>
    friend class RbPointMassFilter;
 
    friend DensityType;  

    template <typename>
    friend class internal::ValuesCheck;

    using ConditionalsType = internal::Conditionals<RestSpace, Covariance_, Scalar, allocator_type, ConditionalType>;

    using IndexAllocator = internal::ReBound<allocator_type, Eigen::Index>;

    using ScalarAllocator = internal::ReBound<allocator_type, Scalar>;

    [[nodiscard]] internal::GridPoints<Grid_>
    points() const
    {
        return {m_grid, cells()};
    }

    [[nodiscard]] GaussianParams<SpaceType, Scalar>
    collapseCells() const
    {
        if constexpr (IsHybrid)
        {
            return internal::collapseGivenHybrids(points(), m_conditionals, logWeights(), get_allocator());
        }
        else
        {
            return internal::collapseGivenPoints(points(), m_conditionals, logWeights());
        }
    }

    void
    assertPosition([[maybe_unused]] Eigen::Index i) const
    {
        EMRYS_FILTER_ASSERT(i >= 0 && i < size() && "the position is out of range");
    }

    template <ExecutionPolicy P, typename Derived>
    [[nodiscard]] std::pair<Scalar, bool>
    reweight(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods, ConditionalsType & posterior)
    {
        return reweightThen(std::forward<P>(policy), logLikelihoods, [&] {
            m_conditionals.swap(posterior);
        });
    }

    template <ExecutionPolicy P, typename Derived>
    [[nodiscard]] std::pair<Scalar, bool>
    reweight(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods)
    {
        return reweightThen(std::forward<P>(policy), logLikelihoods, [] {});
    }

    template <ExecutionPolicy P, typename Derived, typename Place>
    [[nodiscard]] std::pair<Scalar, bool>
    reweightThen(P && policy, Eigen::ArrayBase<Derived> const & logLikelihoods, Place const & place)
    {
        EMRYS_FILTER_ASSERT(logLikelihoods.size() == size() && "give one log-likelihood per live cell");
        m_weightScratch.resize(m_logWeights.size());
        internal::forEachIndex(policy, size(), [&](Eigen::Index i) {
            m_weightScratch[std::size_t(i)] = m_logWeights[std::size_t(i)] + logLikelihoods(i);
        });
        auto const evidence = internal::logTotalWeight<Scalar>(
            policy,
            std::span<Scalar const>{m_weightScratch.data(), m_weightScratch.size()}
        );
        if (!evidence)
        {
            return {
                logSumExp<Scalar>(
                    std::forward<P>(policy),
                    size(),
                    [&](Eigen::Index i) {
                        return m_weightScratch[std::size_t(i)];
                    }
                ),
                false
            };
        }
        m_stamp.touchShape();
        m_logWeights.swap(m_weightScratch);
        place();
        normalizeAndRemoveDead();
        return {*evidence, true};
    }

    template <typename Rule>
    [[nodiscard]] Scalar
    pruneWith(Rule const & rule)
    {
        m_stamp.touchShape();
        auto cut = rule(
            size(),
            [&](Eigen::Index i) {
                return m_logWeights[std::size_t(i)];
            },
            Scalar(0),
            m_weightScratch
        );
        auto removed = LogSumExp<Scalar>{};
        for (Scalar & logWeight : m_logWeights)
        {
            if (cut.removes(logWeight))
            {
                removed.add(logWeight);
                logWeight = -Eigen::NumTraits<Scalar>::infinity();
            }
        }
        normalizeAndRemoveDead();
        return removed.value();
    }

    void
    normalizeAndRemoveDead()
    {
        auto const total = internal::logTotalWeight<Scalar>(m_logWeights);
        EMRYS_FILTER_ASSERT(total && "some live cell must keep a nonzero weight");
        for (Scalar & logWeight : m_logWeights)
        {
            logWeight -= *total;
        }
        m_kept.clear();
        for (Eigen::Index i = 0; i < size(); ++i)
        {
            if (m_logWeights[std::size_t(i)] > -Eigen::NumTraits<Scalar>::infinity())
            {
                m_kept.push_back(i);
            }
        }
        if (m_kept.size() == m_cells.size())
        {
            return;
        }
        for (std::size_t k = 0; k < m_kept.size(); ++k)
        {
            m_cells[k] = m_cells[std::size_t(m_kept[k])];
            m_logWeights[k] = m_logWeights[std::size_t(m_kept[k])];
        }
        m_conditionals.keep(m_kept);
        m_cells.resize(m_kept.size());
        m_logWeights.resize(m_kept.size());
    }
 
    Grid_ m_grid;                                            
    internal::Buffer<Eigen::Index, allocator_type> m_cells;  
    internal::Buffer<Scalar, allocator_type> m_logWeights;   
    ConditionalsType m_conditionals;                         
 
    // Scratch, reused: the cells kept, and the reweighted log-probabilities.
    internal::Buffer<Eigen::Index, allocator_type> m_kept;  
 
    internal::Buffer<Scalar, allocator_type> m_weightScratch;  

    using Child = std::pair<Eigen::Index, Eigen::Index>;

    using ChildAllocator = internal::ReBound<allocator_type, Child>;
 
    internal::Buffer<Child, allocator_type> m_children;  
 
    internal::Stamp m_stamp;  
};

template <typename Grid_, typename Covariance_, typename Component_>
class GridGaussian
{
public:
 
    using ParamsType = GridGaussianParams<Grid_, Covariance_, Component_>;  
    using GridType = Grid_;                                                 
    using Scalar = typename Grid_::Scalar;                                  
    using SpaceType = typename Grid_::SpaceType;                            
    using ConditionalType = typename ParamsType::ConditionalType;           
    using MomentsType = Moments<SpaceType, Scalar>;              
    using allocator_type = typename ParamsType::allocator_type;  
 
    static constexpr bool IsParams = false;  

    [[nodiscard]] Grid_ const &
    grid() const
    {
        return source().grid();
    }

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

    [[nodiscard]] decltype(auto)
    conditional(Eigen::Index i) const
    {
        return source().conditional(i);
    }

    [[nodiscard]] GaussianParams<SpaceType, Scalar>
    collapsed() const
    {
        return source().collapseCells();
    }

    [[nodiscard]] MomentsType
    moments() const
    {
        return collapsed().moments();
    }
 
private:
 
    friend ParamsType;  

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

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

// A hybrid conditional names the component, and its cells have covariances of their own.
template <typename Grid_, typename Component_, typename ModeSpace_, typename Allocator_>
GridGaussianParams(Grid_, HybridParams<Component_, ModeSpace_, Allocator_>, PerPointCovariance = {})
    -> GridGaussianParams<Grid_, PerPointCovariance, Hybrid<Component_, ModeSpace_, Allocator_>>;
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Grid_, typename Covariance_, typename Component_>
void
describe(std::ostream & stream, GridGaussianParams<Grid_, Covariance_, Component_> const & state, Indent indent)
{
    stream << indent << "Rao-Blackwellized grid of ";
    printCount(stream, state.size(), "cell");
    // Every step leaves the cells' weights normalized, so only a cell's distribution can fail.
    auto const density = state.toDensity();
    if (!density)
    {
        stream << '\n';
        describe(stream, density.error(), indent.deeper());
        return;
    }
    auto const collapsed = density->collapsed();
    printDimension(stream, collapsed.valueDimension(), collapsed.differenceDimension());
    printMoments(stream, indent.deeper(), collapsed.moments());
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO