replace SPARSE_CONSTRAINT_JACOBIAN with m.is_sparse

mujoco_warp/_src/constraint.py

@@ -18,7 +18,6 @@
 from mujoco_warp._src import math
 from mujoco_warp._src import support
 from mujoco_warp._src import types
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 from mujoco_warp._src.types import ConstraintType
 from mujoco_warp._src.types import ContactType
 from mujoco_warp._src.types import DisableBit
@@ -2135,7 +2134,7 @@ def make_constraint(m: types.Model, d: types.Data):
     inputs=[d.ne, d.nf, d.nl, d.nefc, efc_nnz],
   )
 
-  if types.SPARSE_CONSTRAINT_JACOBIAN:
+  if m.is_sparse:
     d.contact.efc_address.fill_(-1)
 
   if not (m.opt.disableflags & types.DisableBit.CONSTRAINT):
@@ -2163,7 +2162,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.eq_solref,
           m.eq_solimp,
           m.eq_data,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.eq_connect_adr,
           d.qvel,
           d.eq_active,
@@ -2216,7 +2215,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.eq_solref,
           m.eq_solimp,
           m.eq_data,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.eq_wld_adr,
           d.qvel,
           d.eq_active,
@@ -2262,7 +2261,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.eq_solref,
           m.eq_solimp,
           m.eq_data,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.eq_jnt_adr,
           d.qpos,
           d.qvel,
@@ -2304,7 +2303,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.ten_J_colind,
           m.tendon_length0,
           m.tendon_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.eq_ten_adr,
           d.qvel,
           d.eq_active,
@@ -2332,7 +2331,7 @@ def make_constraint(m: types.Model, d: types.Data):
       )
 
       wp.launch(
-        _equality_flex(SPARSE_CONSTRAINT_JACOBIAN),
+        _equality_flex(m.is_sparse),
         dim=(d.nworld, m.eq_flex_adr.size, m.nflexedge),
         inputs=[
           m.nv,
@@ -2382,7 +2381,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.dof_solimp,
           m.dof_frictionloss,
           m.dof_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           d.qvel,
           d.njmax,
         ],
@@ -2419,7 +2418,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.tendon_solimp_fri,
           m.tendon_frictionloss,
           m.tendon_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           d.qvel,
           d.ten_J,
           d.njmax,
@@ -2459,7 +2458,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.jnt_range,
           m.jnt_margin,
           m.dof_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.jnt_limited_ball_adr,
           d.qpos,
           d.qvel,
@@ -2498,7 +2497,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.jnt_range,
           m.jnt_margin,
           m.dof_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.jnt_limited_slide_hinge_adr,
           d.qpos,
           d.qvel,
@@ -2538,7 +2537,7 @@ def make_constraint(m: types.Model, d: types.Data):
           m.tendon_range,
           m.tendon_margin,
           m.tendon_invweight0,
-          SPARSE_CONSTRAINT_JACOBIAN,
+          m.is_sparse,
           m.tendon_limited_adr,
           d.qvel,
           d.ten_J,
@@ -2584,7 +2583,7 @@ def make_constraint(m: types.Model, d: types.Data):
             m.dof_bodyid,
             m.dof_parentid,
             m.geom_bodyid,
-            SPARSE_CONSTRAINT_JACOBIAN,
+            m.is_sparse,
             d.qvel,
             d.subtree_com,
             d.cdof,
@@ -2638,7 +2637,7 @@ def make_constraint(m: types.Model, d: types.Data):
             m.dof_bodyid,
             m.dof_parentid,
             m.geom_bodyid,
-            SPARSE_CONSTRAINT_JACOBIAN,
+            m.is_sparse,
             d.qvel,
             d.subtree_com,
             d.cdof,

mujoco_warp/_src/constraint_test.py

@@ -26,7 +26,6 @@
 import mujoco_warp as mjw
 from mujoco_warp import ConeType
 from mujoco_warp import test_data
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 
 # tolerance for difference between MuJoCo and MJWarp constraint calculations,
 # mostly due to float precision
@@ -59,7 +58,7 @@ def _assert_efc_eq(mjm, m, d, mjd, nefc, name, nv):
   mjd_sort_indices = np.lexsort((mjd_efc_pos, mjd_efc_type, mjd_efc_vel, mjd_efc_aref, mjd_efc_d))
 
   # convert sparse to dense if necessary
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if m.is_sparse:
     efc_J = np.zeros((nefc, nv))
     mujoco.mju_sparse2dense(
       efc_J,

mujoco_warp/_src/forward_test.py

@@ -28,7 +28,6 @@
 from mujoco_warp import GainType
 from mujoco_warp import IntegratorType
 from mujoco_warp import test_data
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 
 # tolerance for difference between MuJoCo and mjwarp smooth calculations - mostly
 # due to float precision
@@ -399,7 +398,7 @@ def _getattr(arr):
     nefc = d.nefc.numpy()[0]
     if nefc > 0:
       nv = m.nv
-      if SPARSE_CONSTRAINT_JACOBIAN:
+      if m.is_sparse:
         # Reconstruct dense J from sparse representation
         d_efc_J = np.zeros((nefc, nv))
         mujoco.mju_sparse2dense(

mujoco_warp/_src/io.py

@@ -28,7 +28,6 @@
 from mujoco_warp._src import types
 from mujoco_warp._src import warp_util
 from mujoco_warp._src.types import MJ_MINVAL
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 from mujoco_warp._src.types import BiasType
 from mujoco_warp._src.types import TrnType
 from mujoco_warp._src.types import vec10
@@ -712,7 +711,7 @@ def make_data(
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
 
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if is_sparse(mjm):
     efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
     efc.J_rowadr = wp.zeros((nworld, njmax), dtype=int)
     efc.J_colind = wp.zeros((nworld, 1, njmax * mjm.nv), dtype=int)
@@ -913,7 +912,7 @@ def put_data(
 
   efc = types.Constraint(**efc_kwargs)
 
-  if SPARSE_CONSTRAINT_JACOBIAN:
+  if is_sparse(mjm):
     # TODO(team): process efc_J sparsity structure for nv row shift
     efc.J_rownnz = wp.array(np.full((nworld, njmax), mjm.nv, dtype=int), dtype=int)
     efc.J_rowadr = wp.array(
@@ -1151,7 +1150,7 @@ def get_data_into(
     mujoco.mj_factorM(mjm, result)
 
   if nefc > 0:
-    if SPARSE_CONSTRAINT_JACOBIAN:
+    if is_sparse(mjm):
       efc_J = np.zeros((nefc, mjm.nv))
       mujoco.mju_sparse2dense(
         efc_J,

mujoco_warp/_src/io_test.py

@@ -30,7 +30,6 @@
 from mujoco_warp import test_data
 from mujoco_warp._src import warp_util
 from mujoco_warp._src.io import set_length_range
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 
 
 def _assert_eq(a, b, name):
@@ -1613,7 +1612,7 @@ def test_put_data_nefc_zero_dense(self):
     m = mjwarp.put_model(mjm)
     d = mjwarp.put_data(mjm, mjd)
 
-    if SPARSE_CONSTRAINT_JACOBIAN:
+    if m.is_sparse:
       self.assertEqual(d.efc.J.shape[2], d.njmax * m.nv)
     else:
       self.assertEqual(d.efc.J.shape[2], m.nv_pad)

mujoco_warp/_src/smooth.py

@@ -21,7 +21,6 @@
 from mujoco_warp._src import util_misc
 from mujoco_warp._src.types import MJ_MAXVAL
 from mujoco_warp._src.types import MJ_MINVAL
-from mujoco_warp._src.types import SPARSE_CONSTRAINT_JACOBIAN
 from mujoco_warp._src.types import CamLightType
 from mujoco_warp._src.types import ConeType
 from mujoco_warp._src.types import Data
@@ -2622,7 +2621,7 @@ def transmission(m: Model, d: Data):
         d.efc.J_colind,
         d.efc.J,
         d.nacon,
-        SPARSE_CONSTRAINT_JACOBIAN,
+        m.is_sparse,
       ],
       outputs=[d.actuator_moment, ncon],
     )