Move mju_cholFactorCount to engine_util_solve.

PiperOrigin-RevId: 746011856
Change-Id: If9812251420053644f4eca81adb5116d370ee524

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

test/engine/engine_util_sparse_test.cc

@@ -15,7 +15,6 @@
 // Tests for engine/engine_util_sparse.c
 
 #include <array>
-#include <vector>
 
 #include "src/engine/engine_util_sparse.h"
 
@@ -30,11 +29,6 @@ namespace {
 using ::testing::ElementsAre;
 using EngineUtilSparseTest = MujocoTest;
 
-template <typename T>
-std::vector<T> AsVector(const T* array, int n) {
-  return std::vector<T>(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};
@@ -335,14 +329,12 @@ TEST_F(EngineUtilSparseTest, MjuCompressSparse) {
   EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected_minval1, 6));
 }
 
-static constexpr char modelStr[] = R"(<mujoco/>)";
-
 TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {
   //       0 0 0
   // M  =  0 0 0
   //       0 0 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
@@ -388,7 +380,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse2) {
   // M = 1  2 -1
   //     2  2  3
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {2, -1, 1, 2, -1, 2, 2, 2, 3};
@@ -435,7 +427,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse3) {
   // M = 0 3 0
   //     4 0 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 2, 3, 4};
@@ -483,7 +475,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse4) {
   // M = 0 0 3
   //     4 0 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 2, 3, 4};
@@ -532,7 +524,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse5) {
   // M = 0 0 0
   //     2 3 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 4, 2, 3};
@@ -579,7 +571,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse6) {
   // M = 0 2 0
   //     0 0 3
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 2, 2, 3};
@@ -626,7 +618,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse7) {
   // M = 0 3
   //     4 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 2, 3, 4};
@@ -673,7 +665,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse8) {
   // M = 1 0 4
   //     2 3 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 4, 2, 3};
@@ -721,7 +713,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse9) {
   // M = 1 3 4
   //     4 4 4
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 2, 2, 1, 3, 4, 4, 4, 4};
@@ -769,7 +761,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse10) {
   // M = 2 2 2
   //     3 3 3
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 1, 1, 2, 2, 2, 3, 3, 3};
@@ -818,7 +810,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse11) {
   // M = 0 0 0
   //     0 3 3
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 1, 1, 3, 3};
@@ -867,7 +859,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse12) {
   // M = 0 0 0 0
   //     0 0 3 3
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 1, 1, 1, 3, 3};
@@ -918,7 +910,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
   // M = 1 1 0 0 0
   //     1 1 0 0 0
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 1, 1, 1, 1, 1};
@@ -969,7 +961,7 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse13) {
 TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   // M = 1 1 1 1 2 2 2
 
-  mjModel* model = LoadModelFromString(modelStr);
+  mjModel* model = LoadModelFromString("<mujoco/>");
   mjData* data = mj_makeData(model);
 
   mjtNum mat[] = {1, 1, 1, 1, 2, 2, 2};
@@ -1022,78 +1014,6 @@ TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse14) {
   mj_deleteModel(model);
 }
 
-TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
-  mjModel* model = LoadModelFromString(modelStr);
-  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);
-}
-
 TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
   int nr = 2;
   int nc = 3;

test/fixture.h

@@ -108,8 +108,9 @@ std::vector<mjtNum> GetCtrlNoise(const mjModel* m, int nsteps,
 mjtNum CompareModel(const mjModel* m1, const mjModel* m2, std::string& field);
 
 // Returns a vector containing the elements of the array.
-inline std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
+template <typename T>
-  return std::vector<mjtNum>(array, array + n);
+std::vector<T> AsVector(const T* array, int n) {
+  return std::vector<T>(array, array + n);
 }
 
 // Prints a matrix to stderr, useful for debugging.