diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index 2dc9ea58..abb95be0 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -438,22 +438,34 @@ int mju_addChains(int* res, int n, int NV1, int NV2, -// compress layout of sparse matrix -void mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind) { - rowadr[0] = 0; - int adr = rownnz[0]; - for (int r=1; r < nr; r++) { +// compress sparse matrix, remove elements with abs(value) <= minval, return total non-zeros +int mju_compressSparse(mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind, + mjtNum minval) { + int remove_small = minval >= 0; + int adr = 0; + for (int r=0; r < nr; r++) { // save old rowadr, record new - int rowadr1 = rowadr[r]; + int rowadr_old = rowadr[r]; rowadr[r] = adr; - // shift mat and mat_colind - for (int adr1=rowadr1; adr1 < rowadr1+rownnz[r]; adr1++) { - mat[adr] = mat[adr1]; - colind[adr] = colind[adr1]; + // shift mat and colind + int nnz = 0; + int end = rowadr_old + rownnz[r]; + for (int adr_old=rowadr_old; adr_old < end; adr_old++) { + if (remove_small && mju_abs(mat[adr_old]) <= minval) { + continue; + } + if (adr != adr_old) { + mat[adr] = mat[adr_old]; + colind[adr] = colind[adr_old]; + } adr++; + if (remove_small) nnz++; } + if (remove_small) rownnz[r] = nnz; } + + return rowadr[nr-1] + rownnz[nr-1]; } diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 3a93c937..92c2dad0 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -50,9 +50,9 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec MJAPI void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); -// compress layout of sparse matrix -MJAPI void mju_compressSparse(mjtNum* mat, int nr, int nc, - int* rownnz, int* rowadr, int* colind); +// compress sparse matrix, remove elements with abs(value) <= minval, return total non-zeros +MJAPI int mju_compressSparse(mjtNum* mat, int nr, int nc, + int* rownnz, int* rowadr, int* colind, mjtNum minval); // count the number of non-zeros in the sum of two sparse vectors MJAPI int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind); diff --git a/test/benchmark/engine_util_sparse_benchmark_test.cc b/test/benchmark/engine_util_sparse_benchmark_test.cc index bc41edfa..2c2ff6aa 100644 --- a/test/benchmark/engine_util_sparse_benchmark_test.cc +++ b/test/benchmark/engine_util_sparse_benchmark_test.cc @@ -167,7 +167,8 @@ void ABSL_ATTRIBUTE_NOINLINE transposeSparse_baseline( } } - mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind); + mju_compressSparse(res, nc, nr, res_rownnz, res_rowadr, res_colind, + /*minval=*/-1); } int compare_baseline(const int* vec1, diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 505080fe..c1b7cb65 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -300,6 +300,41 @@ TEST_F(EngineUtilSparseTest, MjuTransposeNullMatrix) { EXPECT_THAT(rowadrT, ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)); } +TEST_F(EngineUtilSparseTest, MjuCompressSparse) { + // sparse matrix (uncompressed with spurious values between the rows): + // [[1, 0, 2] + // [0, O, 3] (second zero represented) + mjtNum mat[] = {1, 2, 9, 0, 3}; // spurious 9 value + int colind[] = {0, 2, -1, 1, 2}; // spurious -1 index + int rownnz[] = {2, 2}; + int rowadr[] = {0, 3}; + + mjtNum dense_expected[] = {1, 0, 2, 0, 0, 3}; + mjtNum dense[6]; + mju_sparse2dense(dense, mat, 2, 3, rownnz, rowadr, colind); + EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected, 6)); + + // check that spurious values are removed + int nnz = mju_compressSparse(mat, 2, 3, rownnz, rowadr, colind, + /*minval=*/-1); + EXPECT_EQ(nnz, 4); + mju_sparse2dense(dense, mat, 2, 3, rownnz, rowadr, colind); + EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected, 6)); + + // check that represented zero gets compressed aways with minval=0 + nnz = mju_compressSparse(mat, 2, 3, rownnz, rowadr, colind, /*minval=*/0); + EXPECT_EQ(nnz, 3); + mju_sparse2dense(dense, mat, 2, 3, rownnz, rowadr, colind); + EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected, 6)); + + // check that 1 gets compressed aways with minval=1 + nnz = mju_compressSparse(mat, 2, 3, rownnz, rowadr, colind, /*minval=*/1); + EXPECT_EQ(nnz, 2); + mju_sparse2dense(dense, mat, 2, 3, rownnz, rowadr, colind); + mjtNum dense_expected_minval1[] = {0, 0, 2, 0, 0, 3}; + EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected_minval1, 6)); +} + static constexpr char modelStr[] = R"()"; TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {