jaxsim 0.3.1.dev64__py3-none-any.whl → 0.3.1.dev94__py3-none-any.whl

jaxsim/__init__.py

@@ -33,17 +33,17 @@ def _is_editable() -> bool:
     import pathlib
     import site
 
-    # Get the ModuleSpec of jaxsim
+    # Get the ModuleSpec of jaxsim.
     jaxsim_spec = importlib.util.find_spec(name="jaxsim")
 
     # This can be None. If it's None, assume non-editable installation.
     if jaxsim_spec.origin is None:
         return False
 
-    # Get the folder containing the jaxsim package
+    # Get the folder containing the jaxsim package.
     jaxsim_package_dir = str(pathlib.Path(jaxsim_spec.origin).parent.parent)
 
-    # The installation is editable if the package dir is not in any {site|dist}-packages
+    # The installation is editable if the package dir is not in any {site|dist}-packages.
     return jaxsim_package_dir not in site.getsitepackages()
 
 
@@ -82,10 +82,10 @@ def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
 logging.configure(level=_get_default_logging_level(env_var="JAXSIM_LOGGING_LEVEL"))
 
 
-# Configure JAX
+# Configure JAX.
 _jnp_options()
 
-# Initialize the numpy print options
+# Initialize the numpy print options.
 _np_options()
 
 del _jnp_options

jaxsim/_version.py

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

jaxsim/api/com.py

@@ -1,6 +1,5 @@
 import jax
 import jax.numpy as jnp
-import jaxlie
 
 import jaxsim.api as js
 import jaxsim.math
@@ -28,7 +27,7 @@ def com_position(
 
     W_H_L = js.model.forward_kinematics(model=model, data=data)
     W_H_B = data.base_transform()
-    B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
+    B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
 
     def B_p̃_LCoM(i) -> jtp.Vector:
         m = js.link.mass(model=model, link_index=i)
@@ -179,9 +178,9 @@ def locked_centroidal_spatial_inertia(
         case _:
             raise ValueError(data.velocity_representation)
 
-    B_H_G = jaxlie.SE3.from_matrix(jaxsim.math.Transform.inverse(W_H_B) @ W_H_G)
+    B_H_G = jaxsim.math.Transform.inverse(W_H_B) @ W_H_G
 
-    B_Xv_G = B_H_G.adjoint()
+    B_Xv_G = jaxsim.math.Adjoint.from_transform(transform=B_H_G)
     G_Xf_B = B_Xv_G.transpose()
 
     return G_Xf_B @ B_Mbb_B @ B_Xv_G

jaxsim/api/common.py

@@ -8,10 +8,10 @@ from typing import ContextManager
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
+from jaxsim.math import Adjoint
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 try:
@@ -59,7 +59,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
 
         try:
 
-            # First, we replace the velocity representation
+            # First, we replace the velocity representation.
             with self.mutable_context(
                 mutability=Mutability.MUTABLE_NO_VALIDATION,
                 restore_after_exception=True,
@@ -97,7 +97,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             array: The 6D quantity to convert.
             other_representation: The representation to convert to.
             transform:
-                The `math:W \mathbf{H}_O` transform, where `math:O` is the
+                The :math:`W \mathbf{H}_O` transform, where :math:`O` is the
                 reference frame of the other representation.
             is_force: Whether the quantity is a 6D force or a 6D velocity.
 
@@ -122,11 +122,11 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             case VelRepr.Body:
 
                 if not is_force:
-                    O_Xv_W = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint()
+                    O_Xv_W = Adjoint.from_transform(transform=W_H_O, inverse=True)
                     O_array = O_Xv_W @ W_array
 
                 else:
-                    O_Xf_W = jaxlie.SE3.from_matrix(W_H_O).adjoint().T
+                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).T
                     O_array = O_Xf_W @ W_array
 
                 return O_array
@@ -136,11 +136,11 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
                 W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
 
                 if not is_force:
-                    OW_Xv_W = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint()
+                    OW_Xv_W = Adjoint.from_transform(transform=W_H_OW, inverse=True)
                     OW_array = OW_Xv_W @ W_array
 
                 else:
-                    OW_Xf_W = jaxlie.SE3.from_matrix(W_H_OW).adjoint().transpose()
+                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).T
                     OW_array = OW_Xf_W @ W_array
 
                 return OW_array
@@ -190,11 +190,11 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
                 O_array = array
 
                 if not is_force:
-                    W_Xv_O: jtp.Array = jaxlie.SE3.from_matrix(W_H_O).adjoint()
+                    W_Xv_O: jtp.Array = Adjoint.from_transform(W_H_O)
                     W_array = W_Xv_O @ O_array
 
                 else:
-                    W_Xf_O = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint().T
+                    W_Xf_O = Adjoint.from_transform(transform=W_H_O, inverse=True).T
                     W_array = W_Xf_O @ O_array
 
                 return W_array
@@ -205,11 +205,11 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
                 W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
 
                 if not is_force:
-                    W_Xv_BW: jtp.Array = jaxlie.SE3.from_matrix(W_H_OW).adjoint()
+                    W_Xv_BW: jtp.Array = Adjoint.from_transform(W_H_OW)
                     W_array = W_Xv_BW @ BW_array
 
                 else:
-                    W_Xf_BW = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint().T
+                    W_Xf_BW = Adjoint.from_transform(transform=W_H_OW, inverse=True).T
                     W_array = W_Xf_BW @ BW_array
 
                 return W_array

jaxsim/api/contact.py

@@ -8,7 +8,7 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsParams
+from jaxsim.rbda.contacts.soft import SoftContactsParams
 
 from .common import VelRepr
 
@@ -137,9 +137,17 @@ def collidable_point_dynamics(
     # all collidable points belonging to the robot.
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
+    # Import privately the soft contacts classes.
+    from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
+
     # Build the soft contact model.
     match model.contact_model:
-        case s if isinstance(s, SoftContacts):
+
+        case SoftContacts():
+
+            assert isinstance(model.contact_model, SoftContacts)
+            assert isinstance(data.state.contact, SoftContactsState)
+
             # Build the contact model.
             soft_contacts = SoftContacts(
                 parameters=data.contacts_params, terrain=model.terrain
@@ -337,7 +345,7 @@ def jacobian(
             The output velocity representation of the free-floating jacobian.
 
     Returns:
-        The stacked 6×(6+n) free-floating jacobians of the frames associated to the
+        The stacked :math:`6 \times (6+n)` free-floating jacobians of the frames associated to the
         collidable points.
 
     Note:

jaxsim/api/data.py

@@ -8,7 +8,6 @@ import jax
 import jax.numpy as jnp
 import jax_dataclasses
 import jaxlie
-import numpy as np
 
 import jaxsim.api as js
 import jaxsim.rbda
@@ -390,7 +389,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         ).astype(float)
 
     @jax.jit
-    def base_transform(self) -> jtp.MatrixJax:
+    def base_transform(self) -> jtp.Matrix:
         """
         Get the base transform.
 
@@ -625,9 +624,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_p_B = base_pose[0:3, 3]
 
-        to_wxyz = np.array([3, 0, 1, 2])
-        W_R_B: jaxlie.SO3 = jaxlie.SO3.from_matrix(base_pose[0:3, 0:3])  # noqa
-        W_Q_B = W_R_B.as_quaternion_xyzw()[to_wxyz]
+        W_Q_B = jaxsim.math.Quaternion.from_dcm(dcm=base_pose[0:3, 0:3])
 
         return self.reset_base_position(base_position=W_p_B).reset_base_quaternion(
             base_quaternion=W_Q_B
@@ -815,7 +812,7 @@ def random_model_data(
 
         physics_model_state.base_quaternion = jaxlie.SO3.from_rpy_radians(
             *jax.random.uniform(key=k2, shape=(3,), minval=0, maxval=2 * jnp.pi)
-        ).as_quaternion_xyzw()[np.array([3, 0, 1, 2])]
+        ).wxyz
 
         if model.number_of_joints() > 0:
             physics_model_state.joint_positions = js.joint.random_joint_positions(

jaxsim/api/frame.py

@@ -3,12 +3,11 @@ from typing import Sequence
 
 import jax
 import jax.numpy as jnp
-import jaxlie
 
 import jaxsim.api as js
-import jaxsim.math
 import jaxsim.typing as jtp
 from jaxsim import exceptions
+from jaxsim.math import Adjoint, Transform
 
 from .common import VelRepr
 
@@ -189,7 +188,7 @@ def jacobian(
     frame_index: jtp.IntLike,
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Matrix:
-    """
+    r"""
     Compute the free-floating jacobian of the frame.
 
     Args:
@@ -200,7 +199,7 @@ def jacobian(
             The output velocity representation of the free-floating jacobian.
 
     Returns:
-        The 6×(6+n) free-floating jacobian of the frame.
+        The :math:`6 \times (6+n)` free-floating jacobian of the frame.
 
     Note:
         The input representation of the free-floating jacobian is the active
@@ -228,29 +227,29 @@ def jacobian(
         model=model, data=data, link_index=L, output_vel_repr=VelRepr.Body
     )
 
-    # Adjust the output representation
+    # Adjust the output representation.
     match output_vel_repr:
         case VelRepr.Inertial:
             W_H_L = js.link.transform(model=model, data=data, link_index=L)
-            W_X_L = jaxlie.SE3.from_matrix(W_H_L).adjoint()
+            W_X_L = Adjoint.from_transform(transform=W_H_L)
             W_J_WL = W_X_L @ L_J_WL
             O_J_WL_I = W_J_WL
 
         case VelRepr.Body:
             W_H_L = js.link.transform(model=model, data=data, link_index=L)
             W_H_F = transform(model=model, data=data, frame_index=frame_index)
-            F_H_L = jaxsim.math.Transform.inverse(W_H_F) @ W_H_L
-            F_X_L = jaxlie.SE3.from_matrix(F_H_L).adjoint()
+            F_H_L = Transform.inverse(W_H_F) @ W_H_L
+            F_X_L = Adjoint.from_transform(transform=F_H_L)
             F_J_WL = F_X_L @ L_J_WL
             O_J_WL_I = F_J_WL
 
         case VelRepr.Mixed:
             W_H_L = js.link.transform(model=model, data=data, link_index=L)
             W_H_F = transform(model=model, data=data, frame_index=frame_index)
-            F_H_L = jaxsim.math.Transform.inverse(W_H_F) @ W_H_L
+            F_H_L = Transform.inverse(W_H_F) @ W_H_L
             FW_H_F = W_H_F.at[0:3, 3].set(jnp.zeros(3))
             FW_H_L = FW_H_F @ F_H_L
-            FW_X_L = jaxlie.SE3.from_matrix(FW_H_L).adjoint()
+            FW_X_L = Adjoint.from_transform(transform=FW_H_L)
             FW_J_WL = FW_X_L @ L_J_WL
             O_J_WL_I = FW_J_WL
 

jaxsim/api/kin_dyn_parameters.py

@@ -5,11 +5,10 @@ import dataclasses
 import jax.lax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
-from jaxsim.math import Inertia, JointModel, supported_joint_motion
+from jaxsim.math import Adjoint, Inertia, JointModel, supported_joint_motion
 from jaxsim.parsers.descriptions import JointDescription, ModelDescription
 from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
@@ -168,7 +167,7 @@ class KynDynParameters(JaxsimDataclass):
             for link in ordered_links
             if link.parent is not None
         }
-        parent_array = jnp.array([-1] + list(parent_array_dict.values()), dtype=int)
+        parent_array = jnp.array([-1, *list(parent_array_dict.values())], dtype=int)
 
         # Instead of building the support parent array κ(i) for each link of the model,
         # that has a variable length depending on the number of links connecting the
@@ -432,11 +431,9 @@ class KynDynParameters(JaxsimDataclass):
         # Compute the overall transforms from the parent to the child of each joint by
         # composing all the components of our joint model.
         i_X_λ = jax.vmap(
-            lambda λ_Hi_pre, pre_Hi_suc, suc_Hi_i: jaxlie.SE3.from_matrix(
-                λ_Hi_pre @ pre_Hi_suc @ suc_Hi_i
+            lambda λ_Hi_pre, pre_Hi_suc, suc_Hi_i: Adjoint.from_transform(
+                transform=λ_Hi_pre @ pre_Hi_suc @ suc_Hi_i, inverse=True
             )
-            .inverse()
-            .adjoint()
         )(λ_H_pre, pre_H_suc, suc_H_i)
 
         return i_X_λ, S
@@ -466,12 +463,12 @@ class KynDynParameters(JaxsimDataclass):
     def set_link_inertia(
         self, link_index: int, inertia: jtp.MatrixLike
     ) -> KynDynParameters:
-        """
+        r"""
         Set the inertia tensor of a link.
 
         Args:
             link_index: The index of the link.
-            inertia: The 3×3 inertia tensor of the link.
+            inertia: The :math:`3 \times 3` inertia tensor of the link.
 
         Returns:
             The updated kinematic and dynamic parameters of the model.
@@ -569,7 +566,7 @@ class LinkParameters(JaxsimDataclass):
         index: The index of the link.
         mass: The mass of the link.
         inertia_elements:
-            The unique elements of the 3×3 inertia tensor of the link.
+            The unique elements of the :math:`3 \times 3` inertia tensor of the link.
         center_of_mass:
             The translation :math:`{}^L \mathbf{p}_{\text{CoM}}` between the origin
             of the link frame and the link's center of mass, expressed in the
@@ -588,12 +585,12 @@ class LinkParameters(JaxsimDataclass):
 
     @staticmethod
     def build_from_spatial_inertia(index: jtp.IntLike, M: jtp.Matrix) -> LinkParameters:
-        """
-        Build a LinkParameters object from a 6×6 spatial inertia matrix.
+        r"""
+        Build a LinkParameters object from a :math:`6 \times 6` spatial inertia matrix.
 
         Args:
             index: The index of the link.
-            M: The 6×6 spatial inertia matrix of the link.
+            M: The :math:`6 \times 6` spatial inertia matrix of the link.
 
         Returns:
             The LinkParameters object.
@@ -616,13 +613,13 @@ class LinkParameters(JaxsimDataclass):
     def build_from_inertial_parameters(
         index: jtp.IntLike, m: jtp.FloatLike, I: jtp.MatrixLike, c: jtp.VectorLike
     ) -> LinkParameters:
-        """
+        r"""
         Build a LinkParameters object from the inertial parameters of a link.
 
         Args:
             index: The index of the link.
             m: The mass of the link.
-            I: The 3×3 inertia tensor of the link.
+            I: The :math:`3 \times 3` inertia tensor of the link.
             c: The translation between the link frame and the link's center of mass.
 
         Returns:
@@ -676,14 +673,14 @@ class LinkParameters(JaxsimDataclass):
 
     @staticmethod
     def inertia_tensor(params: LinkParameters) -> jtp.Matrix:
-        """
-        Return the 3×3 inertia tensor of a link.
+        r"""
+        Return the :math:`3 \times 3` inertia tensor of a link.
 
         Args:
             params: The link parameters.
 
         Returns:
-            The 3×3 inertia tensor of the link.
+            The :math:`3 \times 3` inertia tensor of the link.
         """
 
         return LinkParameters.unflatten_inertia_tensor(
@@ -692,14 +689,14 @@ class LinkParameters(JaxsimDataclass):
 
     @staticmethod
     def spatial_inertia(params: LinkParameters) -> jtp.Matrix:
-        """
-        Return the 6×6 spatial inertia matrix of a link.
+        r"""
+        Return the :math:`6 \times 6` spatial inertia matrix of a link.
 
         Args:
             params: The link parameters.
 
         Returns:
-            The 6×6 spatial inertia matrix of the link.
+            The :math:`6 \times 6` spatial inertia matrix of the link.
         """
 
         return Inertia.to_sixd(
@@ -710,11 +707,11 @@ class LinkParameters(JaxsimDataclass):
 
     @staticmethod
     def flatten_inertia_tensor(I: jtp.Matrix) -> jtp.Vector:
-        """
-        Flatten a 3×3 inertia tensor into a vector of unique elements.
+        r"""
+        Flatten a :math:`3 \times 3` inertia tensor into a vector of unique elements.
 
         Args:
-            I: The 3×3 inertia tensor.
+            I: The :math:`3 \times 3` inertia tensor.
 
         Returns:
             The vector of unique elements of the inertia tensor.
@@ -724,14 +721,14 @@ class LinkParameters(JaxsimDataclass):
 
     @staticmethod
     def unflatten_inertia_tensor(inertia_elements: jtp.Vector) -> jtp.Matrix:
-        """
-        Unflatten a vector of unique elements into a 3×3 inertia tensor.
+        r"""
+        Unflatten a vector of unique elements into a :math:`3 \times 3` inertia tensor.
 
         Args:
             inertia_elements: The vector of unique elements of the inertia tensor.
 
         Returns:
-            The 3×3 inertia tensor.
+            The :math:`3 \times 3` inertia tensor.
         """
 
         I = jnp.zeros([3, 3]).at[jnp.triu_indices(3)].set(inertia_elements.squeeze())
@@ -792,7 +789,7 @@ class ContactParameters(JaxsimDataclass):
         )
 
         # Build the ContactParameters object.
-        cp = ContactParameters(point=points, body=link_index_of_points)  # noqa
+        cp = ContactParameters(point=points, body=link_index_of_points)
 
         assert cp.point.shape[1] == 3, cp.point.shape[1]
         assert cp.point.shape[0] == len(cp.body), cp.point.shape[0]

jaxsim/api/link.py

@@ -4,12 +4,12 @@ from typing import Sequence
 import jax
 import jax.numpy as jnp
 import jax.scipy.linalg
-import jaxlie
 
 import jaxsim.api as js
 import jaxsim.rbda
 import jaxsim.typing as jtp
 from jaxsim import exceptions
+from jaxsim.math import Adjoint
 
 from .common import VelRepr
 
@@ -134,7 +134,7 @@ def mass(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> jtp.Float:
 def spatial_inertia(
     model: js.model.JaxSimModel, *, link_index: jtp.IntLike
 ) -> jtp.Matrix:
-    """
+    r"""
     Compute the 6D spatial inertial of the link.
 
     Args:
@@ -142,7 +142,7 @@ def spatial_inertia(
         link_index: The index of the link.
 
     Returns:
-        The 6×6 matrix representing the spatial inertia of the link expressed in
+        The :math:`6 \times 6` matrix representing the spatial inertia of the link expressed in
         the link frame (body-fixed representation).
     """
 
@@ -243,7 +243,7 @@ def jacobian(
     link_index: jtp.IntLike,
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Matrix:
-    """
+    r"""
     Compute the free-floating jacobian of the link.
 
     Args:
@@ -254,7 +254,7 @@ def jacobian(
             The output velocity representation of the free-floating jacobian.
 
     Returns:
-        The 6×(6+n) free-floating jacobian of the link.
+        The :math:`6 \times (6+n)` free-floating jacobian of the link.
 
     Note:
         The input representation of the free-floating jacobian is the active
@@ -287,7 +287,7 @@ def jacobian(
     match data.velocity_representation:
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
-            B_X_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
+            B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
             B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(
                 B_X_W, jnp.eye(model.dofs())
             )
@@ -298,7 +298,7 @@ def jacobian(
         case VelRepr.Mixed:
             W_R_B = data.base_orientation(dcm=True)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
-            B_X_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+            B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
             B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(
                 B_X_BW, jnp.eye(model.dofs())
             )
@@ -312,11 +312,11 @@ def jacobian(
     match output_vel_repr:
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
-            W_X_B = jaxlie.SE3.from_matrix(W_H_B).adjoint()
+            W_X_B = Adjoint.from_transform(transform=W_H_B)
             O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I
 
         case VelRepr.Body:
-            L_X_B = jaxlie.SE3.from_matrix(B_H_L).inverse().adjoint()
+            L_X_B = Adjoint.from_transform(transform=B_H_L, inverse=True)
             L_J_WL_I = L_X_B @ B_J_WL_I
             O_J_WL_I = L_J_WL_I
 
@@ -325,7 +325,7 @@ def jacobian(
             W_H_L = W_H_B @ B_H_L
             LW_H_L = W_H_L.at[0:3, 3].set(jnp.zeros(3))
             LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
-            LW_X_B = jaxlie.SE3.from_matrix(LW_H_B).adjoint()
+            LW_X_B = Adjoint.from_transform(transform=LW_H_B)
             LW_J_WL_I = LW_X_B @ B_J_WL_I
             O_J_WL_I = LW_J_WL_I
 
@@ -393,7 +393,7 @@ def jacobian_derivative(
     link_index: jtp.IntLike,
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Matrix:
-    """
+    r"""
     Compute the derivative of the free-floating jacobian of the link.
 
     Args:
@@ -404,7 +404,7 @@ def jacobian_derivative(
             The output velocity representation of the free-floating jacobian derivative.
 
     Returns:
-        The derivative of the 6×(6+n) free-floating jacobian of the link.
+        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the link.
 
     Note:
         The input representation of the free-floating jacobian derivative is the active