diff --git a/src/dodal/beamlines/i09_2.py b/src/dodal/beamlines/i09_2.py index 99b29b7880b..0b2d4272c54 100644 --- a/src/dodal/beamlines/i09_2.py +++ b/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:") diff --git a/src/dodal/devices/beamlines/i09_2/__init__.py b/src/dodal/devices/beamlines/i09_2/__init__.py new file mode 100644 index 00000000000..5bf0dafaa9e --- /dev/null +++ b/src/dodal/devices/beamlines/i09_2/__init__.py @@ -0,0 +1,7 @@ +from .i09_2_motors import ( + I092SampleManipulator, + PiezoElectricMotor, + PiezoElectricMovableLogic, +) + +__all__ = ["I092SampleManipulator", "PiezoElectricMotor", "PiezoElectricMovableLogic"] diff --git a/src/dodal/devices/beamlines/i09_2/i09_2_motors.py b/src/dodal/devices/beamlines/i09_2/i09_2_motors.py new file mode 100644 index 00000000000..10c2cdf1fff --- /dev/null +++ b/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) diff --git a/src/dodal/devices/beamlines/i10_1/high_field_magnet/high_field_magnet.py b/src/dodal/devices/beamlines/i10_1/high_field_magnet/high_field_magnet.py index ff39102e93b..2720323532d 100644 --- a/src/dodal/devices/beamlines/i10_1/high_field_magnet/high_field_magnet.py +++ b/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, diff --git a/src/dodal/devices/movable.py b/src/dodal/devices/movable.py new file mode 100644 index 00000000000..96c9c7e83b0 --- /dev/null +++ b/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, + ) diff --git a/tests/devices/beamlines/i09_2/__init__.py b/tests/devices/beamlines/i09_2/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/devices/beamlines/i09_2/test_i09_2_motors.py b/tests/devices/beamlines/i09_2/test_i09_2_motors.py new file mode 100644 index 00000000000..e90e566f812 --- /dev/null +++ b/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 diff --git a/tests/devices/beamlines/i10_1/high_field_magnet/test_high_field_magnet.py b/tests/devices/beamlines/i10_1/high_field_magnet/test_high_field_magnet.py index e39f153e9e3..11c6cae3ac9 100644 --- a/tests/devices/beamlines/i10_1/high_field_magnet/test_high_field_magnet.py +++ b/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, diff --git a/tests/devices/test_movable.py b/tests/devices/test_movable.py new file mode 100644 index 00000000000..04ca5daf454 --- /dev/null +++ b/tests/devices/test_movable.py @@ -0,0 +1,129 @@ +import asyncio + +import pytest +from ophyd_async.core import ( + AsyncStatus, + DeviceMock, + set_mock_value, + soft_signal_r_and_setter, + soft_signal_rw, + wait_for_value, +) + +from dodal.devices.movable import MovableWithTolerance + + +class MovableWithToleranceImpl(MovableWithTolerance): + def __init__(self, name: str = ""): + self.custom_tolerance = soft_signal_rw(float) + self.custom_setpoint = soft_signal_rw(float) + self.custom_readback, _ = soft_signal_r_and_setter(float) + super().__init__( + tolerance=self.custom_tolerance, + setpoint=self.custom_setpoint, + readback=self.custom_readback, + name=name, + ) + + +@pytest.fixture +async def movable_with_tolerance() -> MovableWithToleranceImpl: + movable_with_tolerance = MovableWithToleranceImpl("movable_with_tolerance") + # Setup mock to not use default as do not want the setpoint and readback to be in + # sync for tests so we can correctly test threshold signal. + await movable_with_tolerance.connect(mock=DeviceMock()) + return movable_with_tolerance + + +@pytest.mark.parametrize( + "setpoint, readback, tolerance, expected_within_threshold", + [ + (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_movable_with_tolerance_within_threshold( + movable_with_tolerance: MovableWithToleranceImpl, + tolerance: float, + setpoint: float, + readback: float, + expected_within_threshold: bool, +) -> None: + set_mock_value(movable_with_tolerance.custom_tolerance, tolerance) + set_mock_value(movable_with_tolerance.custom_setpoint, setpoint) + set_mock_value(movable_with_tolerance.custom_readback, readback) + assert ( + await movable_with_tolerance.within_tolerance.get_value() + == expected_within_threshold + ) + + +@pytest.mark.parametrize( + "initial_readback, initial_setpoint, initial_within_tolerance", + ((0, 0, True), (0, -10, False)), + ids=("initial_within_tolerance[True]", "initial_within_tolerance[False]"), +) +async def test_movable_with_tolerance_logic_moves_to_setpoint_and_is_done_when_within_tolerance( + initial_readback: float, + initial_setpoint: float, + initial_within_tolerance: bool, + movable_with_tolerance: MovableWithToleranceImpl, +) -> None: + set_mock_value(movable_with_tolerance.custom_tolerance, 0.1) + + set_mock_value(movable_with_tolerance.movable_logic.setpoint, initial_setpoint) + set_mock_value(movable_with_tolerance.movable_logic.readback, initial_readback) + assert ( + await movable_with_tolerance.movable_logic.within_tolerance.get_value() + is initial_within_tolerance + ) + setpoint = 10 + + async with AsyncStatus( + movable_with_tolerance.movable_logic.move(new_position=setpoint, timeout=1) + ) as move_status: + # This prevents a race where the test proceeds before the move coroutine has + # had a chance to execute its first steps. + await wait_for_value( + movable_with_tolerance.movable_logic.setpoint, setpoint, timeout=1 + ) + assert ( + await movable_with_tolerance.movable_logic.setpoint.get_value() == setpoint + ) + # Set some values between initial readback and final setpoint that are outside + # the threshold to test signal is correct and status hasn't completed. + for value in [2.0, 5.0, 9.5, 13.0]: + set_mock_value(movable_with_tolerance.movable_logic.readback, value) + assert await movable_with_tolerance.within_tolerance.get_value() is False + assert move_status.done is False + + # Now move to a value within threshold. + set_mock_value(movable_with_tolerance.movable_logic.readback, 9.91) + assert await movable_with_tolerance.within_tolerance.get_value() is True + + # Wait for the underlying move task to complete so no race condition. + # This ensures the status has fully processed the tolerance condition + # and transitioned to a finished state before asserting `done`. + await asyncio.wait_for(move_status.task, timeout=1) + assert move_status.done is True + + +async def test_movable_with_tolerance_sub_class_signal_names_are_not_renamed( + movable_with_tolerance: MovableWithToleranceImpl, +) -> None: + assert ( + movable_with_tolerance.custom_setpoint.name + == "movable_with_tolerance-custom_setpoint" + ) + assert ( + movable_with_tolerance.custom_tolerance.name + == "movable_with_tolerance-custom_tolerance" + ) + # Readback is the exception as renamed by StandardMovable to be device name + assert movable_with_tolerance.custom_readback.name == "movable_with_tolerance"