diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 34b2892e..d0d14caa 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -455,6 +455,45 @@ void mj_mulJacTVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec) +// multiply Jacobian transpose by vector, for one island +void mj_mulJacTVec_island(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, int island) { + // no island, call regular function + if (island < 0) { + mj_mulJacTVec(m, d, res, vec); + return; + } + + // sizes + int vecnnz = d->island_efcnum[island]; + int resnnz = d->island_dofnum[island]; + + // indices + int* vecind = d->island_efcind + d->island_efcadr[island]; + int* resind = d->island_dofind + d->island_dofadr[island]; + + // sparse Jacobian + if (mj_isSparse(m)) { + for (int i=0; i < resnnz; i++) { + int row = resind[i]; + int JTnnz = d->efc_JT_rownnz[row]; + int JTrowadr = d->efc_JT_rowadr[row]; + int* JTind = d->efc_JT_colind + JTrowadr; + mjtNum* JT = d->efc_JT + JTrowadr; + res[i] = mju_dotSparse2(vec, JT, vecnnz, vecind, JTnnz, JTind); + } + } + + // dense Jacobian + else { + int nefc = d->nefc; + for (int i=0; i < resnnz; i++) { + res[i] = mju_dotSparse(vec, d->efc_JT + nefc*resind[i], vecnnz, vecind); + } + } +} + + + //--------------------- instantiate constraints by type -------------------------------------------- // equality constraints @@ -1772,6 +1811,10 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { // supernodes of JT mju_superSparse(m->nv, d->efc_JT_rowsuper, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind); + } else { + if (mjENABLED(mjENBL_ISLAND)) { + mju_transpose(d->efc_JT, d->efc_J, d->nefc, m->nv); + } } // compute diagApprox diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index afa3447c..3104a339 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -42,6 +42,10 @@ MJAPI void mj_mulJacVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* MJAPI void mj_mulJacVec_island(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, int island); +// multiply Jacobian transposed by vector, for one island +MJAPI void mj_mulJacTVec_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); diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index af2df04c..a0c2a477 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -261,5 +261,67 @@ TEST_F(CoreConstraintTest, MulJacVecIsland) { mj_deleteModel(model); } +TEST_F(CoreConstraintTest, MulJacTVecIsland) { + 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 + mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); + + // 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); + + // allocate vec_nefc, fill with arbitrary values + mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); + for (int i=0; i < data->nefc; i++) { + vec_nefc[i] = 0.2 + 0.3*i; + } + + // multiply by Jacobian: vec_nv = J^T * vec_nefc + mj_mulJacTVec(model, data, vec_nv, vec_nefc); + + // 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_nefci + int* efcind = data->island_efcind + data->island_efcadr[i]; + for (int j=0; j < efcnum; j++) { + vec_nefci[j] = vec_nefc[efcind[j]]; + } + + // multiply by Jacobian, for this island + mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci, i); + + // expect corresponding values to match + int* dofind = data->island_dofind + data->island_dofadr[i]; + for (int j=0; j < dofnum; j++) { + EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[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