Add mjData.dof_islandind, indexing from dofs into islands.

Also fix bug in allocation for constraint islands, added allocation for limit constraints.

PiperOrigin-RevId: 561083128
Change-Id: I9e01d5a5b7454ec3cf60bccadaa0f96f2c0d87d1

doc/includes/references.h

@@ -350,6 +350,7 @@ struct mjData_ {
   int*    island_dofnum;     // number of dofs in island                         (nisland x 1)
   int*    island_dofadr;     // start address in island_dofind                   (nisland x 1)
   int*    island_dofind;     // island dof indices; -1: none                     (nv x 1)
+  int*    dof_islandind;     // dof island indices; -1: none                     (nv x 1)
   int*    efc_island;        // island id of this constraint                     (nefc x 1)
   int*    island_efcnum;     // number of constraints in island                  (nisland x 1)
   int*    island_efcadr;     // start address in island_efcind                   (nisland x 1)

include/mujoco/mjdata.h

@@ -377,6 +377,7 @@ struct mjData_ {
   int*    island_dofnum;     // number of dofs in island                         (nisland x 1)
   int*    island_dofadr;     // start address in island_dofind                   (nisland x 1)
   int*    island_dofind;     // island dof indices; -1: none                     (nv x 1)
+  int*    dof_islandind;     // dof island indices; -1: none                     (nv x 1)
   int*    efc_island;        // island id of this constraint                     (nefc x 1)
   int*    island_efcnum;     // number of constraints in island                  (nisland x 1)
   int*    island_efcadr;     // start address in island_efcind                   (nisland x 1)

include/mujoco/mjxmacro.h

@@ -605,6 +605,7 @@
     X( int,   island_dofnum,      MJ_D(nisland), 1 ) \
     X( int,   island_dofadr,      MJ_D(nisland), 1 ) \
     X( int,   island_dofind,      MJ_M(nv),      1 ) \
+    X( int,   dof_islandind,      MJ_M(nv),      1 ) \
     X( int,   efc_island,         MJ_D(nefc),    1 ) \
     X( int,   island_efcnum,      MJ_D(nisland), 1 ) \
     X( int,   island_efcadr,      MJ_D(nisland), 1 ) \

introspect/structs.py

@@ -4372,6 +4372,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='island dof indices; -1: none                     (nv x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='dof_islandind',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='dof island indices; -1: none                     (nv x 1)',  # pylint: disable=line-too-long
+             ),
              StructFieldDecl(
                  name='efc_island',
                  type=PointerType(

src/engine/engine_island.c

@@ -141,6 +141,7 @@ static int countMaxEdge(const mjModel* m, const mjData* d) {
   int nedge_max = 0;
   nedge_max += 2*d->ncon;  // contact: 2 edges
   nedge_max += 2*d->ne;    // equality: 2 edges
+  nedge_max += d->nl;      // limit: 1 edges (always within same tree)
   nedge_max += d->nf;      // joint friction: 1 edge (always within same tree)
 
   // tendon limits and friction add up to tendon_num edges
@@ -425,13 +426,16 @@ void mj_island(const mjModel* m, mjData* d) {
   // reset island_dofnum
   memset(d->island_dofnum, 0, nisland*sizeof(int));
 
-  // compute dof_islandind
+  // compute dof_islandind, island_dofind
   int num_dof_island = 0;
   for (int i=0; i < nv; i++) {
     int island = d->dof_island[i];
     if (island >= 0) {
-      d->island_dofind[d->island_dofadr[island] + (d->island_dofnum[island]++)] = i;
+      d->island_dofind[d->island_dofadr[island] + d->island_dofnum[island]] = i;
+      d->dof_islandind[i] = d->island_dofnum[island]++;
       num_dof_island++;
+    } else {
+      d->dof_islandind[i] = -1;
     }
   }
 

src/engine/engine_print.c

@@ -1083,6 +1083,12 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
     }
     fprintf(fp, "\n\n");
 
+    fprintf(fp, NAME_FORMAT, "DOF_ISLANDIND");
+    for (int i = 0; i < m->nv; i++) {
+      fprintf(fp, " %d", d->dof_islandind[i]);
+    }
+    fprintf(fp, "\n\n");
+
     fprintf(fp, NAME_FORMAT, "EFC_ISLAND");
     for (int i = 0; i < d->nefc; i++) {
       fprintf(fp, " %d", d->efc_island[i]);

test/engine/engine_island_test.cc

@@ -230,6 +230,11 @@ TEST_F(IslandTest, Abacus) {
   // last index is unassigned since dof 1 is unconstrained
   EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 2, 3, -1));
 
+  // dof 0 constitutes first island
+  // dofs 1 is unassigned
+  // dofs 2, 3 are second island
+  EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, -1, 0, 1));
+
   // island 0 starts at constraint 0
   // island 1 starts at constraint 4
   EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4));
@@ -242,8 +247,7 @@ TEST_F(IslandTest, Abacus) {
   EXPECT_THAT(AsVector(data->efc_island, nefc),
               ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1));
 
-  // linked list for island 0
+  // index lists for islands 0 and 1
-  // linked list for island 1
   EXPECT_THAT(AsVector(data->island_efcind, nefc),
               ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11));
 
@@ -268,6 +272,7 @@ TEST_F(IslandTest, Abacus) {
   EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2, 1));
   EXPECT_THAT(AsVector(data->dof_island, nv), ElementsAre(0, 1, 1, 2));
   EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 1, 2, 3));
+  EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, 0, 1, 0));
   EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4, 8));
   EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 4, 4));
   EXPECT_THAT(AsVector(data->efc_island, nefc),
@@ -317,6 +322,8 @@ TEST_F(IslandTest, DenseSparse) {
             AsVector(data2->dof_island, nv));
   EXPECT_EQ(AsVector(data1->island_dofind, nv),
             AsVector(data2->island_dofind, nv));
+  EXPECT_EQ(AsVector(data1->dof_islandind, nv),
+            AsVector(data2->dof_islandind, nv));
   EXPECT_EQ(AsVector(data1->island_efcadr, nisland),
             AsVector(data2->island_efcadr, nisland));
   EXPECT_EQ(AsVector(data1->island_efcnum, nisland),
@@ -331,5 +338,47 @@ TEST_F(IslandTest, DenseSparse) {
   mj_deleteModel(model);
 }
 
+static const char* const kIlslandEfcPath =
+    "engine/testdata/island/island_efc.xml";
+
+TEST_F(IslandTest, IslandEfc) {
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+
+  while (data->time < 0.3) {
+    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, 2);
+  EXPECT_EQ(data->ne, 1);
+  EXPECT_EQ(data->nf, 1);
+  EXPECT_EQ(data->nl, 1);
+  EXPECT_EQ(nefc, 7);
+  EXPECT_THAT(AsVector(data->dof_island, nv),
+              ElementsAre(0, -1, 0, 0, 0, 1, 0));
+  EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(5, 1));
+  EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 5));
+  EXPECT_THAT(AsVector(data->island_dofind, nv),
+              ElementsAre(0, 2, 3, 4, 6, 5, -1));
+  EXPECT_THAT(AsVector(data->dof_islandind, nv),
+              ElementsAre(0, -1, 1, 2, 3, 0, 4));
+  EXPECT_THAT(AsVector(data->efc_island, nefc),
+              ElementsAre(0, 1, 0, 0, 0, 0, 0));
+  EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(6, 1));
+  EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 6));
+  EXPECT_THAT(AsVector(data->island_efcind, nefc),
+              ElementsAre(0, 2, 3, 4, 5, 6, 1));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco

test/engine/testdata/island/island_efc.xml

@@ -0,0 +1,53 @@
+<!--
+  test model for constraint islanding:
+- one equality
+- one limit
+- one friction
+- one unconstrained dof
+- two islands, one with non-contiguous dofs
+-->
+
+<mujoco>
+  <option>
+    <flag island="enable"/>
+  </option>
+
+  <default>
+    <geom size=".1"/>
+  </default>
+
+  <worldbody>
+    <body>
+      <joint type="slide" axis="0 0 1" range="0 1" limited="true"/>
+      <geom/>
+    </body>
+    <body pos=".25 0 0">
+      <joint type="slide" axis="1 0 0"/>
+      <geom/>
+    </body>
+    <body pos="0 0 0.25">
+      <joint type="slide" axis="0 0 1"/>
+      <geom/>
+      <body pos="0 -.15 0">
+        <joint name="hinge1" axis="0 1 0"/>
+        <geom type="capsule" size="0.03" fromto="0 0 0 -.2 0 0"/>
+        <body pos="-.2 0 0">
+          <joint axis="0 1 0"/>
+          <geom type="capsule" size="0.03" fromto="0 0 0 -.2 0 0"/>
+        </body>
+      </body>
+    </body>
+    <body pos="-.5 0 0">
+      <joint axis="0 1 0" frictionloss=".01"/>
+      <geom type="capsule" size="0.03" fromto="0 0 0 -.2 0 0"/>
+    </body>
+    <body pos="0 0 .5">
+      <joint name="hinge2" axis="0 1 0"/>
+      <geom type="box" size=".08 .02 .08"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <joint joint1="hinge1" joint2="hinge2"/>
+  </equality>
+</mujoco>

unity/Runtime/Bindings/MjBindings.cs

@@ -1721,6 +1721,7 @@ public unsafe struct mjData_ {
   public int* island_dofnum;
   public int* island_dofadr;
   public int* island_dofind;
+  public int* dof_islandind;
   public int* efc_island;
   public int* island_efcnum;
   public int* island_efcadr;