[rotations] add ResourceHolderMixin to fix placement in a number of resources

pylabrobot/liquid_handling/backends/hamilton/STAR.py

@@ -40,7 +40,7 @@
 from pylabrobot.resources.hamilton.hamilton_decks import STAR_SIZE_X, STARLET_SIZE_X
 from pylabrobot.resources.liquid import Liquid
 from pylabrobot.resources.ml_star import HamiltonTip, TipDropMethod, TipPickupMethod, TipSize
-from pylabrobot.resources.utils import get_child_location
+from pylabrobot.resources.resource_holder import get_child_location
 from pylabrobot.utils.linalg import matrix_vector_multiply_3x3
 
 T = TypeVar("T")

pylabrobot/liquid_handling/liquid_handler_tests.py

@@ -1,6 +1,7 @@
 """ Tests for LiquidHandler """
 # pylint: disable=missing-class-docstring
 
+import itertools
 import pytest
 import tempfile
 from typing import Any, Dict, List, Optional, cast
@@ -9,8 +10,11 @@
 
 from pylabrobot.liquid_handling.strictness import Strictness, set_strictness
 from pylabrobot.resources import no_tip_tracking, set_tip_tracking, Liquid
+from pylabrobot.resources.carrier import PlateCarrierSite
 from pylabrobot.resources.errors import HasTipError, NoTipError, CrossContaminationError
 from pylabrobot.resources.volume_tracker import set_volume_tracking, set_cross_contamination_tracking
+from pylabrobot.resources.well import Well
+from pylabrobot.resources.utils import create_ordered_items_2d
 
 from . import backends
 from .liquid_handler import LiquidHandler, OperationCallback
@@ -30,6 +34,7 @@
 from pylabrobot.resources.hamilton import STARLetDeck
 from pylabrobot.resources.ml_star import STF_L, HTF_L
 from .standard import (
+  GripDirection,
   Pickup,
   Drop,
   DropTipRack,
@@ -249,6 +254,50 @@ async def test_move_plate_onto_resource_stack_with_plate(self):
     self.assertEqual(plate2.location.z, 15)
     self.assertEqual(stack.get_absolute_size_z(), 30)
 
+  async def test_move_plate_rotation(self):
+    rotations = [0, 90, 270, 360]
+    grip_directions = [
+        (GripDirection.LEFT, GripDirection.RIGHT),
+        (GripDirection.FRONT, GripDirection.BACK),
+    ]
+    site_classes = [ResourceStack, PlateCarrierSite]
+
+    test_cases = itertools.product(site_classes, rotations, grip_directions)
+
+    for site_class, rotation, (get_direction, put_direction) in test_cases:
+      with self.subTest(stack_type=site_class.__name__, rotation=rotation,
+                        get_direction=get_direction, put_direction=put_direction):
+        if site_class == ResourceStack:
+          site = site_class(name=f"{site_class.__name__.lower()}", direction="z")
+        else:
+          site = site_class(
+            name=f"{site_class.__name__.lower()}",
+            size_x=100,
+            size_y=100,
+            size_z=15,
+            pedestal_size_z=1
+          )
+
+        self.deck.assign_child_resource(site, location=Coordinate(100, 100, 0))
+
+        plate = Plate("plate", size_x=100, size_y=100, size_z=15,
+                      ordered_items=create_ordered_items_2d(
+                        Well, num_items_x=1, num_items_y=1, dx=0, dy=0, dz=0,
+                        item_dx=10, item_dy=10, size_x=10, size_y=10, size_z=10))
+        plate.rotate(z=rotation)
+        site.assign_child_resource(plate)
+        original_center = plate.get_absolute_location(x="c", y="c", z="c")
+        await self.lh.move_plate(plate, site, get_direction=get_direction,
+                                 put_direction=put_direction)
+        new_center = plate.get_absolute_location(x="c", y="c", z="c")
+
+        self.assertEqual(new_center, original_center,
+                         f"Center mismatch for {site_class.__name__}, rotation {rotation}, "
+                         f"get_direction {get_direction}, "
+                         f"put_direction {put_direction}")
+        plate.unassign()
+        self.deck.unassign_child_resource(site)
+
   def test_serialize(self):
     serialized = self.lh.serialize()
     deserialized = LiquidHandler.deserialize(serialized)

pylabrobot/machines/machine.py

@@ -38,6 +38,7 @@ def __init__(
     backend: MachineBackend,
     category: Optional[str] = None,
     model: Optional[str] = None,
+    **_kwargs
   ):
     super().__init__(name=name, size_x=size_x, size_y=size_y, size_z=size_z,
                      category=category, model=model)

pylabrobot/plate_reading/plate_reader.py

@@ -3,14 +3,14 @@
 from pylabrobot.machines.machine import Machine, need_setup_finished
 from pylabrobot.resources import Coordinate, Plate, Resource
 from pylabrobot.plate_reading.backend import PlateReaderBackend
-from pylabrobot.resources.utils import get_child_location
+from pylabrobot.resources.resource_holder import ResourceHolderMixin
 
 
 class NoPlateError(Exception):
   pass
 
 
-class PlateReader(Machine):
+class PlateReader(ResourceHolderMixin, Machine):
   """ The front end for plate readers. Plate readers are devices that can read luminescence,
   absorbance, or fluorescence from a plate.
 
@@ -42,12 +42,12 @@ def __init__(
 
   def assign_child_resource(self, resource: Resource, location: Optional[Coordinate]=None,
                             reassign: bool = True):
-    location = location or get_child_location(resource)
     if len(self.children) >= 1:
       raise ValueError("There already is a plate in the plate reader.")
     if not isinstance(resource, Plate):
       raise ValueError("The resource must be a Plate.")
-    super().assign_child_resource(resource, location=location)
+
+    super().assign_child_resource(resource, location=location, reassign=reassign)
 
   def get_plate(self) -> Plate:
     if len(self.children) == 0:
@@ -87,7 +87,7 @@ async def read_fluorescence(
     emission_wavelength: int,
     focal_height: float
   ) -> List[List[float]]:
-    """ 
+    """
 
     Args:
       excitation_wavelength: The excitation wavelength to read the fluorescence at, in nanometers.

pylabrobot/resources/carrier.py

@@ -3,7 +3,7 @@
 import logging
 from typing import Generic, List, Optional, Type, TypeVar, Union
 
-from pylabrobot.resources.utils import get_child_location
+from pylabrobot.resources.resource_holder import ResourceHolderMixin
 
 from .coordinate import Coordinate
 from .plate import Plate
@@ -14,7 +14,7 @@
 logger = logging.getLogger("pylabrobot")
 
 
-class CarrierSite(Resource):
+class CarrierSite(ResourceHolderMixin, Resource):
   """ A single site within a carrier. """
 
   def __init__(self, name: str, size_x: float, size_y: float, size_z: float,
@@ -30,7 +30,6 @@ def assign_child_resource(
     reassign: bool = True
   ):
     self.resource = resource
-    location = location or get_child_location(resource)
     return super().assign_child_resource(resource, location, reassign)
 
   def unassign_child_resource(self, resource):
@@ -207,7 +206,6 @@ def __init__(self, name: str, size_x: float, size_y: float, size_z: float,
 
   def assign_child_resource(self, resource: Resource, location: Optional[Coordinate] = None,
                             reassign: bool = True):
-    location = location or self._get_child_location(resource)
     if isinstance(resource, ResourceStack):
       if not resource.direction == "z":
         raise ValueError("ResourceStack assigned to PlateCarrierSite must have direction 'z'")
@@ -219,7 +217,7 @@ def assign_child_resource(self, resource: Resource, location: Optional[Coordinat
                       f"resources, not {type(resource)}")
     return super().assign_child_resource(resource, location, reassign)
 
-  def _get_child_location(self, resource: Resource) -> Coordinate:
+  def _get_sinking_depth(self, resource: Resource) -> Coordinate:
     def get_plate_sinking_depth(plate: Plate):
       # Sanity check for equal well clearances / dz
       well_dz_set = {round(well.location.z, 2) for well in plate.get_all_children()
@@ -240,8 +238,10 @@ def get_plate_sinking_depth(plate: Plate):
         z_sinking_depth = get_plate_sinking_depth(first_child)
       resource.register_did_assign_resource_callback(self._update_resource_stack_location)
       self.register_did_unassign_resource_callback(self._deregister_resource_stack_callback)
+    return -Coordinate(z=z_sinking_depth)
 
-    return get_child_location(resource) - Coordinate(z=z_sinking_depth)
+  def get_default_child_location(self, resource: Resource) -> Coordinate:
+    return super().get_default_child_location(resource) + self._get_sinking_depth(resource)
 
   def _update_resource_stack_location(self, resource: Resource):
     """ Callback called when the lowest resource on a ResourceStack changes. Since the location of
@@ -254,7 +254,7 @@ def _update_resource_stack_location(self, resource: Resource):
     resource_stack = resource.parent
     assert isinstance(resource_stack, ResourceStack)
     if resource_stack.children[0] == resource:
-      resource_stack.location = self._get_child_location(resource)
+      resource_stack.location = self._get_sinking_depth(resource)
 
   def _deregister_resource_stack_callback(self, resource: Resource):
     """ Callback called when a ResourceStack (or child) is unassigned from this PlateCarrierSite."""

pylabrobot/resources/itemized_resource.py

@@ -33,7 +33,8 @@ def __init__(self, name: str, size_x: float, size_y: float, size_z: float,
                 ordered_items: Optional[Dict[str, T]] = None,
                 ordering: Optional[List[str]] = None,
                 category: Optional[str] = None,
-                model: Optional[str] = None):
+                model: Optional[str] = None,
+                **_kwargs):
     """ Initialize an itemized resource
 
     Args:

pylabrobot/resources/plate.py

@@ -4,6 +4,8 @@
 
 from typing import Dict, List, Optional, Sequence, Tuple, Union, cast, Literal
 
+from pylabrobot.resources.resource_holder import ResourceHolderMixin
+
 
 from .liquid import Liquid
 from .itemized_resource import ItemizedResource
@@ -44,7 +46,7 @@ def serialize(self) -> dict:
     return {**super().serialize(), "nesting_z_height": self.nesting_z_height}
 
 
-class Plate(ItemizedResource[Well]):
+class Plate(ResourceHolderMixin, ItemizedResource[Well]):
   """ Base class for Plate resources. """
 
   def __init__(

pylabrobot/resources/resource.py

@@ -50,6 +50,7 @@ def __init__(
     rotation: Optional[Rotation] = None,
     category: Optional[str] = None,
     model: Optional[str] = None,
+    **_kwargs
   ):
     self._name = name
     self._size_x = size_x

pylabrobot/resources/resource_holder.py

@@ -0,0 +1,41 @@
+from typing import Optional
+from pylabrobot.resources.coordinate import Coordinate
+from pylabrobot.resources.resource import Resource
+
+def get_child_location(resource: Resource) -> Coordinate:
+  """
+  If a resource is rotated, its rotated around its local origin. This does not always
+  match up with the parent resource's origin. This function calculates the difference
+  between the two origins and calculates the location of the resource correctly.
+  """
+  if not resource.rotation.y == resource.rotation.x == 0:
+    raise ValueError("Resource rotation must be 0 around the x and y axes")
+  if not resource.rotation.z % 90 == 0:
+    raise ValueError("Resource rotation must be a multiple of 90 degrees on the z axis")
+  location = {
+    0.0: Coordinate(x=0, y=0, z=0),
+    90.0: Coordinate(x=resource.get_size_y(), y=0, z=0),
+    180.0: Coordinate(x=resource.get_size_y(), y=resource.get_size_x(), z=0),
+    270.0: Coordinate(x=0, y=resource.get_size_x(), z=0),
+  }[resource.rotation.z % 360]
+  return location
+
+class ResourceHolderMixin:
+  """
+  A mixin class for resources that can hold other resources, like a plate or a lid.
+
+  This applies a linear transformation after the rotation to correctly place the child resource.
+  """
+
+  def get_default_child_location(self, resource: Resource) -> Coordinate:
+    return get_child_location(resource)
+
+  def assign_child_resource(
+    self,
+    resource: Resource,
+    location: Optional[Coordinate] = None,
+    reassign: bool = True
+  ):
+    location = location or self.get_default_child_location(resource)
+    # mypy doesn't play well with the Mixin pattern
+    return super().assign_child_resource(resource, location, reassign) # type: ignore

pylabrobot/resources/resource_stack.py

@@ -1,11 +1,12 @@
 from typing import List, Optional
 
+from pylabrobot.resources.resource_holder import ResourceHolderMixin, get_child_location
 from pylabrobot.resources.resource import Resource
 from pylabrobot.resources.coordinate import Coordinate
 from pylabrobot.resources.plate import Lid, Plate
 
 
-class ResourceStack(Resource):
+class ResourceStack(ResourceHolderMixin, Resource):
   """ ResourceStack represent a group of resources that are stacked together and act as a single
   unit. Stacks can grow be configured to be able to grow in x, y, or z direction. Stacks growing
   in the x direction are from left to right. Stacks growing in the y direction are from front to
@@ -95,26 +96,42 @@ def get_actual_resource_height(resource: Resource) -> float:
       return max(get_actual_resource_height(child) for child in self.children)
     return sum(get_actual_resource_height(child) for child in self.children)
 
-  def assign_child_resource(self, resource: Resource, location: Optional[Coordinate] = None,
-    reassign: bool = False):
+
+  def get_resource_stack_size(self) -> Coordinate:
     if self.direction == "x":
       resource_location = Coordinate(self.get_size_x(), 0, 0)
     elif self.direction == "y":
       resource_location = Coordinate(0, self.get_size_y(), 0)
     elif self.direction == "z":
       resource_location = Coordinate(0, 0, self.get_size_z())
-
-      # special handling for putting a lid on a plate
-      if len(self.children) > 0:
-        top_item = self.get_top_item()
-        if isinstance(resource, Lid) and isinstance(top_item, Plate):
-          resource_location.z -= resource.nesting_z_height
-          top_item.assign_child_resource(resource, location=resource_location)
-          return
     else:
       raise ValueError("self.direction must be one of 'x', 'y', or 'z'")
 
-    super().assign_child_resource(resource, location=resource_location)
+    return resource_location
+
+  def get_default_child_location(self, resource: Resource) -> Coordinate:
+    return super().get_default_child_location(resource) + self.get_resource_stack_size()
+
+  def assign_child_resource(self, resource: Resource, location: Optional[Coordinate] = None,
+    reassign: bool = False):
+
+    # special handling for putting a lid on a plate
+    if len(self.children) > 0:
+      top_item = self.get_top_item()
+      if isinstance(resource, Lid) and isinstance(top_item, Plate):
+        resource_location = self.get_resource_stack_size()
+        resource_location.z -= resource.nesting_z_height
+        top_item.assign_child_resource(
+          resource,
+          location=get_child_location(resource) + resource_location
+        )
+        return
+
+    super().assign_child_resource(
+      resource,
+      location=location,
+      reassign=reassign
+    )
 
   def unassign_child_resource(self, resource: Resource):
     if self.direction == "z" and resource != self.children[-1]: # no floating resources

pylabrobot/resources/utils.py

@@ -199,21 +199,3 @@ def query(
       z=z,
     ))
   return matched
-
-def get_child_location(resource: Resource) -> Coordinate:
-  """ 
-  If a resource is rotated, its rotated around its local origin. This does not always
-  match up with the parent resource's origin. This function calculates the difference
-  between the two origins and calculates the location of the resource correctly.
-  """
-  if not resource.rotation.y == resource.rotation.x == 0:
-    raise ValueError("Resource rotation must be 0 around the x and y axes")
-  if not resource.rotation.z % 90 == 0:
-    raise ValueError("Resource rotation must be a multiple of 90 degrees on the z axis")
-  location = {
-    0.0: Coordinate(x=0, y=0, z=0),
-    90.0: Coordinate(x=resource.get_size_y(), y=0, z=0),
-    180.0: Coordinate(x=resource.get_size_y(), y=resource.get_size_x(), z=0),
-    270.0: Coordinate(x=0, y=resource.get_size_x(), z=0),
-  }[resource.rotation.z % 360]
-  return location

pylabrobot/shaking/shaker.py

@@ -2,11 +2,12 @@
 from typing import Optional
 
 from pylabrobot.machines.machine import Machine
+from pylabrobot.resources.resource_holder import ResourceHolderMixin
 
 from .backend import ShakerBackend
 
 
-class Shaker(Machine):
+class Shaker(ResourceHolderMixin, Machine):
   """ A shaker machine """
 
   def __init__(