chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add mju_dense2sparse to public API.

Browse Source 
PiperOrigin-RevId: 691125900
Change-Id: Id9b7c739ee14bff2317168da258b1a3acbdafdbe
This commit is contained in:
Taylor Howell
2024-10-29 12:47:03 -07:00
committed by Copybara-Service
parent e835f139c8
commit f47840e57c
7 changed files with 109 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+11 -1
View File
@@ -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:
doc/includes/references.h
+2
View File
@@ -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]);
include/mujoco/mujoco.h
+6 -1
View File
@@ -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.
introspect/functions.py
+51 -1
View File
@@ -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(
python/mujoco/bindings_test.py
+17
View File
@@ -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.))
python/mujoco/functions.cc
+19
View File
@@ -1019,6 +1019,25 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mju_transformSpatial>(pymodule);
// Sparse math
DEF_WITH_OMITTED_PY_ARGS(traits::mju_dense2sparse, "nr", "nc", "nnz")(
pymodule,
[](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
Eigen::Ref<EigenVectorI> rownnz, Eigen::Ref<EigenVectorI> rowadr,
Eigen::Ref<EigenVectorI> 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<EigenArrayXX> res,
unity/Runtime/Bindings/MjBindings.cs
+3
View File
@@ -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);