Add mju_fill, mju_symmetrize and mju_eye utility functions.

PiperOrigin-RevId: 484226916
Change-Id: I79593d54e903ddacab0f859bfc6d17eccad7d3a1

src/engine/engine_util_blas.c

@@ -16,7 +16,7 @@
 
 #include <string.h>
 
-#include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
 
 #ifdef mjUSEPLATFORMSIMD
   #if defined(__AVX__) && defined(mjUSEDOUBLE)
@@ -244,6 +244,15 @@ void mju_zero(mjtNum* res, int n) {
 
 
 
+// res = val
+void mju_fill(mjtNum* res, int n, mjtNum val) {
+  for (int i=0; i<n; i++) {
+    res[i] = val;
+  }
+}
+
+
+
 // res = vec
 void mju_copy(mjtNum* res, const mjtNum* vec, int n) {
   if (n>0) {
@@ -718,7 +727,7 @@ mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec
 
 
 
-//------------------------------ matrix-matrix operations ------------------------------------------
+//------------------------------ matrix operations -------------------------------------------------
 
 // transpose matrix
 void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) {
@@ -731,6 +740,30 @@ void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) {
 
 
 
+// symmetrize square matrix M = (M + M')/2
+void mju_symmetrize(mjtNum* mat, int n) {
+  for (int i=0; i<n; i++) {
+    for (int j=0; j<i; j++) {
+      mjtNum tmp = 0.5 * (mat[i*n+j] + mat[j*n+i]);
+      mat[i*n+j] = mat[j*n+i] = tmp;
+    }
+  }
+}
+
+
+
+// identity matrix
+void mju_eye(mjtNum* mat, int n) {
+  mju_zero(mat, n*n);
+  for (int i=0; i<n; i++) {
+    mat[i*(n + 1)] = 1;
+  }
+}
+
+
+
+//------------------------------ matrix-matrix operations ------------------------------------------
+
 // multiply matrices, exploit sparsity of mat1
 void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                    int r1, int c1, int c2) {

src/engine/engine_util_blas.h

@@ -18,7 +18,7 @@
 #include <math.h>
 
 #include <mujoco/mjexport.h>
-#include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -130,6 +130,9 @@ MJAPI mjtNum mju_normalize4(mjtNum vec[4]);
 // res = 0
 MJAPI void mju_zero(mjtNum* res, int n);
 
+// res = val
+MJAPI void mju_fill(mjtNum* res, int n, mjtNum val);
+
 // res = vec
 MJAPI void mju_copy(mjtNum* res, const mjtNum* vec, int n);
 
@@ -184,11 +187,19 @@ MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
 
 
-//------------------------------ matrix-matrix operations ------------------------------------------
+//------------------------------ matrix operations -------------------------------------------------
 
 // transpose matrix
 MJAPI void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc);
 
+// symmetrize square matrix M = (M + M')/2
+MJAPI void mju_symmetrize(mjtNum* mat, int n);
+
+// identity matrix
+MJAPI void mju_eye(mjtNum* mat, int n);
+
+//------------------------------ matrix-matrix operations ------------------------------------------
+
 // multiply matrices
 MJAPI void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
                          int r1, int c1, int c2);