chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move mju_cholFactorCount to engine_util_solve.

Browse Source 
PiperOrigin-RevId: 746011856
Change-Id: If9812251420053644f4eca81adb5116d370ee524
This commit is contained in:
Yuval Tassa
2025-04-10 06:55:37 -07:00
committed by Copybara-Service
parent 3c530976e3
commit 25126e88c7
7 changed files with 150 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_util_solve.c
+54
View File
@@ -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
+4
View File
@@ -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
-52
View File
@@ -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
-3
View File
@@ -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
+75 -1
View File
@@ -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
+14 -94
View File
@@ -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
+3 -2
View File
@@ -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) {
return std::vector<mjtNum>(array, array + n);
template <typename T>
std::vector<T> AsVector(const T* array, int n) {
return std::vector<T>(array, array + n);
}
// Prints a matrix to stderr, useful for debugging.