diff --git a/src/dodal/common/maths.py b/src/dodal/common/maths.py index 4a6d00e48ee..a966c4e66da 100644 --- a/src/dodal/common/maths.py +++ b/src/dodal/common/maths.py @@ -1,4 +1,11 @@ +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Self + import numpy as np +from ophyd_async.core import ( + Array1D, +) def step_to_num(start: float, stop: float, step: float) -> tuple[float, float, int]: @@ -146,3 +153,100 @@ def rotate_clockwise(theta: float, x: float, y: float) -> tuple[float, float]: def rotate_counter_clockwise(theta: float, x: float, y: float) -> tuple[float, float]: return rotate_clockwise(-theta, x, y) + + +MotorOffsetAndPhase = Array1D[np.float32] + + +@dataclass +class AngleWithPhase: + """Represents a point in an absolute rotational space which is defined by a phase where 0<=phase<360 + and an offset from an origin where the absolute coordinate is the sum of the phase and the offset. + + Attributes: + offset: The offset of 0 phase from some other unwrapped rotational coordinate space + phase: The phase in degrees relative to this offset. + """ + + offset: float + phase: float = 0.0 + + def __post_init__(self) -> None: + correction = 360 * (self.phase // 360) + self.offset += correction + self.phase -= correction + + @classmethod + def from_iterable(cls, values: Iterable[float]) -> Self: + """Construct a normalised representation of the offset and phase, such that + 0 <= phase < 360. + + Args: + values (Iterable[float]): the offset and phase as a list or other iterable + """ + offset, phase = values + return cls(offset, phase) + + @classmethod + def wrap(cls, unwrapped: float) -> "AngleWithPhase": + """Construct a representation such that offset = n * 360 and 0 <= phase < 360. + + Args: + unwrapped (float): The unwrapped angle in degrees + """ + offset = AngleWithPhase.offset_from_unwrapped(unwrapped) + return cls(offset, unwrapped - offset) + + def rebase_to(self, other: Self) -> "AngleWithPhase": + """Return this angle with the offset adjusted such that the phases can be compared.""" + correction = other.offset - self.offset + if correction % 360: + return AngleWithPhase.from_iterable( + [self.offset + correction, self.phase - correction] + ) + else: + return self + + def unwrap(self) -> float: + """Generate the unwrapped representation of this angle.""" + return self.offset + self.phase + + def phase_distance(self, phase: float) -> float: + """Determine the shortest distance between this angle and the specified phase. + + Args: + phase (float): The phase angle to compare to + """ + phase = phase % 360 + max_theta = max(self.phase, phase) + min_theta = min(self.phase, phase) + return min(max_theta - min_theta, min_theta + 360 - max_theta) + + @classmethod + def offset_from_unwrapped(cls, unwrapped_deg: float) -> float: + """Obtain the offset from the corresponding wrapped angle in degrees.""" + return round(unwrapped_deg // 360) * 360 + + def nearest_with_phase(self, phase_deg: float) -> "AngleWithPhase": + """Return the nearest angle to this one with the specified phase.""" + phase_deg = phase_deg % 360 + if phase_deg > self.phase: + return ( + AngleWithPhase(self.offset, phase_deg) + if phase_deg - self.phase <= 180 + else AngleWithPhase(self.offset - 360, phase_deg) + ) + else: + return ( + AngleWithPhase(self.offset, phase_deg) + if self.phase - phase_deg <= 180 + else AngleWithPhase(self.offset + 360, phase_deg) + ) + + +def reflect_phase(phase) -> float: + """Convert the phase angle as if the corresponding unwrapped angle were to + be reflected about the origin and then re-wrapped, this corresponds to + converting a clockwise angle to an anti-clockwise angle and vice-versa. + """ + return (360 - phase) % 360 diff --git a/src/dodal/devices/aithre_lasershaping/goniometer.py b/src/dodal/devices/aithre_lasershaping/goniometer.py index ee1c2aec1e1..44f1f757f06 100644 --- a/src/dodal/devices/aithre_lasershaping/goniometer.py +++ b/src/dodal/devices/aithre_lasershaping/goniometer.py @@ -1,9 +1,9 @@ from ophyd_async.epics.motor import Motor -from dodal.devices.motors import XYZOmegaStage, create_axis_perp_to_rotation +from dodal.devices.motors import XYZWrappedOmegaStage, create_axis_perp_to_rotation -class Goniometer(XYZOmegaStage): +class Goniometer(XYZWrappedOmegaStage): """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 diff --git a/src/dodal/devices/motors.py b/src/dodal/devices/motors.py index faecc38a5f8..20bea247e02 100644 --- a/src/dodal/devices/motors.py +++ b/src/dodal/devices/motors.py @@ -14,6 +14,7 @@ from ophyd_async.epics.motor import Motor from dodal.common.maths import rotate_clockwise, rotate_counter_clockwise +from dodal.devices.wrapped_axis import WrappedAxis _X = "X" _Y = "Y" @@ -114,9 +115,27 @@ def __init__( ) -> None: with self.add_children_as_readables(): self.omega = Motor(prefix + omega_infix) + super().__init__(prefix, name, x_infix, y_infix, z_infix) +class XYZWrappedOmegaStage(XYZOmegaStage): + """Four-axis stage with x, y, z linear axes and an omega axis with unrestricted rotation.""" + + def __init__( + self, + prefix: str = "", + name: str = "", + x_infix: str = _X, + y_infix: str = _Y, + z_infix: str = _Z, + omega_infix: str = _OMEGA, + ): + super().__init__(prefix, name, x_infix, y_infix, z_infix, omega_infix) + with self.add_children_as_readables(): + self.wrapped_omega = WrappedAxis(self.omega) + + class XYZAzimuthStage(XYZStage): """Four-axis stage with a standard xyz stage and one axis of rotation: azimuth.""" diff --git a/src/dodal/devices/smargon.py b/src/dodal/devices/smargon.py index 6ca3e9dde9d..b1a2897e385 100644 --- a/src/dodal/devices/smargon.py +++ b/src/dodal/devices/smargon.py @@ -14,7 +14,8 @@ from ophyd_async.epics.core import epics_signal_r, epics_signal_rw from ophyd_async.epics.motor import Motor -from dodal.devices.motors import XYZOmegaStage +from dodal.common.maths import AngleWithPhase +from dodal.devices.motors import XYZWrappedOmegaStage from dodal.devices.util.epics_util import SetWhenEnabled @@ -78,7 +79,7 @@ class CombinedMove(TypedDict, total=False): chi: float | None -class Smargon(XYZOmegaStage, Movable): +class Smargon(XYZWrappedOmegaStage, Movable): """Real motors added to allow stops following pin load (e.g. real_x1.stop() ) X1 and X2 real motors provide compound chi motion as well as the compound X travel, increasing the gap between x1 and x2 changes chi, moving together changes virtual x. @@ -104,6 +105,15 @@ def __init__(self, prefix: str, name: str = ""): super().__init__(prefix, name) + async def _get_target_value(self, motor_name: str, value: float) -> float: + if motor_name == "omega": + current_angle = AngleWithPhase.from_iterable( + await self.wrapped_omega.offset_and_phase.get_value() + ) + return current_angle.nearest_with_phase(value).unwrap() + else: + return value + @AsyncStatus.wrap async def set(self, value: CombinedMove): """This will move all motion together in a deferred move. @@ -114,11 +124,16 @@ async def set(self, value: CombinedMove): only come back after the motion on that axis finished. """ await self.defer_move.set(DeferMoves.ON) + # TODO Hotfix required here until https://github.com/DiamondLightSource/dodal/issues/1998 + # is implemented in separate PR try: finished_moving = [] for motor_name, new_setpoint in value.items(): if new_setpoint is not None and isinstance(new_setpoint, int | float): - axis: Motor = getattr(self, motor_name) + new_setpoint = await self._get_target_value( + motor_name, new_setpoint + ) + axis = getattr(self, motor_name) await axis.check_motor_limit( await axis.user_setpoint.get_value(), new_setpoint ) diff --git a/src/dodal/devices/wrapped_axis.py b/src/dodal/devices/wrapped_axis.py new file mode 100644 index 00000000000..e9c0dfb89a4 --- /dev/null +++ b/src/dodal/devices/wrapped_axis.py @@ -0,0 +1,73 @@ +import numpy as np +from ophyd_async.core import ( + Reference, + StandardReadable, + StandardReadableFormat, + derived_signal_rw, +) +from ophyd_async.epics.motor import Motor + +from dodal.common.maths import AngleWithPhase, MotorOffsetAndPhase + + +class WrappedAxis(StandardReadable): + """This device is a wrapper around a rotational Motor that presents the unwrapped coordinate system of the + underlying motor as a combination of a phase angle and an offset from the motor's origin coordinate. + + Attributes: + phase (float): This is a read-write signal that corresponds to the motor's phase angle, relative to the offset. + The behaviour of the phase signal differs from that of the underlying motor setpoint/readback signal + in the following ways: + * Readback values are constrained to 0 <= omega < 360 + * Write values are normalised to 0 <= omega < 360, and then mapped to the nearest unwrapped angle. The + underlying motor will be moved via the shortest path to the required phase angle, thus the direction + of an unwrapped axis move is not always the same as the direction of a move in the wrapped axis. + * Write values are normalised so for un-normalised writes, the readback will differ. + * Bluesky mvr operations greater of 180 degrees or more will not result in the expected moves. + * set_and_wait_for_value() is unreliable close to the wrap-around (however in the common case this is + where deadband is 0.001 and values are rounded to this level, the default equality comparison is sufficient). + * Sequences of moves on the unwrapped axis that would not result in cumulative motion axis will + result in a cumulative motion in the wrapped axis. e.g. 0 -> 120 -> 240 -> 0 is 0 degrees of real + cumulative motion when performed in the unwrapped case, but is 360 degrees of real motion when performed + on the wrapped axis. + * The reverse is also true 0 -> 240 -> 360 -> is 0 degrees of real motion when executed on the wrapped axis + but 360 degrees when performed in the unwrapped axis. + * The above means that use of phase to set motor position is unsuitable for axes + where the underlying motor rotation is constrained. + offset_and_phase (Array1D[np.float32]): retrieve the offset and phase together, for use when + mapping to the underlying unwrapped axis. These values can be converted and manipulated using the + AngleWithPhase helper class. + """ + + def __init__(self, real_motor: Motor, name=""): + self._real_motor = Reference(real_motor) + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.offset_and_phase = derived_signal_rw( + self._get_motor_offset_and_phase, + self._set_motor_offset_and_phase, + motor_pos=real_motor, + ) + with self.add_children_as_readables(): + self.phase = derived_signal_rw( + self._get_phase, self._set_phase, offset_and_phase=self.offset_and_phase + ) + super().__init__(name=name) + + def _get_motor_offset_and_phase(self, motor_pos: float) -> MotorOffsetAndPhase: + angle = AngleWithPhase.wrap(motor_pos) + return np.array([angle.offset, angle.phase]) + + async def _set_motor_offset_and_phase(self, value: MotorOffsetAndPhase): + await self._real_motor().set(AngleWithPhase.from_iterable(value).unwrap()) + + def _get_phase(self, offset_and_phase: MotorOffsetAndPhase) -> float: + return offset_and_phase[1].item() + + async def _set_phase(self, value: float): + """Set the motor phase to the specified phase value in degrees. + The motor will travel via the shortest distance path. + """ + offset_and_phase = await self.offset_and_phase.get_value() + current_position = AngleWithPhase.from_iterable(offset_and_phase) + target_value = current_position.nearest_with_phase(value).unwrap() + await self._real_motor().set(target_value) diff --git a/tests/common/test_maths.py b/tests/common/test_maths.py index 83179d3ad49..4458f876969 100644 --- a/tests/common/test_maths.py +++ b/tests/common/test_maths.py @@ -5,8 +5,10 @@ from dodal.common import in_micros, step_to_num from dodal.common.maths import ( + AngleWithPhase, Rectangle2D, do_rotation, + reflect_phase, rotate_clockwise, rotate_counter_clockwise, ) @@ -189,3 +191,139 @@ def test_do_rotation_matches_manual_matrix_multiply() -> None: assert x_new == pytest.approx(expected[0]) assert y_new == pytest.approx(expected[1]) + + +@pytest.mark.parametrize( + "offset, phase, expected_offset, expected_phase", + [ + [0, 0, 0, 0], + [360, 0, 360, 0], + [370, 0, 370, 0], + [390, 30, 390, 30], + [360, -30, 0, 330], + [-270, -90, -630, 270], + ], +) +def test_angle_with_phase_construct_from_offset_and_phase( + offset: float, phase: float, expected_offset: float, expected_phase: float +): + angle_with_phase = AngleWithPhase(offset, phase) + assert angle_with_phase.offset == expected_offset + assert angle_with_phase.phase == expected_phase + + +@pytest.mark.parametrize( + "unwrapped, expected_offset, expected_phase", + [ + [0, 0, 0], + [270, 0, 270], + [540, 360, 180], + [-90, -360, 270], + [180, 0, 180], + [360, 360, 0], + ], +) +def test_angle_with_phase_wrap( + unwrapped: float, expected_offset: float, expected_phase: float +): + angle_with_phase = AngleWithPhase.wrap(unwrapped) + assert angle_with_phase.offset == expected_offset + assert angle_with_phase.phase == expected_phase + + +@pytest.mark.parametrize( + "offset, phase, expected_unwrapped", + [ + [0, 0, 0], + [360, 0, 360], + [370, 0, 370], + [390, 30, 420], + [360, -30, 330], + [-270, -90, -360], + ], +) +def test_angle_with_phase_unwrap( + offset: float, phase: float, expected_unwrapped: float +): + angle_with_phase = AngleWithPhase(offset, phase) + assert angle_with_phase.unwrap() == expected_unwrapped + + +@pytest.mark.parametrize( + "offset1, phase1, offset2, phase2, expected_offset, expected_phase", + [ + [0, 0, 360, 0, 0, 0], + [0, 0, 10, 0, -350, 350], + [0, 90, 370, 0, 10, 80], + ], +) +def test_angle_with_phase_rebase( + offset1: float, + phase1: float, + offset2: float, + phase2: float, + expected_offset: float, + expected_phase: float, +): + angle1 = AngleWithPhase(offset1, phase1) + angle2 = AngleWithPhase(offset2, phase2) + rebased = angle1.rebase_to(angle2) + assert rebased.offset == expected_offset + assert rebased.phase == expected_phase + + +@pytest.mark.parametrize( + "phase1, phase2, expected_distance", + [ + [0, 0, 0], + [0, 360, 0], + [0, 330, 30], + [0, 30, 30], + [0, 179, 179], + [0, 180, 180], + [0, 181, 179], + [360, 350, 10], + [355, -90, 85], + ], +) +def test_angle_with_phase_phase_distance( + phase1: float, + phase2: float, + expected_distance: float, +): + angle = AngleWithPhase(0, phase1) + assert angle.phase_distance(phase2) == expected_distance + assert AngleWithPhase(0, phase2).phase_distance(phase1) == expected_distance + + +@pytest.mark.parametrize( + "offset, phase, target_phase, expected_offset, expected_phase", + [ + [0, 0, 0, 0, 0], + [0, 0, 270, -360, 270], + [0, 0, 90, 0, 90], + [360, 179, 181, 360, 181], + [360, 181, 179, 360, 179], + [360, 1, -1, 0, 359], + [360, 1, 359, 0, 359], + ], +) +def test_angle_with_phase_nearest_with_phase( + offset: float, + phase: float, + target_phase: float, + expected_offset: float, + expected_phase: float, +): + angle = AngleWithPhase(offset, phase) + nearest = angle.nearest_with_phase(target_phase) + assert nearest.offset == expected_offset + assert nearest.phase == expected_phase + + +@pytest.mark.parametrize( + "phase, expected_phase", + [[1, 359], [-1, 1], [180, 180], [179, 181], [181, 179], [360, 0], [0, 0]], +) +def test_reflect_phase(phase: float, expected_phase: float): + assert reflect_phase(phase) == expected_phase diff --git a/tests/devices/test_motors.py b/tests/devices/test_motors.py index ee8c66262be..60eefe8693d 100644 --- a/tests/devices/test_motors.py +++ b/tests/devices/test_motors.py @@ -1,8 +1,10 @@ import math -from collections.abc import Generator +from collections.abc import Generator, Iterable +import numpy as np import pytest -from ophyd_async.core import init_devices, set_mock_value +from bluesky import RunEngine +from ophyd_async.core import get_mock_put, init_devices, set_mock_value from ophyd_async.testing import assert_reading, partial_reading from dodal.devices.motors import ( @@ -14,6 +16,7 @@ XYZAzimuthTiltStage, XYZPitchYawRollStage, XYZThetaStage, + XYZWrappedOmegaStage, ) @@ -66,6 +69,18 @@ async def xyzpyr_stage() -> XYZPitchYawRollStage: return xyzpyr_stage +@pytest.fixture() +async def xyz_wrapped_omega_stage( + values_for_rotation: Iterable[float], +) -> XYZWrappedOmegaStage: + input_value, current_real_value, expected_real_value = values_for_rotation + + async with init_devices(mock=True): + xyz_wrapped_omega_stage = XYZWrappedOmegaStage("BL03I-MO-SGON-01:") + set_mock_value(xyz_wrapped_omega_stage.omega.user_readback, current_real_value) + return xyz_wrapped_omega_stage + + async def test_setting_xy_position_table(xyzt_stage: XYZThetaStage): """Test setting x and y positions on the Table using the ophyd_async mock tools.""" await assert_reading( @@ -311,6 +326,25 @@ async def test_reading_six_axis_gonio(six_axis_gonio: SixAxisGonio): ) +@pytest.mark.parametrize("values_for_rotation", [[0, 0, 0]], indirect=True) +async def test_reading_xyz_wrapped_omega_gonio( + xyz_wrapped_omega_stage: XYZWrappedOmegaStage, values_for_rotation: Iterable[float] +): + await assert_reading( + xyz_wrapped_omega_stage, + { + "xyz_wrapped_omega_stage-omega": partial_reading(0.0), + "xyz_wrapped_omega_stage-z": partial_reading(0.0), + "xyz_wrapped_omega_stage-y": partial_reading(0.0), + "xyz_wrapped_omega_stage-x": partial_reading(0.0), + "xyz_wrapped_omega_stage-wrapped_omega-phase": partial_reading(0.0), + "xyz_wrapped_omega_stage-wrapped_omega-offset_and_phase": partial_reading( + np.array([0, 0]) + ), + }, + ) + + @pytest.mark.parametrize( "set_value, omega_set_value, expected_horz, expected_vert", [ @@ -358,3 +392,121 @@ 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_phase", + [ + [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_phase: float, + xyz_wrapped_omega_stage: XYZWrappedOmegaStage, +): + set_mock_value(xyz_wrapped_omega_stage.omega.user_readback, real_value) + ( + offset, + phase, + ) = await xyz_wrapped_omega_stage.wrapped_omega.offset_and_phase.get_value() + assert 0 <= expected_phase < 360 + assert math.isclose(phase, expected_phase, abs_tol=1e-6) + assert math.isclose(phase + offset, real_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: pytest.FixtureRequest, +) -> Generator[tuple[float, float, float], None, None]: + input_value, current_real_value, expected_real_value = request.param + yield input_value, current_real_value, expected_real_value + + +async def test_mod_360_expected_direction_of_rotation_same_as_apparent_for_moves_apparently_less_than_180( + values_for_rotation: Iterable[float], xyz_wrapped_omega_stage: XYZWrappedOmegaStage +): + input_value, current_real_value, expected_real_value = values_for_rotation + stage = xyz_wrapped_omega_stage + + current_readback = await stage.omega.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_expected_actual_movement_never_more_than_180( + values_for_rotation: tuple[float, float, float], + xyz_wrapped_omega_stage: XYZWrappedOmegaStage, +): + 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_unwrap_computes_expected( + values_for_rotation: tuple[float, float, float], + xyz_wrapped_omega_stage: XYZWrappedOmegaStage, + run_engine: RunEngine, +) -> None: + input_value, current_real_value, expected_real_value = values_for_rotation + stage = xyz_wrapped_omega_stage + real_put = get_mock_put(stage.omega.user_setpoint) + await stage.wrapped_omega.phase.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 a11a2f1f85d..73118ff2b7c 100644 --- a/tests/devices/test_smargon.py +++ b/tests/devices/test_smargon.py @@ -85,6 +85,39 @@ async def test_given_center_disp_low_when_stub_offsets_set_to_center_and_moved_t assert await smargon.stub_offsets.to_robot_load.proc.get_value() == 0 +async def test_set_with_omega_outside_smargon_limit( + smargon: Smargon, +): + set_mock_value(smargon.omega.low_limit_travel, -1999) + set_mock_value(smargon.omega.high_limit_travel, 1999) + set_mock_value(smargon.omega.dial_low_limit_travel, -1999) + set_mock_value(smargon.omega.dial_high_limit_travel, 1999) + await smargon.omega.set(1999) + with pytest.raises(MotorLimitsError): + await smargon.set( + CombinedMove( + x=10, + y=20, + z=30, + omega=200, + chi=15, + phi=25, + ) + ) + await smargon.omega.set(-1999) + with pytest.raises(MotorLimitsError): + await smargon.set( + CombinedMove( + x=10, + y=20, + z=30, + omega=160, + chi=15, + phi=25, + ) + ) + + @pytest.mark.parametrize( "test_x, test_y, test_z, test_omega, test_chi, test_phi", [ @@ -94,8 +127,6 @@ async def test_given_center_disp_low_when_stub_offsets_set_to_center_and_moved_t (10, -2000, 30, 5, 15, 25), # y goes beyond lower limit (10, 20, 2000, 5, 15, 25), # z goes beyond upper limit (10, 20, -2000, 5, 15, 25), # z goes beyond lower limit - (10, 20, 30, 2000, 15, 25), # omega goes beyond upper limit - (10, 20, 30, -2000, 15, 25), # omega goes beyond lower limit (10, 20, 30, 5, 2000, 25), # chi goes beyond upper limit (10, 20, 30, 5, -2000, 25), # chi goes beyond lower limit (10, 20, 30, 5, 15, 2000), # phi goes beyond upper limit @@ -109,7 +140,6 @@ async def test_given_set_with_value_outside_motor_limit( smargon.x, smargon.y, smargon.z, - smargon.omega, smargon.chi, smargon.phi, ]: diff --git a/tests/devices/test_wrapped_axis.py b/tests/devices/test_wrapped_axis.py new file mode 100644 index 00000000000..f86f1466798 --- /dev/null +++ b/tests/devices/test_wrapped_axis.py @@ -0,0 +1,110 @@ +import numpy as np +import pytest +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.wrapped_axis import WrappedAxis + + +@pytest.fixture +async def omega() -> Motor: + async with init_devices(mock=True, connect=True): + motor = Motor("BL03I-MO-SGON-01:OMEGA") + return motor + + +@pytest.fixture +def wrapped_omega(omega: Motor): + with init_devices(mock=True): + wrapped_omega = WrappedAxis(omega, "wrapped_omega") + return wrapped_omega + + +@pytest.mark.parametrize( + "initial_unwrapped, phase, expected_omega", + [ + [0, 0, 0], + [0, 180, 180], + [0, 360, 0], + [180, 360, 0], + [181, 360, 360], + [270, 0, 360], + [-270, 0, -360], + [0, 450, 90], + ], +) +async def test_wrapped_axis_set_applies_phase_with_current_set( + omega: Motor, + wrapped_omega: WrappedAxis, + initial_unwrapped: float, + phase: float, + expected_omega: float, +): + set_mock_value(omega.user_readback, initial_unwrapped) + await wrapped_omega.phase.set(phase) + get_mock_put(omega.user_setpoint).assert_called_with(expected_omega) + + +@pytest.mark.parametrize( + "initial_unwrapped, expected_phase", + [ + [0, 0], + [360, 0], + [-360, 0], + [90, 90], + [361, 1], + [-90, 270], + [-270, 90], + [720, 0], + [450, 90], + ], +) +async def test_wrapped_axis_get_returns_current_phase( + omega: Motor, + wrapped_omega: WrappedAxis, + initial_unwrapped: float, + expected_phase: float, +): + set_mock_value(omega.user_readback, initial_unwrapped) + phase = await wrapped_omega.phase.get_value() + assert phase == expected_phase + + +async def test_wrapped_axis_read_returns_phase_and_offset_phase( + omega: Motor, wrapped_omega: WrappedAxis +): + set_mock_value(omega.user_readback, 450) + await assert_reading( + wrapped_omega, + { + "wrapped_omega-offset_and_phase": partial_reading(np.array([360, 90])), + "wrapped_omega-phase": partial_reading(90), + }, + ) + + +@pytest.mark.parametrize( + "offset, phase, expected_unwrapped", + [ + [0, 0, 0], + [360, 0, 360], + [-360, 0, -360], + [0, 90, 90], + [360, 1, 361], + [360, -90, 270], + [0, 270, 270], + [-360, 90, -270], + [720, 0, 720], + [360, 90, 450], + ], +) +async def test_wrapped_axis_set_offset_and_phase( + omega: Motor, + wrapped_omega: WrappedAxis, + offset: float, + phase: float, + expected_unwrapped: float, +): + await wrapped_omega.offset_and_phase.set(np.array([offset, phase])) + get_mock_put(omega.user_setpoint).assert_called_once_with(expected_unwrapped)