Statistics.h

Go to the documentation of this file.

#pragma once

 
#include <utility>
 
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Distributions/SampleSet.h>
#include <EmrysFilter/Spaces/Space.h>
 
namespace Emrys::internal
{

template <typename MatrixType, typename DeviationsDerived, typename WeightsDerived>
[[nodiscard]] MatrixType
weightedCovariance(Eigen::MatrixBase<DeviationsDerived> const & deviations, Eigen::ArrayBase<WeightsDerived> const & w)
{
    using Scalar = typename MatrixType::Scalar;
    using Deviations = MatrixOf<
        Scalar,
        MatrixType::RowsAtCompileTime,
        DeviationsDerived::ColsAtCompileTime,
        MatrixType::MaxRowsAtCompileTime,
        DeviationsDerived::MaxColsAtCompileTime>;
    Deviations const scaled = deviations * w.sqrt().matrix().asDiagonal();
    auto cov = MatrixType{deviations.rows(), deviations.rows()};
    if constexpr (MatrixType::MaxRowsAtCompileTime == 1)
    {
        // rankUpdate() would treat the single-row matrix as a column vector.
        cov.setConstant(scaled.squaredNorm());
    }
    else
    {
        cov.setZero();
        cov.template selfadjointView<Eigen::Lower>().rankUpdate(scaled);
        for (Eigen::Index j = 1; j < cov.cols(); ++j)
        {
            for (Eigen::Index i = 0; i < j; ++i)
            {
                cov(i, j) = cov(j, i);
            }
        }
    }
    return cov;
}

template <typename Derived>
void
mirrorLower(Eigen::MatrixBase<Derived> & cov)
{
    cov.template triangularView<Eigen::StrictlyUpper>() = cov.transpose();
}

template <typename MomentsType, typename Space, typename MeanAt, typename LogWeightAt>
[[nodiscard]] MomentsType
pointMomentsIn(
    Space const & space,
    Eigen::Index count,
    MeanAt meanAt,
    LogWeightAt logWeightAt,
    typename MomentsType::Scalar logTotal
)
{
    using Scalar = typename MomentsType::Scalar;
    using DifferenceVectorType = typename MomentsType::DifferenceVectorType;
    auto const weight = [&](Eigen::Index i) {
        return Eigen::numext::exp(Scalar(logWeightAt(i)) - logTotal);
    };
    auto const skip = [&](Eigen::Index i) {
        return Scalar(logWeightAt(i)) == -Eigen::NumTraits<Scalar>::infinity();
    };
    auto result = MomentsType{valueDimension(space), differenceDimension(space)};
    auto & cov = result.covarianceLower();
    cov.setZero();
    // Eigen's rank update loops at run time whatever the size, and so does a product assigned to a
    // triangle, which Eigen sends to its rank update; at a fixed size the lazy outer product, of
    // the same terms w v_j v_i, is assigned to the triangle coefficient by coefficient with the
    // sizes fixed.
    auto const addCov = [&](DifferenceVectorType const & deviation, Scalar w) {
        if constexpr (DifferenceVectorType::SizeAtCompileTime != Eigen::Dynamic)
        {
            cov.template triangularView<Eigen::Lower>() += deviation.lazyProduct((w * deviation).transpose());
        }
        else
        {
            cov.template selfadjointView<Eigen::Lower>().rankUpdate(deviation, w);
        }
    };
    if constexpr (hasPlainDifferences<Space>)
    {
        result.mean().setZero();
        for (Eigen::Index i = 0; i < count; ++i)
        {
            if (!skip(i))
            {
                result.mean() += weight(i) * meanAt(i);
            }
        }
        for (Eigen::Index i = 0; i < count; ++i)
        {
            if (!skip(i))
            {
                addCov(DifferenceVectorType{meanAt(i) - result.mean()}, weight(i));
            }
        }
        normalizeValues(space, result.mean());
    }
    else
    {
        Eigen::Index reference = 0;
        for (Eigen::Index i = 1; i < count; ++i)
        {
            if (logWeightAt(i) > logWeightAt(reference))
            {
                reference = i;
            }
        }
        result.mean() = meanAt(reference);
        auto shift = DifferenceVectorType{};
        for (int pass = 0; pass < 2; ++pass)
        {
            shift.setZero(differenceDimension(space));
            for (Eigen::Index i = 0; i < count; ++i)
            {
                if (!skip(i))
                {
                    shift += weight(i) * Emrys::minus(space, meanAt(i), result.mean());
                }
            }
            if (pass == 1)
            {
                for (Eigen::Index i = 0; i < count; ++i)
                {
                    if (!skip(i))
                    {
                        addCov(DifferenceVectorType{Emrys::minus(space, meanAt(i), result.mean()) - shift}, weight(i));
                    }
                }
            }
            result.mean() = Emrys::plus(space, result.mean(), shift);
        }
    }
    return result;
}

template <typename MomentsType, typename Space, typename StatesDerived, typename WeightsDerived>
[[nodiscard]] MomentsType
weightedMomentsIn(
    Space const & space,
    Eigen::MatrixBase<StatesDerived> const & states,
    Eigen::ArrayBase<WeightsDerived> const & logWeights,
    typename MomentsType::Scalar logTotal
)
{
    MomentsType result = pointMomentsIn<MomentsType>(
        space,
        states.cols(),
        [&](Eigen::Index i) {
            return states.col(i);
        },
        [&](Eigen::Index i) {
            return logWeights.derived()(i);
        },
        logTotal
    );
    mirrorLower(result.covarianceLower());
    return result;
}

template <typename MomentsType, ExecutionPolicy P, typename Space, typename StatesDerived, typename WeightsDerived>
[[nodiscard]] MomentsType
weightedMomentsIn(
    P && policy,
    Space const & space,
    Eigen::MatrixBase<StatesDerived> const & states,
    Eigen::ArrayBase<WeightsDerived> const & logWeights,
    typename MomentsType::Scalar logTotal
)
{
    if constexpr (isSequential<P>)
    {
        return weightedMomentsIn<MomentsType>(space, states, logWeights, logTotal);
    }
    else
    {
        using Scalar = typename MomentsType::Scalar;
        using VectorType = typename MomentsType::VectorType;
        using DifferenceVectorType = typename MomentsType::DifferenceVectorType;
        using MatrixType = typename MomentsType::MatrixType;
        using Deviations = SampleSet<DifferenceSpace<typename MomentsType::SpaceType>, Scalar>;
        using Weights = Eigen::Array<Scalar, Eigen::Dynamic, 1>;
        Eigen::Index const count = states.cols();
        auto const weightsOf = [&](Eigen::Index begin, Eigen::Index length) {
            return Weights{(logWeights.derived().segment(begin, length) - logTotal).exp()};
        };
        auto const add = [](auto a, auto const & b) {
            a += b;
            return a;
        };
        // The weighted covariance of the differences of the points from center, less shift.
        auto const covarianceAbout = [&](auto const & center, DifferenceVectorType const & shift) {
            return reduceRanges(
                policy,
                count,
                [&](Eigen::Index begin, Eigen::Index length) {
                    Deviations deviations = Emrys::minus(space, states.middleCols(begin, length), center);
                    deviations.colwise() -= shift;
                    return weightedCovariance<MatrixType>(deviations, weightsOf(begin, length));
                },
                add
            );
        };
        auto result = MomentsType{valueDimension(space), differenceDimension(space)};
        if constexpr (hasPlainDifferences<Space>)
        {
            result.mean() = reduceRanges(
                policy,
                count,
                [&](Eigen::Index begin, Eigen::Index length) {
                    return VectorType{states.middleCols(begin, length) * weightsOf(begin, length).matrix()};
                },
                add
            );
            result.covarianceLower() =
                covarianceAbout(result.mean(), DifferenceVectorType::Zero(differenceDimension(space)));
            normalizeValues(space, result.mean());
        }
        else
        {
            // The first of the heaviest points, as Eigen's maxCoeff() finds it.
            struct Heaviest
            {
                Scalar logWeight;    
                Eigen::Index index;  
            };
 
            Heaviest const heaviest = reduceRanges(
                policy,
                count,
                [&](Eigen::Index begin, Eigen::Index length) {
                    Eigen::Index index = 0;
                    Scalar const logWeight = logWeights.derived().segment(begin, length).maxCoeff(&index);
                    return Heaviest{logWeight, begin + index};
                },
                [](Heaviest const & a, Heaviest const & b) {
                    return b.logWeight > a.logWeight ? b : a;
                }
            );
            auto const shiftFrom = [&](VectorType const & center) {
                return reduceRanges(
                    policy,
                    count,
                    [&](Eigen::Index begin, Eigen::Index length) {
                        return DifferenceVectorType{
                            Emrys::minus(space, states.middleCols(begin, length), center)
                            * weightsOf(begin, length).matrix()
                        };
                    },
                    add
                );
            };
            VectorType const reference = states.col(heaviest.index);
            VectorType const center = Emrys::plus(space, reference, shiftFrom(reference));
            DifferenceVectorType const shift = shiftFrom(center);
            result.mean() = Emrys::plus(space, center, shift);
            result.covarianceLower() = covarianceAbout(center, shift);
        }
        return result;
    }
}

template <typename MomentsType, typename Space, typename MeanAt, typename CovAt, typename LogWeightAt>
[[nodiscard]] MomentsType
mixtureMomentsIn(
    Space const & space,
    Eigen::Index count,
    MeanAt meanAt,
    CovAt covAt,
    LogWeightAt logWeightAt,
    typename MomentsType::Scalar logTotal
)
{
    using Scalar = typename MomentsType::Scalar;
    MomentsType result = pointMomentsIn<MomentsType>(space, count, meanAt, logWeightAt, logTotal);
    for (Eigen::Index i = 0; i < count; ++i)
    {
        if (Scalar(logWeightAt(i)) != -Eigen::NumTraits<Scalar>::infinity())
        {
            result.covarianceLower().template triangularView<Eigen::Lower>() +=
                Eigen::numext::exp(Scalar(logWeightAt(i)) - logTotal) * covAt(i);
        }
    }
    mirrorLower(result.covarianceLower());
    return result;
}

template <typename WeightsDerived>
[[nodiscard]] typename WeightsDerived::Scalar
ess(Eigen::ArrayBase<WeightsDerived> const & logWeights, typename WeightsDerived::Scalar logTotal)
{
    using Scalar = typename WeightsDerived::Scalar;
    return Scalar(1) / (logWeights - logTotal).exp().square().sum();
}

template <ExecutionPolicy P, typename WeightsDerived>
[[nodiscard]] typename WeightsDerived::Scalar
ess(P && policy, Eigen::ArrayBase<WeightsDerived> const & logWeights, typename WeightsDerived::Scalar logTotal)
{
    using Scalar = typename WeightsDerived::Scalar;
    if constexpr (isSequential<P>)
    {
        return ess(logWeights, logTotal);
    }
    else
    {
        Scalar const sumOfSquares = reduceRanges(
            std::forward<P>(policy),
            logWeights.size(),
            [&](Eigen::Index begin, Eigen::Index length) {
                return Scalar((logWeights.derived().segment(begin, length) - logTotal).exp().square().sum());
            },
            [](Scalar a, Scalar b) {
                return a + b;
            }
        );
        return Scalar(1) / sumOfSquares;
    }
}
 
}  // namespace Emrys::internal