Add mju_mulVecMatVec, mutiplies a square matrix M by a vector x on both sides. Returns x^T * M * x.

PiperOrigin-RevId: 474292806
Change-Id: I3432469dbe1f02ccf5a13241c7aa12d824cbe034

doc/APIreference.rst

@@ -6040,6 +6040,17 @@ mju_mulMatTVec
 
 Multiply transposed matrix and vector: res = mat' \* vec.
 
+.. _mju_mulVecMatVec:
+
+mju_mulVecMatVec
+~~~~~~~~~~~~~~~~
+
+.. code-block:: C
+
+   mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
+
+Multiply square matrix with vectors on both sides: return vec1' \* mat \* vec2.
+
 mju_transpose
 ~~~~~~~~~~~~~
 

doc/changelog.rst

@@ -25,6 +25,8 @@ General
 
    - The algorithm, introduced in `Tassa et al. 2014 <https://doi.org/10.1109/ICRA.2014.6907001>`_,
      converges after 2-5 Cholesky factorisations, independent of problem size.
+- Added :ref:`mju_mulVecMatVec` to multiply a square matrix :math:`M` with vectors :math:`x` and :math:`y` on both
+  sides. The function returns :math:`x^TMy`.
 
 Version 2.2.2 (September 7, 2022)
 ---------------------------------

include/mujoco/mujoco.h

@@ -911,6 +911,9 @@ MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
 // Multiply transposed matrix and vector: res = mat' * vec.
 MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc);
 
+// Multiply square matrix with vectors on both sides: returns vec1'*mat*vec2.
+MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
+
 // Transpose matrix: res = mat'.
 MJAPI void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc);
 

introspect/functions.py

@@ -5644,6 +5644,36 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc="Multiply transposed matrix and vector: res = mat' * vec.",
      )),
+    ('mju_mulVecMatVec',
+     FunctionDecl(
+         name='mju_mulVecMatVec',
+         return_type=ValueType(name='mjtNum'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='vec1',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='mat',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='vec2',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='n',
+                 type=ValueType(name='int'),
+             ),
+         ),
+         doc="Multiply square matrix with vectors on both sides: returns vec1'*mat*vec2.",  # pylint: disable=line-too-long
+     )),
     ('mju_transpose',
      FunctionDecl(
          name='mju_transpose',

python/mujoco/bindings_test.py

@@ -1006,6 +1006,12 @@ Euler integrator, semi-implicit in velocity.
     rank = mujoco.mju_boxQP(res, r, index, h, g, lower, upper)
     self.assertGreater(rank, -1)
 
+  def test_mju_mul_vec_mat_vec(self):
+    vec1 = np.array([1., 2., 3.])
+    vec2 = np.array([3., 2., 1.])
+    mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
+    self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)
+
   @parameterized.product(flg_html=(False, True), flg_pad=(False, True))
   def test_mj_printSchema(self, flg_html, flg_pad):  # pylint: disable=invalid-name
     # Make sure that mj_printSchema doesn't raise an exception

python/mujoco/functions.cc

@@ -784,6 +784,26 @@ PYBIND11_MODULE(_functions, pymodule) {
         return InterceptMjErrors(::mju_mulMatTVec)(
             res.data(), mat.data(), vec.data(), mat.rows(), mat.cols());
       });
+  DEF_WITH_OMITTED_PY_ARGS(traits::mju_mulVecMatVec, "n")(
+      pymodule,
+      [](Eigen::Ref<const EigenVectorX> vec1,
+         Eigen::Ref<const EigenArrayXX> mat,
+         Eigen::Ref<const EigenVectorX> vec2) {
+        if (vec1.size() != vec2.size()) {
+          throw py::type_error(
+              "size of vec1 should equal the size of vec2");
+        }
+        if (vec1.size() != mat.cols()) {
+          throw py::type_error(
+              "size of vectors should equal the number of columns in mat");
+        }
+        if (vec1.size() != mat.rows()) {
+          throw py::type_error(
+              "size of vectors should equal the number of rows in mat");
+        }
+        return InterceptMjErrors(::mju_mulVecMatVec)(
+            vec1.data(), mat.data(), vec2.data(), vec1.size());
+      });
   DEF_WITH_OMITTED_PY_ARGS(traits::mju_transpose, "nr", "nc")(
       pymodule,
       [](Eigen::Ref<EigenArrayXX> res,

src/engine/engine_solver.c

@@ -190,7 +190,7 @@ static void residual(const mjModel* m, mjData* d, mjtNum* res, int i, int dim, i
 // compute cost change
 static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
                          const mjtNum* res, int dim) {
-  mjtNum delta[6], v[6], change;
+  mjtNum delta[6], change;
 
   // compute change
   if (dim==1) {
@@ -198,8 +198,7 @@ static mjtNum costChange(const mjtNum* A, mjtNum* force, const mjtNum* oldforce,
     change = 0.5*delta[0]*delta[0]*A[0] + delta[0]*res[0];
   } else {
     mju_sub(delta, force, oldforce, dim);
-    mju_mulMatVec(v, A, delta, dim, dim);
+    change = 0.5*mju_mulVecMatVec(delta, A, delta, dim) + mju_dot(delta, res, dim);
-    change = 0.5*mju_dot(delta, v, dim) + mju_dot(delta, res, dim);
   }
 
   // positive change: restore

src/engine/engine_util_blas.c

@@ -685,8 +685,7 @@ mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, const int n) {
 //------------------------------ matrix-vector operations ------------------------------------------
 
 // multiply matrix and vector
-void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
-                   int nr, int nc) {
   for (int r=0; r<nr; r++) {
     res[r] = mju_dot(mat + r*nc, vec, nc);
   }
@@ -695,8 +694,7 @@ void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 
 
 // multiply transposed matrix and vector
-void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) {
-                    int nr, int nc) {
   mjtNum tmp;
   mju_zero(res, nc);
 
@@ -709,6 +707,17 @@ void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 
 
 
+// multiply square matrix with vectors on both sides: return vec1'*mat*vec2
+mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n) {
+  mjtNum res = 0;
+  for (int i=0; i<n; i++) {
+    res += vec1[i] * mju_dot(mat + i*n, vec2, n);
+  }
+  return res;
+}
+
+
+
 //------------------------------ matrix-matrix operations ------------------------------------------
 
 // transpose matrix

src/engine/engine_util_blas.h

@@ -180,6 +180,9 @@ MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
                           int nr, int nc);
 
+// multiply square matrix with vectors on both sides: return vec1'*mat*vec2
+MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n);
+
 
 //------------------------------ matrix-matrix operations ------------------------------------------
 

src/engine/engine_util_solve.c

@@ -934,8 +934,7 @@ int mju_boxQPoption(mjtNum* res, mjtNum* R, int* index,               // outputs
     }
 
     // compute objective: value = 0.5*res'*H*res + res'*g
-    mju_mulMatVec(temp, H, res, n, n);  // TODO(b/246267542): do this in one call
+    value = 0.5 * mju_mulVecMatVec(res, H, res, n) + mju_dot(res, g, n);
-    value = 0.5 * mju_dot(res, temp, n) + mju_dot(res, g, n);
 
     // save last value
     oldvalue = value;
@@ -1056,8 +1055,7 @@ int mju_boxQPoption(mjtNum* res, mjtNum* R, int* index,               // outputs
       }
 
       // new objective value
-      mju_mulMatVec(temp, H, candidate, n, n);
+      value = 0.5 * mju_mulVecMatVec(candidate, H, candidate, n) + mju_dot(candidate, g, n);
-      value = 0.5 * mju_dot(candidate, temp, n) + mju_dot(candidate, g, n);
 
       // increment and break if step is too small
       nstep++;

test/engine/engine_util_blas_test.cc

@@ -41,5 +41,17 @@ TEST_F(EngineUtilBlasTest, MjuDot) {
   EXPECT_EQ(mju_dot(a, b, 7), 7 + 2*6 + 3*5 + 4*4 + 5*3 + 6*2 + 7);
 }
 
+TEST_F(EngineUtilBlasTest, MjuMulVecMatVec) {
+  mjtNum vec1[] = {1, 2, 3};
+  mjtNum vec2[] = {3, 2, 1};
+  mjtNum mat[] = {
+    1, 2, 3,
+    4, 5, 6,
+    7, 8, 9
+  };
+
+  EXPECT_EQ(mju_mulVecMatVec(vec1, mat, vec2, 3), 204);
+}
+
 }  // namespace
 }  // namespace mujoco

test/engine/engine_util_solve_test.cc

@@ -75,10 +75,8 @@ TEST_F(QCQP3Test, DegenerateAMatrix) {
 using BoxQPTest = MujocoTest;
 
 // utility: compute QP objective = 0.5*x'*H*x + x'*g
-mjtNum objective(const mjtNum* x, const mjtNum* H, const mjtNum* g, int n,
+mjtNum objective(const mjtNum* x, const mjtNum* H, const mjtNum* g, int n) {
-                 mjtNum* temp) {
+  return 0.5 * mju_mulVecMatVec(x, H, x, n) + mju_dot(x, g, n);
-  mju_mulMatVec(temp, H, x, n, n);
-  return 0.5 * mju_dot(x, temp, n) + mju_dot(x, g, n);
 }
 
 // utility: test if res is the minimum of a given box-QP problem
@@ -86,11 +84,10 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
                  const mjtNum* lower, const mjtNum* upper) {
   static const mjtNum eps = 1e-4;  // epsilon used for nudging
   bool is_minimum = true;
-  mjtNum* temp = (mjtNum*) mju_malloc(sizeof(mjtNum)*n);
   mjtNum* res_nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*n);
 
   // get solution value
-  mjtNum value = objective(res, H, g, n, temp);
+  mjtNum value = objective(res, H, g, n);
   mjtNum value_nudge;
 
   // compare to nudged solution
@@ -102,7 +99,7 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
     if (lower) {
       res_nudge[i] = mju_max(lower[i], res_nudge[i]);
     }
-    value_nudge = objective(res_nudge, H, g, n, temp);
+    value_nudge = objective(res_nudge, H, g, n);
     if (value_nudge - value < 0) {
       is_minimum = false;
       break;
@@ -113,7 +110,7 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
     if (upper) {
       res_nudge[i] = mju_min(upper[i], res_nudge[i]);
     }
-    value_nudge = objective(res_nudge, H, g, n, temp);
+    value_nudge = objective(res_nudge, H, g, n);
     if (value_nudge - value < 0) {
       is_minimum = false;
       break;
@@ -124,7 +121,6 @@ bool isQPminimum(const mjtNum* res, const mjtNum* H, const mjtNum* g, int n,
   }
 
   mju_free(res_nudge);
-  mju_free(temp);
   return is_minimum;
 }