diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index a2a43c3995..01911f0d57 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1701,7 +1701,7 @@ 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) { +void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) { // if no islands, call mj_solveLD const mjtNum* qLD = d->qLD; const mjtNum* qLDiagInv = d->qLDiagInv; @@ -1710,10 +1710,19 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int return; } - // local constants: general - const int* Madr = m->dof_Madr; - const int* parentid = m->dof_parentid; - const int* simplenum = m->dof_simplenum; + // local copies of key variables + const int* rownnz = d->C_rownnz; + const int* rowadr = d->C_rowadr; + const int* colind = d->C_colind; + const int* diagnum = m->dof_simplenum; + + // temporary: make local CSR version of qLD + int nC = m->nC; + mj_markStack(d); + mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum); + for (int i=0; i < nC; i++) { + qLDs[i] = d->qLD[d->mapM2C[i]]; + } // local constants: island specific int ndof = d->island_dofnum[island]; @@ -1723,18 +1732,12 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int // 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]; + mjtNum x_k; + if (!diagnum[i] && (x_k = x[k])) { + int start = rowadr[i]; + int end = start + rownnz[i] - 1; + for (int adr=end-1; adr >= start; adr--) { + x[islandind[colind[adr]]] -= qLDs[adr] * x_k; } } } @@ -1747,21 +1750,20 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int // 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]; - } + + // skip diagonal rows + if (diagnum[i]) { + continue; + } + + int start = rowadr[i]; + int end = start + rownnz[i] - 1; + for (int adr=end-1; adr >= start; adr--) { + x[k] -= x[islandind[colind[adr]]] * qLDs[adr]; } } + + mj_freeStack(d); } diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index a710416c46..ec8d168d09 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -71,8 +71,9 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i // 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); +// TODO(tassa): Restore mjData const-ness. // 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); +MJAPI void mj_solveM_island(const mjModel* m, 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, diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 684653748a..f23d8bf5ad 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -895,9 +895,9 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) { } - +// TODO(tassa): Restore mjData const-ness. // update grad, Mgrad -static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx) { +static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) { int nv = ctx->nv; const int* dofind = ctx->dofind; diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index b63e57a36d..916fb6256b 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -353,58 +353,63 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) { static const char* const kIlslandEfcPath = "engine/testdata/island/island_efc.xml"; +static const char* const kModelPath = + "testdata/model.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; + for (auto model_path : {kModelPath, kIlslandEfcPath}) { + const std::string xml_path = GetTestDataFilePath(model_path); + 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); - // 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); + if (model->nkey > 0) mj_resetDataKeyframe(model, data, 0); - // simulate for 0.2 seconds - mj_resetData(model, data); - while (data->time < 0.2) { - mj_step(model, data); - } - mj_forward(model, data); + for (int i=0; i < 6; i++) { + mj_step(model, data); + } - // divide by mass matrix: sol = M^-1 * vec - mj_solveM(model, data, res, res, 1); + mj_forward(model, data); - // 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); + // divide by mass matrix: sol = M^-1 * vec + mj_solveM(model, data, res, res, 1); - // 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]]; - } + // 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); + // 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)); + // expect corresponding values to match + for (int j=0; j < dofnum; j++) { + EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14)); + } + + mju_free(res_i); } - mju_free(res_i); + mju_free(res); + mju_free(vec); + mj_deleteData(data); + mj_deleteModel(model); } - - mju_free(res); - mju_free(vec); - mj_deleteData(data); - mj_deleteModel(model); } static const char* const kInertiaPath = "engine/testdata/inertia.xml";