From ea5e00cad8c62a962cd75cf274a69a69d0272da2 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Tue, 29 Aug 2023 07:40:06 -0700 Subject: [PATCH] Move mju_combineSparseCount into engine_util_sparse. PiperOrigin-RevId: 561025741 Change-Id: I0bdc48554928fc4545a38a8cea3bdeb593bd7ecd --- src/engine/engine_core_constraint.c | 27 -------- src/engine/engine_core_constraint.h | 7 -- src/engine/engine_util_sparse.c | 27 ++++++++ src/engine/engine_util_sparse.h | 3 + test/engine/engine_core_constraint_test.cc | 72 --------------------- test/engine/engine_util_sparse_test.cc | 74 ++++++++++++++++++++++ 6 files changed, 104 insertions(+), 106 deletions(-) diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index ff0e8cb0..b8e010b4 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -1351,33 +1351,6 @@ void mj_makeImpedance(const mjModel* m, mjData* d) { //------------------------------------- constraint counting ---------------------------------------- -// count the number of non-zeros in the sum of two sparse vectors -int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) { - int a = 0, b = 0, c_nnz = 0; - - // count c_nnz: nonzero indices common to both a and b - while (a < a_nnz && b < b_nnz) { - // common index, increment everything - if (a_ind[a] == b_ind[b]) { - c_nnz++; - a++; - b++; - } - - // update smallest index - else if (a_ind[a] < b_ind[b]) { - a++; - } else { - b++; - } - } - - // union minus the intersection - return a_nnz + b_nnz - c_nnz; -} - - - // count the non-zero columns in the Jacobian difference of two bodies static int mj_jacDifPairCount(const mjModel* m, int* chain, int b1, int b2, int issparse) { diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index 5ef90e67..f90d1322 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -93,13 +93,6 @@ void mj_diagApprox(const mjModel* m, mjData* d); // compute efc_R, efc_D, efc_KDIP, adjust diagApprox void mj_makeImpedance(const mjModel* m, mjData* d); - -//------------------------- constraint counting - -// count the number of non-zeros in the sum of two sparse vectors -MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind); - - //---------------------------- top-level API for constraint construction --------------------------- // main driver: call all functions above diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index e9925646..a561c4a8 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -207,6 +207,33 @@ static int mju_compare(const int* vec1, const int* vec2, int n) { +// count the number of non-zeros in the sum of two sparse vectors +int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) { + int a = 0, b = 0, c_nnz = 0; + + // count c_nnz: nonzero indices common to both a and b + while (a < a_nnz && b < b_nnz) { + // common index, increment everything + if (a_ind[a] == b_ind[b]) { + c_nnz++; + a++; + b++; + } + + // update smallest index + else if (a_ind[a] < b_ind[b]) { + a++; + } else { + b++; + } + } + + // union minus the intersection + return a_nnz + b_nnz - c_nnz; +} + + + // combine two sparse vectors: dst = a*dst + b*src, return nnz of result int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b, int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind, diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index a4852e63..c864d25d 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -51,6 +51,9 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind); +// count the number of non-zeros in the sum of two sparse vectors +MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind); + // combine two sparse vectors: dst = a*dst + b*src, return nnz of result int mju_combineSparse(mjtNum* dst, const mjtNum* src, mjtNum a, mjtNum b, int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind, diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index 7d47545b..8022f163 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -198,77 +198,5 @@ TEST_F(CoreConstraintTest, JacobianPreAllocate) { } } -TEST_F(CoreConstraintTest, CombineSparseCount) { - { - std::array a_ind{0, 1}; - std::array b_ind{2}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3); - } - - { - std::array a_ind{2}; - std::array b_ind{0, 1}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3); - } - - { - std::array a_ind{0, 1}; - std::array b_ind{2, 3, 4}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); - } - - { - std::array a_ind{5, 6}; - std::array b_ind{1, 3, 8}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); - } - - { - std::array a_ind{1, 2, 3}; - std::array b_ind{0, 4}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); - } - - { - std::array a_ind{1, 4}; - std::array b_ind{2, 3}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); - } - - { - std::array a_ind{0, 1, 3}; - std::array b_ind{0, 3, 4}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); - } - - { - std::array a_ind{1, 3, 5, 6}; - std::array b_ind{1, 3, 5, 6}; - EXPECT_EQ(mju_combineSparseCount( - a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); - } - - EXPECT_EQ(mju_combineSparseCount(0, 0, nullptr, nullptr), 0); - - { - std::array b_ind{1, 2}; - EXPECT_EQ( - mju_combineSparseCount(0, b_ind.size(), nullptr, b_ind.data()), 2); - } - - { - std::array a_ind{0}; - EXPECT_EQ( - mju_combineSparseCount(a_ind.size(), 0, a_ind.data(), nullptr), 1); - } -} - } // namespace } // namespace mujoco diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 0bbde8a3..668fda6f 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -14,6 +14,8 @@ // Tests for engine/engine_util_sparse.c +#include + #include "src/engine/engine_util_sparse.h" #include @@ -43,6 +45,78 @@ TEST_F(EngineUtilSparseTest, MjuDot) { EXPECT_EQ(mju_dotSparse(a, b, 7, i), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7); } +TEST_F(EngineUtilSparseTest, CombineSparseCount) { + { + std::array a_ind{0, 1}; + std::array b_ind{2}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3); + } + + { + std::array a_ind{2}; + std::array b_ind{0, 1}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 3); + } + + { + std::array a_ind{0, 1}; + std::array b_ind{2, 3, 4}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); + } + + { + std::array a_ind{5, 6}; + std::array b_ind{1, 3, 8}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); + } + + { + std::array a_ind{1, 2, 3}; + std::array b_ind{0, 4}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 5); + } + + { + std::array a_ind{1, 4}; + std::array b_ind{2, 3}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); + } + + { + std::array a_ind{0, 1, 3}; + std::array b_ind{0, 3, 4}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); + } + + { + std::array a_ind{1, 3, 5, 6}; + std::array b_ind{1, 3, 5, 6}; + EXPECT_EQ(mju_combineSparseCount( + a_ind.size(), b_ind.size(), a_ind.data(), b_ind.data()), 4); + } + + EXPECT_EQ(mju_combineSparseCount(0, 0, nullptr, nullptr), 0); + + { + std::array b_ind{1, 2}; + EXPECT_EQ( + mju_combineSparseCount(0, b_ind.size(), nullptr, b_ind.data()), 2); + } + + { + std::array a_ind{0}; + EXPECT_EQ( + mju_combineSparseCount(a_ind.size(), 0, a_ind.data(), nullptr), 1); + } +} + TEST_F(EngineUtilSparseTest, MjuTranspose3by3) { // 1 2 0 1 0 0 // 0 1 0 --> 2 1 3