diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 22309c25..c95a3d65 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -386,8 +386,9 @@ void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* // multiply Jacobian by vector, for one island -void mj_mulJacVec_island(const mjModel* m, const mjData* d, - mjtNum* res, const mjtNum* vec, int island) { +// flg_resunc and flg_vecunc denote whether res/vec are uncompressed +void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, + int island, int flg_resunc, int flg_vecunc) { // no island, call regular function if (island < 0) { mj_mulJacVec(m, d, res, vec); @@ -410,7 +411,8 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d, int Jrowadr = d->efc_J_rowadr[row]; int* Jind = d->efc_J_colind + Jrowadr; mjtNum* J = d->efc_J + Jrowadr; - res[i] = mju_dotSparse2(vec, J, vecnnz, vecind, Jnnz, Jind, /*flg_unc2=*/0); + int j = flg_resunc ? row : i; + res[j] = mju_dotSparse2(J, vec, Jnnz, Jind, vecnnz, vecind, flg_vecunc); } } @@ -418,7 +420,9 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d, else { int nv = m->nv; for (int i=0; i < resnnz; i++) { - res[i] = mju_dotSparse(vec, d->efc_J + nv*resind[i], vecnnz, vecind, /*flg_unc1=*/0); + int row = resind[i]; + int j = flg_resunc ? row : i; + res[j] = mju_dotSparse(vec, d->efc_J + nv*row, vecnnz, vecind, flg_vecunc); } } } @@ -447,8 +451,9 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* // multiply Jacobian transpose by vector, for one island -void mj_mulJacTVec_island(const mjModel* m, const mjData* d, - mjtNum* res, const mjtNum* vec, int island) { +// flg_resunc and flg_vecunc denote whether res/vec are uncompressed +void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, + int island, int flg_resunc, int flg_vecunc) { // no island, call regular function if (island < 0) { mj_mulJacTVec(m, d, res, vec); @@ -471,7 +476,8 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d, 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, /*flg_unc2=*/0); + int j = flg_resunc ? row : i; + res[j] = mju_dotSparse2(JT, vec, JTnnz, JTind, vecnnz, vecind, flg_vecunc); } } @@ -479,7 +485,9 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d, 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, /*flg_unc1=*/0); + int row = resind[i]; + int j = flg_resunc ? row : i; + res[j] = mju_dotSparse(vec, d->efc_JT + nefc*row, vecnnz, vecind, flg_vecunc); } } } diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index 2dcc42fb..5c6b353a 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -39,15 +39,15 @@ MJAPI int mj_isDual(const mjModel* m); MJAPI void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); // multiply Jacobian by vector, for one island -MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d, - mjtNum* res, const mjtNum* vec, int island); +MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, + int island, int flg_resunc, int flg_vecunc); // multiply JacobianT by vector MJAPI void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec); // multiply JacobianT by vector, for one island -MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d, - mjtNum* res, const mjtNum* vec, int island); +MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec, + int island, int flg_resunc, int flg_vecunc); //-------------------------- utility functions ----------------------------------------------------- diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index fee8ba7d..32c5fb0b 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -15,6 +15,7 @@ // Tests for engine/engine_core_constraint.c. #include +#include #include #include @@ -225,6 +226,7 @@ TEST_F(CoreConstraintTest, MulJacVecIsland) { // multiply by Jacobian: vec_nefc = J * vec_nv mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); mj_mulJacVec(model, data, vec_nefc, vec_nv); + mjtNum* vec_nefc_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc); // iterate over islands for (int i=0; i < data->nisland; i++) { @@ -234,25 +236,57 @@ TEST_F(CoreConstraintTest, MulJacVecIsland) { int efcnum = data->island_efcnum[i]; mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - // copy values into vec_nvi + // get indices int* dofind = data->island_dofind + data->island_dofadr[i]; + int* efcind = data->island_efcind + data->island_efcadr[i]; + + // copy values into vec_nvi for (int j=0; j < dofnum; j++) { vec_nvi[j] = vec_nv[dofind[j]]; } - // multiply by Jacobian, for this island - mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, i); + // ===== both compressed + int flg_resunc = 0; + int flg_vecunc = 0; + mju_zero(vec_nefci, efcnum); // clear output + mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, + i, flg_resunc, flg_vecunc); // expect corresponding values to match - int* efcind = data->island_efcind + data->island_efcadr[i]; for (int j=0; j < efcnum; j++) { EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); } + // ===== input uncompressed: read from vec_nv + flg_resunc = 0; + flg_vecunc = 1; + mju_zero(vec_nefci, efcnum); // clear output + mj_mulJacVec_island(model, data, vec_nefci, vec_nv, + i, flg_resunc, flg_vecunc); + + // expect corresponding values to match + for (int j=0; j < efcnum; j++) { + EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12)); + } + + // ===== output uncompressed: write to vec_nefc_tmp + flg_resunc = 1; + flg_vecunc = 0; + mju_zero(vec_nefc_tmp, data->nefc); // clear output + mj_mulJacVec_island(model, data, vec_nefc_tmp, vec_nvi, + i, flg_resunc, flg_vecunc); + + // expect corresponding values to match + for (int j=0; j < efcnum; j++) { + EXPECT_THAT(vec_nefc_tmp[efcind[j]], + DoubleNear(vec_nefc[efcind[j]], 1e-12)); + } + mju_free(vec_nvi); mju_free(vec_nefci); } + mju_free(vec_nefc_tmp); mju_free(vec_nefc); } @@ -268,6 +302,7 @@ TEST_F(CoreConstraintTest, MulJacTVecIsland) { // allocate vec_nv mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); + mjtNum* vec_nv_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv); // iterate through dense and sparse for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { @@ -297,27 +332,59 @@ TEST_F(CoreConstraintTest, MulJacTVecIsland) { int efcnum = data->island_efcnum[i]; mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); - // copy values into vec_nefci + // get indices int* efcind = data->island_efcind + data->island_efcadr[i]; + int* dofind = data->island_dofind + data->island_dofadr[i]; + + // copy values into vec_nefci 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); + // ==== both compressed + int flg_resunc = 0; + int flg_vecunc = 0; + mju_zero(vec_nvi, dofnum); // clear output + mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci, + i, flg_resunc, flg_vecunc); // 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)); } + // ===== input uncompressed: read from vec_nefc + flg_resunc = 0; + flg_vecunc = 1; + mju_zero(vec_nvi, dofnum); // clear output + mj_mulJacTVec_island(model, data, vec_nvi, vec_nefc, + i, flg_resunc, flg_vecunc); + + // expect corresponding values to match + for (int j=0; j < dofnum; j++) { + EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12)); + } + + // ===== output uncompressed: write to vec_nv_tmp + flg_resunc = 1; + flg_vecunc = 0; + mju_zero(vec_nv_tmp, model->nv); // clear output + mj_mulJacTVec_island(model, data, vec_nv_tmp, vec_nefci, + i, flg_resunc, flg_vecunc); + + // expect corresponding values to match + for (int j=0; j < dofnum; j++) { + EXPECT_THAT(vec_nv_tmp[dofind[j]], + DoubleNear(vec_nv[dofind[j]], 1e-12)); + } + mju_free(vec_nvi); mju_free(vec_nefci); } mju_free(vec_nefc); } + mju_free(vec_nv_tmp); mju_free(vec_nv); mj_deleteData(data); mj_deleteModel(model);