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__