Fix 3D rotation composition with quaternions

pylabrobot/resources/resource.py

@@ -364,7 +364,9 @@ def get_absolute_location(self, x: str = "l", y: str = "f", z: str = "b") -> Coo
     # carry no location yet still rotate what hangs from them, so the rotation is taken from the
     # whole tree rather than from the chain.
     rotation = chain[0].get_absolute_rotation()
-    matrix = rotation.get_rotation_matrix() if (rotation.x or rotation.y or rotation.z) else None
+    matrix = (
+      None if rotation._quaternion == (1.0, 0.0, 0.0, 0.0) else rotation.get_rotation_matrix()
+    )
     position = cast(Coordinate, chain[0].location)
 
     # 2b. Accumulate each child's offset in its parent's frame
@@ -376,7 +378,7 @@ def get_absolute_location(self, x: str = "l", y: str = "f", z: str = "b") -> Coo
         position += Coordinate(*matrix_vector_multiply_3x3(matrix, anchor.vector())) + Coordinate(
           *matrix_vector_multiply_3x3(matrix, location.vector())
         )
-      if child.rotation.x or child.rotation.y or child.rotation.z:
+      if child.rotation._quaternion != (1.0, 0.0, 0.0, 0.0):
         rotation = rotation + child.rotation
         matrix = rotation.get_rotation_matrix()
 
@@ -891,11 +893,9 @@ def location(self, location: Optional[Coordinate]) -> None:
       self._state_updated()
 
   def rotate(self, x: float = 0, y: float = 0, z: float = 0):
-    """Rotate counter-clockwise by the given number of degrees."""
+    """Rotate counter-clockwise around the parent-coordinate axes by the given degrees."""
 
-    self.rotation.x = (self.rotation.x + x) % 360
-    self.rotation.y = (self.rotation.y + y) % 360
-    self.rotation.z = (self.rotation.z + z) % 360
+    self.rotation._prepend(Rotation(x=x, y=y, z=z))
     # Rotation is part of the resource's state; notify subscribers (e.g. the
     # Visualizer) so they can re-render.
     self._state_updated()

pylabrobot/resources/resource_holder.py

@@ -1,3 +1,4 @@
+import math
 from typing import Any, Mapping, Optional
 
 from pylabrobot.resources.coordinate import Coordinate
@@ -13,14 +14,16 @@ def get_child_location(resource: Resource) -> Coordinate:
   """
   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:
+  z = resource.rotation.z % 360
+  snapped_z = round(z / 90) * 90
+  if not math.isclose(z, snapped_z, rel_tol=0, abs_tol=1e-7):
     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_x(), y=resource.get_size_y(), z=0),
     270.0: Coordinate(x=0, y=resource.get_size_x(), z=0),
-  }[resource.rotation.z % 360]
+  }[snapped_z % 360]
   return location
 
 

pylabrobot/resources/resource_holder_tests.py

@@ -1,5 +1,6 @@
 import unittest
 
+from .coordinate import Coordinate
 from .resource import Resource
 from .resource_holder import ResourceHolder
 
@@ -15,6 +16,20 @@ def test_assign_via_property(self):
     self.assertEqual(self.holder.resource, self.resource)
     self.assertEqual(self.resource.parent, self.holder)
 
+  def test_assign_with_yaw_rounding_error(self):
+    self.resource.rotation.z = 269.999999999999
+
+    self.holder.assign_child_resource(self.resource)
+
+    self.assertEqual(self.resource.location, Coordinate(0, 1, 0))
+    self.assertEqual(self.resource.rotation.z, 269.999999999999)
+
+  def test_reject_misaligned_yaw(self):
+    self.resource.rotation.z = 269.99
+
+    with self.assertRaisesRegex(ValueError, "multiple of 90 degrees"):
+      self.holder.assign_child_resource(self.resource)
+
   def test_over_assignment(self):
     self.holder.resource = self.resource
     with self.assertRaises(ValueError):

pylabrobot/resources/resource_tests.py

@@ -24,7 +24,7 @@
 from pylabrobot.resources.resource import Resource
 from pylabrobot.resources.rotation import Rotation
 from pylabrobot.resources.tip import Tip
-from pylabrobot.utils.linalg import matrix_vector_multiply_3x3
+from pylabrobot.utils.linalg import matrix_multiply_3x3, matrix_vector_multiply_3x3
 
 
 def _make_test_deck() -> Deck:
@@ -438,6 +438,45 @@ def level_by_level(resource: Resource, x="l", y="f", z="b") -> Coordinate:
           for anchors in (("l", "f", "b"), ("c", "c", "c"), ("r", "b", "t")):
             self.assertEqual(node.get_absolute_location(*anchors), level_by_level(node, *anchors))
 
+  def test_rotate_composes_around_fixed_axes(self):
+    resource = Resource("resource", size_x=10, size_y=10, size_z=10, rotation=Rotation(z=90))
+    expected = matrix_multiply_3x3(
+      Rotation(x=90).get_rotation_matrix(),
+      resource.rotation.get_rotation_matrix(),
+    )
+
+    resource.rotate(x=90)
+
+    actual = resource.rotation.get_rotation_matrix()
+    for actual_row, expected_row in zip(actual, expected):
+      for actual_value, expected_value in zip(actual_row, expected_row):
+        self.assertAlmostEqual(actual_value, expected_value)
+
+  def test_rotate_keeps_angles_normalized(self):
+    resource = Resource("resource", size_x=10, size_y=10, size_z=10)
+    rotation = resource.rotation
+
+    resource.rotate(z=270)
+    self.assertIs(resource.rotation, rotation)
+    self.assertEqual(resource.rotation.z, 270)
+
+    resource.rotate(z=90)
+    self.assertEqual(resource.rotation.z, 0)
+
+  def test_absolute_rotation_composes_parent_and_child(self):
+    parent = Resource("parent", size_x=10, size_y=10, size_z=10, rotation=Rotation(x=90))
+    child = Resource("child", size_x=5, size_y=5, size_z=5, rotation=Rotation(z=90))
+    parent.assign_child_resource(child, location=Coordinate.zero())
+    expected = matrix_multiply_3x3(
+      parent.rotation.get_rotation_matrix(),
+      child.rotation.get_rotation_matrix(),
+    )
+
+    actual = child.get_absolute_rotation().get_rotation_matrix()
+    for actual_row, expected_row in zip(actual, expected):
+      for actual_value, expected_value in zip(actual_row, expected_row):
+        self.assertAlmostEqual(actual_value, expected_value)
+
 
 class TestResourceCallback(unittest.TestCase):
   def setUp(self) -> None:

pylabrobot/resources/rotation.py

@@ -1,71 +1,186 @@
 import math
 
 from pylabrobot.serializer import SerializableMixin
-from pylabrobot.utils.linalg import matrix_multiply_3x3
+
+# Quaternions use (w, x, y, z); Euler angles retain the Rz * Ry * Rx convention.
+_Quaternion = tuple[float, float, float, float]
+
+
+def _normalize_quaternion(quaternion: _Quaternion) -> _Quaternion:
+  """Return a unit-length quaternion."""
+  norm = math.sqrt(sum(component * component for component in quaternion))
+  w, x, y, z = quaternion
+  return w / norm, x / norm, y / norm, z / norm
+
+
+def _multiply_quaternions(left: _Quaternion, right: _Quaternion) -> _Quaternion:
+  """Compose two quaternions using the Hamilton product."""
+  left_w, left_x, left_y, left_z = left
+  right_w, right_x, right_y, right_z = right
+  return _normalize_quaternion(
+    (
+      left_w * right_w - left_x * right_x - left_y * right_y - left_z * right_z,
+      left_w * right_x + left_x * right_w + left_y * right_z - left_z * right_y,
+      left_w * right_y - left_x * right_z + left_y * right_w + left_z * right_x,
+      left_w * right_z + left_x * right_y - left_y * right_x + left_z * right_w,
+    )
+  )
+
+
+def _quaternion_from_euler(x: float, y: float, z: float) -> _Quaternion:
+  """Convert roll, pitch, and yaw in degrees to a quaternion."""
+  half_roll = math.radians(x) / 2
+  half_pitch = math.radians(y) / 2
+  half_yaw = math.radians(z) / 2
+  cos_roll, sin_roll = math.cos(half_roll), math.sin(half_roll)
+  cos_pitch, sin_pitch = math.cos(half_pitch), math.sin(half_pitch)
+  cos_yaw, sin_yaw = math.cos(half_yaw), math.sin(half_yaw)
+  return _normalize_quaternion(
+    (
+      cos_roll * cos_pitch * cos_yaw + sin_roll * sin_pitch * sin_yaw,
+      sin_roll * cos_pitch * cos_yaw - cos_roll * sin_pitch * sin_yaw,
+      cos_roll * sin_pitch * cos_yaw + sin_roll * cos_pitch * sin_yaw,
+      cos_roll * cos_pitch * sin_yaw - sin_roll * sin_pitch * cos_yaw,
+    )
+  )
+
+
+def _quaternion_to_euler(quaternion: _Quaternion) -> tuple[float, float, float]:
+  """Convert a quaternion to roll, pitch, and yaw in degrees."""
+  w, x, y, z = quaternion
+  sin_pitch = max(-1.0, min(1.0, 2 * (w * y - z * x)))
+  roll = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
+  pitch = math.asin(sin_pitch)
+  yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))
+  return math.degrees(roll), math.degrees(pitch), math.degrees(yaw)
+
+
+def _nearest_equivalent_angle(angle: float, reference: float) -> float:
+  """Shift an angle by full turns so it stays close to a reference value."""
+  equivalent = angle + 360 * round((reference - angle) / 360)
+  if math.isclose(equivalent, reference, rel_tol=0.0, abs_tol=1e-12):
+    return reference
+  return equivalent
 
 
 class Rotation(SerializableMixin):
   """Represents a 3D rotation."""
 
   def __init__(self, x: float = 0, y: float = 0, z: float = 0):
-    self.x = x  # around x-axis, roll
-    self.y = y  # around y-axis, pitch
-    self.z = z  # around z-axis, yaw
-
-  def get_rotation_matrix(self):
-    # Create rotation matrices for each axis
-    Rz = [
-      [
-        math.cos(math.radians(self.z)),
-        -math.sin(math.radians(self.z)),
-        0,
-      ],
-      [
-        math.sin(math.radians(self.z)),
-        math.cos(math.radians(self.z)),
-        0,
-      ],
-      [0, 0, 1],
-    ]
-    Ry = [
-      [
-        math.cos(math.radians(self.y)),
-        0,
-        math.sin(math.radians(self.y)),
-      ],
-      [0, 1, 0],
-      [
-        -math.sin(math.radians(self.y)),
-        0,
-        math.cos(math.radians(self.y)),
-      ],
+    self._x = x  # around x-axis, roll
+    self._y = y  # around y-axis, pitch
+    self._z = z  # around z-axis, yaw
+    self._quaternion = _quaternion_from_euler(x=x, y=y, z=z)
+
+  @classmethod
+  def _from_quaternion(
+    cls, quaternion: _Quaternion, reference: tuple[float, float, float]
+  ) -> "Rotation":
+    """Create a rotation using the equivalent Euler angles closest to a reference."""
+    quaternion = _normalize_quaternion(quaternion)
+    w, x, y, z = quaternion
+    sin_pitch = 2 * (w * y - z * x)
+    cos_pitch = math.hypot(1 - 2 * (y * y + z * z), 2 * (x * y + w * z))
+    if cos_pitch < 1e-12:
+      coupled_angle = math.degrees(2 * math.atan2(x, w))
+      if sin_pitch > 0:
+        difference = _nearest_equivalent_angle(coupled_angle, reference[0] - reference[2])
+        euler = (
+          (reference[0] + reference[2] + difference) / 2,
+          _nearest_equivalent_angle(90, reference[1]),
+          (reference[0] + reference[2] - difference) / 2,
+        )
+      else:
+        total = _nearest_equivalent_angle(coupled_angle, reference[0] + reference[2])
+        euler = (
+          (reference[0] - reference[2] + total) / 2,
+          _nearest_equivalent_angle(-90, reference[1]),
+          (reference[2] - reference[0] + total) / 2,
+        )
+      rotation = cls(x=euler[0], y=euler[1], z=euler[2])
+      rotation._quaternion = quaternion
+      return rotation
+
+    principal = _quaternion_to_euler(quaternion)
+    alternate = (principal[0] + 180, 180 - principal[1], principal[2] + 180)
+    candidates = [
+      tuple(_nearest_equivalent_angle(angle, target) for angle, target in zip(candidate, reference))
+      for candidate in (principal, alternate)
     ]
-    Rx = [
-      [1, 0, 0],
-      [
-        0,
-        math.cos(math.radians(self.x)),
-        -math.sin(math.radians(self.x)),
-      ],
-      [
-        0,
-        math.sin(math.radians(self.x)),
-        math.cos(math.radians(self.x)),
-      ],
+    x, y, z = min(
+      candidates,
+      key=lambda candidate: sum(
+        (angle - target) ** 2 for angle, target in zip(candidate, reference)
+      ),
+    )
+    rotation = cls(x=x, y=y, z=z)
+    rotation._quaternion = quaternion
+    return rotation
+
+  def get_rotation_matrix(self) -> list[list[float]]:
+    """Return the rotation as a 3x3 matrix."""
+    w, x, y, z = self._quaternion
+    return [
+      [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+      [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+      [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
     ]
-    # Combine rotations: The order of multiplication matters and defines the behavior significantly.
-    # This is a common order: Rz * Ry * Rx
-    return matrix_multiply_3x3(matrix_multiply_3x3(Rz, Ry), Rx)
+
+  def serialize(self) -> dict:
+    """Serialize using the public Euler-angle representation."""
+    return {"x": self.x, "y": self.y, "z": self.z, "type": self.__class__.__name__}
 
   def __str__(self) -> str:
     return f"Rotation(x={self.x}, y={self.y}, z={self.z})"
 
-  def __add__(self, other) -> "Rotation":
-    return Rotation(x=self.x + other.x, y=self.y + other.y, z=self.z + other.z)
+  def __add__(self, other: "Rotation") -> "Rotation":
+    """Compose rotations so the resulting matrix is ``self * other``."""
+    reference = (self.x + other.x, self.y + other.y, self.z + other.z)
+    return self._from_quaternion(
+      _multiply_quaternions(self._quaternion, other._quaternion), reference=reference
+    )
+
+  def _prepend(self, other: "Rotation") -> None:
+    """Apply another rotation in the current reference frame while keeping this instance."""
+    combined = other + self
+    self._x = combined.x % 360
+    self._y = combined.y % 360
+    self._z = combined.z % 360
+    self._quaternion = combined._quaternion
 
   def __repr__(self) -> str:
     return self.__str__()
 
+  @property
+  def x(self) -> float:
+    """Rotation around the x-axis in degrees (roll)."""
+    return self._x
+
+  @x.setter
+  def x(self, value: float) -> None:
+    self._x = value
+    self._quaternion = _quaternion_from_euler(x=self._x, y=self._y, z=self._z)
+
+  @property
+  def y(self) -> float:
+    """Rotation around the y-axis in degrees (pitch)."""
+    return self._y
+
+  @y.setter
+  def y(self, value: float) -> None:
+    self._y = value
+    self._quaternion = _quaternion_from_euler(x=self._x, y=self._y, z=self._z)
+
+  @property
+  def z(self) -> float:
+    """Rotation around the z-axis in degrees (yaw)."""
+    return self._z
+
+  @z.setter
+  def z(self, value: float) -> None:
+    self._z = value
+    self._quaternion = _quaternion_from_euler(x=self._x, y=self._y, z=self._z)
+
   @property
   def roll(self) -> float:
     return self.x