diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 83fdb634..11889667 100644 --- a/src/engine/engine_util_solve.c +++ b/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) { diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 6bbdc6bb..6a772e2c 100644 --- a/src/engine/engine_util_solve.h +++ b/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); diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index abb95be0..d72a9fca 100644 --- a/src/engine/engine_util_sparse.c +++ b/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; -} diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 92c2dad0..08a663b2 100644 --- a/src/engine/engine_util_sparse.h +++ b/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 ------------------------------------------------ diff --git a/test/engine/engine_util_solve_test.cc b/test/engine/engine_util_solve_test.cc index d4bea9a1..aad5b421 100644 --- a/test/engine/engine_util_solve_test.cc +++ b/test/engine/engine_util_solve_test.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include @@ -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(""); + 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 diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index c1b7cb65..94da0d5d 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -15,7 +15,6 @@ // Tests for engine/engine_util_sparse.c #include -#include #include "src/engine/engine_util_sparse.h" @@ -30,11 +29,6 @@ namespace { using ::testing::ElementsAre; using EngineUtilSparseTest = MujocoTest; -template -std::vector AsVector(const T* array, int n) { - return std::vector(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"()"; - TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) { // 0 0 0 // M = 0 0 0 // 0 0 0 - mjModel* model = LoadModelFromString(modelStr); + mjModel* model = LoadModelFromString(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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(""); 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; diff --git a/test/fixture.h b/test/fixture.h index df568943..cbb71d5f 100644 --- a/test/fixture.h +++ b/test/fixture.h @@ -108,8 +108,9 @@ std::vector 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 AsVector(const mjtNum* array, int n) { - return std::vector(array, array + n); +template +std::vector AsVector(const T* array, int n) { + return std::vector(array, array + n); } // Prints a matrix to stderr, useful for debugging.