From aa092f5e629c592785ef9ad00960cc58b708ad7a Mon Sep 17 00:00:00 2001
From: Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
Date: Fri, 5 Jan 2024 18:23:39 +0100
Subject: [PATCH] [wip] Reduce model tree

---
 src/adam/model/model.py                  | 35 +++++++++++++++++++++++
 src/adam/model/std_factories/std_link.py | 23 +++++++++++++++
 src/adam/model/tree.py                   | 36 ++++++++++++++++++++++++
 3 files changed, 94 insertions(+)

diff --git a/src/adam/model/model.py b/src/adam/model/model.py
index 3421cfdf..5e1369c5 100644
--- a/src/adam/model/model.py
+++ b/src/adam/model/model.py
@@ -80,6 +80,41 @@ def build(factory: ModelFactory, joints_name_list: List[str] = None) -> "Model":
             floating_base=floating_base,
         )
 
+    def reduce(self, joints_name_list: List[str]) -> "Model":
+        """reduce the model to a subset of joints
+
+        Args:
+            joints_name_list (List[str]): the list of the joints to keep
+
+        Returns:
+            Model: the reduced model
+        """
+
+        # check if the joints in the list are in the model
+        for joint_str in joints_name_list:
+            if joint_str not in self.joints.keys():
+                raise ValueError(
+                    f"{joint_str} is not in the robot model. Check the joints_name_list"
+                )
+
+        tree = self.tree.reduce(joints_name_list)
+
+        # update nodes dict
+        links = {link.name: link for link in tree.graph}
+        frames = {frame.name: frame for frame in tree.graph}
+        joints = {joint.name: joint for joint in tree.graph}
+
+        return Model(
+            name=self.name,
+            links=links,
+            frames=frames,
+            joints=joints,
+            tree=tree,
+            NDoF=len(joints_name_list),
+            actuated_joints=joints_name_list,
+            floating_base=self.floating_base,
+        )
+
     def get_joints_chain(self, root: str, target: str) -> List[Joint]:
         """generate the joints chains from a link to a link
 
diff --git a/src/adam/model/std_factories/std_link.py b/src/adam/model/std_factories/std_link.py
index cefac67c..6ecefac2 100644
--- a/src/adam/model/std_factories/std_link.py
+++ b/src/adam/model/std_factories/std_link.py
@@ -44,3 +44,26 @@ def homogeneous(self) -> npt.ArrayLike:
             self.inertial.origin.xyz,
             self.inertial.origin.rpy,
         )
+
+    def lump(self, other: "StdLink", relative_transform: npt.ArrayLike) -> "StdLink":
+        """lump two links together
+
+        Args:
+            other (StdLink): the other link
+            relative_transform (npt.ArrayLike): the transform between the two links
+
+        Returns:
+            StdLink: the lumped link
+        """
+        other_inertia = (
+            relative_transform.T @ other.spatial_inertia() @ relative_transform
+        )
+
+        # lump the inertial properties
+        lumped_mass = self.inertial.mass + other.inertial.mass
+        lumped_inertia = self.spatial_inertia() + other_inertia
+
+        self.mass = lumped_mass
+        self.inertia = lumped_inertia
+
+        return self
diff --git a/src/adam/model/tree.py b/src/adam/model/tree.py
index 8d99d328..45c02a25 100644
--- a/src/adam/model/tree.py
+++ b/src/adam/model/tree.py
@@ -71,6 +71,42 @@ def build_tree(links: List[Link], joints: List[Joint]) -> "Tree":
             raise ValueError("The model has more than one root link")
         return Tree(nodes, root_link[0])
 
+    def reduce(self, considered_joint_names: List[str]) -> "Tree":
+        """reduces the tree to the considered joints
+
+        Args:
+            considered_joint_names (List[str]): the list of the considered joints
+
+        Returns:
+            Tree: the reduced tree
+        """
+        new_graph = {}
+
+        # find the nodes that are not connected to the considered joints
+        nodes_to_lump = list(
+            {
+                node.name
+                for node in self.graph.values()
+                for joint in node.arcs
+                if joint.name not in considered_joint_names
+            }
+        )
+
+        # lump the inertial properties
+        for node in self.graph.values():
+            if node.name in nodes_to_lump:
+                lumped_node = node.parent.lump(
+                    other=node.link,
+                    relative_transform=node.parent_arc.spatial_transform(0),
+                )
+
+                # remove the node from the graph
+                self.graph.pop(node.name)
+
+        # TODO: WIP
+
+        return Tree(new_graph, self.root)
+
     def print(self, root) -> str:
         """prints the tree