Add mju_cholFactorNNZ to compute the number of non-zeros per row for sparse Cholesky factorization.

PiperOrigin-RevId: 679553197
Change-Id: I6f92deec347dae97d2e219cfb5376f92e777da1a

src/engine/engine_util_sparse.c

@@ -860,3 +860,43 @@ 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
+// based on ldl_symbolic from 'Algorithm 8xx: a concise sparse Cholesky factorization package'
+int mju_cholFactorNNZ(int* L_rownnz, int* parent, int* flag, const int* rownnz,
+                      const int* rowadr, const int* colind, int n) {
+  // loop over rows in reverse order
+  for (int r = n - 1; r >= 0; r--) {
+    parent[r] = -1;
+    flag[r] = r;
+    L_rownnz[r] = 0;
+    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];
+      if (i > r) {
+        // follow path from i to root of elimination tree, stop at flagged node
+        while (flag[i] != r) {
+          // find parent of i if not yet determined
+          if (parent[i] == -1) {
+            parent[i] = r;
+          }
+          L_rownnz[i]++;
+          flag[i] = r;
+          i = parent[i];
+        }
+      }
+    }
+  }
+
+  // add 1 for diagonal, accumulate sum
+  int sum = 0;
+  for (int r = 0; r < n; r++) {
+    L_rownnz[r]++;
+    sum += L_rownnz[r];
+  }
+
+  // return total non-zeros
+  return sum;
+}

src/engine/engine_util_sparse.h

@@ -108,6 +108,9 @@ MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr,
 // 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, int* parent, int* flag, const int* rownnz,
+                            const int* rowadr, const int* colind, int n);
 
 #ifdef __cplusplus
 }

test/engine/engine_util_sparse_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_util_sparse.c
 
 #include <array>
+#include <vector>
 
 #include "src/engine/engine_util_sparse.h"
 
@@ -29,6 +30,10 @@ namespace {
 using ::testing::ElementsAre;
 using EngineUtilSparseTest = MujocoTest;
 
+std::vector<int> AsVector(const int* array, int n) {
+  return std::vector<int>(array, array + n);
+}
+
 TEST_F(EngineUtilSparseTest, MjuDot) {
   mjtNum a[] = {2,    3,       4,          5,          6,       7,    8};
   mjtNum u[] = {2, 1, 3, 1, 1, 4, 1, 1, 1, 5, 1, 1, 1, 6, 1, 1, 7, 1, 8};
@@ -966,5 +971,83 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   mj_deleteModel(model);
 }
 
+TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
+  // A = [[1, 0],
+  //      [0, 1]]
+  int nA = 2;
+  mjtNum matA[4] = {1, 0, 0, 1};
+  mjtNum sparseA[4];
+  int rownnzA[2];
+  int rowadrA[2];
+  int colindA[4];
+  int rownnzA_factor[2];
+  int parentA[2];
+  int workspaceA[2];
+  mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA);
+  int nnzA = mju_cholFactorNNZ(rownnzA_factor, parentA, workspaceA, rownnzA,
+                               rowadrA, colindA, nA);
+
+  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 sparseB[9];
+  int rownnzB[3];
+  int rowadrB[3];
+  int colindB[9];
+  int rownnzB_factor[3];
+  int parentB[3];
+  int workspaceB[3];
+  mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB);
+  int nnzB = mju_cholFactorNNZ(rownnzB_factor, parentB, workspaceB, rownnzB,
+                               rowadrB, colindB, nB);
+
+  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 sparseC[9];
+  int rownnzC[3];
+  int rowadrC[3];
+  int colindC[9];
+  int rownnzC_factor[3];
+  int parentC[3];
+  int workspaceC[3];
+  mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC);
+  int nnzC = mju_cholFactorNNZ(rownnzC_factor, parentC, workspaceC, rownnzC,
+                               rowadrC, colindC, nC);
+
+  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 sparseD[16];
+  int rownnzD[4];
+  int rowadrD[4];
+  int colindD[16];
+  int rownnzD_factor[4];
+  int parentD[4];
+  int workspaceD[4];
+  mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD);
+  int nnzD = mju_cholFactorNNZ(rownnzD_factor, parentD, workspaceD, rownnzD,
+                               rowadrD, colindD, nD);
+
+  EXPECT_EQ(nnzD, 8);
+  EXPECT_THAT(AsVector(rownnzD_factor, 4), ElementsAre(1, 2, 2, 3));
+}
+
 }  // namespace
 }  // namespace mujoco