delassus-operator-sparse.hpp

//
// Copyright (c) 2024 INRIA
//

#ifndef __pinocchio_algorithm_delassus_operator_sparse_hpp__
#define __pinocchio_algorithm_delassus_operator_sparse_hpp__

#include "pinocchio/algorithm/fwd.hpp"
#include "pinocchio/algorithm/delassus-operator-base.hpp"

namespace pinocchio
{

 namespace internal
 {

 template<typename Derived>
 struct SimplicialCholeskyWrapper : public Derived
 {
 typedef Eigen::SimplicialCholeskyBase<Derived> Base;

 using Base::derived;
 using Base::m_diag;
 using Base::m_info;
 using Base::m_matrix;
 using Base::m_P;
 using Base::m_Pinv;

 template<typename Rhs, typename Dest, typename Temporary>
 void _solve_impl(
 const Eigen::MatrixBase<Rhs> & b,
 Eigen::MatrixBase<Dest> & dest,
 Eigen::MatrixBase<Temporary> & tmp) const
 {
 // eigen_assert(m_factorizationIsOk && "The decomposition is not in a valid state for
 // solving, you must first call either compute() or symbolic()/numeric()");
 // eigen_assert(m_matrix.rows()==b.rows());

 if (m_info != Eigen::Success)
 return;

 if (m_P.size() > 0)
 tmp.noalias() = m_P * b;
 else
 tmp = b;

 if (m_matrix.nonZeros() > 0) // otherwise L==I
 derived().matrixL().solveInPlace(tmp);

 if (m_diag.size() > 0)
 tmp = m_diag.asDiagonal().inverse() * tmp;

 if (m_matrix.nonZeros() > 0) // otherwise U==I
 derived().matrixU().solveInPlace(tmp);

 if (m_P.size() > 0)
 dest.noalias() = m_Pinv * tmp;
 }

 }; // SimplicialCholeskyWrapper

 template<typename SparseCholeskySolver>
 struct getSparseCholeskySolverBase;

 template<typename SparseCholeskySolver> //, typename Base = typename SparseCholeskySolver::Base>
 struct SparseSolveInPlaceMethod;

#ifdef PINOCCHIO_WITH_ACCELERATE_SUPPORT
 template<typename MatrixType, int UpLo, SparseFactorization_t Solver, bool EnforceSquare>
 struct SparseSolveInPlaceMethod<Eigen::AccelerateImpl<MatrixType, UpLo, Solver, EnforceSquare>>
 {
 typedef Eigen::AccelerateImpl<MatrixType, UpLo, Solver, EnforceSquare> SparseCholeskySolver;

 template<typename Rhs, typename Dest, typename Temporary>
 static void run(
 const SparseCholeskySolver & solver,
 const Eigen::MatrixBase<Rhs> & mat,
 const Eigen::MatrixBase<Dest> & dest,
 Eigen::MatrixBase<Temporary> & /*tmp*/)
 {
 dest.const_cast_derived() = solver.solve(mat.derived());
 }
 };
#endif

 template<typename SparseCholeskySolver>
 struct SparseSolveInPlaceMethod
 {
 template<typename Rhs, typename Dest, typename Temporary>
 static void run(
 const SparseCholeskySolver & solver,
 const Eigen::MatrixBase<Rhs> & mat,
 const Eigen::MatrixBase<Dest> & dest,
 Eigen::MatrixBase<Temporary> & tmp)
 {
 static_assert(
 std::is_base_of<
 Eigen::SimplicialCholeskyBase<SparseCholeskySolver>, SparseCholeskySolver>::value,
 "The solver is not a base of SimplicialCholeskyBase.");
 typedef SimplicialCholeskyWrapper<SparseCholeskySolver> CholeskyWrapper;

 const CholeskyWrapper & wrapper = reinterpret_cast<const CholeskyWrapper &>(solver);
 wrapper._solve_impl(mat, dest.const_cast_derived(), tmp.derived());
 }
 };

 } // namespace internal

 template<typename _Scalar, int _Options, class _SparseCholeskyDecomposition>
 struct traits<DelassusOperatorSparseTpl<_Scalar, _Options, _SparseCholeskyDecomposition>>
 {
 typedef _SparseCholeskyDecomposition CholeskyDecomposition;
 typedef typename CholeskyDecomposition::MatrixType SparseMatrix;
 typedef _Scalar Scalar;

 enum
 {
 Options = _Options,
 RowsAtCompileTime = Eigen::Dynamic
 };

 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1, Options> Vector;
 typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, Options> DenseMatrix;
 };

 template<typename _Scalar, int _Options, class SparseCholeskyDecomposition>
 struct DelassusOperatorSparseTpl
 : DelassusOperatorBase<DelassusOperatorSparseTpl<_Scalar, _Options, SparseCholeskyDecomposition>>
 {
 typedef DelassusOperatorSparseTpl Self;
 typedef typename traits<Self>::Scalar Scalar;
 enum
 {
 Options = traits<Self>::Options,
 RowsAtCompileTime = traits<Self>::RowsAtCompileTime
 };

 typedef typename traits<Self>::SparseMatrix SparseMatrix;
 typedef typename traits<Self>::Vector Vector;
 typedef typename traits<Self>::DenseMatrix DenseMatrix;
 typedef SparseCholeskyDecomposition CholeskyDecomposition;
 typedef DelassusOperatorBase<Self> Base;

 template<typename MatrixDerived>
 explicit DelassusOperatorSparseTpl(const Eigen::SparseMatrixBase<MatrixDerived> & mat)
 : Base(mat.rows())
 , delassus_matrix(mat)
 , delassus_matrix_plus_damping(mat)
 , llt(mat)
 , damping(Vector::Zero(mat.rows()))
 , tmp(mat.rows())
 {
 PINOCCHIO_CHECK_ARGUMENT_SIZE(mat.rows(), mat.cols());
 }

 template<typename VectorLike>
 void updateDamping(const Eigen::MatrixBase<VectorLike> & vec)
 {
 for (Eigen::DenseIndex k = 0; k < size(); ++k)
 {
 delassus_matrix_plus_damping.coeffRef(k, k) += -damping[k] + vec[k];
 }
 damping = vec;
 PINOCCHIO_EIGEN_MALLOC_SAVE_STATUS();
 PINOCCHIO_EIGEN_MALLOC_ALLOWED();
 llt.factorize(delassus_matrix_plus_damping);
 PINOCCHIO_EIGEN_MALLOC_RESTORE_STATUS();
 }

 void updateDamping(const Scalar & mu)
 {
 updateDamping(Vector::Constant(size(), mu));
 }

 template<typename MatrixLike>
 void solveInPlace(const Eigen::MatrixBase<MatrixLike> & mat) const
 {
 internal::SparseSolveInPlaceMethod<CholeskyDecomposition>::run(
 llt, mat.derived(), mat.const_cast_derived(), tmp);
 }

 template<typename MatrixLike>
 typename PINOCCHIO_EIGEN_PLAIN_TYPE(MatrixLike)
 solve(const Eigen::MatrixBase<MatrixLike> & mat) const
 {
 typename PINOCCHIO_EIGEN_PLAIN_TYPE(MatrixLike) res(mat);
 solveInPlace(res);
 return res;
 }

 template<typename MatrixDerivedIn, typename MatrixDerivedOut>
 void solve(
 const Eigen::MatrixBase<MatrixDerivedIn> & x,
 const Eigen::MatrixBase<MatrixDerivedOut> & res) const
 {
 res.const_cast_derived() = x;
 llt._solve_impl(x, res.const_cast_derived());
 }

 template<typename MatrixIn, typename MatrixOut>
 void applyOnTheRight(
 const Eigen::MatrixBase<MatrixIn> & x, const Eigen::MatrixBase<MatrixOut> & res_) const
 {
 MatrixOut & res = res_.const_cast_derived();
 res.noalias() = delassus_matrix * x;
 res.array() += damping.array() * x.array();
 }

 template<typename MatrixDerived>
 typename PINOCCHIO_EIGEN_PLAIN_TYPE(MatrixDerived)
 operator*(const Eigen::MatrixBase<MatrixDerived> & x) const
 {
 typedef typename PINOCCHIO_EIGEN_PLAIN_TYPE(MatrixDerived) ReturnType;

 PINOCCHIO_CHECK_ARGUMENT_SIZE(x.rows(), size());
 ReturnType res(x.rows(), x.cols());
 applyOnTheRight(x, res);
 return res;
 }

 Eigen::DenseIndex size() const
 {
 return delassus_matrix.rows();
 }
 Eigen::DenseIndex rows() const
 {
 return delassus_matrix.rows();
 }
 Eigen::DenseIndex cols() const
 {
 return delassus_matrix.cols();
 }

 SparseMatrix matrix() const
 {
 delassus_matrix_plus_damping = delassus_matrix;
 delassus_matrix_plus_damping += damping.asDiagonal();
 return delassus_matrix_plus_damping;
 }

 SparseMatrix inverse() const
 {
 SparseMatrix identity_matrix(size(), size());
 identity_matrix.setIdentity();
 SparseMatrix res = llt.solve(identity_matrix);
 return res;
 }

 protected:
 SparseMatrix delassus_matrix;
 mutable SparseMatrix delassus_matrix_plus_damping;
 CholeskyDecomposition llt;
 Vector damping;
 mutable Vector tmp;

 }; // struct DelassusOperatorSparseTpl

} // namespace pinocchio

#endif // ifndef __pinocchio_algorithm_delassus_operator_sparse_hpp__