Add mju_mulMatTVecSparse to multiply transposed sparse matrix and d…

…ense vector. PiperOrigin-RevId: 688535240 Change-Id: I0bacf1e911cde99014b2dcb0a1434e61a4f41d51
google-deepmind · Oct 22, 2024 · 94e9912 · 94e9912
1 parent 2b0629d
commit 94e9912
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c
@@ -206,6 +206,28 @@ void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 
 
 
+// multiply transposed sparse matrix and dense vector:  res = mat' * vec.
+void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc,
+                          const int* rownnz, const int* rowadr, const int* colind) {
+  // clear res
+  mju_zero(res, nc);
+
+  for (int i=0; i < nr; i++) {
+    int nnz = rownnz[i];
+    int adr = rowadr[i];
+    const int* ind = colind + adr;
+    const mjtNum* row = mat + adr;
+    mjtNum scl = vec[i];
+
+    // add row scaled by the corresponding vector element
+    for (int j=0; j < nnz; j++) {
+      res[ind[j]] += row[j] * scl;
+    }
+  }
+}
+
+
+
 // res = res*scl1 + vec*scl2
 static void mju_addToSclScl(mjtNum* res, const mjtNum* vec, mjtNum scl1, mjtNum scl2, int n) {
 #ifdef mjUSEAVX

diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h
@@ -46,6 +46,10 @@ MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec
                                int nr, const int* rownnz, const int* rowadr,
                                const int* colind, const int* rowsuper);
 
+// multiply transposed sparse matrix and dense vector:  res = mat' * 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);

diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc
@@ -30,8 +30,9 @@ namespace {
 using ::testing::ElementsAre;
 using EngineUtilSparseTest = MujocoTest;
 
-std::vector<int> AsVector(const int* array, int n) {
-  return std::vector<int>(array, array + n);
+template <typename T>
+std::vector<T> AsVector(const T* array, int n) {
+  return std::vector<T>(array, array + n);
 }
 
 TEST_F(EngineUtilSparseTest, MjuDot) {
@@ -1043,5 +1044,25 @@ TEST_F(EngineUtilSparseTest, MjuCholFactorNNZ) {
   mj_deleteModel(model);
 }
 
+TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
+  int nr = 2;
+  int nc = 3;
+  mjtNum mat[] = {1, 2, 0,
+                  0, 3, 4};
+
+  mjtNum mat_sparse[6];
+  int rownnz[2];
+  int rowadr[2];
+  int colind[4];
+  mju_dense2sparse(mat_sparse, mat, nr, nc, rownnz, rowadr, colind);
+
+  // multiply: res = mat' * vec
+  mjtNum vec[] = {5, 6};
+  mjtNum res[3];
+  mju_mulMatTVecSparse(res, mat_sparse, vec, nr, nc, rownnz, rowadr, colind);
+
+  EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
+}
+
 }  // namespace
 }  // namespace mujoco