Add mju_mulSymVecSparse, private engine function

Multiply sparse symmetric matrix (only lower triangle represented) by vector

PiperOrigin-RevId: 752926934
Change-Id: I5aaae1266256aa88aee8c25e242b7a7c51ea8dd7

src/engine/engine_util_sparse.c

@@ -193,6 +193,44 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 
 
 
+// multiply symmetric matrix (only lower triangle represented) by vector:
+//  res = (mat + strict_upper(mat')) * vec
+void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat,
+                         const mjtNum* restrict vec, int n,
+                         const int* restrict rownnz, const int* restrict rowadr,
+                         const int* restrict diagnum, const int* restrict colind) {
+  // clear res
+  mju_zero(res, n);
+
+  // multiply
+  for (int i=0; i < n; i++) {
+    int adr = rowadr[i];
+    int diag = rownnz[i] - 1;
+    const mjtNum* row = mat + adr;
+
+    // diagonal
+    res[i] = row[diag] * vec[i];
+
+    // TODO: consider using SIMD if diagnum[i] >= 4
+
+    // shortcut for diagonal row/column
+    if (diagnum[i]) {
+      continue;
+    }
+
+    // off-diagonals
+    const int* ind = colind + adr;
+    for (int k=0; k < diag; k++) {
+      int j = ind[k];
+      mjtNum val = row[k];
+      res[i] += val * vec[j]; // strict lower
+      res[j] += val * vec[i]; // strict upper
+    }
+  }
+}
+
+
+
 // count the number of non-zeros in the sum of two sparse vectors
 int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) {
   int a = 0, b = 0, c_nnz = 0;

src/engine/engine_util_sparse.h

@@ -50,6 +50,12 @@ 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);
 
+// multiply symmetric matrix (only lower triangle represented) by vector:
+//  res = (mat + strict_upper(mat')) * vec
+MJAPI void mju_mulSymVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n,
+                               const int* rownnz, const int* rowadr, const int* diagnum,
+                               const 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);

test/engine/engine_util_sparse_test.cc

@@ -1034,6 +1034,42 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
   EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
 }
 
+TEST_F(EngineUtilSparseTest, MjuMulSymVecSparse) {
+  constexpr int n = 4;
+  constexpr int nnz = 9;
+
+  mjtNum mat[n*n] = {1,  0,  0,  0,
+                    -1,  2,  0,  0,  // spurious (ignored) -1 at (1, 0)
+                     3,  0,  4,  0,
+                     5,  6,  7,  8};
+
+  // dense, full matrix
+  mjtNum sym[n*n] = {1,  0,  3,  5,
+                     0,  2,  0,  6,
+                     3,  0,  4,  7,
+                     5,  6,  7,  8};
+
+  mjtNum mat_sparse[nnz];
+  int rownnz[n];
+  int rowadr[n];
+  int colind[nnz];
+  mju_dense2sparse(mat_sparse, mat, n, n, rownnz, rowadr, colind, nnz);
+  int diagnum[n] = {0, 1, 0, 0};
+
+  // multiply: res = (mat + strict_upper(mat')) * vec
+  mjtNum vec[n] = {4, 3, 2, 1};
+  mjtNum res[n];
+  mju_mulSymVecSparse(res, mat_sparse, vec, n, rownnz, rowadr, diagnum, colind);
+
+  // dense multiply
+  mjtNum res2[n];
+  mju_mulMatVec(res2, sym, vec, n, n);
+
+  for (int i=0; i < n; i++) {
+    EXPECT_EQ(res[i], res2[i]);
+  }
+}
+
 TEST_F(EngineUtilSparseTest, MjuDenseToSparse) {
   int nr = 2;
   int nc = 2;