Refactor contact forces sum in api.ode.system_velocity_dynamics

Co-authored-by: Alessandro Croci <[email protected]>

src/jaxsim/api/ode.py

@@ -150,14 +150,12 @@ def system_velocity_dynamics(
         # Sum the forces of all collidable points rigidly attached to a body.
         # Since the contact forces W_f_Ci are expressed in the world frame,
         # we don't need any coordinate transformation.
-        W_f_Li_terrain = jax.vmap(
-            lambda nc: (
-                jnp.vstack(
-                    jnp.equal(parent_link_index_of_collidable_points, nc).astype(int)
-                )
-                * W_f_Ci
-            ).sum(axis=0)
-        )(jnp.arange(model.number_of_links()))
+        W_f_Li_terrain = jnp.where(
+            parent_link_index_of_collidable_points[:, jnp.newaxis]
+            == jnp.arange(model.number_of_links()),
+            W_f_Ci,
+            0.0,
+        ).sum(axis=0)
 
     # ====================
     # Enforce joint limits