Add mju_clip, mju_fill, mju_eye and mju_symmetrize to python bindings.

- Also change signatures of `mju_fill` and `mju_symmetrize` to conform to standard.

PiperOrigin-RevId: 485613824
Change-Id: I8dd618bac92ef40aa6c76380fa73302e81968687

doc/APIreference.rst

@@ -5917,7 +5917,7 @@ mju_fill
 
 .. code-block:: C
 
-   void mju_fill(mjtNum* res, int n, mjtNum val);
+   void mju_fill(mjtNum* res, mjtNum val, int n);
 
 Set res = val.
 
@@ -6101,9 +6101,9 @@ mju_symmetrize
 
 .. code-block:: C
 
-   void mju_symmetrize(mjtNum* mat, int n);
+   void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n);
 
-Symmetrize square matrix :math:`M = \frac{1}{2}(M + M^T)`.
+Symmetrize square matrix :math:`R = \frac{1}{2}(M + M^T)`.
 
 .. _mju_eye:
 

include/mujoco/mujoco.h

@@ -872,7 +872,7 @@ MJAPI mjtNum mju_normalize4(mjtNum res[4]);
 MJAPI void mju_zero(mjtNum* res, int n);
 
 // Set res = val.
-MJAPI void mju_fill(mjtNum* res, int n, mjtNum val);
+MJAPI void mju_fill(mjtNum* res, mjtNum val, int n);
 
 // Set res = vec.
 MJAPI void mju_copy(mjtNum* res, const mjtNum* data, int n);
@@ -925,8 +925,8 @@ MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNu
 // Transpose matrix: res = mat'.
 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);
+// Symmetrize square matrix res = (mat + mat')/2.
+MJAPI void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n);
 
 // Set mat to the identity matrix.
 MJAPI void mju_eye(mjtNum* mat, int n);

introspect/functions.py

@@ -5295,14 +5295,14 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
              ),
-             FunctionParameterDecl(
-                 name='n',
-                 type=ValueType(name='int'),
-             ),
              FunctionParameterDecl(
                  name='val',
                  type=ValueType(name='mjtNum'),
              ),
+             FunctionParameterDecl(
+                 name='n',
+                 type=ValueType(name='int'),
+             ),
          ),
          doc='Set res = val.',
      )),
@@ -5756,17 +5756,23 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          return_type=ValueType(name='void'),
          parameters=(
              FunctionParameterDecl(
-                 name='mat',
+                 name='res',
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
              ),
+             FunctionParameterDecl(
+                 name='mat',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
              FunctionParameterDecl(
                  name='n',
                  type=ValueType(name='int'),
              ),
          ),
-         doc="Symmetrize square matrix M = (M + M')/2.",
+         doc="Symmetrize square matrix res = (mat + mat')/2.",
      )),
     ('mju_eye',
      FunctionDecl(

python/mujoco/bindings_test.py

@@ -1049,6 +1049,27 @@ Euler integrator, semi-implicit in velocity.
     rank = mujoco.mju_boxQP(res, r, index, h, g, lower, upper)
     self.assertGreater(rank, -1)
 
+  def test_mju_fill(self):
+    res = np.empty(3, np.float64)
+    mujoco.mju_fill(res, 1.5)
+    np.testing.assert_array_equal(res, np.full(3, 1.5))
+
+  def test_mju_eye(self):
+    eye4 = np.empty((4, 4), np.float64)
+    mujoco.mju_eye(eye4)
+    np.testing.assert_array_equal(eye4, np.eye(4))
+
+  def test_mju_symmetrize(self):
+    mat = np.linspace(0, 1, 16).reshape(4, 4)
+    res = np.empty((4, 4), np.float64)
+    mujoco.mju_symmetrize(res, mat)
+    np.testing.assert_array_equal(res, 0.5*(mat + mat.T))
+
+  def test_mju_clip(self):
+    self.assertEqual(mujoco.mju_clip(1.5, 1.0, 2.0), 1.5)
+    self.assertEqual(mujoco.mju_clip(1.5, 2.0, 3.0), 2.0)
+    self.assertEqual(mujoco.mju_clip(1.5, 0.0, 1.0), 1.0)
+
   def test_mju_mul_vec_mat_vec(self):
     vec1 = np.array([1., 2., 3.])
     vec2 = np.array([3., 2., 1.])

python/mujoco/functions.cc

@@ -632,6 +632,11 @@ PYBIND11_MODULE(_functions, pymodule) {
       [](Eigen::Ref<EigenVectorX> res) {
         return InterceptMjErrors(::mju_zero)(res.data(), res.size());
       });
+  DEF_WITH_OMITTED_PY_ARGS(traits::mju_fill, "n")(
+      pymodule,
+      [](Eigen::Ref<EigenVectorX> res, mjtNum val) {
+        return InterceptMjErrors(::mju_fill)(res.data(), val, res.size());
+      });
   DEF_WITH_OMITTED_PY_ARGS(traits::mju_copy, "n")(
       pymodule,
       [](Eigen::Ref<EigenVectorX> res,
@@ -817,6 +822,27 @@ PYBIND11_MODULE(_functions, pymodule) {
         return InterceptMjErrors(::mju_transpose)(
             res.data(), mat.data(), mat.rows(), mat.cols());
       });
+  DEF_WITH_OMITTED_PY_ARGS(traits::mju_symmetrize, "n")(
+      pymodule,
+      [](Eigen::Ref<EigenArrayXX> res,
+         Eigen::Ref<const EigenArrayXX> mat) {
+        if (mat.cols() != mat.rows()) {
+          throw py::type_error("mat should be square");
+        }
+        if (res.cols() != mat.cols() || res.rows() != mat.rows()) {
+          throw py::type_error("res and mat should have the same shape");
+        }
+        return InterceptMjErrors(::mju_symmetrize)(
+            res.data(), mat.data(), mat.rows());
+      });
+  DEF_WITH_OMITTED_PY_ARGS(traits::mju_eye, "n")(
+      pymodule,
+      [](Eigen::Ref<EigenArrayXX> mat) {
+        if (mat.cols() != mat.rows()) {
+          throw py::type_error("mat should be square");
+        }
+        return InterceptMjErrors(::mju_eye)(mat.data(), mat.rows());
+      });
   DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulMatMat, "r1", "c1", "c2")(
       pymodule,
       [](Eigen::Ref<EigenArrayXX> res,
@@ -1030,6 +1056,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mju_springDamper>(pymodule);
   Def<traits::mju_min>(pymodule);
   Def<traits::mju_max>(pymodule);
+  Def<traits::mju_clip>(pymodule);
   Def<traits::mju_sign>(pymodule);
   Def<traits::mju_round>(pymodule);
   Def<traits::mju_type2Str>(pymodule);

src/engine/engine_util_blas.c

@@ -245,7 +245,7 @@ void mju_zero(mjtNum* res, int n) {
 
 
 // res = val
-void mju_fill(mjtNum* res, int n, mjtNum val) {
+void mju_fill(mjtNum* res, mjtNum val, int n) {
   for (int i=0; i<n; i++) {
     res[i] = val;
   }
@@ -740,12 +740,12 @@ 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) {
+// symmetrize square matrix res = (mat + mat')/2
+void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n) {
   for (int i=0; i<n; i++) {
+    res[i*(n+1)] = mat[i*(n+1)];
     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;
+      res[i*n+j] = res[j*n+i] = 0.5 * (mat[i*n+j] + mat[j*n+i]);
     }
   }
 }

src/engine/engine_util_blas.h

@@ -131,7 +131,7 @@ MJAPI mjtNum mju_normalize4(mjtNum vec[4]);
 MJAPI void mju_zero(mjtNum* res, int n);
 
 // res = val
-MJAPI void mju_fill(mjtNum* res, int n, mjtNum val);
+MJAPI void mju_fill(mjtNum* res, mjtNum val, int n);
 
 // res = vec
 MJAPI void mju_copy(mjtNum* res, const mjtNum* vec, int n);
@@ -192,8 +192,8 @@ MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNu
 // 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);
+// symmetrize square matrix res = (mat + mat')/2
+MJAPI void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n);
 
 // identity matrix
 MJAPI void mju_eye(mjtNum* mat, int n);

test/engine/engine_util_blas_test.cc

@@ -57,7 +57,7 @@ TEST_F(EngineUtilBlasTest, MjuMulVecMatVec) {
 
 TEST_F(EngineUtilBlasTest, MjuFill) {
   mjtNum vec[] = {0, 1, 4};
-  mju_fill(vec, 3, 4.5);
+  mju_fill(vec, 4.5, 3);
 
   EXPECT_EQ(vec[0], 4.5);
   EXPECT_EQ(vec[1], 4.5);
@@ -82,7 +82,14 @@ TEST_F(EngineUtilBlasTest, MjuSymmetrize) {
     1.5, 2,   4,
     2.5, 3,   3
   };
-  mju_symmetrize(mat, 3);
+  mjtNum res[9] = {0};
+  mju_symmetrize(res, mat, 3);
+  EXPECT_THAT(res, ElementsAre(1, 2,   3,
+                               2, 2,   3.5,
+                               3, 3.5, 3));
+
+  // test for the case res==mat
+  mju_symmetrize(mat, mat, 3);
   EXPECT_THAT(mat, ElementsAre(1, 2,   3,
                                2, 2,   3.5,
                                3, 3.5, 3));

unity/Runtime/Bindings/MujocoBindings.cs

@@ -3341,7 +3341,7 @@ public static unsafe extern double mju_normalize4(double* res);
 public static unsafe extern void mju_zero(double* res, int n);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mju_fill(double* res, int n, double val);
+public static unsafe extern void mju_fill(double* res, double val, int n);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mju_copy(double* res, double* data, int n);
@@ -3395,7 +3395,7 @@ public static unsafe extern double mju_mulVecMatVec(double* vec1, double* mat, d
 public static unsafe extern void mju_transpose(double* res, double* mat, int nr, int nc);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mju_symmetrize(double* mat, int n);
+public static unsafe extern void mju_symmetrize(double* res, double* mat, int n);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mju_eye(double* mat, int n);