Preserve resource state across legacy arm pickup and drop

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

@@ -5049,11 +5049,16 @@ 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 detached immediately after pickup. Use its captured
+    # parent-derived rotation rather than asking the detached resource tree.
+    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
 
@@ -5063,15 +5068,16 @@ 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 (
+      # Keep the release width captured before detaching the resource. Recomputing
+      # its absolute size after pickup can select the long-side plate width.
+      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

@@ -2069,11 +2069,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(
@@ -2091,6 +2100,7 @@ async def pick_up_resource(
       self._resource_pickup = None
       raise e
 
+    resource.unassign()
     self._state_updated()
 
   async def move_picked_up_resource(
@@ -2182,8 +2192,12 @@ async def drop_resource(
     # should be and subtract the rotation of the new parent.
 
     # moving from a resource from a rotated parent to a non-rotated parent means child inherits/'houses' the rotation after move
+    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
@@ -2257,6 +2271,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

@@ -28,6 +28,7 @@
   Plate,
   Resource,
   ResourceNotFoundError,
+  ResourceHolder,
   ResourceStack,
   TipRack,
   cor_96_wellplate_360uL_Fb,
@@ -41,6 +42,7 @@
   HasTipError,
   NoTipError,
 )
+from pylabrobot.resources.rotation import Rotation
 from pylabrobot.resources.hamilton import (
   STARLetDeck,
   hamilton_96_tiprack_300uL_filter,
@@ -1231,6 +1233,27 @@ async def test_serialize_state_after_setup(self):
     # 1 arm, no resource picked up
     self.assertEqual(state["arm_state"], {0: None})
 
+  async def test_resource_drop_uses_pose_and_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()
+    rotation_at_pickup = self.plate.get_absolute_rotation()
+    self.assertNotEqual(width_at_pickup, self.plate.get_size_x())
+
+    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.assertEqual(drop.resource_absolute_rotation_at_pickup, rotation_at_pickup)
+    self.assertAlmostEqual(drop.resource_width_at_pickup, width_at_pickup)
+
   async def test_serialize_state_no_head96(self):
     backend = _create_mock_backend(num_channels=8)
     type(backend).head96_installed = PropertyMock(return_value=False)

pylabrobot/legacy/liquid_handling/standard.py

@@ -151,6 +151,10 @@ class ResourcePickup:
   offset: Coordinate
   pickup_distance_from_top: float
   direction: GripDirection
+  # The resource is detached once the grip succeeds, so capture parent-derived
+  # geometry before it is no longer available from the resource tree.
+  resource_absolute_rotation_at_pickup: Optional[Rotation] = None
+  resource_width_at_pickup: Optional[float] = None
 
 
 @dataclass(frozen=True)
@@ -175,6 +179,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]