From 6a6e10d779428922fe6c7480f4412587b30d9e17 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Fri, 27 Sep 2024 05:38:07 -0700 Subject: [PATCH] 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 | 40 +++++++++++++ src/engine/engine_util_sparse.h | 3 + test/engine/engine_util_sparse_test.cc | 83 ++++++++++++++++++++++++++ 3 files changed, 126 insertions(+) diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index c12132d0..da9982f4 100644 --- a/src/engine/engine_util_sparse.c +++ b/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; +} diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index ff3b9d00..b3619dcd 100644 --- a/src/engine/engine_util_sparse.h +++ b/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 } diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 495a45c9..06260c29 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -15,6 +15,7 @@ // Tests for engine/engine_util_sparse.c #include +#include #include "src/engine/engine_util_sparse.h" @@ -29,6 +30,10 @@ namespace { using ::testing::ElementsAre; using EngineUtilSparseTest = MujocoTest; +std::vector AsVector(const int* 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}; @@ -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