From 5a505d2377265ace64df3ed265be9f7fc29efb27 Mon Sep 17 00:00:00 2001
From: Pierre Gergondet <pierre.gergondet@gmail.com>
Date: Mon, 12 Jun 2023 18:23:19 +0900
Subject: [PATCH] Make it possible to select active joints in CoMIncPlaneConstr

---
 src/QPConstr.cpp     | 18 ++++++++++--------
 src/Tasks/QPConstr.h | 15 ++++++++++++++-
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/QPConstr.cpp b/src/QPConstr.cpp
index b7c44e86..e1b203d5 100644
--- a/src/QPConstr.cpp
+++ b/src/QPConstr.cpp
@@ -672,7 +672,8 @@ CoMIncPlaneConstr::PlaneData::PlaneData(int planeId,
 
 CoMIncPlaneConstr::CoMIncPlaneConstr(const std::vector<rbd::MultiBody> & mbs, int robotIndex, double step)
 : robotIndex_(robotIndex), alphaDBegin_(-1), dataVec_(), step_(step), nrVars_(0), nrActivated_(0), activated_(0),
-  jacCoM_(mbs[robotIndex]), AInEq_(), bInEq_()
+  jacCoM_(mbs[static_cast<size_t>(robotIndex)]), selector_(Eigen::VectorXd::Ones(jacCoM_.jacobian().cols())), AInEq_(),
+  bInEq_()
 {
 }
 
@@ -730,8 +731,8 @@ void CoMIncPlaneConstr::reset()
 
 void CoMIncPlaneConstr::updateNrPlanes()
 {
-  AInEq_.setZero(dataVec_.size(), nrVars_);
-  bInEq_.setZero(dataVec_.size());
+  AInEq_.setZero(static_cast<Eigen::DenseIndex>(dataVec_.size()), nrVars_);
+  bInEq_.setZero(static_cast<Eigen::DenseIndex>(dataVec_.size()));
 }
 
 void CoMIncPlaneConstr::updateNrVars(const std::vector<rbd::MultiBody> & /* mbs */, const SolverData & data)
@@ -745,10 +746,9 @@ void CoMIncPlaneConstr::update(const std::vector<rbd::MultiBody> & mbs,
                                const std::vector<rbd::MultiBodyConfig> & mbcs,
                                const SolverData & data)
 {
-  using namespace Eigen;
-
-  const rbd::MultiBody & mb = mbs[robotIndex_];
-  const rbd::MultiBodyConfig & mbc = mbcs[robotIndex_];
+  const rbd::MultiBody & mb = mbs[static_cast<size_t>(robotIndex_)];
+  const rbd::MultiBodyConfig & mbc = mbcs[static_cast<size_t>(robotIndex_)];
+  assert(selector_.size() == mb.nrDof());
 
   Eigen::Vector3d com = rbd::computeCoM(mb, mbc);
 
@@ -770,7 +770,7 @@ void CoMIncPlaneConstr::update(const std::vector<rbd::MultiBody> & mbs,
   nrActivated_ = 0;
   if(!activated_.empty())
   {
-    const MatrixXd & jacComMat = jacCoM_.jacobian(mb, mbc);
+    const Eigen::MatrixXd & jacComMat = jacCoM_.jacobian(mb, mbc);
     Eigen::Vector3d comSpeed = jacCoM_.velocity(mb, mbc);
     Eigen::Vector3d comNormalAcc = jacCoM_.normalAcceleration(mb, mbc, data.normalAccB(robotIndex_));
 
@@ -794,6 +794,8 @@ void CoMIncPlaneConstr::update(const std::vector<rbd::MultiBody> & mbs,
 
       // -dt*normal^T*J_com
       AInEq_.block(nrActivated_, alphaDBegin_, 1, mb.nrDof()).noalias() = -(step_ * d.normal.transpose()) * jacComMat;
+      AInEq_.block(nrActivated_, alphaDBegin_, 1, mb.nrDof()).noalias() =
+          -(step_ * d.normal.transpose()) * jacComMat * selector_.asDiagonal();
 
       // dampers + ddot + dt*normal^T*J*qdot
       bInEq_(nrActivated_) = dampers + distDot + step_ * (d.normal.dot(comNormalAcc) + distDDot);
diff --git a/src/Tasks/QPConstr.h b/src/Tasks/QPConstr.h
index 6e256fff..0a104e94 100644
--- a/src/Tasks/QPConstr.h
+++ b/src/Tasks/QPConstr.h
@@ -521,9 +521,21 @@ class TASKS_DLLAPI CoMIncPlaneConstr : public ConstraintFunction<Inequality>
   /// @return Number of plane.
   std::size_t nrPlanes() const;
 
-  /// Remove all plane.
+  /// Remove all plane
   void reset();
 
+  /** Access the joint selector
+   *
+   * By default the vector selects all joints in the robot
+   */
+  inline Eigen::VectorXd & selector() noexcept { return selector_; }
+
+  /** Access the joint selector (const)
+   *
+   * By default the vector selects all joints in the robot
+   */
+  inline const Eigen::VectorXd & selector() const noexcept { return selector_; }
+
   /// Reallocate A and b matrix.
   void updateNrPlanes();
 
@@ -583,6 +595,7 @@ class TASKS_DLLAPI CoMIncPlaneConstr : public ConstraintFunction<Inequality>
   std::vector<std::size_t> activated_;
 
   rbd::CoMJacobian jacCoM_;
+  Eigen::VectorXd selector_;
   Eigen::MatrixXd AInEq_;
   Eigen::VectorXd bInEq_;
 };