jaxsim 0.2.1.dev62__py3-none-any.whl → 0.2.1.dev70__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev62'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev62')
+__version__ = version = '0.2.1.dev70'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev70')

jaxsim/api/__init__.py

@@ -1,3 +1,13 @@
 from . import common  # isort:skip
 from . import model, data  # isort:skip
-from . import com, contact, joint, kin_dyn_parameters, link, ode, ode_data, references
+from . import (
+    com,
+    contact,
+    frame,
+    joint,
+    kin_dyn_parameters,
+    link,
+    ode,
+    ode_data,
+    references,
+)

jaxsim/api/frame.py

@@ -0,0 +1,221 @@
+import functools
+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 .common import VelRepr
+
+# =======================
+# Index-related functions
+# =======================
+
+
+def idx_of_parent_link(model: js.model.JaxSimModel, *, frame_idx: jtp.IntLike) -> 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.
+
+    Returns:
+        The index of the frame's parent link.
+    """
+
+    # Get the intermediate representation parsed from the model description.
+    ir = model.description.get()
+
+    # 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:
+    """
+    Convert the name of a frame to its index.
+
+    Args:
+        model: The model to consider.
+        frame_name: The name of the frame.
+
+    Returns:
+        The index of the frame.
+    """
+
+    frame_names = np.array([frame.name for frame in model.description.get().frames])
+
+    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
+
+
+def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str:
+    """
+    Convert the index of a frame to its name.
+
+    Args:
+        model: The model to consider.
+        frame_index: The index of the frame.
+
+    Returns:
+        The name of the frame.
+    """
+
+    return model.description.get().frames[frame_index - model.number_of_links()].name
+
+
+@functools.partial(jax.jit, static_argnames=["frame_names"])
+def names_to_idxs(
+    model: js.model.JaxSimModel, *, frame_names: Sequence[str]
+) -> jax.Array:
+    """
+    Convert a sequence of frame names to their corresponding indices.
+
+    Args:
+        model: The model to consider.
+        frame_names: The names of the frames.
+
+    Returns:
+        The indices of the frames.
+    """
+
+    return jnp.array(
+        [name_to_idx(model=model, frame_name=frame_name) for frame_name in frame_names]
+    ).astype(int)
+
+
+def idxs_to_names(
+    model: js.model.JaxSimModel, *, frame_indices: Sequence[jtp.IntLike]
+) -> tuple[str, ...]:
+    """
+    Convert a sequence of frame indices to their corresponding names.
+
+    Args:
+        model: The model to consider.
+        frame_indices: The indices of the frames.
+
+    Returns:
+        The names of the frames.
+    """
+
+    return tuple(
+        idx_to_name(model=model, frame_index=frame_index)
+        for frame_index in frame_indices
+    )
+
+
+# ==========
+# Frame APIs
+# ==========
+
+
+@functools.partial(jax.jit, static_argnames=["frame_index"])
+def transform(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+) -> jtp.Matrix:
+    """
+    Compute the SE(3) transform from the world frame to the specified frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame for which the transform is requested.
+
+    Returns:
+        The 4x4 matrix representing the transform.
+    """
+
+    # Compute the necessary transforms.
+    L = idx_of_parent_link(model=model, frame_idx=frame_index)
+    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.get().frames[frame_index - model.number_of_links()]
+    L_H_F = frame.pose
+
+    # 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"])
+def jacobian(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the free-floating jacobian of the frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian.
+
+    Returns:
+        The 6×(6+n) free-floating jacobian of the frame.
+
+    Note:
+        The input representation of the free-floating jacobian is the active
+        velocity representation.
+    """
+
+    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)
+
+    # Compute the Jacobian of the parent link using body-fixed output representation.
+    L_J_WL = js.link.jacobian(
+        model=model, data=data, link_index=L, output_vel_repr=VelRepr.Body
+    )
+
+    # 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_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_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
+            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_J_WL = FW_X_L @ L_J_WL
+            O_J_WL_I = FW_J_WL
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    return O_J_WL_I

jaxsim/api/model.py

@@ -252,6 +252,16 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.link_names
 
+    def frame_names(self) -> tuple[str, ...]:
+        """
+        Return the names of the links in the model.
+
+        Returns:
+            The names of the links in the model.
+        """
+
+        return tuple([frame.name for frame in self.description.get().frames])
+
 
 # =====================
 # Model post-processing

jaxsim/parsers/descriptions/link.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 import dataclasses
-from typing import List
 
 import jax.numpy as jnp
 import jax_dataclasses
@@ -16,39 +17,46 @@ 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: List of child 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 = 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[list[LinkDescription]] = dataclasses.field(
         default_factory=list, repr=False
     )
 
     def __hash__(self) -> int:
-        return hash(self.__repr__())
-
-    def __eq__(self, other) -> bool:
-        return (
-            self.name == other.name
-            and self.mass == other.mass
-            and (self.inertia == other.inertia).all()
-            and self.index == other.index
-            and self.parent == other.parent
-            and (self.pose == other.pose).all()
-            and self.children == other.children
+
+        return hash(
+            (
+                hash(self.name),
+                hash(float(self.mass)),
+                hash(tuple(self.inertia.flatten().tolist())),
+                hash(int(self.index)),
+                hash(self.parent),
+                hash(tuple(hash(c) for c in self.children)),
+            )
         )
 
+    def __eq__(self, other: LinkDescription) -> bool:
+
+        if not isinstance(other, LinkDescription):
+            return False
+
+        return hash(self) == hash(other)
+
     @property
     def name_and_index(self) -> str:
         """
@@ -61,19 +69,19 @@ 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
         I_removed = link.inertia
 

jaxsim/parsers/descriptions/model.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import dataclasses
 import itertools
-from typing import List
+from typing import Sequence
 
 from jaxsim import logging
 
@@ -13,63 +15,62 @@ from .link import LinkDescription
 @dataclasses.dataclass(frozen=True)
 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: list[CollisionShape] = dataclasses.field(
+        default_factory=list, repr=False, hash=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: list[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
@@ -78,11 +79,13 @@ class ModelDescription(KinematicGraph):
         # Create the object to compute forward kinematics.
         fk = KinematicGraphTransforms(graph=kinematic_graph)
 
-        # Store here the final model collisions
-        final_collisions: List[CollisionShape] = []
+        # Container of the final model's collision shapes.
+        final_collisions: list[CollisionShape] = []
 
-        # Move and express the collision shapes of the removed link to the lumped link
+        # 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)
 
@@ -112,7 +115,7 @@ class ModelDescription(KinematicGraph):
             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
@@ -145,22 +148,20 @@ class ModelDescription(KinematicGraph):
             _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.
+            The 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.
         """
 
         if len(set(considered_joints) - set(self.joint_names())) != 0:
@@ -172,6 +173,7 @@ class ModelDescription(KinematicGraph):
             name=self.name,
             links=list(self.links_dict.values()),
             joints=self.joints,
+            frames=self.frames,
             collisions=self.collision_shapes,
             fixed_base=self.fixed_base,
             base_link_name=list(iter(self))[0].name,
@@ -190,12 +192,8 @@ class ModelDescription(KinematicGraph):
         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():
@@ -211,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():
@@ -233,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]