Update test_simulations.py

Enable a subset of collidable points in SoftContacts, RigidContacts, and RelaxedRigidContacts
ami-iit · Nov 5, 2024 · d14636b · d14636b
1 parent 5e403b8
commit d14636b
Showing 1 changed file with 31 additions and 0 deletions.
diff --git a/tests/test_simulations.py b/tests/test_simulations.py
@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+import numpy as np
 import pytest
 
 import jaxsim.api as js
@@ -211,6 +212,16 @@ def test_simulation_with_soft_contacts(
         model.contact_model = jaxsim.rbda.contacts.SoftContacts.build(
             terrain=model.terrain,
         )
+        # Enable a subset of the collidable points.
+        enabled_collidable_points_mask = np.zeros(
+            len(model.kin_dyn_parameters.contact_parameters.body), dtype=bool
+        )
+        enabled_collidable_points_mask[[0, 1, 2, 3]] = True
+        model.kin_dyn_parameters.contact_parameters.enabled = tuple(
+            enabled_collidable_points_mask.tolist()
+        )
+
+    assert np.sum(model.kin_dyn_parameters.contact_parameters.enabled) == 4
 
     # Initialize the maximum penetration of each collidable point at steady state.
     max_penetration = 0.001
@@ -295,6 +306,16 @@ def test_simulation_with_rigid_contacts(
         model.contact_model = jaxsim.rbda.contacts.RigidContacts.build(
             terrain=model.terrain,
         )
+        # Enable a subset of the collidable points.
+        enabled_collidable_points_mask = np.zeros(
+            len(model.kin_dyn_parameters.contact_parameters.body), dtype=bool
+        )
+        enabled_collidable_points_mask[[0, 1, 2, 3]] = True
+        model.kin_dyn_parameters.contact_parameters.enabled = tuple(
+            enabled_collidable_points_mask.tolist()
+        )
+
+    assert np.sum(model.kin_dyn_parameters.contact_parameters.enabled) == 4
 
     # Initialize the maximum penetration of each collidable point at steady state.
     # This model is rigid, so we expect (almost) no penetration.
@@ -337,6 +358,16 @@ def test_simulation_with_relaxed_rigid_contacts(
         model.contact_model = jaxsim.rbda.contacts.RelaxedRigidContacts.build(
             terrain=model.terrain,
         )
+        # Enable a subset of the collidable points.
+        enabled_collidable_points_mask = np.zeros(
+            len(model.kin_dyn_parameters.contact_parameters.body), dtype=bool
+        )
+        enabled_collidable_points_mask[[0, 1, 2, 3]] = True
+        model.kin_dyn_parameters.contact_parameters.enabled = tuple(
+            enabled_collidable_points_mask.tolist()
+        )
+
+    assert np.sum(model.kin_dyn_parameters.contact_parameters.enabled) == 4
 
     # Initialize the maximum penetration of each collidable point at steady state.
     # This model is quasi-rigid, so we expect (almost) no penetration.