Mx bluesky 1598 mod 360

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)
 
 

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):

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(

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

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}"
+    )

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(