Add private function mj_solveM_island for mass matrix solve with sub indices corresponding to one island.

PiperOrigin-RevId: 561622962
Change-Id: Ib43f03b3dd0c7456c57d4ba45a68e0708bf1c91c

src/engine/engine_core_smooth.c

@@ -1086,7 +1086,7 @@ void mj_factorM(const mjModel* m, mjData* d) {
 
 
 
-// sparse backsubstitution:  x = inv(L'*D*L)*y
+// in-place sparse backsubstitution:  x = inv(L'*D*L)*x
 //  L is in lower triangle of qLD; D is on diagonal of qLD
 //  handle n vectors at once
 void mj_solveLD(const mjModel* m, mjtNum* restrict x, int n,
@@ -1209,6 +1209,66 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
 }
 
 
+// in-place sparse backsubstitution for one island:  x = inv(L'*D*L)*x
+//  L is in lower triangle of qLD; D is on diagonal of qLD
+void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
+  // local constants: general
+  const int* Madr = m->dof_Madr;
+  const int* parentid = m->dof_parentid;
+  const mjtNum* qLD = d->qLD;
+  const mjtNum* qLDiagInv = d->qLDiagInv;
+  const int* simplenum = m->dof_simplenum;
+
+  // local constants: island specific
+  int ndof = d->island_dofnum[island];
+  const int* dofind = d->island_dofind + d->island_dofadr[island];
+  const int* islandind = d->dof_islandind;
+
+  // x <- inv(L') * x; skip simple, exploit sparsity of input vector
+  for (int k=ndof-1; k >= 0; k--) {
+    int i = dofind[k];
+    if (!simplenum[i] && x[k]) {
+      // init
+      int Madr_ij = Madr[i]+1;
+      int j = parentid[i];
+
+      // traverse ancestors backwards
+      // read directly from x[l] since j cannot be a parent of itself
+      while (j >= 0) {
+        x[islandind[j]] -= qLD[Madr_ij++]*x[k];         // x(j) -= L(i,j) * x(i)
+
+        // advance to parent
+        j = parentid[j];
+      }
+    }
+  }
+
+  // x <- inv(D) * x
+  for (int k=ndof-1; k >= 0; k--) {
+    x[k] *= qLDiagInv[dofind[k]];  // x(i) /= L(i,i)
+  }
+
+  // x <- inv(L) * x; skip simple
+  for (int k=0; k < ndof; k++) {
+    int i = dofind[k];
+    if (!simplenum[i]) {
+      // init
+      int Madr_ij = Madr[i]+1;
+      int j = parentid[i];
+
+      // traverse ancestors backwards
+      // write directly in x[i] since i cannot be a parent of itself
+      while (j >= 0) {
+        x[k] -= qLD[Madr_ij++]*x[islandind[j]];             // x(i) -= L(i,j) * x(j)
+
+        // advance to parent
+        j = parentid[j];
+      }
+    }
+  }
+}
+
+
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {

src/engine/engine_core_smooth.h

@@ -59,6 +59,9 @@ MJAPI void mj_solveLD(const mjModel* m, mjtNum* x, int n,
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 
+// sparse backsubstitution for one island:  x = inv(L'*D*L)*x, use factorization in d
+MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
+
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 

test/engine/engine_core_smooth_test.cc

@@ -14,8 +14,11 @@
 
 // Tests for engine/engine_core_smooth.c.
 
+#include "src/engine/engine_core_smooth.h"
+
 #include <string>
 #include <string_view>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -291,5 +294,61 @@ TEST_F(CoreSmoothTest, RefsiteBringsToPose) {
   mj_deleteModel(model);
 }
 
+static const char* const kIlslandEfcPath =
+    "engine/testdata/island/island_efc.xml";
+
+TEST_F(CoreSmoothTest, SolveMIsland) {
+  const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+  int nv = model->nv;
+
+  // allocate vec, fill with arbitrary values, copy to sol
+  mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
+  mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
+  for (int i=0; i < nv; i++) {
+    vec[i] = 0.2 + 0.3*i;
+  }
+  mju_copy(res, vec, nv);
+
+  // simulate for 0.3 seconds
+  mj_resetData(model, data);
+  while (data->time < 0.3) {
+    mj_step(model, data);
+  }
+  mj_forward(model, data);
+
+  // divide by mass matrix: sol = M^-1 * vec
+  mj_solveM(model, data, res, res, 1);
+
+  // iterate over islands
+  for (int i=0; i < data->nisland; i++) {
+    // allocate dof vectors for island
+    int dofnum = data->island_dofnum[i];
+    mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
+
+    // copy values into sol_i
+    int* dofind = data->island_dofind + data->island_dofadr[i];
+    for (int j=0; j < dofnum; j++) {
+      res_i[j] = vec[dofind[j]];
+    }
+
+    // divide by mass matrix, for this island
+    mj_solveM_island(model, data, res_i, i);
+
+    // expect corresponding values to match
+    for (int j=0; j < dofnum; j++) {
+      EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12));
+    }
+
+    mju_free(res_i);
+  }
+
+  mju_free(res);
+  mju_free(vec);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco