jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/api/common.py

@@ -1,17 +1,18 @@
 import abc
 import contextlib
 import dataclasses
+import enum
 import functools
-from typing import ContextManager
+from collections.abc import Callable, Iterator
+from typing import ParamSpec, TypeVar
 
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
-from jaxsim.high_level.common import VelRepr
+from jaxsim.math import Adjoint
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 try:
@@ -20,6 +21,32 @@ except ImportError:
     from typing_extensions import Self
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
+    """Apply a JAX named scope to a function for improved profiling and clarity."""
+
+    @functools.wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        with jax.named_scope(name or fn.__name__):
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
+@enum.unique
+class VelRepr(enum.IntEnum):
+    """
+    Enumeration of all supported 6D velocity representations.
+    """
+
+    Body = enum.auto()
+    Mixed = enum.auto()
+    Inertial = enum.auto()
+
+
 @jax_dataclasses.pytree_dataclass
 class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
     """
@@ -33,7 +60,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
     @contextlib.contextmanager
     def switch_velocity_representation(
         self, velocity_representation: VelRepr
-    ) -> ContextManager[Self]:
+    ) -> Iterator[Self]:
         """
         Context manager to temporarily switch the velocity representation.
 
@@ -48,7 +75,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,
@@ -59,7 +86,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             # We run this in a mutable context with restoration so that any exception
             # occurring, we restore the original object in case it was modified.
             with self.mutable_context(
-                mutability=self._mutability(), restore_after_exception=True
+                mutability=self.mutability(), restore_after_exception=True
             ):
                 yield self
 
@@ -76,16 +103,17 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
         array: jtp.Array,
         other_representation: VelRepr,
         transform: jtp.Matrix,
-        is_force: bool = False,
+        *,
+        is_force: bool,
     ) -> jtp.Array:
-        """
+        r"""
         Convert a 6D quantity from inertial-fixed to another representation.
 
         Args:
             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.
 
@@ -110,11 +138,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
@@ -124,11 +152,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
@@ -142,9 +170,10 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
         array: jtp.Array,
         other_representation: VelRepr,
         transform: jtp.Matrix,
-        is_force: bool = False,
+        *,
+        is_force: bool,
     ) -> jtp.Array:
-        """
+        r"""
         Convert a 6D quantity from another representation to inertial-fixed.
 
         Args:
@@ -177,11 +206,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
@@ -192,11 +221,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