jaxsim 0.3.1.dev4__py3-none-any.whl → 0.3.1.dev21__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.dev4'
-__version_tuple__ = version_tuple = (0, 3, 1, 'dev4')
+__version__ = version = '0.3.1.dev21'
+__version_tuple__ = version_tuple = (0, 3, 1, 'dev21')

jaxsim/api/data.py

@@ -45,12 +45,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
                 hash(self.state),
-                hash(tuple(self.gravity.flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.gravity),
                 hash(self.soft_contacts_params),
-                hash(jnp.atleast_1d(self.time_ns).flatten().tolist()),
+                HashedNumpyArray.hash_of_array(self.time_ns),
             )
         )
 

jaxsim/api/frame.py

@@ -17,7 +17,9 @@ from .common import VelRepr
 # =======================
 
 
-def idx_of_parent_link(model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike) -> int:
+def idx_of_parent_link(
+    model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike
+) -> jtp.Int:
     """
     Get the index of the link to which the frame is rigidly attached.
 
@@ -29,17 +31,13 @@ def idx_of_parent_link(model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike) -
         The index of the frame's parent link.
     """
 
-    # Get the intermediate representation parsed from the model description.
-    ir = model.description
+    return model.kin_dyn_parameters.frame_parameters.body[
+        frame_idx - model.number_of_links()
+    ]
 
-    # Extract the indices of the frame and the link it is attached to.
-    F = ir.frames[frame_idx - model.number_of_links()]
-    L = ir.links_dict[F.parent.name].index
 
-    return int(L)
-
-
-def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> int:
+@functools.partial(jax.jit, static_argnames="frame_name")
+def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
     """
     Convert the name of a frame to its index.
 
@@ -51,13 +49,19 @@ def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> int:
         The index of the frame.
     """
 
-    frame_names = np.array([frame.name for frame in model.description.frames])
+    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 in frame_names:
-        idx_in_list = np.argwhere(frame_names == frame_name)
-        return int(idx_in_list.squeeze().tolist()) + model.number_of_links()
-
-    return -1
+    return jnp.array(-1).astype(int)
 
 
 def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
@@ -72,7 +76,9 @@ def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str
         The name of the frame.
     """
 
-    return model.description.frames[frame_index - model.number_of_links()].name
+    return model.kin_dyn_parameters.frame_parameters.name[
+        frame_index - model.number_of_links()
+    ]
 
 
 @functools.partial(jax.jit, static_argnames=["frame_names"])
@@ -91,7 +97,7 @@ def names_to_idxs(
     """
 
     return jnp.array(
-        [name_to_idx(model=model, frame_name=frame_name) for frame_name in frame_names]
+        [name_to_idx(model=model, frame_name=name) for name in frame_names]
     ).astype(int)
 
 
@@ -109,10 +115,7 @@ def idxs_to_names(
         The names of the frames.
     """
 
-    return tuple(
-        idx_to_name(model=model, frame_index=frame_index)
-        for frame_index in frame_indices
-    )
+    return tuple(idx_to_name(model=model, frame_index=idx) for idx in frame_indices)
 
 
 # ==========
@@ -120,7 +123,7 @@ def idxs_to_names(
 # ==========
 
 
-@functools.partial(jax.jit, static_argnames=["frame_index"])
+@jax.jit
 def transform(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,
@@ -144,14 +147,15 @@ def transform(
     W_H_L = js.link.transform(model=model, data=data, link_index=L)
 
     # Get the static frame pose wrt the parent link.
-    frame = model.description.frames[frame_index - model.number_of_links()]
-    L_H_F = frame.pose
+    L_H_F = model.kin_dyn_parameters.frame_parameters.transform[
+        frame_index - model.number_of_links()
+    ]
 
     # Combine the transforms computing the frame pose.
     return W_H_L @ L_H_F
 
 
-@functools.partial(jax.jit, static_argnames=["frame_index", "output_vel_repr"])
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
     data: js.data.JaxSimModelData,

jaxsim/api/kin_dyn_parameters.py

@@ -6,7 +6,6 @@ import jax.lax
 import jax.numpy as jnp
 import jax_dataclasses
 import jaxlie
-import numpy as np
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
@@ -15,7 +14,7 @@ from jaxsim.parsers.descriptions import JointDescription, ModelDescription
 from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
 
-@jax_dataclasses.pytree_dataclass
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class KynDynParameters(JaxsimDataclass):
     r"""
     Class storing the kinematic and dynamic parameters of a model.
@@ -26,6 +25,7 @@ class KynDynParameters(JaxsimDataclass):
         support_body_array_bool:
             The boolean support parent array :math:`\kappa_{b}(i)` of the model.
         link_parameters: The parameters of the links.
+        frame_parameters: The parameters of the frames.
         contact_parameters: The parameters of the collidable points.
         joint_model: The joint model of the model.
         joint_parameters: The parameters of the joints.
@@ -42,6 +42,9 @@ class KynDynParameters(JaxsimDataclass):
     # Contacts
     contact_parameters: ContactParameters
 
+    # Frames
+    frame_parameters: FrameParameters
+
     # Joints
     joint_model: JointModel
     joint_parameters: JointParameters | None
@@ -141,6 +144,19 @@ class KynDynParameters(JaxsimDataclass):
             model_description=model_description
         )
 
+        # =================
+        # Frames properties
+        # =================
+
+        # Create the object storing the parameters of frames.
+        # Note that, contrarily to LinkParameters and JointsParameters, this object
+        # is not created with vmap. This is because the "name" attribute of the object
+        # must be Static for JIT-related reasons, and tree_map would not consider it
+        # as a leaf.
+        frame_parameters = FrameParameters.build_from(
+            model_description=model_description
+        )
+
         # ===============
         # Tree properties
         # ===============
@@ -206,6 +222,7 @@ class KynDynParameters(JaxsimDataclass):
             joint_model=joint_model,
             joint_parameters=joint_parameters,
             contact_parameters=contact_parameters,
+            frame_parameters=frame_parameters,
         )
 
     def __eq__(self, other: KynDynParameters) -> bool:
@@ -221,7 +238,8 @@ class KynDynParameters(JaxsimDataclass):
             (
                 hash(self.number_of_links()),
                 hash(self.number_of_joints()),
-                hash(tuple(np.atleast_1d(self.parent_array).flatten().tolist())),
+                hash(self.frame_parameters.name),
+                hash(tuple(self.frame_parameters.body.tolist())),
                 hash(self._parent_array),
                 hash(self._support_body_array_bool),
             )
@@ -730,7 +748,7 @@ class ContactParameters(JaxsimDataclass):
             A tuple of integers representing, for each collidable point, the index of
             the body (link) to which it is rigidly attached to.
         point:
-            The translation between the link frame and the collidable point, expressed
+            The translations between the link frame and the collidable point, expressed
             in the coordinates of the parent link frame.
 
     Note:
@@ -773,10 +791,75 @@ class ContactParameters(JaxsimDataclass):
             links_dict[cp.parent_link.name].index for cp in collidable_points
         )
 
-        # Build the GroundContact object.
+        # Build the ContactParameters object.
         cp = ContactParameters(point=points, body=link_index_of_points)  # noqa
 
-        assert cp.point.shape[1] == 3
-        assert cp.point.shape[0] == len(cp.body)
+        assert cp.point.shape[1] == 3, cp.point.shape[1]
+        assert cp.point.shape[0] == len(cp.body), cp.point.shape[0]
 
         return cp
+
+
+@jax_dataclasses.pytree_dataclass
+class FrameParameters(JaxsimDataclass):
+    """
+    Class storing the frame parameters of a model.
+
+    Attributes:
+        name: A tuple of strings defining the frame names.
+        body:
+            A vector of integers representing, for each frame, the index of
+            the body (link) to which it is rigidly attached to.
+        transform: The transforms of the frames w.r.t. their parent link.
+
+    Note:
+        Contrarily to LinkParameters and JointParameters, this class is not meant
+        to be created with vmap. This is because the `name` attribute must be `Static`.
+    """
+
+    name: Static[tuple[str, ...]] = dataclasses.field(default_factory=tuple)
+
+    body: jtp.Vector = dataclasses.field(default_factory=lambda: jnp.array([]))
+
+    transform: jtp.Array = dataclasses.field(default_factory=lambda: jnp.array([]))
+
+    @staticmethod
+    def build_from(model_description: ModelDescription) -> FrameParameters:
+        """
+        Build a FrameParameters object from a model description.
+
+        Args:
+            model_description: The model description to consider.
+
+        Returns:
+            The FrameParameters object.
+        """
+
+        if len(model_description.frames) == 0:
+            return FrameParameters()
+
+        # Extract the frame names.
+        names = tuple(frame.name for frame in model_description.frames)
+
+        # For each frame, extract the index of the link to which it is attached to.
+        parent_link_index_of_frames = tuple(
+            model_description.links_dict[frame.parent.name].index
+            for frame in model_description.frames
+        )
+
+        # For each frame, extract the transform w.r.t. its parent link.
+        transforms = jnp.atleast_3d(
+            jnp.stack([frame.pose for frame in model_description.frames])
+        )
+
+        # Build the FrameParameters object.
+        fp = FrameParameters(
+            name=names,
+            transform=transforms.astype(float),
+            body=jnp.array(parent_link_index_of_frames).astype(int),
+        )
+
+        assert fp.transform.shape[1:] == (4, 4), fp.transform.shape[1:]
+        assert fp.transform.shape[0] == len(fp.body), fp.transform.shape[0]
+
+        return fp

jaxsim/api/model.py

@@ -17,12 +17,12 @@ import jaxsim.api as js
 import jaxsim.parsers.descriptions
 import jaxsim.typing as jtp
 from jaxsim.math import Cross
-from jaxsim.utils import JaxsimDataclass, Mutability
+from jaxsim.utils import JaxsimDataclass, Mutability, wrappers
 
 from .common import VelRepr
 
 
-@jax_dataclasses.pytree_dataclass
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class JaxSimModel(JaxsimDataclass):
     """
     The JaxSim model defining the kinematics and dynamics of a robot.
@@ -31,34 +31,43 @@ class JaxSimModel(JaxsimDataclass):
     model_name: Static[str]
 
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
-        default=jaxsim.terrain.FlatTerrain(), repr=False, compare=False, hash=False
+        default=jaxsim.terrain.FlatTerrain(), repr=False
     )
 
     kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
-        dataclasses.field(default=None, repr=False, compare=False, hash=False)
+        dataclasses.field(default=None, repr=False)
     )
 
     built_from: Static[str | pathlib.Path | rod.Model | None] = dataclasses.field(
-        default=None, repr=False, compare=False, hash=False
+        default=None, repr=False
     )
 
-    description: Static[jaxsim.parsers.descriptions.ModelDescription | None] = (
-        dataclasses.field(default=None, repr=False, compare=False, hash=False)
-    )
+    _description: Static[
+        wrappers.HashlessObject[jaxsim.parsers.descriptions.ModelDescription | None]
+    ] = dataclasses.field(default=None, repr=False)
+
+    @property
+    def description(self) -> jaxsim.parsers.descriptions.ModelDescription:
+        return self._description.get()
 
     def __eq__(self, other: JaxSimModel) -> bool:
 
         if not isinstance(other, JaxSimModel):
             return False
 
-        return hash(self) == hash(other)
+        if self.model_name != other.model_name:
+            return False
+
+        if self.kin_dyn_parameters != other.kin_dyn_parameters:
+            return False
+
+        return True
 
     def __hash__(self) -> int:
 
         return hash(
             (
                 hash(self.model_name),
-                hash(self.description),
                 hash(self.kin_dyn_parameters),
             )
         )
@@ -152,10 +161,10 @@ class JaxSimModel(JaxsimDataclass):
         # 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
+        # Build the model.
         model = JaxSimModel(
             model_name=model_name,
-            description=model_description,
+            _description=wrappers.HashlessObject(obj=model_description),
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
@@ -270,6 +279,10 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.link_names
 
+    # =====================
+    # Frame-related methods
+    # =====================
+
     def frame_names(self) -> tuple[str, ...]:
         """
         Return the names of the links in the model.
@@ -278,7 +291,7 @@ class JaxSimModel(JaxsimDataclass):
             The names of the links in the model.
         """
 
-        return tuple(frame.name for frame in self.description.frames)
+        return self.kin_dyn_parameters.frame_parameters.name
 
 
 # =====================

jaxsim/api/ode_data.py

@@ -283,12 +283,16 @@ class PhysicsModelState(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
-                hash(tuple(jnp.atleast_1d(self.joint_positions.flatten().tolist()))),
-                hash(tuple(jnp.atleast_1d(self.joint_velocities.flatten().tolist()))),
-                hash(tuple(self.base_position.flatten().tolist())),
-                hash(tuple(self.base_quaternion.flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.joint_positions),
+                HashedNumpyArray.hash_of_array(self.joint_velocities),
+                HashedNumpyArray.hash_of_array(self.base_position),
+                HashedNumpyArray.hash_of_array(self.base_quaternion),
+                HashedNumpyArray.hash_of_array(self.base_linear_velocity),
+                HashedNumpyArray.hash_of_array(self.base_angular_velocity),
             )
         )
 
@@ -613,9 +617,9 @@ class SoftContactsState(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
-        return hash(
-            tuple(jnp.atleast_1d(self.tangential_deformation.flatten()).tolist())
-        )
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return HashedNumpyArray.hash_of_array(self.tangential_deformation)
 
     def __eq__(self, other: SoftContactsState) -> bool:
 

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/parsers/descriptions/joint.py

@@ -41,7 +41,7 @@ class JointGenericAxis:
         return hash(self) == hash(other)
 
 
-@jax_dataclasses.pytree_dataclass
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class JointDescription(JaxsimDataclass):
     """
     In-memory description of a robot link.
@@ -95,25 +95,59 @@ class JointDescription(JaxsimDataclass):
                 norm_of_axis = np.linalg.norm(self.axis)
                 self.axis = self.axis / norm_of_axis
 
+    def __eq__(self, other: JointDescription) -> bool:
+
+        if not isinstance(other, JointDescription):
+            return False
+
+        if not (
+            self.name == other.name
+            and self.jtype == other.jtype
+            and self.child == other.child
+            and self.parent == other.parent
+            and self.index == other.index
+            and all(
+                np.allclose(getattr(self, attr), getattr(other, attr))
+                for attr in [
+                    "axis",
+                    "pose",
+                    "friction_static",
+                    "friction_viscous",
+                    "position_limit_damper",
+                    "position_limit_spring",
+                    "position_limit",
+                    "initial_position",
+                    "motor_inertia",
+                    "motor_viscous_friction",
+                    "motor_gear_ratio",
+                ]
+            ),
+        ):
+            return False
+
+        return True
+
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
                 hash(self.name),
-                hash(tuple(self.axis.tolist())),
-                hash(tuple(self.pose.flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.axis),
+                HashedNumpyArray.hash_of_array(self.pose),
                 hash(int(self.jtype)),
                 hash(self.child),
                 hash(self.parent),
                 hash(int(self.index)) if self.index is not None else 0,
-                hash(float(self.friction_static)),
-                hash(float(self.friction_viscous)),
-                hash(float(self.position_limit_damper)),
-                hash(float(self.position_limit_spring)),
-                hash((float(el) for el in self.position_limit)),
-                hash(tuple(np.atleast_1d(self.initial_position).tolist())),
-                hash(float(self.motor_inertia)),
-                hash(float(self.motor_viscous_friction)),
-                hash(float(self.motor_gear_ratio)),
+                HashedNumpyArray.hash_of_array(self.friction_static),
+                HashedNumpyArray.hash_of_array(self.friction_viscous),
+                HashedNumpyArray.hash_of_array(self.position_limit_damper),
+                HashedNumpyArray.hash_of_array(self.position_limit_spring),
+                HashedNumpyArray.hash_of_array(self.position_limit),
+                HashedNumpyArray.hash_of_array(self.initial_position),
+                HashedNumpyArray.hash_of_array(self.motor_inertia),
+                HashedNumpyArray.hash_of_array(self.motor_viscous_friction),
+                HashedNumpyArray.hash_of_array(self.motor_gear_ratio),
             ),
         )

jaxsim/parsers/descriptions/link.py

@@ -12,7 +12,7 @@ import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass
 
 
-@jax_dataclasses.pytree_dataclass
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class LinkDescription(JaxsimDataclass):
     """
     In-memory description of a robot link.
@@ -31,7 +31,7 @@ class LinkDescription(JaxsimDataclass):
     mass: float = dataclasses.field(repr=False)
     inertia: jtp.Matrix = dataclasses.field(repr=False)
     index: int | None = None
-    parent: LinkDescription = dataclasses.field(default=None, repr=False)
+    parent: LinkDescription | None = dataclasses.field(default=None, repr=False)
     pose: jtp.Matrix = dataclasses.field(default_factory=lambda: jnp.eye(4), repr=False)
 
     children: Static[tuple[LinkDescription]] = dataclasses.field(
@@ -40,13 +40,15 @@ class LinkDescription(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
                 hash(self.name),
                 hash(float(self.mass)),
-                hash(tuple(np.atleast_1d(self.inertia).flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.inertia),
                 hash(int(self.index)) if self.index is not None else 0,
-                hash(tuple(np.atleast_1d(self.pose).flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.pose),
                 hash(tuple(self.children)),
                 # Here only using the name to prevent circular recursion:
                 hash(self.parent.name) if self.parent is not None else 0,
@@ -58,7 +60,22 @@ class LinkDescription(JaxsimDataclass):
         if not isinstance(other, LinkDescription):
             return False
 
-        return hash(self) == hash(other)
+        if not (
+            self.name == other.name
+            and np.allclose(self.mass, other.mass)
+            and np.allclose(self.inertia, other.inertia)
+            and self.index == other.index
+            and np.allclose(self.pose, other.pose)
+            and self.children == other.children
+            and (
+                (self.parent is not None and self.parent.name == other.parent.name)
+                if self.parent is not None
+                else other.parent is None
+            ),
+        ):
+            return False
+
+        return True
 
     @property
     def name_and_index(self) -> str:

jaxsim/parsers/descriptions/model.py

@@ -12,7 +12,7 @@ from .joint import JointDescription
 from .link import LinkDescription
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=False)
 class ModelDescription(KinematicGraph):
     """
     Intermediate representation representing the kinematic graph of a robot model.
@@ -28,7 +28,7 @@ class ModelDescription(KinematicGraph):
     fixed_base: bool = True
 
     collision_shapes: tuple[CollisionShape, ...] = dataclasses.field(
-        default_factory=list, repr=False, hash=False
+        default_factory=list, repr=False
     )
 
     @staticmethod
@@ -249,7 +249,17 @@ class ModelDescription(KinematicGraph):
         if not isinstance(other, ModelDescription):
             return False
 
-        return hash(self) == hash(other)
+        if not (
+            self.name == other.name
+            and self.fixed_base == other.fixed_base
+            and self.root == other.root
+            and self.joints == other.joints
+            and self.frames == other.frames
+            and self.root_pose == other.root_pose
+        ):
+            return False
+
+        return True
 
     def __hash__(self) -> int:
 

jaxsim/parsers/kinematic_graph.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import copy
 import dataclasses
 import functools
-from typing import Any, Callable, Iterable, NamedTuple, Sequence
+from typing import Any, Callable, Iterable, Sequence
 
 import numpy as np
 import numpy.typing as npt
@@ -15,7 +15,8 @@ from jaxsim.utils import Mutability
 from . import descriptions
 
 
-class RootPose(NamedTuple):
+@dataclasses.dataclass
+class RootPose:
     """
     Represents the root pose in a kinematic graph.
 
@@ -28,15 +29,20 @@ class RootPose(NamedTuple):
         The root link of the kinematic graph is the base link.
     """
 
-    root_position: npt.NDArray = np.zeros(3)
-    root_quaternion: npt.NDArray = np.array([1.0, 0, 0, 0])
+    root_position: npt.NDArray = dataclasses.field(default_factory=lambda: np.zeros(3))
+
+    root_quaternion: npt.NDArray = dataclasses.field(
+        default_factory=lambda: np.array([1.0, 0, 0, 0])
+    )
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
-                hash(tuple(self.root_position.tolist())),
-                hash(tuple(self.root_quaternion.tolist())),
+                HashedNumpyArray.hash_of_array(self.root_position),
+                HashedNumpyArray.hash_of_array(self.root_quaternion),
             )
         )
 
@@ -45,7 +51,13 @@ class RootPose(NamedTuple):
         if not isinstance(other, RootPose):
             return False
 
-        return hash(self) == hash(other)
+        if not np.allclose(self.root_position, other.root_position):
+            return False
+
+        if not np.allclose(self.root_quaternion, other.root_quaternion):
+            return False
+
+        return True
 
 
 @dataclasses.dataclass(frozen=True)

jaxsim/rbda/soft_contacts.py

@@ -31,11 +31,13 @@ class SoftContactsParams(JaxsimDataclass):
 
     def __hash__(self) -> int:
 
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
         return hash(
             (
-                hash(tuple(jnp.atleast_1d(self.K).flatten().tolist())),
-                hash(tuple(jnp.atleast_1d(self.D).flatten().tolist())),
-                hash(tuple(jnp.atleast_1d(self.mu).flatten().tolist())),
+                HashedNumpyArray.hash_of_array(self.K),
+                HashedNumpyArray.hash_of_array(self.D),
+                HashedNumpyArray.hash_of_array(self.mu),
             )
         )
 

jaxsim/utils/wrappers.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import Generic, TypeVar
+from typing import Callable, Generic, TypeVar
 
 import jax
 import jax_dataclasses
@@ -40,6 +40,33 @@ class HashlessObject(Generic[T]):
         return hash(self) == hash(other)
 
 
+@dataclasses.dataclass
+class CustomHashedObject(Generic[T]):
+    """
+    A class that wraps an object and computes its hash with a custom hash function.
+    """
+
+    obj: T
+
+    hash_function: Callable[[T], int] = dataclasses.field(default=lambda obj: hash(obj))
+
+    def get(self: CustomHashedObject[T]) -> T:
+        return self.obj
+
+    def __hash__(self) -> int:
+
+        return self.hash_function(self.obj)
+
+    def __eq__(self, other: CustomHashedObject[T]) -> bool:
+
+        if not isinstance(other, CustomHashedObject) and isinstance(
+            other.get(), type(self.get())
+        ):
+            return False
+
+        return hash(self) == hash(other)
+
+
 @jax_dataclasses.pytree_dataclass
 class HashedNumpyArray:
     """
@@ -56,6 +83,10 @@ class HashedNumpyArray:
 
     array: jax.Array | npt.NDArray
 
+    precision: float | None = dataclasses.field(
+        default=1e-9, repr=False, compare=False, hash=False
+    )
+
     large_array: jax_dataclasses.Static[bool] = dataclasses.field(
         default=False, repr=False, compare=False, hash=False
     )
@@ -65,7 +96,9 @@ class HashedNumpyArray:
 
     def __hash__(self) -> int:
 
-        return hash(tuple(np.atleast_1d(self.array).flatten().tolist()))
+        return HashedNumpyArray.hash_of_array(
+            array=self.array, precision=self.precision
+        )
 
     def __eq__(self, other: HashedNumpyArray) -> bool:
 
@@ -76,3 +109,37 @@ class HashedNumpyArray:
             return np.array_equal(self.array, other.array)
 
         return hash(self) == hash(other)
+
+    @staticmethod
+    def hash_of_array(
+        array: jax.Array | npt.NDArray, precision: float | None = 1e-9
+    ) -> int:
+        """
+        Calculate the hash of a NumPy array.
+
+        Args:
+            array: The array to hash.
+            precision: Optionally limit the precision over which the hash is computed.
+
+        Returns:
+            The hash of the array.
+        """
+
+        array = np.array(array).flatten()
+
+        array = np.where(array == np.nan, hash(np.nan), array)
+        array = np.where(array == np.inf, hash(np.inf), array)
+        array = np.where(array == -np.inf, hash(-np.inf), array)
+
+        if precision is not None:
+
+            integer1 = (array * precision).astype(int)
+            integer2 = (array - integer1 / precision).astype(int)
+
+            decimal_array = ((array - integer1 * 1e9 - integer2) / precision).astype(
+                int
+            )
+
+            array = np.hstack([integer1, integer2, decimal_array]).astype(int)
+
+        return hash(tuple(array.tolist()))

{jaxsim-0.3.1.dev4.dist-info → jaxsim-0.3.1.dev21.dist-info}/METADATA

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

{jaxsim-0.3.1.dev4.dist-info → jaxsim-0.3.1.dev21.dist-info}/RECORD

@@ -1,19 +1,20 @@
 jaxsim/__init__.py,sha256=xzuTuZrgKdWLqqDzbvqzm2cJrEtAbepOeUqDu7ByVek,2621
-jaxsim/_version.py,sha256=VGLCHIu949eId2wWA41OJLwZQvXRoNWfGQ_EuSkWWtQ,424
+jaxsim/_version.py,sha256=fSekabX0ZEIHwkdp0Sa0iQ2H7hhPzvja3dhu7EFiX4I,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=xfKJz6Rw0YTk-EHCGiT8BFQrs_ggOz01lRi1Qh1mb28,27256
-jaxsim/api/frame.py,sha256=0YXOrGmx3cSQqa4_Ky-n6zyup3I3xvXNEgub-Bc5xUw,6222
+jaxsim/api/data.py,sha256=fkVDBV1tODRYIaRb2N15l34InAcnzNygMGG1KFiIU2w,27307
+jaxsim/api/frame.py,sha256=vSbFHL4WtKPySxunNoZLlM_aDuJXZtf8CSBKku63BAs,6178
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
-jaxsim/api/kin_dyn_parameters.py,sha256=b1e96I8hKU5fh4StLObdVcDpr_6ZglrgD3SRyrqTu18,26203
+jaxsim/api/kin_dyn_parameters.py,sha256=AEpDg9kihbKUN9PA8pNrAruSuWFUC-k_GGxtlcdcDiQ,29215
 jaxsim/api/link.py,sha256=MdMWaMpM5Dj5JHK8uwHZ4zR4Fjq3R4asi2sGTxk1OAs,16647
-jaxsim/api/model.py,sha256=sCx9CcP23A1I_ae4UqTq4Fpq5u0aDki72CqgnR1H50w,59465
+jaxsim/api/model.py,sha256=iuNYsn4xIfX36smmZpwM2O5eftT7ioDQtb6mSUqWu6Q,59759
 jaxsim/api/ode.py,sha256=luTQJsIXUtCp_81dR42X7WrMvwrXtYbyJiqss29v7zA,10786
-jaxsim/api/ode_data.py,sha256=D6FzMkvY_qNuoFEImyp7sxAk-0pJOd3oZeSr9bBTcLk,23089
+jaxsim/api/ode_data.py,sha256=FxUIV5qDNOg_OiOXWs3UrhDgKhGmTKcbHqgr4NX5bv0,23290
 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
@@ -34,12 +35,12 @@ jaxsim/mujoco/loaders.py,sha256=7rjpeJ6_GuitlCty-ZkLhTILQ0GmsFzDMgve-7Gkkh4,2098
 jaxsim/mujoco/model.py,sha256=1KVRjSLOTCuHt53apBPQTnFYJRknlVoKLQaxWsNK8qc,13494
 jaxsim/mujoco/visualizer.py,sha256=PXgQzwetS9mRJYHBknDMLsQ9152FdrSvZuT9xE_dfIQ,5069
 jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
-jaxsim/parsers/kinematic_graph.py,sha256=WdIxntWfxXf67x90oM5KHHXFrSITMwVahqWgcOjYFzc,34730
+jaxsim/parsers/kinematic_graph.py,sha256=1d0JAc3LrGTymaqO9exRsb33-o0Vtgc3cUvNP1YI-0Q,35083
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
 jaxsim/parsers/descriptions/collision.py,sha256=BQeIG-TKi4SVny23w6riDrQ5itC6VRwEMBX6HgAXHxA,3973
-jaxsim/parsers/descriptions/joint.py,sha256=z_nYSS0fdkcaerjUlPX0U1Vn1ArBT0u_XdKjqxG3HcY,3959
-jaxsim/parsers/descriptions/link.py,sha256=QvEE7J6iMQLibpLqlcBV428UA7NMpFFXJwe35GYnjAY,3124
-jaxsim/parsers/descriptions/model.py,sha256=V9nSyCK3mo7680WYMDEx1MTfdDTJzbCGPqAp3qA2XRE,9511
+jaxsim/parsers/descriptions/joint.py,sha256=7qUabpldRKwpGYQLCtQyMKiY47hB78J80DIuzI6bGLc,5186
+jaxsim/parsers/descriptions/link.py,sha256=s0NXGOqmDknX0DYof31TGjVLLUHC9kSwzlGYLcCc03A,3710
+jaxsim/parsers/descriptions/model.py,sha256=vfubtW68CUdgcbCHPcgKy0_BxzKQhhM8ycbCE-dF7Vk,9827
 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
@@ -50,16 +51,16 @@ 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=52zJOF31hFpqoaOednTvi8j_UxhRcdGNjzOPb2v2MPc,11257
+jaxsim/rbda/soft_contacts.py,sha256=0hx9JT4R1X2PPjhZ1EDizBR1gGoCFCtKYu86SeuIvvA,11269
 jaxsim/rbda/utils.py,sha256=zpbFM2Iq8cntku0BFVu9nfEqZhInCWi9D2INT6MFEI8,5003
 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=EJMcblYKUjxw9HJShVf81Ig3pHUJno6Dx6h-RnY--wM,2040
-jaxsim-0.3.1.dev4.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.3.1.dev4.dist-info/METADATA,sha256=VBJIUI8eyYPeIFYpycB-meBbGxcStaW9lqbUyxoxkaM,9738
-jaxsim-0.3.1.dev4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.3.1.dev4.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.3.1.dev4.dist-info/RECORD,,
+jaxsim/utils/wrappers.py,sha256=QIJitSoljrKR_U4T3ewCJPT3DTh-tPZsRsg0t_MH93E,3896
+jaxsim-0.3.1.dev21.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.3.1.dev21.dist-info/METADATA,sha256=wtxQdWa5FFEqYdZx81i-VgNk7DKBY6YMQAXn5_1ctMY,9739
+jaxsim-0.3.1.dev21.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.3.1.dev21.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.3.1.dev21.dist-info/RECORD,,