chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compute diagonal indices in mj_sqrMatTDSparse

Browse Source 
PiperOrigin-RevId: 713714087
Change-Id: Icc8eae74e6e47ba1d12a6e9aa0d774ab006cfe38
This commit is contained in:
Yuval Tassa
2025-01-09 10:16:37 -08:00
committed by Copybara-Service
parent daba00ca9d
commit 00010f5848
6 changed files with 53 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+5 -8
View File
@@ -2174,23 +2174,20 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// construct supernodes // construct supernodes
mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind); mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind);
// AR = JM2 * JM2' // pre-count efc_AR_rownnz, efc_AR_rowadr
mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT, mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT,
rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1); rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
// AR = JM2 * JM2'
int* diagind = mjSTACKALLOC(d, nefc, int);
mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc, mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
rownnzT, rowadrT, colindT, NULL, rownnzT, rowadrT, colindT, NULL,
rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1); rownnz, rowadr, colind, rowsuper, d, diagind);
// add R to diagonal of AR // add R to diagonal of AR
for (int i=0; i < nefc; i++) { for (int i=0; i < nefc; i++) {
for (int j=0; j < d->efc_AR_rownnz[i]; j++) { d->efc_AR[diagind[i]] += d->efc_R[i];
if (i == d->efc_AR_colind[d->efc_AR_rowadr[i]+j]) {
d->efc_AR[d->efc_AR_rowadr[i]+j] += d->efc_R[i];
break;
}
}
} }
} }
src/engine/engine_solver.c
+2 -2
View File
@@ -1424,7 +1424,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d, /*flg_upper=*/0); d, /*diagind=*/NULL);
// add mass matrix: H = J'*D*J + C // add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
@@ -1518,7 +1518,7 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d, /*flg_upper=*/0); d, /*diagind=*/NULL);
// add mass matrix: H = J'*D*J + C // add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
src/engine/engine_util_sparse.c
+9 -3
View File
@@ -724,7 +724,7 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const int* colind, const int* rowsuper, const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, const int* colindT, const int* rowsuperT,
mjData* d, int flg_upper) { mjData* d, int* diagind) {
// allocate space for accumulation buffer and matT // allocate space for accumulation buffer and matT
mj_markStack(d); mj_markStack(d);
@@ -838,8 +838,14 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
} }
// fill upper triangle // diagonal indices requested: fill upper triangle
if (flg_upper) { if (diagind) {
// save diagonal indices
for (int i=0; i < nc; i++) {
diagind[i] = res_rowadr[i] + res_rownnz[i] - 1;
}
// fill upper triangle
for (int i=0; i < nc; i++) { for (int i=0; i < nc; i++) {
int start = res_rowadr[i]; int start = res_rowadr[i];
int end = start + res_rownnz[i] - 1; int end = start + res_rownnz[i] - 1;
src/engine/engine_util_sparse.h
+2 -3
View File
@@ -89,8 +89,7 @@ MJAPI void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
MJAPI void mju_superSparse(int nr, int* rowsuper, MJAPI void mju_superSparse(int nr, int* rowsuper,
const int* rownnz, const int* rowadr, const int* colind); const int* rownnz, const int* rowadr, const int* colind);
// compute sparse M'*diag*M (diag=NULL: compute M'*M), res has uncompressed layout // compute sparse M'*diag*M (diag=NULL: compute M'*M), res_rowadr must be precomputed
// res_rowadr is required to be precomputed
MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const mjtNum* diag, int nr, int nc, const mjtNum* diag, int nr, int nc,
int* res_rownnz, const int* res_rowadr, int* res_colind, int* res_rownnz, const int* res_rowadr, int* res_colind,
@@ -98,7 +97,7 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT
const int* colind, const int* rowsuper, const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, const int* colindT, const int* rowsuperT,
mjData* d, int flg_upper); mjData* d, int* diagind);
// precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr, MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr,
test/benchmark/engine_util_sparse_benchmark_test.cc
+6 -4
View File
@@ -43,7 +43,7 @@ void ABSL_ATTRIBUTE_NOINLINE mju_sqrMatTDSparse_baseline(
int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind, int nr, int nc, int* res_rownnz, int* res_rowadr, int* res_colind,
const int* rownnz, const int* rowadr, const int* colind, const int* rownnz, const int* rowadr, const int* colind,
const int* rowsuper, const int* rownnzT, const int* rowadrT, const int* rowsuper, const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT, mjData* d, int unused) { const int* colindT, const int* rowsuperT, mjData* d, int* unused) {
mj_markStack(d); mj_markStack(d);
int* chain = mj_stackAllocInt(d, 2 * nc); int* chain = mj_stackAllocInt(d, 2 * nc);
mjtNum* buffer = mj_stackAllocNum(d, nc); mjtNum* buffer = mj_stackAllocNum(d, nc);
@@ -435,6 +435,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
int* rownnz = mj_stackAllocInt(d, m->nv); int* rownnz = mj_stackAllocInt(d, m->nv);
int* rowadr = mj_stackAllocInt(d, m->nv); int* rowadr = mj_stackAllocInt(d, m->nv);
int* colind = mj_stackAllocInt(d, m->nv*m->nv); int* colind = mj_stackAllocInt(d, m->nv*m->nv);
int* diagind = mj_stackAllocInt(d, m->nv);
// compute D corresponding to quad states // compute D corresponding to quad states
mjtNum* D = mj_stackAllocNum(d, d->nefc); mjtNum* D = mj_stackAllocNum(d, d->nefc);
@@ -454,7 +455,7 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
d->efc_J_colind, d->efc_J_rowsuper, d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d, d->efc_JT_colind, d->efc_JT_rowsuper, d,
/*flg_upper=*/1); diagind);
// compute H = M + J'*D*J // compute H = M + J'*D*J
mj_addM(m, d, H, rownnz, rowadr, colind); mj_addM(m, d, H, rownnz, rowadr, colind);
@@ -559,6 +560,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
int* rownnz = mj_stackAllocInt(d, m->nv); int* rownnz = mj_stackAllocInt(d, m->nv);
int* rowadr = mj_stackAllocInt(d, m->nv); int* rowadr = mj_stackAllocInt(d, m->nv);
int* colind = mj_stackAllocInt(d, m->nv * m->nv); int* colind = mj_stackAllocInt(d, m->nv * m->nv);
int* diagind = mj_stackAllocInt(d, m->nv);
// compute D corresponding to quad states // compute D corresponding to quad states
mjtNum* D = mj_stackAllocNum(d, d->nefc); mjtNum* D = mj_stackAllocNum(d, d->nefc);
@@ -579,7 +581,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind, func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, /*flg_upper=*/1); d->efc_JT_rowsuper, d, diagind);
} }
} else { } else {
for (auto s : state) { for (auto s : state) {
@@ -592,7 +594,7 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind, H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, /*unused=*/0); d->efc_JT_rowsuper, d, /*unused=*/nullptr);
} }
} }
test/engine/engine_util_sparse_test.cc
+29 -14
View File
@@ -324,6 +324,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
// test precount // test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
@@ -336,7 +337,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0)); EXPECT_THAT(matH, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -369,6 +370,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
// test precount // test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
@@ -382,7 +384,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(12, 0, 12, 0, 6, 3, 12, 3, 14)); EXPECT_THAT(matH, ElementsAre(12, 0, 12, 0, 6, 3, 12, 3, 14));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -415,6 +417,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4}; mjtNum diag[] = {2, 3, 4};
@@ -429,7 +432,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 4, 35, 0, 0, 0, 0)); EXPECT_THAT(matH, ElementsAre(66, 4, 0, 4, 35, 0, 0, 0, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 1, 0, 0, 0, 0)); EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 0, 1, 0, 0, 0, 0));
@@ -462,6 +465,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4}; mjtNum diag[] = {2, 3, 4};
@@ -477,7 +481,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 0, 0, 0, 0, 4, 35, 0)); EXPECT_THAT(matH, ElementsAre(66, 4, 0, 0, 0, 0, 4, 35, 0));
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 0, 0, 0, 0, 2, 0)); EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 0, 0, 0, 0, 2, 0));
@@ -510,6 +514,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
// test precount // test precount
@@ -523,7 +528,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(5, 6, 4, 6, 9, 0, 4, 16, 0)); EXPECT_THAT(matH, ElementsAre(5, 6, 4, 6, 9, 0, 4, 16, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
@@ -556,6 +561,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
// test precount // test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
@@ -568,12 +574,13 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, nullptr, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(1, 2, 0, 4, 0, 0, 2, 13, 0)); EXPECT_THAT(matH, ElementsAre(1, 2, 0, 4, 0, 0, 2, 13, 0));
EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 1, 0, 0, 0, 2, 0)); EXPECT_THAT(colindH, ElementsAre(0, 2, 0, 1, 0, 0, 0, 2, 0));
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2)); EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6));
EXPECT_THAT(diagindH, ElementsAre(0, 3, 7));
mj_deleteData(data); mj_deleteData(data);
mj_deleteModel(model); mj_deleteModel(model);
@@ -601,6 +608,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
int colindH[] = {0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0};
int rownnzH[] = {0, 0}; int rownnzH[] = {0, 0};
int rowadrH[] = {0, 0}; int rowadrH[] = {0, 0};
int diagindH[] = {0, 0};
mjtNum diag[] = {2, 3, 4}; mjtNum diag[] = {2, 3, 4};
@@ -615,7 +623,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
mju_sqrMatTDUncompressedInit(rowadrH, 2); mju_sqrMatTDUncompressedInit(rowadrH, 2);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 2, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 2, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(66, 4, 4, 35)); EXPECT_THAT(matH, ElementsAre(66, 4, 4, 35));
EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 1)); EXPECT_THAT(colindH, ElementsAre(0, 1, 0, 1));
@@ -647,6 +655,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3}; mjtNum diag[] = {2, 3};
@@ -661,7 +670,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 2, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 2, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(14, 18, 8, 18, 27, 0, 8, 32, 0)); EXPECT_THAT(matH, ElementsAre(14, 18, 8, 18, 27, 0, 8, 32, 0));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 0, 0, 2, 0));
@@ -694,6 +703,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {2, 3, 4}; mjtNum diag[] = {2, 3, 4};
@@ -708,7 +718,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
nullptr, data, 1); nullptr, data, diagindH);
EXPECT_THAT(matH, ElementsAre(69, 77, 80, 77, 99, 108, 80, 108, 120)); EXPECT_THAT(matH, ElementsAre(69, 77, 80, 77, 99, 108, 80, 108, 120));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -742,6 +752,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {1, 1, 1}; mjtNum diag[] = {1, 1, 1};
@@ -756,7 +767,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, 1); rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(14, 14, 14, 14, 14, 14, 14, 14, 14)); EXPECT_THAT(matH, ElementsAre(14, 14, 14, 14, 14, 14, 14, 14, 14));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -790,6 +801,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0}; int rownnzH[] = {0, 0, 0};
int rowadrH[] = {0, 0, 0}; int rowadrH[] = {0, 0, 0};
int diagindH[] = {0, 0, 0};
mjtNum diag[] = {1, 1, 1}; mjtNum diag[] = {1, 1, 1};
@@ -804,7 +816,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
mju_sqrMatTDUncompressedInit(rowadrH, 3); mju_sqrMatTDUncompressedInit(rowadrH, 3);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 3, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, 1); rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10)); EXPECT_THAT(matH, ElementsAre(1, 1, 1, 1, 10, 10, 1, 10, 10));
EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2)); EXPECT_THAT(colindH, ElementsAre(0, 1, 2, 0, 1, 2, 0, 1, 2));
@@ -838,6 +850,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; int colindH[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0}; int rownnzH[] = {0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0}; int rowadrH[] = {0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0};
mjtNum diag[] = {1, 1, 1}; mjtNum diag[] = {1, 1, 1};
@@ -852,7 +865,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
mju_sqrMatTDUncompressedInit(rowadrH, 4); mju_sqrMatTDUncompressedInit(rowadrH, 4);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 4, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 4, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, 1); rowsuperT, data, diagindH);
EXPECT_THAT(matH, EXPECT_THAT(matH,
ElementsAre(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 10, 1, 1, 10, 10)); ElementsAre(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 10, 10, 1, 1, 10, 10));
@@ -890,6 +903,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0, 0}; int rownnzH[] = {0, 0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0, 0}; int rowadrH[] = {0, 0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0, 0};
mjtNum diag[] = {1, 1, 1}; mjtNum diag[] = {1, 1, 1};
@@ -904,7 +918,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
mju_sqrMatTDUncompressedInit(rowadrH, 5); mju_sqrMatTDUncompressedInit(rowadrH, 5);
mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 5, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, diag, 3, 5, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, 1); rowsuperT, data, diagindH);
EXPECT_THAT(matH, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, EXPECT_THAT(matH, ElementsAre(3, 3, 0, 0, 0, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0)); 0, 0, 0, 0, 0, 0, 0, 0, 0));
@@ -942,6 +956,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
int rownnzH[] = {0, 0, 0, 0, 0, 0, 0}; int rownnzH[] = {0, 0, 0, 0, 0, 0, 0};
int rowadrH[] = {0, 0, 0, 0, 0, 0, 0}; int rowadrH[] = {0, 0, 0, 0, 0, 0, 0};
int diagindH[] = {0, 0, 0, 0, 0, 0, 0};
// test precount // test precount
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 7, rownnz, rowadr, colind, mju_sqrMatTDSparseCount(rownnzH, rowadrH, 7, rownnz, rowadr, colind,
@@ -954,7 +969,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
mju_sqrMatTDUncompressedInit(rowadrH, 7); mju_sqrMatTDUncompressedInit(rowadrH, 7);
mju_sqrMatTDSparse(matH, mat, matT, nullptr, 1, 7, rownnzH, rowadrH, colindH, mju_sqrMatTDSparse(matH, mat, matT, nullptr, 1, 7, rownnzH, rowadrH, colindH,
rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT, rownnz, rowadr, colind, nullptr, rownnzT, rowadrT, colindT,
rowsuperT, data, 1); rowsuperT, data, diagindH);
EXPECT_THAT( EXPECT_THAT(
matH, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, matH, ElementsAre(1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 1, 1, 2,