jaxsim 0.3.1.dev21__tar.gz → 0.3.1.dev46__tar.gz

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/.pre-commit-config.yaml

@@ -24,3 +24,8 @@ repos:
     rev: v0.3.2
     hooks:
       - id: ruff
+
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.7.1
+    hooks:
+      - id: nbstripout

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/PKG-INFO

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

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/docs/modules/rbda.rst

@@ -13,7 +13,6 @@ This module provides a set of algorithms for rigid body dynamics.
     crba
     forward_kinematics
     jacobian
-    soft_contacts
     utils
 
 Articulated Body Algorithm
@@ -28,6 +27,12 @@ Collision Detection
 .. automodule:: jaxsim.rbda.collidable_points
     :members:
 
+Contact Models
+~~~~~~~~~~~~~~
+
+.. automodule:: jaxsim.rbda.contacts.soft
+    :members:
+
 Composite Rigid Body Algorithm
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -46,12 +51,6 @@ Jacobians
 .. automodule:: jaxsim.rbda.jacobian
     :members:
 
-Soft Contacts
-~~~~~~~~~~~~~
-
-.. automodule:: jaxsim.rbda.soft_contacts
-    :members:
-
 Utilities
 ~~~~~~~~~
 

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/jaxsim/_version.py

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

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/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-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/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-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/jaxsim/api/frame.py

@@ -4,11 +4,11 @@ from typing import Sequence
 import jax
 import jax.numpy as jnp
 import jaxlie
-import numpy as np
 
 import jaxsim.api as js
 import jaxsim.math
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 
 from .common import VelRepr
 
@@ -17,22 +17,32 @@ from .common import VelRepr
 # =======================
 
 
+@jax.jit
 def idx_of_parent_link(
-    model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike
+    model: js.model.JaxSimModel, *, frame_index: jtp.IntLike
 ) -> jtp.Int:
     """
     Get the index of the link to which the frame is rigidly attached.
 
     Args:
         model: The model to consider.
-        frame_idx: The index of the frame.
+        frame_index: The index of the frame.
 
     Returns:
         The index of the frame's parent link.
     """
 
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
     return model.kin_dyn_parameters.frame_parameters.body[
-        frame_idx - model.number_of_links()
+        frame_index - model.number_of_links()
     ]
 
 
@@ -49,19 +59,17 @@ def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
         The index of the frame.
     """
 
-    if frame_name in model.kin_dyn_parameters.frame_parameters.name:
-        return (
-            jnp.array(
-                np.argwhere(
-                    np.array(model.kin_dyn_parameters.frame_parameters.name)
-                    == frame_name
-                )
-            )
-            .squeeze()
-            .astype(int)
-        ) + model.number_of_links()
+    if frame_name not in model.kin_dyn_parameters.frame_parameters.name:
+        raise ValueError(f"Frame '{frame_name}' not found in the model.")
 
-    return jnp.array(-1).astype(int)
+    return (
+        jnp.array(
+            model.number_of_links()
+            + model.kin_dyn_parameters.frame_parameters.name.index(frame_name)
+        )
+        .astype(int)
+        .squeeze()
+    )
 
 
 def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
@@ -76,6 +84,15 @@ def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str
         The name of the frame.
     """
 
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
     return model.kin_dyn_parameters.frame_parameters.name[
         frame_index - model.number_of_links()
     ]
@@ -142,8 +159,17 @@ def transform(
         The 4x4 matrix representing the transform.
     """
 
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
     # Compute the necessary transforms.
-    L = idx_of_parent_link(model=model, frame_idx=frame_index)
+    L = idx_of_parent_link(model=model, frame_index=frame_index)
     W_H_L = js.link.transform(model=model, data=data, link_index=L)
 
     # Get the static frame pose wrt the parent link.
@@ -181,12 +207,21 @@ def jacobian(
         velocity representation.
     """
 
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
     output_vel_repr = (
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
     # Get the index of the parent link.
-    L = idx_of_parent_link(model=model, frame_idx=frame_index)
+    L = idx_of_parent_link(model=model, frame_index=frame_index)
 
     # Compute the Jacobian of the parent link using body-fixed output representation.
     L_J_WL = js.link.jacobian(

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/jaxsim/api/joint.py

@@ -6,6 +6,7 @@ import jax.numpy as jnp
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 
 # =======================
 # Index-related functions
@@ -25,14 +26,18 @@ def name_to_idx(model: js.model.JaxSimModel, *, joint_name: str) -> jtp.Int:
         The index of the joint.
     """
 
-    if joint_name in model.kin_dyn_parameters.joint_model.joint_names:
-        # Note: the index of the joint for RBDAs starts from 1, but
-        # the index for accessing the right element starts from 0.
-        # Therefore, there is a -1.
-        return jnp.array(
+    if joint_name not in model.joint_names():
+        raise ValueError(f"Joint '{joint_name}' not found in the model.")
+
+    # Note: the index of the joint for RBDAs starts from 1, but the index for
+    # accessing the right element starts from 0. Therefore, there is a -1.
+    return (
+        jnp.array(
             model.kin_dyn_parameters.joint_model.joint_names.index(joint_name) - 1
-        ).squeeze()
-    return jnp.array(-1).astype(int)
+        )
+        .astype(int)
+        .squeeze()
+    )
 
 
 def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str:
@@ -47,6 +52,14 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
         The name of the joint.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [joint_index < 0, joint_index >= model.number_of_joints()]
+        ).any(),
+        msg="Invalid joint index '{idx}'",
+        idx=joint_index,
+    )
+
     return model.kin_dyn_parameters.joint_model.joint_names[joint_index + 1]
 
 
@@ -112,6 +125,14 @@ def position_limit(
     if model.number_of_joints() <= 1:
         return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [joint_index < 0, joint_index >= model.number_of_joints()]
+        ).any(),
+        msg="Invalid joint index '{idx}'",
+        idx=joint_index,
+    )
+
     s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min[joint_index]
     s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max[joint_index]
 

{jaxsim-0.3.1.dev21 → jaxsim-0.3.1.dev46}/src/jaxsim/api/link.py

@@ -5,11 +5,11 @@ import jax
 import jax.numpy as jnp
 import jax.scipy.linalg
 import jaxlie
-import numpy as np
 
 import jaxsim.api as js
 import jaxsim.rbda
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 
 from .common import VelRepr
 
@@ -31,15 +31,14 @@ def name_to_idx(model: js.model.JaxSimModel, *, link_name: str) -> jtp.Int:
         The index of the link.
     """
 
-    if link_name in model.kin_dyn_parameters.link_names:
-        return (
-            jnp.array(
-                np.argwhere(np.array(model.kin_dyn_parameters.link_names) == link_name)
-            )
-            .squeeze()
-            .astype(int)
-        )
-    return jnp.array(-1).astype(int)
+    if link_name not in model.link_names():
+        raise ValueError(f"Link '{link_name}' not found in the model.")
+
+    return (
+        jnp.array(model.kin_dyn_parameters.link_names.index(link_name))
+        .astype(int)
+        .squeeze()
+    )
 
 
 def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
@@ -54,6 +53,14 @@ def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
         The name of the link.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     return model.kin_dyn_parameters.link_names[link_index]
 
 
@@ -112,6 +119,14 @@ def mass(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> jtp.Float:
         The mass of the link.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     return model.kin_dyn_parameters.link_parameters.mass[link_index].astype(float)
 
 
@@ -131,6 +146,14 @@ def spatial_inertia(
         the link frame (body-fixed representation).
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     link_parameters = jax.tree_util.tree_map(
         lambda l: l[link_index], model.kin_dyn_parameters.link_parameters
     )
@@ -157,6 +180,14 @@ def transform(
         The 4x4 matrix representing the transform.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     return js.model.forward_kinematics(model=model, data=data)[link_index]
 
 
@@ -230,6 +261,14 @@ def jacobian(
         velocity representation.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     output_vel_repr = (
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
@@ -318,6 +357,14 @@ def velocity(
         The 6D velocity of the link in the specified velocity representation.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     output_vel_repr = (
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
@@ -364,6 +411,14 @@ def jacobian_derivative(
         velocity representation.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     output_vel_repr = (
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
@@ -538,6 +593,14 @@ def bias_acceleration(
         The 6D bias acceleration of the link.
     """
 
+    exceptions.raise_value_error_if(
+        condition=jnp.array(
+            [link_index < 0, link_index >= model.number_of_links()]
+        ).any(),
+        msg="Invalid link index '{idx}'",
+        idx=link_index,
+    )
+
     # Compute the bias acceleration of all links in the active representation.
     O_v̇_WL = js.model.link_bias_accelerations(model=model, data=data)[link_index]
     return O_v̇_WL