StudentT.h

Go to the documentation of this file.

#pragma once

 
#include <cmath>
#include <concepts>
#include <expected>
#include <optional>
#include <random>
#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/Print.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Moments.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
 
namespace Emrys
{

enum class StudentTError
{
    NotPositiveDefinite,          
    NonFiniteScale,               
    NotPositiveDegreesOfFreedom,  
    NonFiniteDegreesOfFreedom,    
};

template <Space Space_, RealScalar Scalar_ = double>
class StudentT
{
public:

    using ScaleGaussianType = GaussianParams<Space_, Scalar_>;
 
    using MomentsType = typename ScaleGaussianType::MomentsType;  
    using Scalar = Scalar_;                                       
    using Index = Eigen::Index;                                   
    using SpaceType = Space_;                                     
    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;

    template <MatrixExpression Location, typename ScaleDerived>
    [[nodiscard]] static std::expected<StudentT, StudentTError>
    create(Space_ space, Location && location, Eigen::MatrixBase<ScaleDerived> const & scale, Scalar dof)
    {
        if (auto const error = checkDegreesOfFreedom(dof))
        {
            return std::unexpected{*error};
        }
        auto gaussian = ScaleGaussianType{
            std::move(space),
            std::forward<Location>(location),
            MatrixType{scale.template selfadjointView<Eigen::Lower>()}
        };
        if (auto result = gaussian.factorize(); !result)
        {
            return std::unexpected{
                result.error() == GaussianError::NonFinite ? StudentTError::NonFiniteScale
                                                           : StudentTError::NotPositiveDefinite
            };
        }
        return StudentT{std::move(gaussian), dof};
    }

    template <MatrixExpression Location, typename ScaleDerived>
    [[nodiscard]] static std::expected<StudentT, StudentTError>
    create(Location && location, Eigen::MatrixBase<ScaleDerived> const & scale, Scalar dof)
    {
        return create(Space_{}, std::forward<Location>(location), scale, dof);
    }

    [[nodiscard]] static std::expected<StudentT, StudentTError>
    create(ScaleGaussianType const & gaussian, Scalar dof)
    {
        return create(gaussian.space(), gaussian.mean(), gaussian.covarianceLower(), dof);
    }

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

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

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

    [[nodiscard]] VectorType const &
    location() const
    {
        return m_gaussian.mean();
    }

    [[nodiscard]] MatrixType const &
    scale() const
    {
        return m_gaussian.covarianceLower();
    }

    [[nodiscard]] Scalar
    dof() const
    {
        return m_dof;
    }

    [[nodiscard]] typename ScaleGaussianType::DensityType
    scaleDensity() const
    {
        auto const density = m_gaussian.toDensity();
        return *density;
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        return logDensityAtSquaredDistance(scaleDensity().squaredMahalanobisDistance(x));
    }

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

    [[nodiscard]] std::expected<MomentsType, MomentsError>
    moments() const
    {
        if (!(m_dof > Scalar(2)))
        {
            return std::unexpected{MomentsError::Infinite};
        }
        return MomentsType{location(), MatrixType{(m_dof / (m_dof - Scalar(2))) * scale()}};
    }

    [[nodiscard]] Scalar
    logDensityAtSquaredDistance(Scalar squaredDistance) const
    {
        return m_logNormalizer
             - Scalar(0.5) * (m_dof + internal::scalarFromCount<Scalar>(differenceDimension()))
                   * logOnePlus<Scalar>(squaredDistance / m_dof);
    }

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

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

    template <typename GivenPath, typename Derived>
    [[nodiscard]] auto
    conditional(Eigen::MatrixBase<Derived> const & values) const
    {
        return internal::tagConditional<GivenPath>(
            space(),
            [&](auto const & t, auto const & g, auto subspace, auto const & givenSpace) {
                return conditioned(t, g, values, std::move(subspace), givenSpace);
            }
        );
    }
 
private:

    [[nodiscard]] static std::optional<StudentTError>
    checkDegreesOfFreedom(Scalar dof)
    {
        if (!isFinite(dof))
        {
            return StudentTError::NonFiniteDegreesOfFreedom;
        }
        if (!(dof > Scalar(0)))
        {
            return StudentTError::NotPositiveDegreesOfFreedom;
        }
        return std::nullopt;
    }

    template <typename Target, typename SubSpace>
    [[nodiscard]] StudentT<SubSpace, Scalar>
    restricted(Target const & target, SubSpace subspace) const
    {
        using Result = StudentT<SubSpace, Scalar>;
        auto result = Result::create(
            std::move(subspace),
            location()(target.values),
            typename Result::MatrixType{internal::lowerSymmetricBlock(scale(), target.differences, target.differences)},
            m_dof
        );
        EMRYS_FILTER_ASSERT(result.has_value() && "a principal block of a positive definite scale should factorize");
        return std::move(*result);
    }

    template <typename Target, typename Given, typename Derived, typename SubSpace, typename GivenSpace>
    [[nodiscard]] std::expected<StudentT<SubSpace, Scalar>, StudentTError>
    conditioned(
        Target const & target,
        Given const & given,
        Eigen::MatrixBase<Derived> const & values,
        SubSpace subspace,
        GivenSpace const & givenSpace
    ) const
    {
        using Result = StudentT<SubSpace, Scalar>;
        auto gaussian =
            internal::conditionMomentsWithEvidence(location(), scale(), target, given, values, givenSpace, subspace);
        if (!gaussian)
        {
            return std::expected<Result, StudentTError>{std::unexpect, StudentTError::NotPositiveDefinite};
        }
        Scalar const givenSize = Scalar(given.differences.size());
        Scalar const stretch = (m_dof + gaussian->squaredDistance) / (m_dof + givenSize);
        return Result::create(
            std::move(subspace),
            std::move(gaussian->moments.mean),
            typename Result::MatrixType{stretch * gaussian->moments.cov},
            m_dof + givenSize
        );
    }

    StudentT(ScaleGaussianType gaussian, Scalar dof)
        : m_gaussian{std::move(gaussian)}
        , m_dof{dof}
    {
        Scalar const d = internal::scalarFromCount<Scalar>(differenceDimension());
        Scalar const logDeterminant = scaleDensity().logDeterminant();
        m_logNormalizer = logGamma<Scalar>(Scalar(0.5) * (m_dof + d)) - logGamma<Scalar>(Scalar(0.5) * m_dof)
                        - Scalar(0.5) * d * Eigen::numext::log(Scalar(m_dof * pi<Scalar>()))
                        - Scalar(0.5) * logDeterminant;
    }
 
    ScaleGaussianType m_gaussian;  
    Scalar m_dof;                  
    Scalar m_logNormalizer{};      
};
 
namespace internal
{

template <typename Space, typename Scalar = double>
struct StudentTDraws
{
    SampleSet<Space, Scalar> samples;  
    Eigen::Array<Scalar, Eigen::Dynamic, 1>
        squaredDistance;  
};

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] StudentTDraws<Space, Scalar>
drawStudentT(StudentT<Space, Scalar> const & t, Eigen::Index n, G & rng)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const z = standardNormal<Normals>(t.differenceDimension(), n, rng);
    auto chiSquared = std::chi_squared_distribution<Scalar>{t.dof()};
    auto stretch = Eigen::Array<Scalar, Eigen::Dynamic, 1>{n};  // sqrt(ν / u) for each draw
    for (Eigen::Index j = 0; j < n; ++j)
    {
        stretch(j) = Eigen::numext::sqrt(t.dof() / chiSquared(rng));
    }
    Normals offsets = multiply(asTriangular<Eigen::Lower>(t.scaleDensity().llt().matrixLLT()), z);
    for (Eigen::Index j = 0; j < n; ++j)
    {
        offsets.col(j) *= stretch(j);
    }
    return {
        Emrys::plus(t.space(), t.location(), offsets),
        z.colwise().squaredNorm().transpose().array() * stretch.square()
    };
}

template <typename Space, typename Scalar = double>
struct StudentTDraw
{
    typename Moments<Space, Scalar>::VectorType sample;  
    Scalar squaredDistance = 0;                          
};

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] StudentTDraw<Space, Scalar>
drawOneStudentT(StudentT<Space, Scalar> const & t, G & rng)
{
    using Difference = DifferenceVector<Space, Scalar>;
    Difference const z = standardNormalVector<Difference>(t.differenceDimension(), rng);
    Scalar const stretch = Eigen::numext::sqrt(t.dof() / std::chi_squared_distribution<Scalar>{t.dof()}(rng));
    Difference offset = multiply(asTriangular<Eigen::Lower>(t.scaleDensity().llt().matrixLLT()), z);
    offset *= stretch;
    return {
        typename Moments<Space, Scalar>::VectorType{Emrys::plus(t.space(), t.location(), offset)},
        z.squaredNorm() * stretch * stretch
    };
}
 
}  // namespace internal

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleSet<Space, Scalar>
drawSamples(StudentT<Space, Scalar> const & t, Eigen::Index n, G & rng)
{
    auto draws = internal::drawStudentT(t, n, rng);
    return std::move(draws.samples);
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<Space, Scalar>::VectorType
drawSample(StudentT<Space, Scalar> const & t, G & rng)
{
    return internal::drawOneStudentT(t, rng).sample;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(StudentT<Space, Scalar> const & t, Eigen::Index n, G & rng)
{
    auto draws = internal::drawStudentT(t, n, rng);
    auto result = SamplesWithLogDensity<Space, Scalar>{std::move(draws.samples), {}};
    result.logDensity.resize(n);
    if constexpr (isLinear<Space>)
    {
        for (Eigen::Index j = 0; j < n; ++j)
        {
            result.logDensity(j) = t.logDensityAtSquaredDistance(draws.squaredDistance(j));
        }
    }
    else
    {
        for (Eigen::Index j = 0; j < n; ++j)
        {
            result.logDensity(j) = t.logDensity(result.samples.col(j));
        }
    }
    return result;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleWithLogDensity<Space, Scalar>
drawSampleWithLogDensity(StudentT<Space, Scalar> const & t, G & rng)
{
    auto draw = internal::drawOneStudentT(t, rng);
    auto result = SampleWithLogDensity<Space, Scalar>{std::move(draw.sample), 0};
    if constexpr (isLinear<Space>)
    {
        result.logDensity = t.logDensityAtSquaredDistance(draw.squaredDistance);
    }
    else
    {
        result.logDensity = t.logDensity(result.sample);
    }
    return result;
}
 
}  // namespace Emrys
 
#ifndef EMRYS_FILTER_NO_IO
 
namespace Emrys
{
 
namespace internal
{

inline void
describe(std::ostream & stream, StudentTError error, Indent indent)
{
    stream << indent;
    switch (error)
    {
        case StudentTError::NotPositiveDefinite:
            stream << "the scale matrix is not positive definite";
            break;
        case StudentTError::NonFiniteScale:
            stream << "the scale matrix is not finite";
            break;
        case StudentTError::NotPositiveDegreesOfFreedom:
            stream << "the degrees of freedom are not positive";
            break;
        case StudentTError::NonFiniteDegreesOfFreedom:
            stream << "the degrees of freedom are not finite";
            break;
    }
}
 
}  // namespace internal
 
}  // namespace Emrys
 
#endif  // EMRYS_FILTER_NO_IO