From 9fd186ac2bce133eb042fb9a6c116d26a7ca8276 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 31 Aug 2023 06:41:55 -0700
Subject: [PATCH] Add private function `mj_mulM_island` for mass matrix
 multiplication with sub indices corresponding to one island.

PiperOrigin-RevId: 561639405
Change-Id: I3777cea51ef22f6847b3cf3e17bdac5cf621c019
---
 src/engine/engine_support.c        | 52 +++++++++++++++++++++++++++
 src/engine/engine_support.h        |  4 +++
 test/engine/engine_support_test.cc | 58 ++++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+)

diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index 7e57a21a..c58e5e06 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -884,6 +884,58 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
 
 
 
+// multiply vector by inertia matrix for one dof island
+void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, int island) {
+  // if no island, call regular function
+  if (island < 0) {
+    mj_mulM(m, d, res, vec);
+    return;
+  }
+
+  // local constants: general
+  const mjtNum* M = d->qM;
+  const int* Madr = m->dof_Madr;
+  const int* parentid = m->dof_parentid;
+  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;
+
+  mju_zero(res, ndof);
+
+  for (int k=0; k < ndof; k++) {
+    // address in full dof vector
+    int i = dofind[k];
+
+    // address in M
+    int adr = Madr[i];
+
+    // diagonal
+    res[k] = M[adr]*vec[k];
+
+    // simple dof: continue
+    if (simplenum[i]) {
+      continue;
+    }
+
+    // off-diagonal
+    int j = parentid[i];
+    while (j >= 0) {
+      adr++;
+      int l = islandind[j];
+      res[k] += M[adr]*vec[l];
+      res[l] += M[adr]*vec[k];
+
+      // advance to parent
+      j = parentid[j];
+    }
+  }
+}
+
+
+
 // multiply vector by M^(1/2)
 void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
   int adr, nv = m->nv;
diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h
index ce601a06..3b10fbe7 100644
--- a/src/engine/engine_support.h
+++ b/src/engine/engine_support.h
@@ -108,6 +108,10 @@ MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
 // multiply vector by inertia matrix
 MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
 
+// multiply vector by inertia matrix for one dof island
+MJAPI void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
+                          int island);
+
 // multiply vector by (inertia matrix)^(1/2)
 MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
 
diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc
index b2a1a59a..a33b3f28 100644
--- a/test/engine/engine_support_test.cc
+++ b/test/engine/engine_support_test.cc
@@ -14,6 +14,8 @@
 
 // Tests for engine/engine_support.c.
 
+#include "src/engine/engine_support.h"
+
 #include <random>
 #include <string>
 #include <vector>
@@ -453,5 +455,61 @@ TEST_F(AddMTest, DenseSameAsSparse) {
   mj_deleteModel(m);
 }
 
+static const char* const kIlslandEfcPath =
+    "engine/testdata/island/island_efc.xml";
+
+TEST_F(SupportTest, MulMIsland) {
+  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, fill with arbitrary values
+  mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv);
+  for (int i=0; i < model->nv; i++) {
+    vec[i] = 0.2 + 0.3*i;
+  }
+
+  // simulate for 0.3 seconds
+  mj_resetData(model, data);
+  while (data->time < 0.3) {
+    mj_step(model, data);
+  }
+  mj_forward(model, data);
+
+  // multiply by Mass matrix: Mvec = M * vec
+  mjtNum* Mvec = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc);
+  mj_mulM(model, data, Mvec, vec);
+
+  // iterate over islands
+  for (int i=0; i < data->nisland; i++) {
+    // allocate dof vectors for island
+    int dofnum = data->island_dofnum[i];
+    mjtNum* vec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
+    mjtNum* Mvec_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
+
+    // copy values into vec_i
+    int* dofind = data->island_dofind + data->island_dofadr[i];
+    for (int j=0; j < dofnum; j++) {
+      vec_i[j] = vec[dofind[j]];
+    }
+
+    // multiply by Jacobian, for this island
+    mj_mulM_island(model, data, Mvec_i, vec_i, i);
+
+    // expect corresponding values to match
+    for (int j=0; j < dofnum; j++) {
+      EXPECT_THAT(Mvec_i[j], DoubleNear(Mvec[dofind[j]], 1e-12));
+    }
+
+    mju_free(vec_i);
+    mju_free(Mvec_i);
+  }
+
+  mju_free(Mvec);
+  mju_free(vec);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco