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

src/dodal/devices/aithre_lasershaping/goniometer.py

@@ -1,9 +1,8 @@
-import asyncio
-import math
-
-from ophyd_async.core import StandardReadable, derived_signal_rw
+from ophyd_async.core import StandardReadable
 from ophyd_async.epics.motor import Motor
 
+from dodal.devices.motors import create_axis_perp_to_rotation
+
 
 class Goniometer(StandardReadable):
     """The Aithre lab goniometer and the XYZ stage it sits on.
@@ -24,26 +23,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,38 +1,72 @@
-from ophyd_async.core import StandardReadable
+import asyncio
+import math
+
+from ophyd_async.core import StandardReadable, derived_signal_rw
 from ophyd_async.epics.motor import Motor
 
 
-class XYZPositioner(StandardReadable):
-    """
+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.
 
-    Standard ophyd_async xyz motor stage, by combining 3 Motors,
-    with added infix for extra flexibliy to allow different axes other than x,y,z.
-
-    Parameters
-    ----------
-    prefix:
-        EPICS PV (Common part up to and including :).
-    name:
-        name for the stage.
-    infix:
-        EPICS PV, default is the ("X", "Y", "Z").
-    Notes
-    -----
-    Example usage::
-        async with init_devices():
-            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:")
-    Or::
-        with init_devices():
-            xyz_stage = XYZPositioner("BLXX-MO-STAGE-XX:", infix = ("A", "B", "C"))
+    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,
+    )
+
+
+class XYZPositioner(StandardReadable):
     def __init__(
         self,
         prefix: str,
         name: str = "",
         infix: tuple[str, str, str] = ("X", "Y", "Z"),
     ):
+        """Standard ophyd_async xyz motor stage, by combining 3 Motors,
+        with added infix for extra flexibility to allow different axes other than x,y,z.
+
+        Args:
+            prefix: EPICS PV (Common part up to and including :).
+            name: name for the device.
+            infix: EPICS PV suffix, default is the ("X", "Y", "Z").
+
+        """
         with self.add_children_as_readables():
             self.x = Motor(prefix + infix[0])
             self.y = Motor(prefix + infix[1])
@@ -54,8 +88,21 @@ def __init__(
             "OMEGA",
         ),
     ):
+        """Six-axis goniometer with a standard xyz stage and three axes of rotation:
+        kappa, phi and omega.
+
+        Args:
+            prefix: EPICS PV (Common part up to and including :).
+            name: name for the device.
+            infix: EPICS PV suffix, default is the ("X", "Y", "Z", "KAPPA", "PHI", "OMEGA").
+        """
         with self.add_children_as_readables():
             self.kappa = Motor(prefix + infix[3])
             self.phi = Motor(prefix + infix[4])
             self.omega = Motor(prefix + infix[5])
+
         super().__init__(name=name, prefix=prefix, infix=infix[0:3])
+
+        self.vertical_in_lab_space = create_axis_perp_to_rotation(
+            self.omega, self.y, self.z
+        )

tests/devices/unit_tests/test_motors.py

@@ -1,3 +1,4 @@
+import math
 from unittest.mock import ANY
 
 import pytest
@@ -6,12 +7,17 @@
 from ophyd_async.testing import assert_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
 
@@ -50,5 +56,59 @@ async def test_reading_six_axis_gonio(six_axis_gonio: SixAxisGonio):
                 "timestamp": ANY,
                 "alarm_severity": 0,
             },
+            "gonio-_horizontal_stage_axis": {
+                "alarm_severity": 0,
+                "timestamp": ANY,
+                "value": 0.0,
+            },
+            "gonio-_vertical_stage_axis": {
+                "alarm_severity": 0,
+                "timestamp": ANY,
+                "value": 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