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

Standardize names of sparse fill-in pre-counting functions

Browse Source 
PiperOrigin-RevId: 712835830
Change-Id: I8e90dfa52af56ede917e1fd90b8d551898c244e5
This commit is contained in:
Yuval Tassa
2025-01-07 02:38:50 -08:00
committed by Copybara-Service
parent 0be7e6a1f6
commit 674d227050
5 changed files with 75 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+1 -1
View File
@@ -2175,7 +2175,7 @@ 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,
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,
src/engine/engine_solver.c
+4 -4
View File
@@ -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,
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);
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)
src/engine/engine_util_sparse.c
+11 -7
View File
@@ -625,7 +625,7 @@ 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,
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) {
@@ -865,10 +865,10 @@ 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
// 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'
// note: reads pattern from upper triangle
int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const int* colind,
// 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);
@@ -880,12 +880,17 @@ 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) {
// 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
@@ -900,7 +905,6 @@ int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const
}
}
}
}
mj_freeStack(d);
src/engine/engine_util_sparse.h
+4 -4
View File
@@ -99,7 +99,7 @@ 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,
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);
@@ -107,9 +107,9 @@ MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
// 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 ------------------------------------------------
test/engine/engine_util_sparse_test.cc
+18 -18
View File
@@ -326,7 +326,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
int rowadrH[] = {0, 0, 0};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
@@ -371,7 +371,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
int rowadrH[] = {0, 0, 0};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
@@ -419,7 +419,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0));
@@ -467,7 +467,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 0, 2));
@@ -513,7 +513,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
@@ -558,7 +558,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
int rowadrH[] = {0, 0, 0};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 1, 2));
@@ -605,7 +605,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 2, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 2, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 2));
@@ -651,7 +651,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
mjtNum diag[] = {2, 3};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 2, 2));
@@ -698,7 +698,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
mjtNum diag[] = {2, 3, 4};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, nullptr, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
@@ -746,7 +746,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
@@ -794,7 +794,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 3, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(3, 3, 3));
@@ -842,7 +842,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 4, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 4, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(4, 4, 4, 4));
@@ -894,7 +894,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
mjtNum diag[] = {1, 1, 1};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 5, rownnz, rowadr, colind,
mju_sqrMatTDSparseCount(rownnzH, rowadrH, 5, rownnz, rowadr, colind,
rownnzT, rowadrT, colindT, rowsuperT, data, 1);
EXPECT_THAT(rownnzH, ElementsAre(2, 2, 0, 0, 0));
@@ -944,7 +944,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
int rowadrH[] = {0, 0, 0, 0, 0, 0, 0};
// test precount
mju_sqrMatTDSparseInit(rownnzH, rowadrH, 7, rownnz, rowadr, colind,
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));
@@ -985,7 +985,7 @@ 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,
int nnzA = mju_cholFactorCount(rownnzA_factor,
rownnzA, rowadrA, colindA, nA, d);
EXPECT_EQ(nnzA, 2);
@@ -1001,7 +1001,7 @@ 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,
int nnzB = mju_cholFactorCount(rownnzB_factor,
rownnzB, rowadrB, colindB, nB, d);
EXPECT_EQ(nnzB, 5);
@@ -1017,7 +1017,7 @@ 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,
int nnzC = mju_cholFactorCount(rownnzC_factor,
rownnzC, rowadrC, colindC, nC, d);
EXPECT_EQ(nnzC, 4);
@@ -1034,7 +1034,7 @@ 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,
int nnzD = mju_cholFactorCount(rownnzD_factor,
rownnzD, rowadrD, colindD, nD, d);
EXPECT_EQ(nnzD, 8);