From 7742f6280323bd22f1cbaf7299dd2494882ef14b Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 29 Apr 2025 16:32:49 -0700
Subject: [PATCH] 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        | 38 ++++++++++++++++++++++++++
 src/engine/engine_util_sparse.h        |  6 ++++
 test/engine/engine_util_sparse_test.cc | 36 ++++++++++++++++++++++++
 3 files changed, 80 insertions(+)

diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c
index 23c038ec..00fb99d9 100644
--- a/src/engine/engine_util_sparse.c
+++ b/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;
diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h
index 1f75739b..1246abc1 100644
--- a/src/engine/engine_util_sparse.h
+++ b/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);
diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc
index 41231c2c..a11d18e3 100644
--- a/test/engine/engine_util_sparse_test.cc
+++ b/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;