diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 3c90f032..65a704ac 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -160,13 +160,12 @@ static int countMaxEdge(const mjModel* m, const mjData* d) { // return id of next tree in Jacobian row i that is different from tree, -1 if not found -// write the index of the found tree to *index if given -// start search from *index if given, otherwise 0 -// if J is (dense/sparse) *index is the (column/nonzro) index, respectively +// start search from *index +// write the index of the found tree to *index +// if J is (dense/sparse) *index is the (column/nonzero) index, respectively static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *index) { int tree_next = -1; - int j0 = index ? *index : 0; // start searching at *index if given, otherwise 0 - int j; // loop variable, saved to *index + int j; // local loop variable, saved to *index // sparse if (mj_isSparse(m)) { @@ -174,7 +173,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind int* colind = d->efc_J_colind + d->efc_J_rowadr[i]; // loop over remaining nonzeros, look for different tree - for (j=j0; j < rownnz; j++) { + for (j=(*index); j < rownnz; j++) { int tree_j = m->dof_treeid[colind[j]]; if (tree_j != tree) { // found different tree @@ -189,7 +188,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind int nv = m->nv; // scan row, look for different tree - for (j=j0; j < nv; j++) { + for (j=(*index); j < nv; j++) { if (d->efc_J[nv*i + j]) { int tree_j = m->dof_treeid[j]; if (tree_j != tree) { @@ -202,13 +201,94 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind } // save last index - if (index) *index = j; + *index = j; return tree_next; } +// find first and possibly second nonegative tree ids in Jacobian row i +// if row i is special-cased (no more trees), return -1 +// otherwise call treeNext, starting scan at index 0, return index +static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) { + int efc_type = d->efc_type[i]; + int efc_id = d->efc_id[i]; + + // clear outputs + tree[0] = -1; + tree[1] = -1; + + // ==== fast handling of special cases + + // joint friction + if (efc_type == mjCNSTR_FRICTION_DOF) { + tree[0] = m->dof_treeid[efc_id]; + return -1; + } + + // joint limit + if (efc_type == mjCNSTR_LIMIT_JOINT) { + tree[0] = m->dof_treeid[m->jnt_dofadr[efc_id]]; + return -1; + } + + // contact + if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS || + efc_type == mjCNSTR_CONTACT_PYRAMIDAL || + efc_type == mjCNSTR_CONTACT_ELLIPTIC) { + tree[0] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]]; + tree[1] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]]; + + // handle static bodies + if (tree[0] < 0) { + if (tree[1] < 0) { + mjERROR("contact %d is between two static bodies", efc_id); // SHOULD NOT OCCUR + } else { + int tmp = tree[0]; + tree[0] = tree[1]; + tree[1] = tmp; + } + } + + return -1; + } + + // connect or weld constraints + if (efc_type == mjCNSTR_EQUALITY) { + mjtEq eq_type = m->eq_type[efc_id]; + if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) { + tree[0] = m->body_treeid[m->eq_obj1id[efc_id]]; + tree[1] = m->body_treeid[m->eq_obj2id[efc_id]]; + + // handle static bodies + if (tree[0] < 0) { + if (tree[1] < 0) { + mjERROR("equality %d is between two static bodies", efc_id); // SHOULD NOT OCCUR + } else { + int tmp = tree[0]; + tree[0] = tree[1]; + tree[1] = tmp; + } + } + + return -1; + } + } + + // ==== generic case: scan Jacobian + int index = 0; + tree[0] = treeNext(m, d, -1, i, &index); + + if (tree[0] < 0) { + mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR + } + + return index; +} + + + // add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge // return current number of edges static int addEdge(int* treenedge, int* edge, int nedge, int tree1, int tree2, int nedge_max) { @@ -273,7 +353,6 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg int nefc = d->nefc; int efc_type = -1; int efc_id = -1; - int tree1, tree2; // clear treenedge memset(treenedge, 0, m->ntree*sizeof(int)); @@ -287,60 +366,34 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg efc_type = d->efc_type[i]; efc_id = d->efc_id[i]; - // ==== fast handling of special cases + int tree[2]; + int index = treeFirst(m, d, tree, i); + int tree1 = tree[0]; + int tree2 = tree[1]; - // joint friction - if (efc_type == mjCNSTR_FRICTION_DOF) { - tree1 = m->dof_treeid[efc_id]; - nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); + // no more edges to find, add and continue + if (index == -1) { + nedge = addEdge(treenedge, edge, nedge, tree1, tree2 == -1 ? tree1 : tree2, nedge_max); continue; } - // joint limit - if (efc_type == mjCNSTR_LIMIT_JOINT) { - tree1 = m->dof_treeid[m->jnt_dofadr[efc_id]]; - nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); - continue; - } + // possibly more edges, scan Jacobian row + else { + tree2 = treeNext(m, d, tree1, i, &index); - // contact - if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS || - efc_type == mjCNSTR_CONTACT_PYRAMIDAL || - efc_type == mjCNSTR_CONTACT_ELLIPTIC) { - tree1 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]]; - tree2 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]]; - nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); - continue; - } - - // connect or weld constraints - if (efc_type == mjCNSTR_EQUALITY) { - mjtEq eq_type = m->eq_type[efc_id]; - if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) { - tree1 = m->body_treeid[m->eq_obj1id[efc_id]]; - tree2 = m->body_treeid[m->eq_obj2id[efc_id]]; + if (tree2 == -1) { + // 1 tree found: add self-edge + nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); + } else { + // 2 trees found: add edge, keep scanning and adding until no more trees nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); - continue; - } - } - - // ==== generic case: scan Jacobian - int index = 0; - tree1 = treeNext(m, d, -1, i, &index); - tree2 = treeNext(m, d, tree1, i, &index); - - if (tree2 == -1) { - // 1 tree found: add self-edge - nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max); - } else { - // 2 trees found: add edge, keep scanning and adding until no more trees - nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); - int tree3 = treeNext(m, d, tree2, i, &index); - while (tree3 > -1 && tree3 != tree2) { - tree1 = tree2; - tree2 = tree3; - nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); - tree3 = treeNext(m, d, tree2, i, &index); + int tree3 = treeNext(m, d, tree2, i, &index); + while (tree3 > -1 && tree3 != tree2) { + tree1 = tree2; + tree2 = tree3; + nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max); + tree3 = treeNext(m, d, tree2, i, &index); + } } } } @@ -452,7 +505,9 @@ void mj_island(const mjModel* m, mjData* d) { // compute efc_island, island_efcnum memset(d->island_efcnum, 0, nisland*sizeof(int)); for (int i=0; i < nefc; i++) { - int island = tree_island[treeNext(m, d, -1, i, NULL)]; + int tree[2]; + treeFirst(m, d, tree, i); + int island = tree_island[tree[0]]; d->efc_island[i] = island; d->island_efcnum[island]++; } diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 2e74bfe3..58027551 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -350,33 +350,33 @@ TEST_F(IslandTest, IslandEfc) { mj_step(model, data); } - // sizes - int nv = model->nv; - int nefc = data->nefc; - int nisland = data->nisland; - // expect island structure to correspond to comment at top of xml - EXPECT_EQ(nisland, 3); - EXPECT_EQ(data->ne, 1); + EXPECT_EQ(data->nisland, 4); + EXPECT_EQ(data->ne, 4); EXPECT_EQ(data->nf, 2); EXPECT_EQ(data->nl, 1); - EXPECT_EQ(nefc, 24); - EXPECT_THAT(AsVector(data->dof_island, nv), - ElementsAre(0, -1, 0, 0, 0, 1, 2, 0, 1, 1, 1, 1, 1, 1)); - EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(5, 7, 1)); - EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 5, 12)); - EXPECT_THAT(AsVector(data->island_dofind, nv), - ElementsAre(0, 2, 3, 4, 7, 5, 8, 9, 10, 11, 12, 13, 6, -1)); - EXPECT_THAT(AsVector(data->dof_islandind, nv), - ElementsAre(0, -1, 1, 2, 3, 0, 0, 4, 1, 2, 3, 4, 5, 6)); - EXPECT_THAT(AsVector(data->efc_island, nefc), - ElementsAre(0, 1, 2, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1)); - EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(6, 17, 1)); - EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 6, 23)); - EXPECT_THAT(AsVector(data->island_efcind, nefc), - ElementsAre(0, 3, 4, 5, 6, 7, 1, 8, 9, 10, 11, 12, 13, 14, 15, 16, - 17, 18, 19, 20, 21, 22, 23, 2)); + EXPECT_EQ(data->nefc, 27); + + mj_deleteData(data); + mj_deleteModel(model); +} + +TEST_F(IslandTest, IslandEfcElliptic) { + const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + + model->opt.cone = mjCONE_ELLIPTIC; + while (data->time < 0.2) { + mj_step(model, data); + } + mj_forward(model, data); + + EXPECT_EQ(data->nisland, 4); + EXPECT_EQ(data->ne, 4); + EXPECT_EQ(data->nf, 2); + EXPECT_EQ(data->nl, 1); + EXPECT_EQ(data->nefc, 22); mj_deleteData(data); mj_deleteModel(model); diff --git a/test/engine/testdata/island/island_efc.xml b/test/engine/testdata/island/island_efc.xml index 3f30d9ec..b500f57c 100644 --- a/test/engine/testdata/island/island_efc.xml +++ b/test/engine/testdata/island/island_efc.xml @@ -1,10 +1,10 @@ @@ -64,9 +64,15 @@ + + + + + +