LinearOperator.h

Go to the documentation of this file.

 
#ifndef PBAT_MATH_LINEAR_OPERATOR_H
#define PBAT_MATH_LINEAR_OPERATOR_H
 
#include <concepts>
#include <exception>
#include <pbat/Aliases.h>
#include <pbat/profiling/Profiling.h>
#include <tuple>
 
namespace pbat {
namespace math {

template <class T>
concept CLinearOperator = requires(T t)
{
    {
        t.OutputDimensions()
    } -> std::convertible_to<int>;
    {
        t.InputDimensions()
    } -> std::convertible_to<int>;
    {t.Apply(VectorX{}, std::declval<VectorX&>())};
    {t.Apply(MatrixX{}, std::declval<MatrixX&>())};
    {
        t.ToMatrix()
    } -> std::convertible_to<CSCMatrix>;
};
 
template <CLinearOperator... TLinearOperators>
class LinearOperator;
 
} // namespace math
} // namespace pbat
 
namespace Eigen {
namespace internal {
 
template <pbat::math::CLinearOperator... TLinearOperators>
struct traits<pbat::math::LinearOperator<TLinearOperators...>>
{
    using Scalar       = pbat::Scalar;
    using StorageIndex = pbat::CSCMatrix::StorageIndex;
    using StorageKind  = Sparse;
    using XprKind      = MatrixXpr;
    enum {
        RowsAtCompileTime    = Dynamic,
        ColsAtCompileTime    = Dynamic,
        MaxRowsAtCompileTime = Dynamic,
        MaxColsAtCompileTime = Dynamic,
        Flags                = 0
    };
};
 
} // namespace internal
} // namespace Eigen
 
namespace pbat {
namespace math {

template <CLinearOperator... TLinearOperators>
class LinearOperator : public Eigen::EigenBase<LinearOperator<TLinearOperators...>>
{
  public:
    using SelfType = LinearOperator<TLinearOperators...>; 
    using BaseType = Eigen::EigenBase<SelfType>;          
 
    using Scalar           = pbat::Scalar;                     
    using RealScalar       = pbat::Scalar;                     
    using StorageIndex     = typename CSCMatrix::StorageIndex; 
    using NestedExpression = SelfType;                         
    enum {
        ColsAtCompileTime    = Eigen::Dynamic,
        MaxColsAtCompileTime = Eigen::Dynamic,
        RowsAtCompileTime    = Eigen::Dynamic,
        MaxRowsAtCompileTime = Eigen::Dynamic,
        IsRowMajor           = false
    };

    LinearOperator(TLinearOperators const&... inOps);
 
    SelfType& operator=(SelfType const&) = delete;

    template <class TDerivedIn, class TDerivedOut>
    void Apply(Eigen::MatrixBase<TDerivedIn> const& x, Eigen::DenseBase<TDerivedOut>& y) const;

    CSCMatrix ToMatrix() const;

    pbat::Index OutputDimensions() const;
    pbat::Index InputDimensions() const;
 
    // For Eigen compatibility

    BaseType::Index rows() const { return static_cast<BaseType::Index>(OutputDimensions()); }
    BaseType::Index cols() const { return static_cast<BaseType::Index>(InputDimensions()); }

    template <class Rhs>
    Eigen::Product<SelfType, Rhs, Eigen::AliasFreeProduct>
    operator*(Eigen::MatrixBase<Rhs> const& x) const
    {
        return Eigen::Product<SelfType, Rhs, Eigen::AliasFreeProduct>(*this, x.derived());
    }
 
  private:
    std::tuple<TLinearOperators const&...> ops;
};

template <CLinearOperator... TLinearOperators>
LinearOperator<TLinearOperators...> ComposeLinearOperators(TLinearOperators const&... inOps)
{
    return LinearOperator(inOps...);
}
 
template <CLinearOperator... TLinearOperators>
inline LinearOperator<TLinearOperators...>::LinearOperator(TLinearOperators const&... inOps)
    : ops(std::make_tuple(std::cref(inOps)...))
{
    bool const bInputDimensionsMatch = std::apply(
        [this](auto... op) -> bool {
            return ((InputDimensions() == op.InputDimensions()) and ...);
        },
        ops);
    bool const bOutputDimensionsMatch = std::apply(
        [this](auto... op) -> bool {
            return ((OutputDimensions() == op.OutputDimensions()) and ...);
        },
        ops);
    if (not(bInputDimensionsMatch and bOutputDimensionsMatch))
    {
        throw std::invalid_argument(
            "Dimensionality mismatch found in CompositeLinearOperator's linear operators.");
    }
}
 
template <CLinearOperator... TLinearOperators>
template <class TDerivedIn, class TDerivedOut>
inline void LinearOperator<TLinearOperators...>::Apply(
    Eigen::MatrixBase<TDerivedIn> const& x,
    Eigen::DenseBase<TDerivedOut>& y) const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.math.LinearOperator.Apply");
    std::apply([&](auto... op) { (op.Apply(x, y), ...); }, ops);
}
 
template <CLinearOperator... TLinearOperators>
inline CSCMatrix LinearOperator<TLinearOperators...>::ToMatrix() const
{
    PBAT_PROFILE_NAMED_SCOPE("pbat.math.LinearOperator.ToMatrix");
    CSCMatrix M(OutputDimensions(), InputDimensions());
    std::apply([&](auto... op) { ((M += op.ToMatrix()), ...); }, ops);
    return M;
}
 
template <CLinearOperator... TLinearOperators>
inline Index LinearOperator<TLinearOperators...>::OutputDimensions() const
{
    return std::get<0>(ops).OutputDimensions();
}
 
template <CLinearOperator... TLinearOperators>
inline Index LinearOperator<TLinearOperators...>::InputDimensions() const
{
    return std::get<0>(ops).InputDimensions();
}
 
} // namespace math
} // namespace pbat
 
namespace Eigen {
namespace internal {

template <pbat::math::CLinearOperator Lhs, typename Rhs, int ProductType>
struct generic_product_impl<Lhs, Rhs, SparseShape, DenseShape, ProductType>
    : generic_product_impl_base<Lhs, Rhs, generic_product_impl<Lhs, Rhs>>
{
    typedef typename Product<Lhs, Rhs>::Scalar Scalar;
 
    template <typename Dst>
    static void scaleAndAddTo(Dst& dst, Lhs const& lhs, Rhs const& rhs, Scalar const& alpha)
    {
        lhs.Apply(alpha * rhs, dst);
    }
};
 
} // namespace internal
} // namespace Eigen
 
#endif // PBAT_MATH_LINEAR_OPERATOR_H