Merge pull request #123 from Tigul/world-copy

Removed JSON serializer from World deepcopy for better perfromance

Author: Tigul

src/semantic_digital_twin/world.py

@@ -3,7 +3,7 @@
 
 import inspect
 import logging
-from copy import deepcopy
+from copy import deepcopy, copy
 from dataclasses import dataclass, field
 from enum import IntEnum
 from functools import wraps, lru_cache, cached_property
@@ -1772,7 +1772,6 @@ def __deepcopy__(self, memo):
         new_world = World(name=self.name)
         memo[me_id] = new_world
 
-        tracker = KinematicStructureEntityKwargsTracker.from_world(new_world)
         with new_world.modify_world():
             for body in self.bodies:
                 new_body = Body(
@@ -1790,9 +1789,7 @@ def __deepcopy__(self, memo):
                 new_world.add_degree_of_freedom(new_dof)
                 new_world.state[dof.name] = self.state[dof.name].data
             for connection in self.connections:
-                new_connection = SubclassJSONSerializer.from_json(
-                    connection.to_json(), **tracker.create_kwargs()
-                )
+                new_connection = connection.copy_for_world(new_world)
                 new_world.add_connection(new_connection)
         return new_world
 

src/semantic_digital_twin/world_description/connections.py

@@ -309,6 +309,24 @@ def jerk(self, value: float) -> None:
         self._world.state[self.raw_dof.name].jerk = value / self.multiplier
         self._world.notify_state_change()
 
+    def copy_for_world(self, world: World):
+        (
+            other_parent,
+            other_child,
+            parent_T_connection_expression,
+        ) = self._find_references_in_world(world)
+
+        return self.__class__(
+            name=PrefixedName(self.name.name, self.name.prefix),
+            parent=other_parent,
+            child=other_child,
+            parent_T_connection_expression=parent_T_connection_expression,
+            dof_name=PrefixedName(self.dof_name.name, self.dof_name.prefix),
+            axis=self.axis,
+            multiplier=self.multiplier,
+            offset=self.offset,
+        )
+
 
 @dataclass(eq=False)
 class PrismaticConnection(ActiveConnection1DOF):
@@ -557,6 +575,32 @@ def origin(
         self._world.state[self.qw.name].position = orientation[3]
         self._world.notify_state_change()
 
+    def copy_for_world(self, world: World) -> Connection6DoF:
+        """
+        Copies this 6DoF connection for another world. Returns a new connection with references to the given world.
+        :param world: The world to copy this connection for.
+        :return: A copy of this connection for the given world.
+        """
+        (
+            other_parent,
+            other_child,
+            parent_T_connection_expression,
+        ) = self._find_references_in_world(world)
+
+        return Connection6DoF(
+            name=deepcopy(self.name),
+            parent=other_parent,
+            child=other_child,
+            parent_T_connection_expression=parent_T_connection_expression,
+            x_name=deepcopy(self.x_name),
+            y_name=deepcopy(self.y_name),
+            z_name=deepcopy(self.z_name),
+            qx_name=deepcopy(self.qx_name),
+            qy_name=deepcopy(self.qy_name),
+            qz_name=deepcopy(self.qz_name),
+            qw_name=deepcopy(self.qw_name),
+        )
+
 
 @dataclass(eq=False)
 class OmniDrive(ActiveConnection, HasUpdateState):
@@ -821,3 +865,30 @@ def has_hardware_interface(self, value: bool) -> None:
         self.x_velocity.has_hardware_interface = value
         self.y_velocity.has_hardware_interface = value
         self.yaw.has_hardware_interface = value
+
+    def copy_for_world(self, world: World) -> OmniDrive:
+        """
+        Copies this OmniDriveConnection for the provided world. This finds the references for the parent and child in
+        the new world and returns a new connection with references to the new parent and child.
+        :param world: The world where the connection is copied.
+        :return: The connection with references to the new parent and child.
+        """
+        (
+            other_parent,
+            other_child,
+            parent_T_connection_expression,
+        ) = self._find_references_in_world(world)
+
+        return OmniDrive(
+            name=deepcopy(self.name),
+            parent=other_parent,
+            child=other_child,
+            parent_T_connection_expression=parent_T_connection_expression,
+            x_name=deepcopy(self.x_name),
+            y_name=deepcopy(self.y_name),
+            roll_name=deepcopy(self.roll_name),
+            pitch_name=deepcopy(self.pitch_name),
+            yaw_name=deepcopy(self.yaw_name),
+            x_velocity_name=deepcopy(self.x_velocity_name),
+            y_velocity_name=deepcopy(self.y_velocity_name),
+        )

src/semantic_digital_twin/world_description/world_entity.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import inspect
+from copy import copy, deepcopy
+
 import itertools
 from abc import ABC, abstractmethod
 from collections import deque
@@ -793,7 +795,9 @@ class Connection(WorldEntity, SubclassJSONSerializer):
     """
 
     parent_T_connection_expression: TransformationMatrix = field(default=None)
-    _connection_T_child_expression: TransformationMatrix = field(default=None)
+    _connection_T_child_expression: TransformationMatrix = field(
+        default=None, init=False
+    )
     """
     The origin expression of a connection is split into 2 transforms:
     1. parent_T_connection describes the pose of the connection and is always constant.
@@ -955,6 +959,51 @@ def create_with_dofs(
             "ConnectionWithDofs.create_with_dofs is not implemented."
         )
 
+    def _find_references_in_world(
+        self, world: World
+    ) -> Tuple[
+        KinematicStructureEntity, KinematicStructureEntity, TransformationMatrix
+    ]:
+        """
+        Finds the reference frames to this connection in the given world and returns them as usable objects.
+        :param world: Reference to the world where the reference frames are searched.
+        :return: The other parent and child and new connection expressions with correct reference frames.
+        """
+        other_parent = world.get_kinematic_structure_entity_by_name(self.parent.name)
+        other_child = world.get_kinematic_structure_entity_by_name(self.child.name)
+
+        parent_T_connection_expression = deepcopy(self.parent_T_connection_expression)
+        parent_T_connection_expression.reference_frame = (
+            world.get_kinematic_structure_entity_by_name(
+                parent_T_connection_expression.reference_frame.name
+            )
+        )
+        return (
+            other_parent,
+            other_child,
+            parent_T_connection_expression,
+        )
+
+    def copy_for_world(self, world: World) -> Self:
+        """
+        Copies this connection to the given world the parent and child references are updated to the new world as well
+        as the references from the expression.
+        :param world: World in which the connection should be copied.
+        :return: The copied connection.
+        """
+        (
+            other_parent,
+            other_child,
+            parent_T_connection_expression,
+        ) = self._find_references_in_world(world)
+
+        return self.__class__(
+            other_parent,
+            other_child,
+            parent_T_connection_expression=parent_T_connection_expression,
+            name=PrefixedName(self.name.name, prefix=self.name.prefix),
+        )
+
 
 GenericConnection = TypeVar("GenericConnection", bound=Connection)
 

test/test_worlds/test_world.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 import pytest
+from numpy.testing import assert_raises
+
 from semantic_digital_twin.semantic_annotations.semantic_annotations import Handle
 
 from semantic_digital_twin.spatial_types import Vector3
@@ -647,7 +649,7 @@ def test_copy_pr2_world_connection_origin(pr2_world):
         pr2_body = pr2_world.get_kinematic_structure_entity_by_name(body.name)
         pr2_copy_body = pr2_copy.get_kinematic_structure_entity_by_name(body.name)
         np.testing.assert_array_almost_equal(
-            pr2_body.global_pose.to_np(), pr2_copy_body.global_pose.to_np()
+            pr2_body.global_pose.to_np(), pr2_copy_body.global_pose.to_np(), decimal=4
         )
 
 
@@ -676,6 +678,35 @@ def test_copy_pr2(pr2_world):
     ).global_pose.to_np()[2, 3] == pytest.approx(1.472, abs=1e-3)
 
 
+def test_copy_connections(pr2_world):
+    pr2_copy = deepcopy(pr2_world)
+    for connection in pr2_world.connections:
+        pr2_copy_connection = pr2_copy.get_connection_by_name(connection.name)
+        assert connection.name == pr2_copy_connection.name
+        np.testing.assert_array_almost_equal(
+            connection.origin.to_np(), pr2_copy_connection.origin.to_np(), decimal=3
+        )
+    pr2_copy.state[
+        pr2_copy.get_degree_of_freedom_by_name("torso_lift_joint").name
+    ].position = 0.3
+    pr2_copy.notify_state_change()
+
+    assert_raises(
+        AssertionError,
+        np.testing.assert_array_almost_equal,
+        pr2_world.get_connection_by_name("torso_lift_joint").origin.to_np(),
+        pr2_copy.get_connection_by_name("torso_lift_joint").origin.to_np(),
+    )
+
+
+def test_copy_two_times(pr2_world):
+    pr2_copy = deepcopy(pr2_world)
+    pr2_copy_2 = deepcopy(pr2_copy)
+    for connection in pr2_world.connections:
+        pr2_copy_connection = pr2_copy_2.get_connection_by_name(connection.name)
+        assert connection.name == pr2_copy_connection.name
+
+
 def test_add_entity_with_duplicate_name(world_setup):
     world, l1, l2, bf, r1, r2 = world_setup
     body_duplicate = Body(name=PrefixedName("l1"))