jaxsim 0.2.1.dev17__py3-none-any.whl → 0.2.1.dev38__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.dev17'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev17')
+__version__ = version = '0.2.1.dev38'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev38')

jaxsim/api/data.py

@@ -417,7 +417,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
         r"""
         Get the generalized position
-        :math:`\\mathbf{q} = ({}^W \\mathbf{H}_B, \\mathbf{s}) \\in \text{SO}(3) \times \\mathbb{R}^n`.
+        :math:`\mathbf{q} = ({}^W \mathbf{H}_B, \mathbf{s}) \in \text{SO}(3) \times \mathbb{R}^n`.
 
         Returns:
             A tuple containing the base transform and the joint positions.
@@ -429,7 +429,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     def generalized_velocity(self) -> jtp.Vector:
         r"""
         Get the generalized velocity
-        :math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\\, \boldsymbol{\\omega}_{W,B};\\, \\mathbf{s}) \\in \\mathbb{R}^{6+n}`
+        :math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\, \boldsymbol{\omega}_{W,B};\, \mathbf{s}) \in \mathbb{R}^{6+n}`
 
         Returns:
             The generalized velocity in the active representation.

jaxsim/api/model.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import copy
 import dataclasses
 import functools
 import pathlib
-from typing import Any
+from typing import Any, Sequence
 
 import jax
 import jax.numpy as jnp
@@ -55,7 +56,7 @@ class JaxSimModel(JaxsimDataclass):
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
         is_urdf: bool | None = None,
-        considered_joints: list[str] | None = None,
+        considered_joints: Sequence[str] | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from a model description.
@@ -257,24 +258,50 @@ class JaxSimModel(JaxsimDataclass):
 # =====================
 
 
-def reduce(model: JaxSimModel, considered_joints: tuple[str, ...]) -> JaxSimModel:
+def reduce(
+    model: JaxSimModel,
+    considered_joints: tuple[str, ...],
+    locked_joint_positions: dict[str, jtp.Float] | None = None,
+) -> JaxSimModel:
     """
     Reduce the model by lumping together the links connected by removed joints.
 
     Args:
         model: The model to reduce.
         considered_joints: The sequence of joints to consider.
-
-    Note:
-        If considered_joints contains joints not existing in the model, the method
-        will raise an exception. If considered_joints is empty, the method will
-        return a copy of the input model.
+        locked_joint_positions:
+            A dictionary containing the positions of the joints to be considered
+            in the reduction process. The removed joints in the reduced model
+            will have their position locked to their value in this dictionary.
+            If a joint is not part of the dictionary, its position is set to zero.
     """
 
+    locked_joint_positions = (
+        locked_joint_positions if locked_joint_positions is not None else {}
+    )
+
+    # If locked joints are passed, make sure that they are valid.
+    if not set(locked_joint_positions).issubset(model.joint_names()):
+        new_joints = set(model.joint_names()) - set(locked_joint_positions)
+        raise ValueError(f"Passed joints not existing in the model: {new_joints}")
+
+    # Copy the model description with a deep copy of the joints.
+    intermediate_description = dataclasses.replace(
+        model.description.get(), joints=copy.deepcopy(model.description.get().joints)
+    )
+
+    # Update the initial position of the joints.
+    # This is necessary to compute the correct pose of the link pairs connected
+    # to removed joints.
+    for joint_name in set(model.joint_names()) - set(considered_joints):
+        j = intermediate_description.joints_dict[joint_name]
+        with j.mutable_context():
+            j.initial_position = float(locked_joint_positions.get(joint_name, 0.0))
+
     # Reduce the model description.
-    # If considered_joints contains joints not existing in the model, the method
-    # will raise an exception.
-    reduced_intermediate_description = model.description.obj.reduce(
+    # If `considered_joints` contains joints not existing in the model,
+    # the method will raise an exception.
+    reduced_intermediate_description = intermediate_description.reduce(
         considered_joints=list(considered_joints)
     )
 

jaxsim/math/joint_model.py

@@ -15,6 +15,7 @@ from jaxsim.parsers.descriptions import (
     JointType,
     ModelDescription,
 )
+from jaxsim.parsers.kinematic_graph import KinematicGraphTransforms
 
 from .rotation import Rotation
 
@@ -87,21 +88,19 @@ class JointModel:
         # w.r.t. the implicit __model__ SDF frame is not the identity).
         suc_H_i = suc_H_i.at[0].set(ordered_links[0].pose)
 
+        # Create the object to compute forward kinematics.
+        fk = KinematicGraphTransforms(graph=description)
+
         # Compute the parent-to-predecessor and successor-to-child transforms for
         # each joint belonging to the model.
         # Note that the joint indices starts from i=1 given our joint model,
         # therefore the entries at index 0 are not updated.
         for joint in ordered_joints:
             λ_H_pre = λ_H_pre.at[joint.index].set(
-                description.relative_transform(
-                    relative_to=joint.parent.name,
-                    name=joint.name,
-                )
+                fk.relative_transform(relative_to=joint.parent.name, name=joint.name)
             )
             suc_H_i = suc_H_i.at[joint.index].set(
-                description.relative_transform(
-                    relative_to=joint.name, name=joint.child.name
-                )
+                fk.relative_transform(relative_to=joint.name, name=joint.child.name)
             )
 
         # Define the DoFs of the base link.
@@ -243,16 +242,16 @@ def supported_joint_motion(
     """
 
     if isinstance(joint_type, JointType):
-        code = joint_type
+        type_enum = joint_type
     elif isinstance(joint_type, JointDescriptor):
-        code = joint_type.code
+        type_enum = joint_type.joint_type
     else:
         raise ValueError(joint_type)
 
     # Prepare the joint position
     s = jnp.array(joint_position).astype(float)
 
-    match code:
+    match type_enum:
 
         case JointType.R:
             joint_type: JointGenericAxis
@@ -276,58 +275,8 @@ def supported_joint_motion(
             S = jnp.vstack(jnp.hstack([joint_type.axis.squeeze(), jnp.zeros(3)]))
 
         case JointType.F:
-            raise ValueError("Fixed joints shouldn't be here")
-
-        case JointType.Rx:
-
-            pre_H_suc = jaxlie.SE3.from_rotation(
-                rotation=jaxlie.SO3.from_x_radians(theta=s)
-            )
-
-            S = jnp.vstack([0, 0, 0, 1.0, 0, 0])
-
-        case JointType.Ry:
-
-            pre_H_suc = jaxlie.SE3.from_rotation(
-                rotation=jaxlie.SO3.from_y_radians(theta=s)
-            )
-
-            S = jnp.vstack([0, 0, 0, 0, 1.0, 0])
-
-        case JointType.Rz:
-
-            pre_H_suc = jaxlie.SE3.from_rotation(
-                rotation=jaxlie.SO3.from_z_radians(theta=s)
-            )
-
-            S = jnp.vstack([0, 0, 0, 0, 0, 1.0])
-
-        case JointType.Px:
-
-            pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
-                rotation=jaxlie.SO3.identity(),
-                translation=jnp.array([s, 0.0, 0.0]),
-            )
-
-            S = jnp.vstack([1.0, 0, 0, 0, 0, 0])
-
-        case JointType.Py:
-
-            pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
-                rotation=jaxlie.SO3.identity(),
-                translation=jnp.array([0.0, s, 0.0]),
-            )
-
-            S = jnp.vstack([0, 1.0, 0, 0, 0, 0])
-
-        case JointType.Pz:
-
-            pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
-                rotation=jaxlie.SO3.identity(),
-                translation=jnp.array([0.0, 0.0, s]),
-            )
-
-            S = jnp.vstack([0, 0, 1.0, 0, 0, 0])
+            pre_H_suc = jaxlie.SE3.identity()
+            S = jnp.zeros(shape=(6, 1))
 
         case _:
             raise ValueError(joint_type)

jaxsim/parsers/descriptions/joint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import dataclasses
 import enum
 from typing import Tuple, Union
@@ -6,79 +8,60 @@ import jax_dataclasses
 import numpy as np
 import numpy.typing as npt
 
+import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 from .link import LinkDescription
 
 
+@enum.unique
 class JointType(enum.IntEnum):
     """
-    Enumeration of joint types for robot joints.
-
-    Args:
-        F: Fixed joint (no movement).
-        R: Revolute joint (rotation).
-        P: Prismatic joint (translation).
-        Rx: Revolute joint with rotation about the X-axis.
-        Ry: Revolute joint with rotation about the Y-axis.
-        Rz: Revolute joint with rotation about the Z-axis.
-        Px: Prismatic joint with translation along the X-axis.
-        Py: Prismatic joint with translation along the Y-axis.
-        Pz: Prismatic joint with translation along the Z-axis.
+    Type of supported joints.
     """
 
-    F = enum.auto()  # Fixed
-    R = enum.auto()  # Revolute
-    P = enum.auto()  # Prismatic
+    @staticmethod
+    def _generate_next_value_(name, start, count, last_values):
+        # Start auto Enum value from 0 instead of 1
+        return count
+
+    #: Fixed joint.
+    F = enum.auto()
 
-    # Revolute joints, single axis
-    Rx = enum.auto()
-    Ry = enum.auto()
-    Rz = enum.auto()
+    #: Revolute joint (1 DoF around axis).
+    R = enum.auto()
 
-    # Prismatic joints, single axis
-    Px = enum.auto()
-    Py = enum.auto()
-    Pz = enum.auto()
+    #: Prismatic joint (1 DoF along axis).
+    P = enum.auto()
 
 
-@dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class JointDescriptor:
     """
-    Description of a joint type with a specific code.
-
-    Args:
-        code (JointType): The code representing the joint type.
-
+    Base class for joint types requiring to store additional metadata.
     """
 
-    code: JointType
+    #: The joint type.
+    joint_type: JointType
 
-    def __hash__(self) -> int:
-        return hash(self.__repr__())
 
-
-@dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class JointGenericAxis(JointDescriptor):
     """
-    Description of a joint type with a generic axis.
-
-    Attributes:
-        axis (npt.NDArray): The axis of rotation or translation for the joint.
-
+    A joint requiring the specification of a 3D axis.
     """
 
-    axis: npt.NDArray
+    #: The axis of rotation or translation of the joint (must have norm 1).
+    axis: jtp.Vector
 
-    def __post_init__(self):
-        if np.allclose(self.axis, 0.0):
-            raise ValueError(self.axis)
+    def __hash__(self) -> int:
+        return hash((self.joint_type, tuple(np.array(self.axis).tolist())))
 
-    def __eq__(self, other):
-        return super().__eq__(other) and np.allclose(self.axis, other.axis)
+    def __eq__(self, other: JointGenericAxis) -> bool:
+        if not isinstance(other, JointGenericAxis):
+            return False
 
-    def __hash__(self) -> int:
-        return hash(self.__repr__())
+        return hash(self) == hash(other)
 
 
 @jax_dataclasses.pytree_dataclass

jaxsim/parsers/descriptions/model.py

@@ -4,7 +4,7 @@ from typing import List
 
 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
@@ -75,6 +75,9 @@ class ModelDescription(KinematicGraph):
                 considered_joints=considered_joints
             )
 
+        # Create the object to compute forward kinematics.
+        fk = KinematicGraphTransforms(graph=kinematic_graph)
+
         # Store here the final model collisions
         final_collisions: List[CollisionShape] = []
 
@@ -121,7 +124,7 @@ 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,
                     ),
@@ -139,7 +142,9 @@ class ModelDescription(KinematicGraph):
             root=kinematic_graph.root,
             joints=kinematic_graph.joints,
             frames=kinematic_graph.frames,
+            _joints_removed=kinematic_graph._joints_removed,
         )
+
         assert kinematic_graph.root.name == base_link_name, kinematic_graph.root.name
 
         return model
@@ -158,15 +163,12 @@ class ModelDescription(KinematicGraph):
             ValueError: If the specified joints are not part of the model.
         """
 
-        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,
@@ -177,6 +179,12 @@ class ModelDescription(KinematicGraph):
             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.

jaxsim/parsers/kinematic_graph.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import dataclasses
 import functools
@@ -9,6 +11,7 @@ from typing import (
     List,
     NamedTuple,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -41,7 +44,7 @@ class RootPose(NamedTuple):
 
 
 @dataclasses.dataclass(frozen=True)
-class KinematicGraph:
+class KinematicGraph(Sequence[descriptions.LinkDescription]):
     """
     Represents a kinematic graph of links and joints.
 
@@ -76,6 +79,12 @@ class KinematicGraph:
         repr=False, compare=False, default_factory=dict
     )
 
+    # Private attribute storing the unconnected joints from the parsed model and
+    # the joints removed after model reduction.
+    _joints_removed: list[descriptions.JointDescription] = dataclasses.field(
+        default_factory=list, repr=False, compare=False
+    )
+
     @functools.cached_property
     def links_dict(self) -> Dict[str, descriptions.LinkDescription]:
         return {l.name: l for l in iter(self)}
@@ -153,15 +162,24 @@ class KinematicGraph:
         # Couple links and joints and create the graph of links.
         # Note that the pose of the frames is not updated; it's the caller's
         # responsibility to update their pose if they want to use them.
-        graph_root_node, graph_joints, graph_frames = KinematicGraph.create_graph(
-            links=links, joints=joints, root_link_name=root_link_name
+        graph_root_node, graph_joints, graph_frames, unconnected_joints = (
+            KinematicGraph.create_graph(
+                links=links, joints=joints, root_link_name=root_link_name
+            )
         )
 
         for frame in graph_frames:
             logging.warning(msg=f"Ignoring unconnected link / frame: '{frame.name}'")
 
+        for joint in unconnected_joints:
+            logging.warning(msg=f"Ignoring unconnected joint: '{joint.name}'")
+
         return KinematicGraph(
-            root=graph_root_node, joints=graph_joints, frames=[], root_pose=root_pose
+            root=graph_root_node,
+            joints=graph_joints,
+            frames=[],
+            root_pose=root_pose,
+            _joints_removed=unconnected_joints,
         )
 
     @staticmethod
@@ -173,9 +191,10 @@ class KinematicGraph:
         descriptions.LinkDescription,
         List[descriptions.JointDescription],
         List[descriptions.LinkDescription],
+        list[descriptions.JointDescription],
     ]:
         """
-        Create a kinematic graph from lists of links and joints.
+        Create a kinematic graph from the lists of parsed links and joints.
 
         Args:
             links (List[descriptions.LinkDescription]): A list of link descriptions.
@@ -183,8 +202,9 @@ class KinematicGraph:
             root_link_name (str): The name of the root link.
 
         Returns:
-            Tuple[descriptions.LinkDescription, List[descriptions.JointDescription], List[descriptions.LinkDescription]]:
-                A tuple containing the root link, list of joints, and list of frames in the graph.
+            A tuple containing the root node with the full kinematic graph as child nodes,
+            the list of joints associated to graph nodes, the list of frames rigidly
+            attached to graph nodes, and the list of joints not part of the graph.
         """
 
         # Create a dict that maps link name to the link, for easy retrieval
@@ -246,18 +266,25 @@ class KinematicGraph:
             links_dict[root_link_name].mutable(mutable=False),
             list(set(joints) - set(removed_joints)),
             frames,
+            list(set(removed_joints)),
         )
 
-    def reduce(self, considered_joints: List[str]) -> "KinematicGraph":
+    def reduce(self, considered_joints: List[str]) -> KinematicGraph:
         """
-        Reduce the kinematic graph by removing specified joints and lumping the mass and inertia of removed links into their parent links.
+        Reduce the kinematic graph by removing unspecified joints.
+
+        When a joint is removed, the mass and inertia of its child link are lumped
+        with those of its parent link, obtaining a new link that combines the two.
+        The description of the removed joint specifies the default angle (usually 0)
+        that is considered when the joint is removed.
 
         Args:
-            considered_joints (List[str]): A list of joint names to consider.
+            considered_joints: A list of joint names to consider.
 
         Returns:
-            KinematicGraph: The reduced kinematic graph.
+            The reduced kinematic graph.
         """
+
         # The current object represents the complete kinematic graph
         full_graph = self
 
@@ -281,6 +308,9 @@ class KinematicGraph:
         links_dict = copy.deepcopy(full_graph.links_dict)
         joints_dict = copy.deepcopy(full_graph.joints_dict)
 
+        # Create the object to compute forward kinematics.
+        fk = KinematicGraphTransforms(graph=full_graph)
+
         # The following steps are implemented below in order to create the reduced graph:
         #
         # 1. Lump the mass of the removed links into their parent
@@ -329,7 +359,7 @@ class KinematicGraph:
             # Lump the link
             lumped_link = parent_of_link_to_remove.lump_with(
                 link=link_to_remove,
-                lumped_H_removed=full_graph.relative_transform(
+                lumped_H_removed=fk.relative_transform(
                     relative_to=parent_of_link_to_remove.name, name=link_to_remove.name
                 ),
             )
@@ -370,7 +400,7 @@ class KinematicGraph:
             # Update the pose. Note that after the lumping process, the dict entry
             # links_dict[joint.parent.name] contains the final lumped link
             with joint.mutable_context(mutability=Mutability.MUTABLE):
-                joint.pose = full_graph.relative_transform(
+                joint.pose = fk.relative_transform(
                     relative_to=links_dict[joint.parent.name].name, name=joint.name
                 )
             with joint.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
@@ -396,18 +426,25 @@ class KinematicGraph:
 
         # Create the reduced graph data. We pass the full list of links so that those
         # that are not part of the graph will be returned as frames.
-        reduced_root_node, reduced_joints, reduced_frames = KinematicGraph.create_graph(
-            links=list(full_graph_links_dict.values()),
-            joints=[joints_dict[joint_name] for joint_name in considered_joints],
-            root_link_name=full_graph.root.name,
+        reduced_root_node, reduced_joints, reduced_frames, unconnected_joints = (
+            KinematicGraph.create_graph(
+                links=list(full_graph_links_dict.values()),
+                joints=[joints_dict[joint_name] for joint_name in considered_joints],
+                root_link_name=full_graph.root.name,
+            )
         )
 
         # Create the reduced graph
         reduced_graph = KinematicGraph(
             root=reduced_root_node,
             joints=reduced_joints,
-            frames=reduced_frames,
+            frames=self.frames + reduced_frames,
             root_pose=full_graph.root_pose,
+            _joints_removed=(
+                self._joints_removed
+                + unconnected_joints
+                + [joints_dict[name] for name in joint_names_to_remove]
+            ),
         )
 
         # ================================================================
@@ -424,7 +461,7 @@ class KinematicGraph:
 
             # Update the connection of the frame
             frame.parent = new_parent_link
-            frame.pose = full_graph.relative_transform(
+            frame.pose = fk.relative_transform(
                 relative_to=new_parent_link.name, name=frame.name
             )
 
@@ -462,65 +499,6 @@ class KinematicGraph:
         """
         return list(self.frames_dict.keys())
 
-    def transform(self, name: str) -> npt.NDArray:
-        """
-        Compute the transformation matrix for a given link, joint, or frame.
-
-        Args:
-            name (str): The name of the link, joint, or frame.
-
-        Returns:
-            npt.NDArray: The transformation matrix.
-        """
-        if name in self.transform_cache:
-            return self.transform_cache[name]
-
-        if name in self.joint_names():
-            joint = self.joints_dict[name]
-
-            if joint.initial_position != 0.0:
-                msg = f"Ignoring unsupported initial position of joint '{name}'"
-                logging.warning(msg=msg)
-
-            transform = self.transform(name=joint.parent.name) @ joint.pose
-            self.transform_cache[name] = transform
-            return self.transform_cache[name]
-
-        if name in self.link_names():
-            link = self.links_dict[name]
-
-            if link.name == self.root.name:
-                return link.pose
-
-            parent_joint = self.joints_connection_dict[(link.parent.name, link.name)]
-            transform = self.transform(name=parent_joint.name) @ link.pose
-            self.transform_cache[name] = transform
-            return self.transform_cache[name]
-
-        # It can only be a plain frame
-        if name not in self.frame_names():
-            raise ValueError(name)
-
-        frame = self.frames_dict[name]
-        transform = self.transform(name=frame.parent.name) @ frame.pose
-        self.transform_cache[name] = transform
-        return self.transform_cache[name]
-
-    def relative_transform(self, relative_to: str, name: str) -> npt.NDArray:
-        """
-        Compute the relative transformation matrix between two elements in the kinematic graph.
-
-        Args:
-            relative_to (str): The name of the reference element.
-            name (str): The name of the element to compute the relative transformation for.
-
-        Returns:
-            npt.NDArray: The relative transformation matrix.
-        """
-        return np.linalg.inv(self.transform(name=relative_to)) @ self.transform(
-            name=name
-        )
-
     def print_tree(self) -> None:
         """
         Print the tree structure of the kinematic graph.
@@ -574,6 +552,10 @@ class KinematicGraph:
 
                 yield child
 
+    # =================
+    # Sequence protocol
+    # =================
+
     def __iter__(self) -> Iterable[descriptions.LinkDescription]:
         yield from KinematicGraph.breadth_first_search(root=self.root)
 
@@ -606,3 +588,188 @@ class KinematicGraph:
             return list(iter(self))[key]
 
         raise TypeError(type(key).__name__)
+
+    def count(self, value: descriptions.LinkDescription) -> int:
+        return list(iter(self)).count(value)
+
+    def index(
+        self, value: descriptions.LinkDescription, start: int = 0, stop: int = -1
+    ) -> int:
+        return list(iter(self)).index(value, start, stop)
+
+
+# ====================
+# Other useful classes
+# ====================
+
+
+@dataclasses.dataclass(frozen=True)
+class KinematicGraphTransforms:
+
+    graph: KinematicGraph
+
+    _transform_cache: dict[str, npt.NDArray] = dataclasses.field(
+        default_factory=dict, init=False, repr=False, compare=False
+    )
+
+    _initial_joint_positions: dict[str, float] = dataclasses.field(
+        init=False, repr=False, compare=False
+    )
+
+    def __post_init__(self) -> None:
+
+        super().__setattr__(
+            "_initial_joint_positions",
+            {joint.name: joint.initial_position for joint in self.graph.joints},
+        )
+
+    @property
+    def initial_joint_positions(self) -> npt.NDArray:
+
+        return np.atleast_1d(
+            np.array(list(self._initial_joint_positions.values()))
+        ).astype(float)
+
+    @initial_joint_positions.setter
+    def initial_joint_positions(
+        self,
+        positions: npt.NDArray | Sequence,
+        joint_names: Sequence[str] | None = None,
+    ) -> None:
+
+        joint_names = (
+            joint_names
+            if joint_names is not None
+            else list(self._initial_joint_positions.keys())
+        )
+
+        s = np.atleast_1d(np.array(positions).squeeze())
+
+        if s.size != len(joint_names):
+            raise ValueError(s.size, len(joint_names))
+
+        for joint_name in joint_names:
+            if joint_name not in self._initial_joint_positions:
+                raise ValueError(joint_name)
+
+        # Clear transform cache.
+        self._transform_cache.clear()
+
+        # Update initial joint positions.
+        for joint_name, position in zip(joint_names, s):
+            self._initial_joint_positions[joint_name] = position
+
+    def transform(self, name: str) -> npt.NDArray:
+        """
+        Compute the SE(3) transform of elements belonging to the kinematic graph.
+
+        Args:
+            name: The name of a link, a joint, or a frame.
+
+        Returns:
+            The 4x4 transform matrix of the element w.r.t. the model frame.
+        """
+
+        # If the transform was already computed, return it.
+        if name in self._transform_cache:
+            return self._transform_cache[name]
+
+        # If the name is a joint, compute M_H_J transform.
+        if name in self.graph.joint_names():
+
+            # Get the joint.
+            joint = self.graph.joints_dict[name]
+
+            # Get the transform of the parent link.
+            M_H_L = self.transform(name=joint.parent.name)
+
+            # Rename the pose of the predecessor joint frame w.r.t. its parent link.
+            L_H_pre = joint.pose
+
+            # Compute the joint transform from the predecessor to the successor frame.
+            pre_H_J = self.pre_H_suc(
+                joint_type=joint.jtype,
+                joint_position=self._initial_joint_positions[joint.name],
+            )
+
+            # Compute the M_H_J transform.
+            self._transform_cache[name] = M_H_L @ L_H_pre @ pre_H_J
+            return self._transform_cache[name]
+
+        # If the name is a link, compute M_H_L transform.
+        if name in self.graph.link_names():
+
+            # Get the link.
+            link = self.graph.links_dict[name]
+
+            # Handle the pose between the __model__ frame and the root link.
+            if link.name == self.graph.root.name:
+                M_H_B = link.pose
+                return M_H_B
+
+            # Get the joint between the link and its parent.
+            parent_joint = self.graph.joints_connection_dict[
+                (link.parent.name, link.name)
+            ]
+
+            # Get the transform of the parent joint.
+            M_H_J = self.transform(name=parent_joint.name)
+
+            # Rename the pose of the link w.r.t. its parent joint.
+            J_H_L = link.pose
+
+            # Compute the M_H_L transform.
+            self._transform_cache[name] = M_H_J @ J_H_L
+            return self._transform_cache[name]
+
+        # It can only be a plain frame.
+        if name not in self.graph.frame_names():
+            raise ValueError(name)
+
+        # Get the frame.
+        frame = self.graph.frames_dict[name]
+
+        # Get the transform of the parent link.
+        M_H_L = self.transform(name=frame.parent.name)
+
+        # Rename the pose of the frame w.r.t. its parent link.
+        L_H_F = frame.pose
+
+        # Compute the M_H_F transform.
+        self._transform_cache[name] = M_H_L @ L_H_F
+        return self._transform_cache[name]
+
+    def relative_transform(self, relative_to: str, name: str) -> npt.NDArray:
+        """
+        Compute the SE(3) relative transform of elements belonging to the kinematic graph.
+
+        Args:
+            relative_to: The name of the reference element.
+            name: The name of a link, a joint, or a frame.
+
+        Returns:
+            The 4x4 transform matrix of the element w.r.t. the desired frame.
+        """
+
+        import jaxsim.math
+
+        M_H_target = self.transform(name=name)
+        M_H_R = self.transform(name=relative_to)
+
+        # Compute the relative transform R_H_target, where R is the reference frame,
+        # and i the frame of the desired link|joint|frame.
+        return np.array(jaxsim.math.Transform.inverse(M_H_R)) @ M_H_target
+
+    @staticmethod
+    def pre_H_suc(
+        joint_type: descriptions.JointType | descriptions.JointDescriptor,
+        joint_position: float | None = None,
+    ) -> npt.NDArray:
+
+        import jaxsim.math
+
+        return np.array(
+            jaxsim.math.supported_joint_motion(
+                joint_type=joint_type, joint_position=joint_position
+            )[0]
+        )

jaxsim/parsers/rod/parser.py

@@ -134,7 +134,7 @@ def extract_model_data(
                 name=j.name,
                 parent=world_link,
                 child=links_dict[j.child],
-                jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
+                jtype=utils.joint_to_joint_type(joint=j),
                 axis=(
                     np.array(j.axis.xyz.xyz)
                     if j.axis is not None
@@ -201,7 +201,7 @@ def extract_model_data(
             name=j.name,
             parent=links_dict[j.parent],
             child=links_dict[j.child],
-            jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
+            jtype=utils.joint_to_joint_type(joint=j),
             axis=(
                 np.array(j.axis.xyz.xyz)
                 if j.axis is not None

jaxsim/parsers/rod/utils.py

@@ -1,5 +1,4 @@
 import os
-from typing import Union
 
 import jaxlie
 import numpy as np
@@ -60,57 +59,43 @@ def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
     return M_L.astype(dtype=float)
 
 
-def axis_to_jtype(
-    axis: rod.Axis, type: str
-) -> Union[descriptions.JointType, descriptions.JointDescriptor]:
+def joint_to_joint_type(
+    joint: rod.Joint,
+) -> descriptions.JointType | descriptions.JointDescriptor:
     """
-    Convert an SDF axis to a joint type.
+    Extract the joint type from an SDF joint.
 
     Args:
-        axis: The SDF axis.
-        type: The SDF joint type.
+        joint: The parsed SDF joint.
 
     Returns:
         The corresponding joint type description.
     """
 
-    if type == "fixed":
+    axis = joint.axis
+    joint_type = joint.type
+
+    if joint_type == "fixed":
         return descriptions.JointType.F
 
     if not (axis.xyz is not None and axis.xyz.xyz is not None):
         raise ValueError("Failed to read axis xyz data")
 
-    axis_xyz = np.array(axis.xyz.xyz)
-
-    if np.allclose(axis_xyz, [1, 0, 0]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Rx
-
-    if np.allclose(axis_xyz, [0, 1, 0]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Ry
-
-    if np.allclose(axis_xyz, [0, 0, 1]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Rz
-
-    if np.allclose(axis_xyz, [1, 0, 0]) and type == "prismatic":
-        return descriptions.JointType.Px
-
-    if np.allclose(axis_xyz, [0, 1, 0]) and type == "prismatic":
-        return descriptions.JointType.Py
-
-    if np.allclose(axis_xyz, [0, 0, 1]) and type == "prismatic":
-        return descriptions.JointType.Pz
+    # Make sure that the axis is a unary vector.
+    axis_xyz = np.array(axis.xyz.xyz).astype(float)
+    axis_xyz = axis_xyz / np.linalg.norm(axis_xyz)
 
-    if type == "revolute":
+    if joint_type in {"revolute", "continuous"}:
         return descriptions.JointGenericAxis(
-            code=descriptions.JointType.R, axis=np.array(axis_xyz, dtype=float)
+            joint_type=descriptions.JointType.R, axis=axis_xyz
         )
 
-    if type == "prismatic":
+    if joint_type == "prismatic":
         return descriptions.JointGenericAxis(
-            code=descriptions.JointType.P, axis=np.array(axis_xyz, dtype=float)
+            joint_type=descriptions.JointType.P, axis=axis_xyz
         )
 
-    raise ValueError("Joint not supported", axis_xyz, type)
+    raise ValueError("Joint not supported", axis_xyz, joint_type)
 
 
 def create_box_collision(

jaxsim/utils/jaxsim_dataclass.py

@@ -50,7 +50,9 @@ class JaxsimDataclass(abc.ABC):
 
     @contextlib.contextmanager
     def mutable_context(
-        self: Self, mutability: Mutability, restore_after_exception: bool = True
+        self: Self,
+        mutability: Mutability = Mutability.MUTABLE,
+        restore_after_exception: bool = True,
     ) -> Iterator[Self]:
         """
         Context manager to temporarily change the mutability of the object.
@@ -86,7 +88,7 @@ class JaxsimDataclass(abc.ABC):
                 setattr(self, f.name, getattr(self_copy, f.name))
 
         try:
-            self.set_mutability(mutability)
+            self.set_mutability(mutability=mutability)
             yield self
 
             if mutability is not Mutability.MUTABLE_NO_VALIDATION:

{jaxsim-0.2.1.dev17.dist-info → jaxsim-0.2.1.dev38.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev17
+Version: 0.2.1.dev38
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it
@@ -43,7 +43,7 @@ Provides-Extra: all
 Requires-Dist: black[jupyter] ~=24.0 ; extra == 'all'
 Requires-Dist: isort ; extra == 'all'
 Requires-Dist: pre-commit ; extra == 'all'
-Requires-Dist: idyntree ; extra == 'all'
+Requires-Dist: idyntree >=12.2.1 ; extra == 'all'
 Requires-Dist: pytest >=6.0 ; extra == 'all'
 Requires-Dist: pytest-icdiff ; extra == 'all'
 Requires-Dist: robot-descriptions ; extra == 'all'
@@ -55,7 +55,7 @@ Requires-Dist: black[jupyter] ~=24.0 ; extra == 'style'
 Requires-Dist: isort ; extra == 'style'
 Requires-Dist: pre-commit ; extra == 'style'
 Provides-Extra: testing
-Requires-Dist: idyntree ; extra == 'testing'
+Requires-Dist: idyntree >=12.2.1 ; extra == 'testing'
 Requires-Dist: pytest >=6.0 ; extra == 'testing'
 Requires-Dist: pytest-icdiff ; extra == 'testing'
 Requires-Dist: robot-descriptions ; extra == 'testing'

{jaxsim-0.2.1.dev17.dist-info → jaxsim-0.2.1.dev38.dist-info}/RECORD

@@ -1,16 +1,16 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=E6AF09CLGzT7X3HBCAUGv8Fgy4Lo-CauUZTQ3aY6D5k,426
+jaxsim/_version.py,sha256=Cl_kuTOEO0aVEiwkm9tH1bjrDIciBhU4mo1dsIZmCdM,426
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=fNTCPUeDfOAizRd4RsW3Epv0sLTu0KJGoFRSEsi75VM,162
 jaxsim/api/com.py,sha256=Qtm_6qpiK4WtDVn6JMUHa8DySgBl9CjgKCybJqZ58Lc,7379
 jaxsim/api/common.py,sha256=DV-WZG28sikXopNv458aYvpLjmiAtFr5LRscOwXusuk,6640
 jaxsim/api/contact.py,sha256=Ve4ZOWkLEBRgK3KhtICxKY7YzsxYvc3lO-pPRBjqSnY,8659
-jaxsim/api/data.py,sha256=cPFy6TTm0rNpDk1O0ZxOpvBwWVPu6V14XLRMeC7x81Q,26786
+jaxsim/api/data.py,sha256=GXSCfq4_PsRZuUwsmcl18e2Ppdtu4ykYs0oWih8J2ZI,26775
 jaxsim/api/joint.py,sha256=q31Kp3Cqv-yTcxijjzbj_QADFnGQyjb2al9fYZtzedo,4763
 jaxsim/api/kin_dyn_parameters.py,sha256=G4mtSi8fElYe0yttLgsxSOPf7vcK-yqTu06Aa5SSrYg,26012
 jaxsim/api/link.py,sha256=rypTwkMf9HJ5UuAtHRJh0LqqdJWcLKTtTjWcjduEsF0,9842
-jaxsim/api/model.py,sha256=UwjZctpQR_3F8_Cm-NmFQb14Wtw_vscs3Yg-ULpQyrs,52269
+jaxsim/api/model.py,sha256=Ci7hl4M1gbnHNQAO5GLjxOdvJe8w3dhvGB_7x2gdpaU,53540
 jaxsim/api/ode.py,sha256=6l-6i2YHagsQvR8Ac-_fmO6P0hBVT6NkHhwXnrdITEg,9785
 jaxsim/api/ode_data.py,sha256=dwRFVPJ30XMmdUbPXEu7YxsQ97jZP4L4fd5ZzhrO5Ys,22184
 jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
@@ -22,7 +22,7 @@ jaxsim/math/__init__.py,sha256=inJ9nRFkqstuGa8OyFkfWVudo5U9Ug4WgDBuKva8AIA,337
 jaxsim/math/adjoint.py,sha256=DT21izjVW497GRrgNfx8tv0ZeWW5QncWMGMhI0acUNw,4425
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
 jaxsim/math/inertia.py,sha256=UAB7ym4gXFanejcs_ovZMpteHCc6poWYmt-mLmd5hhk,1640
-jaxsim/math/joint_model.py,sha256=LKLB26VMz6vx9JLdFUWhGyrElYFEQV-bJiQO5kaZUGY,10896
+jaxsim/math/joint_model.py,sha256=LSZM94oLpdc0z8DJr9lrN10PHYyScjvMZ1BVORUMkhk,9618
 jaxsim/math/quaternion.py,sha256=X9b8jHf0QemKUjIZSnXRJc3DdMr42CBhBy_mi9_X_AM,5068
 jaxsim/math/rotation.py,sha256=Z90daUjGpuNEVLfWB3SVtM9EtwAIaneVj9A9UpWXqhA,2182
 jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
@@ -33,15 +33,15 @@ jaxsim/mujoco/loaders.py,sha256=-z6TIsems7DDogueDBwYLS2JuoWalNfv6sV1F45H63k,2102
 jaxsim/mujoco/model.py,sha256=5-4KTbEbU19zjrSuvUVdLo3noWxTvlCNsFIs3rQTNDY,13506
 jaxsim/mujoco/visualizer.py,sha256=YlteqcCbeB1B6saAHKBz1IJad3N5Rp163reZrzKLAys,5065
 jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
-jaxsim/parsers/kinematic_graph.py,sha256=_ZJr-u1m64qPCeY6BHTI56dSOWhOGfbg2HN49XbTXhA,23784
+jaxsim/parsers/kinematic_graph.py,sha256=pzjY-2O9QDY5wUXHAhuTyWHs1Te238gkzlHnyKeA5T0,28935
 jaxsim/parsers/descriptions/__init__.py,sha256=EbTfnrK3oCxA3pNv--YUwllJ6uICENvFgAdRbYtS9ts,238
 jaxsim/parsers/descriptions/collision.py,sha256=HUWwuRgI9KznY29FFw1_zU3bGigDEezrcPOJSxSJGNU,3382
-jaxsim/parsers/descriptions/joint.py,sha256=hpH0ANvIhbEQk-NGRmWIvPv3lXW385TBIMWNgz5rzM4,4106
+jaxsim/parsers/descriptions/joint.py,sha256=vozAx4sr3q-9IroN90iLP0BPpY_5BJ6VME3DKPrBIoY,3600
 jaxsim/parsers/descriptions/link.py,sha256=hqLLitrAXnouia6ULk1BPOIEfRxrXwHmoPsi306IZW8,2859
-jaxsim/parsers/descriptions/model.py,sha256=wenuDrjoBf6prkzm9WyYT0nFWc0l6WBpKNjLoRUDPxo,8937
+jaxsim/parsers/descriptions/model.py,sha256=hapbpO1VpC7f5ELkGH1ZtM2I3dpmYLz8Em5cb3DdiaA,9264
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=mFi1baSJte6EMmWLpVjVuCpicfAAF48aFUzoKzYzPpo,12555
-jaxsim/parsers/rod/utils.py,sha256=xoxNdb4IgOePFFnblNo6UrgYCRm53mB3No3yXM6qlpw,6471
+jaxsim/parsers/rod/parser.py,sha256=83s8Z33LSS1ddYjjm0lOqf_fQPBgeFnaayuMgc2KRXk,12533
+jaxsim/parsers/rod/utils.py,sha256=SJg9t-nnItmo_gfxkfQ6uILNYN_BddcAHckyQkTDAjY,5902
 jaxsim/rbda/__init__.py,sha256=HLwxeU-IxaRpFGUCSQv-LDv20JHTt3Xj7ELiRbRieS8,319
 jaxsim/rbda/aba.py,sha256=0OoCzHhf1v-qqr1y5PIrD7_mPwAlid0fjXxUrIa5E_s,9118
 jaxsim/rbda/collidable_points.py,sha256=4ZNJbEj2nEi15jBLR-GNbdaqKgkN58FBgqd_TXupEgg,4948
@@ -55,10 +55,10 @@ jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
 jaxsim/terrain/terrain.py,sha256=q0xkWqEShVq-p1j2abTLZq8sEhjyJwquxQKm80PaHhM,2161
 jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
 jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
-jaxsim/utils/jaxsim_dataclass.py,sha256=T8nZolb563mYJD7YvTmYW-z-hpkEiaK3cDbYHqwtqRc,11347
+jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim-0.2.1.dev17.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev17.dist-info/METADATA,sha256=ZuAZDrgEcUOjS01EwlNvXEUG2D04PJqtR2wvC6QnXeA,9726
-jaxsim-0.2.1.dev17.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev17.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev17.dist-info/RECORD,,
+jaxsim-0.2.1.dev38.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev38.dist-info/METADATA,sha256=63Xh0podujOMHsBQHfMyNCFYIXQPovXz6I6oWQ-IGd4,9744
+jaxsim-0.2.1.dev38.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev38.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev38.dist-info/RECORD,,