Categorical.h

Go to the documentation of this file.

#pragma once

 
#include <array>
#include <concepts>
#include <cstddef>
#include <expected>
#include <random>
#include <utility>
#include <vector>
 
#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/Stamp.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Params.h>
#include <EmrysFilter/Spaces/Discrete.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{
 
template <DiscreteSpace Space_, RealScalar Scalar_>
class CategoricalParams;
 
template <DiscreteSpace Space_, RealScalar Scalar_>
class Categorical;

enum class ChainError
{
    Reducible,  
};
 
namespace internal
{

template <typename Scalar, typename Matrix, typename DwellTimes, typename ExitWeights>
[[nodiscard]] Matrix
dwellGenerator(Eigen::MatrixBase<DwellTimes> const & dwellTimes, Eigen::MatrixBase<ExitWeights> const & exitWeights)
{
    Eigen::Index const n = exitWeights.rows();
    EMRYS_FILTER_ASSERT(
        dwellTimes.size() == n && exitWeights.cols() == n && "give a dwell time and a row of exit weights per value"
    );
    auto generator = Matrix{};
    generator.resize(n, n);
    for (Eigen::Index i = 0; i < n; ++i)
    {
        Scalar const dwell = Scalar(dwellTimes(i));
        EMRYS_FILTER_ASSERT(dwell > Scalar(0) && "a dwell time must be positive");
        EMRYS_FILTER_ASSERT(exitWeights(i, i) == 0 && "a value's weight of exiting to itself must be zero");
        EMRYS_FILTER_ASSERT((exitWeights.row(i).array() >= 0).all() && "exit weights must be nonnegative");
        Scalar const leave = Scalar(1) / dwell;
        Scalar const total = Scalar(exitWeights.row(i).sum());
        EMRYS_FILTER_ASSERT(
            (leave == Scalar(0) || total > Scalar(0)) && "a value that is left needs a value to exit to"
        );
        if (leave == Scalar(0))
        {
            generator.row(i).setZero();
        }
        else
        {
            generator.row(i) = exitWeights.row(i).template cast<Scalar>() * Scalar(leave / total);
        }
        generator(i, i) = -leave;
    }
    return generator;
}

template <typename Scalar, typename Array, typename Matrix>
[[nodiscard]] std::expected<Array, ChainError>
stationaryProbabilities(Matrix kernel)
{
    Eigen::Index const n = kernel.rows();
    EMRYS_FILTER_ASSERT(n > 0 && "a chain has at least one value");
    for (Eigen::Index k = n - 1; k > 0; --k)
    {
        // Watched only at the values before k, the chain moves from i to j directly or through
        // k, which it leaves for an earlier value with the total of row k before it.
        Scalar const leave = kernel.row(k).head(k).sum();
        if (!(leave > Scalar(0)))
        {
            return std::unexpected{ChainError::Reducible};
        }
        kernel.col(k).head(k) /= leave;
        kernel.topLeftCorner(k, k).noalias() += kernel.col(k).head(k) * kernel.row(k).head(k);
    }
    Array probabilities = Array::Zero(n);
    probabilities(0) = Scalar(1);
    for (Eigen::Index k = 1; k < n; ++k)
    {
        Scalar const flow = probabilities.head(k).matrix().dot(kernel.col(k).head(k));
        if (!(flow > Scalar(0)))
        {
            return std::unexpected{ChainError::Reducible};
        }
        probabilities(k) = flow;
    }
    probabilities /= probabilities.sum();
    return probabilities;
}

template <typename P, typename Scalar, DiscreteSpace Space, typename Values>
[[nodiscard]] auto
categoricalMarginal(Space const & space, Values const & logValues)
{
    auto const [stride, axis] = locateAxisByPath<P>(space);
    using Axis = std::remove_cvref_t<decltype(*axis)>;
    using AxisArray = typename CategoricalParams<Axis, Scalar>::ArrayType;
    Eigen::Index const count = axis->cardinality();
    // A LogSumExp per value of the axis, kept as its running maximum and sum in the axis's own
    // storage, so a bounded axis allocates nothing.
    AxisArray maxima = AxisArray::Constant(count, -Eigen::NumTraits<Scalar>::infinity());
    AxisArray sums = AxisArray::Zero(count);
    for (Eigen::Index v = 0; v < logValues.size(); ++v)
    {
        Eigen::Index const i = (v / stride) % count;
        Scalar const x = logValues(v);
        if (x == -Eigen::NumTraits<Scalar>::infinity())
        {
            continue;
        }
        if (x <= maxima(i))
        {
            sums(i) += Eigen::numext::exp(x - maxima(i));
        }
        else
        {
            sums(i) = sums(i) * Eigen::numext::exp(maxima(i) - x) + Scalar(1);
            maxima(i) = x;
        }
    }
    for (Eigen::Index i = 0; i < count; ++i)
    {
        maxima(i) =
            sums(i) == Scalar(0) ? -Eigen::NumTraits<Scalar>::infinity() : maxima(i) + Eigen::numext::log(sums(i));
    }
    return CategoricalParams<Axis, Scalar>{*axis, maxima};
}

template <typename Scalar, DiscreteSpace Rest, typename Values>
[[nodiscard]] CategoricalParams<Rest, Scalar>
categoricalConditional(Rest rest, ValueLocation const & location, Values const & logValues)
{
    using RestArray = typename CategoricalParams<Rest, Scalar>::ArrayType;
    auto selected = RestArray{};
    selected.resize(rest.cardinality());
    Eigen::Index next = 0;
    for (Eigen::Index v = 0; v < logValues.size(); ++v)
    {
        if (location.matches(v))
        {
            selected(next++) = logValues(v);
        }
    }
    EMRYS_FILTER_ASSERT(next == selected.size());
    return CategoricalParams<Rest, Scalar>{std::move(rest), selected};
}

template <typename Scalar, typename Array, typename Kernel>
[[nodiscard]] Array
movedLogWeights(Array const & logWeights, Kernel const & kernel)
{
    // Every discrete space has a value, so there is a largest log-weight.
    Scalar const largest = logWeights.maxCoeff();
    Scalar const shift = isFinite(largest) ? largest : Scalar(0);
    Array const weights = (logWeights - shift).exp();
    Array moved = multiply(kernel.transpose(), weights.matrix()).array();
    return moved.log() + shift;
}
 
}  // namespace internal

template <DiscreteSpace Space_, RealScalar Scalar_ = double>
class CategoricalParams
{
public:
 
    using Scalar = Scalar_;                                               
    using SpaceType = Space_;                                             
    using DensityType = Categorical<Space_, Scalar_>;                     
    static constexpr int CardinalityAtCompileTime = Space_::Cardinality;  
    static constexpr int MaxCardinalityAtCompileTime = Space_::MaxCardinality;  
    static constexpr bool IsParams = true;                                      

    using ArrayType =
        Eigen::Array<Scalar_, CardinalityAtCompileTime, 1, Eigen::ColMajor, MaxCardinalityAtCompileTime, 1>;

    explicit CategoricalParams(Space_ space = {})
        : m_space{std::move(space)}
        , m_logWeights{ArrayType::Constant(m_space.cardinality(), -Eigen::numext::log(Scalar(m_space.cardinality())))}
    {}

    template <typename Derived>
    CategoricalParams(Space_ space, Eigen::DenseBase<Derived> const & logWeights)
        : m_space{std::move(space)}
        , m_logWeights{logWeights.derived().array().template cast<Scalar>()}
    {
        EMRYS_FILTER_ASSERT(m_logWeights.size() == m_space.cardinality() && "give one log-weight per value");
    }

    template <typename Derived>
    [[nodiscard]] static CategoricalParams
    fromWeights(Space_ space, Eigen::DenseBase<Derived> const & weights)
    {
        return CategoricalParams{std::move(space), weights.derived().array().template cast<Scalar>().log()};
    }

    template <typename Derived>
        requires FixedSizeSpace<Space_>
    [[nodiscard]] static CategoricalParams
    fromWeights(Eigen::DenseBase<Derived> const & weights)
    {
        return fromWeights(Space_{}, weights);
    }

    template <typename... Weights>
        requires FixedSizeSpace<Space_> && (CardinalityAtCompileTime != Eigen::Dynamic) && (sizeof...(Weights) > 0)
              && (std::convertible_to<Weights, Scalar_> && ...)
    [[nodiscard]] static CategoricalParams
    fromWeights(Weights... weights)
    {
        static_assert(sizeof...(Weights) == CardinalityAtCompileTime, "give one weight per value");
        auto array = ArrayType{};
        Eigen::Index i = 0;
        ((array(i++) = Scalar_(weights)), ...);
        return fromWeights(Space_{}, array);
    }

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

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

    [[nodiscard]] ArrayType const &
    logWeights() const
    {
        return m_logWeights;
    }

    [[nodiscard]] Scalar
    logWeight(Eigen::Index value) const
    {
        EMRYS_FILTER_ASSERT(value >= 0 && value < cardinality() && "the value is out of range");
        return m_logWeights(value);
    }

    [[nodiscard]] std::expected<Scalar, WeightError>
    logTotalWeight() const
    {
        return internal::logTotalWeight<Scalar>(m_logWeights);
    }

    template <typename Derived>
    void
    setLogWeights(Eigen::DenseBase<Derived> const & logWeights)
    {
        EMRYS_FILTER_ASSERT(logWeights.size() == cardinality() && "give one log-weight per value");
        m_stamp.touchValues();
        m_logWeights = logWeights.derived().array().template cast<Scalar>();
    }

    template <typename Derived>
    void
    setWeights(Eigen::DenseBase<Derived> const & weights)
    {
        setLogWeights(weights.derived().array().template cast<Scalar>().log());
    }

    void
    setLogWeight(Eigen::Index value, Scalar logWeight)
    {
        EMRYS_FILTER_ASSERT(value >= 0 && value < cardinality() && "the value is out of range");
        m_stamp.touchValues();
        m_logWeights(value) = logWeight;
    }

    void
    setWeight(Eigen::Index value, Scalar weight)
    {
        setLogWeight(value, Scalar(Eigen::numext::log(weight)));
    }

    template <typename... Values>
        requires(sizeof...(Values) > 0)
    void
    setLogWeight(Scalar logWeight)
    {
        setLogWeight(valueOf<Values...>(m_space), logWeight);
    }

    template <typename... Values>
        requires(sizeof...(Values) > 0)
    void
    setWeight(Scalar weight)
    {
        setWeight(valueOf<Values...>(m_space), weight);
    }

    template <typename Derived>
    void
    reweight(Eigen::DenseBase<Derived> const & logLikelihoods)
    {
        EMRYS_FILTER_ASSERT(logLikelihoods.size() == cardinality() && "give one log-likelihood per value");
        m_stamp.touchValues();
        m_logWeights += logLikelihoods.derived().array().template cast<Scalar>();
    }

    template <typename Derived>
    [[nodiscard]] std::expected<Scalar, WeightError>
    update(Eigen::DenseBase<Derived> const & logLikelihoods)
    {
        EMRYS_FILTER_ASSERT(logLikelihoods.size() == cardinality() && "give one log-likelihood per value");
        auto const prior = logTotalWeight();
        if (!prior)
        {
            return std::unexpected{prior.error()};
        }
        ArrayType posterior = m_logWeights + logLikelihoods.derived().array().template cast<Scalar>();
        auto const total = internal::logTotalWeight<Scalar>(posterior);
        if (!total)
        {
            return std::unexpected{total.error()};
        }
        m_stamp.touchValues();
        m_logWeights = posterior - *total;
        return *total - *prior;
    }

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

    template <typename P>
    [[nodiscard]] auto
    marginal() const
    {
        return internal::categoricalMarginal<P, Scalar>(m_space, m_logWeights);
    }

    template <typename P>
    [[nodiscard]] auto
    conditional(Eigen::Index value) const
    {
        auto const [stride, axis] = internal::locateAxisByPath<P>(m_space);
        EMRYS_FILTER_ASSERT(value >= 0 && value < axis->cardinality() && "the value is out of range");
        return internal::categoricalConditional<Scalar>(
            internal::complementAxes<P>(m_space),
            internal::ValueLocation{stride, axis->cardinality(), value},
            m_logWeights
        );
    }

    template <typename Value>
    [[nodiscard]] auto
    conditional() const
    {
        return internal::categoricalConditional<Scalar>(
            internal::complementAxes<internal::ValueAxisPath<Value, Space_>>(m_space),
            internal::locateValuePath<Value>(m_space),
            m_logWeights
        );
    }

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

    void
    toDensity() const && = delete;
 
private:
 
    friend DensityType;  

    template <typename>
    friend class internal::ValuesCheck;
 
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Space_ m_space;  
    ArrayType m_logWeights;                         
    internal::Stamp m_stamp;                        
};

template <DiscreteSpace Space, typename Derived>
CategoricalParams(Space, Eigen::DenseBase<Derived> const &) -> CategoricalParams<Space, typename Derived::Scalar>;

 
namespace internal
{

template <typename Space, typename Scalar>
[[nodiscard]] auto
logProbabilitiesOf(CategoricalParams<Space, Scalar> const & params)
{
    return params.logWeights() - params.logTotalWeight().value_or(Eigen::NumTraits<Scalar>::quiet_NaN());
}
 
}  // namespace internal

template <DiscreteSpace Space_, RealScalar Scalar_ = double>
class Categorical
{
public:
 
    using ParamsType = CategoricalParams<Space_, Scalar_>;  
    using Scalar = Scalar_;                                 
    using SpaceType = Space_;                               
    using ArrayType = typename ParamsType::ArrayType;       
    static constexpr int CardinalityAtCompileTime =
        ParamsType::CardinalityAtCompileTime;  
    static constexpr int MaxCardinalityAtCompileTime =
        ParamsType::MaxCardinalityAtCompileTime;  
    static constexpr bool IsParams = false;       

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

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

    [[nodiscard]] auto
    logProbabilities() const
    {
        return source().m_logWeights - m_logTotal;
    }

    [[nodiscard]] auto
    probabilities() const
    {
        return logProbabilities().exp();
    }

    [[nodiscard]] Scalar
    logProbability(Eigen::Index value) const
    {
        return source().logWeight(value) - m_logTotal;
    }

    [[nodiscard]] Scalar
    probability(Eigen::Index value) const
    {
        return Eigen::numext::exp(logProbability(value));
    }

    template <typename... Values>
        requires(sizeof...(Values) > 0)
    [[nodiscard]] Scalar
    logProbability() const
    {
        ParamsType const & params = source();
        auto const locations = std::array<internal::ValueLocation, sizeof...(Values)>{
            internal::locateValuePath<Values>(params.m_space)...
        };
        auto total = LogSumExp<Scalar>{};
        for (Eigen::Index v = 0; v < params.cardinality(); ++v)
        {
            bool matches = true;
            for (auto const & location : locations)
            {
                matches = matches && location.matches(v);
            }
            if (matches)
            {
                total.add(params.m_logWeights(v));
            }
        }
        return total.value() - m_logTotal;
    }

    template <typename... Values>
        requires(sizeof...(Values) > 0)
    [[nodiscard]] Scalar
    probability() const
    {
        return Eigen::numext::exp(logProbability<Values...>());
    }

    [[nodiscard]] Scalar
    entropy() const
    {
        ArrayType const & logWeights = source().m_logWeights;
        Scalar result = 0;
        for (Eigen::Index v = 0; v < logWeights.size(); ++v)
        {
            if (logWeights(v) != -Eigen::NumTraits<Scalar>::infinity())
            {
                Scalar const logProbability = logWeights(v) - m_logTotal;
                result -= Eigen::numext::exp(logProbability) * logProbability;
            }
        }
        return Eigen::numext::maxi(result, Scalar(0));
    }

    [[nodiscard]] Eigen::Index
    mostProbable() const
    {
        ArrayType const & logWeights = source().m_logWeights;
        Eigen::Index best = 0;
        for (Eigen::Index v = 1; v < logWeights.size(); ++v)
        {
            if (logWeights(v) > logWeights(best))
            {
                best = v;
            }
        }
        return best;
    }

    template <typename P>
    [[nodiscard]] auto
    marginal() const
    {
        return internal::categoricalMarginal<P, Scalar>(space(), logProbabilities());
    }

    template <typename P>
    [[nodiscard]] auto
    conditional(Eigen::Index value) const
    {
        return source().template conditional<P>(value);
    }

    template <typename Value>
    [[nodiscard]] auto
    conditional() const
    {
        return source().template conditional<Value>();
    }
 
private:
 
    friend ParamsType;  

    Categorical(ParamsType const & source, Scalar logTotal)
        : m_check{source}
        , m_logTotal{logTotal}
    {}

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

template <DiscreteSpace Space_, RealScalar Scalar_ = double>
class MarkovChain
{
public:
 
    using Scalar = Scalar_;                                               
    using SpaceType = Space_;                                             
    static constexpr int CardinalityAtCompileTime = Space_::Cardinality;  
    static constexpr int MaxCardinalityAtCompileTime = Space_::MaxCardinality;  

    using MatrixType = Eigen::Matrix<
        Scalar_,
        CardinalityAtCompileTime,
        CardinalityAtCompileTime,
        Eigen::RowMajor,
        MaxCardinalityAtCompileTime,
        MaxCardinalityAtCompileTime>;

    explicit MarkovChain(Space_ space = {})
        : m_space{std::move(space)}
        , m_transitions{MatrixType::Identity(m_space.cardinality(), m_space.cardinality())}
    {}

    template <MatrixExpression Transitions>
    MarkovChain(Space_ space, Transitions && transitions)
        : m_space{std::move(space)}
        , m_transitions{internal::castScalar<Scalar>(std::forward<Transitions>(transitions))}
    {
        EMRYS_FILTER_ASSERT(
            m_transitions.rows() == m_space.cardinality() && m_transitions.cols() == m_space.cardinality()
            && "the transition matrix has a row and a column per value"
        );
        EMRYS_FILTER_ASSERT(
            (m_transitions.array() >= Scalar(0)).all() && "transition probabilities must be nonnegative"
        );
        EMRYS_FILTER_ASSERT(
            ((m_transitions.rowwise().sum().array() - Scalar(1)).abs() <= Eigen::NumTraits<Scalar>::dummy_precision())
                .all()
            && "each row of the transition matrix must sum to one"
        );
    }

    template <MatrixExpression Transitions>
        requires FixedSizeSpace<Space_>
    explicit MarkovChain(Transitions && transitions)
        : MarkovChain{Space_{}, std::forward<Transitions>(transitions)}
    {}

    template <DiscreteSpace... Axes>
        requires(
            internal::IsDiscreteProduct<Space_>
            && sizeof...(Axes) == std::tuple_size_v<internal::DiscreteProductAxes<Space_>>
        )
    MarkovChain(Space_ space, MarkovChain<Axes, Scalar_> const &... chains)
        : m_space{std::move(space)}
    {
        static_assert(
            std::is_same_v<std::tuple<Axes...>, internal::DiscreteProductAxes<Space_>>,
            "give one chain per axis of the product, in order"
        );
        Eigen::Index const n = m_space.cardinality();
        m_transitions = MatrixType::Ones(n, n);
        Eigen::Index stride = n;
        (
            [&] {
                Eigen::Index const count = chains.space().cardinality();
                stride /= count;
                for (Eigen::Index i = 0; i < n; ++i)
                {
                    for (Eigen::Index j = 0; j < n; ++j)
                    {
                        m_transitions(i, j) *= chains.matrix()((i / stride) % count, (j / stride) % count);
                    }
                }
            }(),
            ...
        );
        EMRYS_FILTER_ASSERT(stride == 1 && "each chain must have its axis's cardinality");
    }

    template <DiscreteSpace... Axes>
        requires(
            FixedSizeSpace<Space_> && internal::IsDiscreteProduct<Space_>
            && sizeof...(Axes) == std::tuple_size_v<internal::DiscreteProductAxes<Space_>>
        )
    explicit MarkovChain(MarkovChain<Axes, Scalar_> const &... chains)
        : MarkovChain{Space_{}, chains...}
    {}

    [[nodiscard]] static MarkovChain
    persistent(Scalar stay, Space_ space = {})
    {
        EMRYS_FILTER_ASSERT(stay >= Scalar(0) && stay <= Scalar(1) && "the probability must be in [0, 1]");
        Eigen::Index const n = space.cardinality();
        auto result = MarkovChain{std::move(space)};
        if (n > 1)
        {
            result.m_transitions.setConstant((Scalar(1) - stay) / Scalar(n - 1));
            result.m_transitions.diagonal().setConstant(stay);
        }
        return result;
    }

    [[nodiscard]] static MarkovChain
    memoryless(Categorical<Space_, Scalar_> const & next)
    {
        auto result = MarkovChain{next.space()};
        result.m_transitions = next.probabilities().matrix().transpose().replicate(next.cardinality(), 1);
        return result;
    }

    template <typename DwellTimes, typename ExitWeights>
    [[nodiscard]] static MarkovChain
    fromDwellTimes(
        Eigen::MatrixBase<DwellTimes> const & dwellTimes,
        Eigen::MatrixBase<ExitWeights> const & exitWeights,
        Space_ space = {}
    )
    {
        EMRYS_FILTER_ASSERT((dwellTimes.array() >= 1).all() && "a dwell time is at least one step");
        MatrixType transitions = internal::dwellGenerator<Scalar, MatrixType>(dwellTimes, exitWeights);
        transitions.diagonal().array() += Scalar(1);
        EMRYS_FILTER_ASSERT(transitions.rows() == space.cardinality() && "give a dwell time per value");
        return MarkovChain{std::move(space), std::move(transitions)};
    }

    template <typename DwellTimes>
    [[nodiscard]] static MarkovChain
    fromDwellTimes(Eigen::MatrixBase<DwellTimes> const & dwellTimes, Space_ space = {})
    {
        Eigen::Index const n = space.cardinality();
        return fromDwellTimes(dwellTimes, MatrixType::Ones(n, n) - MatrixType::Identity(n, n), std::move(space));
    }

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

    [[nodiscard]] MatrixType const &
    matrix() const
    {
        return m_transitions;
    }

    [[nodiscard]] Scalar
    transition(Eigen::Index from, Eigen::Index to) const
    {
        return m_transitions(from, to);
    }

    [[nodiscard]] CategoricalParams<Space_, Scalar_>
    next(Eigen::Index from) const
    {
        EMRYS_FILTER_ASSERT(from >= 0 && from < m_space.cardinality() && "the value is not in the chain's space");
        return CategoricalParams<Space_, Scalar_>{m_space, m_transitions.row(from).transpose().array().log()};
    }
 
private:
 
    Space_ m_space;            
    MatrixType m_transitions;  
};

template <DiscreteSpace Space, MatrixExpression Transitions>
MarkovChain(Space, Transitions &&) -> MarkovChain<Space, typename std::remove_cvref_t<Transitions>::Scalar>;


template <typename Scalar, typename Space>
[[nodiscard]] CategoricalParams<Space, Scalar>
predict(CategoricalParams<Space, Scalar> const & prior, MarkovChain<Space, Scalar> const & chain)
{
    EMRYS_FILTER_ASSERT(prior.space() == chain.space() && "the chain is on a different space");
    return CategoricalParams<Space, Scalar>{
        prior.space(),
        internal::movedLogWeights<Scalar>(prior.logWeights(), chain.matrix())
    };
}

template <typename Scalar, typename Space>
[[nodiscard]] std::expected<CategoricalParams<Space, Scalar>, ChainError>
stationary(MarkovChain<Space, Scalar> const & chain)
{
    using ArrayType = typename CategoricalParams<Space, Scalar>::ArrayType;
    std::expected<ArrayType, ChainError> const probabilities =
        internal::stationaryProbabilities<Scalar, ArrayType>(chain.matrix());
    if (!probabilities)
    {
        return std::unexpected{probabilities.error()};
    }
    return CategoricalParams<Space, Scalar>{chain.space(), probabilities->log()};
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] std::vector<Eigen::Index>
drawSamples(Categorical<Space, Scalar> const & distribution, Eigen::Index n, G & rng)
{
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    EMRYS_FILTER_ASSERT(n >= 0);
    typename Categorical<Space, Scalar>::ArrayType const probabilities = distribution.probabilities();
    auto draw =
        std::discrete_distribution<Eigen::Index>{probabilities.data(), probabilities.data() + probabilities.size()};
    std::size_t const count = std::size_t(n);
    auto samples = std::vector<Eigen::Index>(count);
    for (auto & sample : samples)
    {
        sample = draw(rng);
    }
    return samples;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] Eigen::Index
drawSample(Categorical<Space, Scalar> const & distribution, G & rng)
{
    auto const probabilities = distribution.probabilities();
    return internal::drawIndex<Scalar>(
        probabilities.size(),
        [&](Eigen::Index i) {
            return probabilities(i);
        },
        rng
    );
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

template <typename Space_, typename Scalar_>
void
describe(std::ostream & stream, Categorical<Space_, Scalar_> const & categorical, Indent indent)
{
    typename Categorical<Space_, Scalar_>::ArrayType const probabilities = categorical.probabilities();
    stream << indent << "categorical over ";
    printCount(stream, probabilities.size(), "value");
    stream << '\n';
    printRow(stream, indent.deeper(), "probabilities", probabilities);
}

template <typename Space_, typename Scalar_>
void
describe(std::ostream & stream, CategoricalParams<Space_, Scalar_> const & params, Indent indent)
{
    if (auto const density = params.toDensity())
    {
        describe(stream, *density, indent);
        return;
    }
    stream << indent << "categorical over ";
    printCount(stream, params.cardinality(), "value");
    stream << ", whose weights are not normalizable\n";
    printRow(stream, indent.deeper(), "log-weights", params.logWeights());
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO