Improve comments in mju_cholFactorNNZ, clarify that this function rea…

…ds from the upper triangle of the source matrix. PiperOrigin-RevId: 686113654 Change-Id: I029f9279d177f4cb9ff60396a26d64821aee3a6c
google-deepmind · Oct 15, 2024 · 0ffaeb4 · 0ffaeb4
1 parent 096853e
commit 0ffaeb4
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 27 deletions.
diff --git 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
@@ -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];