Hybrid.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <cstddef>
#include <expected>
#include <memory>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Core/Allocation.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/Component.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Discrete.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{
 
template <internal::ComponentOrInfoForm Component_, DiscreteSpace ModeSpace_, typename Allocator_ = DefaultAllocator>
class Hybrid;

template <internal::ComponentOrInfoForm Component_, DiscreteSpace ModeSpace_, typename Allocator_ = DefaultAllocator>
class HybridParams
{
    static_assert(
        Component<Component_>,
        "a Gaussian in an information form is not a hybrid's conditional, 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 Scalar = typename ComponentType::Scalar;                   
    using SpaceType = typename ComponentType::SpaceType;             
    using ModeSpaceType = ModeSpace_;                                
    using ModesType = CategoricalParams<ModeSpace_, Scalar>;         
    using MomentsType = typename ComponentType::MomentsType;         
    using allocator_type = Allocator_;                               
    using DensityType = Hybrid<Component_, ModeSpace_, Allocator_>;  

    using ConditionalsType = internal::Buffer<ComponentType, 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;  

    HybridParams(ModesType modes, ComponentType const & conditional, Allocator_ const & allocator = {})
        : m_modes{std::move(modes)}
        , m_conditionals(std::size_t(m_modes.cardinality()), conditional, ComponentAllocator(allocator))
    {}

    HybridParams(ModesType modes, ConditionalsType conditionals)
        : m_modes{std::move(modes)}
        , m_conditionals{std::move(conditionals)}
    {
        check();
    }

    template <typename Conditionals>
    HybridParams(ModesType modes, Conditionals const & conditionals, Allocator_ const & allocator)
        : m_modes{std::move(modes)}
        , m_conditionals(conditionals.begin(), conditionals.end(), ComponentAllocator(allocator))
    {
        check();
    }

    HybridParams(HybridParams const & other, Allocator_ const & allocator)
        : m_modes{other.m_modes}
        , m_conditionals{other.m_conditionals, ComponentAllocator(allocator)}
    {}

    HybridParams(HybridParams && other, Allocator_ const & allocator)
        : m_modes{std::move(other.m_modes)}
        , m_conditionals{std::move(other.m_conditionals), ComponentAllocator(allocator)}
    {
        other.m_stamp.touchShape();
    }
 
    HybridParams(HybridParams const &) = default;      
    HybridParams(HybridParams &&) noexcept = default;  

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

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

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

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

    [[nodiscard]] ModeSpace_ const &
    modeSpace() const
    {
        return m_modes.space();
    }

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

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

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

    [[nodiscard]] ModesType const &
    modes() const
    {
        return m_modes;
    }

    void
    setModes(ModesType modes)
    {
        EMRYS_FILTER_ASSERT(modes.space() == modeSpace() && "the modes are on a different space");
        m_stamp.touchValues();
        m_modes = std::move(modes);
    }

    [[nodiscard]] ComponentType const &
    conditional(Eigen::Index mode) const &
    {
        EMRYS_FILTER_ASSERT(mode >= 0 && mode < size() && "the mode is out of range");
        return m_conditionals[std::size_t(mode)];
    }

    void
    conditional(Eigen::Index mode) const && = delete;

    void
    setConditional(Eigen::Index mode, ComponentType const & component)
    {
        conditionalToSet(mode, component) = component;
    }

    void
    setConditional(Eigen::Index mode, ComponentType && component)
    {
        conditionalToSet(mode, component) = std::move(component);
    }

    template <typename F>
        requires std::invocable<F &, ComponentType &>
    decltype(auto)
    updateConditional(Eigen::Index mode, F && f)
    {
        EMRYS_FILTER_ASSERT(mode >= 0 && mode < size() && "the mode is out of range");
        m_stamp.touchValues();
        return internal::keepingSpace(m_conditionals[std::size_t(mode)], f);
    }

    template <typename... Values, typename... Paths, std::size_t... N>
        requires(sizeof...(Values) + sizeof...(Paths) > 0)
    [[nodiscard]] ComponentType const &
    conditional(AxisIndex<Paths, N> const &... indices) const &
    {
        return conditional(valueOf<Values...>(modeSpace(), indices...));
    }

    template <typename... Values, typename... Paths, std::size_t... N>
        requires(sizeof...(Values) + sizeof...(Paths) > 0)
    void
    conditional(AxisIndex<Paths, N> const &...) const && = delete;

    void
    swapConditionals(ConditionalsType & conditionals)
    {
        EMRYS_FILTER_ASSERT(Eigen::Index(conditionals.size()) == size() && "give one component per mode");
        EMRYS_FILTER_ASSERT(
            conditionals.get_allocator() == m_conditionals.get_allocator()
            && "the components must come from an equal allocator"
        );
        m_stamp.touchShape();
        m_conditionals.swap(conditionals);
    }

    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(Eigen::Index element) const &
    {
        return internal::elementMarginal<MixtureMarginal, Path>(*this, element);
    }

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

    [[nodiscard]] MixtureParams<Component_, Allocator_>
    toMixture() const
    {
        auto mixture = MixtureParams<Component_, Allocator_>{space(), get_allocator()};
        mixture.reserve(size());
        for (Eigen::Index m = 0; m < size(); ++m)
        {
            mixture.add(m_modes.logWeight(m), m_conditionals[std::size_t(m)]);
        }
        return mixture;
    }

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

    template <typename R>
    [[nodiscard]] std::expected<void, NumericError>
    factorize(R const & repair)
    {
        m_stamp.touchValues();
        if (auto const total = m_modes.logTotalWeight(); !total)
        {
            return std::unexpected{total.error()};
        }
        // Every repair is tried on a copy of a component that needs one before any is kept. Repairs
        // are rare, so the copies are too; a component that factorizes is cached for the loop after.
        for (auto const & conditional : m_conditionals)
        {
            if (!conditional.factorize())
            {
                auto trial = ComponentType{conditional};
                if (auto result = trial.factorize(repair); !result)
                {
                    return std::unexpected{result.error()};
                }
            }
        }
        for (auto & conditional : m_conditionals)
        {
            if (auto result = conditional.factorize(repair); !result)
            {
                return std::unexpected{result.error()};
            }
        }
        return {};
    }

    [[nodiscard]] std::expected<DensityType, NumericError>
    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, NumericError>
    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:

    template <typename>
    friend class internal::ValuesCheck;

    template <typename>
    friend class internal::ShapeCheck;
 
    using ComponentAllocator = internal::ReBound<Allocator_, ComponentType>;  

    [[nodiscard]] ComponentType &
    conditionalToSet(Eigen::Index mode, [[maybe_unused]] ComponentType const & component)
    {
        EMRYS_FILTER_ASSERT(mode >= 0 && mode < size() && "the mode is out of range");
        EMRYS_FILTER_ASSERT(component.space() == space() && "the component must be in the hybrid's space");
        m_stamp.touchValues();
        return m_conditionals[std::size_t(mode)];
    }

    void
    check() const
    {
        EMRYS_FILTER_ASSERT(
            Eigen::Index(m_conditionals.size()) == m_modes.cardinality() && "give one component per mode"
        );
        for (auto const & conditional : m_conditionals)
        {
            EMRYS_FILTER_ASSERT(
                conditional.space() == m_conditionals.front().space()
                && "every mode's component must be on the same space"
            );
        }
    }
 
    ModesType m_modes;                
    ConditionalsType m_conditionals;  
    internal::Stamp m_stamp;          
};

template <internal::ComponentOrInfoForm Component_, DiscreteSpace ModeSpace_, typename Allocator_>
class Hybrid
{
public:
 
    using ParamsType = HybridParams<Component_, ModeSpace_, Allocator_>;  
    using ComponentType = typename ParamsType::ComponentType;             
    using Scalar = typename ParamsType::Scalar;                           
    using SpaceType = typename ParamsType::SpaceType;                     
    using ModeSpaceType = ModeSpace_;                                     
    using ModesType = typename ParamsType::ModesType;                     
    using MomentsType = typename ParamsType::MomentsType;                 
    using allocator_type = Allocator_;                                    

    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]] SpaceType const &
    space() const
    {
        return source().space();
    }

    [[nodiscard]] ModeSpace_ const &
    modeSpace() const
    {
        return source().modeSpace();
    }

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

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

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

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

    [[nodiscard]] ModesType const &
    modes() const
    {
        return source().modes();
    }

    [[nodiscard]] ComponentType const &
    conditional(Eigen::Index mode) const
    {
        return source().conditional(mode);
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        return ComponentTraits<ParamsType>::logDensity(source(), x);
    }

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

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

    [[nodiscard]] ComponentType
    collapsed() const
    {
        return ComponentTraits<ParamsType>::collapseModes(source());
    }

    template <typename P>
    [[nodiscard]] auto
    collapse() const
    {
        ParamsType const & hybrid = source();
        auto const [stride, axis] = internal::locateAxisByPath<P>(modeSpace());
        using Axis = std::remove_cvref_t<decltype(*axis)>;
        using Result = HybridParams<Component_, Axis, Allocator_>;
        using LogWeights = typename ModesType::ArrayType;
        Eigen::Index const count = axis->cardinality();
        typename Result::ConditionalsType grouped(internal::ReBound<Allocator_, ComponentType>(hybrid.get_allocator()));
        grouped.reserve(std::size_t(count));
        auto const componentAt = [&](Eigen::Index m) -> ComponentType const & {
            return hybrid.conditional(m);
        };
        for (Eigen::Index value = 0; value < count; ++value)
        {
            LogWeights logWeights = hybrid.modes().logWeights();
            for (Eigen::Index m = 0; m < size(); ++m)
            {
                if ((m / stride) % count != value)
                {
                    logWeights(m) = -Eigen::NumTraits<Scalar>::infinity();
                }
            }
            // A value of probability zero keeps the components' plain average, which is never used.
            if (!internal::logTotalWeight<Scalar>(logWeights))
            {
                logWeights = LogWeights::Constant(size(), -Eigen::NumTraits<Scalar>::infinity());
                for (Eigen::Index m = 0; m < size(); ++m)
                {
                    if ((m / stride) % count == value)
                    {
                        logWeights(m) = Scalar(0);
                    }
                }
            }
            grouped.push_back(ComponentTraits<ComponentType>::collapse(size(), componentAt, [&](Eigen::Index m) {
                return logWeights(m);
            }));
        }
        return Result{hybrid.modes().template marginal<P>(), std::move(grouped)};
    }
 
private:
 
    friend ParamsType;  

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

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

template <typename Scalar, typename ModeSpace, typename C>
    requires Component<Density<C>> && std::same_as<Scalar, typename C::Scalar>
HybridParams(CategoricalParams<ModeSpace, Scalar>, C) -> HybridParams<Density<C>, ModeSpace>;
 
template <typename Scalar, typename ModeSpace, typename C, typename Allocator>
    requires Component<Density<C>> && std::same_as<Scalar, typename C::Scalar>
HybridParams(CategoricalParams<ModeSpace, Scalar>, C, Allocator) -> HybridParams<Density<C>, ModeSpace, Allocator>;
 
template <typename Scalar, typename ModeSpace, typename C, typename ComponentAllocator>
    requires Component<Density<C>> && std::same_as<Scalar, typename C::Scalar>
HybridParams(CategoricalParams<ModeSpace, Scalar>, std::vector<C, ComponentAllocator>)
    -> HybridParams<Density<C>, ModeSpace, internal::ReBound<ComponentAllocator, std::byte>>;

template <Space Space_, DiscreteSpace ModeSpace_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
using HybridGaussian = Hybrid<Gaussian<Space_, Scalar_>, ModeSpace_, Allocator_>;

template <Space Space_, DiscreteSpace ModeSpace_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
using HybridGaussianParams = HybridParams<Gaussian<Space_, Scalar_>, ModeSpace_, Allocator_>;

template <Component Component_, DiscreteSpace ModeSpace_, typename Allocator_>
struct ComponentTraits<HybridParams<Component_, ModeSpace_, Allocator_>>
{
    using Type = HybridParams<Component_, ModeSpace_, Allocator_>;  
    using Scalar = typename Component_::Scalar;                     

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

    template <typename Derived>
    [[nodiscard]] static Scalar
    logDensity(Type const & hybrid, Eigen::MatrixBase<Derived> const & x)
    {
        auto sum = LogSumExp<Scalar>{};
        auto total = LogSumExp<Scalar>{};
        for (Eigen::Index m = 0; m < hybrid.size(); ++m)
        {
            Scalar const logWeight = hybrid.modes().logWeight(m);
            if (logWeight == -Eigen::NumTraits<Scalar>::infinity())
            {
                continue;
            }
            total.add(logWeight);
            sum.add(logWeight + ComponentTraits<Params<Component_>>::logDensity(hybrid.conditional(m), x));
        }
        return sum.value() - total.value();
    }

    template <typename ComponentAt, typename LogWeightAt>
    [[nodiscard]] static Type
    collapse(Eigen::Index count, ComponentAt componentAt, LogWeightAt logWeightAt)
    {
        EMRYS_FILTER_ASSERT(count > 0 && "a collapse needs a component");
        Type const & first = componentAt(0);
        Eigen::Index const modes = first.size();
        auto logModes = typename Type::ModesType::ArrayType{};
        logModes.resize(modes);
        for (Eigen::Index m = 0; m < modes; ++m)
        {
            auto sum = LogSumExp<Scalar>{};
            for (Eigen::Index c = 0; c < count; ++c)
            {
                EMRYS_FILTER_ASSERT(
                    componentAt(c).modeSpace() == first.modeSpace() && "every hybrid must be on one mode space"
                );
                sum.add(Scalar(logWeightAt(c)) + modeLogProbability(componentAt(c), m));
            }
            logModes(m) = sum.value();
        }
        typename Type::ConditionalsType conditionals(
            internal::ReBound<Allocator_, Params<Component_>>(first.get_allocator())
        );
        conditionals.reserve(std::size_t(modes));
        for (Eigen::Index m = 0; m < modes; ++m)
        {
            if (logModes(m) == -Eigen::NumTraits<Scalar>::infinity())
            {
                conditionals.push_back(first.conditional(m));
                continue;
            }
            conditionals.push_back(
                ComponentTraits<Params<Component_>>::collapse(
                    count,
                    [&](Eigen::Index c) -> Params<Component_> const & {
                        return componentAt(c).conditional(m);
                    },
                    [&](Eigen::Index c) {
                        return Scalar(logWeightAt(c)) + modeLogProbability(componentAt(c), m);
                    }
                )
            );
        }
        auto collapsedModes = typename Type::ModesType{first.modeSpace(), logModes};
        // Weights that fail to normalize are kept as they are, for the next factorize() to report.
        static_cast<void>(collapsedModes.normalize());
        return Type{std::move(collapsedModes), std::move(conditionals)};
    }

    [[nodiscard]] static Params<Component_>
    collapseModes(Type const & hybrid)
    {
        return ComponentTraits<Params<Component_>>::collapse(
            hybrid.size(),
            [&](Eigen::Index m) -> Params<Component_> const & {
                return hybrid.conditional(m);
            },
            [&](Eigen::Index m) {
                return hybrid.modes().logWeight(m);
            }
        );
    }

    [[nodiscard]] static typename Type::MomentsType
    moments(Type const & hybrid)
    {
        return ComponentTraits<Params<Component_>>::moments(collapseModes(hybrid));
    }

    [[nodiscard]] static Scalar
    modeLogProbability(Type const & hybrid, Eigen::Index m)
    {
        auto const total = hybrid.modes().logTotalWeight();
        EMRYS_FILTER_ASSERT(total && "a hybrid's mode weights must define a distribution to collapse it");
        return Scalar(hybrid.modes().logWeight(m) - *total);
    }
};
 
namespace internal
{

template <Component Component_, DiscreteSpace ModeSpace_, typename Allocator_>
struct ElementsOf<HybridParams<Component_, ModeSpace_, Allocator_>>
{
    using Type = HybridParams<Component_, ModeSpace_, Allocator_>;  

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

    [[nodiscard]] static auto
    logWeight(Type const & hybrid, Eigen::Index m)
    {
        return hybrid.modes().logWeight(m);
    }

    [[nodiscard]] static auto const &
    component(Type const & hybrid, Eigen::Index m)
    {
        return hybrid.conditional(m);
    }

    template <typename Marginal>
    [[nodiscard]] static auto
    copy(Type const & hybrid, Marginal const & marginal, Allocator_ const & allocator)
    {
        using Result = CopyType<typename Marginal::ComponentMarginalType::ParamsType>;
        auto conditionals =
            typename Result::ConditionalsType(ReBound<Allocator_, typename Result::ComponentType>(allocator));
        conditionals.reserve(std::size_t(marginal.size()));
        for (auto const & element : marginal)
        {
            conditionals.push_back(element.component.copy());
        }
        return Result{hybrid.modes(), std::move(conditionals)};
    }
};
 
}  // namespace internal
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Component_, typename ModeSpace_, typename Allocator_>
void
describe(std::ostream & stream, HybridParams<Component_, ModeSpace_, Allocator_> const & hybrid, Indent indent)
{
    stream << indent << "hybrid of ";
    printCount(stream, hybrid.size(), "mode");
    printDimension(stream, hybrid.valueDimension(), hybrid.differenceDimension());
    auto const total = hybrid.modes().logTotalWeight();
    if (total)
    {
        printMoments(
            stream,
            indent.deeper(),
            ComponentTraits<HybridParams<Component_, ModeSpace_, Allocator_>>::moments(hybrid)
        );
    }
    else
    {
        stream << '\n';
        describe(stream, total.error(), indent.deeper());
    }
    for (Eigen::Index mode = 0; mode < hybrid.size(); ++mode)
    {
        auto const logWeight = hybrid.modes().logWeight(mode);
        stream << '\n' << indent.deeper() << "mode " << mode;
        if (total)
        {
            stream << ", probability " << Eigen::numext::exp(logWeight - *total) << ":\n";
        }
        else
        {
            stream << ", log-weight " << logWeight << ":\n";
        }
        describe(stream, hybrid.conditional(mode), indent.deeper().deeper());
    }
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO