Gamma.h

Go to the documentation of this file.

#pragma once

 
#include <cmath>
#include <concepts>
#include <expected>
#include <limits>
#include <optional>
#include <random>
#include <array>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Print.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Spaces/Composite.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys
{

enum class GammaError
{
    NotPositive,  
    NonFinite,    
};
 
namespace internal
{
 
template <typename S>
consteval bool
holdsHalfLine();

template <typename Blocks>
struct AllHoldHalfLine;

template <typename... Blocks>
struct AllHoldHalfLine<std::tuple<Blocks...>> : std::bool_constant<(holdsHalfLine<Blocks>() && ...)>
{};

template <typename S>
consteval bool
holdsHalfLine()
{
    if constexpr (IsComposite<S>)
    {
        return AllHoldHalfLine<CompositeBlocks<S>>::value;
    }
    else if constexpr (IsRepeated<S>)
    {
        return holdsHalfLine<RepeatedBlock<S>>();
    }
    else
    {
        constexpr auto support = leafSupport<S>();
        return support && !support->periodic && support->lower <= 0
            && support->upper == std::numeric_limits<double>::infinity();
    }
}
 
}  // namespace internal

template <Space Space_, RealScalar Scalar_ = double>
class Gamma
{
public:
 
    using Scalar = Scalar_;                               
    using Index = Eigen::Index;                           
    using SpaceType = Space_;                             
    using MomentsType = Moments<Space_, Scalar_>;         
    using VectorType = typename MomentsType::VectorType;  
    using MatrixType = typename MomentsType::MatrixType;  

    static constexpr int ValueDimensionAtCompileTime = MomentsType::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = MomentsType::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = MomentsType::DifferenceDimensionAtCompileTime;

    static constexpr int MaxDifferenceDimensionAtCompileTime = MomentsType::MaxDifferenceDimensionAtCompileTime;
 
    static_assert(
        internal::hasEqualSizes<Space_>,
        "a gamma distribution needs a space whose values and differences have the same coordinates"
    );
    static_assert(
        internal::holdsHalfLine<Space_>(),
        "a gamma distribution needs every coordinate to contain [0, inf) without folding it"
    );

    template <typename Shape, typename Scale>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    create(Space_ space, Shape const & shape, Scale const & scale)
    {
        Index const n = parameterSize(shape, scale, Emrys::valueDimension(space));
        internal::fitValueDimension(space, n);
        VectorType shapes = broadcast(shape, n);
        VectorType scales = broadcast(scale, n);
        if (auto const error = check(shapes, scales))
        {
            return std::unexpected{*error};
        }
        return Gamma{std::move(space), std::move(shapes), std::move(scales)};
    }

    template <typename Shape, typename Scale>
        requires internal::SizedByLength<Space_>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    create(Shape const & shape, Scale const & scale)
    {
        return create(Space_{}, shape, scale);
    }

    template <typename Scale>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    exponential(Space_ space, Scale const & scale)
    {
        return create(std::move(space), Scalar(1), scale);
    }

    template <typename Scale>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    exponential(Scale const & scale)
    {
        return exponential(Space_{}, scale);
    }

    template <typename Dof>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    chiSquared(Space_ space, Dof const & dof)
    {
        Index const n = parameterSize(dof, Scalar(2), Emrys::valueDimension(space));
        return create(std::move(space), VectorType{Scalar(0.5) * broadcast(dof, n)}, Scalar(2));
    }

    template <typename Dof>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    chiSquared(Dof const & dof)
    {
        return chiSquared(Space_{}, dof);
    }

    template <typename Mean, typename StdDev>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    fromMeanAndStdDev(Space_ space, Mean const & mean, StdDev const & stdDev)
    {
        Index const n = parameterSize(mean, stdDev, Emrys::valueDimension(space));
        VectorType const mu = broadcast(mean, n);
        VectorType const sigma = broadcast(stdDev, n);
        if (!allFinite(mu) || !allFinite(sigma))
        {
            return std::unexpected{GammaError::NonFinite};
        }
        if (!(mu.array() > Scalar(0)).all() || !(sigma.array() > Scalar(0)).all())
        {
            return std::unexpected{GammaError::NotPositive};
        }
        return create(
            std::move(space),
            VectorType{(mu.array() / sigma.array()).square().matrix()},
            VectorType{(sigma.array().square() / mu.array()).matrix()}
        );
    }

    template <typename Mean, typename StdDev>
    [[nodiscard]] static std::expected<Gamma, GammaError>
    fromMeanAndStdDev(Mean const & mean, StdDev const & stdDev)
    {
        return fromMeanAndStdDev(Space_{}, mean, stdDev);
    }

    [[nodiscard]] Index
    valueDimension() const
    {
        return m_shape.size();
    }

    [[nodiscard]] Index
    differenceDimension() const
    {
        return m_shape.size();
    }

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

    [[nodiscard]] VectorType const &
    shape() const
    {
        return m_shape;
    }

    [[nodiscard]] VectorType const &
    scale() const
    {
        return m_scale;
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        static_assert(Derived::IsVectorAtCompileTime, "logDensity() expects a vector");
        EMRYS_FILTER_ASSERT(x.size() == valueDimension());
        Scalar result = m_logNormalizer;
        for (Index i = 0; i < valueDimension(); ++i)
        {
            Scalar const xi = x(i);
            if (xi < Scalar(0))
            {
                return -std::numeric_limits<Scalar>::infinity();
            }
            // Skip the power term for shape 1, where it is 0 · log 0 at the origin.
            if (m_shape(i) != Scalar(1))
            {
                result += (m_shape(i) - Scalar(1)) * Eigen::numext::log(xi);
            }
            result -= xi / m_scale(i);
        }
        return result;
    }

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

    [[nodiscard]] MomentsType
    moments() const
    {
        auto result = MomentsType{valueDimension(), differenceDimension()};
        result.mean() = m_shape.cwiseProduct(m_scale);
        result.covarianceLower() = result.mean().cwiseProduct(m_scale).asDiagonal();
        return result;
    }

    template <typename Path>
    [[nodiscard]] auto
    marginal() const
    {
        return internal::marginalOverPath<Path>(m_space, [&](auto const & target, auto space) {
            return restricted(target.values, std::move(space));
        });
    }

    template <typename Path>
    [[nodiscard]] auto
    marginal(Index element) const
    {
        return internal::marginalOverElement<Path>(m_space, element, [&](auto const & target, auto space) {
            return restricted(target.values, std::move(space));
        });
    }

    template <typename GivenPath, typename Derived>
    [[nodiscard]] auto
    conditional(Eigen::MatrixBase<Derived> const & values) const
    {
        return internal::tagConditional<GivenPath>(
            m_space,
            [&](auto const & t, auto const & g, auto space, auto const &) {
                EMRYS_FILTER_ASSERT(values.size() == g.values.size());
                return restricted(t.values, std::move(space));
            }
        );
    }
 
private:

    template <Space, RealScalar>
    friend class Gamma;

    Gamma(Space_ space, VectorType shape, VectorType scale)
        : m_space{std::move(space)}
        , m_shape{std::move(shape)}
        , m_scale{std::move(scale)}
    {
        for (Index i = 0; i < m_shape.size(); ++i)
        {
            m_logNormalizer -= logGamma<Scalar>(m_shape(i)) + m_shape(i) * Eigen::numext::log(m_scale(i));
        }
    }

    [[nodiscard]] static std::optional<GammaError>
    check(VectorType const & shape, VectorType const & scale)
    {
        if (!allFinite(shape) || !allFinite(scale))
        {
            return GammaError::NonFinite;
        }
        if (!(shape.array() > Scalar(0)).all() || !(scale.array() > Scalar(0)).all())
        {
            return GammaError::NotPositive;
        }
        return std::nullopt;
    }

    template <typename Target, typename SubSpace>
    [[nodiscard]] Gamma<SubSpace, Scalar>
    restricted(Target const & target, SubSpace space) const
    {
        using ResultVector = typename Gamma<SubSpace, Scalar>::VectorType;
        return Gamma<SubSpace, Scalar>{std::move(space), ResultVector{m_shape(target)}, ResultVector{m_scale(target)}};
    }

    template <typename P>
    static constexpr bool isScalar = std::is_arithmetic_v<P>;

    template <typename A, typename B>
    [[nodiscard]] static Index
    parameterSize(A const & a, B const & b, Index spaceSize)
    {
        if constexpr (!isScalar<A>)
        {
            if constexpr (!isScalar<B>)
            {
                EMRYS_FILTER_ASSERT(a.size() == b.size() && "the parameters differ in size");
            }
            return a.size();
        }
        else if constexpr (!isScalar<B>)
        {
            return b.size();
        }
        else
        {
            return spaceSize;
        }
    }

    template <typename P>
    [[nodiscard]] static VectorType
    broadcast(P const & p, Index n)
    {
        if constexpr (isScalar<P>)
        {
            return VectorType::Constant(n, Scalar(p));
        }
        else
        {
            EMRYS_FILTER_ASSERT(p.size() == n);
            return VectorType{p};
        }
    }
 
    EMRYS_FILTER_NO_UNIQUE_ADDRESS Space_ m_space;  
    VectorType m_shape;                             
    VectorType m_scale;                             
    Scalar m_logNormalizer = 0;                     
};

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleSet<Space, Scalar>
drawSamples(Gamma<Space, Scalar> const & gamma, Eigen::Index n, G & rng)
{
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    EMRYS_FILTER_ASSERT(n >= 0);
    using Distribution = std::gamma_distribution<Scalar>;
    constexpr int maxDimension = Gamma<Space, Scalar>::MaxValueDimensionAtCompileTime;
    // One distribution per coordinate, inline for a bounded space (see \ref Allocation).
    using Distributions = std::conditional_t<
        maxDimension == Eigen::Dynamic,
        std::vector<Distribution>,
        std::array<Distribution, std::size_t(maxDimension == Eigen::Dynamic ? 0 : maxDimension)>>;
    Eigen::Index const d = gamma.valueDimension();
    auto distributions = Distributions();
    if constexpr (maxDimension == Eigen::Dynamic)
    {
        distributions.resize(std::size_t(d));
    }
    for (Eigen::Index i = 0; i < d; ++i)
    {
        distributions[std::size_t(i)] = Distribution{gamma.shape()(i), gamma.scale()(i)};
    }
    auto samples = SampleSet<Space, Scalar>{d, n};
    for (Eigen::Index j = 0; j < n; ++j)
    {
        for (Eigen::Index i = 0; i < d; ++i)
        {
            samples(i, j) = distributions[std::size_t(i)](rng);
        }
    }
    return samples;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<Space, Scalar>::VectorType
drawSample(Gamma<Space, Scalar> const & gamma, G & rng)
{
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    auto sample = typename Moments<Space, Scalar>::VectorType{};
    sample.resize(gamma.valueDimension());
    for (Eigen::Index i = 0; i < sample.size(); ++i)
    {
        sample(i) = std::gamma_distribution<Scalar>{gamma.shape()(i), gamma.scale()(i)}(rng);
    }
    return sample;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(Gamma<Space, Scalar> const & gamma, Eigen::Index n, G & rng)
{
    auto result = SamplesWithLogDensity<Space, Scalar>{drawSamples(gamma, n, rng), {}};
    result.logDensity.resize(n);
    for (Eigen::Index j = 0; j < n; ++j)
    {
        result.logDensity(j) = gamma.logDensity(result.samples.col(j));
    }
    return result;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleWithLogDensity<Space, Scalar>
drawSampleWithLogDensity(Gamma<Space, Scalar> const & gamma, G & rng)
{
    auto result = SampleWithLogDensity<Space, Scalar>{drawSample(gamma, rng), 0};
    result.logDensity = gamma.logDensity(result.sample);
    return result;
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

inline void
describe(std::ostream & stream, GammaError error, Indent indent)
{
    stream << indent;
    switch (error)
    {
        case GammaError::NotPositive:
            stream << "a shape or scale is not positive";
            break;
        case GammaError::NonFinite:
            stream << "a shape or scale is not finite";
            break;
    }
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO