Add i09-2 stage using StandardMovable

src/dodal/beamlines/i09_2.py

@@ -1,12 +1,14 @@
 from dodal.beamlines.i09_2_shared import devices as i09_2_shared_devices
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
 from dodal.device_manager import DeviceManager
+from dodal.devices.beamlines.i09_2.i09_2_motors import I092SampleManipulator
 from dodal.devices.synchrotron import Synchrotron
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
 BL = get_beamline_name("i09-2")
-PREFIX = BeamlinePrefix(BL, suffix="J")
+J_PREFIX = BeamlinePrefix(BL, suffix="J")
+K_PREFIX = BeamlinePrefix(BL, suffix="K")
 set_log_beamline(BL)
 set_utils_beamline(BL)
 
@@ -17,3 +19,8 @@
 @devices.factory()
 def synchrotron() -> Synchrotron:
     return Synchrotron()
+
+
+@devices.factory()
+def sm() -> I092SampleManipulator:
+    return I092SampleManipulator(f"{K_PREFIX.beamline_prefix}-MO-SM-01:PI:")

src/dodal/devices/beamlines/i09_2/__init__.py

@@ -0,0 +1,7 @@
+from .i09_2_motors import (
+    I092SampleManipulator,
+    PiezoElectricMotor,
+    PiezoElectricMovableLogic,
+)
+
+__all__ = ["I092SampleManipulator", "PiezoElectricMotor", "PiezoElectricMovableLogic"]

src/dodal/devices/beamlines/i09_2/i09_2_motors.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+from functools import cached_property
+
+from ophyd_async.core import (
+    SignalW,
+    StandardReadable,
+    StandardReadableFormat,
+    soft_signal_rw,
+)
+from ophyd_async.epics.core import epics_signal_r, epics_signal_rw, epics_signal_w
+from ophyd_async.epics.motor import Motor
+
+from dodal.devices.movable import MovableWithTolerance, MovableWithToleranceLogic
+
+
+@dataclass
+class PiezoElectricMovableLogic(MovableWithToleranceLogic):
+    motor_stop: SignalW[int]
+
+    async def stop(self) -> None:
+        await self.motor_stop.set(1)
+
+
+class PiezoElectricMotor(MovableWithTolerance):
+    """A piezoelectric positioning stage with configurable move tolerance.
+
+    This device exposes EPICS signals for readback, setpoint, and motion stop commands.
+    Motion completion is determined by comparing the readback and setpoint positions
+    using a configurable tolerance.
+    """
+
+    def __init__(self, prefix: str, tolerance: float = 0.01, name: str = ""):
+        with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
+            self.user_readback = epics_signal_r(float, prefix + ":POS:RD")
+
+        self.user_setpoint = epics_signal_rw(float, prefix + ":MOV:RD")
+        self.tolerance = soft_signal_rw(float, initial_value=tolerance)
+        self.motor_stop = epics_signal_w(int, prefix + ":HLT:WR.PROC")
+        super().__init__(
+            tolerance=self.tolerance,
+            setpoint=self.user_setpoint,
+            readback=self.user_readback,
+            name=name,
+        )
+
+    @cached_property
+    def movable_logic(self) -> PiezoElectricMovableLogic:
+        return PiezoElectricMovableLogic(
+            readback=self.user_readback,
+            setpoint=self.user_setpoint,
+            within_tolerance=self.within_tolerance,
+            motor_stop=self.motor_stop,
+        )
+
+
+class I092SampleManipulator(StandardReadable):
+    def __init__(self, prefix: str, name: str = ""):
+        with self.add_children_as_readables():
+            self.x1 = PiezoElectricMotor(prefix + "X1")
+            self.x2 = PiezoElectricMotor(prefix + "X2")
+            self.x3 = PiezoElectricMotor(prefix + "X3")
+            self.y = PiezoElectricMotor(prefix + "Y")
+            self.z1 = PiezoElectricMotor(prefix + "Z1")
+            self.z2 = PiezoElectricMotor(prefix + "Z2")
+
+            self.xc = Motor(prefix + "X")
+            self.zc = Motor(prefix + "Z")
+
+        super().__init__(name)

src/dodal/devices/beamlines/i10_1/high_field_magnet/high_field_magnet.py

@@ -1,32 +1,23 @@
-from __future__ import annotations
-
 from dataclasses import dataclass
 from functools import cached_property
 
-from bluesky.protocols import (
-    Flyable,
-    Preparable,
-)
+from bluesky.protocols import Flyable, Preparable
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     AsyncStatus,
-    MovableLogic,
-    SignalR,
     SignalRW,
-    StandardMovable,
-    StandardReadable,
     StandardReadableFormat,
     StrictEnum,
     SubsetEnum,
     WatchableAsyncStatus,
-    derived_signal_r,
     error_if_none,
-    set_and_wait_for_other_value,
     soft_signal_rw,
 )
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw, epics_signal_w
 from pydantic import BaseModel, Field
 
+from dodal.devices.movable import MovableWithTolerance, MovableWithToleranceLogic
+
 
 class HighFieldMangetSweepTypes(StrictEnum):
     FAST = "Fast"
@@ -56,11 +47,9 @@ class FlyMagInfo(BaseModel):
 
 
 @dataclass
-class ToleranceMovableLogic(MovableLogic[float]):
-    tolerance: SignalRW[float]
+class HighFieldMagnetMovableLogic(MovableWithToleranceLogic):
     speed: SignalRW[float]
     acc_time: SignalRW[float]
-    within_tolerance: SignalR[bool]
 
     async def stop(self):
         current_val = await self.readback.get_value()
@@ -81,22 +70,8 @@ async def calculate_timeout(
             raise ValueError(msg) from error
         return timeout
 
-    async def move(self, new_position: float, timeout: float | None) -> None:
-        await set_and_wait_for_other_value(
-            set_signal=self.setpoint,
-            set_value=new_position,
-            match_signal=self.within_tolerance,
-            match_value=True,
-            timeout=timeout,
-        )
-
 
-class HighFieldMagnet(
-    StandardMovable,
-    StandardReadable,
-    Flyable,
-    Preparable,
-):
+class HighFieldMagnet(MovableWithTolerance, Flyable, Preparable):
     def __init__(
         self, prefix: str, field_tolerance: float = 0.01, name: str = ""
     ) -> None:
@@ -117,6 +92,7 @@ def __init__(
             )
             self.ramp_up_time = soft_signal_rw(datatype=float, initial_value=1.0)
             self.field_tolerance = soft_signal_rw(float, initial_value=field_tolerance)
+
         with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
             self.user_readback = epics_signal_r(float, prefix + "RBV:DEMANDFIELD")
 
@@ -129,31 +105,22 @@ def __init__(
             read_pv=prefix + "RBV:SETPOINTFIELD",
             write_pv=prefix + "SET:SETPOINTFIELD",
         )
-
-        self.within_tolerance = derived_signal_r(
-            raw_to_derived=self._within_tolerance,
-            setpoint=self.user_setpoint,
-            readback=self.user_readback,
-            tolerance=self.field_tolerance,
-        )
         self._set_success = True
         self._fly_info: FlyMagInfo | None = None
         self._fly_status: WatchableAsyncStatus | None = None
 
-        super().__init__(name=name)
-
-    def _within_tolerance(
-        self, setpoint: float, readback: float, tolerance: float
-    ) -> bool:
-        """Check if the readback is within the tolerance of the setpoint."""
-        return abs(setpoint - readback) < abs(tolerance)
+        super().__init__(
+            tolerance=self.field_tolerance,
+            setpoint=self.user_setpoint,
+            readback=self.user_readback,
+            name=name,
+        )
 
     @cached_property
-    def movable_logic(self) -> MovableLogic:
-        return ToleranceMovableLogic(
+    def movable_logic(self) -> HighFieldMagnetMovableLogic:
+        return HighFieldMagnetMovableLogic(
             setpoint=self.user_setpoint,
             readback=self.user_readback,
-            tolerance=self.field_tolerance,
             within_tolerance=self.within_tolerance,
             speed=self.sweep_rate,
             acc_time=self.ramp_up_time,

src/dodal/devices/movable.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+from functools import cached_property
+
+from ophyd_async.core import (
+    MovableLogic,
+    Reference,
+    SignalR,
+    SignalRW,
+    StandardMovable,
+    StandardReadable,
+    derived_signal_r,
+    wait_for_value,
+)
+
+
+@dataclass
+class MovableWithToleranceLogic(MovableLogic[float]):
+    within_tolerance: SignalR[bool]
+
+    async def move(self, new_position: float, timeout: float | None) -> None:
+        await self.setpoint.set(new_position, timeout=timeout)
+        await wait_for_value(self.setpoint, new_position, timeout=timeout)
+        # Once setpoint is at new position, we can check for tolerance signal to see if
+        # true now as the within_tolerance window has updated to the new setpoint
+        # position. Now it doesn't matter if motor steps are small or large.
+        await wait_for_value(self.within_tolerance, True, timeout=timeout)
+
+
+def _within_tolerance_read(setpoint: float, readback: float, tolerance: float) -> bool:
+    return abs(setpoint - readback) < abs(tolerance)
+
+
+class MovableWithTolerance(StandardMovable[float], StandardReadable):
+    """Movable with a signal to configure the tolerance of when the device is done
+    moving if it the readback and setpoint difference is within the tolerance.
+    """
+
+    def __init__(
+        self,
+        tolerance: SignalR[float],
+        setpoint: SignalRW[float],
+        readback: SignalR[float],
+        name: str = "",
+    ):
+        # Use reference so sub classes still have flexibility to name signals to what
+        # they want.
+        self._setpoint_ref = Reference(setpoint)
+        self._readback_ref = Reference(readback)
+        self.within_tolerance = derived_signal_r(
+            _within_tolerance_read,
+            tolerance=tolerance,
+            setpoint=setpoint,
+            readback=readback,
+        )
+        super().__init__(name)
+
+    @cached_property
+    def movable_logic(self) -> MovableWithToleranceLogic:
+        return MovableWithToleranceLogic(
+            readback=self._readback_ref(),
+            setpoint=self._setpoint_ref(),
+            within_tolerance=self.within_tolerance,
+        )

tests/devices/beamlines/i09_2/test_i09_2_motors.py

@@ -0,0 +1,49 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from ophyd_async.core import init_devices
+from ophyd_async.testing import assert_reading, partial_reading
+
+from dodal.devices.beamlines.i09_2 import I092SampleManipulator, PiezoElectricMotor
+
+
+@pytest.fixture
+def sm() -> I092SampleManipulator:
+    with init_devices(mock=True):
+        sm = I092SampleManipulator("TEST:")
+    return sm
+
+
+async def test_sm_read(sm: I092SampleManipulator) -> None:
+    await assert_reading(
+        sm,
+        {
+            "sm-x1": partial_reading(0),
+            "sm-x2": partial_reading(0),
+            "sm-x3": partial_reading(0),
+            "sm-y": partial_reading(0),
+            "sm-z1": partial_reading(0),
+            "sm-z2": partial_reading(0),
+            "sm-xc": partial_reading(0),
+            "sm-zc": partial_reading(0),
+        },
+    )
+
+
+@pytest.fixture
+async def piezo_motor() -> PiezoElectricMotor:
+    with init_devices(mock=True):
+        piezo_motor = PiezoElectricMotor("TEST:")
+    return piezo_motor
+
+
+async def test_piezo_motor_stop(piezo_motor: PiezoElectricMotor) -> None:
+    piezo_motor.motor_stop.set = AsyncMock()
+    await piezo_motor.stop()
+    piezo_motor.motor_stop.set.assert_awaited_once_with(1)
+
+
+async def test_piezo_motor_set(piezo_motor: PiezoElectricMotor) -> None:
+    await piezo_motor.set(4)
+    assert await piezo_motor.user_readback.get_value() == 4
+    assert await piezo_motor.user_setpoint.get_value() == 4

tests/devices/beamlines/i10_1/high_field_magnet/test_high_field_magnet.py

@@ -147,32 +147,9 @@ async def test_read(high_field_magnet: HighFieldMagnet):
     )
 
 
-@pytest.mark.parametrize(
-    "setpoint,readback,tolerance,expected",
-    [
-        (10.0, 10.005, 0.01, True),  # test positive tolerance
-        (10.0, 9.995, 0.01, True),
-        (10.0, 10.02, 0.01, False),
-        (10.0, 0.9, 0.01, False),
-        (10.0, 9.995, -0.01, True),  # test negative tolerance
-        (10.0, 10.0005, -0.01, True),
-        (10.0, 9.98, -0.01, False),
-        (10.0, 10.1, -0.01, False),
-    ],
-)
-async def test_tolerance_logic_within_tolerance(
-    high_field_magnet: HighFieldMagnet, setpoint, readback, tolerance, expected
-):
-    result = high_field_magnet._within_tolerance(
-        setpoint=setpoint, readback=readback, tolerance=tolerance
-    )
-    assert result is expected
-
-
 async def test_tolerance_logic_stop_clears_set_success_and_restores_setpoint(
     high_field_magnet: HighFieldMagnet,
 ):
-    set_mock_value(high_field_magnet.movable_logic.readback, 7.5)
     set_mock_value(high_field_magnet.movable_logic.readback, 1.5)
     await high_field_magnet.stop()
     assert high_field_magnet._set_success is False
@@ -189,18 +166,6 @@ async def test_tolerance_logic_calculate_timeout_with_zero_speed(
         )
 
 
-async def test_tolerance_logic_move(high_field_magnet: HighFieldMagnet):
-    set_mock_value(high_field_magnet.movable_logic.readback, 0.0)
-    move_task = high_field_magnet.movable_logic.move(new_position=10.0, timeout=5.0)
-    for value in [2.0, 5.0, 8.0, 9.5, 13.0]:
-        set_mock_value(high_field_magnet.movable_logic.readback, value)
-        await asyncio.sleep(0.0)
-        assert await high_field_magnet.within_tolerance.get_value() is False
-    set_mock_value(high_field_magnet.movable_logic.readback, 9.91)
-    await move_task
-    assert await high_field_magnet.within_tolerance.get_value() is True
-
-
 def test_run_engine_scan(
     run_engine: RunEngine,
     high_field_magnet: HighFieldMagnet,