Add mju_addToSymSparse

- Internal engine function to add a symmetric sparse matrix to a dense matrix.
- Also minor refactors to related functions.

PiperOrigin-RevId: 757954336
Change-Id: I8d8a48bcbbd5d6c618ae097a87df2bc1e6f5ef1d

src/engine/engine_support.c

@@ -1089,20 +1089,18 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
 
 // add inertia matrix to sparse destination matrix
 void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
-                   int* rownnz, int* rowadr, int* colind, mjtNum* M,
+                   int* rownnz, int* rowadr, int* colind, const mjtNum* M,
-                   int* M_rownnz, int* M_rowadr, int* M_colind) {
+                   const int* M_rownnz, const int* M_rowadr, const int* M_colind) {
   int nv = m->nv;
 
   mj_markStack(d);
+  mjtNum* buf_val = mjSTACKALLOC(d, nv, mjtNum);
   int* buf_ind = mjSTACKALLOC(d, nv, int);
-  mjtNum* sparse_buf = mjSTACKALLOC(d, nv, mjtNum);
 
-  // add to destination
+  mju_addToMatSparse(dst, rownnz, rowadr, colind, nv,
-  for (int i=0; i < nv; i++) {
+                     M, M_rownnz, M_rowadr, M_colind,
-    rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
+                     buf_val, buf_ind);
-                                  rownnz[i], M_rownnz[i], colind + rowadr[i],
+
-                                  M_colind + M_rowadr[i], sparse_buf, buf_ind);
-  }
   mj_freeStack(d);
 }
 

src/engine/engine_support.h

@@ -137,8 +137,8 @@ MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
 
 // add inertia matrix to sparse destination matrix
 MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
-                         int* rownnz, int* rowadr, int* colind, mjtNum* M,
+                         int* rownnz, int* rowadr, int* colind, const mjtNum* M,
-                         int* M_rownnz, int* M_rowadr, int* M_colind);
+                         const int* M_rownnz, const int* M_rowadr, const int* M_colind);
 
 // add inertia matrix to dense destination matrix
 MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);

src/engine/engine_util_blas.c

@@ -843,9 +843,9 @@ void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
 }
 
 
-
+// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional
-// compute M'*diag*M (diag=NULL: compute M'*M)
+void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag,
-void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
+                       int nr, int nc, int flg_upper) {
   mjtNum tmp;
 
   // half of MatMat routine: only lower triangle
@@ -870,15 +870,23 @@ void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, in
     }
   }
 
-  // make symmetric
+  // flg_upper is set: make symmetric
-  for (int i=0; i < nc; i++) {
+  if (flg_upper) {
-    for (int j=i+1; j < nc; j++) {
+    for (int i=0; i < nc; i++) {
-      res[i*nc+j] = res[j*nc+i];
+      for (int j=i+1; j < nc; j++) {
+        res[i*nc+j] = res[j*nc+i];
+      }
     }
   }
 }
 
 
+// compute M'*diag*M (diag=NULL: compute M'*M)
+void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc) {
+  mju_sqrMatTD_impl(res, mat, diag, nr, nc, /*flg_upper=*/ 1);
+}
+
+
 
 // multiply matrices, first argument transposed
 void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,

src/engine/engine_util_blas.h

@@ -223,6 +223,10 @@ MJAPI void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
 MJAPI void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                           int r1, int c1, int c2);
 
+// compute M'*diag*M (diag=NULL: compute M'*M), upper triangle optional
+void mju_sqrMatTD_impl(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc,
+                       int flg_upper);
+
 // compute M'*diag*M (diag=NULL: compute M'*M)
 MJAPI void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc);
 

src/engine/engine_util_sparse.c

@@ -186,6 +186,41 @@ void mju_mulMatTVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 }
 
 
+// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
+void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
+                        const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
+                        const int* M_colind,
+                        mjtNum* buf_val, int* buf_ind) {
+  for (int i=0; i < nr; i++) {
+    rownnz[i] = mju_combineSparse(dst + rowadr[i], M + M_rowadr[i], 1, 1,
+                                  rownnz[i], M_rownnz[i], colind + rowadr[i],
+                                  M_colind + M_rowadr[i], buf_val, buf_ind);
+  }
+}
+
+
+// add symmetric matrix (lower triangle) to dense matrix, upper triangle optional
+void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n,
+                        const int* rownnz, const int* rowadr, const int* colind, int flg_upper) {
+  for (int i=0; i < n; i++) {
+    int start = rowadr[i];
+    int end = start + rownnz[i];
+    for (int adr=start; adr < end; adr++) {
+      mjtNum val = mat[adr];
+      int j = colind[adr];
+
+      // lower + diagonal
+      res[i*n + j] += val;
+
+      // strict upper
+      if (flg_upper && j < i) {
+        res[j*n + i] += val;
+      }
+    }
+  }
+}
+
+
 
 // multiply symmetric matrix (only lower triangle represented) by vector:
 //  res = (mat + strict_upper(mat')) * vec
@@ -214,7 +249,7 @@ void mju_mulSymVecSparse(mjtNum* restrict res, const mjtNum* restrict mat,
 
     // off-diagonals
     const int* ind = colind + adr;
-    for (int k=0; k < diag; k++) {
+    for (int k=diag-1; k >= 0; k--) {
       int j = ind[k];
       mjtNum val = row[k];
       res[i] += val * vec[j]; // strict lower

src/engine/engine_util_sparse.h

@@ -50,6 +50,17 @@ 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);
 
+// add sparse matrix M to sparse destination matrix, requires pre-allocated buffers
+MJAPI void mju_addToMatSparse(mjtNum* dst, int* rownnz, int* rowadr, int* colind, int nr,
+                              const mjtNum* M, const int* M_rownnz, const int* M_rowadr,
+                              const int* M_colind,
+                              mjtNum* buf_val, int* buf_ind);
+
+// add symmetric matrix (only lower triangle represented) to dense matrix
+MJAPI void mju_addToSymSparse(mjtNum* res, const mjtNum* mat, int n,
+                              const int* rownnz, const int* rowadr, const int* colind,
+                              int flg_upper);
+
 // 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,

test/engine/engine_util_sparse_test.cc

@@ -1036,6 +1036,40 @@ TEST_F(EngineUtilSparseTest, MjuMulMatTVec) {
   EXPECT_THAT(AsVector(res, 3), ElementsAre(5, 28, 24));
 }
 
+TEST_F(EngineUtilSparseTest, MjuAddToSymSparse) {
+  //     1 2 4
+  // M = 2 3 0
+  //     4 0 5
+
+  // only lower triangle represented
+  mjtNum mat[] = {1, 2, 3, 4, 5};
+  int colind[] = {0, 0, 1, 0, 2};
+  int rownnz[] = {1, 2, 2};
+  int rowadr[] = {0, 1, 3};
+
+  //     0 0 0
+  // A = 5 4 2
+  //     4 3 2
+  mjtNum A[] = {0, 0, 0,
+                5, 4, 2,
+                4, 3, 2};
+
+  mju_addToSymSparse(A, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/1);
+  EXPECT_THAT(AsVector(A, 9), ElementsAre(1, 2, 4,
+                                          7, 7, 2,
+                                          8, 3, 7));
+
+  // same as A
+  mjtNum B[] = {0, 0, 0,
+                5, 4, 2,
+                4, 3, 2};
+
+  mju_addToSymSparse(B, mat, 3, rownnz, rowadr, colind, /*flg_upper=*/0);
+  EXPECT_THAT(AsVector(B, 9), ElementsAre(1, 0, 0,
+                                          7, 7, 2,
+                                          8, 3, 7));
+}
+
 TEST_F(EngineUtilSparseTest, MjuMulSymVecSparse) {
   constexpr int n = 4;
   constexpr int nnz = 9;