Multi-arm support and conditional head96 in LiquidHandler & Backends

pylabrobot/liquid_handling/backends/backend.py

@@ -37,6 +37,8 @@ class LiquidHandlerBackend(MachineBackend, metaclass=ABCMeta):
   def __init__(self):
     super().__init__()
     self.setup_finished = False
+    self.num_arms: int = 0
+    self.core96_head_installed: bool = False
     self._deck: Optional[Deck] = None
     self._head: Optional[Dict[int, TipTracker]] = None
     self._head96: Optional[Dict[int, TipTracker]] = None

pylabrobot/liquid_handling/backends/chatterbox.py

@@ -43,6 +43,8 @@ def __init__(self, num_channels: int = 8):
     """Initialize a chatter box backend."""
     super().__init__()
     self._num_channels = num_channels
+    self.num_arms = 1
+    self.core96_head_installed = True
 
   async def setup(self):
     await super().setup()

pylabrobot/liquid_handling/backends/hamilton/STAR_backend.py

@@ -1187,9 +1187,9 @@ def __init__(
       serial_number=serial_number,
     )
 
-    self.iswap_installed: Optional[bool] = None
-    self.autoload_installed: Optional[bool] = None
-    self.core96_head_installed: Optional[bool] = None
+    self.iswap_installed: bool = False
+    self.autoload_installed: bool = False
+    self.core96_head_installed: bool = False
 
     self._iswap_parked: Optional[bool] = None
     self._num_channels: Optional[int] = None
@@ -1448,6 +1448,7 @@ async def setup(
     self.autoload_installed = autoload_configuration_byte == "1"
     self.core96_head_installed = left_x_drive_configuration_byte_1[2] == "1"
     self.iswap_installed = left_x_drive_configuration_byte_1[1] == "1"
+    self.num_arms = 1 if self.iswap_installed else 0
     self._head96_information: Optional[Head96Information] = None
 
     initialized = await self.request_instrument_initialization_status()

pylabrobot/liquid_handling/backends/hamilton/STAR_chatterbox.py

@@ -10,17 +10,24 @@
 class STARChatterboxBackend(STARBackend):
   """Chatterbox backend for 'STAR'"""
 
-  def __init__(self, num_channels: int = 8, core96_head_installed: bool = True):
+  def __init__(
+    self,
+    num_channels: int = 8,
+    core96_head_installed: bool = True,
+    iswap_installed: bool = True,
+  ):
     """Initialize a chatter box backend.
 
     Args:
       num_channels: Number of pipetting channels (default: 8)
       core96_head_installed: Whether the CoRe 96 head is installed (default: True)
+      iswap_installed: Whether the iSWAP robotic arm is installed (default: True)
     """
     super().__init__()
     self._num_channels = num_channels
     self._iswap_parked = True
     self._core96_head_installed = core96_head_installed
+    self._iswap_installed = iswap_installed
 
   async def setup(
     self,
@@ -53,6 +60,7 @@ async def setup(
     # Use bitwise operations to check specific bits
     self.iswap_installed = bool(xl_value & 0b10)  # Check bit 1
     self.core96_head_installed = bool(xl_value & 0b100)  # Check bit 2
+    self.num_arms = 1 if self.iswap_installed else 0
 
     # Parse autoload from kb configuration byte
     configuration_data1 = bin(conf["kb"]).split("b")[-1].zfill(8)
@@ -132,12 +140,13 @@ async def request_extended_configuration(self):
     """
     # Calculate xl byte based on installed modules
     # Bit 0: (reserved)
-    # Bit 1: iSWAP (always True in this mock)
+    # Bit 1: iSWAP (based on __init__ parameter)
     # Bit 2: 96-head (based on __init__ parameter)
-    xl_value = 0b10  # iSWAP installed (bit 1)
+    xl_value = 0
+    if self._iswap_installed:
+      xl_value |= 0b10  # Add iSWAP (bit 1)
     if self._core96_head_installed:
       xl_value |= 0b100  # Add 96-head (bit 2)
-    # Result: xl = 6 (0b110) if 96-head installed, 2 (0b10) if not
 
     self._extended_conf = {
       "ka": 65537,

pylabrobot/liquid_handling/backends/hamilton/STAR_tests.py

@@ -178,6 +178,7 @@ def __init__(self):
   async def setup(self) -> None:  # type: ignore
     self._num_channels = 8
     self.iswap_installed = True
+    self.num_arms = 1
     self.core96_head_installed = True
     self._core_parked = True
 
@@ -262,6 +263,7 @@ def __init__(self, name: str):
     self.STAR._num_channels = 8
     self.STAR.core96_head_installed = True
     self.STAR.iswap_installed = True
+    self.STAR.num_arms = 1
     self.STAR.setup = unittest.mock.AsyncMock()
     self.STAR._core_parked = True
     self.STAR._iswap_parked = True
@@ -1093,6 +1095,7 @@ async def asyncSetUp(self):
     self.STAR._num_channels = 8
     self.STAR.core96_head_installed = True
     self.STAR.iswap_installed = True
+    self.STAR.num_arms = 1
     self.STAR.setup = unittest.mock.AsyncMock()
     self.STAR._core_parked = True
     self.STAR._iswap_parked = True
@@ -1222,6 +1225,7 @@ async def asyncSetUp(self):
     self.star._num_channels = 8
     self.star.core96_head_installed = True
     self.star.iswap_installed = True
+    self.star.num_arms = 1
     self.star.setup = unittest.mock.AsyncMock()
     self.star._core_parked = True
     self.star._iswap_parked = True
@@ -1418,6 +1422,7 @@ async def asyncSetUp(self):
     self.backend._num_channels = 8
     self.backend.core96_head_installed = True
     self.backend.iswap_installed = True
+    self.backend.num_arms = 1
     self.backend.setup = unittest.mock.AsyncMock()
     self.backend._core_parked = True
     self.backend._iswap_parked = True

pylabrobot/liquid_handling/backends/hamilton/vantage_tests.py

@@ -217,6 +217,7 @@ async def setup(self) -> None:  # type: ignore
     self.setup_finished = True
     self._num_channels = 8
     self.iswap_installed = True
+    self.num_arms = 1
     self.core96_head_installed = True
 
   async def send_command(

pylabrobot/liquid_handling/backends/serializing_backend.py

@@ -30,6 +30,8 @@ class SerializingBackend(LiquidHandlerBackend, metaclass=ABCMeta):
   def __init__(self, num_channels: int):
     LiquidHandlerBackend.__init__(self)
     self._num_channels = num_channels
+    self.num_arms = 1
+    self.core96_head_installed = True
 
   @property
   def num_channels(self) -> int:

pylabrobot/liquid_handling/liquid_handler.py

@@ -145,7 +145,16 @@ def __init__(
     self.location = Coordinate.zero()
     super().assign_child_resource(deck, location=deck.location or Coordinate.zero())
 
-    self._resource_pickup: Optional[ResourcePickup] = None
+    self._resource_pickups: Dict[int, Optional[ResourcePickup]] = {}
+
+  @property
+  def _resource_pickup(self) -> Optional[ResourcePickup]:
+    """Backward-compatible access to the first arm's pickup state."""
+    return self._resource_pickups.get(0)
+
+  @_resource_pickup.setter
+  def _resource_pickup(self, value: Optional[ResourcePickup]) -> None:
+    self._resource_pickups[0] = value
 
   async def setup(self, **backend_kwargs):
     """Prepare the robot for use."""
@@ -158,16 +167,58 @@ async def setup(self, **backend_kwargs):
     await super().setup(**backend_kwargs)
 
     self.head = {c: TipTracker(thing=f"Channel {c}") for c in range(self.backend.num_channels)}
-    self.head96 = {c: TipTracker(thing=f"Channel {c}") for c in range(96)}
 
-    self._resource_pickup = None
+    self.head96 = (
+      {c: TipTracker(thing=f"Channel {c}") for c in range(96)}
+      if self.backend.core96_head_installed
+      else {}
+    )
+
+    self.backend.set_heads(head=self.head, head96=self.head96 or None)
+
+    for tracker in self.head.values():
+      tracker.register_callback(self._state_updated)
+    for tracker in self.head96.values():
+      tracker.register_callback(self._state_updated)
+
+    self._resource_pickups = {a: None for a in range(self.backend.num_arms)}
 
   def serialize_state(self) -> Dict[str, Any]:
     """Serialize the state of this liquid handler. Use :meth:`~Resource.serialize_all_states` to
     serialize the state of the liquid handler and all children (the deck)."""
 
     head_state = {channel: tracker.serialize() for channel, tracker in self.head.items()}
-    return {"head_state": head_state}
+    head96_state = (
+      {channel: tracker.serialize() for channel, tracker in self.head96.items()}
+      if self.head96
+      else None
+    )
+    arm_state: Optional[Dict[int, Any]]
+    if self._resource_pickups:
+      arm_state = {}
+      for arm_id, pickup in self._resource_pickups.items():
+        if pickup is not None:
+          res = pickup.resource
+          arm_entry: Dict[str, Any] = {
+            "has_resource": True,
+            "resource_name": res.name,
+            "resource_type": type(res).__name__,
+            "direction": pickup.direction.name,
+            "pickup_distance_from_top": pickup.pickup_distance_from_top,
+            "size_x": res.get_size_x(),
+            "size_y": res.get_size_y(),
+            "size_z": res.get_size_z(),
+          }
+          if hasattr(res, "num_items_x"):
+            arm_entry["num_items_x"] = res.num_items_x
+          if hasattr(res, "num_items_y"):
+            arm_entry["num_items_y"] = res.num_items_y
+          arm_state[arm_id] = arm_entry
+        else:
+          arm_state[arm_id] = None
+    else:
+      arm_state = None
+    return {"head_state": head_state, "head96_state": head96_state, "arm_state": arm_state}
 
   def load_state(self, state: Dict[str, Any]):
     """Load the liquid handler state from a file. Use :meth:`~Resource.load_all_state` to load the
@@ -177,6 +228,14 @@ def load_state(self, state: Dict[str, Any]):
     for channel, tracker_state in head_state.items():
       self.head[channel].load_state(tracker_state)
 
+    head96_state = state.get("head96_state", {})
+    if head96_state and self.head96:
+      for channel, tracker_state in head96_state.items():
+        self.head96[channel].load_state(tracker_state)
+
+    # arm_state is informational only (read via serialize_state); no load needed since
+    # _resource_pickup is set/cleared by pick_up_resource/drop_resource at runtime.
+
   def update_head_state(self, state: Dict[int, Optional[Tip]]):
     """Update the state of the liquid handler head.
 
@@ -1964,6 +2023,9 @@ async def pick_up_resource(
       direction=direction,
     )
 
+    if self.setup_finished and not self._resource_pickups:
+      raise RuntimeError("No robotic arm is installed on this liquid handler.")
+
     if pickup_distance_from_top is None:
       if resource.preferred_pickup_location is not None:
         logger.debug(
@@ -2001,6 +2063,8 @@ async def pick_up_resource(
       self._resource_pickup = None
       raise e
 
+    self._state_updated()
+
   async def move_picked_up_resource(
     self,
     to: Coordinate,
@@ -2166,6 +2230,7 @@ async def drop_resource(
     result = await self.backend.drop_resource(drop=drop, **backend_kwargs)
 
     self._resource_pickup = None
+    self._state_updated()
 
     # we rotate the resource on top of its original rotation. So in order to set the new rotation,
     # we have to subtract its current rotation.

pylabrobot/liquid_handling/liquid_handler_tests.py

@@ -66,6 +66,8 @@ def _create_mock_backend(num_channels: int = 8):
   """Create a mock LiquidHandlerBackend with the specified number of channels."""
   mock = unittest.mock.create_autospec(LiquidHandlerBackend, instance=True)
   type(mock).num_channels = PropertyMock(return_value=num_channels)
+  mock.num_arms = 1
+  mock.core96_head_installed = True
   mock.can_pick_up_tip.return_value = True
   return mock
 

pylabrobot/resources/tip_tracker.py

@@ -111,6 +111,10 @@ def remove_tip(self, commit: bool = False) -> None:
   def commit(self) -> None:
     """Commit the pending operations."""
     self._tip = self._pending_tip
+    # Propagate state-update callback to the tip's volume tracker so that
+    # aspirate/dispense volume changes trigger a visualizer refresh.
+    if self._tip is not None and self._callback is not None:
+      self._tip.tracker.register_callback(self._callback)
     if self._callback is not None:
       self._callback()
 

pylabrobot/server/liquid_handling_api_tests.py

@@ -26,6 +26,8 @@ def _create_mock_backend(num_channels: int = 8):
   """Create a mock LiquidHandlerBackend with the specified number of channels."""
   mock = unittest.mock.create_autospec(LiquidHandlerBackend, instance=True)
   type(mock).num_channels = PropertyMock(return_value=num_channels)
+  mock.num_arms = 1
+  mock.core96_head_installed = True
   mock.can_pick_up_tip.return_value = True
   return mock