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.sampz, self.sampy, self.omega
         )
         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,7 +1,57 @@
-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
 
 
+def create_axis_perp_to_rotation(
+    parallel_to_0: Motor, parallel_to_minus_90: Motor, rotation: Motor
+):
+    """Create a new derived signal that moves perpendicular to a rotation axis.
+
+    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:
+        parallel_to_0 (Motor): this is the axis that, when the sample is at 0 deg rotation,
+                               a move here is entirely parallel with the derived axis.
+        parallel_to_minus_90 (Motor): this is the axis that, when the sample is at 90 deg
+                                      rotation, a move here is entirely parallel with the
+                                      derived axis.
+        rotation (Motor): this is the rotation axis of the sample.
+    """
+    # By convention use y/z internally as that is what is used on most beamlines but the
+    # function is actually indifferent to this
+    y_mot = parallel_to_0
+    z_mot = parallel_to_minus_90
+
+    def _get(z_val: float, y_val: float, rot_value: float) -> float:
+        y_component = y_val * math.cos(math.radians(rot_value))
+        z_component = z_val * math.sin(math.radians(rot_value))
+        return z_component + y_component
+
+    async def _set(vertical_value: float) -> None:
+        rot_value = await rotation.user_readback.get_value()
+        y_component = vertical_value * math.cos(math.radians(rot_value))
+        z_component = vertical_value * math.sin(math.radians(rot_value))
+        await asyncio.gather(
+            y_mot.set(y_component),
+            z_mot.set(z_component),
+            rotation.set(rot_value),
+        )
+
+    return derived_signal_rw(
+        _get,
+        _set,
+        y_val=y_mot,
+        z_val=z_mot,
+        rot_value=rotation,
+    )
+
+
 class XYZPositioner(StandardReadable):
     """
 
@@ -41,6 +91,31 @@ def __init__(
 
 
 class SixAxisGonio(XYZPositioner):
+    """
+
+    Six-axis goniometer with a standard xyz stage and three axes of rotation: kappa, phi
+    and omega.
+
+        Parameters
+    ----------
+    prefix:
+        EPICS PV (Common part up to and including :).
+    name:
+        name for the stage.
+    infix:
+        EPICS PV, default is the ("X", "Y", "Z", "KAPPA", "PHI", "OMEGA").
+    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", \
+            "KAPPA", "PHI", "OMEGA"))
+
+    """
+
     def __init__(
         self,
         prefix: str,
@@ -58,4 +133,9 @@ def __init__(
             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.y, self.z, self.omega
+        )

tests/devices/unit_tests/test_motors.py

@@ -1,3 +1,4 @@
+import math
 from unittest.mock import ANY
 
 import pytest
@@ -6,12 +7,29 @@
 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
+
+
+@pytest.fixture
+def y_up_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 +68,90 @@ 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_value, omega_set_value, expected_horz, expected_vert",
+    [
+        [2, 60, 1, math.sqrt(3)],
+        [-10, 390, -5 * math.sqrt(3), -5],
+        [0.5, -135, -math.sqrt(2) / 4, -math.sqrt(2) / 4],
+        [1, 0, 1, 0],
+    ],
+)
+async def test_j_signal_for_rotated_axes(
+    y_up_six_axis_gonio: SixAxisGonio,
+    set_value: float,
+    omega_set_value: float,
+    expected_horz: float,
+    expected_vert: float,
+):
+    await y_up_six_axis_gonio.omega.set(omega_set_value)
+    await y_up_six_axis_gonio.vertical_in_lab_space.set(set_value)
+
+    assert await y_up_six_axis_gonio.y.user_readback.get_value() == pytest.approx(
+        expected_horz
+    )
+    assert await y_up_six_axis_gonio.z.user_readback.get_value() == pytest.approx(
+        expected_vert
+    )
+
+    await y_up_six_axis_gonio.vertical_in_lab_space.set(set_value * 2)
+    assert await y_up_six_axis_gonio.y.user_readback.get_value() == pytest.approx(
+        expected_horz * 2
+    )
+    assert await y_up_six_axis_gonio.z.user_readback.get_value() == pytest.approx(
+        expected_vert * 2
+    )
+
+
+async def test_get_j(y_up_six_axis_gonio: SixAxisGonio):
+    await y_up_six_axis_gonio.vertical_in_lab_space.set(5)
+    assert await y_up_six_axis_gonio.vertical_in_lab_space.get_value() == 5