Add private function mj_mulJacVec_island for Jacobian multiplication with sub indices corresponding to one island.

PiperOrigin-RevId: 561648540
Change-Id: I8080a620807e824b40b5d6ba484beb06d081d444

src/engine/engine_core_constraint.c

@@ -395,6 +395,45 @@ void mj_mulJacVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec) {
 
 
 
+// multiply Jacobian by vector, for one island
+void mj_mulJacVec_island(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, int island) {
+  // no island, call regular function
+  if (island < 0) {
+    mj_mulJacVec(m, d, res, vec);
+    return;
+  }
+
+  // sizes
+  int vecnnz = d->island_dofnum[island];
+  int resnnz = d->island_efcnum[island];
+
+  // indices
+  int* vecind = d->island_dofind + d->island_dofadr[island];
+  int* resind = d->island_efcind + d->island_efcadr[island];
+
+  // sparse Jacobian
+  if (mj_isSparse(m)) {
+    for (int i=0; i < resnnz; i++) {
+      int row = resind[i];
+      int Jnnz = d->efc_J_rownnz[row];
+      int Jrowadr = d->efc_J_rowadr[row];
+      int* Jind = d->efc_J_colind + Jrowadr;
+      mjtNum* J = d->efc_J + Jrowadr;
+      res[i] = mju_dotSparse2(vec, J, vecnnz, vecind, Jnnz, Jind);
+    }
+  }
+
+  // dense Jacobian
+  else {
+    int nv = m->nv;
+    for (int i=0; i < resnnz; i++) {
+      res[i] = mju_dotSparse(vec, d->efc_J + nv*resind[i], vecnnz, vecind);
+    }
+  }
+}
+
+
+
 // multiply JacobianT by vector
 void mj_mulJacTVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec) {
   // exit if no constraints

src/engine/engine_core_constraint.h

@@ -38,6 +38,10 @@ MJAPI int mj_isDual(const mjModel* m);
 // multiply Jacobian by vector
 MJAPI void mj_mulJacVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec);
 
+// multiply Jacobian by vector, for one island
+MJAPI void mj_mulJacVec_island(const mjModel* m, mjData* d,
+                               mjtNum* res, const mjtNum* vec, int island);
+
 // multiply JacobianT by vector
 MJAPI void mj_mulJacTVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec);
 

test/engine/engine_core_constraint_test.cc

@@ -14,7 +14,6 @@
 
 // Tests for engine/engine_core_constraint.c.
 
-#include <array>
 #include <cstddef>
 #include <string>
 #include <vector>
@@ -198,5 +197,69 @@ TEST_F(CoreConstraintTest, JacobianPreAllocate) {
   }
 }
 
+static const char* const kIlslandEfcPath =
+    "engine/testdata/island/island_efc.xml";
+
+TEST_F(CoreConstraintTest, MulJacVecIsland) {
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+
+  // allocate vec_nv, fill with arbitrary values
+  mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv);
+  for (int i=0; i < model->nv; i++) {
+    vec_nv[i] = 0.2 + 0.3*i;
+  }
+
+  // iterate through dense and sparse
+  for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) {
+    model->opt.jacobian = sparsity;
+
+    // simulate for 0.3 seconds
+    mj_resetData(model, data);
+    while (data->time < 0.3) {
+      mj_step(model, data);
+    }
+    mj_forward(model, data);
+
+    // multiply by Jacobian: vec_nefc = J * vec_nv
+    mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc);
+    mj_mulJacVec(model, data, vec_nefc, vec_nv);
+
+    // iterate over islands
+    for (int i=0; i < data->nisland; i++) {
+      // allocate dof and efc vectors for island
+      int dofnum = data->island_dofnum[i];
+      mjtNum* vec_nvi = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
+      int efcnum = data->island_efcnum[i];
+      mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum);
+
+      // copy values into vec_nvi
+      int* dofind = data->island_dofind + data->island_dofadr[i];
+      for (int j=0; j < dofnum; j++) {
+        vec_nvi[j] = vec_nv[dofind[j]];
+      }
+
+      // multiply by Jacobian, for this island
+      mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, i);
+
+      // expect corresponding values to match
+      int* efcind = data->island_efcind + data->island_efcadr[i];
+      for (int j=0; j < efcnum; j++) {
+        EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12));
+      }
+
+      mju_free(vec_nvi);
+      mju_free(vec_nefci);
+    }
+
+    mju_free(vec_nefc);
+  }
+
+  mju_free(vec_nv);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco