Skip to content

Commit

Permalink
Expose mju_sparse2dense to public API.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 688985153
Change-Id: I9559f1d47363d0afada80c86abdc729cd1910b63
  • Loading branch information
yuvaltassa authored and copybara-github committed Oct 23, 2024
1 parent b598d79 commit e14aebf
Show file tree
Hide file tree
Showing 9 changed files with 103 additions and 2 deletions.
1 change: 1 addition & 0 deletions doc/APIreference/APIfunctions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ API function can be classified as:
- **Math**
- Aliases for C :ref:`standard math<Standardmath>` functions.
- :ref:`Vector math<Vectormath>`.
- :ref:`Sparse math<Sparsemath>`.
- :ref:`Quaternions<Quaternions>`.
- :ref:`Pose transformations<Poses>`.
- :ref:`Matrix decompositions and solvers<Decompositions>`.
Expand Down
13 changes: 13 additions & 0 deletions doc/APIreference/functions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -3440,6 +3440,19 @@ mju_transformSpatial
Coordinate transform of 6D motion or force vector in rotation:translation format.
rotnew2old is 3-by-3, NULL means no rotation; flg_force specifies force or motion type.

.. _Sparsemath:

Sparse math
^^^^^^^^^^^
.. _mju_sparse2dense:

mju_sparse2dense
~~~~~~~~~~~~~~~~

.. mujoco-include:: mju_sparse2dense

Convert matrix from sparse to dense.

.. _Quaternions:

Quaternions
Expand Down
2 changes: 2 additions & 0 deletions doc/includes/references.h
Original file line number Diff line number Diff line change
Expand Up @@ -3467,6 +3467,8 @@ void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, in
void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_force,
const mjtNum newpos[3], const mjtNum oldpos[3],
const mjtNum rotnew2old[9]);
void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);
void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]);
void mju_negQuat(mjtNum res[4], const mjtNum quat[4]);
void mju_mulQuat(mjtNum res[4], const mjtNum quat1[4], const mjtNum quat2[4]);
Expand Down
7 changes: 7 additions & 0 deletions include/mujoco/mujoco.h
Original file line number Diff line number Diff line change
Expand Up @@ -1081,6 +1081,13 @@ MJAPI void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_forc
const mjtNum rotnew2old[9]);


//---------------------------------- Sparse math ---------------------------------------------------

// Convert matrix from sparse to dense.
MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);


//---------------------------------- Quaternions ---------------------------------------------------

// Rotate vector by quaternion.
Expand Down
46 changes: 46 additions & 0 deletions introspect/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -7062,6 +7062,52 @@
),
doc='Coordinate transform of 6D motion or force vector in rotation:translation format. rotnew2old is 3-by-3, NULL means no rotation; flg_force specifies force or motion type.', # pylint: disable=line-too-long
)),
('mju_sparse2dense',
FunctionDecl(
name='mju_sparse2dense',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='res',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='mat',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='nr',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='nc',
type=ValueType(name='int'),
),
FunctionParameterDecl(
name='rownnz',
type=PointerType(
inner_type=ValueType(name='int', is_const=True),
),
),
FunctionParameterDecl(
name='rowadr',
type=PointerType(
inner_type=ValueType(name='int', is_const=True),
),
),
FunctionParameterDecl(
name='colind',
type=PointerType(
inner_type=ValueType(name='int', is_const=True),
),
),
),
doc='Convert matrix from sparse to dense.',
)),
('mju_rotVecQuat',
FunctionDecl(
name='mju_rotVecQuat',
Expand Down
10 changes: 10 additions & 0 deletions python/mujoco/bindings_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -1337,6 +1337,16 @@ def test_mju_mul_vec_mat_vec(self):
mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)

def test_mju_sparse_to_dense(self):
expected = np.array([[0., 1., 0.], [2., 0., 3.]])
mat = np.array((1., 2., 3.))
rownnz = np.array([1, 2])
rowadr = np.array([0, 1])
colind = np.array([1, 0, 2])
res = np.zeros((2, 3))
mujoco.mju_sparse2dense(res, mat, rownnz, rowadr, colind)
np.testing.assert_array_equal(res, expected)

def test_mju_euler_to_quat(self):
quat = np.zeros(4)
euler = np.array([0, np.pi/2, 0])
Expand Down
19 changes: 19 additions & 0 deletions python/mujoco/functions.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1021,6 +1021,25 @@ PYBIND11_MODULE(_functions, pymodule) {
});
Def<traits::mju_transformSpatial>(pymodule);

// Sparse math
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sparse2dense, "nr", "nc")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorI> rownnz,
Eigen::Ref<const EigenVectorI> rowadr,
Eigen::Ref<const EigenVectorI> colind) {
if (res.rows() != rownnz.size()) {
throw py::type_error("#rows in res should equal size of rownnz");
}
if (res.rows() != rowadr.size()) {
throw py::type_error("#rows in res should equal size of rowadr");
}
return ::mju_sparse2dense(res.data(), mat.data(), res.rows(),
res.cols(), rownnz.data(), rowadr.data(),
colind.data());
});

// Quaternions
Def<traits::mju_rotVecQuat>(pymodule);
Def<traits::mju_negQuat>(pymodule);
Expand Down
4 changes: 2 additions & 2 deletions src/engine/engine_util_sparse.h
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ MJAPI void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
int* rownnz, int* rowadr, int* colind);

// convert matrix from sparse to dense
void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);
MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
const int* rownnz, const int* rowadr, const int* colind);

// multiply sparse matrix and dense vector: res = mat * vec
MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
Expand Down
3 changes: 3 additions & 0 deletions unity/Runtime/Bindings/MjBindings.cs
Original file line number Diff line number Diff line change
Expand Up @@ -7220,6 +7220,9 @@ public unsafe struct mjvSceneState_ {
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_transformSpatial(double* res, double* vec, int flg_force, double* newpos, double* oldpos, double* rotnew2old);

[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_sparse2dense(double* res, double* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind);

[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_rotVecQuat(double* res, double* vec, double* quat);

Expand Down

0 comments on commit e14aebf

Please sign in to comment.