From a063e2f0ef5b3314ab09b7d30333f2a7001e4277 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 2 Aug 2023 08:42:25 -0700 Subject: [PATCH] Add private functions and associated tests related to constraint island discovery. PiperOrigin-RevId: 553156672 Change-Id: Ifd73e3df443f7955359c5243d5f7756069335c05 --- src/engine/engine_core_constraint.c | 129 +++++++++++- src/engine/engine_core_constraint.h | 7 + src/engine/engine_util_sparse.h | 4 +- test/engine/engine_core_constraint_test.cc | 233 ++++++++++++++++++++- 4 files changed, 369 insertions(+), 4 deletions(-) diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 433f50fc..709843dd 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -23,6 +23,7 @@ #include #include #include "engine/engine_array_safety.h" +#include "engine/engine_crossplatform.h" #include "engine/engine_core_smooth.h" #include "engine/engine_io.h" #include "engine/engine_support.h" @@ -1486,7 +1487,7 @@ static inline int mj_nl(const mjModel* m, const mjData* d, int *nnz) { margin = m->jnt_margin[i]; - // slider and hinge joint limits can be bilateral, check both side + // slider and hinge joint limits can be bilateral, check both sides if (m->jnt_type[i] == mjJNT_SLIDE || m->jnt_type[i] == mjJNT_HINGE) { value = d->qpos[m->jnt_qposadr[i]]; for (side=-1; side <= 1; side+=2) { @@ -2095,3 +2096,129 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, *cost = s; } } + + + +//---------------------------- constraint islands -------------------------------------------------- + +// comparison function for lexicographic edge sorting +quicksortfunc(edgecompare, context, edge0, edge1) { + int* e0 = (int*)edge0; + int* e1 = (int*)edge1; + int v00 = e0[0]; + int v10 = e1[0]; + + if (v00 < v10) { + return -1; + } + + if (v00 == v10) { + int v01 = e0[1]; + int v11 = e1[1]; + + if (v01 < v11) { + return -1; + } + + if (v01 == v11) { + return 0; + } + } + + return 1; +} + + + +// construct sparse matrix from unsorted edge array, return number of nonzeros +int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr) { + if (!ne) { + return 0; + } + + // sort edges + mjQUICKSORT(edge, ne, 2*sizeof(int), edgecompare, NULL); + + // construct sparse + int nnz = 0; // number of nonzeros + int e = 0; // current edge + for (int r=0; r < nr; r++) { + // init row + rownnz[r] = 0; + rowadr[r] = nnz; + + // copy values while making unique and checking indices + while (e < ne && edge[2*e] == r) { + int v0 = edge[2*e]; + int v1 = edge[2*e + 1]; + + // skip if duplicate + if (rownnz[r] && v0 == edge[2*e - 2] && v1 == edge[2*e - 1]) { + e++; + continue; + } + + // check for invalid indices + if (v0 < 0 || v0 >= nr) mju_error("invalid row index %d in edge %d", v0, e); + if (v1 < 0 || v1 >= nr) mju_error("invalid column index %d in edge %d", v1, e); + + // copy column index, increment nnz, e, rownnz + colind[nnz++] = edge[2*(e++) + 1]; + rownnz[r]++; + } + } + + return nnz; +} + + +// find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix +// arguments: +// island (nr) - island index assigned to vertex, -1 if vertex has no edges +// nr - number of rows/columns of adjacency matrix +// rownnz (nr) - matrix row nonzeros +// rowadr (nr) - matrix row addresses +// colind (nnz) - matrix column indices +// stack (nnz) - stack space +// returns number of islands +int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, + int* stack) { + // initialize island count, set ids to -1 + int nisland = 0; + for (int i=0; i < nr; i++) island[i] = -1; + + // iterate over vertices, discover islands + for (int i=0; i < nr; i++) { + // vertex already in island or singleton with no edges: skip + if (island[i] != -1 || !rownnz[i]) { + continue; + } + + // push i onto stack + int nstack = 0; + stack[nstack++] = i; + + // DFS traversal of island + while (nstack) { + // pop v from stack + int v = stack[--nstack]; + + // if v is already assigned, continue + if (island[v] != -1) { + continue; + } + + // assign v to current island + island[v] = nisland; + + // push adjacent vertices onto stack + memcpy(stack + nstack, colind + rowadr[v], rownnz[v]*sizeof(int)); + nstack += rownnz[v]; + } + + // island is filled: increment nisland + nisland++; + } + + return nisland; +} diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index e9ca636e..0d357bbc 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -114,6 +114,13 @@ MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d); // optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian); + +// construct sparse matrix from unsorted edge array, return number of nonzeros +MJAPI int mj_edge2Sparse(int* rownnz, int* rowadr, int* colind, int* edge, int ne, int nr); + +MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, + int* scratch); + #ifdef __cplusplus } #endif diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 3da92b8c..9bc43c19 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -35,8 +35,8 @@ mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2, const int nnz2, const int* ind2); // convert matrix from dense to sparse -void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, - int* rownnz, int* rowadr, int* colind); +MJAPI void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, + int* rownnz, int* rowadr, int* colind); // convert matrix from sparse to dense void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index cd64ae37..f28d6dff 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -24,12 +24,15 @@ #include #include "src/engine/engine_core_constraint.h" #include "src/engine/engine_support.h" +#include "src/engine/engine_util_sparse.h" #include "test/fixture.h" namespace mujoco { namespace { using ::testing::DoubleNear; +using ::testing::ElementsAre; +using ::testing::ElementsAreArray; using ::testing::Pointwise; using CoreConstraintTest = MujocoTest; @@ -83,7 +86,7 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) { ASSERT_THAT(model, testing::NotNull()); ASSERT_EQ(model->nq, 7); ASSERT_EQ(model->nv, 6); - static const int nv = 6; // for increased readabilty + static const int nv = 6; // for increased readability mjData* data = mj_makeData(model); // arbitrary initial values for the ball and hinge joints @@ -269,5 +272,233 @@ TEST_F(CoreConstraintTest, CombineSparseCount) { } } +TEST_F(CoreConstraintTest, EdgeToSparse4) { + // unsorted edges, with duplication + constexpr int ne = 6; + constexpr int nr = 5; + int edge[2*ne] = { + 1, 1, + 0, 0, + 0, 1, + 3, 2, + 1, 1, + 0, 0 + }; + + int rownnz[nr]; + int rowadr[nr]; + int colind[ne]; + + int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr); + + constexpr int expected_nnz = 4; + + EXPECT_EQ(nnz, expected_nnz); + + EXPECT_THAT(rownnz, ElementsAre(2, 1, 0, 1, 0)); + EXPECT_THAT(rowadr, ElementsAre(0, 2, 3, 3, 4)); + + int expected_colind[expected_nnz] = {0, 1, 1, 2}; + EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz)); +} + +TEST_F(CoreConstraintTest, EdgeToSparse2) { + // unsorted edges, with duplication + constexpr int ne = 4; + constexpr int nr = 5; + int edge[2*ne] = { + 3, 4, + 1, 1, + 3, 4, + 1, 1 + }; + + int rownnz[nr]; + int rowadr[nr]; + int colind[ne]; + + int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr); + + constexpr int expected_nnz = 2; + + EXPECT_EQ(nnz, expected_nnz); + + EXPECT_THAT(rownnz, ElementsAre(0, 1, 0, 1, 0)); + EXPECT_THAT(rowadr, ElementsAre(0, 0, 1, 1, 2)); + + int expected_colind[expected_nnz] = {1, 4}; + EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz)); +} + +TEST_F(CoreConstraintTest, EdgeToSparse3) { + // unsorted edges, with duplication + constexpr int ne = 3; + constexpr int nr = 1; + int edge[2*ne] = { + 0, 0, + 0, 0, + 0, 0 + }; + + int rownnz[nr]; + int rowadr[nr]; + int colind[ne]; + + int nnz = mj_edge2Sparse(rownnz, rowadr, colind, edge, ne, nr); + + constexpr int expected_nnz = 1; + + EXPECT_EQ(nnz, expected_nnz); + + EXPECT_THAT(rownnz, ElementsAre(1)); + EXPECT_THAT(rowadr, ElementsAre(0)); + + int expected_colind[expected_nnz] = {0}; + EXPECT_THAT(expected_colind, ElementsAreArray(colind, expected_nnz)); +} + + +TEST_F(CoreConstraintTest, FloodFillSingleton) { + // adjacency matrix for the graph 0 1 2 + // U U + // (3 singletons, 0 and 2 have self-edges) + mjtNum mat[9] = { + 1, 0, 0, + 0, 0, 0, + 0, 0, 1 + }; + constexpr int nr = 3; + constexpr int nnz = 2; + int rownnz[nr]; + int rowadr[nr]; + int colind[nnz]; + mjtNum res[nnz]; // unused + mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind); + + // outputs / scratch + int island[nr]; + int scratch[2*nr]; + + // flood fill + int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, scratch); + + EXPECT_EQ(nisland, 2); + EXPECT_THAT(island, ElementsAre(0, -1, 1)); +} + +TEST_F(CoreConstraintTest, FloodFill1) { + // adjacency matrix for the graph 0 - 1 - 2 + mjtNum mat[9] = { + 0, 1, 0, + 1, 0, 1, + 0, 1, 0 + }; + constexpr int nr = 3; + constexpr int nnz = 4; + int rownnz[nr]; + int rowadr[nr]; + int colind[nnz]; + mjtNum res[nnz]; // unused + mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind); + + // outputs / stack + int island[nr]; + int stack[nnz]; + + int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack); + + EXPECT_EQ(nisland, 1); + EXPECT_THAT(island, ElementsAre(0, 0, 0)); +} + +TEST_F(CoreConstraintTest, FloodFill2) { + // adjacency matrix for the graph 6 – 1 – 4 0 – 3 – 5 – 2 + mjtNum mat[49] = { + 0, 0, 0, 1, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 1, + 0, 0, 0, 0, 0, 1, 0, + 1, 0, 0, 0, 0, 1, 0, + 0, 1, 0, 0, 0, 0, 0, + 0, 0, 1, 1, 0, 0, 0, + 0, 1, 0, 0, 0, 0, 0, + }; + constexpr int nr = 7; + constexpr int nnz = 10; + int rownnz[nr]; + int rowadr[nr]; + int colind[nnz]; + mjtNum res[nnz]; // unused + mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind); + + // outputs / stack + int island[nr]; + int stack[nnz]; + + int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack); + + EXPECT_EQ(nisland, 2); + EXPECT_THAT(island, ElementsAre(0, 1, 0, 0, 1, 0, 1)); +} + +TEST_F(CoreConstraintTest, FloodFill3a) { + // adjacency matrix for the graph 0 2 1 – 3 + // U + mjtNum mat[16] = { + 0, 0, 0, 0, + 0, 0, 0, 1, + 0, 0, 1, 0, + 0, 1, 0, 0, + }; + constexpr int nr = 4; + constexpr int nnz = 3; + int rownnz[nr]; + int rowadr[nr]; + int colind[nnz]; + mjtNum res[nnz]; // unused + mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind); + + // outputs / stack + int island[nr]; + int stack[nnz]; + + int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack); + + EXPECT_EQ(nisland, 2); + EXPECT_THAT(island, ElementsAre(-1, 0, 1, 0)); +} + +TEST_F(CoreConstraintTest, FloodFill3b) { + /* + adjacency matrix for the graph 1 – 2 3 4 – 5 + U | \ | + 0 – 6 + */ + mjtNum mat[49] = { + 0, 0, 0, 0, 1, 0, 1, + 0, 1, 1, 0, 0, 0, 0, + 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, + 1, 0, 0, 0, 0, 1, 1, + 0, 0, 0, 0, 1, 0, 1, + 1, 0, 0, 0, 1, 1, 0, + }; + constexpr int nr = 7; + constexpr int nnz = 13; + int rownnz[nr]; + int rowadr[nr]; + int colind[nnz]; + mjtNum res[nnz]; // unused + mju_dense2sparse(res, mat, nr, nr, rownnz, rowadr, colind); + + // outputs / stack + int island[nr]; + int stack[nnz]; + + int nisland = mj_floodFill(island, nr, rownnz, rowadr, colind, stack); + + EXPECT_EQ(nisland, 2); + EXPECT_THAT(island, ElementsAre(0, 1, 1, -1, 0, 0, 0)); +} + } // namespace } // namespace mujoco