diff --git a/src/dodal/devices/motors.py b/src/dodal/devices/motors.py index fc84bcc5a2..510e0ef37b 100644 --- a/src/dodal/devices/motors.py +++ b/src/dodal/devices/motors.py @@ -1,8 +1,19 @@ import asyncio import math from abc import ABC - -from ophyd_async.core import SignalRW, StandardReadable, derived_signal_rw +from functools import partial + +from bluesky.protocols import Locatable, Location +from ophyd_async.core import ( + CALCULATE_TIMEOUT, + AsyncStatus, + CalculatableTimeout, + SignalRW, + StandardReadable, + StandardReadableFormat, + derived_signal_r, + derived_signal_rw, +) from ophyd_async.epics.motor import Motor from dodal.common.maths import rotate_clockwise, rotate_counter_clockwise @@ -89,6 +100,89 @@ def __init__( super().__init__(prefix, name, x_infix, y_infix, z_infix) +class ModMotor(StandardReadable, Locatable[float]): + def __init__(self, real_motor: Motor, name=""): + self._real_motor = real_motor + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.user_readback = derived_signal_r( + self._get_mod_360, real_value=self._real_motor.user_readback + ) + with self.add_children_as_readables(): + self.user_setpoint = derived_signal_rw( + self._get_mod_360, + partial(self._set_mod_360, self._real_motor.user_setpoint), + real_value=self._real_motor.user_setpoint, + ) + self.velocity = derived_signal_rw( + self._identity, + partial(self._set_value, self._real_motor.velocity), + value=self._real_motor.velocity, + ) + self.max_velocity = derived_signal_r( + self._identity, value=self._real_motor.max_velocity + ) + self.acceleration_time = derived_signal_r( + self._identity, value=self._real_motor.acceleration_time + ) + + super().__init__(name=name) + + def set_name(self, name: str, *, child_name_separator: str | None = None) -> None: + """Set name of the motor and its children.""" + super().set_name(name, child_name_separator=child_name_separator) + # Readback should be named the same as its parent in read() + self.user_readback.set_name(name) + + def _identity(self, value: float) -> float: + return value + + async def check_motor_limit(self, abs_start_pos: float, abs_end_pos: float): + offset = await self._current_offset() + await self._real_motor.check_motor_limit( + abs_start_pos + offset, abs_end_pos + offset + ) + + async def _set_value(self, real_signal: SignalRW, modded_value: float): + await real_signal.set(modded_value) + + def _get_mod_360(self, real_value: float) -> float: + return real_value - self._offset_from_real_value(real_value) + + def _offset_from_real_value(self, real_value: float) -> float: + return round(real_value // 360) * 360 + + async def _current_offset(self): + real_value = await self._real_motor.user_readback.get_value() + return self._offset_from_real_value(real_value) + + async def _set_mod_360(self, signal_to_set: SignalRW, input_value: float): + real_value = await self._real_motor.user_readback.get_value() + offset = self._offset_from_real_value(real_value) + target_value = input_value + offset + while target_value - real_value >= 180: + target_value -= 360 + else: + while target_value - real_value <= -180: + target_value += 360 + + await signal_to_set.set(target_value) + + async def locate(self) -> Location[float]: + """Return the current setpoint and readback of the motor.""" + setpoint, readback = await asyncio.gather( + self.user_setpoint.get_value(), self.user_readback.get_value() + ) + return Location(setpoint=setpoint, readback=readback) + + @AsyncStatus.wrap + async def set( + self, value: float, timeout: CalculatableTimeout = CALCULATE_TIMEOUT + ) -> Location[float]: + offset = await self._current_offset() + await self._real_motor.set(value + offset, timeout) + return await self.locate() + + class XYZOmegaStage(XYZStage): """Four-axis stage with a standard xyz stage and one axis of rotation: omega.""" @@ -100,9 +194,14 @@ def __init__( y_infix: str = _Y, z_infix: str = _Z, omega_infix: str = _OMEGA, + mod_360_motor: bool = False, ) -> None: with self.add_children_as_readables(): - self.omega = Motor(prefix + omega_infix) + real_motor = Motor(prefix + omega_infix) + if mod_360_motor: + self.omega = ModMotor(real_motor) + else: + self.omega = real_motor super().__init__(prefix, name, x_infix, y_infix, z_infix) diff --git a/src/dodal/devices/smargon.py b/src/dodal/devices/smargon.py index 6ca3e9dde9..e2cfcf9cef 100644 --- a/src/dodal/devices/smargon.py +++ b/src/dodal/devices/smargon.py @@ -102,7 +102,7 @@ def __init__(self, prefix: str, name: str = ""): self.defer_move = epics_signal_rw(DeferMoves, prefix + "CS1:DeferMoves") - super().__init__(prefix, name) + super().__init__(prefix, name, mod_360_motor=True) @AsyncStatus.wrap async def set(self, value: CombinedMove): diff --git a/tests/common/beamlines/test_beamline_utils.py b/tests/common/beamlines/test_beamline_utils.py index ec58f06677..44997d0cd4 100644 --- a/tests/common/beamlines/test_beamline_utils.py +++ b/tests/common/beamlines/test_beamline_utils.py @@ -82,7 +82,7 @@ def test_instantiate_v2_function_fake_makes_fake(): Smargon, "smargon", "", True, True, None ) assert isinstance(fake_smargon, StandardReadable) - assert fake_smargon.omega.user_setpoint.source.startswith("mock+ca") + assert fake_smargon.omega._real_motor.user_setpoint.source.startswith("mock+ca") @pytest.mark.parametrize( diff --git a/tests/devices/oav/test_oav_utils.py b/tests/devices/oav/test_oav_utils.py index a1295e8363..50c931f639 100644 --- a/tests/devices/oav/test_oav_utils.py +++ b/tests/devices/oav/test_oav_utils.py @@ -85,7 +85,7 @@ async def test_values_for_move_so_that_beam_is_at_pixel( run_engine: RunEngine, ): set_mock_value(oav.zoom_controller.level, zoom_level) - set_mock_value(smargon.omega.user_readback, angle) + set_mock_value(smargon.omega._real_motor.user_readback, angle) pos = run_engine( get_move_required_so_that_beam_is_at_pixel(smargon, pixel_to_move_to, oav) ).plan_result # type: ignore diff --git a/tests/devices/test_motors.py b/tests/devices/test_motors.py index daa8335af5..20ced64893 100644 --- a/tests/devices/test_motors.py +++ b/tests/devices/test_motors.py @@ -2,10 +2,12 @@ from collections.abc import Generator import pytest -from ophyd_async.core import init_devices, set_mock_value +from ophyd_async.core import get_mock_put, init_devices, set_mock_value +from ophyd_async.epics.motor import Motor from ophyd_async.testing import assert_reading, partial_reading from dodal.devices.motors import ( + ModMotor, SixAxisGonio, XThetaStage, XYStage, @@ -358,3 +360,119 @@ async def test_vertical_in_lab_space_for_default_axes( 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 + + +@pytest.mark.parametrize( + "real_value, expected_value", + [ + [0, 0], + [0.001, 0.001], + [360, 0], + [-359.999, 0.001], + [-360, 0], + [-180.001, 179.999], + [-179.999, 180.001], + [-0.001, 359.999], + [180.001, 180.001], + [360.001, 0.001], + [719.999, 359.999], + [720.001, 0.001], + [10000 * 360 + 0.001, 0.001], + [-10000 * 360 - 0.001, 359.999], + ], +) +async def test_mod_360_read(real_value: float, expected_value: float): + motor = ModMotor(Motor("BL03I-MO-SGON-01:OMEGA")) + await motor.connect(mock=True) + set_mock_value(motor._real_motor.user_readback, real_value) + readback_value = await motor.user_readback.get_value() + assert math.isclose(readback_value, expected_value, abs_tol=1e-6) + + +@pytest.fixture( + params=[ + [0, 0, 0], + [0.001, 0, 0.001], + [-5, 0, -5], + [360, 0, 0], + [-359.999, 0, 0.001], + [-360, 0, 0], + [-180.001, 0, 179.999], + [-179.999, 0, -179.999], + [-0.001, 0, -0.001], + [0, 90, 0], + [0.001, 90, 0.001], + [360, 90, 0], + [-359.999, 90, 0.001], + [-360, 90, 0], + [-180.001, 90, 179.999], + [-179.999, 90, 180.001], + [-0.001, 90, -0.001], + [0, 90, 0], + [0.001, 450, 360.001], + [360, 450, 360], + [-359.999, 450, 360.001], + [-360, 450, 360], + [-180.001, 450, 539.999], + [-179.999, 450, 540.001], + [-0.001, 450, 359.999], + [180.001, -270, -179.999], + [360.001, -270, -359.999], + [719.999, -270, -360.001], + [-720.001, 270, 359.999], + [10000 * 360 + 0.001, 360000 - 90, 360000.001], + [-10000 * 360 - 0.001, -360000 + 90, -360000.001], + ], + ids=lambda values: f"input={values[0]}, current={values[1]}, expected={values[2]}", +) +def values_for_rotation(request): + input_value, current_real_value, expected_real_value = request.param + yield input_value, current_real_value, expected_real_value + + +@pytest.fixture() +async def motor_in_initial_state(values_for_rotation): + input_value, current_real_value, expected_real_value = values_for_rotation + + motor = ModMotor(Motor("BL03I-MO-SGON-01:OMEGA")) + await motor.connect(mock=True) + set_mock_value(motor._real_motor.user_readback, current_real_value) + return motor + + +async def test_mod_360_write_actual_direction_of_rotation_same_as_apparent_for_moves_apparently_less_than_180( + values_for_rotation, motor_in_initial_state +): + input_value, current_real_value, expected_real_value = values_for_rotation + motor = motor_in_initial_state + + current_readback = await motor.user_readback.get_value() + is_no_move = expected_real_value == current_real_value + is_ge_than_180 = abs(input_value - current_readback) >= 180 + motion_in_same_direction_as_apparent = (input_value < current_readback) == ( + expected_real_value < current_real_value + ) + assert is_no_move or is_ge_than_180 or motion_in_same_direction_as_apparent, ( + f"expected ({input_value} - {current_real_value}) same sign as ({expected_real_value} - {current_real_value})" + ) + + +async def test_mod_360_write_actual_movement_never_more_than_180( + values_for_rotation, motor_in_initial_state +): + input_value, current_real_value, expected_real_value = values_for_rotation + assert abs(expected_real_value - current_real_value) <= 180 + + +async def test_mod_360_put_proceeds_as_expected( + values_for_rotation, motor_in_initial_state +): + input_value, current_real_value, expected_real_value = values_for_rotation + motor = motor_in_initial_state + real_put = get_mock_put(motor._real_motor.user_setpoint) + await motor.user_setpoint.set(input_value) + real_put.assert_called_once() + actual = real_put.mock_calls[0].args[0] + assert math.isclose(actual, expected_real_value), ( + f"expected {expected_real_value} but was {actual}" + ) diff --git a/tests/devices/test_smargon.py b/tests/devices/test_smargon.py index e6d4834cd7..688769a658 100644 --- a/tests/devices/test_smargon.py +++ b/tests/devices/test_smargon.py @@ -109,7 +109,7 @@ async def test_given_set_with_value_outside_motor_limit( smargon.x, smargon.y, smargon.z, - smargon.omega, + # smargon.omega, smargon.chi, smargon.phi, ]: @@ -118,6 +118,11 @@ async def test_given_set_with_value_outside_motor_limit( set_mock_value(motor.dial_low_limit_travel, -1999) set_mock_value(motor.dial_high_limit_travel, 1999) + set_mock_value(smargon.omega._real_motor.low_limit_travel, -1999) + set_mock_value(smargon.omega._real_motor.high_limit_travel, 1999) + set_mock_value(smargon.omega._real_motor.dial_low_limit_travel, -1999) + set_mock_value(smargon.omega._real_motor.dial_high_limit_travel, 1999) + with pytest.raises(MotorLimitsError): await smargon.set( CombinedMove(