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
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,