Add FK-based STARBackend.iswap_request_pose() + iswap_request_joint_state()

pylabrobot/liquid_handling/backends/hamilton/STAR_backend.py

@@ -34,6 +34,7 @@
   from typing import Concatenate, ParamSpec
 
 from pylabrobot import audio
+from pylabrobot.arms.standard import CartesianCoords
 from pylabrobot.heating_shaking.hamilton_backend import HamiltonHeaterShakerInterface
 from pylabrobot.liquid_handling.backends.hamilton.base import (
   HamiltonLiquidHandler,
@@ -1322,6 +1323,36 @@ class Head96Information:
 class STARBackend(HamiltonLiquidHandler, HamiltonHeaterShakerInterface):
   """Interface for the Hamilton STARBackend."""
 
+  class iSWAPAxis(enum.IntEnum):
+    """Axis index for `iswap_request_joint_state` dicts.
+
+    Units are axis-implicit (matches PF400's `Dict[int, float]` keyed by `PFAxis`):
+    prismatic axes (X/Y/Z, GRIPPER) are in mm; revolute axes (ROTATION, WRIST)
+    are in degrees. Z is the rotation-drive-bottom Z (sits 13 mm above the
+    gripper finger plane).
+    """
+
+    X = 1  # X-arm carriage (rotation-drive X in deck coords)
+    Y = 2  # Y carriage at rotation drive
+    Z = 3  # Z carriage at rotation drive (rotation-drive-bottom Z, deck coords).
+    # NOT the grip-center Z - that sits ~13 mm below this plane and only
+    # appears in iswap_request_pose().location.z.
+    ROTATION = 4  # W joint 1, signed from calibrated FRONT (deg)
+    WRIST = 5  # T joint 2, signed from motor zero (deg)
+    GRIPPER = 6  # gripper jaw opening width
+
+    @property
+    def is_in_kinematic_chain(self) -> bool:
+      """Whether this axis enters forward kinematics. The gripper is an
+      addressable actuator but not a chain member - it changes what is held,
+      not where the gripper frame is."""
+      return self is not STARBackend.iSWAPAxis.GRIPPER
+
+    @property
+    def is_revolute(self) -> bool:
+      """True for revolute (rotary, deg) axes; False for prismatic (linear, mm)."""
+      return self in (STARBackend.iSWAPAxis.ROTATION, STARBackend.iSWAPAxis.WRIST)
+
   PIP_X_MIN_WITH_LEFT_SIDE_PANEL: float = 320.0
   HEAD96_X_MIN_WITH_LEFT_SIDE_PANEL: float = 0.0
 
@@ -1375,6 +1406,8 @@ def __init__(
     self._iswap_wrist_drive_predefined_increments: Optional[
       Dict[STARBackend.WristDriveOrientation, int]
     ] = None
+    self._iswap_link_1_length: Optional[float] = None
+    self._iswap_link_2_length: Optional[float] = None
     self.core_adjustment = Coordinate.zero()
     self._unsafe = UnSafe(self)
 
@@ -1790,6 +1823,9 @@ async def set_up_iswap():
           STARBackend.WristDriveOrientation.REVERSE: wrist_predefined["reverse"],
         }
 
+        self._iswap_link_1_length = await self.iswap_request_link_1_length()
+        self._iswap_link_2_length = await self.iswap_request_link_2_length()
+
     async def set_up_core96_head():
       if self.extended_conf.left_x_drive.core_96_head_installed and not skip_core96_head:
         # Initialize 96-head
@@ -10735,6 +10771,120 @@ async def iswap_wrist_drive_rotate_to_angle(
       wrist_current_limit=current_limit,
     )
 
+  # -----------------------------------------------------------------------
+  # iSWAP: Forward Kinematics
+  # -----------------------------------------------------------------------
+
+  @staticmethod
+  def _iswap_fk(
+    joints: Dict["STARBackend.iSWAPAxis", float],
+    link_1_length: float,
+    link_2_length: float,
+    wrist_straight_angle: float,
+  ) -> CartesianCoords:
+    """Pure forward-kinematics map: joint state -> gripper pose (no I/O).
+
+    Takes an `iSWAPAxis`-keyed joint dict and the calibrated link lengths /
+    STRAIGHT angle, returns just the gripper-frame pose (grip-center location
+    + yaw). No intermediate frames in the return; callers that need the wrist
+    XY can recompute trivially from joints + L1.
+
+    Sign convention follows right-hand rule about +Z (CCW positive looking
+    down). Yaw is the deck-frame direction of link 2:
+    `alpha_2 = (W - 90) + (T - T_STRAIGHT)`, with 0 deg = +x deck-right.
+    Z: rotation-drive-bottom Z minus `iswap_rotation_drive_z_offset_above_finger_mm`
+    (13 mm).
+    """
+    rotation_drive_angle = joints[STARBackend.iSWAPAxis.ROTATION]
+    wrist_drive_angle = joints[STARBackend.iSWAPAxis.WRIST]
+    base_x = joints[STARBackend.iSWAPAxis.X]
+    base_y = joints[STARBackend.iSWAPAxis.Y]
+    base_z = joints[STARBackend.iSWAPAxis.Z]
+
+    link_1_deck_angle = rotation_drive_angle - 90.0
+    link_2_deck_angle = link_1_deck_angle + (wrist_drive_angle - wrist_straight_angle)
+
+    alpha_1_rad = math.radians(link_1_deck_angle)
+    alpha_2_rad = math.radians(link_2_deck_angle)
+
+    grip_x = base_x + link_1_length * math.cos(alpha_1_rad) + link_2_length * math.cos(alpha_2_rad)
+    grip_y = base_y + link_1_length * math.sin(alpha_1_rad) + link_2_length * math.sin(alpha_2_rad)
+    grip_z = base_z - STARBackend.iswap_rotation_drive_z_offset_above_finger_mm
+
+    return CartesianCoords(
+      location=Coordinate(x=grip_x, y=grip_y, z=grip_z),
+      rotation=Rotation(z=link_2_deck_angle),
+    )
+
+  async def iswap_request_joint_state(self) -> Dict[int, float]:
+    """Snapshot of the iSWAP's current joint state, keyed by `iSWAPAxis`.
+
+    Composes the per-axis request methods into a single read-through dict.
+    Units are axis-implicit (see `iSWAPAxis`): X/Y/Z/GRIPPER in mm, ROTATION
+    and WRIST in degrees.
+
+    Raises:
+      RuntimeError: if iSWAP is not installed or if `setup()` has not run
+        (the rotation-drive angle reader needs the predefined-stop table).
+    """
+    if not self.extended_conf.left_x_drive.iswap_installed:
+      raise RuntimeError("iSWAP is not installed")
+
+    return {
+      STARBackend.iSWAPAxis.X: await self.iswap_rotation_drive_request_x(),
+      STARBackend.iSWAPAxis.Y: await self.iswap_rotation_drive_request_y(),
+      STARBackend.iSWAPAxis.Z: await self.iswap_rotation_drive_request_z(),
+      STARBackend.iSWAPAxis.ROTATION: await self.iswap_rotation_drive_request_angle(),
+      STARBackend.iSWAPAxis.WRIST: await self.iswap_wrist_drive_request_angle(),
+      STARBackend.iSWAPAxis.GRIPPER: await self.iswap_gripper_request_width(),
+    }
+
+  async def iswap_request_pose(self) -> CartesianCoords:
+    """Compute the iSWAP gripper pose via forward kinematics.
+
+    FK-based alternative to `request_iswap_position` (C0 QG), which is
+    firmware-state-dependent and only returns correct values after certain
+    preceding commands. Reads the joint state directly via
+    `iswap_request_joint_state` and runs `_iswap_fk` against the link lengths
+    cached during `setup()`.
+
+    Returns:
+      `CartesianCoords` with `location` = grip-center deck coordinates (mm) and
+      `rotation.z` = gripper yaw (deg, deck-frame, 0 = +x; `rotation.x`/`.y` = 0
+      since the gripper plane stays parallel to the deck).
+
+    Raises:
+      RuntimeError: if iSWAP is not installed or if `setup()` has not populated
+        the wrist STRAIGHT calibration / cached link lengths.
+    """
+    if (
+      self._iswap_wrist_drive_predefined_increments is None
+      or self._iswap_link_1_length is None
+      or self._iswap_link_2_length is None
+    ):
+      raise RuntimeError(
+        "iSWAP setup state not loaded (wrist predefined positions / link lengths); "
+        "ensure the iSWAP is installed and `setup()` has run."
+      )
+
+    wrist_straight_increments = self._iswap_wrist_drive_predefined_increments[
+      STARBackend.WristDriveOrientation.STRAIGHT
+    ]
+    wrist_straight_angle = STARBackend._iswap_wrist_drive_increments_to_angle(
+      wrist_straight_increments
+    )
+
+    joints = {
+      STARBackend.iSWAPAxis(k): v for k, v in (await self.iswap_request_joint_state()).items()
+    }
+
+    return STARBackend._iswap_fk(
+      joints=joints,
+      link_1_length=self._iswap_link_1_length,
+      link_2_length=self._iswap_link_2_length,
+      wrist_straight_angle=wrist_straight_angle,
+    )
+
   # -----------------------------------------------------------------------
   # iSWAP: Combined Rotation-Wrist Moves
   # -----------------------------------------------------------------------

pylabrobot/liquid_handling/backends/hamilton/STAR_tests.py

@@ -5,6 +5,7 @@
 import unittest.mock
 from typing import Literal, cast
 
+from pylabrobot.arms.standard import CartesianCoords
 from pylabrobot.liquid_handling import LiquidHandler
 from pylabrobot.liquid_handling.standard import GripDirection, Pickup
 from pylabrobot.plate_reading import PlateReader
@@ -147,6 +148,157 @@ def _any_write_and_read_command_call(cmd):
   )
 
 
+class TestiSWAPForwardKinematics(unittest.TestCase):
+  """Geometry of `STARBackend._iswap_fk` (pure FK, no I/O).
+
+  Verifies the canonical (W, T) configurations against the docstring examples
+  in `request_iswap_wrist_drive_orientation`:
+    - W=FRONT (0)  + T=STRAIGHT  -> arm extends in -y (front of deck)
+    - W=LEFT (-90) + T=STRAIGHT  -> arm extends in -x (left of deck)
+    - W=RIGHT(+90) + T=STRAIGHT  -> arm extends in +x (right of deck)
+    - W=FRONT (0)  + T=RIGHT     -> arm tip ends up to the left (-x) of rot. drive
+  Uses factory-default link lengths (138 mm each) and the factory-default
+  STRAIGHT angle (~-45 deg). Asserts only on the public pose contract
+  (`location` + `rotation.z`).
+  """
+
+  L1 = 138.0
+  L2 = 138.0
+  T_STRAIGHT = -45.0  # factory-default STRAIGHT calibration (EEPROM-dependent in practice)
+  Z_OFFSET = STARBackend.iswap_rotation_drive_z_offset_above_finger_mm
+  BASE_X, BASE_Y, BASE_Z = 100.0, 500.0, 200.0
+  # Gripper jaw width: hardware range ~71-134 mm (drive min/max increments).
+  # FK doesn't read this axis, but the joint dict represents full state, so
+  # use a realistic mid-range plate-grip value.
+  GRIPPER_WIDTH = 90.0
+
+  def _fk(self, w: float, t: float) -> CartesianCoords:
+    joints = {
+      STARBackend.iSWAPAxis.X: self.BASE_X,
+      STARBackend.iSWAPAxis.Y: self.BASE_Y,
+      STARBackend.iSWAPAxis.Z: self.BASE_Z,
+      STARBackend.iSWAPAxis.ROTATION: w,
+      STARBackend.iSWAPAxis.WRIST: t,
+      STARBackend.iSWAPAxis.GRIPPER: self.GRIPPER_WIDTH,
+    }
+    return STARBackend._iswap_fk(
+      joints=joints,
+      link_1_length=self.L1,
+      link_2_length=self.L2,
+      wrist_straight_angle=self.T_STRAIGHT,
+    )
+
+  def test_front_straight_extends_in_minus_y(self):
+    pose = self._fk(w=0.0, t=self.T_STRAIGHT)
+    self.assertAlmostEqual(pose.location.x, self.BASE_X, places=6)
+    self.assertAlmostEqual(pose.location.y, self.BASE_Y - (self.L1 + self.L2), places=6)
+    self.assertAlmostEqual(pose.location.z, self.BASE_Z - self.Z_OFFSET, places=6)
+    self.assertAlmostEqual(pose.rotation.z, -90.0, places=6)
+
+  def test_left_straight_extends_in_minus_x(self):
+    pose = self._fk(w=-90.0, t=self.T_STRAIGHT)
+    self.assertAlmostEqual(pose.location.x, self.BASE_X - (self.L1 + self.L2), places=6)
+    self.assertAlmostEqual(pose.location.y, self.BASE_Y, places=6)
+    self.assertAlmostEqual(pose.rotation.z, -180.0, places=6)
+
+  def test_front_right_extends_in_minus_x(self):
+    """W=FRONT + T=RIGHT (-135 deg motor) -> gripper points to deck-left."""
+    pose = self._fk(w=0.0, t=-135.0)
+    # Link 1 points -y from base; link 2 is bent right (CW by 90 deg) -> points -x.
+    self.assertAlmostEqual(pose.location.x, self.BASE_X - self.L2, places=6)
+    self.assertAlmostEqual(pose.location.y, self.BASE_Y - self.L1, places=6)
+    self.assertAlmostEqual(pose.rotation.z, -180.0, places=6)
+
+  def test_reverse_folds_arm_back_onto_base_xy(self):
+    """T=REVERSE (+135) folds link 2 back 180 deg from link 1 -> tip XY = base XY."""
+    pose = self._fk(w=0.0, t=+135.0)
+    self.assertAlmostEqual(pose.location.x, self.BASE_X, places=6)
+    self.assertAlmostEqual(pose.location.y, self.BASE_Y, places=6)
+
+
+class TestiSWAPAxisPredicates(unittest.TestCase):
+  """Predicates on `STARBackend.iSWAPAxis` classify axes by kinematic role / unit."""
+
+  def test_is_in_kinematic_chain(self):
+    Axis = STARBackend.iSWAPAxis
+    for a in (Axis.X, Axis.Y, Axis.Z, Axis.ROTATION, Axis.WRIST):
+      self.assertTrue(a.is_in_kinematic_chain, f"{a.name} should be in the chain")
+    self.assertFalse(Axis.GRIPPER.is_in_kinematic_chain, "GRIPPER should NOT be in the chain")
+
+
+class TestiSWAPRequestJointState(unittest.IsolatedAsyncioTestCase):
+  """`iswap_request_joint_state` composes the per-axis request methods into one dict."""
+
+  def _make_backend(self) -> STARBackend:
+    b = STARBackend()
+    b._extended_conf = _DEFAULT_EXTENDED_CONFIGURATION
+    b.iswap_rotation_drive_request_x = unittest.mock.AsyncMock(return_value=100.0)
+    b.iswap_rotation_drive_request_y = unittest.mock.AsyncMock(return_value=500.0)
+    b.iswap_rotation_drive_request_z = unittest.mock.AsyncMock(return_value=200.0)
+    b.iswap_rotation_drive_request_angle = unittest.mock.AsyncMock(return_value=0.0)
+    b.iswap_wrist_drive_request_angle = unittest.mock.AsyncMock(return_value=-45.0)
+    b.iswap_gripper_request_width = unittest.mock.AsyncMock(return_value=90.0)
+    return b
+
+  async def test_returns_full_axis_dict(self):
+    b = self._make_backend()
+    joints = await b.iswap_request_joint_state()
+    Axis = STARBackend.iSWAPAxis
+    self.assertEqual(
+      joints,
+      {
+        Axis.X: 100.0,
+        Axis.Y: 500.0,
+        Axis.Z: 200.0,
+        Axis.ROTATION: 0.0,
+        Axis.WRIST: -45.0,
+        Axis.GRIPPER: 90.0,
+      },
+    )
+
+
+class TestiSWAPRequestPose(unittest.IsolatedAsyncioTestCase):
+  """`iswap_request_pose` reads joints + runs FK against the cached link lengths."""
+
+  def _make_backend(self) -> STARBackend:
+    b = STARBackend()
+    b._extended_conf = _DEFAULT_EXTENDED_CONFIGURATION
+    b._iswap_link_1_length = 138.0
+    b._iswap_link_2_length = 138.0
+    b._iswap_wrist_drive_predefined_increments = {
+      STARBackend.WristDriveOrientation.RIGHT: -26577,
+      STARBackend.WristDriveOrientation.STRAIGHT: -8859,
+      STARBackend.WristDriveOrientation.LEFT: 8859,
+      STARBackend.WristDriveOrientation.REVERSE: 26577,
+    }
+    b.iswap_rotation_drive_request_x = unittest.mock.AsyncMock(return_value=100.0)
+    b.iswap_rotation_drive_request_y = unittest.mock.AsyncMock(return_value=500.0)
+    b.iswap_rotation_drive_request_z = unittest.mock.AsyncMock(return_value=200.0)
+    b.iswap_rotation_drive_request_angle = unittest.mock.AsyncMock(return_value=0.0)
+    b.iswap_wrist_drive_request_angle = unittest.mock.AsyncMock(return_value=-45.0)
+    b.iswap_gripper_request_width = unittest.mock.AsyncMock(return_value=90.0)
+    return b
+
+  async def test_front_straight_pose(self):
+    """Canonical W=0 / T=-45 / base=(100, 500, 200) -> grip ~ (100, 224, 187), yaw ~ -90°.
+
+    Verifies the full I/O path (per-axis reads -> joint state -> FK -> pose). EEPROM
+    STRAIGHT (-8859 incr) maps to ~-45.0007 deg, so the canonical -90° yaw lands at
+    -89.999° and grip x picks up a ~0.002 mm offset - this is intentional per-machine
+    calibration drift, not an FK bug.
+    """
+    b = self._make_backend()
+    pose = await b.iswap_request_pose()
+    self.assertIsInstance(pose, CartesianCoords)
+    self.assertAlmostEqual(pose.location.x, 100.0, places=2)
+    self.assertAlmostEqual(pose.location.y, 224.0, places=3)
+    self.assertAlmostEqual(pose.location.z, 187.0, places=6)
+    self.assertAlmostEqual(pose.rotation.z, -90.0, places=2)
+    # rotation.x / y are always 0 - the gripper plane stays parallel to the deck.
+    self.assertEqual(pose.rotation.x, 0.0)
+    self.assertEqual(pose.rotation.y, 0.0)
+
+
 class TestSTARUSBComms(unittest.IsolatedAsyncioTestCase):
   """Test that USB data is parsed correctly."""