tridiagonal-matrix.hpp

//
// Copyright (c) 2024 INRIA
//
 
#ifndef __pinocchio_math_tridiagonal_matrix_hpp__
#define __pinocchio_math_tridiagonal_matrix_hpp__
 
#include "pinocchio/fwd.hpp"
#include <Eigen/Dense>
 
namespace pinocchio
{
  template<typename Scalar, int Options = 0>
  struct TridiagonalSymmetricMatrixTpl;
 
  template<typename _Scalar, int _Options>
  struct traits<TridiagonalSymmetricMatrixTpl<_Scalar, _Options>>
  {
    typedef _Scalar Scalar;
    enum
    {
      Options = _Options
    };
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Options> PlainMatrixType;
  };
 
  template<typename TridiagonalSymmetricMatrix, typename MatrixDerived>
  struct TridiagonalSymmetricMatrixApplyOnTheRightReturnType;
 
  template<typename TridiagonalSymmetricMatrix, typename MatrixDerived>
  struct traits<
    TridiagonalSymmetricMatrixApplyOnTheRightReturnType<TridiagonalSymmetricMatrix, MatrixDerived>>
  : public traits<TridiagonalSymmetricMatrix>
  {
  };
 
  template<typename MatrixDerived, typename TridiagonalSymmetricMatrix>
  struct TridiagonalSymmetricMatrixApplyOnTheLeftReturnType;
 
  template<typename MatrixDerived, typename TridiagonalSymmetricMatrix>
  struct traits<
    TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<MatrixDerived, TridiagonalSymmetricMatrix>>
  : public traits<TridiagonalSymmetricMatrix>
  {
  };
 
  template<typename TridiagonalSymmetricMatrix>
  struct TridiagonalSymmetricMatrixInverse;
 
  template<typename TridiagonalSymmetricMatrix>
  struct traits<TridiagonalSymmetricMatrixInverse<TridiagonalSymmetricMatrix>>
  : public traits<TridiagonalSymmetricMatrix>
  {
  };
 
  template<typename TridiagonalSymmetricMatrixInverse, typename MatrixDerived>
  struct TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType;
 
  template<typename TridiagonalSymmetricMatrixInverse, typename MatrixDerived>
  struct traits<TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<
    TridiagonalSymmetricMatrixInverse,
    MatrixDerived>> : public traits<TridiagonalSymmetricMatrixInverse>
  {
  };
} // namespace pinocchio
 
namespace Eigen
{
  namespace internal
  {
 
    template<typename Scalar, int Options>
    struct traits<pinocchio::TridiagonalSymmetricMatrixTpl<Scalar, Options>>
    : public traits<typename pinocchio::traits<
        pinocchio::TridiagonalSymmetricMatrixTpl<Scalar, Options>>::PlainMatrixType>
    {
      typedef pinocchio::traits<pinocchio::TridiagonalSymmetricMatrixTpl<Scalar, Options>> Base;
      typedef typename Base::PlainMatrixType ReturnType;
      enum
      {
        Flags = 0
      };
    };
 
    template<typename TridiagonalSymmetricMatrix, typename MatrixDerived>
    struct traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheRightReturnType<
      TridiagonalSymmetricMatrix,
      MatrixDerived>>
    : public traits<
        typename pinocchio::traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheRightReturnType<
          TridiagonalSymmetricMatrix,
          MatrixDerived>>::PlainMatrixType>
    {
      typedef pinocchio::traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheRightReturnType<
        TridiagonalSymmetricMatrix,
        MatrixDerived>>
        Base;
      typedef typename Base::PlainMatrixType ReturnType;
      enum
      {
        Flags = 0
      };
    };
 
    template<typename MatrixDerived, typename TridiagonalSymmetricMatrix>
    struct traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<
      MatrixDerived,
      TridiagonalSymmetricMatrix>>
    : public traits<
        typename pinocchio::traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<
          MatrixDerived,
          TridiagonalSymmetricMatrix>>::PlainMatrixType>
    {
      typedef pinocchio::traits<pinocchio::TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<
        MatrixDerived,
        TridiagonalSymmetricMatrix>>
        Base;
      typedef typename Base::PlainMatrixType ReturnType;
      enum
      {
        Flags = 0
      };
    };
 
    template<typename TridiagonalSymmetricMatrix>
    struct traits<pinocchio::TridiagonalSymmetricMatrixInverse<TridiagonalSymmetricMatrix>>
    : public traits<TridiagonalSymmetricMatrix>
    {
    };
 
    template<typename TridiagonalSymmetricMatrixInverse, typename MatrixDerived>
    struct traits<pinocchio::TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<
      TridiagonalSymmetricMatrixInverse,
      MatrixDerived>>
    : public traits<typename pinocchio::traits<
        pinocchio::TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<
          TridiagonalSymmetricMatrixInverse,
          MatrixDerived>>::PlainMatrixType>
    {
      typedef pinocchio::traits<
        pinocchio::TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<
          TridiagonalSymmetricMatrixInverse,
          MatrixDerived>>
        Base;
      typedef typename Base::PlainMatrixType ReturnType;
      enum
      {
        Flags = 0
      };
    };
 
  } // namespace internal
} // namespace Eigen
 
namespace pinocchio
{
 
  template<typename _Scalar, int _Options>
  struct TridiagonalSymmetricMatrixTpl
  : public Eigen::ReturnByValue<TridiagonalSymmetricMatrixTpl<_Scalar, _Options>>
  {
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    typedef TridiagonalSymmetricMatrixTpl Self;
    typedef _Scalar Scalar;
    enum
    {
      Options = _Options
    };
 
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Options> CoeffVectorType;
    typedef typename traits<Self>::PlainMatrixType PlainMatrixType;
 
    explicit TridiagonalSymmetricMatrixTpl(const Eigen::DenseIndex size)
    : m_size(size)
    , m_diagonal(size)
    , m_sub_diagonal(size - 1)
    {
      assert(size > 0 && "size should be greater than 0.");
    }
 
    bool operator==(const TridiagonalSymmetricMatrixTpl & other) const
    {
      if (this == &other)
        return true;
      return diagonal() == other.diagonal() && subDiagonal() == other.subDiagonal();
    }
 
    bool operator!=(const TridiagonalSymmetricMatrixTpl & other) const
    {
      return !(*this == other);
    }
 
    TridiagonalSymmetricMatrixInverse<Self> inverse() const
    {
      return TridiagonalSymmetricMatrixInverse<Self>(*this);
    }
 
    CoeffVectorType & diagonal()
    {
      return m_diagonal;
    }
    const CoeffVectorType & diagonal() const
    {
      return m_diagonal;
    }
    CoeffVectorType & subDiagonal()
    {
      return m_sub_diagonal;
    }
    const CoeffVectorType & subDiagonal() const
    {
      return m_sub_diagonal;
    }
 
    void setIdentity()
    {
      diagonal().setOnes();
      subDiagonal().setZero();
    }
 
    bool isIdentity(const Scalar prec = Eigen::NumTraits<Scalar>::dummy_precision()) const
    {
      return subDiagonal().isZero(prec) && diagonal().isOnes(prec);
    }
 
    void setZero()
    {
      diagonal().setZero();
      subDiagonal().setZero();
    }
 
    bool isZero(const Scalar prec = Eigen::NumTraits<Scalar>::dummy_precision()) const
    {
      return subDiagonal().isZero(prec) && diagonal().isZero(prec);
    }
 
    void setRandom()
    {
      diagonal().setRandom();
      subDiagonal().setRandom();
    }
 
    bool isDiagonal(const Scalar prec = Eigen::NumTraits<Scalar>::dummy_precision()) const
    {
      return subDiagonal().isZero(prec);
    }
 
    template<typename VectorCoeffType>
    void setDiagonal(const Eigen::MatrixBase<VectorCoeffType> & diagonal_coefficients)
    {
      PINOCCHIO_CHECK_ARGUMENT_SIZE(diagonal_coefficients.size(), cols());
      static_assert(
        VectorCoeffType::IsVectorAtCompileTime,
        "VectorCoeffType should be a vector type at compile time");
 
      diagonal() = diagonal_coefficients;
      subDiagonal().setZero();
    }
 
    EIGEN_CONSTEXPR Eigen::Index rows() const EIGEN_NOEXCEPT
    {
      return m_size;
    }
    EIGEN_CONSTEXPR Eigen::Index cols() const EIGEN_NOEXCEPT
    {
      return m_size;
    }
 
    PlainMatrixType matrix() const
    {
      return PlainMatrixType(*this);
    }
 
    template<typename ResultType>
    inline void evalTo(ResultType & result) const
    {
      result.setZero();
      result.template diagonal<1>() = subDiagonal().conjugate();
      result.diagonal() = diagonal();
      result.template diagonal<-1>() = subDiagonal();
    }
 
    template<typename MatrixDerived>
    TridiagonalSymmetricMatrixApplyOnTheRightReturnType<Self, MatrixDerived>
    applyOnTheRight(const Eigen::MatrixBase<MatrixDerived> & mat) const
    {
      typedef TridiagonalSymmetricMatrixApplyOnTheRightReturnType<Self, MatrixDerived> ReturnType;
      return ReturnType(*this, mat.derived());
    }
 
    template<typename MatrixDerived>
    TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<MatrixDerived, Self>
    applyOnTheLeft(const Eigen::MatrixBase<MatrixDerived> & mat) const
    {
      typedef TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<MatrixDerived, Self> ReturnType;
      return ReturnType(mat.derived(), *this);
    }
 
    template<typename MatrixDerived>
    inline TridiagonalSymmetricMatrixApplyOnTheRightReturnType<Self, MatrixDerived>
    operator*(const Eigen::MatrixBase<MatrixDerived> & mat) const
    {
      return this->applyOnTheRight(mat.derived());
    }
 
  protected:
    Eigen::DenseIndex m_size;
    CoeffVectorType m_diagonal;
    CoeffVectorType m_sub_diagonal;
  };
 
  template<typename LhsMatrixType, typename S, int O>
  TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<
    LhsMatrixType,
    TridiagonalSymmetricMatrixTpl<S, O>>
  operator*(
    const Eigen::MatrixBase<LhsMatrixType> & lhs, const TridiagonalSymmetricMatrixTpl<S, O> & rhs)
  {
    return rhs.applyOnTheLeft(lhs);
  }
 
  template<typename TridiagonalSymmetricMatrix, typename RhsMatrixType>
  struct TridiagonalSymmetricMatrixApplyOnTheRightReturnType
  : public Eigen::ReturnByValue<TridiagonalSymmetricMatrixApplyOnTheRightReturnType<
      TridiagonalSymmetricMatrix,
      RhsMatrixType>>
  {
    typedef TridiagonalSymmetricMatrixApplyOnTheRightReturnType Self;
    typedef typename traits<Self>::PlainMatrixType PlainMatrixType;
 
    TridiagonalSymmetricMatrixApplyOnTheRightReturnType(
      const TridiagonalSymmetricMatrix & lhs, const RhsMatrixType & rhs)
    : m_lhs(lhs)
    , m_rhs(rhs)
    {
    }
 
    template<typename ResultType>
    inline void evalTo(ResultType & result) const
    {
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.rows(), rows());
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.cols(), cols());
 
      assert(cols() >= 1);
      assert(rows() >= 1);
 
      const Eigen::DenseIndex reduced_size = cols() - 1;
      // Main diagonal
      result.noalias() = m_lhs.diagonal().asDiagonal() * m_rhs;
      // Upper diagonal
      result.topRows(reduced_size).noalias() +=
        m_lhs.subDiagonal().conjugate().asDiagonal() * m_rhs.bottomRows(reduced_size);
      // Sub diagonal
      result.bottomRows(reduced_size).noalias() +=
        m_lhs.subDiagonal().asDiagonal() * m_rhs.topRows(reduced_size);
    }
 
    EIGEN_CONSTEXPR Eigen::Index rows() const EIGEN_NOEXCEPT
    {
      return m_lhs.rows();
    }
    EIGEN_CONSTEXPR Eigen::Index cols() const EIGEN_NOEXCEPT
    {
      return m_rhs.cols();
    }
 
  protected:
    const TridiagonalSymmetricMatrix & m_lhs;
    const RhsMatrixType & m_rhs;
  };
 
  template<typename LhsMatrixType, typename TridiagonalSymmetricMatrix>
  struct TridiagonalSymmetricMatrixApplyOnTheLeftReturnType
  : public Eigen::ReturnByValue<
      TridiagonalSymmetricMatrixApplyOnTheLeftReturnType<LhsMatrixType, TridiagonalSymmetricMatrix>>
  {
    typedef TridiagonalSymmetricMatrixApplyOnTheLeftReturnType Self;
    typedef typename traits<Self>::PlainMatrixType PlainMatrixType;
 
    TridiagonalSymmetricMatrixApplyOnTheLeftReturnType(
      const LhsMatrixType & lhs, const TridiagonalSymmetricMatrix & rhs)
    : m_lhs(lhs)
    , m_rhs(rhs)
    {
    }
 
    template<typename ResultType>
    inline void evalTo(ResultType & result) const
    {
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.rows(), rows());
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.cols(), cols());
 
      assert(cols() >= 1);
      assert(rows() >= 1);
 
      const Eigen::DenseIndex reduced_size = cols() - 1;
      // Main diagonal
      result.noalias() = m_lhs * m_rhs.diagonal().asDiagonal();
      // Upper diagonal
      result.rightCols(reduced_size).noalias() +=
        m_lhs.leftCols(reduced_size) * m_rhs.subDiagonal().conjugate().asDiagonal();
      // Sub diagonal
      result.leftCols(reduced_size).noalias() +=
        m_lhs.rightCols(reduced_size) * m_rhs.subDiagonal().asDiagonal();
    }
 
    EIGEN_CONSTEXPR Eigen::Index rows() const EIGEN_NOEXCEPT
    {
      return m_lhs.rows();
    }
    EIGEN_CONSTEXPR Eigen::Index cols() const EIGEN_NOEXCEPT
    {
      return m_rhs.cols();
    }
 
  protected:
    const LhsMatrixType & m_lhs;
    const TridiagonalSymmetricMatrix & m_rhs;
  };
 
  template<typename _TridiagonalSymmetricMatrix>
  struct TridiagonalSymmetricMatrixInverse
  : public Eigen::ReturnByValue<TridiagonalSymmetricMatrixInverse<_TridiagonalSymmetricMatrix>>
  {
    typedef _TridiagonalSymmetricMatrix TridiagonalSymmetricMatrix;
    typedef TridiagonalSymmetricMatrixInverse Self;
    typedef typename TridiagonalSymmetricMatrix::Scalar Scalar;
    enum
    {
      Options = TridiagonalSymmetricMatrix::Options
    };
 
    typedef typename TridiagonalSymmetricMatrix::CoeffVectorType CoeffVectorType;
    typedef typename traits<Self>::PlainMatrixType PlainMatrixType;
 
    explicit TridiagonalSymmetricMatrixInverse(
      const TridiagonalSymmetricMatrix & tridiagonal_matrix)
    : tridiagonal_matrix(tridiagonal_matrix)
    , m_size(tridiagonal_matrix.rows())
    , m_diagonal(m_size)
    , m_sub_diagonal(m_size - 1)
    {
      eval();
    }
 
    const TridiagonalSymmetricMatrix & inverse() const
    {
      return tridiagonal_matrix;
    }
 
    template<typename MatrixDerived>
    TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<Self, MatrixDerived>
    applyOnTheRight(const Eigen::MatrixBase<MatrixDerived> & mat) const
    {
      typedef TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<Self, MatrixDerived>
        ReturnType;
      return ReturnType(*this, mat.derived());
    }
 
    template<typename MatrixDerived>
    inline TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<Self, MatrixDerived>
    operator*(const Eigen::MatrixBase<MatrixDerived> & mat) const
    {
      return this->applyOnTheRight(mat.derived());
    }
 
    template<typename ResultType>
    inline void evalTo(ResultType & result) const
    {
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.rows(), rows());
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.cols(), cols());
 
      assert(cols() >= 1);
      assert(rows() >= 1);
 
      const auto & b = m_diagonal;
      const auto & c = tridiagonal_matrix.subDiagonal();
      const auto & w = m_sub_diagonal;
 
      // Forward sweep
      result.setIdentity();
      for (Eigen::DenseIndex i = 1; i < m_size; ++i)
      {
        result.row(i).head(i) -= w[i - 1] * result.row(i - 1).head(i);
      }
 
      // Backward sweep
      result.row(m_size - 1) /= b[m_size - 1];
      for (Eigen::DenseIndex i = m_size - 2; i >= 0; --i)
      {
        result.row(i) -= c[i] * result.row(i + 1);
        result.row(i) /= b[i];
      }
    }
 
    EIGEN_CONSTEXPR Eigen::Index rows() const EIGEN_NOEXCEPT
    {
      return m_size;
    }
    EIGEN_CONSTEXPR Eigen::Index cols() const EIGEN_NOEXCEPT
    {
      return m_size;
    }
 
  protected:
    template<typename T, typename MatrixDerived>
    friend struct TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType;
 
    void eval()
    {
      m_diagonal = tridiagonal_matrix.diagonal();
      m_sub_diagonal = tridiagonal_matrix.subDiagonal();
      auto & w = m_sub_diagonal;
      auto & b = m_diagonal;
      const auto & c = tridiagonal_matrix.subDiagonal();
      for (Eigen::DenseIndex i = 1; i < m_size; ++i)
      {
        w.coeffRef(i - 1) /= b[i - 1];
        m_diagonal.coeffRef(i) -= w[i - 1] * c[i - 1];
      }
    }
 
    const TridiagonalSymmetricMatrix & tridiagonal_matrix;
    Eigen::DenseIndex m_size;
    CoeffVectorType m_diagonal;
    CoeffVectorType m_sub_diagonal;
  };
 
  template<typename TridiagonalSymmetricMatrixInverse, typename RhsMatrixType>
  struct TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType
  : public Eigen::ReturnByValue<TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType<
      TridiagonalSymmetricMatrixInverse,
      RhsMatrixType>>
  {
    typedef TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType Self;
    typedef typename traits<Self>::PlainMatrixType PlainMatrixType;
 
    TridiagonalSymmetricMatrixInverseApplyOnTheRightReturnType(
      const TridiagonalSymmetricMatrixInverse & lhs, const RhsMatrixType & rhs)
    : m_lhs(lhs)
    , m_rhs(rhs)
    {
    }
 
    template<typename ResultType>
    inline void evalTo(ResultType & result) const
    {
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.rows(), rows());
      PINOCCHIO_CHECK_ARGUMENT_SIZE(result.cols(), cols());
 
      assert(cols() >= 1);
      assert(rows() >= 1);
 
      const Eigen::DenseIndex size = m_lhs.rows();
      const auto & b = m_lhs.m_diagonal;
      const auto & c = m_lhs.tridiagonal_matrix.subDiagonal();
      const auto & w = m_lhs.m_sub_diagonal;
 
      // Forward sweep
      result = m_rhs;
      for (Eigen::DenseIndex i = 1; i < size; ++i)
      {
        result.row(i) -= w[i - 1] * result.row(i - 1);
      }
 
      // Backward sweep
      result.row(size - 1) /= b[size - 1];
      for (Eigen::DenseIndex i = size - 2; i >= 0; --i)
      {
        result.row(i) -= c[i] * result.row(i + 1);
        result.row(i) /= b[i];
      }
    }
 
    EIGEN_CONSTEXPR Eigen::Index rows() const EIGEN_NOEXCEPT
    {
      return m_lhs.rows();
    }
    EIGEN_CONSTEXPR Eigen::Index cols() const EIGEN_NOEXCEPT
    {
      return m_rhs.cols();
    }
 
  protected:
    const TridiagonalSymmetricMatrixInverse & m_lhs;
    const RhsMatrixType & m_rhs;
  };
} // namespace pinocchio
 
#endif // #ifndef __pinocchio_math_tridiagonal_matrix_hpp__