Sampling.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <concepts>
#include <random>
 
#include <Eigen/Cholesky>
 
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Distributions/SampleSet.h>
 
namespace Emrys
{

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

template <typename Space, typename Scalar = double>
struct SampleWithLogDensity
{
    typename Moments<Space, Scalar>::VectorType sample;  
    Scalar logDensity = 0;                               
};
 
namespace internal
{

template <typename Matrix, std::uniform_random_bit_generator G>
[[nodiscard]] Matrix
standardNormal(Eigen::Index rows, Eigen::Index cols, G & rng)
{
    using Scalar = typename Matrix::Scalar;
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    auto z = Matrix{rows, cols};
    auto normal = std::normal_distribution<Scalar>{};
    std::generate(z.data(), z.data() + z.size(), [&] {
        return normal(rng);
    });
    return z;
}

template <typename Vector, std::uniform_random_bit_generator G>
[[nodiscard]] Vector
standardNormalVector(Eigen::Index rows, G & rng)
{
    static_assert(std::floating_point<typename Vector::Scalar>, "sampling requires a standard floating-point scalar");
    auto z = Vector{};
    z.resize(rows);
    auto normal = std::normal_distribution<typename Vector::Scalar>{};
    std::generate(z.data(), z.data() + z.size(), [&] {
        return normal(rng);
    });
    return z;
}

template <typename Derived, std::uniform_random_bit_generator G, typename Scalar>
void
fillStandardNormal(Derived && z, G & rng, std::normal_distribution<Scalar> & normal)
{
    static_assert(std::same_as<typename std::remove_cvref_t<Derived>::Scalar, Scalar>);
    std::generate(z.data(), z.data() + z.size(), [&] {
        return normal(rng);
    });
}

template <typename Derived, std::uniform_random_bit_generator G>
void
fillStandardNormal(Derived && z, G & rng)
{
    using Scalar = typename std::remove_cvref_t<Derived>::Scalar;
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    auto normal = std::normal_distribution<Scalar>{};
    fillStandardNormal(std::forward<Derived>(z), rng, normal);
}

template <typename Matrix, std::uniform_random_bit_generator G>
[[nodiscard]] Matrix
standardUniform(Eigen::Index rows, Eigen::Index cols, G & rng)
{
    using Scalar = typename Matrix::Scalar;
    static_assert(std::floating_point<Scalar>, "sampling requires a standard floating-point scalar");
    auto u = Matrix{rows, cols};
    auto uniform = std::uniform_real_distribution<Scalar>{};
    std::generate(u.data(), u.data() + u.size(), [&] {
        return uniform(rng);
    });
    return u;
}
 
}  // namespace internal

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleSet<Space, Scalar>
drawSamples(Gaussian<Space, Scalar> const & gaussian, Eigen::Index n, G & rng)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const z = internal::standardNormal<Normals>(gaussian.differenceDimension(), n, rng);
    return SampleSet<Space, Scalar>{Emrys::plus(
        gaussian.space(),
        gaussian.mean(),
        internal::multiply(internal::asTriangular<Eigen::Lower>(gaussian.llt().matrixLLT()), z)
    )};
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<Space, Scalar>::VectorType
drawSample(Gaussian<Space, Scalar> const & gaussian, G & rng)
{
    using Difference = DifferenceVector<Space, Scalar>;
    Difference const z = internal::standardNormalVector<Difference>(gaussian.differenceDimension(), rng);
    return typename Moments<Space, Scalar>::VectorType{Emrys::plus(
        gaussian.space(),
        gaussian.mean(),
        Difference{internal::multiply(internal::asTriangular<Eigen::Lower>(gaussian.llt().matrixLLT()), z)}
    )};
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(Gaussian<Space, Scalar> const & gaussian, Eigen::Index n, G & rng)
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const z = internal::standardNormal<Normals>(gaussian.differenceDimension(), n, rng);
    auto result = SamplesWithLogDensity<Space, Scalar>{
        Emrys::plus(
            gaussian.space(),
            gaussian.mean(),
            internal::multiply(internal::asTriangular<Eigen::Lower>(gaussian.llt().matrixLLT()), z)
        ),
        {}
    };
    if constexpr (isLinear<Space>)
    {
        result.logDensity = gaussian.logPeak() - Scalar(0.5) * z.colwise().squaredNorm().transpose().array();
    }
    else
    {
        result.logDensity.resize(n);
        for (Eigen::Index j = 0; j < n; ++j)
        {
            result.logDensity(j) = gaussian.logDensity(result.samples.col(j));
        }
    }
    return result;
}

template <typename Scalar, typename Space, std::uniform_random_bit_generator G>
[[nodiscard]] SampleWithLogDensity<Space, Scalar>
drawSampleWithLogDensity(Gaussian<Space, Scalar> const & gaussian, G & rng)
{
    using Difference = DifferenceVector<Space, Scalar>;
    Difference const z = internal::standardNormalVector<Difference>(gaussian.differenceDimension(), rng);
    auto result = SampleWithLogDensity<Space, Scalar>{
        typename Moments<Space, Scalar>::VectorType{Emrys::plus(
            gaussian.space(),
            gaussian.mean(),
            Difference{internal::multiply(internal::asTriangular<Eigen::Lower>(gaussian.llt().matrixLLT()), z)}
        )},
        0
    };
    if constexpr (isLinear<Space>)
    {
        result.logDensity = gaussian.logPeak() - Scalar(0.5) * z.squaredNorm();
    }
    else
    {
        result.logDensity = gaussian.logDensity(result.sample);
    }
    return result;
}
 
}  // namespace Emrys