From 2feb8b0e95a00671d71123ba758e8781f32a34bd Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 22 Aug 2024 10:29:50 +0200 Subject: [PATCH] Cache commonly used properties --- src/rod/kinematics/kinematic_tree.py | 21 ++++++++++------- src/rod/kinematics/tree_transforms.py | 6 ++--- src/rod/sdf/link.py | 1 + src/rod/sdf/model.py | 17 ++++++++++---- src/rod/urdf/exporter.py | 22 ++++++++--------- src/rod/utils/frame_convention.py | 34 +++++++++++++-------------- src/rod/utils/resolve_frames.py | 8 +++---- tests/test_urdf_exporter.py | 2 +- tests/test_urdf_parsing.py | 16 ++++++------- 9 files changed, 69 insertions(+), 58 deletions(-) diff --git a/src/rod/kinematics/kinematic_tree.py b/src/rod/kinematics/kinematic_tree.py index 168237c..b956544 100644 --- a/src/rod/kinematics/kinematic_tree.py +++ b/src/rod/kinematics/kinematic_tree.py @@ -38,12 +38,15 @@ def __post_init__(self): self.joints.sort(key=lambda j: j.index) self.frames.sort(key=lambda f: f.index) + @functools.cached_property def link_names(self) -> list[str]: return [node.name() for node in self] + @functools.cached_property def frame_names(self) -> list[str]: return [frame.name() for frame in self.frames] + @functools.cached_property def joint_names(self) -> list[str]: return [joint.name() for joint in self.joints] @@ -75,7 +78,7 @@ def build(model: rod.Model, is_top_level: bool = True) -> KinematicTree: # Create a dict mapping link names to tree nodes, for easy retrieval. nodes_links_dict: dict[str, DirectedTreeNode] = { # Add one node for each link of the model - **{link.name: DirectedTreeNode(_source=link) for link in model.links()}, + **{link.name: DirectedTreeNode(_source=link) for link in model.links}, # Add special world node, that will become a frame later TreeFrame.WORLD: DirectedTreeNode( _source=rod.Link( @@ -103,7 +106,7 @@ def build(model: rod.Model, is_top_level: bool = True) -> KinematicTree: # Create a dict mapping frame names to frame nodes, for easy retrieval. nodes_frames_dict: dict[str, TreeFrame] = { # Add a frame node for each frame in the model - **{frame.name: TreeFrame(_source=frame) for frame in model.frames()}, + **{frame.name: TreeFrame(_source=frame) for frame in model.frames}, # Add implicit frames used in the SDF specification (__model__). # The following frames are attached to the first link found in the model # description and never moved, so that all elements expressing their pose @@ -123,7 +126,7 @@ def build(model: rod.Model, is_top_level: bool = True) -> KinematicTree: ) == (len(nodes_links_dict) + len(nodes_frames_dict)) # Use joints to connect nodes by defining their parent and children - for joint in model.joints(): + for joint in model.joints: if joint.child == TreeFrame.WORLD: msg = f"A joint cannot have '{TreeFrame.WORLD}' as child" raise RuntimeError(msg) @@ -161,30 +164,30 @@ def build(model: rod.Model, is_top_level: bool = True) -> KinematicTree: # Get all the joints part of the kinematic tree ... joints_in_tree_names = [ j.name - for j in model.joints() + for j in model.joints if {j.parent, j.child}.issubset(all_node_names_in_tree) ] - joints_in_tree = [j for j in model.joints() if j.name in joints_in_tree_names] + joints_in_tree = [j for j in model.joints if j.name in joints_in_tree_names] # ... and those that are not joints_not_in_tree = [ - j for j in model.joints() if j.name not in joints_in_tree_names + j for j in model.joints if j.name not in joints_in_tree_names ] # A valid rod.Model does not have any dangling link and any unconnected joints. # Here we check that the rod.Model contains a valid tree representation. found_num_extra_joints = len(joints_not_in_tree) - expected_num_extra_joints = 1 if model.is_fixed_base() else 0 + expected_num_extra_joints = 1 if model.is_fixed_base else 0 if found_num_extra_joints != expected_num_extra_joints: - if model.is_fixed_base() and found_num_extra_joints == 0: + if model.is_fixed_base and found_num_extra_joints == 0: raise RuntimeError("Failed to find joint connecting the model to world") unexpected_joint_names = [j.name for j in joints_not_in_tree] raise RuntimeError(f"Found unexpected joints: {unexpected_joint_names}") # Handle connection to world of fixed-base models - if model.is_fixed_base(): + if model.is_fixed_base: assert len(joints_not_in_tree) == 1 world_to_base_joint = joints_not_in_tree[0] diff --git a/src/rod/kinematics/tree_transforms.py b/src/rod/kinematics/tree_transforms.py index 00f4e8a..627c63e 100644 --- a/src/rod/kinematics/tree_transforms.py +++ b/src/rod/kinematics/tree_transforms.py @@ -50,7 +50,7 @@ def _compute_transform(self, name: str) -> npt.NDArray: assert relative_to in {None, ""}, (relative_to, name) return self.kinematic_tree.model.pose.transform() - case name if name in self.kinematic_tree.joint_names(): + case name if name in self.kinematic_tree.joint_names: edge = self.kinematic_tree.joints_dict[name] assert edge.name() == name @@ -66,7 +66,7 @@ def _compute_transform(self, name: str) -> npt.NDArray: return W_H_E - case name if name in self.kinematic_tree.link_names(): + case name if name in self.kinematic_tree.link_names: element = self.kinematic_tree.links_dict[name] @@ -81,7 +81,7 @@ def _compute_transform(self, name: str) -> npt.NDArray: W_H_L = W_H_x @ x_H_L return W_H_L - case name if name in self.kinematic_tree.frame_names(): + case name if name in self.kinematic_tree.frame_names: element = self.kinematic_tree.frames_dict[name] diff --git a/src/rod/sdf/link.py b/src/rod/sdf/link.py index c685b55..e17c99b 100644 --- a/src/rod/sdf/link.py +++ b/src/rod/sdf/link.py @@ -1,4 +1,5 @@ import dataclasses + import mashumaro import numpy as np import numpy.typing as npt diff --git a/src/rod/sdf/model.py b/src/rod/sdf/model.py index b02b208..86d2ad1 100644 --- a/src/rod/sdf/model.py +++ b/src/rod/sdf/model.py @@ -1,6 +1,8 @@ from __future__ import annotations import dataclasses +import functools + import mashumaro import rod @@ -62,24 +64,26 @@ class Model(Element): joint: Joint | list[Joint] | None = dataclasses.field(default=None) + @functools.cached_property def is_fixed_base(self) -> bool: - joints_having_world_parent = [j for j in self.joints() if j.parent == "world"] + joints_having_world_parent = [j for j in self.joints if j.parent == "world"] assert len(joints_having_world_parent) in {0, 1} return len(joints_having_world_parent) > 0 def get_canonical_link(self) -> str: - if len(self.models()) != 0: + if len(self.models) != 0: msg = "Model composition is not yet supported." msg += " The returned canonical link could be wrong." logging.warning(msg=msg) if self.canonical_link is not None: - assert self.canonical_link in {l.name for l in self.links()} + assert self.canonical_link in {l.name for l in self.links} return self.canonical_link - return self.links()[0].name + return self.links[0].name + @functools.cached_property def models(self) -> list[Model]: if self.model is None: return [] @@ -90,6 +94,7 @@ def models(self) -> list[Model]: assert isinstance(self.model, list) return self.model + @functools.cached_property def frames(self) -> list[Frame]: if self.frame is None: return [] @@ -100,6 +105,7 @@ def frames(self) -> list[Frame]: assert isinstance(self.frame, list) return self.frame + @functools.cached_property def links(self) -> list[Link]: if self.link is None: return [] @@ -110,6 +116,7 @@ def links(self) -> list[Link]: assert isinstance(self.link, list), type(self.link) return self.link + @functools.cached_property def joints(self) -> list[Joint]: if self.joint is None: return [] @@ -137,7 +144,7 @@ def add_frame(self, frame: Frame) -> None: def resolve_uris(self) -> None: from rod.utils import resolve_uris - for link in self.links(): + for link in self.links: for visual in link.visuals(): resolve_uris.resolve_geometry_uris(geometry=visual.geometry) diff --git a/src/rod/urdf/exporter.py b/src/rod/urdf/exporter.py index a1ba3ed..9570975 100644 --- a/src/rod/urdf/exporter.py +++ b/src/rod/urdf/exporter.py @@ -81,7 +81,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: model.resolve_frames(is_top_level=True, explicit_frames=False) # Model composition is not supported, ignoring sub-models - if len(model.models()) > 0: + if len(model.models) > 0: msg = f"Ignoring unsupported sub-models of model '{model.name}'" logging.warning(msg=msg) @@ -94,7 +94,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # If the model pose is not zero, warn that it will be ignored. # In fact, the pose wrt world of the canonical link (base) will be used instead. if ( - model.is_fixed_base() + model.is_fixed_base and model.pose is not None and not np.allclose(model.pose.pose, np.zeros(6)) ): @@ -103,7 +103,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # Get the canonical link of the model logging.debug(f"Detected '{model.get_canonical_link()}' as root link") - canonical_link: rod.Link = {l.name: l for l in model.links()}[ + canonical_link: rod.Link = {l.name: l for l in model.links}[ model.get_canonical_link() ] @@ -113,7 +113,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # of a model, instead in URDF this reference is represented by the root link # (that is, by definition, the SDF canonical link). if ( - not model.is_fixed_base() + not model.is_fixed_base and canonical_link.pose is not None and not np.allclose(canonical_link.pose.pose, np.zeros(6)) ): @@ -141,7 +141,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # Since URDF does not support plain frames as SDF, we convert all frames # to (fixed_joint->dummy_link) sequences - for frame in model.frames(): + for frame in model.frames: # New dummy link with same name of the frame dummy_link = { @@ -203,7 +203,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # If it is a boolean, automatically populate the list with all fixed joints. if gazebo_preserve_fixed_joints is True: gazebo_preserve_fixed_joints = [ - j.name for j in model.joints() if j.type == "fixed" + j.name for j in model.joints if j.type == "fixed" ] if gazebo_preserve_fixed_joints is False: @@ -214,7 +214,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # Check that all fixed joints to preserve are actually present in the model. for fixed_joint_name in gazebo_preserve_fixed_joints: logging.debug(f"Preserving fixed joint '{fixed_joint_name}'") - all_model_joint_names = {j.name for j in model.joints()} + all_model_joint_names = {j.name for j in model.joints} if fixed_joint_name not in all_model_joint_names: raise RuntimeError(f"Joint '{fixed_joint_name}' not found in the model") @@ -223,7 +223,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # =================== # In URDF, links are directly attached to the frame of their parent joint - for link in model.links(): + for link in model.links: if link.pose is not None and not np.allclose(link.pose.pose, np.zeros(6)): msg = "Ignoring non-trivial pose of link '{name}'" logging.warning(msg.format(name=link.name)) @@ -237,7 +237,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: "robot": { **{"@name": model.name}, # http://wiki.ros.org/urdf/XML/link - "link": ([world_link.to_dict()] if model.is_fixed_base() else []) + "link": ([world_link.to_dict()] if model.is_fixed_base else []) + [ { "@name": l.name, @@ -292,7 +292,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: for c in l.collisions() ], } - for l in model.links() + for l in model.links ] # Add the extra links resulting from the frame->dummy_link conversion + extra_links_from_frames, @@ -372,7 +372,7 @@ def to_urdf_string(self, sdf: rod.Sdf | rod.Model) -> str: # mimic: does not have any SDF corresponding element # safety_controller: does not have any SDF corresponding element } - for j in model.joints() + for j in model.joints if j.type in UrdfExporter.SupportedSdfJointTypes ] # Add the extra joints resulting from the frame->link conversion diff --git a/src/rod/utils/frame_convention.py b/src/rod/utils/frame_convention.py index 697b607..f5aa30a 100644 --- a/src/rod/utils/frame_convention.py +++ b/src/rod/utils/frame_convention.py @@ -38,7 +38,7 @@ def switch_frame_convention( # Update the //frame/attached_to attribute of all frames so that they are # directly attached to links. if attach_frames_to_links: - for frame in model.frames(): + for frame in model.frames: # Find the link to which the frame is attached to following recursively # the //frame/attached_to attribute. parent_link = find_parent_link_of_frame(frame=frame, model=model) @@ -95,7 +95,7 @@ def switch_frame_convention( visual_name_to_parent_link = { visual_name: parent_link for d in [ - {v.name: link for v in link.visuals()} for link in model.links() + {v.name: link for v in link.visuals()} for link in model.links ] for visual_name, parent_link in d.items() } @@ -103,7 +103,7 @@ def switch_frame_convention( collision_name_to_parent_link = { collision_name: parent_link for d in [ - {c.name: link for c in link.collisions()} for link in model.links() + {c.name: link for c in link.collisions()} for link in model.links ] for collision_name, parent_link in d.items() } @@ -124,7 +124,7 @@ def switch_frame_convention( visual_name_to_parent_link = { visual_name: parent_link for d in [ - {v.name: link for v in link.visuals()} for link in model.links() + {v.name: link for v in link.visuals()} for link in model.links ] for visual_name, parent_link in d.items() } @@ -132,20 +132,20 @@ def switch_frame_convention( collision_name_to_parent_link = { collision_name: parent_link for d in [ - {c.name: link for c in link.collisions()} for link in model.links() + {c.name: link for c in link.collisions()} for link in model.links ] for collision_name, parent_link in d.items() } link_name_to_parent_joint_names = defaultdict(list) - for j in model.joints(): + for j in model.joints: if j.child != model.get_canonical_link(): link_name_to_parent_joint_names[j.child].append(j.name) else: # The pose of the canonical link is used to define the origin of # the URDF joint connecting the world to the robot - assert model.is_fixed_base() + assert model.is_fixed_base link_name_to_parent_joint_names[j.child].append("world") reference_frame_model = lambda m: "world" @@ -158,8 +158,8 @@ def switch_frame_convention( c.name ].name - if model.is_fixed_base(): - canonical_link = {l.name: l for l in model.links()}[ + if model.is_fixed_base: + canonical_link = {l.name: l for l in model.links}[ model.get_canonical_link() ] reference_frame_link_canonical = reference_frame_links(l=canonical_link) @@ -191,7 +191,7 @@ def switch_frame_convention( ) # Adjust the reference frames of all sub-models - for sub_model in model.models(): + for sub_model in model.models: logging.info( "Model composition not yet supported, ignoring '{}/{}'".format( model.name, sub_model.name @@ -199,7 +199,7 @@ def switch_frame_convention( ) # Adjust the reference frames of all joints - for joint in model.joints(): + for joint in model.joints: x_H_joint = joint.pose.transform() target_H_x = kin.relative_transform( relative_to=reference_frame_joints(j=joint), @@ -212,7 +212,7 @@ def switch_frame_convention( ) # Adjust the reference frames of all frames - for frame in model.frames(): + for frame in model.frames: x_H_frame = frame.pose.transform() target_H_x = kin.relative_transform( relative_to=reference_frame_frames(f=frame), @@ -225,7 +225,7 @@ def switch_frame_convention( ) # Adjust the reference frames of all links - for link in model.links(): + for link in model.links: relative_to = ( reference_frame_links(l=link) if link.name != model.get_canonical_link() @@ -283,10 +283,10 @@ def switch_frame_convention( def find_parent_link_of_frame(frame: rod.Frame, model: rod.Model) -> str: - links_dict = {l.name: l for l in model.links()} - frames_dict = {f.name: f for f in model.frames()} - joints_dict = {j.name: j for j in model.joints()} - sub_models_dict = {m.name: m for m in model.models()} + links_dict = {l.name: l for l in model.links} + frames_dict = {f.name: f for f in model.frames} + joints_dict = {j.name: j for j in model.joints} + sub_models_dict = {m.name: m for m in model.models} assert isinstance(frame, rod.Frame) diff --git a/src/rod/utils/resolve_frames.py b/src/rod/utils/resolve_frames.py index 48b11c9..9dd0490 100644 --- a/src/rod/utils/resolve_frames.py +++ b/src/rod/utils/resolve_frames.py @@ -69,7 +69,7 @@ def resolve_model_frames( else: update_element(element=model, default_relative_to="world") - for frame in model.frames(): + for frame in model.frames: update_element( element=frame, default_relative_to=( @@ -79,7 +79,7 @@ def resolve_model_frames( ) # Update the links and its children elements - for link in model.links(): + for link in model.links: update_element(element=link, default_relative_to=["__model__", model.name]) update_element(element=link.inertial, default_relative_to=link.name) @@ -91,11 +91,11 @@ def resolve_model_frames( update_element(element=collision, default_relative_to=link.name) # Update the joints - for joint in model.joints(): + for joint in model.joints: update_element(element=joint, default_relative_to=joint.child) # Update the sub-models - for sub_model in model.models(): + for sub_model in model.models: resolve_model_frames( model=sub_model, is_top_level=False, explicit_frames=explicit_frames ) diff --git a/tests/test_urdf_exporter.py b/tests/test_urdf_exporter.py index 281e779..258ec73 100644 --- a/tests/test_urdf_exporter.py +++ b/tests/test_urdf_exporter.py @@ -36,7 +36,7 @@ def test_urdf_exporter(robot: Robot) -> None: mdl_loader_exported = idt.ModelLoader() # Get the joint serialization from ROD. - joint_names = [j.name for j in sdf.model.joints() if j.type != "fixed"] + joint_names = [j.name for j in sdf.model.joints if j.type != "fixed"] # Load the original URDF. assert mdl_loader_original.loadReducedModelFromString( diff --git a/tests/test_urdf_parsing.py b/tests/test_urdf_parsing.py index 8ee3e94..d9a268a 100644 --- a/tests/test_urdf_parsing.py +++ b/tests/test_urdf_parsing.py @@ -62,20 +62,20 @@ def test_urdf_parsing(robot: Robot) -> None: ) # Check the canonical link - assert rod_sdf.model.get_canonical_link() == idt_model_full.getLinkName( + assert rod_sdf.model.canonical_link == idt_model_full.getLinkName( idt_model_full.getDefaultBaseLink() ) # Extract data from the ROD model - link_names = [l.name for l in rod_sdf.model.links()] - joint_names = [j.name for j in rod_sdf.model.joints()] - joint_names_dofs = [j.name for j in rod_sdf.model.joints() if j.type != "fixed"] + link_names = [l.name for l in rod_sdf.model.links] + joint_names = [j.name for j in rod_sdf.model.joints] + joint_names_dofs = [j.name for j in rod_sdf.model.joints if j.type != "fixed"] # Remove the world joint if the model is fixed-base - if rod_sdf.model.is_fixed_base(): + if rod_sdf.model.is_fixed_base: world_joints = [ j.name - for j in rod_sdf.model.joints() + for j in rod_sdf.model.joints if j.type == "fixed" and j.parent == "world" ] assert len(world_joints) == 1 @@ -85,7 +85,7 @@ def test_urdf_parsing(robot: Robot) -> None: # New sdformat versions, after lumping fixed joints, create a new frame called # "_fixed_joint" frame_names = [ - f.name for f in rod_sdf.model.frames() if not f.name.endswith("_fixed_joint") + f.name for f in rod_sdf.model.frames if not f.name.endswith("_fixed_joint") ] # Check the number of DoFs (all non-fixed joints) @@ -110,5 +110,5 @@ def test_urdf_parsing(robot: Robot) -> None: assert set(joint_names_idt) == set(joint_names) assert set(frame_names) - set(frame_names_idt) == set() - total_mass = sum([l.inertial.mass for l in rod_sdf.model.links()]) + total_mass = sum([l.inertial.mass for l in rod_sdf.model.links]) assert total_mass == pytest.approx(idt_model_full.getTotalMass())