Conditionals.h

Go to the documentation of this file.

#pragma once

 
#include <cmath>
#include <cstddef>
#include <expected>
#include <numbers>
#include <optional>
#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/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Weights.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/Sampling.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{

struct SharedCovariance
{};

struct PerPointCovariance
{};
 
namespace internal
{

template <typename C>
concept CovarianceForm = std::is_same_v<C, SharedCovariance> || std::is_same_v<C, PerPointCovariance>;

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

template <typename Component_, typename ModeSpace_, typename Allocator_>
inline constexpr bool isGaussianHybrid<HybridParams<Component_, ModeSpace_, Allocator_>> = std::is_same_v<
    Params<Component_>,
    GaussianParams<typename Params<Component_>::SpaceType, typename Params<Component_>::Scalar>>;

template <typename Space, typename P>
struct ComplementOf
{
    using type = decltype(complementSubspace<P>(std::declval<Space const &>()));
};

template <typename Space, typename... Tags>
struct ComplementOf<Space, TagPath<Tags...>>
{
    using type = decltype(complementSubspace<Tags...>(std::declval<Space const &>()));
};

template <typename P, typename Space>
[[nodiscard]] auto
complementOf(Space const & space)
{
    return withPathTags<P>([&]<typename... Tags>() {
        return complementSubspace<Tags...>(space);
    });
}

template <typename Space_, typename Path_>
class BlockSplit
{
public:

    using PathType = Path_;

    using PinnedSpace = PathTag<Path_>;

    using RestSpace = typename ComplementOf<Space_, Path_>::type;

    using RestIndices = IndexVector<valueDimensionAtCompileTime<RestSpace>, maxValueDimensionAtCompileTime<RestSpace>>;

    using RestDifferenceIndices = IndexVector<
        valueDimensionAtCompileTime<DifferenceSpace<RestSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<RestSpace>>>;

    explicit BlockSplit(Space_ const & space)
        : m_pinned{pathBlock<Path_>(space)}
        , m_rest{complementOf<Path_>(space)}
        , m_rows{blockRows<Path_>(space)}
        , m_differenceRows{differenceBlockRows<Path_>(space)}
        , m_size{valueDimension(space)}
        , m_restIndices{outside<RestIndices>(m_rows, m_size)}
        , m_restDifferenceIndices{outside<RestDifferenceIndices>(m_differenceRows, differenceDimension(space))}
    {}

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

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

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

    [[nodiscard]] RestIndices const &
    restIndices() const
    {
        return m_restIndices;
    }

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

    [[nodiscard]] RestDifferenceIndices const &
    restDifferenceIndices() const
    {
        return m_restDifferenceIndices;
    }

    template <typename X, typename P>
    [[nodiscard]] auto
    assemble(X const & rest, P const & pinned) const
    {
        using T = ScalarOf<X>;
        EMRYS_FILTER_ASSERT(
            rest.size() == m_restIndices.size() && pinned.size() == m_rows.size && "the blocks do not fit the space"
        );
        Eigen::Matrix<
            T,
            valueDimensionAtCompileTime<Space_>,
            1,
            Eigen::ColMajor,
            maxValueDimensionAtCompileTime<Space_>,
            1>
            full;
        full.resize(m_size);
        full(m_restIndices) = rest;
        middleRowsOf<valueDimensionAtCompileTime<PinnedSpace>>(full, m_rows.start, m_rows.size) =
            pinned.template cast<T>();
        return full;
    }
 
private:

    template <typename Indices>
    [[nodiscard]] static Indices
    outside(BlockRows rows, Eigen::Index n)
    {
        auto result = Indices{};
        result.resize(n - rows.size);
        Eigen::Index k = 0;
        for (Eigen::Index i = 0; i < n; ++i)
        {
            if (i < rows.start || i >= rows.start + rows.size)
            {
                result(k++) = i;
            }
        }
        return result;
    }
 
    PinnedSpace m_pinned;                           
    RestSpace m_rest;                               
    BlockRows m_rows;                               
    BlockRows m_differenceRows;                     
    Eigen::Index m_size;                            
    RestIndices m_restIndices;                      
    RestDifferenceIndices m_restDifferenceIndices;  
};

template <
    typename Space,
    typename Covariance,
    typename Scalar,
    typename Allocator = DefaultAllocator,
    typename Element = GaussianParams<Space, Scalar>>
class Conditionals;

template <typename Space, typename Scalar_, typename Allocator>
class Conditionals<Space, SharedCovariance, Scalar_, Allocator, GaussianParams<Space, Scalar_>>
{
public:

    using Scalar = Scalar_;

    using ConditionalType = GaussianParams<Space, Scalar>;

    using MatrixType = typename ConditionalType::MatrixType;

    using MeansType = SampleSet<Space, Scalar>;
 
    static constexpr bool Shared = true;  

    Conditionals(ConditionalType const & conditional, Eigen::Index count, Allocator const & allocator = {})
        : m_space{conditional.space()}
        , m_rows{conditional.valueDimension()}
        , m_count{count}
        , m_means(ScalarAllocator(allocator))
        , m_scratch(ScalarAllocator(allocator))
        , m_cov{conditional.covarianceLower()}
    {
        m_means.resize(std::size_t(m_rows * m_count));
        means().colwise() = conditional.mean();
    }

    Conditionals(Conditionals const & other, Allocator const & allocator)
        : m_space{other.m_space}
        , m_rows{other.m_rows}
        , m_count{other.m_count}
        , m_means(other.m_means, ScalarAllocator(allocator))
        , m_scratch(ScalarAllocator(allocator))
        , m_cov{other.m_cov}
    {}

    Conditionals(Conditionals const & other)
        : Conditionals{other, other.m_means.get_allocator()}
    {}
 
    Conditionals(Conditionals &&) noexcept = default;  

    Conditionals &
    operator=(Conditionals const & other)
    {
        if (this == &other)
        {
            return *this;
        }
        m_space = other.m_space;
        m_rows = other.m_rows;
        m_count = other.m_count;
        m_means.assign(other.m_means.begin(), other.m_means.end());
        m_cov = other.m_cov;
        return *this;
    }

    Conditionals &
    operator=(Conditionals &&) noexcept = default;
 
    ~Conditionals() = default;

    void
    resizeLike(Conditionals const & other)
    {
        m_space = other.m_space;
        m_rows = other.m_rows;
        m_count = other.m_count;
        m_means.resize(std::size_t(m_rows * m_count));
    }

    void
    swap(Conditionals & other) noexcept
    {
        std::swap(m_space, other.m_space);
        std::swap(m_rows, other.m_rows);
        std::swap(m_count, other.m_count);
        m_means.swap(other.m_means);
        m_scratch.swap(other.m_scratch);
        std::swap(m_cov, other.m_cov);
    }

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

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

    [[nodiscard]] auto
    mean(Eigen::Index i) const
    {
        return means().col(i);
    }

    [[nodiscard]] MatrixType const &
    covarianceLower(Eigen::Index /*i*/) const
    {
        return m_cov;
    }

    [[nodiscard]] ConditionalType
    conditional(Eigen::Index i) const
    {
        return ConditionalType{m_space, mean(i), m_cov};
    }

    [[nodiscard]] std::expected<void, NumericError>
    factorize() const
    {
        if constexpr (MatrixType::RowsAtCompileTime == 0)
        {
            return {};
        }
        else
        {
            ConditionalType const first = conditional(0);
            return first.factorize();
        }
    }

    [[nodiscard]] ConstSampleSetMap<Space, Scalar>
    means() const
    {
        return {m_means.data(), m_rows, m_count};
    }

    [[nodiscard]] SampleSetMap<Space, Scalar>
    means()
    {
        return {m_means.data(), m_rows, m_count};
    }

    [[nodiscard]] MatrixType const &
    sharedCovariance() const
    {
        return m_cov;
    }

    [[nodiscard]] MatrixType &
    sharedCovariance()
    {
        return m_cov;
    }

    void
    reserve(Eigen::Index count)
    {
        m_means.reserve(std::size_t(m_rows * count));
        m_scratch.reserve(std::size_t(m_rows * count));
    }

    void
    select(std::span<Eigen::Index const> from)
    {
        select(execution::sequential, from);
    }

    template <ExecutionPolicy P>
    void
    select(P && policy, std::span<Eigen::Index const> from)
    {
        Eigen::Index const count = Eigen::Index(from.size());
        m_scratch.resize(std::size_t(m_rows * count));
        gather(std::forward<P>(policy), from, means(), SampleSetMap<Space, Scalar>{m_scratch.data(), m_rows, count});
        m_means.swap(m_scratch);
        m_count = count;
    }

    void
    keep(std::span<Eigen::Index const> kept)
    {
        auto all = means();
        for (std::size_t k = 0; k < kept.size(); ++k)
        {
            if (kept[k] != Eigen::Index(k))
            {
                all.col(Eigen::Index(k)) = all.col(kept[k]);
            }
        }
        m_count = Eigen::Index(kept.size());
        m_means.resize(std::size_t(m_rows * m_count));
    }
 
private:

    using ScalarAllocator = ReBound<Allocator, Scalar>;
 
    Space m_space;                        
    Eigen::Index m_rows;                  
    Eigen::Index m_count;                 
    Buffer<Scalar, Allocator> m_means;    
    Buffer<Scalar, Allocator> m_scratch;  
    MatrixType m_cov;                     
};

template <typename Space, typename Scalar_, typename Allocator, typename Element>
class Conditionals<Space, PerPointCovariance, Scalar_, Allocator, Element>
{
    static constexpr bool IsGaussian = std::is_same_v<Element, GaussianParams<Space, Scalar_>>;
 
public:

    using Scalar = Scalar_;

    using ConditionalType = Element;

    using MatrixType = typename Moments<Space, Scalar>::MatrixType;
 
    static constexpr bool Shared = false;  

    Conditionals(ConditionalType const & conditional, Eigen::Index count, Allocator const & allocator = {})
        : m_conditionals(std::size_t(count), conditional, GaussianAllocator(allocator))
        , m_scratch(GaussianAllocator(allocator))
    {}

    Conditionals(Conditionals const & other, Allocator const & allocator)
        : m_conditionals(other.m_conditionals, GaussianAllocator(allocator))
        , m_scratch(GaussianAllocator(allocator))
    {}

    Conditionals(Conditionals const & other)
        : Conditionals{other, other.m_conditionals.get_allocator()}
    {}
 
    Conditionals(Conditionals &&) noexcept = default;  

    Conditionals &
    operator=(Conditionals const & other)
    {
        if (this != &other)
        {
            m_conditionals.assign(other.m_conditionals.begin(), other.m_conditionals.end());
        }
        return *this;
    }

    Conditionals &
    operator=(Conditionals &&) noexcept = default;
 
    ~Conditionals() = default;

    void
    resizeLike(Conditionals const & other)
        requires IsGaussian
    {
        m_conditionals.resize(other.m_conditionals.size());
    }

    void
    swap(Conditionals & other) noexcept
    {
        m_conditionals.swap(other.m_conditionals);
        m_scratch.swap(other.m_scratch);
    }

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

    [[nodiscard]] Space const &
    space() const
    {
        return m_conditionals.front().space();
    }

    [[nodiscard]] decltype(auto)
    mean(Eigen::Index i) const
        requires IsGaussian
    {
        return m_conditionals[std::size_t(i)].mean();
    }

    [[nodiscard]] MatrixType const &
    covarianceLower(Eigen::Index i) const
        requires IsGaussian
    {
        return m_conditionals[std::size_t(i)].covarianceLower();
    }

    [[nodiscard]] ConditionalType const &
    conditional(Eigen::Index i) const
    {
        return m_conditionals[std::size_t(i)];
    }

    [[nodiscard]] ConditionalType &
    conditional(Eigen::Index i)
    {
        return m_conditionals[std::size_t(i)];
    }

    [[nodiscard]] std::expected<void, NumericError>
    factorize() const
    {
        if constexpr (MatrixType::RowsAtCompileTime != 0)
        {
            for (auto const & conditional : m_conditionals)
            {
                if (auto result = conditional.factorize(); !result)
                {
                    return std::unexpected{result.error()};
                }
            }
        }
        return {};
    }

    void
    reserve(Eigen::Index count)
    {
        m_conditionals.reserve(std::size_t(count));
        m_scratch.reserve(std::size_t(count));
    }

    void
    select(std::span<Eigen::Index const> from)
    {
        select(execution::sequential, from);
    }

    template <ExecutionPolicy P>
    void
    select(P && policy, std::span<Eigen::Index const> from)
    {
        if (from.size() < m_scratch.size())
        {
            m_scratch.erase(m_scratch.begin() + std::ptrdiff_t(from.size()), m_scratch.end());
        }
        while (m_scratch.size() < from.size())
        {
            m_scratch.push_back(m_conditionals.front());
        }
        if constexpr (isSequential<P>)
        {
            for (std::size_t j = 0; j < from.size(); ++j)
            {
                m_scratch[j] = m_conditionals[std::size_t(from[j])];
            }
        }
        else
        {
            forEachIndex(std::forward<P>(policy), Eigen::Index(from.size()), [&](Eigen::Index j) {
                m_scratch[std::size_t(j)] = m_conditionals[std::size_t(from[std::size_t(j)])];
            });
        }
        m_conditionals.swap(m_scratch);
    }

    void
    keep(std::span<Eigen::Index const> kept)
    {
        for (std::size_t k = 0; k < kept.size(); ++k)
        {
            if (kept[k] != Eigen::Index(k))
            {
                m_conditionals[k] = m_conditionals[std::size_t(kept[k])];
            }
        }
        m_conditionals.erase(m_conditionals.begin() + std::ptrdiff_t(kept.size()), m_conditionals.end());
    }
 
private:

    using GaussianAllocator = ReBound<Allocator, ConditionalType>;
 
    Buffer<ConditionalType, Allocator> m_conditionals;  
    Buffer<ConditionalType, Allocator> m_scratch;  
};

template <typename Points, typename F>
[[nodiscard]] auto
withPoint(Points const & points, Eigen::Index i, F && f)
{
    using Scalar = typename Points::Scalar;
    using PinnedVector = typename Moments<typename Points::PinnedSpace, Scalar>::VectorType;
    PinnedVector center = points.center(i);
    if constexpr (Points::HasKernels)
    {
        if (auto const * kernel = points.kernel(i))
        {
            return f(
                GaussianParams<typename Points::PinnedSpace, Scalar>{points.pinnedSpace(), std::move(center), *kernel}
            );
        }
    }
    return f(center);
}

template <typename Points, typename Store, typename WeightsDerived>
[[nodiscard]] GaussianParams<typename Points::SpaceType, typename Points::Scalar>
collapseGivenPoints(
    Points const & points,
    Store const & conditionals,
    Eigen::ArrayBase<WeightsDerived> const & logWeights
)
{
    using Scalar = typename Points::Scalar;
    using SpaceType = typename Points::SpaceType;
    using MomentsType = Moments<SpaceType, Scalar>;
    EMRYS_FILTER_ASSERT(
        logWeights.size() == points.size() && conditionals.size() == points.size()
        && "give one weight and one Gaussian per point"
    );
    BlockRows const rows = points.pinnedRows();
    auto const & rest = points.restIndices();
    BlockRows const differenceRows = points.pinnedDifferenceRows();
    auto const & restDifferences = points.restDifferenceIndices();
    auto const logTotal = logTotalWeight<Scalar>(logWeights);
    EMRYS_FILTER_ASSERT(logTotal && "the weights must be finite and not all zero");
    // The whole state at point i: its center in the block's rows, its Gaussian's mean in the rest's.
    constexpr int PinnedValues = valueDimensionAtCompileTime<typename Points::PinnedSpace>;
    constexpr int PinnedDifferences = valueDimensionAtCompileTime<DifferenceSpace<typename Points::PinnedSpace>>;
    auto const meanAt = [&](Eigen::Index i) {
        auto mean = typename MomentsType::VectorType{};
        mean.resize(valueDimension(points.space()));
        middleRowsOf<PinnedValues>(mean, rows.start, rows.size) = points.center(i);
        mean(rest) = conditionals.mean(i);
        return mean;
    };
    MomentsType moments = pointMomentsIn<MomentsType>(
        points.space(),
        points.size(),
        meanAt,
        [&](Eigen::Index i) {
            return logWeights.derived()(i);
        },
        *logTotal
    );
    auto & cov = moments.covarianceLower();
    auto pinned = blockOf<PinnedDifferences, PinnedDifferences>(
        cov,
        differenceRows.start,
        differenceRows.start,
        differenceRows.size,
        differenceRows.size
    );
    for (Eigen::Index i = 0; i < points.size(); ++i)
    {
        Scalar const logWeight = logWeights.derived()(i);
        if (logWeight == -Eigen::NumTraits<Scalar>::infinity())
        {
            continue;
        }
        Scalar const weight = Eigen::numext::exp(logWeight - *logTotal);
        if constexpr (Points::HasKernels)
        {
            if (auto const * kernel = points.kernel(i))
            {
                pinned += weight * *kernel;
            }
        }
        if constexpr (!Store::Shared)
        {
            cov(restDifferences, restDifferences) += weight * conditionals.covarianceLower(i);
        }
    }
    if constexpr (Store::Shared)
    {
        cov(restDifferences, restDifferences) += conditionals.sharedCovariance();
    }
    mirrorLower(cov);
    return GaussianParams<SpaceType, Scalar>{points.space(), std::move(moments)};
}

template <typename Points, typename Store, typename WeightsDerived, typename Allocator>
[[nodiscard]] GaussianParams<typename Points::SpaceType, typename Points::Scalar>
collapseGivenHybrids(
    Points const & points,
    Store const & hybrids,
    Eigen::ArrayBase<WeightsDerived> const & logWeights,
    Allocator const & allocator
)
{
    using Scalar = typename Points::Scalar;
    using Traits = ComponentTraits<typename Store::ConditionalType>;
    using Gaussians = Conditionals<typename Points::RestSpace, PerPointCovariance, Scalar, Allocator>;
    auto gaussians = Gaussians{Traits::collapseModes(hybrids.conditional(0)), hybrids.size(), allocator};
    for (Eigen::Index i = 1; i < hybrids.size(); ++i)
    {
        gaussians.conditional(i) = Traits::collapseModes(hybrids.conditional(i));
    }
    return collapseGivenPoints(points, gaussians, logWeights);
}

template <typename Store, typename WeightsDerived>
[[nodiscard]] auto
summedModes(Store const & hybrids, Eigen::ArrayBase<WeightsDerived> const & logWeights)
{
    using ModesType = typename Store::ConditionalType::ModesType;
    using Scalar = typename ModesType::Scalar;
    auto const & modeSpace = hybrids.conditional(0).modeSpace();
    Eigen::Index const count = modeSpace.cardinality();
    auto logP = typename ModesType::ArrayType{};
    logP.resize(count);
    for (Eigen::Index m = 0; m < count; ++m)
    {
        auto total = LogSumExp<Scalar>{};
        for (Eigen::Index i = 0; i < hybrids.size(); ++i)
        {
            auto const & modes = hybrids.conditional(i).modes();
            // A point's mode weights are proportional to its probabilities, so each is normalized.
            total.add(Scalar(logWeights.derived()(i)) + modes.logWeight(m) - modes.logTotalWeight().value());
        }
        logP(m) = total.value();
    }
    auto modes = ModesType{modeSpace, logP};
    [[maybe_unused]] auto const normalized = modes.normalize();
    EMRYS_FILTER_ASSERT(normalized && "the points' weights must define a distribution");
    return modes;
}

template <typename Factor, typename Gain, typename Cov>
struct PinnedConditioning
{
    Factor factor;  
    Gain gain;      
    Cov cov;        
};

template <typename Split, typename CovDerived>
[[nodiscard]] auto
pinnedConditioning(Split const & split, Eigen::MatrixBase<CovDerived> const & cov)
{
    using Scalar = typename CovDerived::Scalar;
    using PinnedMoments = Moments<typename Split::PinnedSpace, Scalar>;
    using RestMoments = Moments<typename Split::RestSpace, Scalar>;
    using PinnedMatrix = typename PinnedMoments::MatrixType;
    using RestMatrix = typename RestMoments::MatrixType;
    // Typed by the blocks' bounds: evaluating an indexed view would give an unbounded matrix.
    using PinnedRestMatrix = MatrixOf<
        Scalar,
        PinnedMoments::DifferenceDimensionAtCompileTime,
        RestMoments::DifferenceDimensionAtCompileTime,
        PinnedMoments::MaxDifferenceDimensionAtCompileTime,
        RestMoments::MaxDifferenceDimensionAtCompileTime>;
    using GainMatrix = MatrixOf<
        Scalar,
        RestMoments::DifferenceDimensionAtCompileTime,
        PinnedMoments::DifferenceDimensionAtCompileTime,
        RestMoments::MaxDifferenceDimensionAtCompileTime,
        PinnedMoments::MaxDifferenceDimensionAtCompileTime>;
    BlockRows const rows = split.pinnedDifferenceRows();
    auto const & rest = split.restDifferenceIndices();
    constexpr int PinnedRows = PinnedMoments::DifferenceDimensionAtCompileTime;
    PinnedMatrix const pinned = blockOf<PinnedRows, PinnedRows>(cov, rows.start, rows.start, rows.size, rows.size);
    PinnedRestMatrix const pinnedRest = cov(Eigen::seqN(rows.start, rows.size), rest);
    auto const ldlt = Eigen::LDLT<PinnedMatrix>{pinned};
    auto gain = transposedSolve<GainMatrix>(ldlt, pinnedRest);
    RestMatrix restCov = cov(rest, rest);
    restCov -= multiply(gain, pinnedRest);
    restCov = (Scalar(0.5) * (restCov + restCov.transpose())).eval();
    return PinnedConditioning{defaultRoot(pinned), std::move(gain), std::move(restCov)};
}

template <
    typename Split,
    typename Conditioning,
    typename MeanDerived,
    typename ZDerived,
    typename PinnedOut,
    typename RestOut>
void
drawPinned(
    Split const & split,
    Conditioning const & conditioning,
    Eigen::MatrixBase<MeanDerived> const & mean,
    Eigen::MatrixBase<ZDerived> const & z,
    PinnedOut && pinned,
    RestOut && restMean
)
{
    BlockRows const rows = split.pinnedRows();
    auto const offset = multiply(conditioning.factor, z).eval();
    constexpr int PinnedValues = valueDimensionAtCompileTime<typename Split::PinnedSpace>;
    pinned = Emrys::plus(split.pinnedSpace(), middleRowsOf<PinnedValues>(mean, rows.start, rows.size), offset);
    restMean = Emrys::plus(split.restSpace(), mean(split.restIndices()), multiply(conditioning.gain, offset));
}
 
}  // namespace internal
 
}  // namespace Emrys