From f4e7fa97af84dcd83ce46acdc0bf67eb2d8424d8 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 14 Sep 2022 08:04:40 -0700 Subject: [PATCH] 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 | 11 ++++++++++ doc/changelog.rst | 2 ++ include/mujoco/mujoco.h | 3 +++ introspect/functions.py | 30 +++++++++++++++++++++++++++ python/mujoco/bindings_test.py | 6 ++++++ python/mujoco/functions.cc | 20 ++++++++++++++++++ src/engine/engine_solver.c | 5 ++--- src/engine/engine_util_blas.c | 17 +++++++++++---- src/engine/engine_util_blas.h | 3 +++ src/engine/engine_util_solve.c | 6 ++---- test/engine/engine_util_blas_test.cc | 12 +++++++++++ test/engine/engine_util_solve_test.cc | 14 +++++-------- 12 files changed, 109 insertions(+), 20 deletions(-) diff --git a/doc/APIreference.rst b/doc/APIreference.rst index bdddec4e..61864af4 100644 --- a/doc/APIreference.rst +++ b/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 ~~~~~~~~~~~~~ diff --git a/doc/changelog.rst b/doc/changelog.rst index 368b90eb..bc81bf37 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -25,6 +25,8 @@ General - The algorithm, introduced in `Tassa et al. 2014 `_, 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) --------------------------------- diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 922092a3..41e464c6 100644 --- a/include/mujoco/mujoco.h +++ b/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); diff --git a/introspect/functions.py b/introspect/functions.py index 6214ff7d..5ffb6f50 100644 --- a/introspect/functions.py +++ b/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', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 2c2f95e4..c0bf2f26 100644 --- a/python/mujoco/bindings_test.py +++ b/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 diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 14b78b35..e0445b12 100644 --- a/python/mujoco/functions.cc +++ b/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 vec1, + Eigen::Ref mat, + Eigen::Ref 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 res, diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 4afe922d..80a6c1fa 100644 --- a/src/engine/engine_solver.c +++ b/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_dot(delta, v, dim) + mju_dot(delta, res, dim); + change = 0.5*mju_mulVecMatVec(delta, A, delta, dim) + mju_dot(delta, res, dim); } // positive change: restore diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c index d686a925..769177ae 100644 --- a/src/engine/engine_util_blas.c +++ b/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, - int nr, int nc) { +void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc) { for (int r=0; r