Update arm resource state after pickup

pylabrobot/capabilities/arms/arm.py

@@ -30,6 +30,10 @@ class _PickedUpState:
   offset: Coordinate
   pickup_distance_from_bottom: float
   resource_width: float
+  # Preserve the resource pose at pickup time. Once pickup succeeds, the
+  # resource is unassigned from its source so the source is immediately free in
+  # the model; later drop math must not depend on the old parent tree.
+  resource_absolute_rotation_at_pickup: Rotation
   rotation: Rotation = Rotation()
 
 
@@ -152,11 +156,9 @@ def _compute_end_effector_location(
     to_location: Coordinate,
     offset: Coordinate,
     pickup_distance_from_bottom: float,
-    rotation_applied_by_move: float,
+    resource_absolute_rotation_after_move: Rotation,
   ) -> Coordinate:
-    center = resource.center().rotated(
-      Rotation(z=resource.get_absolute_rotation().z + rotation_applied_by_move)
-    )
+    center = resource.center().rotated(resource_absolute_rotation_after_move)
     loc = to_location + center + offset
     return Coordinate(
       loc.x,
@@ -234,14 +236,18 @@ def _compute_drop(
     offset: Coordinate,
     pickup_distance_from_bottom: float,
     rotation_applied_by_move: float = 0,
+    resource_absolute_rotation_at_pickup: Optional[Rotation] = None,
   ) -> Tuple[Coordinate, float]:
-    resource_absolute_rotation_after_move = (
-      resource.get_absolute_rotation().z + rotation_applied_by_move
+    resource_absolute_rotation_at_pickup = (
+      resource_absolute_rotation_at_pickup or resource.get_absolute_rotation()
+    )
+    resource_absolute_rotation_after_move = resource_absolute_rotation_at_pickup + Rotation(
+      z=rotation_applied_by_move
     )
     dest_rotation = (
       destination.get_absolute_rotation().z if not isinstance(destination, Coordinate) else 0
     )
-    resource_rotation_wrt_destination = resource_absolute_rotation_after_move - dest_rotation
+    resource_rotation_wrt_destination = resource_absolute_rotation_after_move.z - dest_rotation
     resource_rotation_wrt_destination_wrt_local = (
       resource_rotation_wrt_destination - resource.rotation.z
     )
@@ -257,7 +263,11 @@ def _compute_drop(
       resource, destination, resource_rotation_wrt_destination_wrt_local
     )
     location = self._compute_end_effector_location(
-      resource, to_location, offset, pickup_distance_from_bottom, rotation_applied_by_move
+      resource,
+      to_location,
+      offset,
+      pickup_distance_from_bottom,
+      resource_absolute_rotation_after_move,
     )
     return location, resource_rotation_wrt_destination
 
@@ -445,13 +455,16 @@ async def pick_up_resource(
       resource, offset, pickup_distance_from_bottom
     )
     resource_width = self._resource_width(resource)
+    resource_absolute_rotation_at_pickup = resource.get_absolute_rotation()
     await self.pick_up_at_location(location, resource_width, backend_params)
     self._picked_up = _PickedUpState(
       resource=resource,
       offset=offset,
       pickup_distance_from_bottom=pickup_distance_from_bottom,
       resource_width=resource_width,
+      resource_absolute_rotation_at_pickup=resource_absolute_rotation_at_pickup,
     )
+    resource.unassign()
     self._state_updated()
 
   async def drop_at_location(
@@ -480,6 +493,9 @@ async def drop_resource(
       destination=destination,
       offset=offset,
       pickup_distance_from_bottom=self._picked_up.pickup_distance_from_bottom,
+      resource_absolute_rotation_at_pickup=(
+        self._picked_up.resource_absolute_rotation_at_pickup
+      ),
     )
     await self.drop_at_location(location, backend_params)
     self._finalize_drop(resource, destination, rotation)

pylabrobot/capabilities/arms/arm_tests.py

@@ -8,6 +8,7 @@
 )
 from pylabrobot.capabilities.arms.orientable_arm import OrientableGripperArm
 from pylabrobot.resources import Coordinate, Resource, ResourceHolder
+from pylabrobot.resources.rotation import Rotation
 
 
 def _assert_location(test, call, x, y, z, places=1):
@@ -68,6 +69,8 @@ async def test_pick_up_resource(self):
     _assert_location(self, call, 165, 145, 58)
     # default grip_axis="x" → resource_width is X size = 120
     self.assertAlmostEqual(call.kwargs["resource_width"], 120)
+    self.assertIsNone(self.plate.parent)
+    self.assertIsNone(self.site_a.resource)
 
   async def test_drop_resource(self):
     await self.arm.pick_up_resource(self.plate, pickup_distance_from_bottom=8)
@@ -191,6 +194,8 @@ async def test_pick_up_front(self):
     self.assertAlmostEqual(call.kwargs["direction"], 270.0)
     # "front" grips along the X axis → X width = 120
     self.assertAlmostEqual(call.kwargs["resource_width"], 120)
+    self.assertIsNone(self.plate.parent)
+    self.assertIsNone(self.site_a.resource)
 
   async def test_pick_up_right(self):
     await self.arm.pick_up_resource(self.plate, pickup_distance_from_bottom=8, direction="right")
@@ -223,6 +228,16 @@ async def test_move_plate(self):
     _assert_location(self, drop_call, 160, 340, 58)
     self.assertEqual(self.plate.parent.name, "site_b")
 
+  async def test_drop_after_pickup_uses_captured_source_rotation(self):
+    self.site_a.rotation = Rotation(z=90)
+    await self.arm.pick_up_resource(self.plate, pickup_distance_from_bottom=8, direction="front")
+    self.assertIsNone(self.plate.parent)
+
+    await self.arm.drop_resource(self.site_b, direction="front")
+
+    self.assertEqual(self.plate.parent.name, "site_b")
+    self.assertAlmostEqual(self.plate.get_absolute_rotation().z % 360, 90)
+
 
 class TestGripperArmAbstract(unittest.TestCase):
   def test_cannot_instantiate_abstract_base(self):

pylabrobot/capabilities/arms/articulated_arm.py

@@ -51,14 +51,17 @@ async def pick_up_resource(
       resource, offset, pickup_distance_from_bottom
     )
     resource_width = self._resource_width_for_rotation(resource, rotation)
+    resource_absolute_rotation_at_pickup = resource.get_absolute_rotation()
     await self.pick_up_at_location(location, resource_width, rotation, backend_params)
     self._picked_up = _PickedUpState(
       resource=resource,
       offset=offset,
       pickup_distance_from_bottom=pickup_distance_from_bottom,
       resource_width=resource_width,
+      resource_absolute_rotation_at_pickup=resource_absolute_rotation_at_pickup,
       rotation=rotation,
     )
+    resource.unassign()
     self._state_updated()
 
   async def drop_at_location(
@@ -95,6 +98,9 @@ async def drop_resource(
       offset=offset,
       pickup_distance_from_bottom=self._picked_up.pickup_distance_from_bottom,
       rotation_applied_by_move=rotation_applied_by_move,
+      resource_absolute_rotation_at_pickup=(
+        self._picked_up.resource_absolute_rotation_at_pickup
+      ),
     )
     await self.drop_at_location(location, rotation, backend_params)
     self._finalize_drop(resource, destination, final_rotation)

pylabrobot/capabilities/arms/orientable_arm.py

@@ -85,14 +85,17 @@ async def pick_up_resource(
     )
     dir_degrees = _resolve_direction(direction)
     resource_width = self._resource_width_for_direction(resource, dir_degrees)
+    resource_absolute_rotation_at_pickup = resource.get_absolute_rotation()
     await self.pick_up_at_location(location, resource_width, dir_degrees, backend_params)
     self._picked_up = _PickedUpState(
       resource=resource,
       offset=offset,
       pickup_distance_from_bottom=pickup_distance_from_bottom,
       resource_width=resource_width,
+      resource_absolute_rotation_at_pickup=resource_absolute_rotation_at_pickup,
       rotation=Rotation(z=dir_degrees),
     )
+    resource.unassign()
     self._state_updated()
 
   async def drop_at_location(
@@ -129,6 +132,9 @@ async def drop_resource(
       offset=offset,
       pickup_distance_from_bottom=self._picked_up.pickup_distance_from_bottom,
       rotation_applied_by_move=rotation_applied_by_move,
+      resource_absolute_rotation_at_pickup=(
+        self._picked_up.resource_absolute_rotation_at_pickup
+      ),
     )
     await self.drop_at_location(location, drop_dir, backend_params)
     self._finalize_drop(resource, destination, rotation)

pylabrobot/legacy/liquid_handling/backends/hamilton/STAR_backend.py

@@ -2859,11 +2859,17 @@ async def drop_resource(
     # This means that the center vector has to be rotated from the child local space by the
     # new child absolute rotation. The moved resource's rotation will be the original child
     # rotation plus the rotation applied by the movement.
-    # The resource is moved by drop.rotation
-    # The new resource absolute location is
-    # drop.resource.get_absolute_rotation().z + drop.rotation
+    # The resource is moved by drop.rotation. The source may already have been
+    # unassigned after pickup, so use the absolute rotation captured at pickup
+    # rather than asking the resource tree for its current parent-derived pose.
+    resource_absolute_rotation_at_pickup = (
+      drop.resource_absolute_rotation_at_pickup or drop.resource.get_absolute_rotation()
+    )
+    resource_absolute_rotation_after_move = resource_absolute_rotation_at_pickup + Rotation(
+      z=drop.rotation
+    )
     center_in_absolute_space = drop.resource.center().rotated(
-      Rotation(z=drop.resource.get_absolute_rotation().z + drop.rotation)
+      resource_absolute_rotation_after_move
     )
     x, y, z = drop.destination + center_in_absolute_space + drop.offset
 
@@ -2873,15 +2879,17 @@ async def drop_resource(
       )
       z_position_at_the_command_end = z_position_at_the_command_end or self._iswap.traversal_height
       assert (
-        drop.resource.get_absolute_rotation().x == 0
-        and drop.resource.get_absolute_rotation().y == 0
+        resource_absolute_rotation_at_pickup.x == 0
+        and resource_absolute_rotation_at_pickup.y == 0
       )
-      assert drop.resource.get_absolute_rotation().z % 90 == 0
+      assert resource_absolute_rotation_at_pickup.z % 90 == 0
 
-      # Use the pickup direction to determine how wide the plate is gripped.
-      # Note that the plate is still in the original orientation at this point,
-      # so get_absolute_size_{x,y}() will return the size of the plate in the original orientation.
-      if (
+      # Use the width captured at pickup time. The resource may already be
+      # unassigned while it is held, so recomputing absolute size here can lose
+      # source/carrier rotation context and release the gripper too wide.
+      if drop.resource_width_at_pickup is not None:
+        plate_width = drop.resource_width_at_pickup
+      elif (
         drop.pickup_direction == GripDirection.FRONT or drop.pickup_direction == GripDirection.BACK
       ):
         plate_width = drop.resource.get_absolute_size_x()

pylabrobot/legacy/liquid_handling/liquid_handler.py

@@ -1171,11 +1171,20 @@ async def pick_up_resource(
     if self._resource_pickup is not None:
       raise RuntimeError(f"Resource {self._resource_pickup.resource.name} already picked up")
 
+    if backend_kwargs.get("plate_width") is not None:
+      resource_width_at_pickup = backend_kwargs["plate_width"]
+    elif direction in (GripDirection.FRONT, GripDirection.BACK):
+      resource_width_at_pickup = resource.get_absolute_size_x()
+    else:
+      resource_width_at_pickup = resource.get_absolute_size_y()
+
     self._resource_pickup = ResourcePickup(
       resource=resource,
       offset=offset,
       pickup_distance_from_top=pickup_distance_from_top,
       direction=direction,
+      resource_absolute_rotation_at_pickup=resource.get_absolute_rotation(),
+      resource_width_at_pickup=resource_width_at_pickup,
     )
 
     extras = self._check_args(
@@ -1193,6 +1202,7 @@ async def pick_up_resource(
       self._resource_pickup = None
       raise e
 
+    resource.unassign()
     self._state_updated()
 
   async def move_picked_up_resource(
@@ -1269,8 +1279,12 @@ async def drop_resource(
     ):
       rotation_applied_by_move = 270
 
+    resource_absolute_rotation_at_pickup = (
+      self._resource_pickup.resource_absolute_rotation_at_pickup
+      or resource.get_absolute_rotation()
+    )
     resource_absolute_rotation_after_move = (
-      resource.get_absolute_rotation().z + rotation_applied_by_move
+      resource_absolute_rotation_at_pickup.z + rotation_applied_by_move
     )
     destination_rotation = (
       destination.get_absolute_rotation().z if not isinstance(destination, Coordinate) else 0
@@ -1329,6 +1343,8 @@ async def drop_resource(
       pickup_direction=self._resource_pickup.direction,
       direction=direction,
       rotation=rotation_applied_by_move,
+      resource_absolute_rotation_at_pickup=resource_absolute_rotation_at_pickup,
+      resource_width_at_pickup=self._resource_pickup.resource_width_at_pickup,
     )
 
     result = await self.backend.drop_resource(drop=drop, **backend_kwargs)

pylabrobot/legacy/liquid_handling/liquid_handler_tests.py

@@ -27,13 +27,15 @@
   Lid,
   Plate,
   ResourceNotFoundError,
+  ResourceHolder,
   ResourceStack,
   TipRack,
   nest_1_troughplate_195000uL_Vb,
   no_tip_tracking,
   set_tip_tracking,
 )
 from pylabrobot.resources.carrier import PlateHolder
+from pylabrobot.resources.rotation import Rotation
 from pylabrobot.resources.errors import (
   HasTipError,
   NoTipError,
@@ -1130,6 +1132,46 @@ async def asyncSetUp(self):
     self.deck.assign_child_resource(self.plate, location=Coordinate(100, 100, 0))
     await self.lh.setup()
 
+  async def test_resource_pickup_unassigns_source_after_success(self):
+    parent = self.plate.parent
+    await self.lh.pick_up_resource(self.plate)
+
+    self.assertIsNone(self.plate.parent)
+    self.assertNotIn(self.plate, parent.children)
+    self.assertIs(self.lh.get_picked_up_resource(), self.plate)
+
+  async def test_resource_drop_after_pickup_assigns_destination(self):
+    destination = ResourceHolder("destination", size_x=130, size_y=90, size_z=0)
+    self.deck.assign_child_resource(destination, location=Coordinate(300, 100, 0))
+
+    await self.lh.pick_up_resource(self.plate)
+    self.assertIsNone(self.plate.parent)
+
+    await self.lh.drop_resource(destination)
+
+    self.assertIs(destination.resource, self.plate)
+    self.assertIsNone(self.lh.get_picked_up_resource())
+
+  async def test_resource_drop_uses_width_captured_at_pickup(self):
+    source = ResourceHolder("rotated_source", size_x=130, size_y=90, size_z=0)
+    source.rotation = Rotation(z=90)
+    destination = ResourceHolder("destination", size_x=130, size_y=90, size_z=0)
+    self.plate.unassign()
+    self.deck.assign_child_resource(source, location=Coordinate(100, 100, 0))
+    self.deck.assign_child_resource(destination, location=Coordinate(300, 100, 0))
+    source.assign_child_resource(self.plate, location=Coordinate.zero())
+
+    width_at_pickup = self.plate.get_absolute_size_x()
+    width_after_unassign = self.plate.get_size_x()
+    self.assertNotEqual(width_at_pickup, width_after_unassign)
+
+    await self.lh.pick_up_resource(self.plate, direction=GripDirection.FRONT)
+    self.assertIsNone(self.plate.parent)
+    await self.lh.drop_resource(destination, direction=GripDirection.FRONT)
+
+    drop = self.backend.drop_resource.call_args.kwargs["drop"]
+    self.assertAlmostEqual(drop.resource_width_at_pickup, width_at_pickup)
+
   async def test_serialize_state_after_setup(self):
     state = self.lh.serialize_state()
     self.assertIn("head_state", state)

pylabrobot/legacy/liquid_handling/standard.py

@@ -140,6 +140,8 @@ class ResourcePickup:
   offset: Coordinate
   pickup_distance_from_top: float
   direction: GripDirection
+  resource_absolute_rotation_at_pickup: Optional[Rotation] = None
+  resource_width_at_pickup: Optional[float] = None
 
 
 @dataclass(frozen=True)
@@ -164,6 +166,8 @@ class ResourceDrop:
   pickup_direction: GripDirection
   direction: GripDirection
   rotation: float
+  resource_absolute_rotation_at_pickup: Optional[Rotation] = None
+  resource_width_at_pickup: Optional[float] = None
 
 
 PipettingOp = Union[Pickup, Drop, SingleChannelAspiration, SingleChannelDispense]