diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 878bbf00..4ef49615 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -3414,6 +3414,16 @@ rotnew2old is 3-by-3, NULL means no rotation; flg_force specifies force or motio Sparse math ^^^^^^^^^^^ +.. _mju_dense2sparse: + +mju_dense2sparse +~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_dense2sparse + +Convert matrix from dense to sparse. + nnz is size of res and colind, return 1 if too small, 0 otherwise. + .. _mju_sparse2dense: mju_sparse2dense @@ -4372,7 +4382,7 @@ mjs_setFrame .. mujoco-include:: mjs_setFrame -Set element's enlcosing frame. +Set element's enclosing frame. .. _mjs_resolveOrientation: diff --git a/doc/includes/references.h b/doc/includes/references.h index 9a777763..6c0c1d41 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3464,6 +3464,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]); +int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, + int* rownnz, int* rowadr, int* colind, int nnz); 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]); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index fae4600f..54b807e9 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1069,6 +1069,11 @@ MJAPI void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_forc //---------------------------------- Sparse math --------------------------------------------------- +// Convert matrix from dense to sparse. +// nnz is size of res and colind, return 1 if too small, 0 otherwise. +MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, + int* rownnz, int* rowadr, int* colind, int nnz); + // 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); @@ -1611,7 +1616,7 @@ MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size); // Set element's default. MJAPI void mjs_setDefault(mjsElement* element, mjsDefault* def); -// Set element's enlcosing frame. +// Set element's enclosing frame. MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame); // Resolve alternative orientations to quat, return error if any. diff --git a/introspect/functions.py b/introspect/functions.py index 2afb3798..6c3a5b9e 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -6958,6 +6958,56 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), 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_dense2sparse', + FunctionDecl( + name='mju_dense2sparse', + return_type=ValueType(name='int'), + 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'), + ), + ), + FunctionParameterDecl( + name='rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + ), + FunctionParameterDecl( + name='colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + ), + FunctionParameterDecl( + name='nnz', + type=ValueType(name='int'), + ), + ), + doc='Convert matrix from dense to sparse. nnz is size of res and colind, return 1 if too small, 0 otherwise.', # pylint: disable=line-too-long + )), ('mju_sparse2dense', FunctionDecl( name='mju_sparse2dense', @@ -10236,7 +10286,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc="Set element's enlcosing frame.", + doc="Set element's enclosing frame.", )), ('mjs_resolveOrientation', FunctionDecl( diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 5448a609..ffdfb75b 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -1337,6 +1337,23 @@ Euler integrator, semi-implicit in velocity. mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]) self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.) + def test_mju_dense_to_sparse(self): + mat = np.array([[0., 1., 0.], [2., 0., 3.]]) + expected_vals = np.array([1., 2., 3.]) + expected_rownnz = np.array([1, 2]) + expected_rowadr = np.array([0, 1]) + expected_colind = np.array([1, 0, 2]) + vals = np.zeros(3) + row_nnz = np.zeros(2, np.int32) + row_adr = np.zeros(2, np.int32) + col_ind = np.zeros(3, np.int32) + status = mujoco.mju_dense2sparse(vals, mat, row_nnz, row_adr, col_ind) + np.testing.assert_equal(status, 0) + np.testing.assert_array_equal(vals, expected_vals) + np.testing.assert_array_equal(row_nnz, expected_rownnz) + np.testing.assert_array_equal(row_adr, expected_rowadr) + np.testing.assert_array_equal(col_ind, expected_colind) + def test_mju_sparse_to_dense(self): expected = np.array([[0., 1., 0.], [2., 0., 3.]]) mat = np.array((1., 2., 3.)) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 47bf5314..5bb00ba9 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -1019,6 +1019,25 @@ PYBIND11_MODULE(_functions, pymodule) { Def(pymodule); // Sparse math + DEF_WITH_OMITTED_PY_ARGS(traits::mju_dense2sparse, "nr", "nc", "nnz")( + pymodule, + [](Eigen::Ref res, Eigen::Ref mat, + Eigen::Ref rownnz, Eigen::Ref rowadr, + Eigen::Ref colind) { + if (mat.rows() != rownnz.size()) { + throw py::type_error("#rows in mat should equal size of rownnz"); + } + if (mat.rows() != rowadr.size()) { + throw py::type_error("#rows in mat should equal size of rowadr"); + } + if (res.size() != colind.size()) { + throw py::type_error("#size of res should equal size of colind"); + } + return ::mju_dense2sparse(res.data(), mat.data(), mat.rows(), + mat.cols(), rownnz.data(), rowadr.data(), + colind.data(), res.size()); + }); + DEF_WITH_OMITTED_PY_ARGS(traits::mju_sparse2dense, "nr", "nc")( pymodule, [](Eigen::Ref res, diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 51a13917..e963ab1c 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -7230,6 +7230,9 @@ public static unsafe extern void mju_sqrMatTD(double* res, double* mat, double* [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 int mju_dense2sparse(double* res, double* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind, int nnz); + [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);