Move mju_cholFactorCount to engine_util_solve.

PiperOrigin-RevId: 746011856
Change-Id: If9812251420053644f4eca81adb5116d370ee524

Author: yuvaltassa

src/engine/engine_util_solve.c

@@ -191,6 +191,60 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
 
 
 
+// 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);
+
+  // loop over rows in reverse order
+  for (int r = n - 1; r >= 0; r--) {
+    parent[r] = -1;
+    flag[r] = r;
+    L_rownnz[r] = 1;  // start with 1 for diagonal
+
+    // 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];
+
+      // 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];
+      }
+    }
+  }
+
+  mj_freeStack(d);
+
+  // sum up all row non-zeros
+  int nnz = 0;
+  for (int r = 0; r < n; r++) {
+    nnz += L_rownnz[r];
+  }
+
+  return nnz;
+}
+
+
+
 // sparse reverse-order Cholesky solve
 void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
                          const int* rownnz, const int* rowadr, const int* colind) {

src/engine/engine_util_solve.h

@@ -38,6 +38,10 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
                          int* rownnz, const int* rowadr, int* colind,
                          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);
+
 // sparse reverse-order Cholesky solve
 void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
                          const int* rownnz, const int* rowadr, const int* colind);

src/engine/engine_util_sparse.c

@@ -795,55 +795,3 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
 
   mj_freeStack(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);
-
-  // loop over rows in reverse order
-  for (int r = n - 1; r >= 0; r--) {
-    parent[r] = -1;
-    flag[r] = r;
-    L_rownnz[r] = 1;  // start with 1 for diagonal
-
-    // 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];
-
-      // 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];
-      }
-    }
-  }
-
-  mj_freeStack(d);
-
-  // sum up all row non-zeros
-  int nnz = 0;
-  for (int r = 0; r < n; r++) {
-    nnz += L_rownnz[r];
-  }
-
-  return nnz;
-}

src/engine/engine_util_sparse.h

@@ -103,9 +103,6 @@ MJAPI void mju_sqrMatTDSparseCount(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);
 
-// 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_solve_test.cc

@@ -21,7 +21,6 @@
 #include <random>
 #include <iomanip>
 #include <string>
-#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -36,6 +35,7 @@ namespace {
 using ::testing::DoubleEq;
 using ::testing::Pointwise;
 using ::testing::DoubleNear;
+using ::testing::ElementsAre;
 using ::std::string;
 using ::std::setw;
 using QCQP2Test = MujocoTest;
@@ -649,5 +649,79 @@ TEST_F(BandMatrixTest, Solve) {
   }
 }
 
+using EngineUtilSolveTest = MujocoTest;
+
+TEST_F(EngineUtilSolveTest, MjuCholFactorNNZ) {
+  mjModel* model = LoadModelFromString("<mujoco/>");
+  mjData* d = mj_makeData(model);
+
+  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];
+  mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA, 4);
+  int nnzA = mju_cholFactorCount(rownnzA_factor,
+                                 rownnzA, rowadrA, colindA, nA, d);
+
+  EXPECT_EQ(nnzA, 2);
+  EXPECT_THAT(AsVector(rownnzA_factor, 2), ElementsAre(1, 1));
+
+  int nB = 3;
+  mjtNum matB[9] = {10, 1, 0,
+                    0, 10, 1,
+                    0, 0, 10};
+  mjtNum sparseB[9];
+  int rownnzB[3];
+  int rowadrB[3];
+  int colindB[9];
+  int rownnzB_factor[3];
+  mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB, 9);
+  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));
+
+  int nC = 3;
+  mjtNum matC[9] = {10, 1, 0,
+                    0, 10, 0,
+                    0, 0, 10};
+  mjtNum sparseC[9];
+  int rownnzC[3];
+  int rowadrC[3];
+  int colindC[9];
+  int rownnzC_factor[3];
+  mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC, 9);
+  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));
+
+  int nD = 4;
+  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];
+  int colindD[16];
+  int rownnzD_factor[4];
+  mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD, 16);
+  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));
+
+  mj_deleteData(d);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco