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

jaxsim/parsers/descriptions/link.py

@@ -1,53 +1,81 @@
+from __future__ import annotations
+
 import dataclasses
-from typing import List
 
 import jax.numpy as jnp
 import jax_dataclasses
+import numpy as np
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
-from jaxsim.sixd import se3
+from jaxsim.math import Adjoint
 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.
 
     Attributes:
-        name (str): The name of the link.
-        mass (float): The mass of the link.
-        inertia (jtp.Matrix): The inertia matrix of the link.
-        index (Optional[int]): An optional index for the link.
-        parent (Optional[LinkDescription]): The parent link of this link.
-        pose (jtp.Matrix): The pose transformation matrix of the link.
-        children (List[LinkDescription]): List of child links.
+        name: The name of the link.
+        mass: The mass of the link.
+        inertia: The inertia tensor of the link.
+        index: An optional index for the link (it gets automatically assigned).
+        parent: The parent link of this link.
+        pose: The pose transformation matrix of the link.
+        children: The children links.
     """
 
     name: Static[str]
-    mass: float
-    inertia: jtp.Matrix
+    mass: float = dataclasses.field(repr=False)
+    inertia: jtp.Matrix = dataclasses.field(repr=False)
     index: int | None = None
-    parent: Static["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[List["LinkDescription"]] = dataclasses.field(
+
+    children: Static[tuple[LinkDescription]] = dataclasses.field(
         default_factory=list, repr=False
     )
 
     def __hash__(self) -> int:
-        return hash(self.__repr__())
 
-    def __eq__(self, other) -> bool:
-        return (
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                hash(self.name),
+                hash(float(self.mass)),
+                HashedNumpyArray.hash_of_array(self.inertia),
+                hash(int(self.index)) if self.index is not None else 0,
+                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,
+            )
+        )
+
+    def __eq__(self, other: LinkDescription) -> bool:
+
+        if not isinstance(other, LinkDescription):
+            return False
+
+        if not (
             self.name == other.name
-            and self.mass == other.mass
-            and (self.inertia == other.inertia).all()
+            and np.allclose(self.mass, other.mass)
+            and np.allclose(self.inertia, other.inertia)
             and self.index == other.index
-            and self.parent == other.parent
-            and (self.pose == other.pose).all()
+            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:
@@ -61,25 +89,24 @@ class LinkDescription(JaxsimDataclass):
         return f"#{self.index}_<{self.name}>"
 
     def lump_with(
-        self, link: "LinkDescription", lumped_H_removed: jtp.Matrix
-    ) -> "LinkDescription":
+        self, link: LinkDescription, lumped_H_removed: jtp.Matrix
+    ) -> LinkDescription:
         """
         Combine the current link with another link, preserving mass and inertia.
 
         Args:
-            link (LinkDescription): The link to combine with.
-            lumped_H_removed (jtp.Matrix): The transformation matrix between the two links.
+            link: The link to combine with.
+            lumped_H_removed: The transformation matrix between the two links.
 
         Returns:
-            LinkDescription: The combined link.
-
+            The combined link.
         """
-        # Get the 6D inertia of the link to remove
+
+        # Get the 6D inertia of the link to remove.
         I_removed = link.inertia
 
         # Create the SE3 object. Note the inverse.
-        r_H_l = se3.SE3.from_matrix(lumped_H_removed).inverse()
-        r_X_l = r_H_l.adjoint()
+        r_X_l = Adjoint.from_transform(transform=lumped_H_removed, inverse=True)
 
         # Move the inertia
         I_removed_in_lumped_frame = r_X_l.transpose() @ I_removed @ r_X_l

jaxsim/parsers/descriptions/model.py

@@ -1,87 +1,93 @@
+from __future__ import annotations
+
 import dataclasses
 import itertools
-from typing import List
+from collections.abc import Sequence
 
 from jaxsim import logging
 
-from ..kinematic_graph import KinematicGraph, RootPose
+from ..kinematic_graph import KinematicGraph, KinematicGraphTransforms, RootPose
 from .collision import CollidablePoint, CollisionShape
 from .joint import JointDescription
 from .link import LinkDescription
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True, eq=False, unsafe_hash=False)
 class ModelDescription(KinematicGraph):
     """
-    Description of a robotic model including links, joints, and collision shapes.
-
-    Args:
-        name (str): The name of the model.
-        fixed_base (bool): Indicates whether the model has a fixed base.
-        collision_shapes (List[CollisionShape]): List of collision shapes associated with the model.
+    Intermediate representation representing the kinematic graph of a robot model.
 
     Attributes:
-        name (str): The name of the model.
-        fixed_base (bool): Indicates whether the model has a fixed base.
-        collision_shapes (List[CollisionShape]): List of collision shapes associated with the model.
+        name: The name of the model.
+        fixed_base: Whether the model is either fixed-base or floating-base.
+        collision_shapes: List of collision shapes associated with the model.
     """
 
     name: str = None
+
     fixed_base: bool = True
-    collision_shapes: List[CollisionShape] = dataclasses.field(default_factory=list)
+
+    collision_shapes: tuple[CollisionShape, ...] = dataclasses.field(
+        default_factory=list, repr=False
+    )
 
     @staticmethod
     def build_model_from(
         name: str,
-        links: List[LinkDescription],
-        joints: List[JointDescription],
-        collisions: List[CollisionShape] = (),
+        links: list[LinkDescription],
+        joints: list[JointDescription],
+        frames: list[LinkDescription] | None = None,
+        collisions: tuple[CollisionShape, ...] = (),
         fixed_base: bool = False,
         base_link_name: str | None = None,
-        considered_joints: List[str] | None = None,
+        considered_joints: Sequence[str] | None = None,
         model_pose: RootPose = RootPose(),
-    ) -> "ModelDescription":
+    ) -> ModelDescription:
         """
         Build a model description from provided components.
 
         Args:
-            name (str): The name of the model.
-            links (List[LinkDescription]): List of link descriptions.
-            joints (List[JointDescription]): List of joint descriptions.
-            collisions (List[CollisionShape]): List of collision shapes associated with the model.
-            fixed_base (bool): Indicates whether the model has a fixed base.
-            base_link_name (str): Name of the base link.
-            considered_joints (List[str]): List of joint names to consider.
-            model_pose (RootPose): Pose of the model's root.
+            name: The name of the model.
+            links: List of link descriptions.
+            joints: List of joint descriptions.
+            frames: List of frame descriptions.
+            collisions: List of collision shapes associated with the model.
+            fixed_base: Indicates whether the model has a fixed base.
+            base_link_name: Name of the base link (i.e. the root of the kinematic tree).
+            considered_joints: List of joint names to consider (by default all joints).
+            model_pose: Pose of the model's root (by default an identity transform).
 
         Returns:
-            ModelDescription: A ModelDescription instance representing the model.
-
-        Raises:
-            ValueError: If invalid or missing input data.
+            A ModelDescription instance representing the model.
         """
 
-        # Create the full kinematic graph
+        # Create the full kinematic graph.
         kinematic_graph = KinematicGraph.build_from(
             links=links,
             joints=joints,
+            frames=frames,
             root_link_name=base_link_name,
             root_pose=model_pose,
         )
 
-        # Reduce the graph if needed
+        # Reduce the graph if needed.
         if considered_joints is not None:
             kinematic_graph = kinematic_graph.reduce(
                 considered_joints=considered_joints
             )
 
-        # Store here the final model collisions
-        final_collisions: List[CollisionShape] = []
+        # Create the object to compute forward kinematics.
+        fk = KinematicGraphTransforms(graph=kinematic_graph)
 
-        # Move and express the collision shapes of the removed link to the lumped link
+        # Container of the final model's collision shapes.
+        final_collisions: list[CollisionShape] = []
+
+        # Move and express the collision shapes of removed links to the resulting
+        # lumped link that replace the combination of the removed link and its parent.
         for collision_shape in collisions:
+
             # Get all the collidable points of the shape
-            coll_points = list(collision_shape.collidable_points)
+            coll_points = tuple(collision_shape.collidable_points)
 
             # Assume they have an unique parent link
             if not len(set({cp.parent_link.name for cp in coll_points})) == 1:
@@ -105,11 +111,11 @@ class ModelDescription(KinematicGraph):
                 continue
 
             # Create a new collision shape
-            new_collision_shape = CollisionShape(collidable_points=[])
+            new_collision_shape = CollisionShape(collidable_points=())
             final_collisions.append(new_collision_shape)
 
             # If the frame was found, update the collidable points' pose and add them
-            #  to the new collision shape
+            # to the new collision shape.
             for cp in collision_shape.collidable_points:
                 # Find the link that is part of the (reduced) model in which the
                 # collision shape's parent was lumped into
@@ -121,73 +127,73 @@ class ModelDescription(KinematicGraph):
                 # relative pose
                 moved_cp = cp.change_link(
                     new_link=real_parent_link_of_shape,
-                    new_H_old=kinematic_graph.relative_transform(
+                    new_H_old=fk.relative_transform(
                         relative_to=real_parent_link_of_shape.name,
                         name=cp.parent_link.name,
                     ),
                 )
 
-                # Store the updated collision
-                new_collision_shape.collidable_points.append(moved_cp)
+                # Store the updated collision.
+                new_collision_shape.collidable_points += (moved_cp,)
 
         # Build the model
         model = ModelDescription(
             name=name,
             root_pose=kinematic_graph.root_pose,
             fixed_base=fixed_base,
-            collision_shapes=final_collisions,
+            collision_shapes=tuple(final_collisions),
             root=kinematic_graph.root,
             joints=kinematic_graph.joints,
             frames=kinematic_graph.frames,
+            _joints_removed=kinematic_graph.joints_removed,
         )
+
+        # Check that the root link of kinematic graph is the desired base link.
         assert kinematic_graph.root.name == base_link_name, kinematic_graph.root.name
 
         return model
 
-    def reduce(self, considered_joints: List[str]) -> "ModelDescription":
+    def reduce(self, considered_joints: Sequence[str]) -> ModelDescription:
         """
         Reduce the model by removing specified joints.
 
         Args:
-            considered_joints (List[str]): List of joint names to consider.
+            considered_joints: Sequence of joint names to consider.
 
         Returns:
-            ModelDescription: A reduced ModelDescription instance.
-
-        Raises:
-            ValueError: If the specified joints are not part of the model.
+            A `ModelDescription` instance that only includes the considered joints.
         """
 
-        msg = "The model reduction logic assumes that removed joints have zero angles"
-        logging.info(msg=msg)
-
         if len(set(considered_joints) - set(self.joint_names())) != 0:
             extra_joints = set(considered_joints) - set(self.joint_names())
             msg = f"Found joints not part of the model: {extra_joints}"
             raise ValueError(msg)
 
-        return ModelDescription.build_model_from(
+        reduced_model_description = ModelDescription.build_model_from(
             name=self.name,
             links=list(self.links_dict.values()),
             joints=self.joints,
-            collisions=self.collision_shapes,
+            frames=self.frames,
+            collisions=tuple(self.collision_shapes),
             fixed_base=self.fixed_base,
-            base_link_name=list(iter(self))[0].name,
+            base_link_name=next(iter(self)).name,
             model_pose=self.root_pose,
             considered_joints=considered_joints,
         )
 
+        # Include the unconnected/removed joints from the original model.
+        for joint in self.joints_removed:
+            reduced_model_description.joints_removed.append(joint)
+
+        return reduced_model_description
+
     def update_collision_shape_of_link(self, link_name: str, enabled: bool) -> None:
         """
         Enable or disable collision shapes associated with a link.
 
         Args:
-            link_name (str): Name of the link.
-            enabled (bool): Enable or disable collision shapes associated with the link.
-
-        Raises:
-            ValueError: If the link name is not found in the model.
-
+            link_name: The name of the link.
+            enabled: Enable or disable collision shapes associated with the link.
         """
 
         if link_name not in self.link_names():
@@ -203,14 +209,10 @@ class ModelDescription(KinematicGraph):
         Get the collision shape associated with a specific link.
 
         Args:
-            link_name (str): Name of the link.
+            link_name: The name of the link.
 
         Returns:
-            CollisionShape: The collision shape associated with the link.
-
-        Raises:
-            ValueError: If the link name is not found in the model.
-
+            The collision shape associated with the link.
         """
 
         if link_name not in self.link_names():
@@ -225,14 +227,15 @@ class ModelDescription(KinematicGraph):
             ]
         )
 
-    def all_enabled_collidable_points(self) -> List[CollidablePoint]:
+    def all_enabled_collidable_points(self) -> list[CollidablePoint]:
         """
         Get all enabled collidable points in the model.
 
         Returns:
-            List[CollidablePoint]: A list of all enabled collidable points.
+            The list of all enabled collidable points.
 
         """
+
         # Get iterator of all collidable points
         all_collidable_points = itertools.chain.from_iterable(
             [shape.collidable_points for shape in self.collision_shapes]
@@ -240,3 +243,33 @@ class ModelDescription(KinematicGraph):
 
         # Return enabled collidable points
         return [cp for cp in all_collidable_points if cp.enabled]
+
+    def __eq__(self, other: ModelDescription) -> bool:
+
+        if not isinstance(other, ModelDescription):
+            return False
+
+        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:
+
+        return hash(
+            (
+                hash(self.name),
+                hash(self.fixed_base),
+                hash(self.root),
+                hash(tuple(self.joints)),
+                hash(tuple(self.frames)),
+                hash(self.root_pose),
+            )
+        )