Batched.h

Go to the documentation of this file.

#pragma once

 
#include <array>
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include <unsupported/Eigen/AutoDiff>
 
#include <EmrysFilter/Core/Forwarding.h>
#include <EmrysFilter/Core/Macros.h>
#include <EmrysFilter/Models/Methods.h>
#include <EmrysFilter/Spaces/Composite.h>
#include <EmrysFilter/Spaces/Space.h>
#include <EmrysFilter/Spaces/TagPath.h>
 
namespace Emrys
{

namespace jacobians
{

struct PerColumn
{};

struct Shared
{};

[[nodiscard]] constexpr PerColumn
perColumn()
{
    return {};
}

[[nodiscard]] constexpr Shared
shared()
{
    return {};
}
 
}  // namespace jacobians
 
namespace internal
{

struct NoJacobians
{};

template <typename T>
inline constexpr bool isBatchedMark = false;

template <>
inline constexpr bool isBatchedMark<NoJacobians>{true};
 
template <>
inline constexpr bool isBatchedMark<jacobians::PerColumn>{true};
 
template <>
inline constexpr bool isBatchedMark<jacobians::Shared>{true};

 
}  // namespace internal

template <typename T>
concept JacobianMark =
    internal::isBatchedMark<std::remove_cvref_t<T>> && !std::is_same_v<std::remove_cvref_t<T>, internal::NoJacobians>;

template <typename F, typename Mark = internal::NoJacobians>
struct Batched
{
    static_assert(internal::isBatchedMark<Mark>);

    using FunctionType = F;
    using MarkType = Mark;

    static constexpr bool GivesJacobians = !std::is_same_v<Mark, internal::NoJacobians>;

    static constexpr bool PerColumn = std::is_same_v<Mark, jacobians::PerColumn>;
 
    F function;  

    template <typename In, typename Out, MatrixExpression X, typename... Rest>
    [[nodiscard]] auto
    operator()(In const & in, Out const & out, X const & x, Rest const &... rest) const;
};
 
namespace internal
{

template <typename T>
inline constexpr bool isBatched = false;

template <typename F, typename Mark>
inline constexpr bool isBatched<Batched<F, Mark>>{true};


template <typename F>
inline constexpr bool takesBatches = isBatched<F> || requires { requires F::TakesBatches; };

template <typename F>
struct BatchMarkOf
{
    using type = typename F::MarkType;  
};

template <typename F, typename Mark>
struct BatchMarkOf<Batched<F, Mark>>
{
    using type = Mark;  
};


template <typename F>
using BatchMark = typename BatchMarkOf<F>::type;

template <typename F, bool = takesBatches<F>>
struct BatchMarkOrNoneOf
{
    using type = NoJacobians;  
};

template <typename F>
struct BatchMarkOrNoneOf<F, true>
{
    using type = BatchMark<F>;  
};


template <typename F>
using BatchMarkOrNone = typename BatchMarkOrNoneOf<F>::type;

template <typename F>
[[nodiscard]] auto const &
batchBody(F const & f)
{
    if constexpr (isBatched<F>)
    {
        return f.function;
    }
    else
    {
        return f;
    }
}

template <typename F>
using BatchBody = std::remove_cvref_t<decltype(batchBody(std::declval<F const &>()))>;
 
}  // namespace internal

template <typename F>
    requires(!internal::isBatched<std::remove_cvref_t<F>> && !MatrixExpression<F>)
[[nodiscard]] Batched<std::remove_cvref_t<F>>
batched(F && function)
{
    return {std::forward<F>(function)};
}

template <typename F, JacobianMark Mark>
    requires(!internal::isBatched<std::remove_cvref_t<F>> && !MatrixExpression<F>)
[[nodiscard]] Batched<std::remove_cvref_t<F>, std::remove_cvref_t<Mark>>
batched(F && function, Mark)
{
    return {std::forward<F>(function)};
}
 
namespace internal
{

template <typename Name, typename Blocks>
inline constexpr std::size_t blockPosition = 0;

template <typename Name, typename... Blocks>
inline constexpr std::size_t blockPosition<Name, std::tuple<Blocks...>> = [] {
    constexpr std::array<bool, sizeof...(Blocks) + 1> same{std::is_same_v<Name, Blocks>..., true};
    std::size_t k = 0;
    while (!same[k])
    {
        ++k;
    }
    return k;
}();


template <typename In>
using InputBlocks = typename decltype([] {
    if constexpr (IsComposite<In>)
    {
        return std::type_identity<CompositeBlocks<In>>{};
    }
    else
    {
        return std::type_identity<std::tuple<In>>{};
    }
}())::type;

template <typename Name, typename In>
[[nodiscard]] BlockRows
inputBlockRows(In const & in)
{
    if constexpr (IsComposite<In>)
    {
        return differenceBlockRows<Name>(in);
    }
    else
    {
        return BlockRows{0, differenceDimension(in)};
    }
}
 
}  // namespace internal

template <typename In, typename Out, typename Scalar_, typename Mark>
class BatchedOutput
{
    using Blocks = internal::InputBlocks<In>;
 
    static constexpr std::size_t BlockCount = std::tuple_size_v<Blocks>;  
 
public:

    using Scalar = Scalar_;

    using ValuesType = Eigen::Map<internal::MatrixOf<
        Scalar_,
        valueDimensionAtCompileTime<Out>,
        Eigen::Dynamic,
        maxValueDimensionAtCompileTime<Out>,
        Eigen::Dynamic>>;

    using JacobiansType = Eigen::Map<internal::MatrixOf<
        Scalar_,
        valueDimensionAtCompileTime<DifferenceSpace<Out>>,
        Eigen::Dynamic,
        maxValueDimensionAtCompileTime<DifferenceSpace<Out>>,
        Eigen::Dynamic>>;

    BatchedOutput(
        ValuesType values,
        JacobiansType jacobians,
        In const & in,
        std::array<bool, BlockCount> const & wanted
    )
        : m_values{values}
        , m_jacobians{jacobians}
        , m_in{in}
    {
        Eigen::Index offset = 0;
        [&]<std::size_t... K>(std::index_sequence<K...>) {
            ((m_offsets[K] = wanted[K] ? offset : -1,
              offset += wanted[K] ? internal::inputBlockRows<std::tuple_element_t<K, Blocks>>(in).size : 0),
             ...);
        }(std::make_index_sequence<BlockCount>{});
        m_wantedColumns = offset;
        EMRYS_FILTER_ASSERT(
            m_jacobians.cols() == (std::is_same_v<Mark, jacobians::PerColumn> ? values.cols() : 1) * m_wantedColumns
            && "the Jacobians' storage holds the wanted blocks' columns for each column, or once"
        );
    }

    [[nodiscard]] ValuesType &
    values()
    {
        return m_values;
    }

    [[nodiscard]] ValuesType const &
    values() const
    {
        return m_values;
    }

    template <typename Name>
    [[nodiscard]] bool
    wantsJacobian() const
    {
        checkName<Name>();
        return m_offsets[internal::blockPosition<Name, Blocks>] >= 0;
    }

    [[nodiscard]] bool
    wantsJacobians() const
    {
        return !std::is_same_v<Mark, internal::NoJacobians> && m_wantedColumns == differenceDimension(m_in);
    }

    [[nodiscard]] auto
    jacobian(Eigen::Index j)
    {
        checkWhole();
        static_assert(
            !std::is_same_v<Mark, jacobians::Shared>,
            "under jacobians::shared(), a batched function writes its whole Jacobian once for every column, jacobian()"
        );
        EMRYS_FILTER_ASSERT(j >= 0 && j < m_values.cols() && "the column must be one of the batch");
        return wholeColumns(j * m_wantedColumns);
    }

    [[nodiscard]] auto
    jacobian()
    {
        checkWhole();
        static_assert(
            !std::is_same_v<Mark, jacobians::PerColumn>,
            "under jacobians::perColumn(), a batched function writes its whole Jacobian once per column, jacobian(j)"
        );
        return wholeColumns(0);
    }

    template <typename Name>
    [[nodiscard]] auto
    jacobian(Eigen::Index j)
    {
        checkJacobian<Name>();
        static_assert(
            !std::is_same_v<Mark, jacobians::Shared>,
            "a batched function marked jacobians::shared() writes one Jacobian for every column, jacobian<Name>()"
        );
        EMRYS_FILTER_ASSERT(j >= 0 && j < m_values.cols() && "the column must be one of the batch");
        return columnsOf<Name>(j * m_wantedColumns);
    }

    template <typename Name>
    [[nodiscard]] auto
    jacobian()
    {
        checkJacobian<Name>();
        static_assert(
            !std::is_same_v<Mark, jacobians::PerColumn>,
            "a batched function marked jacobians::perColumn() writes one Jacobian per column, jacobian<Name>(j)"
        );
        return columnsOf<Name>(0);
    }

    template <typename Out2, typename In2>
    [[nodiscard]] BatchedOutput<In2, Out2, Scalar_, Mark>
    renamed(In2 const & in)
    {
        using Renamed = BatchedOutput<In2, Out2, Scalar_, Mark>;
        static_assert(std::tuple_size_v<internal::InputBlocks<In2>> == BlockCount);
        auto wanted = std::array<bool, BlockCount>{};
        for (std::size_t k = 0; k < BlockCount; ++k)
        {
            wanted[k] = m_offsets[k] >= 0;
        }
        return Renamed{
            typename Renamed::ValuesType{m_values.data(), m_values.rows(), m_values.cols()},
            typename Renamed::JacobiansType{m_jacobians.data(), m_jacobians.rows(), m_jacobians.cols()},
            in,
            wanted
        };
    }
 
private:

    template <typename Name>
    static consteval void
    checkName()
    {
        static_assert(
            internal::blockPosition<Name, Blocks> < BlockCount,
            "a batched function writes a Jacobian only for a block of its input, by the name reads() gives it"
        );
    }

    template <typename Name>
    static consteval void
    checkJacobian()
    {
        static_assert(
            !std::is_same_v<Mark, internal::NoJacobians>,
            "a batched function marked without Jacobians writes none: mark it batched(f, jacobians::perColumn()) "
            "or batched(f, jacobians::shared())"
        );
        checkName<Name>();
    }

    template <typename Name>
    [[nodiscard]] auto
    columnsOf(Eigen::Index start)
    {
        Eigen::Index const offset = m_offsets[internal::blockPosition<Name, Blocks>];
        EMRYS_FILTER_ASSERT(offset >= 0 && "write only the Jacobians wantsJacobian() says are wanted");
        constexpr int Cols = valueDimensionAtCompileTime<DifferenceSpace<Name>>;
        return internal::middleColsOf<Cols>(m_jacobians, start + offset, internal::inputBlockRows<Name>(m_in).size);
    }

    static consteval void
    checkWhole()
    {
        static_assert(
            !std::is_same_v<Mark, internal::NoJacobians>,
            "a batched function marked without Jacobians writes none: mark it batched(f, jacobians::perColumn()) "
            "or batched(f, jacobians::shared())"
        );
    }

    [[nodiscard]] auto
    wholeColumns(Eigen::Index start)
    {
        EMRYS_FILTER_ASSERT(
            m_wantedColumns == differenceDimension(m_in)
            && "write the whole Jacobian only when wantsJacobians() says every block is wanted"
        );
        constexpr int Cols = valueDimensionAtCompileTime<DifferenceSpace<In>>;
        return internal::middleColsOf<Cols>(m_jacobians, start, m_wantedColumns);
    }
 
    ValuesType m_values;                               
    JacobiansType m_jacobians;                         
    In const & m_in;                                   
    std::array<Eigen::Index, BlockCount> m_offsets{};  
    Eigen::Index m_wantedColumns = 0;                  
};
 
template <typename F, typename Mark>
template <typename In, typename Out, MatrixExpression X, typename... Rest>
auto
Batched<F, Mark>::operator()(In const & in, Out const & out, X const & x, Rest const &... rest) const
{
    static_assert(
        Space<Out>,
        "a batched function returns nothing to size its output by, so its model names the output with writes() or "
        "outputSpace()"
    );
    using T = typename X::Scalar;
    using Output = BatchedOutput<In, Out, T, Mark>;
    using Value = internal::MatrixOf<T, valueDimensionAtCompileTime<Out>, 1, maxValueDimensionAtCompileTime<Out>, 1>;
    using Jacobian = internal::MatrixOf<
        T,
        valueDimensionAtCompileTime<DifferenceSpace<Out>>,
        valueDimensionAtCompileTime<DifferenceSpace<In>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<Out>>,
        maxValueDimensionAtCompileTime<DifferenceSpace<In>>>;
    EMRYS_FILTER_ASSERT(x.rows() == valueDimension(in) && x.cols() == 1 && "the point is a vector in the input space");
    auto value = Value{Value::Zero(valueDimension(out))};
    auto jacobian = Jacobian{};
    auto wanted = std::array<bool, std::tuple_size_v<internal::InputBlocks<In>>>{};
    if constexpr (GivesJacobians)
    {
        jacobian.setZero(differenceDimension(out), differenceDimension(in));
        wanted.fill(true);
    }
    auto y = Output{
        typename Output::ValuesType{value.data(), value.rows(), 1},
        typename Output::JacobiansType{jacobian.data(), differenceDimension(out), GivesJacobians ? jacobian.cols() : 0},
        in,
        wanted
    };
    internal::callForm(function, in, out, x, rest..., y);
    if constexpr (GivesJacobians)
    {
        return withJacobian(std::move(value), std::move(jacobian));
    }
    else
    {
        return value;
    }
}
 
namespace internal
{

template <typename F, typename In, typename Out, typename X, typename Y, typename... Rest>
void
callBatched(F const & batched, In const & in, Out const & out, X const & x, Y & y, Rest const &... rest)
{
    static_assert(takesBatches<F>);
    using T = typename X::Scalar;
    using Output = BatchedOutput<In, Out, T, BatchMark<F>>;
    EMRYS_FILTER_ASSERT(
        x.rows() == valueDimension(in) && y.rows() == valueDimension(out) && y.cols() == x.cols()
        && "the batch writes a column in the output space per column in the input space"
    );
    auto output = Output{
        typename Output::ValuesType{y.data(), y.rows(), y.cols()},
        typename Output::JacobiansType{nullptr, differenceDimension(out), 0},
        in,
        {}
    };
    callForm(batchBody(batched), in, out, x, rest..., output);
}

template <typename Name, typename In>
[[nodiscard]] BlockRows
inputValueBlockRows(In const & in)
{
    if constexpr (IsComposite<In>)
    {
        return blockRows<Name>(in);
    }
    else
    {
        return BlockRows{0, valueDimension(in)};
    }
}

template <typename F, typename In, typename Out, typename Dual, typename... Rest>
inline constexpr bool batchTakesDuals = functionForm<
                                            F,
                                            In,
                                            Out,
                                            MatrixOf<
                                                Dual,
                                                valueDimensionAtCompileTime<In>,
                                                Eigen::Dynamic,
                                                maxValueDimensionAtCompileTime<In>,
                                                Eigen::Dynamic>,
                                            Rest const &...,
                                            BatchedOutput<In, Out, Dual, NoJacobians> &>
                                     != 0;

template <typename F, typename In, typename Out, typename X, std::size_t K, typename Y, typename J, typename... Rest>
void
callBatchedWithJacobians(
    F const & batched,
    In const & in,
    Out const & out,
    X const & x,
    std::array<bool, K> const & wanted,
    Y & y,
    J & jacobians,
    Rest const &... rest
)
{
    static_assert(takesBatches<F>);
    using Mark = BatchMark<F>;
    using T = typename X::Scalar;
    using Blocks = InputBlocks<In>;
    static_assert(K == std::tuple_size_v<Blocks>);
    Eigen::Index const count = x.cols();
    Eigen::Index const m = differenceDimension(out);
    Eigen::Index w = 0;
    [&]<std::size_t... B>(std::index_sequence<B...>) {
        ((w += wanted[B] ? inputBlockRows<std::tuple_element_t<B, Blocks>>(in).size : 0), ...);
    }(std::make_index_sequence<K>{});
    EMRYS_FILTER_ASSERT(
        x.rows() == valueDimension(in) && y.rows() == valueDimension(out) && y.cols() == count && jacobians.rows() == m
        && jacobians.cols() == w * count
        && "the batch writes a value and a Jacobian over the wanted blocks per column of its input"
    );
    if constexpr (!std::is_same_v<Mark, NoJacobians>)
    {
        using Output = BatchedOutput<In, Out, T, Mark>;
        constexpr bool PerColumn = std::is_same_v<Mark, jacobians::PerColumn>;
        // A shared Jacobian is written once, beside the values, and then copied to every column.
        using Shared = MatrixOf<
            T,
            valueDimensionAtCompileTime<DifferenceSpace<Out>>,
            Eigen::Dynamic,
            maxValueDimensionAtCompileTime<DifferenceSpace<Out>>,
            maxValueDimensionAtCompileTime<DifferenceSpace<In>>>;
        auto shared = Shared{};
        if constexpr (!PerColumn)
        {
            shared.resize(m, w);
        }
        auto output = Output{
            typename Output::ValuesType{y.data(), y.rows(), count},
            typename Output::JacobiansType{PerColumn ? jacobians.data() : shared.data(), m, PerColumn ? w * count : w},
            in,
            wanted
        };
        callForm(batchBody(batched), in, out, x, rest..., output);
        if constexpr (!PerColumn)
        {
            for (Eigen::Index j = 0; j < count; ++j)
            {
                jacobians.middleCols(j * w, w) = shared;
            }
        }
    }
    else
    {
        using Derivatives = Eigen::
            Matrix<T, Eigen::Dynamic, 1, Eigen::ColMajor, maxValueDimensionAtCompileTime<DifferenceSpace<In>>, 1>;
        using Dual = Eigen::AutoDiffScalar<Derivatives>;
        static_assert(
            batchTakesDuals<BatchBody<F>, In, Out, Dual, Rest...>,
            "a batched function whose Jacobian a method uses is differentiated in dual numbers, so it takes X as "
            "auto const & and writes Y in X's scalar; or it writes its Jacobians, batched(f, jacobians::perColumn()) "
            "or batched(f, jacobians::shared())"
        );
        using DualX = MatrixOf<
            Dual,
            valueDimensionAtCompileTime<In>,
            X::ColsAtCompileTime,
            maxValueDimensionAtCompileTime<In>,
            X::MaxColsAtCompileTime>;
        using DualY = MatrixOf<
            Dual,
            valueDimensionAtCompileTime<Out>,
            Y::ColsAtCompileTime,
            maxValueDimensionAtCompileTime<Out>,
            Y::MaxColsAtCompileTime>;
        auto xd = DualX{};
        xd.resize(x.rows(), count);
        for (Eigen::Index j = 0; j < count; ++j)
        {
            for (Eigen::Index r = 0; r < x.rows(); ++r)
            {
                xd(r, j) = Dual{x(r, j), Derivatives::Zero(w)};
            }
        }
        Eigen::Index offset = 0;
        [&]<std::size_t... B>(std::index_sequence<B...>) {
            (
                [&] {
                    using Name = std::tuple_element_t<B, Blocks>;
                    if (wanted[B])
                    {
                        BlockRows const values = inputValueBlockRows<Name>(in);
                        BlockRows const differences = inputBlockRows<Name>(in);
                        EMRYS_FILTER_ASSERT(
                            isLinear<Name> && values.size == differences.size
                            && "a block differentiated in a batch is linear, its values its differences"
                        );
                        for (Eigen::Index r = 0; r < values.size; ++r)
                        {
                            for (Eigen::Index j = 0; j < count; ++j)
                            {
                                xd(values.start + r, j).derivatives()(offset + r) = T(1);
                            }
                        }
                        offset += differences.size;
                    }
                }(),
                ...
            );
        }(std::make_index_sequence<K>{});
        auto yd = DualY{};
        yd.resize(y.rows(), count);
        callBatched(batched, in, out, xd, yd, rest...);
        using Deviation = MatrixOf<
            Dual,
            valueDimensionAtCompileTime<DifferenceSpace<Out>>,
            1,
            maxValueDimensionAtCompileTime<DifferenceSpace<Out>>,
            1>;
        for (Eigen::Index j = 0; j < count; ++j)
        {
            for (Eigen::Index r = 0; r < y.rows(); ++r)
            {
                y(r, j) = yd(r, j).value();
            }
            Deviation const deviation = [&] {
                if constexpr (HasMinus<Out>)
                {
                    return Deviation{Emrys::minus(out, yd.col(j), y.col(j).template cast<Dual>())};
                }
                else
                {
                    return Deviation{yd.col(j)};
                }
            }();
            for (Eigen::Index r = 0; r < m; ++r)
            {
                // An output that does not depend on the wanted blocks may carry no derivatives.
                if (deviation(r).derivatives().size() > 0)
                {
                    jacobians.block(r, j * w, 1, w) = deviation(r).derivatives().transpose();
                }
                else
                {
                    jacobians.block(r, j * w, 1, w).setZero();
                }
            }
        }
    }
}
 
}  // namespace internal
 
}  // namespace Emrys