diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 9b4affe7..b8a667d6 100644 --- a/src/engine/engine_core_smooth.c +++ b/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) { diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index edaa4133..48776db5 100644 --- a/src/engine/engine_core_smooth.h +++ b/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); diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 7f672263..6d123d1a 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/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 #include +#include #include #include @@ -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