jaxsim 0.3.1.dev17__py3-none-any.whl → 0.3.1.dev40__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.1.dev17'
-__version_tuple__ = version_tuple = (0, 3, 1, 'dev17')
+__version__ = version = '0.3.1.dev40'
+__version_tuple__ = version_tuple = (0, 3, 1, 'dev40')

jaxsim/api/contact.py

@@ -1,11 +1,14 @@
+from __future__ import annotations
+
 import functools
 
 import jax
 import jax.numpy as jnp
 
 import jaxsim.api as js
-import jaxsim.rbda
+import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsParams
 
 from .common import VelRepr
 
@@ -135,17 +138,23 @@ def collidable_point_dynamics(
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
     # Build the soft contact model.
-    soft_contacts = jaxsim.rbda.SoftContacts(
-        parameters=data.soft_contacts_params, terrain=model.terrain
-    )
+    match model.contact_model:
+        case s if isinstance(s, SoftContacts):
+            # Build the contact model.
+            soft_contacts = SoftContacts(
+                parameters=data.contacts_params, terrain=model.terrain
+            )
+
+            # Compute the 6D force expressed in the inertial frame and applied to each
+            # collidable point, and the corresponding material deformation rate.
+            # Note that the material deformation rate is always returned in the mixed frame
+            # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
+            W_f_Ci, (CW_ṁ,) = jax.vmap(soft_contacts.compute_contact_forces)(
+                W_p_Ci, W_ṗ_Ci, data.state.contact.tangential_deformation
+            )
 
-    # Compute the 6D force expressed in the inertial frame and applied to each
-    # collidable point, and the corresponding material deformation rate.
-    # Note that the material deformation rate is always returned in the mixed frame
-    # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
-    W_f_Ci, CW_ṁ = jax.vmap(soft_contacts.contact_model)(
-        W_p_Ci, W_ṗ_Ci, data.state.soft_contacts.tangential_deformation
-    )
+        case _:
+            raise ValueError("Invalid contact model {}".format(model.contact_model))
 
     # Convert the 6D forces to the active representation.
     f_Ci = jax.vmap(
@@ -213,7 +222,7 @@ def estimate_good_soft_contacts_parameters(
     number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
     damping_ratio: jtp.FloatLike = 1.0,
     max_penetration: jtp.FloatLike | None = None,
-) -> jaxsim.rbda.soft_contacts.SoftContactsParams:
+) -> SoftContactsParams:
     """
     Estimate good soft contacts parameters for the given model.
 
@@ -237,13 +246,14 @@ def estimate_good_soft_contacts_parameters(
         The user is encouraged to fine-tune the parameters based on the
         specific application.
     """
+    from jaxsim.rbda.contacts.soft import SoftContactsParams
 
     def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
         """"""
 
         zero_data = js.data.JaxSimModelData.build(
             model=model,
-            soft_contacts_params=jaxsim.rbda.soft_contacts.SoftContactsParams(),
+            contacts_params=SoftContactsParams(),
         )
 
         W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
@@ -262,15 +272,13 @@ def estimate_good_soft_contacts_parameters(
 
     nc = number_of_active_collidable_points_steady_state
 
-    sc_parameters = (
-        jaxsim.rbda.soft_contacts.SoftContactsParams.build_default_from_jaxsim_model(
-            model=model,
-            standard_gravity=standard_gravity,
-            static_friction_coefficient=static_friction_coefficient,
-            max_penetration=max_δ,
-            number_of_active_collidable_points_steady_state=nc,
-            damping_ratio=damping_ratio,
-        )
+    sc_parameters = SoftContactsParams.build_default_from_jaxsim_model(
+        model=model,
+        standard_gravity=standard_gravity,
+        static_friction_coefficient=static_friction_coefficient,
+        max_penetration=max_δ,
+        number_of_active_collidable_points_steady_state=nc,
+        damping_ratio=damping_ratio,
     )
 
     return sc_parameters

jaxsim/api/data.py

@@ -14,6 +14,7 @@ import jaxsim.api as js
 import jaxsim.rbda
 import jaxsim.typing as jtp
 from jaxsim.math import Quaternion
+from jaxsim.rbda.contacts.soft import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
 
@@ -37,7 +38,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     gravity: jtp.Array
 
-    soft_contacts_params: jaxsim.rbda.SoftContactsParams = dataclasses.field(repr=False)
+    contacts_params: jaxsim.rbda.ContactsParams = dataclasses.field(repr=False)
 
     time_ns: jtp.Int = dataclasses.field(
         default_factory=lambda: jnp.array(0, dtype=jnp.uint64)
@@ -51,8 +52,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             (
                 hash(self.state),
                 HashedNumpyArray.hash_of_array(self.gravity),
-                hash(self.soft_contacts_params),
                 HashedNumpyArray.hash_of_array(self.time_ns),
+                hash(self.contacts_params),
             )
         )
 
@@ -112,8 +113,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         base_angular_velocity: jtp.Vector | None = None,
         joint_velocities: jtp.Vector | None = None,
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
-        soft_contacts_state: js.ode_data.SoftContactsState | None = None,
-        soft_contacts_params: jaxsim.rbda.SoftContactsParams | None = None,
+        contact: jaxsim.rbda.ContactsState | None = None,
+        contacts_params: jaxsim.rbda.ContactsParams | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
         time: jtp.FloatLike | None = None,
     ) -> JaxSimModelData:
@@ -131,8 +132,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
             standard_gravity: The standard gravity constant.
-            soft_contacts_state: The state of the soft contacts.
-            soft_contacts_params: The parameters of the soft contacts.
+            contact: The state of the soft contacts.
+            contacts_params: The parameters of the soft contacts.
             velocity_representation: The velocity representation to use.
             time: The time at which the state is created.
 
@@ -178,13 +179,16 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             else jnp.array(0, dtype=jnp.uint64)
         )
 
-        soft_contacts_params = (
-            soft_contacts_params
-            if soft_contacts_params is not None
-            else js.contact.estimate_good_soft_contacts_parameters(
-                model=model, standard_gravity=standard_gravity
+        if isinstance(model.contact_model, SoftContacts):
+            contacts_params = (
+                contacts_params
+                if contacts_params is not None
+                else js.contact.estimate_good_soft_contacts_parameters(
+                    model=model, standard_gravity=standard_gravity
+                )
             )
-        )
+        else:
+            contacts_params = model.contact_model.parameters
 
         W_H_B = jaxlie.SE3.from_rotation_and_translation(
             translation=base_position,
@@ -209,8 +213,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_angular_velocity=v_WB[3:6].astype(float),
             joint_velocities=joint_velocities.astype(float),
             tangential_deformation=(
-                soft_contacts_state.tangential_deformation
-                if soft_contacts_state is not None
+                contact.tangential_deformation
+                if contact is not None and isinstance(model.contact_model, SoftContacts)
                 else None
             ),
         )
@@ -222,7 +226,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             time_ns=time_ns,
             state=ode_state,
             gravity=gravity.astype(float),
-            soft_contacts_params=soft_contacts_params,
+            contacts_params=contacts_params,
             velocity_representation=velocity_representation,
         )
 
@@ -652,7 +656,10 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.reset_base_velocity(
             base_velocity=jnp.hstack(
-                [linear_velocity.squeeze(), self.base_velocity()[3:6]]
+                [
+                    linear_velocity.squeeze(),
+                    self.base_velocity()[3:6],
+                ]
             ),
             velocity_representation=velocity_representation,
         )
@@ -680,7 +687,10 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.reset_base_velocity(
             base_velocity=jnp.hstack(
-                [self.base_velocity()[0:3], angular_velocity.squeeze()]
+                [
+                    self.base_velocity()[0:3],
+                    angular_velocity.squeeze(),
+                ]
             ),
             velocity_representation=velocity_representation,
         )

jaxsim/api/model.py

@@ -34,6 +34,10 @@ class JaxSimModel(JaxsimDataclass):
         default=jaxsim.terrain.FlatTerrain(), repr=False
     )
 
+    contact_model: jaxsim.rbda.ContactModel | None = dataclasses.field(
+        default=None, repr=False
+    )
+
     kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
         dataclasses.field(default=None, repr=False)
     )
@@ -69,6 +73,7 @@ class JaxSimModel(JaxsimDataclass):
             (
                 hash(self.model_name),
                 hash(self.kin_dyn_parameters),
+                hash(self.contact_model),
             )
         )
 
@@ -82,6 +87,7 @@ class JaxSimModel(JaxsimDataclass):
         model_name: str | None = None,
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
+        contact_model: jaxsim.rbda.ContactModel | None = None,
         is_urdf: bool | None = None,
         considered_joints: Sequence[str] | None = None,
     ) -> JaxSimModel:
@@ -127,6 +133,7 @@ class JaxSimModel(JaxsimDataclass):
             model_description=intermediate_description,
             model_name=model_name,
             terrain=terrain,
+            contact_model=contact_model,
         )
 
         # Store the origin of the model, in case downstream logic needs it
@@ -141,6 +148,7 @@ class JaxSimModel(JaxsimDataclass):
         model_name: str | None = None,
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
+        contact_model: jaxsim.rbda.ContactModel | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from an intermediate model description.
@@ -153,22 +161,30 @@ class JaxSimModel(JaxsimDataclass):
                 The optional name of the model overriding the physics model name.
             terrain:
                 The optional terrain to consider.
+            contact_model:
+                The optional contact model to consider. If None, the soft contact model is used.
 
         Returns:
             The built Model object.
         """
+        from jaxsim.rbda.contacts.soft import SoftContacts
 
         # Set the model name (if not provided, use the one from the model description)
         model_name = model_name if model_name is not None else model_description.name
 
-        # Build the model.
+        # Set the terrain (if not provided, use the default flat terrain)
+        terrain = terrain or JaxSimModel.__dataclass_fields__["terrain"].default
+        contact_model = contact_model or SoftContacts(terrain=terrain)
+
+        # Build the model
         model = JaxSimModel(
             model_name=model_name,
             _description=wrappers.HashlessObject(obj=model_description),
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
-            terrain=terrain or JaxSimModel.__dataclass_fields__["terrain"].default,
+            terrain=terrain,
+            contact_model=contact_model,
         )
 
         return model
@@ -350,6 +366,7 @@ def reduce(
         model_description=reduced_intermediate_description,
         model_name=model.name(),
         terrain=model.terrain,
+        contact_model=model.contact_model,
     )
 
     # Store the origin of the model, in case downstream logic needs it

jaxsim/api/ode.py

@@ -132,7 +132,7 @@ def system_velocity_dynamics(
     W_f_Ci = None
 
     # Initialize the derivative of the tangential deformation ṁ ∈ ℝ^{n_c × 3}.
-    ṁ = jnp.zeros_like(data.state.soft_contacts.tangential_deformation).astype(float)
+    ṁ = jnp.zeros_like(data.state.contact.tangential_deformation).astype(float)
 
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
         # Compute the 6D forces applied to each collidable point and the

jaxsim/api/ode_data.py

@@ -5,6 +5,8 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim.rbda import ContactsState
+from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 from jaxsim.utils import JaxsimDataclass
 
 # =============================================================================
@@ -116,11 +118,11 @@ class ODEState(JaxsimDataclass):
 
     Attributes:
         physics_model: The state of the physics model.
-        soft_contacts: The state of the soft-contacts model.
+        contact: The state of the contacts model.
     """
 
     physics_model: PhysicsModelState
-    soft_contacts: SoftContactsState
+    contact: ContactsState
 
     @staticmethod
     def build_from_jaxsim_model(
@@ -158,6 +160,20 @@ class ODEState(JaxsimDataclass):
             `JaxSimModel` and initialized to zero.
         """
 
+        # Get the contact model from the `JaxSimModel`
+        match model.contact_model:
+            case SoftContacts():
+                contact = SoftContactsState.build_from_jaxsim_model(
+                    model=model,
+                    **(
+                        dict(tangential_deformation=tangential_deformation)
+                        if tangential_deformation is not None
+                        else dict()
+                    ),
+                )
+            case _:
+                raise ValueError("Unable to determine contact state class prefix.")
+
         return ODEState.build(
             model=model,
             physics_model_state=PhysicsModelState.build_from_jaxsim_model(
@@ -169,24 +185,21 @@ class ODEState(JaxsimDataclass):
                 base_linear_velocity=base_linear_velocity,
                 base_angular_velocity=base_angular_velocity,
             ),
-            soft_contacts_state=SoftContactsState.build_from_jaxsim_model(
-                model=model,
-                tangential_deformation=tangential_deformation,
-            ),
+            contact=contact,
         )
 
     @staticmethod
     def build(
         physics_model_state: PhysicsModelState | None = None,
-        soft_contacts_state: SoftContactsState | None = None,
+        contact: ContactsState | None = None,
         model: js.model.JaxSimModel | None = None,
     ) -> ODEState:
         """
-        Build an `ODEState` from a `PhysicsModelState` and a `SoftContactsState`.
+        Build an `ODEState` from a `PhysicsModelState` and a `ContactsState`.
 
         Args:
             physics_model_state: The state of the physics model.
-            soft_contacts_state: The state of the soft-contacts model.
+            contact: The state of the contacts model.
             model: The `JaxSimModel` associated with the ODE state.
 
         Returns:
@@ -199,15 +212,16 @@ class ODEState(JaxsimDataclass):
             else PhysicsModelState.zero(model=model)
         )
 
-        soft_contacts_state = (
-            soft_contacts_state
-            if soft_contacts_state is not None
-            else SoftContactsState.zero(model=model)
-        )
+        # Get the contact model from the `JaxSimModel`
+        match contact:
+            case SoftContactsState():
+                pass
+            case None:
+                contact = SoftContactsState.zero(model=model)
+            case _:
+                raise ValueError("Unable to determine contact state class prefix.")
 
-        return ODEState(
-            physics_model=physics_model_state, soft_contacts=soft_contacts_state
-        )
+        return ODEState(physics_model=physics_model_state, contact=contact)
 
     @staticmethod
     def zero(model: js.model.JaxSimModel) -> ODEState:
@@ -236,9 +250,7 @@ class ODEState(JaxsimDataclass):
             `True` if the ODE state is valid for the given model, `False` otherwise.
         """
 
-        return self.physics_model.valid(model=model) and self.soft_contacts.valid(
-            model=model
-        )
+        return self.physics_model.valid(model=model) and self.contact.valid(model=model)
 
 
 # ==================================================
@@ -595,135 +607,3 @@ class PhysicsModelInput(JaxsimDataclass):
             return False
 
         return True
-
-
-# ===========================================
-# Define the state of the soft-contacts model
-# ===========================================
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContactsState(JaxsimDataclass):
-    """
-    Class storing the state of the soft contacts model.
-
-    Attributes:
-        tangential_deformation:
-            The matrix of 3D tangential material deformations corresponding to
-            each collidable point.
-    """
-
-    tangential_deformation: jtp.Matrix
-
-    def __hash__(self) -> int:
-
-        from jaxsim.utils.wrappers import HashedNumpyArray
-
-        return HashedNumpyArray.hash_of_array(self.tangential_deformation)
-
-    def __eq__(self, other: SoftContactsState) -> bool:
-
-        if not isinstance(other, SoftContactsState):
-            return False
-
-        return hash(self) == hash(other)
-
-    @staticmethod
-    def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
-        tangential_deformation: jtp.Matrix | None = None,
-    ) -> SoftContactsState:
-        """
-        Build a `SoftContactsState` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the soft contacts state.
-            tangential_deformation: The matrix of 3D tangential material deformations.
-
-        Returns:
-            The `SoftContactsState` built from the `JaxSimModel`.
-
-        Note:
-            If any of the state components are not provided, they are built from the
-            `JaxSimModel` and initialized to zero.
-        """
-
-        return SoftContactsState.build(
-            tangential_deformation=tangential_deformation,
-            number_of_collidable_points=len(
-                model.kin_dyn_parameters.contact_parameters.body
-            ),
-        )
-
-    @staticmethod
-    def build(
-        tangential_deformation: jtp.Matrix | None = None,
-        number_of_collidable_points: int | None = None,
-    ) -> SoftContactsState:
-        """
-        Create a `SoftContactsState`.
-
-        Args:
-            tangential_deformation:
-                The matrix of 3D tangential material deformations corresponding to
-                each collidable point.
-            number_of_collidable_points: The number of collidable points.
-
-        Returns:
-            A `SoftContactsState` instance.
-        """
-
-        tangential_deformation = (
-            tangential_deformation
-            if tangential_deformation is not None
-            else jnp.zeros(shape=(number_of_collidable_points, 3))
-        )
-
-        if tangential_deformation.shape[1] != 3:
-            raise RuntimeError("The tangential deformation matrix must have 3 columns.")
-
-        if (
-            number_of_collidable_points is not None
-            and tangential_deformation.shape[0] != number_of_collidable_points
-        ):
-            msg = "The number of collidable points must match the number of rows "
-            msg += "in the tangential deformation matrix."
-            raise RuntimeError(msg)
-
-        return SoftContactsState(
-            tangential_deformation=jnp.array(tangential_deformation).astype(float)
-        )
-
-    @staticmethod
-    def zero(model: js.model.JaxSimModel) -> SoftContactsState:
-        """
-        Build a zero `SoftContactsState` from a `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` associated with the soft contacts state.
-
-        Returns:
-            A zero `SoftContactsState` instance.
-        """
-
-        return SoftContactsState.build_from_jaxsim_model(model=model)
-
-    def valid(self, model: js.model.JaxSimModel) -> bool:
-        """
-        Check if the `SoftContactsState` is valid for a given `JaxSimModel`.
-
-        Args:
-            model: The `JaxSimModel` to validate the `SoftContactsState` against.
-
-        Returns:
-            `True` if the soft contacts state is valid for the given `JaxSimModel`,
-            `False` otherwise.
-        """
-
-        shape = self.tangential_deformation.shape
-        expected = (len(model.kin_dyn_parameters.contact_parameters.body), 3)
-
-        if shape != expected:
-            return False
-
-        return True

jaxsim/exceptions.py

@@ -0,0 +1,63 @@
+import jax
+
+
+def raise_if(
+    condition: bool | jax.Array, exception: type, msg: str, *args, **kwargs
+) -> None:
+    """
+    Raise a host-side exception if a condition is met. Useful in jit-compiled functions.
+
+    Args:
+        condition:
+            The boolean condition of the evaluated expression that triggers
+            the exception during runtime.
+        exception: The type of exception to raise.
+        msg:
+            The message to display when the exception is raised. The message can be a
+            format string (fmt), whose fields are filled with the args and kwargs.
+    """
+
+    # Check early that the format string is well-formed.
+    try:
+        _ = msg.format(*args, **kwargs)
+    except Exception as e:
+        msg = "Error in formatting exception message with args={} and kwargs={}"
+        raise ValueError(msg.format(args, kwargs)) from e
+
+    def _raise_exception(condition: bool, *args, **kwargs) -> None:
+        """The function called by the JAX callback."""
+
+        if condition:
+            raise exception(msg.format(*args, **kwargs))
+
+    def _callback(args, kwargs) -> None:
+        """The function that calls the JAX callback, executed only when needed."""
+
+        jax.debug.callback(_raise_exception, condition, *args, **kwargs)
+
+    # Since running a callable on the host is expensive, we prevent its execution
+    # if the condition is False with a low-level conditional expression.
+    def _run_callback_only_if_condition_is_true(*args, **kwargs) -> None:
+        return jax.lax.cond(
+            condition,
+            _callback,
+            lambda args, kwargs: None,
+            args,
+            kwargs,
+        )
+
+    return _run_callback_only_if_condition_is_true(*args, **kwargs)
+
+
+def raise_runtime_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+
+    return raise_if(condition, RuntimeError, msg, *args, **kwargs)
+
+
+def raise_value_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+
+    return raise_if(condition, ValueError, msg, *args, **kwargs)

jaxsim/rbda/__init__.py

@@ -1,5 +1,6 @@
 from .aba import aba
 from .collidable_points import collidable_points_pos_vel
+from .contacts.common import ContactModel, ContactsParams, ContactsState
 from .crba import crba
 from .forward_kinematics import forward_kinematics, forward_kinematics_model
 from .jacobian import (
@@ -8,4 +9,3 @@ from .jacobian import (
     jacobian_full_doubly_left,
 )
 from .rnea import rnea
-from .soft_contacts import SoftContacts, SoftContactsParams

jaxsim/rbda/contacts/common.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import abc
+from typing import Any
+
+import jaxsim.terrain
+import jaxsim.typing as jtp
+
+
+class ContactsState(abc.ABC):
+    """
+    Abstract class storing the state of the contacts model.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def build(cls, **kwargs) -> ContactsState:
+        """
+        Build the contact state object.
+
+        Returns:
+            The contact state object.
+        """
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def zero(cls, **kwargs) -> ContactsState:
+        """
+        Build a zero contact state.
+
+        Returns:
+            The zero contact state.
+        """
+        pass
+
+    @abc.abstractmethod
+    def valid(self, **kwargs) -> bool:
+        """
+        Check if the contacts state is valid.
+        """
+        pass
+
+
+class ContactsParams(abc.ABC):
+    """
+    Abstract class representing the parameters of a contact model.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def build(cls) -> ContactsParams:
+        """
+        Create a `ContactsParams` instance with specified parameters.
+        Returns:
+            The `ContactsParams` instance.
+        """
+        pass
+
+    @abc.abstractmethod
+    def valid(self, *args, **kwargs) -> bool:
+        """
+        Check if the parameters are valid.
+        Returns:
+            True if the parameters are valid, False otherwise.
+        """
+        pass
+
+
+class ContactModel(abc.ABC):
+    """
+    Abstract class representing a contact model.
+
+    Attributes:
+        parameters: The parameters of the contact model.
+        terrain: The terrain model.
+    """
+
+    parameters: ContactsParams
+    terrain: jaxsim.terrain.Terrain
+
+    @abc.abstractmethod
+    def compute_contact_forces(
+        self,
+        position: jtp.Vector,
+        velocity: jtp.Vector,
+        **kwargs,
+    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
+        """
+        Compute the contact forces.
+
+        Args:
+            position: The position of the collidable point w.r.t. the world frame.
+            velocity:
+                The linear velocity of the collidable point (linear component of the mixed 6D velocity).
+
+        Returns:
+            A tuple containing as first element the computed 6D contact force applied to the contact point and expressed in the world frame,
+            and as second element a tuple of optional additional information.
+        """
+        pass

jaxsim/rbda/{soft_contacts.py → contacts/soft.py}

@@ -10,11 +10,12 @@ import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.math import Skew, StandardGravity
 from jaxsim.terrain import FlatTerrain, Terrain
-from jaxsim.utils import JaxsimDataclass
+
+from .common import ContactModel, ContactsParams, ContactsState
 
 
 @jax_dataclasses.pytree_dataclass
-class SoftContactsParams(JaxsimDataclass):
+class SoftContactsParams(ContactsParams):
     """Parameters of the soft contacts model."""
 
     K: jtp.Float = dataclasses.field(
@@ -127,9 +128,23 @@ class SoftContactsParams(JaxsimDataclass):
 
         return SoftContactsParams.build(K=K, D=D, mu=μc)
 
+    def valid(self) -> bool:
+        """
+        Check if the parameters are valid.
+
+        Returns:
+            `True` if the parameters are valid, `False` otherwise.
+        """
+
+        return (
+            jnp.all(self.K >= 0.0)
+            and jnp.all(self.D >= 0.0)
+            and jnp.all(self.mu >= 0.0)
+        )
+
 
 @jax_dataclasses.pytree_dataclass
-class SoftContacts:
+class SoftContacts(ContactModel):
     """Soft contacts model."""
 
     parameters: SoftContactsParams = dataclasses.field(
@@ -138,12 +153,12 @@ class SoftContacts:
 
     terrain: Terrain = dataclasses.field(default_factory=FlatTerrain)
 
-    def contact_model(
+    def compute_contact_forces(
         self,
         position: jtp.Vector,
         velocity: jtp.Vector,
         tangential_deformation: jtp.Vector,
-    ) -> tuple[jtp.Vector, jtp.Vector]:
+    ) -> tuple[jtp.Vector, tuple[jtp.Vector, None]]:
         """
         Compute the contact forces and material deformation rate.
 
@@ -222,7 +237,7 @@ class SoftContacts:
             # Compute lin-ang 6D forces (inertial representation)
             W_f = W_Xf_CW @ CW_f
 
-            return W_f, ṁ
+            return W_f, (ṁ,)
 
         # =========================
         # Compute tangential forces
@@ -240,7 +255,7 @@ class SoftContacts:
             active_contact = pz < self.terrain.height(x=px, y=py)
 
             def above_terrain():
-                return jnp.zeros(6), ṁ
+                return jnp.zeros(6), (ṁ,)
 
             def below_terrain():
                 # Decompose the velocity in normal and tangential components
@@ -296,9 +311,9 @@ class SoftContacts:
                 W_f = W_Xf_CW @ CW_f
 
                 # Return the 6D force in the world frame and the deformation derivative
-                return W_f, ṁ
+                return W_f, (ṁ,)
 
-            # (W_f, ṁ)
+            # (W_f, (ṁ,))
             return jax.lax.cond(
                 pred=active_contact,
                 true_fun=lambda _: below_terrain(),
@@ -313,3 +328,128 @@ class SoftContacts:
             false_fun=lambda _: with_friction(),
             operand=None,
         )
+
+
+@jax_dataclasses.pytree_dataclass
+class SoftContactsState(ContactsState):
+    """
+    Class storing the state of the soft contacts model.
+
+    Attributes:
+        tangential_deformation:
+            The matrix of 3D tangential material deformations corresponding to
+            each collidable point.
+    """
+
+    tangential_deformation: jtp.Matrix
+
+    def __hash__(self) -> int:
+        return hash(
+            tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
+        )
+
+    def __eq__(self, other: SoftContactsState) -> bool:
+        if not isinstance(other, SoftContactsState):
+            return False
+
+        return hash(self) == hash(other)
+
+    @staticmethod
+    def build_from_jaxsim_model(
+        model: js.model.JaxSimModel | None = None,
+        tangential_deformation: jtp.Matrix | None = None,
+    ) -> SoftContactsState:
+        """
+        Build a `SoftContactsState` from a `JaxSimModel`.
+
+        Args:
+            model: The `JaxSimModel` associated with the soft contacts state.
+            tangential_deformation: The matrix of 3D tangential material deformations.
+
+        Returns:
+            The `SoftContactsState` built from the `JaxSimModel`.
+
+        Note:
+            If any of the state components are not provided, they are built from the
+            `JaxSimModel` and initialized to zero.
+        """
+
+        return SoftContactsState.build(
+            tangential_deformation=tangential_deformation,
+            number_of_collidable_points=len(
+                model.kin_dyn_parameters.contact_parameters.body
+            ),
+        )
+
+    @staticmethod
+    def build(
+        tangential_deformation: jtp.Matrix | None = None,
+        number_of_collidable_points: int | None = None,
+    ) -> SoftContactsState:
+        """
+        Create a `SoftContactsState`.
+
+        Args:
+            tangential_deformation:
+                The matrix of 3D tangential material deformations corresponding to
+                each collidable point.
+            number_of_collidable_points: The number of collidable points.
+
+        Returns:
+            A `SoftContactsState` instance.
+        """
+
+        tangential_deformation = (
+            tangential_deformation
+            if tangential_deformation is not None
+            else jnp.zeros(shape=(number_of_collidable_points, 3))
+        )
+
+        if tangential_deformation.shape[1] != 3:
+            raise RuntimeError("The tangential deformation matrix must have 3 columns.")
+
+        if (
+            number_of_collidable_points is not None
+            and tangential_deformation.shape[0] != number_of_collidable_points
+        ):
+            msg = "The number of collidable points must match the number of rows "
+            msg += "in the tangential deformation matrix."
+            raise RuntimeError(msg)
+
+        return SoftContactsState(
+            tangential_deformation=jnp.array(tangential_deformation).astype(float)
+        )
+
+    @staticmethod
+    def zero(model: js.model.JaxSimModel) -> SoftContactsState:
+        """
+        Build a zero `SoftContactsState` from a `JaxSimModel`.
+
+        Args:
+            model: The `JaxSimModel` associated with the soft contacts state.
+
+        Returns:
+            A zero `SoftContactsState` instance.
+        """
+
+        return SoftContactsState.build_from_jaxsim_model(model=model)
+
+    def valid(self, model: js.model.JaxSimModel) -> bool:
+        """
+        Check if the `SoftContactsState` is valid for a given `JaxSimModel`.
+
+        Args:
+            model: The `JaxSimModel` to validate the `SoftContactsState` against.
+
+        Returns:
+            `True` if the soft contacts state is valid for the given `JaxSimModel`,
+            `False` otherwise.
+        """
+
+        shape = self.tangential_deformation.shape
+        expected = (len(model.kin_dyn_parameters.contact_parameters.body), 3)
+
+        if shape != expected:
+            return False
+
+        return True

{jaxsim-0.3.1.dev17.dist-info → jaxsim-0.3.1.dev40.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.3.1.dev17
+Version: 0.3.1.dev40
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.3.1.dev17.dist-info → jaxsim-0.3.1.dev40.dist-info}/RECORD

@@ -1,19 +1,20 @@
 jaxsim/__init__.py,sha256=xzuTuZrgKdWLqqDzbvqzm2cJrEtAbepOeUqDu7ByVek,2621
-jaxsim/_version.py,sha256=EQQfkY5WXMHFjdRnYAQqABGWC0VK4dlpuNh_wr1KxYA,426
+jaxsim/_version.py,sha256=JHeHgaRnZLSH9QDd807jRnlOgX9sn4w_CGqGSxA8lL0,426
+jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=cl7HHQCeP3mHmtF6EuQZcCjGvDmc_AryMWntP_lRBGg,722
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=Yof6otFi-mLWAs1rqjmeNJTOWIH9gn7BdU5EIjiL6Ts,13481
 jaxsim/api/common.py,sha256=bqQ__pIQZbh-j8rkoHUkYHAgGiJnDzjHG-q4Ny0OOYQ,6646
-jaxsim/api/contact.py,sha256=79kcdq7C1_kWgxd1QWBabBhIPkwWEVLk-Fiz9kh-4so,12800
-jaxsim/api/data.py,sha256=fkVDBV1tODRYIaRb2N15l34InAcnzNygMGG1KFiIU2w,27307
+jaxsim/api/contact.py,sha256=soB28vqmzUwE6CN36TU4keASWZoSWE2_zhJLXA8yw2E,13132
+jaxsim/api/data.py,sha256=oAJ2suPeQLQZGpHZi98g6UZp1VcoDtuqT_aZBpynA30,27582
 jaxsim/api/frame.py,sha256=vSbFHL4WtKPySxunNoZLlM_aDuJXZtf8CSBKku63BAs,6178
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
 jaxsim/api/kin_dyn_parameters.py,sha256=AEpDg9kihbKUN9PA8pNrAruSuWFUC-k_GGxtlcdcDiQ,29215
 jaxsim/api/link.py,sha256=MdMWaMpM5Dj5JHK8uwHZ4zR4Fjq3R4asi2sGTxk1OAs,16647
-jaxsim/api/model.py,sha256=iuNYsn4xIfX36smmZpwM2O5eftT7ioDQtb6mSUqWu6Q,59759
-jaxsim/api/ode.py,sha256=luTQJsIXUtCp_81dR42X7WrMvwrXtYbyJiqss29v7zA,10786
-jaxsim/api/ode_data.py,sha256=FxUIV5qDNOg_OiOXWs3UrhDgKhGmTKcbHqgr4NX5bv0,23290
+jaxsim/api/model.py,sha256=HAnrlgPDl5CCZQzQ84AfjC_DZjmrCzBKEDodE6hyLf8,60518
+jaxsim/api/ode.py,sha256=xQL53ppnKweMQWRNm5gGR8FTjqRVzds8WKg9js9k5TA,10780
+jaxsim/api/ode_data.py,sha256=Sa2i1zZhqyQqIGv1jarTmmU-W9HhTw-DErs12kFA1GA,19737
 jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=9HXRVFo95Mpt6RcVhBrOfvOO7mDxqbkXeg_lKUibEFY,20693
@@ -43,23 +44,25 @@ jaxsim/parsers/descriptions/model.py,sha256=vfubtW68CUdgcbCHPcgKy0_BxzKQhhM8ycbC
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
 jaxsim/parsers/rod/parser.py,sha256=4COuhkAYv4-GIpCqvkXEJWpDEQczEkBM3KwpqX48Rek,13514
 jaxsim/parsers/rod/utils.py,sha256=KSjgy6WsmTrD5HZEA2x8hOBSRU4bUGOOHzxKkeFO5r8,5721
-jaxsim/rbda/__init__.py,sha256=MqEZwzu8SHPAlIFHmSXmCjehuOJGRX58OrBVAbBVMwg,374
+jaxsim/rbda/__init__.py,sha256=H7DhXpxkPOi9lpUvg31IMHFfRafke1UoJLc5GQIdyhA,387
 jaxsim/rbda/aba.py,sha256=0OoCzHhf1v-qqr1y5PIrD7_mPwAlid0fjXxUrIa5E_s,9118
 jaxsim/rbda/collidable_points.py,sha256=4ZNJbEj2nEi15jBLR-GNbdaqKgkN58FBgqd_TXupEgg,4948
 jaxsim/rbda/crba.py,sha256=awsWEQXLE0UPEXIcZCVsAqBEPjyahMNzY9ux6nE1l-s,4739
 jaxsim/rbda/forward_kinematics.py,sha256=94W7TUXvZjMb-99CyYR8pObuxIYYX9B_dtRZqsNcThs,3418
 jaxsim/rbda/jacobian.py,sha256=M79bGir-2w_iJ2GurYhOGgMfJnp7ZMOCW6AeeWKK8iM,10745
 jaxsim/rbda/rnea.py,sha256=DjwkvXQVUSUclM3Uy3UPZ2tao91R5dGd4o7TsS2qObI,7650
-jaxsim/rbda/soft_contacts.py,sha256=0hx9JT4R1X2PPjhZ1EDizBR1gGoCFCtKYu86SeuIvvA,11269
 jaxsim/rbda/utils.py,sha256=zpbFM2Iq8cntku0BFVu9nfEqZhInCWi9D2INT6MFEI8,5003
+jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jaxsim/rbda/contacts/common.py,sha256=iMKLP30Qft9eGTiHo2iY-UoACJjg1JphA9_pW8wRdjc,2410
+jaxsim/rbda/contacts/soft.py,sha256=TvjGrKFmk6IIip-D3WLcOr9hWjlmF11-ULPkAqJKTZY,15601
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=UXQCt7TCkq6GkM8bOZu44pNTpf-FZWiKN6VE4kb4kFk,2342
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=QIJitSoljrKR_U4T3ewCJPT3DTh-tPZsRsg0t_MH93E,3896
-jaxsim-0.3.1.dev17.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.3.1.dev17.dist-info/METADATA,sha256=zRsMl96hDJt919NgrEuxkhye1S8X20bi_nWdPzJiptU,9739
-jaxsim-0.3.1.dev17.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.3.1.dev17.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.3.1.dev17.dist-info/RECORD,,
+jaxsim-0.3.1.dev40.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.3.1.dev40.dist-info/METADATA,sha256=ds7zFF0BWvBs8TLnHcQ54-diuMpka_Y5NKlF-vDK2nA,9739
+jaxsim-0.3.1.dev40.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.3.1.dev40.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.3.1.dev40.dist-info/RECORD,,