diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 10120294..c1313cd6 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2175,8 +2175,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { mju_superSparse(nefc, rowsuper, rownnz, rowadr, colind); // AR = JM2 * JM2' - mju_sqrMatTDSparseInit(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT, - rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1); + mju_sqrMatTDSparseCount(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT, + rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1); mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc, d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 39ca9b30..7f0fe182 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -1400,10 +1400,10 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { } // initialize Hessian rowadr, rownnz - mju_sqrMatTDSparseInit(ctx->H_rownnz, ctx->H_rowadr, nv, - d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, - d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper, - d, /*flg_upper=*/0); + mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, + d->efc_JT_rowsuper, d, /*flg_upper=*/0); // add nC to Hessian total nonzeros (unavoidable overcounting since H_colind is still unknown) ctx->nH = m->nC + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1]; @@ -1440,7 +1440,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) { ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind); // count total and row non-zeros of reverse-Cholesky factor L - ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, HT_rownnz, HT_rowadr, HT_colind, nv, d); + ctx->nL = mju_cholFactorCount(ctx->L_rownnz, HT_rownnz, HT_rowadr, HT_colind, nv, d); mj_freeStack(d); // compute L row adresses: rowadr = cumsum(rownnz) diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 444edbf2..e4e1c1c9 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -625,10 +625,10 @@ void mju_superSparse(int nr, int* rowsuper, // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse -void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr, - const int* rownnz, const int* rowadr, const int* colind, - const int* rownnzT, const int* rowadrT, const int* colindT, - const int* rowsuperT, mjData* d, int flg_upper) { +void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr, + const int* rownnz, const int* rowadr, const int* colind, + const int* rownnzT, const int* rowadrT, const int* colindT, + const int* rowsuperT, mjData* d, int flg_upper) { mj_markStack(d); int* chain = mjSTACKALLOC(d, 2*nr, int); int nchain = 0; @@ -865,11 +865,11 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT, mj_freeStack(d); } -// compute row non-zeros of reverse-Cholesky factor L, return total non-zeros -// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package' -// note: reads pattern from upper triangle -int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind, - int n, mjData* d) { +// precount row non-zeros of reverse-Cholesky factor L, return total non-zeros +// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package' +// reads pattern from upper triangle +int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind, + int n, mjData* d) { mj_markStack(d); int* parent = mjSTACKALLOC(d, n, int); int* flag = mjSTACKALLOC(d, n, int); @@ -880,24 +880,28 @@ int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const flag[r] = r; L_rownnz[r] = 1; // start with 1 for diagonal - // loop over non-zero columns + // loop over non-zero columns of upper triangle int start = rowadr[r]; int end = start + rownnz[r]; for (int c = start; c < end; c++) { int i = colind[c]; - if (i > r) { - // traverse from i to ancestor, stop when row is flagged - while (flag[i] != r) { - // if not yet set, set parent to current row - if (parent[i] == -1) { - parent[i] = r; - } - // increment non-zeros, flag row i, advance to parent - L_rownnz[i]++; - flag[i] = r; - i = parent[i]; + // skip lower triangle + if (i <= r) { + continue; + } + + // traverse from i to ancestor, stop when row is flagged + while (flag[i] != r) { + // if not yet set, set parent to current row + if (parent[i] == -1) { + parent[i] = r; } + + // increment non-zeros, flag row i, advance to parent + L_rownnz[i]++; + flag[i] = r; + i = parent[i]; } } } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 62cf6890..04c37610 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -99,17 +99,17 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT mjData* d, int flg_upper); // precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse -MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr, - const int* rownnz, const int* rowadr, const int* colind, - const int* rownnzT, const int* rowadrT, const int* colindT, - const int* rowsuperT, mjData* d, int flg_upper); +MJAPI void mju_sqrMatTDSparseCount(int* res_rownnz, int* res_rowadr, int nr, + const int* rownnz, const int* rowadr, const int* colind, + const int* rownnzT, const int* rowadrT, const int* colindT, + const int* rowsuperT, mjData* d, int flg_upper); // precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc); -// compute row non-zeros of reverse-Cholesky factor L, return total -MJAPI int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind, - int n, mjData* d); +// precount row non-zeros of reverse-Cholesky factor L, return total +MJAPI int mju_cholFactorCount(int* L_rownnz, const int* rownnz, const int* rowadr, + const int* colind, int n, mjData* d); // ------------------------------ inlined functions ------------------------------------------------ diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 37da5155..d1057b22 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -326,8 +326,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) { int rowadrH[] = {0, 0, 0}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); @@ -371,8 +371,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) { int rowadrH[] = {0, 0, 0}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); @@ -419,8 +419,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) { mjtNum diag[] = {2, 3, 4}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0)); EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4)); @@ -467,8 +467,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) { // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(2, 0, 2)); EXPECT_THAT(rowadrH, ElementsAre(0, 2, 2)); @@ -513,8 +513,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) { // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5)); @@ -558,8 +558,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) { int rowadrH[] = {0, 0, 0}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2)); EXPECT_THAT(rowadrH, ElementsAre(0, 2, 3)); @@ -605,8 +605,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) { mjtNum diag[] = {2, 3, 4}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 2, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 2, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(2, 2)); EXPECT_THAT(rowadrH, ElementsAre(0, 2)); @@ -651,8 +651,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) { mjtNum diag[] = {2, 3}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 5)); @@ -698,8 +698,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) { mjtNum diag[] = {2, 3, 4}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, nullptr, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, nullptr, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); @@ -746,8 +746,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) { mjtNum diag[] = {1, 1, 1}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, rowsuperT, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, rowsuperT, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); @@ -794,8 +794,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) { mjtNum diag[] = {1, 1, 1}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, rowsuperT, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, rowsuperT, data, 1); EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3)); EXPECT_THAT(rowadrH, ElementsAre(0, 3, 6)); @@ -842,8 +842,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) { mjtNum diag[] = {1, 1, 1}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 4, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, rowsuperT, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 4, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, rowsuperT, data, 1); EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4)); EXPECT_THAT(rowadrH, ElementsAre(0, 4, 8, 12)); @@ -894,8 +894,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) { mjtNum diag[] = {1, 1, 1}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 5, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, rowsuperT, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 5, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, rowsuperT, data, 1); EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0, 0, 0)); EXPECT_THAT(rowadrH, ElementsAre(0, 2, 4, 4, 4)); @@ -944,8 +944,8 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) { int rowadrH[] = {0, 0, 0, 0, 0, 0, 0}; // test precount - mju_sqrMatTDSparseInit(rownnzH, rowadrH, 7, rownnz, rowadr, colind, - rownnzT, rowadrT, colindT, rowsuperT, data, 1); + mju_sqrMatTDSparseCount(rownnzH, rowadrH, 7, rownnz, rowadr, colind, + rownnzT, rowadrT, colindT, rowsuperT, data, 1); EXPECT_THAT(rownnzH, ElementsAre(7, 7, 7, 7, 7, 7, 7)); EXPECT_THAT(rowadrH, ElementsAre(0, 7, 14, 21, 28, 35, 42)); @@ -985,8 +985,8 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { int colindA[4]; int rownnzA_factor[2]; mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA, 4); - int nnzA = mju_cholFactorNNZ(rownnzA_factor, - rownnzA, rowadrA, colindA, nA, d); + int nnzA = mju_cholFactorCount(rownnzA_factor, + rownnzA, rowadrA, colindA, nA, d); EXPECT_EQ(nnzA, 2); EXPECT_THAT(AsVector(rownnzA_factor, 2), ElementsAre(1, 1)); @@ -1001,8 +1001,8 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { int colindB[9]; int rownnzB_factor[3]; mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB, 9); - int nnzB = mju_cholFactorNNZ(rownnzB_factor, - rownnzB, rowadrB, colindB, nB, d); + int nnzB = mju_cholFactorCount(rownnzB_factor, + rownnzB, rowadrB, colindB, nB, d); EXPECT_EQ(nnzB, 5); EXPECT_THAT(AsVector(rownnzB_factor, 3), ElementsAre(1, 2, 2)); @@ -1017,8 +1017,8 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { int colindC[9]; int rownnzC_factor[3]; mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC, 9); - int nnzC = mju_cholFactorNNZ(rownnzC_factor, - rownnzC, rowadrC, colindC, nC, d); + int nnzC = mju_cholFactorCount(rownnzC_factor, + rownnzC, rowadrC, colindC, nC, d); EXPECT_EQ(nnzC, 4); EXPECT_THAT(AsVector(rownnzC_factor, 3), ElementsAre(1, 2, 1)); @@ -1034,8 +1034,8 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { int colindD[16]; int rownnzD_factor[4]; mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD, 16); - int nnzD = mju_cholFactorNNZ(rownnzD_factor, - rownnzD, rowadrD, colindD, nD, d); + int nnzD = mju_cholFactorCount(rownnzD_factor, + rownnzD, rowadrD, colindD, nD, d); EXPECT_EQ(nnzD, 8); EXPECT_THAT(AsVector(rownnzD_factor, 4), ElementsAre(1, 2, 2, 3));