Delete sparsity structures from mjData, use the ones in mjModel

PiperOrigin-RevId: 795067950
Change-Id: If4f52e40ac49f7c6986b33ef9c8b534199dfd21a

mjx/mujoco/mjx/_src/io.py

@@ -764,19 +764,6 @@ def _make_data_c(
       'ten_velocity': (m.ntendon, float_),
       'actuator_velocity': (m.nu, float_),
       'plugin_data': (get(m, 'nplugin'), np.uint64),
-      'B_rownnz': (m.nbody, np.int32),
-      'B_rowadr': (m.nbody, np.int32),
-      'B_colind': (m.nB, np.int32),
-      'M_rownnz': (m.nv, np.int32),
-      'M_rowadr': (m.nv, np.int32),
-      'M_colind': (m.nC, np.int32),
-      'mapM2M': (m.nC, np.int32),
-      'D_rownnz': (m.nv, np.int32),
-      'D_rowadr': (m.nv, np.int32),
-      'D_diag': (m.nv, np.int32),
-      'D_colind': (m.nD, np.int32),
-      'mapM2D': (m.nD, np.int32),
-      'mapD2M': (m.nM, np.int32),
       'qDeriv': (m.nD, float_),
       'qLU': (m.nD, float_),
       'qfrc_spring': (m.nv, float_),
@@ -1497,7 +1484,7 @@ def _get_data_into(
 
     # TODO(taylorhowell): remove mapping once qM is deprecated
     # map inertia (sparse) to reduced inertia (compressed sparse) representation
-    result_i.M[:] = result_i.qM[result_i.mapM2M]
+    result_i.M[:] = result_i.qM[m.mapM2M]
 
     # recalculate qLD and qLDiagInv as MJX and MuJoCo have different
     # representations of the Cholesky decomposition.

mjx/mujoco/mjx/_src/smooth_test.py

@@ -94,7 +94,7 @@ class SmoothTest(absltest.TestCase):
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
     qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
     for i in range(m.nC):
-      qLDLegacy[d.mapM2M[i]] = d.qLD[i]
+      qLDLegacy[m.mapM2M[i]] = d.qLD[i]
     _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
     _assert_attr_eq(d, dx._impl, 'qLDiagInv')
     # com_vel

mjx/mujoco/mjx/_src/types.py

@@ -599,6 +599,19 @@ class ModelC(PyTreeNode):
   sensor_plugin: jax.Array
   plugin: jax.Array
   plugin_stateadr: jax.Array
+  B_rownnz: jax.Array  # pylint:disable=invalid-name
+  B_rowadr: jax.Array  # pylint:disable=invalid-name
+  B_colind: jax.Array  # pylint:disable=invalid-name
+  M_rownnz: jax.Array  # pylint:disable=invalid-name
+  M_rowadr: jax.Array  # pylint:disable=invalid-name
+  M_colind: jax.Array  # pylint:disable=invalid-name
+  mapM2M: jax.Array  # pylint:disable=invalid-name
+  D_rownnz: jax.Array  # pylint:disable=invalid-name
+  D_rowadr: jax.Array  # pylint:disable=invalid-name
+  D_diag: jax.Array  # pylint:disable=invalid-name
+  D_colind: jax.Array  # pylint:disable=invalid-name
+  mapM2D: jax.Array  # pylint:disable=invalid-name
+  mapD2M: jax.Array  # pylint:disable=invalid-name
 
 
 class ModelJAX(PyTreeNode):
@@ -1020,19 +1033,6 @@ class DataC(PyTreeNode):
   plugin_data: jax.Array
   qH: jax.Array  # pylint:disable=invalid-name
   qHDiagInv: jax.Array  # pylint:disable=invalid-name
-  B_rownnz: jax.Array  # pylint:disable=invalid-name
-  B_rowadr: jax.Array  # pylint:disable=invalid-name
-  B_colind: jax.Array  # pylint:disable=invalid-name
-  M_rownnz: jax.Array  # pylint:disable=invalid-name
-  M_rowadr: jax.Array  # pylint:disable=invalid-name
-  M_colind: jax.Array  # pylint:disable=invalid-name
-  mapM2M: jax.Array  # pylint:disable=invalid-name
-  D_rownnz: jax.Array  # pylint:disable=invalid-name
-  D_rowadr: jax.Array  # pylint:disable=invalid-name
-  D_diag: jax.Array  # pylint:disable=invalid-name
-  D_colind: jax.Array  # pylint:disable=invalid-name
-  mapM2D: jax.Array  # pylint:disable=invalid-name
-  mapD2M: jax.Array  # pylint:disable=invalid-name
   qDeriv: jax.Array  # pylint:disable=invalid-name
   qLU: jax.Array  # pylint:disable=invalid-name
   qfrc_spring: jax.Array

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -142,9 +142,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   # calculate some fields that cannot be easily computed inline
   nlsp = mjm.opt.ls_iterations  # TODO(team): how to set nlsp?
 
-  # unfortunately we must create Data in order to get some model fields like M_rownnz
-  mjd = mujoco.MjData(mjm)
-
   # dof lower triangle row and column indices (used in solver)
   dof_tri_row, dof_tri_col = np.tril_indices(mjm.nv)
 
@@ -182,11 +179,11 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
 
   for k in range(mjm.nv):
     # skip diagonal rows
-    if mjd.M_rownnz[k] == 1:
+    if mjm.M_rownnz[k] == 1:
       continue
     dof_depth[k] = dof_depth[mjm.dof_parentid[k]] + 1
     i = mjm.dof_parentid[k]
-    diag_k = mjd.M_rowadr[k] + mjd.M_rownnz[k] - 1
+    diag_k = mjm.M_rowadr[k] + mjm.M_rownnz[k] - 1
     Madr_ki = diag_k - 1
     while i > -1:
       qLD_updates.setdefault(dof_depth[i], []).append((i, k, Madr_ki))
@@ -487,10 +484,10 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     qM_mulm_j=wp.array(qM_mulm_j, dtype=int),
     qM_madr_ij=wp.array(qM_madr_ij, dtype=int),
     qLD_updates=qLD_updates,
-    M_rownnz=wp.array(mjd.M_rownnz, dtype=int),
-    M_rowadr=wp.array(mjd.M_rowadr, dtype=int),
-    M_colind=wp.array(mjd.M_colind, dtype=int),
-    mapM2M=wp.array(mjd.mapM2M, dtype=int),
+    M_rownnz=wp.array(mjm.M_rownnz, dtype=int),
+    M_rowadr=wp.array(mjm.M_rowadr, dtype=int),
+    M_colind=wp.array(mjm.M_colind, dtype=int),
+    mapM2M=wp.array(mjm.mapM2M, dtype=int),
     qM_tiles=qM_tiles,
     body_tree=body_tree,
     body_parentid=wp.array(mjm.body_parentid, dtype=int),