Emission.h

Go to the documentation of this file.

#pragma once

 
#include <concepts>
#include <expected>
#include <type_traits>
#include <utility>
 
#include <EmrysFilter/Core/Forwarding.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Time.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Categorical.h>
#include <EmrysFilter/Distributions/ContinuousTimeMarkovChain.h>
#include <EmrysFilter/Spaces/Discrete.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

template <DiscreteSpace Space, typename Scalar>
using EmissionBuffer = Eigen::Ref<Eigen::Array<Scalar, Space::Cardinality, 1>>;

template <typename E, typename Space>
concept EmissionModel =
    DiscreteSpace<Space>
    && requires(E const & emission, Space const & space, EmissionBuffer<Space, typename E::Scalar> out) {
           requires RealScalar<typename E::Scalar>;
           emission.logLikelihoods(space, out);
       };

template <DiscreteSpace Space, RealScalar Scalar>
struct CategoricalUpdate
{
    CategoricalParams<Space, Scalar> posterior;  
    Scalar logEvidence;                          
};
 
namespace internal
{

template <typename Scalar, typename Space, EmissionModel<Space> E>
[[nodiscard]] typename CategoricalParams<Space, typename E::Scalar>::ArrayType
emissionLogLikelihoods(CategoricalParams<Space, Scalar> const & state, E const & emission)
{
    using Buffer = typename CategoricalParams<Space, typename E::Scalar>::ArrayType;
    auto logLikelihoods = Buffer(state.cardinality());
    emission.logLikelihoods(state.space(), logLikelihoods);
    return logLikelihoods;
}
 
}  // namespace internal

template <typename Scalar, typename Space, EmissionModel<Space> E>
[[nodiscard]] std::expected<CategoricalUpdate<Space, Scalar>, WeightError>
update(CategoricalParams<Space, Scalar> const & prior, E const & emission)
{
    auto result = CategoricalUpdate<Space, Scalar>{prior, Scalar(0)};
    auto const evidence = result.posterior.update(internal::emissionLogLikelihoods(prior, emission));
    if (!evidence)
    {
        return std::unexpected{evidence.error()};
    }
    result.logEvidence = *evidence;
    return result;
}

template <DiscreteSpace Space_, DiscreteSpace ObservationSpace_, RealScalar Scalar_ = double>
class EmissionTable
{
public:
 
    using Scalar = Scalar_;                          
    using SpaceType = Space_;                        
    using ObservationSpaceType = ObservationSpace_;  

    using MatrixType = Eigen::Matrix<
        Scalar_,
        Space_::Cardinality,
        ObservationSpace_::Cardinality,
        Eigen::RowMajor,
        Space_::MaxCardinality,
        ObservationSpace_::MaxCardinality>;

    class Observation
    {
    public:
 
        using Scalar = Scalar_;  

        Observation(EmissionTable const & table, Eigen::Index observed)
            : m_table{&table}
            , m_observed{observed}
        {
            EMRYS_FILTER_ASSERT(
                observed >= 0 && observed < table.m_matrix.cols() && "the observed value is out of range"
            );
        }

        [[nodiscard]] Eigen::Index
        observed() const
        {
            return m_observed;
        }

        void
        logLikelihoods(Space_ const & space, EmissionBuffer<Space_, Scalar_> out) const
        {
            EMRYS_FILTER_ASSERT(space == m_table->m_space && "the table is over a different space");
            EMRYS_FILTER_ASSERT(out.size() == m_table->m_matrix.rows());
            out = m_table->m_matrix.col(m_observed).array().log();
        }
 
    private:
 
        EmissionTable const * m_table;  
        Eigen::Index m_observed;        
    };

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

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

    [[nodiscard]] ObservationSpace_ const &
    observationSpace() const
    {
        return m_observations;
    }

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

    [[nodiscard]] Observation
    observe(Eigen::Index observed) const &
    {
        return Observation{*this, observed};
    }

    void
    observe(Eigen::Index observed) const && = delete;
 
private:
 
    Space_ m_space;                    
    ObservationSpace_ m_observations;  
    MatrixType m_matrix;               
};

template <DiscreteSpace Space, DiscreteSpace ObservationSpace, MatrixExpression Matrix>
EmissionTable(Space, ObservationSpace, Matrix &&)
    -> EmissionTable<Space, ObservationSpace, typename std::remove_cvref_t<Matrix>::Scalar>;


template <typename Observation, typename DistributionOf>
class DensityEmission
{
public:

    using Scalar = std::remove_cvref_t<decltype(std::declval<DistributionOf const &>()(Eigen::Index{})
                                                    .logDensity(std::declval<Observation const &>()))>;

    DensityEmission(Observation observation, DistributionOf distributionOf)
        : m_observation{std::move(observation)}
        , m_distributionOf{std::move(distributionOf)}
    {}

    [[nodiscard]] Observation const &
    observation() const
    {
        return m_observation;
    }

    template <DiscreteSpace Space>
    void
    logLikelihoods(Space const & space, EmissionBuffer<Space, Scalar> out) const
    {
        EMRYS_FILTER_ASSERT(out.size() == space.cardinality());
        for (Eigen::Index v = 0; v < out.size(); ++v)
        {
            out(v) = m_distributionOf(v).logDensity(m_observation);
        }
    }
 
private:
 
    Observation m_observation;        
    DistributionOf m_distributionOf;  
};
 
namespace internal
{

template <DiscreteSpace Space, RealScalar Scalar>
class ValueRates
{
public:

    using ArrayType = typename CategoricalParams<Space, Scalar>::ArrayType;

    template <typename Derived>
    explicit ValueRates(Eigen::DenseBase<Derived> const & rates)
        : m_rates{rates.derived().array().template cast<Scalar>()}
    {
        EMRYS_FILTER_ASSERT((m_rates >= Scalar(0)).all() && "rates must be nonnegative");
    }

    [[nodiscard]] ArrayType const &
    rates() const
    {
        return m_rates;
    }
 
private:
 
    ArrayType m_rates;  
};
 
}  // namespace internal

template <DiscreteSpace Space, RealScalar Scalar_ = double>
class PoissonEmission : public internal::ValueRates<Space, Scalar_>
{
public:
 
    using Scalar = Scalar_;  

    template <typename Derived, std::integral Count>
        requires(!std::same_as<Count, bool>)
    PoissonEmission(Eigen::DenseBase<Derived> const & rates, Count count, TimeStep<Scalar_> auto const & dt)
        : internal::ValueRates<Space, Scalar_>{rates}
        , m_count{Eigen::Index(count)}
        , m_dt{internal::toScalarTime<Scalar_>(dt)}
    {
        EMRYS_FILTER_ASSERT(count >= 0 && "the count must be nonnegative");
        EMRYS_FILTER_ASSERT(m_dt > 0 && "the interval must be positive");
    }

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

    [[nodiscard]] Scalar_
    interval() const
    {
        return m_dt;
    }

    void
    logLikelihoods(Space const & space, EmissionBuffer<Space, Scalar> out) const
    {
        auto const & rates = this->rates();
        EMRYS_FILTER_ASSERT(out.size() == space.cardinality() && rates.size() == out.size());
        Scalar const n = Scalar(m_count);
        Scalar const logFactorial = logGamma<Scalar>(n + Scalar(1));
        for (Eigen::Index v = 0; v < out.size(); ++v)
        {
            Scalar const mean = rates(v) * m_dt;
            if (mean == Scalar(0))
            {
                out(v) = m_count == 0 ? Scalar(0) : -Eigen::NumTraits<Scalar>::infinity();
            }
            else
            {
                out(v) = n * Eigen::numext::log(mean) - mean - logFactorial;
            }
        }
    }
 
private:
 
    Eigen::Index m_count;  
    Scalar_ m_dt;          
};

template <DiscreteSpace Space, RealScalar Scalar_ = double>
class ExponentialEmission : public internal::ValueRates<Space, Scalar_>
{
public:
 
    using Scalar = Scalar_;  

    template <typename Derived>
    ExponentialEmission(Eigen::DenseBase<Derived> const & rates, TimeStep<Scalar_> auto const & tau)
        : internal::ValueRates<Space, Scalar_>{rates}
        , m_tau{internal::toScalarTime<Scalar_>(tau)}
    {}

    [[nodiscard]] Scalar_
    interval() const
    {
        return m_tau;
    }

    void
    logLikelihoods(Space const & space, EmissionBuffer<Space, Scalar> out) const
    {
        auto const & rates = this->rates();
        EMRYS_FILTER_ASSERT(out.size() == space.cardinality() && rates.size() == out.size());
        for (Eigen::Index v = 0; v < out.size(); ++v)
        {
            out(v) = Eigen::numext::log(rates(v)) - rates(v) * m_tau;
        }
    }
 
private:
 
    Scalar_ m_tau;  
};

template <DiscreteSpace Space, RealScalar Scalar_ = double>
class ArrivalInterval : public internal::ValueRates<Space, Scalar_>
{
public:
 
    using Scalar = Scalar_;  

    template <typename Derived>
    explicit ArrivalInterval(Eigen::DenseBase<Derived> const & rates, bool endsInArrival = true)
        : internal::ValueRates<Space, Scalar_>{rates}
        , m_endsInArrival{endsInArrival}
    {}

    [[nodiscard]] bool
    endsInArrival() const
    {
        return m_endsInArrival;
    }
 
private:
 
    bool m_endsInArrival;  
};
 
namespace internal
{

template <typename Space, typename Scalar>
struct ArrivalKernel
{
    typename ContinuousTimeMarkovChain<Space, Scalar>::MatrixType kernel;  
    Scalar logScale;  
};

template <typename Space, typename Scalar>
[[nodiscard]] ArrivalKernel<Space, Scalar>
arrivalKernel(
    ContinuousTimeMarkovChain<Space, Scalar> const & chain,
    Scalar dt,
    ArrivalInterval<Space, Scalar> const & interval
)
{
    EMRYS_FILTER_ASSERT(dt >= 0 && "the interval must be nonnegative");
    using MatrixType = typename ContinuousTimeMarkovChain<Space, Scalar>::MatrixType;
    auto const & rates = interval.rates();
    EMRYS_FILTER_ASSERT(rates.size() == chain.space().cardinality() && "give one rate per value");
    Scalar const shift = rates.minCoeff();
    MatrixType exponent = chain.generator() * dt;
    exponent.diagonal().array() -= (rates - shift) * dt;
    MatrixType kernel = exponent.exp();
    kernel = kernel.cwiseMax(Scalar(0));
    if (interval.endsInArrival())
    {
        kernel = kernel * rates.matrix().asDiagonal();
    }
    return {std::move(kernel), -shift * dt};
}

template <typename Space, typename Scalar, typename Kernel>
[[nodiscard]] std::expected<CategoricalUpdate<Space, Scalar>, WeightError>
moveByKernel(CategoricalParams<Space, Scalar> const & prior, Kernel const & kernel, Scalar logScale)
{
    using ArrayType = typename CategoricalParams<Space, Scalar>::ArrayType;
    auto const priorTotal = prior.logTotalWeight();
    if (!priorTotal)
    {
        return std::unexpected{priorTotal.error()};
    }
    ArrayType const logMoved = movedLogWeights<Scalar>(prior.logWeights(), kernel);
    auto const total = logTotalWeight<Scalar>(logMoved);
    if (!total)
    {
        return std::unexpected{total.error()};
    }
    return CategoricalUpdate<Space, Scalar>{
        CategoricalParams<Space, Scalar>{prior.space(), logMoved - *total},
        *total - *priorTotal + logScale
    };
}
 
}  // namespace internal

template <typename Scalar, typename Space>
[[nodiscard]] std::expected<CategoricalUpdate<Space, Scalar>, WeightError>
predict(
    CategoricalParams<Space, Scalar> const & prior,
    ContinuousTimeMarkovChain<Space, Scalar> const & chain,
    TimeStep<Scalar> auto const & dt,
    ArrivalInterval<Space, Scalar> const & interval
)
{
    EMRYS_FILTER_ASSERT(prior.space() == chain.space() && "the chain is on a different space");
    auto const step = internal::arrivalKernel(chain, internal::toScalarTime<Scalar>(dt), interval);
    return internal::moveByKernel(prior, step.kernel, step.logScale);
}
 
}  // namespace Emrys