mju_compressSparse(): add option to compress away small elements.

PiperOrigin-RevId: 726010576
Change-Id: Ib6458037a5cf9a9d82abb364b7dda7aab2f36807

Author: yuvaltassa

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];
 }
 
 

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);

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,

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"(<mujoco/>)";
 
 TEST_F(EngineUtilSparseTest, MjuSqrMatTDSparse1) {