Add mju_fill, mju_symmetrize and mju_eye utility functions.

PiperOrigin-RevId: 484226916
Change-Id: I79593d54e903ddacab0f859bfc6d17eccad7d3a1
This commit is contained in:
Yuval Tassa
2022-10-27 05:27:46 -07:00
committed by Copybara-Service
parent 507b573763
commit f905c7fb7f
8 changed files with 202 additions and 6 deletions
+33
View File
@@ -5909,6 +5909,17 @@ mju_zero
Set res = 0.
.. _mju_fill:
mju_fill
~~~~~~~~
.. code-block:: C
void mju_fill(mjtNum* res, int n, mjtNum val);
Set res = val.
mju_copy
~~~~~~~~
@@ -6082,6 +6093,28 @@ mju_transpose
Transpose matrix: res = mat'.
.. _mju_symmetrize:
mju_symmetrize
~~~~~~~~~~~~~~
.. code-block:: C
void mju_symmetrize(mjtNum* mat, int n);
Symmetrize square matrix :math:`M = \frac{1}{2}(M + M^T)`.
.. _mju_eye:
mju_eye
~~~~~~~
.. code-block:: C
void mju_eye(mjtNum* mat, int n);
Set mat to the identity matrix.
.. _mju_mulMatMat:
mju_mulMatMat
+8
View File
@@ -3,6 +3,14 @@ Changelog
=========
Upcoming version (not yet released)
-----------------------------------
General
^^^^^^^
- Added :ref:`mju_fill`, :ref:`mju_symmetrize` and :ref:`mju_eye` utility functions.
Version 2.3.0 (October 18, 2022)
--------------------------------
+10 -1
View File
@@ -871,6 +871,9 @@ MJAPI mjtNum mju_normalize4(mjtNum res[4]);
// Set res = 0.
MJAPI void mju_zero(mjtNum* res, int n);
// Set res = val.
MJAPI void mju_fill(mjtNum* res, int n, mjtNum val);
// Set res = vec.
MJAPI void mju_copy(mjtNum* res, const mjtNum* data, int n);
@@ -916,12 +919,18 @@ 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.
// 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);
// Symmetrize square matrix M = (M + M')/2.
MJAPI void mju_symmetrize(mjtNum* mat, int n);
// Set mat to the identity matrix.
MJAPI void mju_eye(mjtNum* mat, int n);
// Multiply matrices: res = mat1 * mat2.
MJAPI void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
int r1, int c1, int c2);
+59 -1
View File
@@ -5284,6 +5284,28 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Set res = 0.',
)),
('mju_fill',
FunctionDecl(
name='mju_fill',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='n',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='val',
type=ValueType(name='mjtNum'),
),
),
doc='Set res = val.',
)),
('mju_copy',
FunctionDecl(
name='mju_copy',
@@ -5698,7 +5720,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
type=ValueType(name='int'),
),
),
doc="Multiply square matrix with vectors on both sides: returns vec1'*mat*vec2.", # pylint: disable=line-too-long
doc="Multiply square matrix with vectors on both sides: returns vec1' * mat * vec2.", # pylint: disable=line-too-long
)),
('mju_transpose',
FunctionDecl(
@@ -5728,6 +5750,42 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc="Transpose matrix: res = mat'.",
)),
('mju_symmetrize',
FunctionDecl(
name='mju_symmetrize',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='n',
type=ValueType(name='int'),
),
),
doc="Symmetrize square matrix M = (M + M')/2.",
)),
('mju_eye',
FunctionDecl(
name='mju_eye',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='n',
type=ValueType(name='int'),
),
),
doc='Set mat to the identity matrix.',
)),
('mju_mulMatMat',
FunctionDecl(
name='mju_mulMatMat',
+35 -2
View File
@@ -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) {
+13 -2
View File
@@ -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);
+35
View File
@@ -24,6 +24,8 @@
namespace mujoco {
namespace {
using ::testing::ElementsAre;
using EngineUtilBlasTest = MujocoTest;
TEST_F(EngineUtilBlasTest, MjuDot) {
@@ -53,5 +55,38 @@ TEST_F(EngineUtilBlasTest, MjuMulVecMatVec) {
EXPECT_EQ(mju_mulVecMatVec(vec1, mat, vec2, 3), 204);
}
TEST_F(EngineUtilBlasTest, MjuFill) {
mjtNum vec[] = {0, 1, 4};
mju_fill(vec, 3, 4.5);
EXPECT_EQ(vec[0], 4.5);
EXPECT_EQ(vec[1], 4.5);
EXPECT_EQ(vec[2], 4.5);
}
TEST_F(EngineUtilBlasTest, MjuEye) {
mjtNum mat1[1];
mju_eye(mat1, 1);
EXPECT_EQ(mat1[0], 1);
mjtNum mat3[9];
mju_eye(mat3, 3);
EXPECT_THAT(mat3, ElementsAre(1, 0, 0,
0, 1, 0,
0, 0, 1));
}
TEST_F(EngineUtilBlasTest, MjuSymmetrize) {
mjtNum mat[] = {
1, 2.5, 3.5,
1.5, 2, 4,
2.5, 3, 3
};
mju_symmetrize(mat, 3);
EXPECT_THAT(mat, ElementsAre(1, 2, 3,
2, 2, 3.5,
3, 3.5, 3));
}
} // namespace
} // namespace mujoco
+9
View File
@@ -3337,6 +3337,9 @@ public static unsafe extern double mju_normalize4(double* res);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
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);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_copy(double* res, double* data, int n);
@@ -3388,6 +3391,12 @@ public static unsafe extern double mju_mulVecMatVec(double* vec1, double* mat, d
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
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);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_eye(double* mat, int n);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_mulMatMat(double* res, double* mat1, double* mat2, int r1, int c1, int c2);