Likelihood.h

Go to the documentation of this file.

#pragma once

 
#include <array>
#include <concepts>
#include <cstddef>
#include <expected>
#include <optional>
#include <span>
#include <tuple>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
 
#include <Eigen/Cholesky>
 
#include <EmrysFilter/Core/Allocation.h>
#include <EmrysFilter/Core/Execution.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Core/Weights.h>
#include <EmrysFilter/Distributions/Gaussian.h>
#include <EmrysFilter/Distributions/Mixture.h>
#include <EmrysFilter/Distributions/Product.h>
#include <EmrysFilter/Models/MeasurementModel.h>
#include <EmrysFilter/Models/MeasurementReads.h>
#include <EmrysFilter/Models/Transform.h>
#include <EmrysFilter/Spaces/Composite.h>
#include <EmrysFilter/Spaces/Coordinates.h>
#include <EmrysFilter/Spaces/Indexing.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{
 
namespace internal
{

template <typename Scalar, typename Noise, typename Residual, typename Omega>
[[nodiscard]] Scalar
widenedLogDensity(Noise const & noise, Residual const & residual, Omega const & omega)
{
    if (!allFinite(omega))
    {
        return Eigen::NumTraits<Scalar>::quiet_NaN();
    }
    auto const component = [&](auto const & gaussian) {
        using Gaussian = std::remove_cvref_t<decltype(gaussian)>;
        using Widened = GaussianParams<typename Gaussian::SpaceType, Scalar>;
        using Matrix = typename Widened::MatrixType;
        Matrix cov = symmetricFromLower<Matrix>(gaussian.covarianceLower());
        cov += omega;
        auto const widened = Widened{
            gaussian.space(),
            typename Widened::VectorType{gaussian.mean().template cast<Scalar>()},
            std::move(cov)
        };
        auto const density = widened.toDensity();
        EMRYS_FILTER_ASSERT(density && "the measurement noise must factorize");
        return Scalar(density->logDensity(residual));
    };
    if constexpr (isGaussianMixture<Noise>)
    {
        auto const total = noise.logTotalWeight();
        EMRYS_FILTER_ASSERT(total && "the noise mixture's weights must be valid");
        auto sum = LogSumExp<Scalar>{};
        for (auto const & entry : noise)
        {
            sum.add(Scalar(entry.logWeight) - Scalar(*total) + component(entry.component));
        }
        return sum.value();
    }
    else
    {
        return component(noise);
    }
}

template <ExecutionPolicy P, typename R, typename RowsOf, typename Scalar, typename Scratch>
void
fillBatchedLogLikelihoods(
    P && policy,
    R const & resolved,
    Eigen::Index count,
    RowsOf const & rowsOf,
    Eigen::Index start,
    Scalar * result,
    Scratch & scratch
)
{
    using InMatrix = typename R::template InMatrix<Scalar>;
    Eigen::Index const rows = valueDimension(resolved.in());
    scratch.resize(std::size_t(rows * count));
    forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
        auto x = Eigen::Map<InMatrix>{scratch.data() + rows * begin, rows, length};
        resolved.template gatherColumnsFromRows<Scalar>(
            x,
            [&](Eigen::Index j) {
                return rowsOf(begin + j);
            },
            start
        );
        resolved.batchedLogLikelihoods(
            Eigen::Map<InMatrix const>{x.data(), rows, length},
            Eigen::Map<Eigen::Matrix<Scalar, 1, Eigen::Dynamic>>{result + begin, length}
        );
    });
}

template <ExecutionPolicy P, WeighsPoints M, typename Space, typename Derived, typename Result, typename Scratch>
void
fillLogLikelihoods(
    P && policy,
    M const & measurement,
    Space const & space,
    Eigen::MatrixBase<Derived> const & states,
    Result && result,
    Scratch & scratch
)
{
    using Scalar = typename Derived::Scalar;
    EMRYS_FILTER_ASSERT(states.rows() == valueDimension(space) && "the states do not fit the space");
    EMRYS_FILTER_ASSERT(result.size() == states.cols() && "give one entry per state");
    static_assert(
        !hasConsider<M>,
        "consider() is a Kalman update's; a filter that weights points moves no block it does not condition "
        "(decisions.md [consider-states])"
    );
    auto const resolved = resolveMeasurement(measurement, space);
    using R = std::remove_cvref_t<decltype(resolved)>;
    if constexpr (R::BatchedLikelihood)
    {
        fillBatchedLogLikelihoods(
            std::forward<P>(policy),
            resolved,
            states.cols(),
            [&](Eigen::Index j) {
                return states.col(j);
            },
            Eigen::Index{0},
            result.data(),
            scratch
        );
    }
    else if constexpr (R::HasLogLikelihood)
    {
        forEachIndex(std::forward<P>(policy), states.cols(), [&](Eigen::Index j) {
            result(j) = Scalar(resolved.logLikelihoodAt(states.col(j)));
        });
    }
    else
    {
        static_assert(
            R::HasAdditiveNoise,
            "a measurement model with no additive noise has no likelihood of its own at a point; give it "
            "logLikelihood(g), or name the inputs it is affine in with affineIn() beside additiveNoise()"
        );
        static_assert(
            R::inputsAffine(),
            "the likelihood at a point of noise that passes through the function is exact only where the model is "
            "affine in it, so name the input in affineIn(), or give the model logLikelihood(g) (decisions.md "
            "[likelihood-measurements])"
        );
        auto const & space_ = resolved.space();
        if constexpr (R::InputCount == 0 && R::BatchedFunction)
        {
            // The function once per batch, its outputs after its inputs in the scratch.
            auto const & noise = resolved.noise().distribution();
            auto const density = noise.toDensity();
            EMRYS_FILTER_ASSERT(density && "the measurement noise must factorize");
            using InMatrix = typename R::template InMatrix<Scalar>;
            using OutMatrix = typename R::template OutMatrix<Scalar>;
            Eigen::Index const count = states.cols();
            Eigen::Index const inRows = valueDimension(resolved.in());
            Eigen::Index const outRows = valueDimension(space_);
            scratch.resize(std::size_t((inRows + outRows) * count));
            forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
                auto x = Eigen::Map<InMatrix>{scratch.data() + inRows * begin, inRows, length};
                auto y = Eigen::Map<OutMatrix>{scratch.data() + inRows * count + outRows * begin, outRows, length};
                resolved.template gatherColumnsFromRows<Scalar>(
                    x,
                    [&](Eigen::Index j) {
                        return states.col(begin + j);
                    },
                    Eigen::Index{0}
                );
                resolved.evaluateBatch(Eigen::Map<InMatrix const>{x.data(), inRows, length}, y);
                for (Eigen::Index j = 0; j < length; ++j)
                {
                    result(begin + j) = density->logDensity(observedResidual<Scalar>(resolved, space_, y.col(j)));
                }
            });
        }
        else if constexpr (R::InputCount == 0)
        {
            auto const & noise = resolved.noise().distribution();
            auto const density = noise.toDensity();
            EMRYS_FILTER_ASSERT(density && "the measurement noise must factorize");
            forEachIndex(std::forward<P>(policy), states.cols(), [&](Eigen::Index j) {
                auto const predicted = valueOf(resolved.evaluateAtState(states.col(j))).eval();
                result(j) = density->logDensity(observedResidual<Scalar>(resolved, space_, predicted));
            });
        }
        else if constexpr (R::BatchedFunction)
        {
            // The function and its Jacobian in the inputs once per batch, each input at its mean, the
            // inputs' covariance carried through that Jacobian at each point: what the linear method
            // gives below, exact where the function is affine in them.
            using Out = typename R::Out;
            using InMatrix = typename R::template InMatrix<Scalar>;
            using OutMatrix = typename R::template OutMatrix<Scalar>;
            using JacobianMatrix = MatrixOf<
                Scalar,
                valueDimensionAtCompileTime<DifferenceSpace<Out>>,
                Eigen::Dynamic,
                maxValueDimensionAtCompileTime<DifferenceSpace<Out>>,
                Eigen::Dynamic>;
            using Omega = typename Moments<Out, Scalar>::MatrixType;
            auto const inputs = resolved.inputs();
            Eigen::Index w = 0;
            std::apply(
                [&](auto const &... input) {
                    ((w += inputOf(input).differenceDimension()), ...);
                },
                inputs
            );
            Eigen::Index const count = states.cols();
            Eigen::Index const inRows = valueDimension(resolved.in());
            Eigen::Index const outRows = valueDimension(space_);
            Eigen::Index const m = differenceDimension(space_);
            scratch.resize(std::size_t((inRows + outRows + m * w) * count));
            Scalar * const outputs = scratch.data() + inRows * count;
            Scalar * const jacobianData = outputs + outRows * count;
            forEachChunk(std::forward<P>(policy), count, [&](Eigen::Index begin, Eigen::Index length) {
                auto x = Eigen::Map<InMatrix>{scratch.data() + inRows * begin, inRows, length};
                auto y = Eigen::Map<OutMatrix>{outputs + outRows * begin, outRows, length};
                auto jacobians = Eigen::Map<JacobianMatrix>{jacobianData + m * w * begin, m, w * length};
                resolved.template gatherColumnsFromRows<Scalar>(
                    x,
                    [&](Eigen::Index j) {
                        return states.col(begin + j);
                    },
                    Eigen::Index{0}
                );
                resolved.evaluateBatchWithInputJacobians(
                    Eigen::Map<InMatrix const>{x.data(), inRows, length},
                    y,
                    jacobians
                );
                for (Eigen::Index j = 0; j < length; ++j)
                {
                    Omega omega = Omega::Zero(m, m);
                    Eigen::Index offset = j * w;
                    std::apply(
                        [&](auto const &... argument) {
                            (
                                [&](auto const & input) {
                                    using Input = std::remove_cvref_t<decltype(input)>;
                                    using Cov = MatrixOf<
                                        Scalar,
                                        Input::DifferenceDimensionAtCompileTime,
                                        Input::DifferenceDimensionAtCompileTime,
                                        Input::MaxDifferenceDimensionAtCompileTime,
                                        Input::MaxDifferenceDimensionAtCompileTime>;
                                    Eigen::Index const d = input.differenceDimension();
                                    auto const block =
                                        middleColsOf<Input::DifferenceDimensionAtCompileTime>(jacobians, offset, d);
                                    Cov const cov = symmetricFromLower<Cov>(covarianceLowerOf(input));
                                    omega.noalias() += multiply(multiply(block, cov), block.transpose());
                                    offset += d;
                                }(inputOf(argument)),
                                ...
                            );
                        },
                        inputs
                    );
                    omega = (Scalar(0.5) * (omega + omega.transpose())).eval();
                    result(begin + j) = widenedLogDensity<Scalar>(
                        resolved.noise().distribution(),
                        observedResidual<Scalar>(resolved, space_, y.col(j)),
                        omega
                    );
                }
            });
        }
        else
        {
            // The inputs' covariance carried through the function's Jacobian at the point, by the linear
            // method, which is exact where the function is affine in them.
            auto const & reads = resolved.reads();
            using Reads = typename R::Reads;
            using ReadPoint = PointInput<typename Reads::JointSpace, Scalar>;
            auto const f = ResolvedFunction{resolved};
            forEachIndex(std::forward<P>(policy), states.cols(), [&](Eigen::Index j) {
                auto const read =
                    ReadPoint{reads.space(), typename ReadPoint::VectorType{states.col(j)(reads.indices())}};
                result(j) = linearizeMeasurement(
                    resolved,
                    methods::linearized(),
                    f,
                    read,
                    [&](auto const & problem, auto const & linearized) {
                        using Problem = std::remove_cvref_t<decltype(problem)>;
                        using Omega = typename Problem::OutputMatrix;
                        Omega const omega = inputsCovariance<Omega>(resolved, problem, linearized.jacobian);
                        return widenedLogDensity<Scalar>(
                            resolved.noise().distribution(),
                            observedResidual<Scalar>(resolved, space_, linearized.mean),
                            omega
                        );
                    }
                );
            });
        }
    }
}

template <ExecutionPolicy P, WeighsPoints M, typename Space, typename Derived, typename Result>
void
fillLogLikelihoods(
    P && policy,
    M const & measurement,
    Space const & space,
    Eigen::MatrixBase<Derived> const & states,
    Result && result
)
{
    auto scratch = Buffer<typename Derived::Scalar, DefaultAllocator>();
    fillLogLikelihoods(std::forward<P>(policy), measurement, space, states, std::forward<Result>(result), scratch);
}

template <WeighsPoints M, typename Space, typename Derived, typename Result>
void
fillLogLikelihoods(
    M const & measurement,
    Space const & space,
    Eigen::MatrixBase<Derived> const & states,
    Result && result
)
{
    fillLogLikelihoods(execution::sequential, measurement, space, states, std::forward<Result>(result));
}
 
}  // namespace internal

template <internal::WeighsPoints M, typename Space, typename Derived>
[[nodiscard]] Eigen::Array<typename Derived::Scalar, Eigen::Dynamic, 1>
logLikelihoods(M const & measurement, Space const & space, Eigen::MatrixBase<Derived> const & states)
{
    return logLikelihoods(execution::sequential, measurement, space, states);
}

template <ExecutionPolicy P, internal::WeighsPoints M, typename Space, typename Derived>
[[nodiscard]] Eigen::Array<typename Derived::Scalar, Eigen::Dynamic, 1>
logLikelihoods(P && policy, M const & measurement, Space const & space, Eigen::MatrixBase<Derived> const & states)
{
    auto result = Eigen::Array<typename Derived::Scalar, Eigen::Dynamic, 1>{states.cols()};
    internal::fillLogLikelihoods(std::forward<P>(policy), measurement, space, states, result);
    return result;
}

template <internal::WeighsPoints M, typename Space, typename Derived>
[[nodiscard]] typename Derived::Scalar
logLikelihood(M const & measurement, Space const & space, Eigen::MatrixBase<Derived> const & x)
{
    EMRYS_FILTER_ASSERT(x.cols() == 1 && "the state point must be a column vector");
    return logLikelihoods(measurement, space, x)(0);
}

template <FixedSizeSpace S, internal::WeighsPoints M, typename Derived>
    requires Space<S>
[[nodiscard]] Eigen::Array<typename Derived::Scalar, Eigen::Dynamic, 1>
logLikelihoods(M const & measurement, Eigen::MatrixBase<Derived> const & states)
{
    return logLikelihoods(execution::sequential, measurement, S{}, states);
}

template <FixedSizeSpace S, ExecutionPolicy P, internal::WeighsPoints M, typename Derived>
    requires Space<S>
[[nodiscard]] Eigen::Array<typename Derived::Scalar, Eigen::Dynamic, 1>
logLikelihoods(P && policy, M const & measurement, Eigen::MatrixBase<Derived> const & states)
{
    return logLikelihoods(std::forward<P>(policy), measurement, S{}, states);
}

template <FixedSizeSpace S, internal::WeighsPoints M, typename Derived>
    requires Space<S>
[[nodiscard]] typename Derived::Scalar
logLikelihood(M const & measurement, Eigen::MatrixBase<Derived> const & x)
{
    return logLikelihood(measurement, S{}, x);
}
 
}  // namespace Emrys