Add generic calculation for translating from sample space to lab space and use for i23 OAV

src/dodal/devices/aithre_lasershaping/goniometer.py

@@ -1,10 +1,6 @@
-import asyncio
-import math
-
-from ophyd_async.core import derived_signal_rw
 from ophyd_async.epics.motor import Motor
 
-from dodal.devices.motors import XYZStage
+from dodal.devices.motors import XYZStage, create_axis_perp_to_rotation
 
 
 class Goniometer(XYZStage):
@@ -23,26 +19,7 @@ def __init__(self, prefix: str, name: str = "") -> None:
         self.sampy = Motor(prefix + "SAMPY")
         self.sampz = Motor(prefix + "SAMPZ")
         self.omega = Motor(prefix + "OMEGA")
-        self.vertical_position = derived_signal_rw(
-            self._get,
-            self._set,
-            sampy=self.sampy,
-            sampz=self.sampz,
-            omega=self.omega,
+        self.vertical_position = create_axis_perp_to_rotation(
+            self.omega, self.sampz, self.sampy
         )
         super().__init__(name)
-
-    def _get(self, sampz: float, sampy: float, omega: float) -> float:
-        z_component = sampz * math.cos(math.radians(omega))
-        y_component = sampy * math.sin(math.radians(omega))
-        return z_component + y_component
-
-    async def _set(self, value: float) -> None:
-        omega = await self.omega.user_readback.get_value()
-        z_component = value * math.cos(math.radians(omega))
-        y_component = value * math.sin(math.radians(omega))
-        await asyncio.gather(
-            self.sampy.set(y_component),
-            self.sampz.set(z_component),
-            self.omega.set(omega),
-        )

src/dodal/devices/motors.py

@@ -1,6 +1,8 @@
+import asyncio
+import math
 from abc import ABC
 
-from ophyd_async.core import StandardReadable
+from ophyd_async.core import StandardReadable, derived_signal_rw
 from ophyd_async.epics.motor import Motor
 
 _X, _Y, _Z = "X", "Y", "Z"
@@ -101,12 +103,19 @@ def __init__(
         phi_infix: str = "PHI",
         omega_infix: str = "OMEGA",
     ):
+        """Six-axis goniometer with a standard xyz stage and three axes of rotation:
+        kappa, phi and omega.
+        """
         with self.add_children_as_readables():
             self.kappa = Motor(prefix + kappa_infix)
             self.phi = Motor(prefix + phi_infix)
             self.omega = Motor(prefix + omega_infix)
         super().__init__(prefix, name, x_infix, y_infix, z_infix)
 
+        self.vertical_in_lab_space = create_axis_perp_to_rotation(
+            self.omega, self.y, self.z
+        )
+
 
 class YZStage(Stage):
     def __init__(
@@ -116,3 +125,49 @@ def __init__(
             self.y = Motor(prefix + y_infix)
             self.z = Motor(prefix + z_infix)
         super().__init__(name)
+
+
+def create_axis_perp_to_rotation(motor_theta: Motor, motor_i: Motor, motor_j: Motor):
+    """Given a signal that controls a motor in a rotation axis and two other
+    signals controlling motors on a pair of orthogonal axes, these axes being in the
+    rotating frame of reference created by the first axis, create a derived signal
+    that is a projection of the two axes in the non-rotating frame of reference.
+
+    The projection is onto the axis defined by i when the rotation angle is 0 and
+    defined by j when the angle is at 90.
+
+    The usual use case for this is translating from sample space to lab space. For
+    example, if you have a sample that is mounted on a goniometer to the right hand side
+    of an OAV view this can provide an axis that will move the sample up/down in that
+    view regardless of the omega orientation of the sample.
+
+    Args:
+        motor_theta (Motor): this is the rotation axis of the sample.
+        motor_i (Motor): this is the axis that, when the sample is at 0 deg rotation,
+                         a move here is entirely parallel with the derived axis.
+        motor_j (Motor): this is the axis that, when the sample is at 90 deg rotation,
+                         a move here is entirely parallel with the derived axis.
+    """
+
+    def _get(j_val: float, i_val: float, rot_value: float) -> float:
+        i_component = i_val * math.cos(math.radians(rot_value))
+        j_component = j_val * math.sin(math.radians(rot_value))
+        return i_component + j_component
+
+    async def _set(vertical_value: float) -> None:
+        rot_value = await motor_theta.user_readback.get_value()
+        i_component = vertical_value * math.cos(math.radians(rot_value))
+        j_component = vertical_value * math.sin(math.radians(rot_value))
+        await asyncio.gather(
+            motor_i.set(i_component),
+            motor_j.set(j_component),
+            motor_theta.set(rot_value),
+        )
+
+    return derived_signal_rw(
+        _get,
+        _set,
+        i_val=motor_i,
+        j_val=motor_j,
+        rot_value=motor_theta,
+    )

tests/devices/unit_tests/test_motors.py

@@ -1,15 +1,22 @@
+import math
+
 import pytest
 from bluesky.run_engine import RunEngine
 from ophyd_async.core import init_devices
 from ophyd_async.testing import assert_reading, partial_reading
 
 from dodal.devices.motors import SixAxisGonio
+from dodal.devices.util.test_utils import patch_motor
 
 
 @pytest.fixture
-def six_axis_gonio(RE: RunEngine):
+def six_axis_gonio(RE: RunEngine) -> SixAxisGonio:
     with init_devices(mock=True):
         gonio = SixAxisGonio("")
+    patch_motor(gonio.omega)
+    patch_motor(gonio.z)
+    patch_motor(gonio.y)
+    patch_motor(gonio.x)
 
     return gonio
 
@@ -26,3 +33,52 @@ async def test_reading_six_axis_gonio(six_axis_gonio: SixAxisGonio):
             "gonio-x": partial_reading(0.0),
         },
     )
+
+
+@pytest.mark.parametrize(
+    "set_value, omega_set_value, expected_horz, expected_vert",
+    [
+        [2, 60, math.sqrt(3), 1],
+        [-10, 390, -5, -5 * math.sqrt(3)],
+        [0.5, -135, -math.sqrt(2) / 4, -math.sqrt(2) / 4],
+        [1, 0, 0, 1],
+    ],
+)
+async def test_vertical_in_lab_space_for_default_axes(
+    six_axis_gonio: SixAxisGonio,
+    set_value: float,
+    omega_set_value: float,
+    expected_horz: float,
+    expected_vert: float,
+):
+    await six_axis_gonio.omega.set(omega_set_value)
+    await six_axis_gonio.vertical_in_lab_space.set(set_value)
+
+    assert await six_axis_gonio.z.user_readback.get_value() == pytest.approx(
+        expected_horz
+    )
+    assert await six_axis_gonio.y.user_readback.get_value() == pytest.approx(
+        expected_vert
+    )
+
+    await six_axis_gonio.vertical_in_lab_space.set(set_value * 2)
+    assert await six_axis_gonio.z.user_readback.get_value() == pytest.approx(
+        expected_horz * 2
+    )
+    assert await six_axis_gonio.y.user_readback.get_value() == pytest.approx(
+        expected_vert * 2
+    )
+
+
+@pytest.mark.parametrize(
+    "set_point",
+    [
+        -5,
+        0,
+        100,
+        0.7654,
+    ],
+)
+async def test_lab_vertical_round_trip(six_axis_gonio: SixAxisGonio, set_point: float):
+    await six_axis_gonio.vertical_in_lab_space.set(set_point)
+    assert await six_axis_gonio.vertical_in_lab_space.get_value() == set_point