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

Add mju_cholFactorNNZ to compute the number of non-zeros per row for sparse Cholesky factorization.

Browse Source 
PiperOrigin-RevId: 679553197
Change-Id: I6f92deec347dae97d2e219cfb5376f92e777da1a
This commit is contained in:
Taylor Howell
2024-09-27 05:38:07 -07:00
committed by Copybara-Service
parent bb6da503dc
commit 6a6e10d779
3 changed files with 126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/engine/engine_util_sparse_test.cc
+83
View File
@@ -15,6 +15,7 @@
// Tests for engine/engine_util_sparse.c
#include <array>
#include <vector>
#include "src/engine/engine_util_sparse.h"
@@ -29,6 +30,10 @@ namespace {
using ::testing::ElementsAre;
using EngineUtilSparseTest = MujocoTest;
std::vector<int> AsVector(const int* array, int n) {
return std::vector<int>(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