adding getTermsSize to constraintCollection

Author: farbod-farshidian

ocs2_core/include/ocs2_core/constraint/StateConstraintCollection.h

@@ -49,8 +49,11 @@ class StateConstraintCollection : public Collection<StateConstraint> {
   ~StateConstraintCollection() override = default;
   StateConstraintCollection* clone() const override;
 
-  /** Returns the number of active constraints at given time. */
-  virtual size_t getNumConstraints(scalar_t time) const;
+  /** Returns the number of active constraints at a given time. */
+  size_t getNumConstraints(scalar_t time) const;
+
+  /** Returns the number of active constraints at a given time for each term. If a term is inactive, its size is zero. */
+  size_array_t getTermsSize(scalar_t time) const;
 
   /** Get an array of all constraints */
   virtual vector_array_t getValue(scalar_t time, const vector_t& state, const PreComputation& preComp) const;

ocs2_core/include/ocs2_core/constraint/StateInputConstraintCollection.h

@@ -49,8 +49,11 @@ class StateInputConstraintCollection : public Collection<StateInputConstraint> {
   ~StateInputConstraintCollection() override = default;
   StateInputConstraintCollection* clone() const override;
 
-  /** Returns the number of active constraints at given time. */
-  virtual size_t getNumConstraints(scalar_t time) const;
+  /** Returns the number of active constraints at a given time. */
+  size_t getNumConstraints(scalar_t time) const;
+
+  /** Returns the number of active constraints at a given time for each term. If a term is inactive, its size is zero. */
+  size_array_t getTermsSize(scalar_t time) const;
 
   /** Get an array of all constraints */
   virtual vector_array_t getValue(scalar_t time, const vector_t& state, const vector_t& input, const PreComputation& preComp) const;

ocs2_core/src/constraint/StateConstraintCollection.cpp

@@ -48,17 +48,27 @@ StateConstraintCollection* StateConstraintCollection::clone() const {
 /******************************************************************************************************/
 size_t StateConstraintCollection::getNumConstraints(scalar_t time) const {
   size_t numConstraints = 0;
-
-  // accumulate number of constraints for each constraintTerm
   for (const auto& constraintTerm : this->terms_) {
     if (constraintTerm->isActive(time)) {
       numConstraints += constraintTerm->getNumConstraints(time);
     }
   }
-
   return numConstraints;
 }
 
+/******************************************************************************************************/
+/******************************************************************************************************/
+/******************************************************************************************************/
+size_array_t StateConstraintCollection::getTermsSize(scalar_t time) const {
+  size_array_t termsSize(this->terms_.size(), 0);
+  for (size_t i = 0; i < this->terms_.size(); ++i) {
+    if (this->terms_[i]->isActive(time)) {
+      termsSize[i] = this->terms_[i]->getNumConstraints(time);
+    }
+  }
+  return termsSize;
+}
+
 /******************************************************************************************************/
 /******************************************************************************************************/
 /******************************************************************************************************/

ocs2_core/src/constraint/StateInputConstraintCollection.cpp

@@ -49,17 +49,27 @@ StateInputConstraintCollection* StateInputConstraintCollection::clone() const {
 /******************************************************************************************************/
 size_t StateInputConstraintCollection::getNumConstraints(scalar_t time) const {
   size_t numConstraints = 0;
-
-  // accumulate number of constraints for each constraintTerm
   for (const auto& constraintTerm : this->terms_) {
     if (constraintTerm->isActive(time)) {
       numConstraints += constraintTerm->getNumConstraints(time);
     }
   }
-
   return numConstraints;
 }
 
+/******************************************************************************************************/
+/******************************************************************************************************/
+/******************************************************************************************************/
+size_array_t StateInputConstraintCollection::getTermsSize(scalar_t time) const {
+  size_array_t termsSize(this->terms_.size(), 0);
+  for (size_t i = 0; i < this->terms_.size(); ++i) {
+    if (this->terms_[i]->isActive(time)) {
+      termsSize[i] = this->terms_[i]->getNumConstraints(time);
+    }
+  }
+  return termsSize;
+}
+
 /******************************************************************************************************/
 /******************************************************************************************************/
 /******************************************************************************************************/

ocs2_core/test/constraint/testConstraintCollection.cpp

@@ -27,6 +27,8 @@ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************************************************************/
 
+#include <numeric>
+
 #include <gtest/gtest.h>
 
 #include <ocs2_core/constraint/StateConstraintCollection.h>
@@ -56,6 +58,37 @@ TEST(TestConstraintCollection, numberOfConstraints) {
   EXPECT_EQ(constraintCollection.getNumConstraints(0.0), addedConstraints);
 }
 
+TEST(TestConstraintCollection, termsSize) {
+  ocs2::StateInputConstraintCollection constraintCollection;
+
+  // Initially we have zero constraints for all types
+  auto termsSize = constraintCollection.getTermsSize(0.0);
+  EXPECT_EQ(termsSize.size(), 0);
+  EXPECT_EQ(std::accumulate(termsSize.begin(), termsSize.end(), 0), 0);
+
+  // Add 2 Linear inequality constraint term, which has 2 constraints
+  constraintCollection.add("Constraint1", std::make_unique<TestDummyConstraint>());
+  constraintCollection.add("Constraint2", std::make_unique<TestDummyConstraint>());
+  constraintCollection.add("Constraint3", std::make_unique<TestDummyConstraint>());
+  auto& constraint1 = constraintCollection.get<TestDummyConstraint>("Constraint1");
+  const size_t constraint1Size = constraint1.getNumConstraints(0.0);
+
+  // Check the right constraint size
+  termsSize = constraintCollection.getTermsSize(0.0);
+  EXPECT_EQ(termsSize.size(), 3);
+  if (termsSize.size() == 3) {
+    EXPECT_EQ(termsSize[0], constraint1Size);
+    EXPECT_EQ(termsSize[1], constraint1Size);
+    EXPECT_EQ(termsSize[2], constraint1Size);
+  }
+
+  // Deactivate constraint1
+  constraint1.setActivity(false);
+  termsSize = constraintCollection.getTermsSize(0.0);
+  EXPECT_EQ(termsSize.size(), 3);
+  EXPECT_EQ(std::accumulate(termsSize.begin(), termsSize.end(), 0), 2 * constraint1Size);
+}
+
 TEST(TestConstraintCollection, activatingConstraints) {
   ocs2::StateInputConstraintCollection constraintCollection;