Apply suggestions from code review

Co-authored-by: Filippo Luca Ferretti <[email protected]>

src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -166,13 +166,13 @@ def __eq__(self, other: RelaxedRigidContactsState) -> bool:
     def build(cls: type[Self]) -> Self:
         """Create a `RelaxedRigidContactsState` instance"""
 
-        return RelaxedRigidContactsState()
+        return cls()
 
     @classmethod
-    def zero(cls: type[Self], *, model: js.model.JaxSimModel) -> Self:
+    def zero(cls: type[Self]) -> Self:
         """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
 
-        return RelaxedRigidContactsState.build()
+        return cls.build()
 
     def valid(self, *, model: js.model.JaxSimModel) -> jtp.BoolLike:
         return True

src/jaxsim/rbda/contacts/rigid.py

@@ -63,7 +63,7 @@ def build(
     ) -> Self:
         """Create a `RigidContactParams` instance"""
 
-        return RigidContactsParams(
+        return cls(
             mu=mu or cls.__dataclass_fields__["mu"].default,
             K=K or cls.__dataclass_fields__["K"].default,
             D=D or cls.__dataclass_fields__["D"].default,
@@ -89,13 +89,13 @@ def __eq__(self, other: RigidContactsState) -> bool:
     def build(cls: type[Self]) -> Self:
         """Create a `RigidContactsState` instance"""
 
-        return RigidContactsState()
+        return cls()
 
     @classmethod
     def zero(cls: type[Self]) -> Self:
         """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
 
-        return RigidContactsState.build()
+        return cls.build()
 
     def valid(self) -> jtp.BoolLike:
         return True

src/jaxsim/rbda/contacts/soft.py

@@ -178,7 +178,7 @@ def valid(self) -> jtp.BoolLike:
         return jnp.hstack(
             [
                 self.K >= 0.0,
-                self.D >= 0,
+                self.D >= 0.0,
                 self.mu >= 0.0,
                 self.p >= 0.0,
                 self.q >= 0.0,
@@ -310,9 +310,11 @@ def compute_contact_forces(
         contact_status += (δ <= 0).astype(int)
 
         # Select the right material deformation rate depending on the contact status.
-        ṁ = jax.lax.switch(
-            index=contact_status,
-            branches=(lambda: ṁ_slipping, lambda: ṁ_sticking, lambda: ṁ_no_contact),
+        ṁ = jax.lax.select_n(
+            ṁ_slipping,
+            ṁ_sticking,
+            ṁ_no_contact,
+            which=contact_status,
         )
 
         # ==========================================