Add nnz argument for size of allocated memory to mju_dense2sparse and check if this number is too small for number of non-zeros.

PiperOrigin-RevId: 690545825
Change-Id: I0e31cb907a1151cb2e4766d5f0cdfecb5629a2a6
This commit is contained in:
Taylor Howell
2024-10-28 03:46:00 -07:00
committed by Copybara-Service
parent 1c424644dd
commit a36f2cccb6
4 changed files with 63 additions and 15 deletions
+42 -6
View File
@@ -984,7 +984,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrA[2];
int colindA[4];
int rownnzA_factor[2];
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA);
mju_dense2sparse(sparseA, matA, nA, nA, rownnzA, rowadrA, colindA, 4);
int nnzA = mju_cholFactorNNZ(rownnzA_factor,
rownnzA, rowadrA, colindA, nA, d);
@@ -1000,7 +1000,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrB[3];
int colindB[9];
int rownnzB_factor[3];
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB);
mju_dense2sparse(sparseB, matB, nB, nB, rownnzB, rowadrB, colindB, 9);
int nnzB = mju_cholFactorNNZ(rownnzB_factor,
rownnzB, rowadrB, colindB, nB, d);
@@ -1016,7 +1016,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrC[3];
int colindC[9];
int rownnzC_factor[3];
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC);
mju_dense2sparse(sparseC, matC, nC, nC, rownnzC, rowadrC, colindC, 9);
int nnzC = mju_cholFactorNNZ(rownnzC_factor,
rownnzC, rowadrC, colindC, nC, d);
@@ -1033,7 +1033,7 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
int rowadrD[4];
int colindD[16];
int rownnzD_factor[4];
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD);
mju_dense2sparse(sparseD, matD, nD, nD, rownnzD, rowadrD, colindD, 16);
int nnzD = mju_cholFactorNNZ(rownnzD_factor,
rownnzD, rowadrD, colindD, nD, d);
@@ -1050,11 +1050,11 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
mjtNum mat[] = {1, 2, 0,
0, 3, 4};
mjtNum mat_sparse[6];
mjtNum mat_sparse[4];
int rownnz[2];
int rowadr[2];
int colind[4];
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind);
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 4);
// multiply: res = mat' * vec
mjtNum vec[] = {5, 6};
@@ -1064,5 +1064,41 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
}
TEST_F(EngineUtilSparseTest, MjuDenseToSparse) {
int nr = 2;
int nc = 2;
mjtNum mat[] = {1, 2,
0, 3};
mjtNum mat_sparse[4];
int rownnz[2];
int rowadr[2];
int colind[4];
// nnz == number of non-zeros
int status3 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 3);
EXPECT_EQ(status3, 0);
// nnz > number of non-zeros
int status4 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 4);
EXPECT_EQ(status4, 0);
// nnz < number of non-zeros
int status2 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 2);
EXPECT_EQ(status2, 1);
// nnz == 0
int status0 =
mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind, 0);
EXPECT_EQ(status0, 1);
}
} // namespace
} // namespace mujoco