Vertical position derived signal for I23 Aithre goniometer

src/dodal/devices/aithre_lasershaping/goniometer.py

@@ -1,9 +1,21 @@
-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 Goniometer(StandardReadable):
-    """Goniometer and the stages it sits on"""
+    """The Aithre lab goniometer and the XYZ stage it sits on.
+
+    `x`, `y` and `z` control the axes of the positioner at the base, while `sampy` and
+    `sampz` control the positioner of the sample. `omega` is the rotation about the
+    x-axis (along the length of the sample holder).
+
+    The `vertical_position` signal refers to the height of the sample from the point of
+    view of the OAV and setting this value moves the sample vertically in the OAV plane
+    regardless of the current rotation.
+    """
 
     def __init__(self, prefix: str, name: str = "") -> None:
         self.x = Motor(prefix + "X")
@@ -12,4 +24,26 @@ 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,
+        )
         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),
+        )

tests/devices/aithre_lasershaping/test_goniometer.py

@@ -0,0 +1,61 @@
+import math
+
+import pytest
+from bluesky.run_engine import RunEngine
+from ophyd_async.core import init_devices
+
+from dodal.beamlines import aithre
+from dodal.devices.aithre_lasershaping.goniometer import Goniometer
+from dodal.devices.util.test_utils import patch_motor
+
+
+@pytest.fixture
+def goniometer(RE: RunEngine) -> Goniometer:
+    with init_devices(mock=True):
+        gonio = aithre.goniometer(connect_immediately=True, mock=True)
+    patch_motor(gonio.omega)
+    patch_motor(gonio.sampy)
+    patch_motor(gonio.sampz)
+    return gonio
+
+
+@pytest.mark.parametrize(
+    "vertical_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],
+        [-1.5, 90, 0, -1.5],
+    ],
+)
+async def test_vertical_signal_set(
+    goniometer: Goniometer,
+    vertical_set_value: float,
+    omega_set_value: float,
+    expected_horz: float,
+    expected_vert: float,
+):
+    await goniometer.omega.set(omega_set_value)
+    await goniometer.vertical_position.set(vertical_set_value)
+
+    assert await goniometer.sampz.user_readback.get_value() == pytest.approx(
+        expected_horz
+    )
+    assert await goniometer.sampy.user_readback.get_value() == pytest.approx(
+        expected_vert
+    )
+
+    await goniometer.vertical_position.set(vertical_set_value * 2)
+    assert await goniometer.sampz.user_readback.get_value() == pytest.approx(
+        expected_horz * 2
+    )
+    assert await goniometer.sampy.user_readback.get_value() == pytest.approx(
+        expected_vert * 2
+    )
+
+
+@pytest.mark.parametrize("set_value", [-5, 2.7, 0])
+async def test_vertical_position_get(goniometer: Goniometer, set_value: float):
+    await goniometer.vertical_position.set(set_value)
+    assert await goniometer.vertical_position.get_value() == set_value