Indexing.h

Go to the documentation of this file.

#pragma once

 
#include <algorithm>
#include <array>
#include <cstddef>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <Eigen/Cholesky>
 
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Core/Math.h>
#include <EmrysFilter/Core/Multiply.h>
#include <EmrysFilter/Core/Solve.h>
#include <EmrysFilter/Spaces/Reset.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/Subspace.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys::internal
{

template <int Size, int MaxSize>
using IndexVector = Eigen::Matrix<Eigen::Index, Size, 1, Eigen::ColMajor, MaxSize, 1>;

template <typename Indices, int SourceSize, int SourceMaxSize>
struct IndexedSizeOf
{
private:

    using Probe = Eigen::Matrix<int, SourceSize, 1, Eigen::ColMajor, SourceMaxSize, 1>;

    using Expr = std::remove_cvref_t<decltype(std::declval<Probe const &>()(std::declval<Indices const &>()))>;
 
    static_assert(Expr::ColsAtCompileTime == 1, "marginal and conditional indices must select a set, not one index");
 
public:
 
    static constexpr int size = Expr::RowsAtCompileTime;  

    static constexpr int maxSize = size != Eigen::Dynamic                       ? size
                                 : Expr::MaxRowsAtCompileTime != Eigen::Dynamic ? int(Expr::MaxRowsAtCompileTime)
                                                                                : SourceMaxSize;
 
    using type = IndexVector<size, maxSize>;  
};

template <typename Indices, int SourceSize, int SourceMaxSize>
using IndexedSize = typename IndexedSizeOf<Indices, SourceSize, SourceMaxSize>::type;

template <int SourceSize, int SourceMaxSize, typename Indices>
[[nodiscard]] IndexedSize<Indices, SourceSize, SourceMaxSize>
resolveIndices(Eigen::Index n, Indices const & indices)
{
    using Identity = Eigen::Matrix<Eigen::Index, SourceSize, 1, Eigen::ColMajor, SourceMaxSize, 1>;
    Identity const identity = Identity::LinSpaced(n, 0, n - 1);
    return identity(indices);
}

template <int Size, int MaxSize>
[[nodiscard]] IndexVector<Size, MaxSize>
consecutiveIndices(Eigen::Index start, Eigen::Index count)
{
    return IndexVector<Size, MaxSize>::LinSpaced(count, start, start + count - 1);
}

template <typename... Vectors>
[[nodiscard]] bool
distinctIndices(Vectors const &... indices)
{
    bool distinct = true;
    std::size_t position = 0;
    auto const seenBefore = [&](Eigen::Index index, std::size_t before) {
        std::size_t k = 0;
        bool seen = false;
        auto const scan = [&](auto const & vector) {
            for (Eigen::Index const other : vector)
            {
                if (k++ < before && other == index)
                {
                    seen = true;
                }
            }
        };
        (scan(indices), ...);
        return seen;
    };
    auto const check = [&](auto const & vector) {
        for (Eigen::Index const index : vector)
        {
            if (seenBefore(index, position))
            {
                distinct = false;
            }
            ++position;
        }
    };
    (check(indices), ...);
    return distinct;
}

template <int SourceSize, int SourceMaxSize, int GivenSize, int GivenMaxSize>
[[nodiscard]] auto
complementIndices(Eigen::Index n, IndexVector<GivenSize, GivenMaxSize> const & given)
{
    auto taken = Eigen::Matrix<bool, SourceSize, 1, Eigen::ColMajor, SourceMaxSize, 1>{};
    taken.setConstant(n, false);
    for (Eigen::Index const i : given)
    {
        taken(i) = true;
    }
    constexpr int RestSize =
        SourceSize != Eigen::Dynamic && GivenSize != Eigen::Dynamic ? SourceSize - GivenSize : Eigen::Dynamic;
    auto result = IndexVector<RestSize, RestSize == Eigen::Dynamic ? SourceMaxSize : RestSize>{};
    result.resize(n - given.size());
    Eigen::Index out = 0;
    for (Eigen::Index i = 0; i < n; ++i)
    {
        if (!taken(i))
        {
            result(out++) = i;
        }
    }
    return result;
}

template <typename Values, typename Differences>
struct Coordinates
{
    Values values;            
    Differences differences;  
};

template <typename Space, typename ValueIndices, typename DifferenceIndices>
[[nodiscard]] auto
resolveCoordinates(Space const & space, ValueIndices const & values, DifferenceIndices const & differences)
{
    return Coordinates{
        resolveIndices<valueDimensionAtCompileTime<Space>, maxValueDimensionAtCompileTime<Space>>(
            valueDimension(space),
            asIndexSet(values)
        ),
        resolveIndices<
            valueDimensionAtCompileTime<DifferenceSpace<Space>>,
            maxValueDimensionAtCompileTime<DifferenceSpace<Space>>>(differenceDimension(space), asIndexSet(differences))
    };
}

template <typename Path, typename Space>
[[nodiscard]] auto
pathCoordinates(Space const & space)
{
    return resolveCoordinates(space, idx<Path>(space), diffIdx<Path>(space));
}

template <typename IndexVector_, typename DifferenceIndexVector_, typename SubSpace>
struct ResolvedSelection
{
    IndexVector_ indices;                      
    DifferenceIndexVector_ differenceIndices;  
    SubSpace space;                            
};

template <typename Resolved>
struct SharedSelection
{
    Resolved const * resolved;  
};

template <typename Selection>
inline constexpr bool isSharedSelection = false;

template <typename Resolved>
inline constexpr bool isSharedSelection<SharedSelection<Resolved>> = true;

template <typename Source, typename Selection, typename SubSpace>
class SelectionHolder
{
public:

    using ResolvedType = ResolvedSelection<
        IndexedSize<
            decltype(Selection::values),
            Source::ValueDimensionAtCompileTime,
            Source::MaxValueDimensionAtCompileTime>,
        IndexedSize<
            decltype(Selection::differences),
            Source::DifferenceDimensionAtCompileTime,
            Source::MaxDifferenceDimensionAtCompileTime>,
        SubSpace>;

    SelectionHolder(Source const & source, Selection const & selection, SubSpace space)
        : m_resolved{
              resolveIndices<Source::ValueDimensionAtCompileTime, Source::MaxValueDimensionAtCompileTime>(
                  source.valueDimension(),
                  selection.values
              ),
              resolveIndices<Source::DifferenceDimensionAtCompileTime, Source::MaxDifferenceDimensionAtCompileTime>(
                  source.differenceDimension(),
                  selection.differences
              ),
              std::move(space)
          }
    {
        EMRYS_FILTER_ASSERT(
            distinctIndices(m_resolved.indices) && distinctIndices(m_resolved.differenceIndices)
            && "marginal indices must be distinct"
        );
    }

    [[nodiscard]] ResolvedType const &
    get() const
    {
        return m_resolved;
    }
 
private:
 
    ResolvedType m_resolved;  
};

template <typename Source, typename Resolved, typename SubSpace>
class SelectionHolder<Source, SharedSelection<Resolved>, SubSpace>
{
public:
 
    using ResolvedType = Resolved;  

    explicit SelectionHolder(SharedSelection<Resolved> shared)
        : m_resolved{shared.resolved}
    {}

    [[nodiscard]] ResolvedType const &
    get() const
    {
        return *m_resolved;
    }
 
private:
 
    Resolved const * m_resolved;  
};

template <typename GivenPath, typename Space, typename Condition>
[[nodiscard]] auto
tagConditional(Space const & space, Condition condition)
{
    constexpr int N = valueDimensionAtCompileTime<Space>;
    constexpr int NMax = maxValueDimensionAtCompileTime<Space>;
    constexpr int D = valueDimensionAtCompileTime<DifferenceSpace<Space>>;
    constexpr int DMax = maxValueDimensionAtCompileTime<DifferenceSpace<Space>>;
    Eigen::Index const n = valueDimension(space);
    Eigen::Index const d = differenceDimension(space);
    auto const given = pathCoordinates<GivenPath>(space);
    auto const target =
        Coordinates{complementIndices<N, NMax>(n, given.values), complementIndices<D, DMax>(d, given.differences)};
    return condition(target, given, pathComplement<GivenPath>(space), pathSubspace<GivenPath>(space));
}

template <typename Path, typename Space, typename Make>
[[nodiscard]] auto
marginalOverPath(Space const & space, Make make)
{
    return make(pathCoordinates<Path>(space), pathSubspace<Path>(space));
}

template <typename Path, typename Space, typename Make>
[[nodiscard]] auto
marginalOverElement(Space const & space, Eigen::Index element, Make make)
{
    return make(
        resolveCoordinates(space, idx<Path>(space, element), diffIdx<Path>(space, element)),
        pathElementSubspace<Path>(space)
    );
}

template <typename Matrix, typename Derived>
void
insertRows(Matrix & matrix, Eigen::Index at, Eigen::MatrixBase<Derived> const & rows)
{
    Eigen::Index const n = matrix.rows();
    Eigen::Index const k = rows.rows();
    EMRYS_FILTER_ASSERT(at >= 0 && at <= n && rows.cols() == matrix.cols());
    auto result = Matrix{};
    result.resize(n + k, matrix.cols());
    result.topRows(at) = matrix.topRows(at);
    result.middleRows(at, k) = rows;
    result.bottomRows(n - at) = matrix.bottomRows(n - at);
    matrix = std::move(result);
}

template <typename Matrix>
void
eraseRows(Matrix & matrix, Eigen::Index at, Eigen::Index count)
{
    Eigen::Index const n = matrix.rows();
    EMRYS_FILTER_ASSERT(at >= 0 && count >= 0 && at + count <= n);
    auto result = Matrix{};
    result.resize(n - count, matrix.cols());
    result.topRows(at) = matrix.topRows(at);
    result.bottomRows(n - at - count) = matrix.bottomRows(n - at - count);
    matrix = std::move(result);
}

template <typename Matrix, typename BlockDerived, typename CrossDerived>
void
insertSymmetric(
    Matrix & matrix,
    Eigen::Index at,
    Eigen::MatrixBase<BlockDerived> const & block,
    Eigen::MatrixBase<CrossDerived> const & cross
)
{
    Eigen::Index const n = matrix.rows();
    Eigen::Index const k = block.rows();
    Eigen::Index const after = n - at;
    EMRYS_FILTER_ASSERT(at >= 0 && at <= n && block.cols() == k && cross.rows() == k && cross.cols() == n);
    auto result = Matrix{};
    result.resize(n + k, n + k);
    result.topLeftCorner(at, at) = matrix.topLeftCorner(at, at);
    result.topRightCorner(at, after) = matrix.topRightCorner(at, after);
    result.bottomLeftCorner(after, at) = matrix.bottomLeftCorner(after, at);
    result.bottomRightCorner(after, after) = matrix.bottomRightCorner(after, after);
    result.block(at, at, k, k) = block;
    result.block(at, 0, k, at) = cross.leftCols(at);
    result.block(at, at + k, k, after) = cross.rightCols(after);
    result.block(0, at, at, k) = cross.leftCols(at).transpose();
    result.block(at + k, at, after, k) = cross.rightCols(after).transpose();
    matrix = std::move(result);
}

template <typename Matrix>
void
eraseSymmetric(Matrix & matrix, Eigen::Index at, Eigen::Index count)
{
    Eigen::Index const n = matrix.rows();
    Eigen::Index const after = n - at - count;
    EMRYS_FILTER_ASSERT(at >= 0 && count >= 0 && after >= 0);
    auto result = Matrix{};
    result.resize(n - count, n - count);
    result.topLeftCorner(at, at) = matrix.topLeftCorner(at, at);
    result.topRightCorner(at, after) = matrix.topRightCorner(at, after);
    result.bottomLeftCorner(after, at) = matrix.bottomLeftCorner(after, at);
    result.bottomRightCorner(after, after) = matrix.bottomRightCorner(after, after);
    matrix = std::move(result);
}

template <typename Result, typename Derived>
[[nodiscard]] Result
symmetricFromLower(Eigen::MatrixBase<Derived> const & matrix)
{
    if constexpr (std::is_base_of_v<Eigen::PlainObjectBase<Derived>, Derived>)
    {
        Result result = matrix.template selfadjointView<Eigen::Lower>();
        return result;
    }
    else
    {
        Result result = matrix;
        result.template triangularView<Eigen::StrictlyUpper>() = result.transpose();
        return result;
    }
}

template <int Rows, int MaxRows, int Cols, int MaxCols, typename Derived>
[[nodiscard]] auto
lowerSymmetricBlock(
    Eigen::MatrixBase<Derived> const & matrix,
    IndexVector<Rows, MaxRows> const & rows,
    IndexVector<Cols, MaxCols> const & cols
)
{
    using Scalar = typename Derived::Scalar;
    using Result = MatrixOf<Scalar, Rows, Cols, MaxRows, MaxCols>;
    return Result::NullaryExpr(
        rows.size(),
        cols.size(),
        [&matrix = matrix.derived(), &rows, &cols](Eigen::Index i, Eigen::Index j) {
            Eigen::Index const r = rows(i);
            Eigen::Index const c = cols(j);
            return r >= c ? matrix(r, c) : matrix(c, r);
        }
    );
}

template <int Size, int MaxSize, typename Derived, typename Indices>
[[nodiscard]] auto
lowerSymmetricSelection(Eigen::MatrixBase<Derived> const & matrix, Indices const & indices)
{
    if constexpr (std::is_integral_v<Indices>)
    {
        return matrix.derived().coeff(indices, indices);
    }
    else
    {
        using Scalar = typename Derived::Scalar;
        using Resolved = IndexedSize<Indices, Size, MaxSize>;
        using Result = MatrixOf<
            Scalar,
            Resolved::RowsAtCompileTime,
            Resolved::RowsAtCompileTime,
            Resolved::MaxRowsAtCompileTime,
            Resolved::MaxRowsAtCompileTime>;
        Resolved resolved = resolveIndices<Size, MaxSize>(matrix.rows(), indices);
        Eigen::Index const n = resolved.size();
        return Result::NullaryExpr(
            n,
            n,
            [&matrix = matrix.derived(), resolved = std::move(resolved)](Eigen::Index i, Eigen::Index j) {
                Eigen::Index const r = resolved(i);
                Eigen::Index const c = resolved(j);
                return r >= c ? matrix(r, c) : matrix(c, r);
            }
        );
    }
}

template <int Size, int MaxSize, typename Derived, typename Rows, typename Cols>
[[nodiscard]] auto
lowerSymmetricCross(Eigen::MatrixBase<Derived> const & matrix, Rows const & rows, Cols const & cols)
{
    if constexpr (std::is_integral_v<Rows> && std::is_integral_v<Cols>)
    {
        return rows >= cols ? matrix.derived().coeff(rows, cols) : matrix.derived().coeff(cols, rows);
    }
    else
    {
        auto const resolvedRows = resolveIndices<Size, MaxSize>(matrix.rows(), asIndexSet(rows));
        auto const resolvedCols = resolveIndices<Size, MaxSize>(matrix.rows(), asIndexSet(cols));
        using RowIndices = std::remove_cvref_t<decltype(resolvedRows)>;
        using ColIndices = std::remove_cvref_t<decltype(resolvedCols)>;
        using Result = MatrixOf<
            typename Derived::Scalar,
            RowIndices::RowsAtCompileTime,
            ColIndices::RowsAtCompileTime,
            RowIndices::MaxRowsAtCompileTime,
            ColIndices::MaxRowsAtCompileTime>;
        return Result{lowerSymmetricBlock(matrix, resolvedRows, resolvedCols)};
    }
}

template <typename Scalar, typename TargetSpace>
struct ConditionedMoments
{
    Eigen::Matrix<
        Scalar,
        valueDimensionAtCompileTime<TargetSpace>,
        1,
        Eigen::ColMajor,
        maxValueDimensionAtCompileTime<TargetSpace>,
        1>
        mean;

    MatrixOf<
        Scalar,
        valueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>,
        valueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>>
        cov;
};

template <typename Scalar, typename TargetSpace>
struct ConditionedMomentsWithEvidence
{
    ConditionedMoments<Scalar, TargetSpace> moments;  
    Scalar logLikelihood;    
    Scalar squaredDistance;  
};

template <typename Scalar, typename GivenSpace>
using GivenMatrixOf = MatrixOf<
    Scalar,
    valueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>,
    valueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>,
    maxValueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>,
    maxValueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>>;

template <typename Scalar, typename GivenSpace>
using GivenVectorOf = Eigen::Matrix<
    Scalar,
    valueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>,
    1,
    Eigen::ColMajor,
    maxValueDimensionAtCompileTime<DifferenceSpace<GivenSpace>>,
    1>;

template <
    typename Factorization,
    typename MeanDerived,
    typename CovDerived,
    typename ResidualDerived,
    typename Target,
    typename Given,
    typename TargetSpace>
[[nodiscard]] ConditionedMoments<typename MeanDerived::Scalar, TargetSpace>
conditionMomentsBy(
    Factorization const & factorization,
    Eigen::MatrixBase<MeanDerived> const & mean,
    Eigen::MatrixBase<CovDerived> const & cov,
    Target const & target,
    Given const & given,
    Eigen::MatrixBase<ResidualDerived> const & residual,
    TargetSpace const & targetSpace
)
{
    using Scalar = typename MeanDerived::Scalar;
    constexpr int GivenSize = ResidualDerived::RowsAtCompileTime;
    constexpr int GivenMaxSize = ResidualDerived::MaxRowsAtCompileTime;
    constexpr int TargetSize = valueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>;
    constexpr int TargetMaxSize = maxValueDimensionAtCompileTime<DifferenceSpace<TargetSpace>>;
    using CrossMatrix = MatrixOf<Scalar, TargetSize, GivenSize, TargetMaxSize, GivenMaxSize>;
 
    auto result = ConditionedMoments<Scalar, TargetSpace>{};
    // With nothing to condition, only the likelihood remains; MSVC can't form an empty gain.
    if constexpr (TargetSize != 0)
    {
        CrossMatrix const cross = lowerSymmetricBlock(cov, target.differences, given.differences);
        CrossMatrix const gain = transposedSolve<CrossMatrix>(factorization, cross.transpose());
        result.cov = lowerSymmetricBlock(cov, target.differences, target.differences);
        result.cov.noalias() -= multiply(gain, cross.transpose());
        if constexpr (hasIdentityReset<TargetSpace>)
        {
            result.mean = Emrys::plus(targetSpace, mean.derived()(target.values), multiply(gain, residual));
        }
        else
        {
            using TargetVector = decltype(result.mean);
            using TargetDifference = Eigen::Matrix<Scalar, TargetSize, 1, Eigen::ColMajor, TargetMaxSize, 1>;
            TargetVector const start = mean.derived()(target.values);
            TargetDifference const step = multiply(gain, residual);
            result.mean = Emrys::plus(targetSpace, start, step);
            resetCovarianceAt(targetSpace, start, step, result.mean, result.cov);
        }
    }
    return result;
}

template <
    typename MeanDerived,
    typename CovDerived,
    typename ValuesDerived,
    typename Target,
    typename Given,
    typename GivenSpace,
    typename TargetSpace>
[[nodiscard]] ConditionedMoments<typename MeanDerived::Scalar, TargetSpace>
conditionMoments(
    Eigen::MatrixBase<MeanDerived> const & mean,
    Eigen::MatrixBase<CovDerived> const & cov,
    Target const & target,
    Given const & given,
    Eigen::MatrixBase<ValuesDerived> const & values,
    GivenSpace const & givenSpace,
    TargetSpace const & targetSpace
)
{
    using Scalar = typename MeanDerived::Scalar;
    EMRYS_FILTER_ASSERT(values.size() == given.values.size());
    GivenVectorOf<Scalar, GivenSpace> const residual = Emrys::minus(givenSpace, values, mean.derived()(given.values));
    auto const ldlt =
        Eigen::LDLT<GivenMatrixOf<Scalar, GivenSpace>>{lowerSymmetricBlock(cov, given.differences, given.differences)};
    return conditionMomentsBy(ldlt, mean, cov, target, given, residual, targetSpace);
}

template <
    typename MeanDerived,
    typename CovDerived,
    typename ValuesDerived,
    typename Target,
    typename Given,
    typename GivenSpace,
    typename TargetSpace>
[[nodiscard]] std::optional<ConditionedMomentsWithEvidence<typename MeanDerived::Scalar, TargetSpace>>
conditionMomentsWithEvidence(
    Eigen::MatrixBase<MeanDerived> const & mean,
    Eigen::MatrixBase<CovDerived> const & cov,
    Target const & target,
    Given const & given,
    Eigen::MatrixBase<ValuesDerived> const & values,
    GivenSpace const & givenSpace,
    TargetSpace const & targetSpace
)
{
    using Scalar = typename MeanDerived::Scalar;
    using GivenVector = GivenVectorOf<Scalar, GivenSpace>;
    EMRYS_FILTER_ASSERT(values.size() == given.values.size());
    auto const llt =
        Eigen::LLT<GivenMatrixOf<Scalar, GivenSpace>>{lowerSymmetricBlock(cov, given.differences, given.differences)};
    if (llt.info() != Eigen::Success)
    {
        return std::nullopt;
    }
    GivenVector const residual = Emrys::minus(givenSpace, values, mean.derived()(given.values));
    GivenVector whitened = residual;
    llt.matrixL().solveInPlace(whitened);
 
    Scalar const log2Pi = Emrys::log2Pi<Scalar>();
    Scalar const logDeterminant = Scalar(2) * llt.matrixLLT().diagonal().array().log().sum();
    Scalar const squaredDistance = whitened.squaredNorm();
    return ConditionedMomentsWithEvidence<Scalar, TargetSpace>{
        conditionMomentsBy(llt, mean, cov, target, given, residual, targetSpace),
        Scalar(-0.5) * (Scalar(given.differences.size()) * log2Pi + logDeterminant + squaredDistance),
        squaredDistance
    };
}
 
}  // namespace Emrys::internal