diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 815a57c2..cc2c7fd4 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -857,8 +857,9 @@ 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 +// 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) { mj_markStack(d); @@ -870,18 +871,21 @@ int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const parent[r] = -1; flag[r] = r; L_rownnz[r] = 1; // start with 1 for diagonal + + // loop over non-zero columns int start = rowadr[r]; int end = start + rownnz[r]; - // loop over non-zero columns - for (int p = start; p < end; p++) { - int i = colind[p]; + for (int c = start; c < end; c++) { + int i = colind[c]; if (i > r) { - // follow path from i to root of elimination tree, stop at flagged node + // traverse from i to ancestor, stop when row is flagged while (flag[i] != r) { - // find parent of i if not yet determined + // 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]; @@ -892,12 +896,11 @@ int mju_cholFactorNNZ(int* L_rownnz, const int* rownnz, const int* rowadr, const mj_freeStack(d); - // accumulate sum - int sum = 0; + // sum up all row non-zeros + int nnz = 0; for (int r = 0; r < n; r++) { - sum += L_rownnz[r]; + nnz += L_rownnz[r]; } - // return total non-zeros - return sum; + return nnz; } diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 16f80ca5..dfdadd82 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -975,10 +975,9 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { mjModel* model = LoadModelFromString(modelStr); mjData* d = mj_makeData(model); - // A = [[1, 0], - // [0, 1]] int nA = 2; - mjtNum matA[4] = {1, 0, 0, 1}; + mjtNum matA[4] = {1, 0, + 0, 1}; mjtNum sparseA[4]; int rownnzA[2]; int rowadrA[2]; @@ -991,11 +990,10 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { EXPECT_EQ(nnzA, 2); EXPECT_THAT(AsVector(rownnzA_factor, 2), ElementsAre(1, 1)); - // B = [[10, 1, 0], - // [1, 10, 1], - // [0, 1, 10]] int nB = 3; - mjtNum matB[9] = {10, 1, 0, 1, 10, 1, 0, 1, 10}; + mjtNum matB[9] = {10, 1, 0, + 0, 10, 1, + 0, 0, 10}; mjtNum sparseB[9]; int rownnzB[3]; int rowadrB[3]; @@ -1008,11 +1006,10 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { EXPECT_EQ(nnzB, 5); EXPECT_THAT(AsVector(rownnzB_factor, 3), ElementsAre(1, 2, 2)); - // C = [[10, 1, 0], - // [1, 10, 0], - // [0, 0, 10]] int nC = 3; - mjtNum matC[9] = {10, 1, 0, 1, 10, 0, 0, 0, 10}; + mjtNum matC[9] = {10, 1, 0, + 0, 10, 0, + 0, 0, 10}; mjtNum sparseC[9]; int rownnzC[3]; int rowadrC[3]; @@ -1025,12 +1022,11 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) { EXPECT_EQ(nnzC, 4); EXPECT_THAT(AsVector(rownnzC_factor, 3), ElementsAre(1, 2, 1)); - // D = [[10, 1, 2, 3], - // [1, 10, 0, 0], - // [2, 0, 10, 1], - // [3, 0, 1, 10]] int nD = 4; - mjtNum matD[16] = {10, 1, 2, 3, 1, 10, 0, 0, 2, 0, 10, 1, 3, 0, 1, 10}; + mjtNum matD[16] = {10, 1, 2, 3, + 0, 10, 0, 0, + 0, 0, 10, 1, + 0, 0, 0, 10}; mjtNum sparseD[16]; int rownnzD[4]; int rowadrD[4];