KernelDensity.h

Go to the documentation of this file.

#pragma once

 
#include <cstddef>
#include <expected>
#include <random>
#include <span>
#include <utility>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Forwarding.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Resampling.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Stamp.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/ParticleSet.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Distributions/Sampling.h>
#include <EmrysFilter/Distributions/Statistics.h>
#include <EmrysFilter/Spaces/Indexing.h>
 
namespace Emrys
{

template <Space Space_, RealScalar Scalar_ = double, typename Allocator_ = DefaultAllocator>
class KernelDensity
{
public:
 
    using SpaceType = Space_;                                     
    using ParamsType = ParticleSet<Space_, Scalar_, Allocator_>;  
    using allocator_type = Allocator_;                            

    static constexpr int ValueDimensionAtCompileTime = ParamsType::ValueDimensionAtCompileTime;

    static constexpr int MaxValueDimensionAtCompileTime = ParamsType::MaxValueDimensionAtCompileTime;

    static constexpr int DifferenceDimensionAtCompileTime = ParamsType::DifferenceDimensionAtCompileTime;

    static constexpr int MaxDifferenceDimensionAtCompileTime = ParamsType::MaxDifferenceDimensionAtCompileTime;

    using KernelType =
        GaussianParams<Linear<DifferenceDimensionAtCompileTime, MaxDifferenceDimensionAtCompileTime>, Scalar_>;
 
    using MomentsType = typename ParamsType::MomentsType;                    
    using Scalar = typename ParamsType::Scalar;                              
    using Index = typename ParamsType::Index;                                
    using VectorType = typename ParamsType::VectorType;                      
    using DifferenceVectorType = typename ParamsType::DifferenceVectorType;  
    using MatrixType = typename ParamsType::MatrixType;                      
    using SampleSetType = typename ParamsType::SampleSetType;                

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

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

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

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

    [[nodiscard]] ParamsType const &
    particles() const
    {
        return m_particles.get();
    }

    [[nodiscard]] Scalar
    logTotalWeight() const
    {
        return m_logTotalWeight;
    }

    [[nodiscard]] MatrixType const &
    bandwidth() const
    {
        return m_kernel.covarianceLower();
    }

    [[nodiscard]] typename KernelType::DensityType
    kernel() const
    {
        auto const density = m_kernel.toDensity();
        return *density;
    }

    [[nodiscard]] MomentsType
    moments() const
    {
        return moments(execution::sequential);
    }

    template <ExecutionPolicy P>
    [[nodiscard]] MomentsType
    moments(P && policy) const
    {
        MomentsType result = internal::weightedMomentsIn<MomentsType>(
            std::forward<P>(policy),
            space(),
            particles().states(),
            particles().logWeights(),
            m_logTotalWeight
        );
        result.covarianceLower() += m_kernel.covarianceLower().template selfadjointView<Eigen::Lower>().toDenseMatrix();
        return result;
    }

    template <typename Derived>
    [[nodiscard]] Scalar
    logDensity(Eigen::MatrixBase<Derived> const & x) const
    {
        return logDensity(execution::sequential, x);
    }

    template <ExecutionPolicy P, typename Derived>
    [[nodiscard]] Scalar
    logDensity(P && policy, Eigen::MatrixBase<Derived> const & x) const
    {
        static_assert(Derived::IsVectorAtCompileTime, "logDensity() expects a vector");
        EMRYS_FILTER_ASSERT(x.size() == valueDimension());
        auto const & llt = internal::FactorizedGaussian::llt(m_kernel);
        auto const & set = particles();
        LogSumExp<Scalar> const sum = internal::reduceRanges(
            std::forward<P>(policy),
            size(),
            [&](Index begin, Index length) {
                auto range = LogSumExp<Scalar>{};
                for (Index i = begin; i < begin + length; ++i)
                {
                    DifferenceVectorType residual = Emrys::minus(set.space(), set.state(i), x);
                    llt.matrixL().solveInPlace(residual);
                    range.add(set.logWeight(i) - Scalar(0.5) * residual.squaredNorm());
                }
                return range;
            },
            [](LogSumExp<Scalar> a, LogSumExp<Scalar> const & b) {
                a.add(b);
                return a;
            }
        );
        return sum.value() - m_logTotalWeight + internal::FactorizedGaussian::logPeak(m_kernel);
    }

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

    template <MatrixExpression Bandwidth>
    [[nodiscard]] static std::expected<KernelDensity, GaussianError>
    create(ParamsType const & particles, Scalar logTotalWeight, Bandwidth && bandwidth)
    {
        auto kernel =
            KernelType{DifferenceVectorType::Zero(particles.differenceDimension()), std::forward<Bandwidth>(bandwidth)};
        if (auto result = kernel.factorize(); !result)
        {
            return std::unexpected{result.error()};
        }
        return KernelDensity{particles, logTotalWeight, std::move(kernel)};
    }

    KernelDensity(ParamsType const & particles, Scalar logTotalWeight, KernelType kernel)
        : m_particles{particles}
        , m_logTotalWeight{logTotalWeight}
        , m_kernel{std::move(kernel)}
    {}
 
    internal::ValuesCheck<ParamsType> m_particles;  
    Scalar m_logTotalWeight;                        
    KernelType m_kernel;                            
};

template <
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SampleSet<Space, Scalar>
drawSamples(KernelDensity<Space, Scalar, Allocator> const & kde, Eigen::Index n, G & rng, R const & resampler = {})
{
    EMRYS_FILTER_ASSERT(n >= 0);
    using SampleSetType = SampleSet<Space, Scalar>;
    auto const & particles = kde.particles();
    auto weights = internal::bufferOf<Scalar>(particles.get_allocator());
    weights.resize(std::size_t(particles.size()));
    for (Eigen::Index i = 0; i < particles.size(); ++i)
    {
        weights[std::size_t(i)] = Eigen::numext::exp(particles.logWeight(i) - kde.logTotalWeight());
    }
    auto ancestors = internal::bufferOf<Eigen::Index>(particles.get_allocator());
    ancestors.resize(std::size_t(n));
    resampler(std::span<Scalar const>{weights}, std::span<Eigen::Index>{ancestors}, rng);
 
    using Normals = SampleSet<DifferenceSpace<Space>, Scalar>;
    Normals const offsets = internal::multiply(
        internal::asTriangular<Eigen::Lower>(kde.kernel().llt().matrixLLT()),
        internal::standardNormal<Normals>(kde.differenceDimension(), n, rng)
    );
    auto samples = SampleSetType{kde.valueDimension(), n};
    for (Eigen::Index j = 0; j < n; ++j)
    {
        samples.col(j) = Emrys::plus(kde.space(), particles.state(ancestors[std::size_t(j)]), offsets.col(j));
    }
    return samples;
}

template <typename Scalar, typename Space, typename Allocator, std::uniform_random_bit_generator G>
[[nodiscard]] typename Moments<Space, Scalar>::VectorType
drawSample(KernelDensity<Space, Scalar, Allocator> const & kde, G & rng)
{
    using Difference = DifferenceVector<Space, Scalar>;
    auto const & particles = kde.particles();
    Eigen::Index const picked = internal::drawIndex<Scalar>(
        particles.size(),
        [&](Eigen::Index k) {
            return Eigen::numext::exp(Scalar(particles.logWeight(k) - kde.logTotalWeight()));
        },
        rng
    );
    Difference const z = internal::standardNormalVector<Difference>(kde.differenceDimension(), rng);
    return typename Moments<Space, Scalar>::VectorType{Emrys::plus(
        kde.space(),
        particles.state(picked),
        Difference{internal::multiply(internal::asTriangular<Eigen::Lower>(kde.kernel().llt().matrixLLT()), z)}
    )};
}

template <
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(
    KernelDensity<Space, Scalar, Allocator> const & kde,
    Eigen::Index n,
    G & rng,
    R const & resampler = {}
)
{
    return drawSamplesWithLogDensity(execution::sequential, kde, n, rng, resampler);
}

template <
    ExecutionPolicy P,
    typename Scalar,
    typename Space,
    typename Allocator,
    std::uniform_random_bit_generator G,
    typename R = resamplers::Multinomial>
    requires Resampler<R, Scalar, G>
[[nodiscard]] SamplesWithLogDensity<Space, Scalar>
drawSamplesWithLogDensity(
    P && policy,
    KernelDensity<Space, Scalar, Allocator> const & kde,
    Eigen::Index n,
    G & rng,
    R const & resampler = {}
)
{
    auto result = SamplesWithLogDensity<Space, Scalar>{drawSamples(kde, n, rng, resampler), {}};
    result.logDensity.resize(n);
    internal::forEachIndex(std::forward<P>(policy), n, [&](Eigen::Index j) {
        result.logDensity(j) = kde.logDensity(result.samples.col(j));
    });
    return result;
}

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