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

jaxsim/parsers/kinematic_graph.py

@@ -1,263 +1,397 @@
+from __future__ import annotations
+
 import copy
 import dataclasses
 import functools
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Tuple,
-    Union,
-)
+from collections.abc import Callable, Iterable, Iterator, Sequence
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
 
+import jaxsim.utils
 from jaxsim import logging
 from jaxsim.utils import Mutability
 
-from . import descriptions
+from .descriptions.joint import JointDescription, JointType
+from .descriptions.link import LinkDescription
 
 
-class RootPose(NamedTuple):
+@dataclasses.dataclass
+class RootPose:
     """
     Represents the root pose in a kinematic graph.
 
     Attributes:
-        root_position (npt.NDArray): A NumPy array of shape (3,) representing the root's position.
-        root_quaternion (npt.NDArray): A NumPy array of shape (4,) representing the root's quaternion.
+        root_position: The 3D position of the root link of the graph.
+        root_quaternion:
+            The quaternion representing the rotation of the root link of the graph.
+
+    Note:
+        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(
+            (
+                HashedNumpyArray.hash_of_array(self.root_position),
+                HashedNumpyArray.hash_of_array(self.root_quaternion),
+            )
+        )
+
+    def __eq__(self, other: RootPose) -> bool:
+
+        if not isinstance(other, RootPose):
+            return False
+
+        if not np.allclose(self.root_position, other.root_position):
+            return False
+
+        if not np.allclose(self.root_quaternion, other.root_quaternion):
+            return False
 
-    def __eq__(self, other):
-        return (self.root_position == other.root_position).all() and (
-            self.root_quaternion == other.root_quaternion
-        ).all()
+        return True
 
 
 @dataclasses.dataclass(frozen=True)
-class KinematicGraph:
+class KinematicGraph(Sequence[LinkDescription]):
     """
-    Represents a kinematic graph of links and joints.
-
-    Args:
-        root (descriptions.LinkDescription): The root link of the kinematic graph.
-        frames (List[descriptions.LinkDescription]): A list of frame links in the graph.
-        joints (List[descriptions.JointDescription]): A list of joint descriptions in the graph.
-        root_pose (RootPose): The root pose of the graph.
-        transform_cache (Dict[str, npt.NDArray]): A dictionary to cache transformation matrices.
-        extra_info (Dict[str, Any]): Additional information associated with the graph.
+    Class storing a kinematic graph having links as nodes and joints as edges.
 
     Attributes:
-        links_dict (Dict[str, descriptions.LinkDescription]): A dictionary mapping link names to link descriptions.
-        frames_dict (Dict[str, descriptions.LinkDescription]): A dictionary mapping frame names to frame link descriptions.
-        joints_dict (Dict[str, descriptions.JointDescription]): A dictionary mapping joint names to joint descriptions.
-        joints_connection_dict (Dict[Tuple[str, str], descriptions.JointDescription]): A dictionary mapping pairs of parent and child link names to joint descriptions.
+        root: The root node of the kinematic graph.
+        frames: List of frames rigidly attached to the graph nodes.
+        joints: List of joints connecting the graph nodes.
+        root_pose: The pose of the kinematic graph's root.
     """
 
-    root: descriptions.LinkDescription
-    frames: List[descriptions.LinkDescription] = dataclasses.field(default_factory=list)
-    joints: List[descriptions.JointDescription] = dataclasses.field(
-        default_factory=list
+    root: LinkDescription
+    frames: list[LinkDescription] = dataclasses.field(
+        default_factory=list, hash=False, compare=False
+    )
+    joints: list[JointDescription] = dataclasses.field(
+        default_factory=list, hash=False, compare=False
     )
 
     root_pose: RootPose = dataclasses.field(default_factory=RootPose)
 
-    transform_cache: Dict[str, npt.NDArray] = dataclasses.field(
-        repr=False, init=False, compare=False, default_factory=dict
+    # Private attribute storing optional additional info.
+    _extra_info: dict[str, Any] = dataclasses.field(
+        default_factory=dict, repr=False, hash=False, compare=False
     )
 
-    extra_info: Dict[str, Any] = dataclasses.field(
-        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[JointDescription] = dataclasses.field(
+        default_factory=list, repr=False, hash=False, compare=False
     )
 
     @functools.cached_property
-    def links_dict(self) -> Dict[str, descriptions.LinkDescription]:
+    def links_dict(self) -> dict[str, LinkDescription]:
+        """
+        Get a dictionary of links indexed by their name.
+        """
         return {l.name: l for l in iter(self)}
 
     @functools.cached_property
-    def frames_dict(self) -> Dict[str, descriptions.LinkDescription]:
+    def frames_dict(self) -> dict[str, LinkDescription]:
+        """
+        Get a dictionary of frames indexed by their name.
+        """
         return {f.name: f for f in self.frames}
 
     @functools.cached_property
-    def joints_dict(self) -> Dict[str, descriptions.JointDescription]:
+    def joints_dict(self) -> dict[str, JointDescription]:
+        """
+        Get a dictionary of joints indexed by their name.
+        """
         return {j.name: j for j in self.joints}
 
     @functools.cached_property
     def joints_connection_dict(
         self,
-    ) -> Dict[Tuple[str, str], descriptions.JointDescription]:
-        return {(j.parent.name, j.child.name): j for j in self.joints}
-
-    def __post_init__(self):
+    ) -> dict[tuple[str, str], JointDescription]:
         """
-        Post-initialization method to set various properties and validate the kinematic graph.
+        Get a dictionary of joints indexed by the tuple (parent, child) link names.
         """
-        # Assign the link index traversing the graph with BFS.
-        # Here we assume the model is fixed-base, therefore the base link will
-        # have index 0. We will deal with the floating base in a later stage,
-        # when this Model object is converted to the physics model.
+        return {(j.parent.name, j.child.name): j for j in self.joints}
+
+    def __post_init__(self) -> None:
+
+        # Assign the link index by traversing the graph with BFS.
+        # Here we assume the model being fixed-base, therefore the base link will
+        # have index 0. We will deal with the floating base in a later stage.
         for index, link in enumerate(self):
             link.mutable(validate=False).index = index
 
-        # Order frames with their name
+        # Get the names of the links, frames, and joints.
+        link_names = [l.name for l in self]
+        frame_names = [f.name for f in self.frames]
+        joint_names = [j.name for j in self.joints]
+
+        # Make sure that they are unique.
+        assert len(link_names) == len(set(link_names))
+        assert len(frame_names) == len(set(frame_names))
+        assert len(joint_names) == len(set(joint_names))
+        assert set(link_names).isdisjoint(set(frame_names))
+        assert set(link_names).isdisjoint(set(joint_names))
+
+        # Order frames with their name.
         super().__setattr__("frames", sorted(self.frames, key=lambda f: f.name))
 
         # Assign the frame index following the name-based indexing.
-        # Also here, we assume the model is fixed-base, therefore the first frame will
-        # have last_link_idx + 1. These frames are not part of the physics model.
+        # We assume the model being fixed-base, therefore the first frame will
+        # have last_link_idx + 1.
         for index, frame in enumerate(self.frames):
-            frame.index = index + len(self.link_names())
+            with frame.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+                frame.index = int(index + len(self.link_names()))
 
-        # Number joints so that their index matches their child link index
+        # Number joints so that their index matches their child link index.
+        # Therefore, the first joint has index 1.
         links_dict = {l.name: l for l in iter(self)}
         for joint in self.joints:
             with joint.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
                 joint.index = links_dict[joint.child.name].index
 
-        # Check that joint indices are unique
+        # Check that joint indices are unique.
         assert len([j.index for j in self.joints]) == len(
             {j.index for j in self.joints}
         )
 
-        # Order joints with their indices
+        # Order joints with their indices.
         super().__setattr__("joints", sorted(self.joints, key=lambda j: j.index))
 
     @staticmethod
     def build_from(
-        links: List[descriptions.LinkDescription],
-        joints: List[descriptions.JointDescription],
+        links: list[LinkDescription],
+        joints: list[JointDescription],
+        frames: list[LinkDescription] | None = None,
         root_link_name: str | None = None,
         root_pose: RootPose = RootPose(),
-    ) -> "KinematicGraph":
+    ) -> KinematicGraph:
         """
-        Build a KinematicGraph from a list of links and joints.
+        Build a KinematicGraph from links, joints, and frames.
 
         Args:
-            links (List[descriptions.LinkDescription]): A list of link descriptions.
-            joints (List[descriptions.JointDescription]): A list of joint descriptions.
-            root_link_name (str, optional): The name of the root link. If not provided, it's assumed to be the first link's name.
-            root_pose (RootPose, optional): The root pose of the kinematic graph.
+            links: A list of link descriptions.
+            joints: A list of joint descriptions.
+            frames: A list of frame descriptions.
+            root_link_name:
+                The name of the root link. If not provided, it's assumed to be the
+                first link's name.
+            root_pose: The root pose of the kinematic graph.
 
         Returns:
-            KinematicGraph: The constructed kinematic graph.
+            The resulting kinematic graph.
         """
+
+        # Consider the first link as the root link if not provided.
         if root_link_name is None:
             root_link_name = links[0].name
+            logging.debug(msg=f"Assuming '{root_link_name}' as the root link")
 
         # 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
+        # Note that the pose of the frames is not updated; it is 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_links,
+            unconnected_joints,
+            unconnected_frames,
+        ) = KinematicGraph._create_graph(
+            links=links, joints=joints, root_link_name=root_link_name, frames=frames
         )
 
-        for frame in graph_frames:
-            logging.warning(msg=f"Ignoring unconnected link / frame: '{frame.name}'")
+        for link in unconnected_links:
+            logging.warning(msg=f"Ignoring unconnected link: '{link.name}'")
+
+        for joint in unconnected_joints:
+            logging.warning(msg=f"Ignoring unconnected joint: '{joint.name}'")
+
+        for frame in unconnected_frames:
+            logging.warning(msg=f"Ignoring unconnected frame: '{frame.name}'")
 
         return KinematicGraph(
-            root=graph_root_node, joints=graph_joints, frames=[], root_pose=root_pose
+            root=graph_root_node,
+            joints=graph_joints,
+            frames=graph_frames,
+            root_pose=root_pose,
+            _joints_removed=unconnected_joints,
         )
 
     @staticmethod
-    def create_graph(
-        links: List[descriptions.LinkDescription],
-        joints: List[descriptions.JointDescription],
+    def _create_graph(
+        links: list[LinkDescription],
+        joints: list[JointDescription],
         root_link_name: str,
-    ) -> Tuple[
-        descriptions.LinkDescription,
-        List[descriptions.JointDescription],
-        List[descriptions.LinkDescription],
+        frames: list[LinkDescription] | None = None,
+    ) -> tuple[
+        LinkDescription,
+        list[JointDescription],
+        list[LinkDescription],
+        list[LinkDescription],
+        list[JointDescription],
+        list[LinkDescription],
     ]:
         """
-        Create a kinematic graph from lists of links and joints.
+        Low-level creator of kinematic graph components.
 
         Args:
-            links (List[descriptions.LinkDescription]): A list of link descriptions.
-            joints (List[descriptions.JointDescription]): A list of joint descriptions.
-            root_link_name (str): The name of the root link.
+            links: A list of parsed link descriptions.
+            joints: A list of parsed joint descriptions.
+            root_link_name: The name of the root link used as root node of the graph.
+            frames: A list of parsed frame descriptions.
 
         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 of the graph (defining the entire kinematic
+            tree by iterating on its child nodes), the list of joints representing the
+            actual graph edges, the list of frames rigidly attached to the graph nodes,
+            the list of unconnected links, the list of unconnected joints, and the list
+            of unconnected frames.
         """
 
-        # Create a dict that maps link name to the link, for easy retrieval
-        links_dict: Dict[str, descriptions.LinkDescription] = {
+        # Create a dictionary that maps the link name to the link, for easy retrieval.
+        links_dict: dict[str, LinkDescription] = {
             l.name: l.mutable(validate=False) for l in links
         }
 
+        # Create an empty list of frames if not provided.
+        frames = frames if frames is not None else []
+
+        # Create a dictionary that maps the frame name to the frame, for easy retrieval.
+        frames_dict = {frame.name: frame for frame in frames}
+
+        # Check that our parser correctly resolved the frame's parent to be a link.
+        for frame in frames:
+            assert frame.parent.name != "", frame
+            assert frame.parent.name is not None, frame
+            assert frame.parent.name != "__model__", frame
+            assert frame.parent.name not in frames_dict, frame
+
+        # ===========================================================
+        # Populate the kinematic graph with links, joints, and frames
+        # ===========================================================
+
+        # Check the existence of the root link.
         if root_link_name not in links_dict:
             raise ValueError(root_link_name)
 
-        # Reset the connections of the root link
+        # Reset the connections of the root link.
         for link in links_dict.values():
-            link.children = []
+            link.children = tuple()
 
-        # Couple links and joints creating the final kinematic graph
+        # Couple links and joints creating the kinematic graph.
         for joint in joints:
-            # Get the parent and child links of the joint
+
+            # Get the parent and child links of the joint.
             parent_link = links_dict[joint.parent.name]
             child_link = links_dict[joint.child.name]
 
             assert child_link.name == joint.child.name
             assert parent_link.name == joint.parent.name
 
-            # Assign link parent
+            # Assign link's parent.
             child_link.parent = parent_link
 
-            # Assign link children and make sure they are unique
+            # Assign link's children and make sure they are unique.
             if child_link.name not in {l.name for l in parent_link.children}:
-                parent_link.children.append(child_link)
+                with parent_link.mutable_context(Mutability.MUTABLE_NO_VALIDATION):
+                    parent_link.children = (*parent_link.children, child_link)
 
-        # Collect all the links of the kinematic graph
+        # Collect all the links of the kinematic graph.
         all_links_in_graph = list(
             KinematicGraph.breadth_first_search(root=links_dict[root_link_name])
         )
+
+        # Get the names of all links in the kinematic graph.
         all_link_names_in_graph = [l.name for l in all_links_in_graph]
 
-        # Collect all the joints not part of the kinematic graph
-        removed_joints = [
-            j
-            for j in joints
-            if not {j.parent.name, j.child.name}.issubset(all_link_names_in_graph)
+        # Collect all the joints of the kinematic graph.
+        all_joints_in_graph = [
+            joint
+            for joint in joints
+            if joint.parent.name in all_link_names_in_graph
+            and joint.child.name in all_link_names_in_graph
         ]
 
-        for removed_joint in removed_joints:
-            msg = "Joint '{}' has been removed for the graph because unconnected"
-            logging.info(msg=msg.format(removed_joint.name))
+        # Get the names of all joints in the kinematic graph.
+        all_joint_names_in_graph = [j.name for j in all_joints_in_graph]
 
-        # Store as frames all the links that are not part of the kinematic graph
-        frames = list(set(links) - set(all_links_in_graph))
+        # Collect all the frames of the kinematic graph.
+        # Note: our parser ensures that the parent of a frame is not another frame.
+        all_frames_in_graph = [
+            frame for frame in frames if frame.parent.name in all_link_names_in_graph
+        ]
 
-        # Update the frames. In particular, reset their children. The other properties
-        # are kept as they are, and it's caller responsibility to update them if needed.
-        for frame in frames:
-            frame.children = []
-            msg = f"Link '{frame.name}' became a frame"
-            logging.info(msg=msg)
+        # Get the names of all frames in the kinematic graph.
+        all_frames_names_in_graph = [f.name for f in all_frames_in_graph]
+
+        # ============================
+        # Collect unconnected elements
+        # ============================
+
+        # Collect all the joints that are not part of the kinematic graph.
+        removed_joints = [j for j in joints if j.name not in all_joint_names_in_graph]
+
+        for joint in removed_joints:
+            msg = "Joint '{}' is unconnected and it will be removed"
+            logging.debug(msg=msg.format(joint.name))
+
+        # Collect all the links that are not part of the kinematic graph.
+        unconnected_links = [l for l in links if l.name not in all_link_names_in_graph]
+
+        # Update the unconnected links by removing their children. The other properties
+        # are left untouched, it's caller responsibility to post-process them if needed.
+        for link in unconnected_links:
+            link.children = tuple()
+            msg = "Link '{}' won't be part of the kinematic graph because unconnected"
+            logging.debug(msg=msg.format(link.name))
+
+        # Collect all the frames that are not part of the kinematic graph.
+        unconnected_frames = [
+            f for f in frames if f.name not in all_frames_names_in_graph
+        ]
+
+        for frame in unconnected_frames:
+            msg = "Frame '{}' won't be part of the kinematic graph because unconnected"
+            logging.debug(msg=msg.format(frame.name))
 
         return (
             links_dict[root_link_name].mutable(mutable=False),
             list(set(joints) - set(removed_joints)),
-            frames,
+            all_frames_in_graph,
+            unconnected_links,
+            list(set(removed_joints)),
+            unconnected_frames,
         )
 
-    def reduce(self, considered_joints: List[str]) -> "KinematicGraph":
+    def reduce(self, considered_joints: Sequence[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
 
@@ -268,11 +402,11 @@ class KinematicGraph:
 
         # Return early if there is no action to take
         if len(joint_names_to_remove) == 0:
-            logging.info(f"The kinematic graph doesn't need to be reduced")
+            logging.info("The kinematic graph doesn't need to be reduced")
             return copy.deepcopy(self)
 
         # Check if all considered joints are part of the full kinematic graph
-        if len(set(considered_joints) - set(j.name for j in full_graph.joints)) != 0:
+        if len(set(considered_joints) - {j.name for j in full_graph.joints}) != 0:
             extra_j = set(considered_joints) - {j.name for j in full_graph.joints}
             msg = f"Not all joints to consider are part of the graph ({{{extra_j}}})"
             raise ValueError(msg)
@@ -281,6 +415,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
@@ -320,7 +457,7 @@ class KinematicGraph:
                 msg.format(
                     link_to_remove.name,
                     self.joints_connection_dict[
-                        (parent_of_link_to_remove.name, link_to_remove.name)
+                        parent_of_link_to_remove.name, link_to_remove.name
                     ].name,
                     parent_of_link_to_remove.name,
                 )
@@ -329,14 +466,14 @@ 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
                 ),
             )
 
             # Pop the original two links from the dictionary...
-            links_dict.pop(link_to_remove.name)
-            links_dict.pop(parent_of_link_to_remove.name)
+            _ = links_dict.pop(link_to_remove.name)
+            _ = links_dict.pop(parent_of_link_to_remove.name)
 
             # ... and insert the lumped link (having the same name of the parent)
             links_dict[lumped_link.name] = lumped_link
@@ -346,11 +483,13 @@ class KinematicGraph:
             links_dict[link_to_remove.name] = lumped_link
 
             # As a consequence of the back-insertion, we need to adjust the resulting
-            # lumped link of links that have been removed previously
+            # lumped link of links that have been removed previously.
+            # Note: in the dictionary, only items whose key is not matching value.name
+            #       are links that have been removed.
             for previously_removed_link_name in {
-                k
-                for k, v in links_dict.items()
-                if k != v.name and v.name == link_to_remove.name
+                link_name
+                for link_name, link in links_dict.items()
+                if link_name != link.name and link.name == link_to_remove.name
             }:
                 links_dict[previously_removed_link_name] = lumped_link
 
@@ -370,7 +509,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,135 +535,114 @@ 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(
+        (
+            reduced_root_node,
+            reduced_joints,
+            reduced_frames,
+            unconnected_links,
+            unconnected_joints,
+            unconnected_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,
         )
 
-        # Create the reduced graph
+        assert {f.name for f in self.frames}.isdisjoint(
+            {f.name for f in unconnected_frames + reduced_frames}
+        )
+
+        for link in unconnected_links:
+            logging.debug(msg=f"Link '{link.name}' is unconnected and became a frame")
+
+        # Create the reduced graph.
         reduced_graph = KinematicGraph(
             root=reduced_root_node,
             joints=reduced_joints,
-            frames=reduced_frames,
+            frames=self.frames + unconnected_links + 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]
+            ),
         )
 
         # ================================================================
         # 4. Resolve the pose of the frames wrt their reduced graph parent
         # ================================================================
 
-        # Update frames properties using the transforms from the full graph
-        for frame in reduced_graph.frames:
-            # Get the link in which the removed link was lumped into
-            new_parent_link = links_dict[frame.name]
+        # Build a new object to compute FK on the reduced graph.
+        fk_reduced = KinematicGraphTransforms(graph=reduced_graph)
 
-            msg = f"New parent of frame '{frame.name}' is '{new_parent_link.name}'"
-            logging.info(msg)
+        # We need to adjust the pose of the frames since their parent link
+        # could have been removed by the reduction process.
+        for frame in reduced_graph.frames:
 
-            # Update the connection of the frame
-            frame.parent = new_parent_link
-            frame.pose = full_graph.relative_transform(
-                relative_to=new_parent_link.name, name=frame.name
+            # Always find the real parent link of the frame
+            name_of_new_parent_link = fk_reduced.find_parent_link_of_frame(
+                name=frame.name
             )
+            assert name_of_new_parent_link in reduced_graph, name_of_new_parent_link
 
-            # Update frame data
-            frame.mass = 0.0
-            frame.inertia = np.zeros_like(frame.inertia)
+            # Notify the user if the parent link has changed.
+            if name_of_new_parent_link != frame.parent.name:
+                msg = "New parent of frame '{}' is '{}'"
+                logging.debug(msg=msg.format(frame.name, name_of_new_parent_link))
 
-        # Return the reduced graph
+            # Always recompute the pose of the frame, and set zero inertial params.
+            with frame.mutable_context(jaxsim.utils.Mutability.MUTABLE_NO_VALIDATION):
+
+                # Update kinematic parameters of the frame.
+                # Note that here we compute the transform using the FK object of the
+                # full model, so that we are sure that the kinematic is not altered.
+                frame.pose = fk.relative_transform(
+                    relative_to=name_of_new_parent_link, name=frame.name
+                )
+
+                # Update the parent link such that the pose is expressed in its frame.
+                frame.parent = reduced_graph.links_dict[name_of_new_parent_link]
+
+                # Update dynamic parameters of the frame.
+                frame.mass = 0.0
+                frame.inertia = np.zeros_like(frame.inertia)
+
+        # Return the reduced graph.
         return reduced_graph
 
-    def link_names(self) -> List[str]:
+    def link_names(self) -> list[str]:
         """
-        Get the names of all links in the kinematic graph.
+        Get the names of all links in the kinematic graph (i.e. the nodes).
 
         Returns:
-            List[str]: A list of link names.
+            The list of link names.
         """
         return list(self.links_dict.keys())
 
-    def joint_names(self) -> List[str]:
+    def joint_names(self) -> list[str]:
         """
-        Get the names of all joints in the kinematic graph.
+        Get the names of all joints in the kinematic graph (i.e. the edges).
 
         Returns:
-            List[str]: A list of joint names.
+            The list of joint names.
         """
         return list(self.joints_dict.keys())
 
-    def frame_names(self) -> List[str]:
+    def frame_names(self) -> list[str]:
         """
         Get the names of all frames in the kinematic graph.
 
         Returns:
-            List[str]: A list of frame names.
+            The list of frame names.
         """
-        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
-        )
+        return list(self.frames_dict.keys())
 
     def print_tree(self) -> None:
         """
         Print the tree structure of the kinematic graph.
         """
+
         import pptree
 
         root_node = self.root
@@ -536,24 +654,37 @@ class KinematicGraph:
             horizontal=True,
         )
 
+    @property
+    def joints_removed(self) -> list[JointDescription]:
+        """
+        Get the list of joints removed during the graph reduction.
+
+        Returns:
+            The list of removed joints.
+        """
+
+        return self._joints_removed
+
     @staticmethod
     def breadth_first_search(
-        root: descriptions.LinkDescription,
-        sort_children: Optional[Callable[[Any], Any]] = lambda link: link.name,
-    ) -> Iterable[descriptions.LinkDescription]:
+        root: LinkDescription,
+        sort_children: Callable[[Any], Any] | None = lambda link: link.name,
+    ) -> Iterable[LinkDescription]:
         """
         Perform a breadth-first search (BFS) traversal of the kinematic graph.
 
         Args:
-            root (descriptions.LinkDescription): The root link for BFS.
-            sort_children (Optional[Callable[[Any], Any]]): A function to sort children of a node.
+            root: The root link for BFS.
+            sort_children: A function to sort children of a node.
 
         Yields:
-            Iterable[descriptions.LinkDescription]: An iterable of link descriptions.
+            The links in the kinematic graph in BFS order.
         """
+
+        # Initialize the queue with the root node.
         queue = [root]
 
-        # We assume that nodes have unique names, and mark a link as visited using
+        # We assume that nodes have unique names and mark a link as visited using
         # its name. This speeds up considerably object comparison.
         visited = []
         visited.append(root.name)
@@ -561,11 +692,14 @@ class KinematicGraph:
         yield root
 
         while len(queue) > 0:
+
+            # Extract the first element of the queue.
             l = queue.pop(0)
 
             # Note: sorting the links with their name so that the order of children
-            # insertion does not matter when assigning the link index
+            # insertion does not matter when assigning the link index.
             for child in sorted(l.children, key=sort_children):
+
                 if child.name in visited:
                     continue
 
@@ -574,25 +708,29 @@ class KinematicGraph:
 
                 yield child
 
-    def __iter__(self) -> Iterable[descriptions.LinkDescription]:
+    # =================
+    # Sequence protocol
+    # =================
+
+    def __iter__(self) -> Iterator[LinkDescription]:
         yield from KinematicGraph.breadth_first_search(root=self.root)
 
-    def __reversed__(self) -> Iterable[descriptions.LinkDescription]:
+    def __reversed__(self) -> Iterable[LinkDescription]:
         yield from reversed(list(iter(self)))
 
     def __len__(self) -> int:
         return len(list(iter(self)))
 
-    def __contains__(self, item: Union[str, descriptions.LinkDescription]) -> bool:
+    def __contains__(self, item: str | LinkDescription) -> bool:
         if isinstance(item, str):
             return item in self.link_names()
 
-        if isinstance(item, descriptions.LinkDescription):
+        if isinstance(item, LinkDescription):
             return item in set(iter(self))
 
         raise TypeError(type(item).__name__)
 
-    def __getitem__(self, key: Union[int, str]) -> descriptions.LinkDescription:
+    def __getitem__(self, key: int | str) -> LinkDescription:
         if isinstance(key, str):
             if key not in self.link_names():
                 raise KeyError(key)
@@ -606,3 +744,242 @@ class KinematicGraph:
             return list(iter(self))[key]
 
         raise TypeError(type(key).__name__)
+
+    def count(self, value: LinkDescription) -> int:
+        """
+        Count the occurrences of a link in the kinematic graph.
+        """
+        return list(iter(self)).count(value)
+
+    def index(self, value: LinkDescription, start: int = 0, stop: int = -1) -> int:
+        """
+        Find the index of a link in the kinematic graph.
+        """
+        return list(iter(self)).index(value, start, stop)
+
+
+# ====================
+# Other useful classes
+# ====================
+
+
+@dataclasses.dataclass(frozen=True)
+class KinematicGraphTransforms:
+    """
+    Class to compute forward kinematics on a kinematic graph.
+
+    Attributes:
+        graph: The kinematic graph on which to compute forward kinematics.
+    """
+
+    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:
+        """
+        Get the initial joint positions of the kinematic graph.
+        """
+
+        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, strict=True):
+            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]
+            assert joint.name == 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_axis=joint.axis,
+                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: JointType,
+        joint_axis: npt.NDArray,
+        joint_position: float | None = None,
+    ) -> npt.NDArray:
+        """
+        Compute the SE(3) transform from the predecessor to the successor frame.
+
+        Args:
+            joint_type: The type of the joint.
+            joint_axis: The axis of the joint.
+            joint_position: The position of the joint.
+
+        Returns:
+            The 4x4 transform matrix from the predecessor to the successor frame.
+        """
+
+        import jaxsim.math
+
+        return np.array(
+            jaxsim.math.supported_joint_motion(joint_type, joint_position, joint_axis)[
+                0
+            ]
+        )
+
+    def find_parent_link_of_frame(self, name: str) -> str:
+        """
+        Find the parent link of a frame.
+
+        Args:
+            name: The name of the frame.
+
+        Returns:
+            The name of the parent link of the frame.
+        """
+
+        try:
+            frame = self.graph.frames_dict[name]
+        except KeyError as e:
+            raise ValueError(f"Frame '{name}' not found in the kinematic graph") from e
+
+        match frame.parent.name:
+            case parent_name if parent_name in self.graph.links_dict:
+                return parent_name
+
+            case parent_name if parent_name in self.graph.frames_dict:
+                return self.find_parent_link_of_frame(name=parent_name)
+
+            case _:
+                msg = f"Failed to find parent element of frame '{name}' with name '{frame.parent.name}'"
+                raise RuntimeError(msg)