diff --git a/src/dodal/beamlines/i06_1.py b/src/dodal/beamlines/i06_1.py index b98a1d6a9e9..86b99c26c75 100644 --- a/src/dodal/beamlines/i06_1.py +++ b/src/dodal/beamlines/i06_1.py @@ -2,6 +2,12 @@ from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline from dodal.device_manager import DeviceManager from dodal.devices.beamlines.i06_1 import DiffractionDichroism +from dodal.devices.beamlines.i06_1.magnet import ( + SuperConductingMagnetController, +) +from dodal.devices.beamlines.i06_1.magnet.superconducting_magnet import ( + SuperConductingMagnet, +) from dodal.devices.motors import XYThetaStage from dodal.devices.temperture_controller import Lakeshore336 from dodal.log import set_beamline as set_log_beamline @@ -17,12 +23,12 @@ @devices.factory() -def diff_cooling_temperature_controller() -> Lakeshore336: +def ls336_cooling() -> Lakeshore336: return Lakeshore336(prefix=f"{PREFIX.beamline_prefix}-EA-TCTRL-02:") @devices.factory() -def diff_heating_temperature_controller() -> Lakeshore336: +def ls336_heating() -> Lakeshore336: return Lakeshore336(prefix=f"{PREFIX.beamline_prefix}-EA-TCTRL-03:") @@ -34,3 +40,11 @@ def xabs() -> XYThetaStage: @devices.factory() def dd() -> DiffractionDichroism: return DiffractionDichroism(f"{PREFIX.beamline_prefix}-EA-DDIFF-01:") + + +@devices.factory() +def scmc() -> SuperConductingMagnet: + controller = SuperConductingMagnetController( + prefix=f"{PREFIX.beamline_prefix}-EA", + ) + return SuperConductingMagnet(controller) diff --git a/src/dodal/devices/beamlines/i06_1/magnet/__init__.py b/src/dodal/devices/beamlines/i06_1/magnet/__init__.py new file mode 100644 index 00000000000..b1ddbc977bc --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnet/__init__.py @@ -0,0 +1,31 @@ +from .coordinates import ( + MagnetPosition, + MagnetRequest, + MagnetSphericalPosition, + MagnetSphericalRequest, +) +from .movement import MagnetPositionError +from .superconducting_magnet import ( + FlyVectorMagnetInfo, + MagnetAxis, + MagnetLimitStatus, + MagnetMode, + MagnetRampStatus, + SuperConductingMagnet, + SuperConductingMagnetController, +) + +__all__ = [ + "MagnetPosition", + "MagnetSphericalPosition", + "MagnetSphericalRequest", + "MagnetRequest", + "MagnetPositionError", + "FlyVectorMagnetInfo", + "MagnetAxis", + "MagnetLimitStatus", + "MagnetMode", + "MagnetRampStatus", + "SuperConductingMagnetController", + "SuperConductingMagnet", +] diff --git a/src/dodal/devices/beamlines/i06_1/magnet/coordinates.py b/src/dodal/devices/beamlines/i06_1/magnet/coordinates.py new file mode 100644 index 00000000000..3c1aba32134 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnet/coordinates.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass + + +@dataclass(kw_only=True) +class MagnetPosition: + x: float + y: float + z: float + + @staticmethod + def calc_theta(x: float, z: float) -> float: + theta = math.degrees(math.atan2(-x, z)) + return theta % 360 + + @staticmethod + def calc_rho(x: float, y: float, z: float) -> float: + return math.hypot(x, y, z) + + @staticmethod + def calc_phi(x: float, y: float, z: float) -> float: + return math.degrees(math.atan2(math.hypot(x, z), y)) + + def to_spherical(self) -> MagnetSphericalPosition: + return MagnetSphericalPosition( + rho=self.calc_rho(self.x, self.y, self.z), + theta=self.calc_theta(self.x, self.z), + phi=self.calc_phi(self.x, self.y, self.z), + ) + + @property + def field_magnitude(self) -> float: + return math.hypot(self.x, self.y, self.z) + + def to_request_pos(self) -> MagnetRequest: + return MagnetRequest(x=self.x, y=self.y, z=self.z) + + +@dataclass(kw_only=True) +class MagnetSphericalPosition: + rho: float + theta: float # degrees + phi: float # degrees + + def to_cartesian(self) -> MagnetPosition: + theta = math.radians(self.theta) + phi = math.radians(self.phi) + return MagnetPosition( + x=-self.rho * math.sin(phi) * math.sin(theta), + y=self.rho * math.cos(phi), + z=self.rho * math.sin(phi) * math.cos(theta), + ) + + def to_request_pos(self) -> MagnetSphericalRequest: + return MagnetSphericalRequest(rho=self.rho, theta=self.theta, phi=self.phi) + + +@dataclass(frozen=True, kw_only=True) +class MagnetRequest: + x: float | None = None + y: float | None = None + z: float | None = None + + def __post_init__(self): + if self.x is None and self.y is None and self.z is None: + raise ValueError("At least one of x, y, or z must be specified.") + + def resolve_pos(self, current: MagnetPosition) -> MagnetPosition: + """Resolve unspecified coordinates using the current position.""" + return MagnetPosition( + x=current.x if self.x is None else self.x, + y=current.y if self.y is None else self.y, + z=current.z if self.z is None else self.z, + ) + + +@dataclass(frozen=True, kw_only=True) +class MagnetSphericalRequest: + rho: float | None = None + theta: float | None = None + phi: float | None = None + + def __post_init__(self): + if self.rho is None and self.theta is None and self.phi is None: + raise ValueError("At least one of rho, theta, or phi must be specified.") + + def resolve_pos(self, current: MagnetSphericalPosition) -> MagnetSphericalPosition: + """Resolve unspecified coordinates using the current position.""" + return MagnetSphericalPosition( + rho=current.rho if self.rho is None else self.rho, + theta=current.theta if self.theta is None else self.theta, + phi=current.phi if self.phi is None else self.phi, + ) diff --git a/src/dodal/devices/beamlines/i06_1/magnet/enums.py b/src/dodal/devices/beamlines/i06_1/magnet/enums.py new file mode 100644 index 00000000000..dbd7505925b --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnet/enums.py @@ -0,0 +1,34 @@ +from ophyd_async.core import StrictEnum + + +class MagnetMode(StrictEnum): + UNIAXIAL_X = "UNIAXIAL_X" + UNIAXIAL_Y = "UNIAXIAL_Y" + UNIAXIAL_Z = "UNIAXIAL_Z" + SPHERICAL = "SPHERICAL" + PLANAR_XZ = "PLANAR_XZ" + QUADRANT_XY = "QUADRANT_XY" + CUBIC = "CUBIC" + UNDEFINED = "UNDEFINED" + + @property + def axis_alias(self) -> str: + match self: + case MagnetMode.UNIAXIAL_X: + return "x" + case MagnetMode.UNIAXIAL_Y: + return "y" + case MagnetMode.UNIAXIAL_Z: + return "z" + case _: + raise ValueError(f"{self} is not a valid uniaxial.") + + +class MagnetRampStatus(StrictEnum): + RAMP_MADE = "RAMP MADE" + RAMPING = "RAMPING" + + +class MagnetLimitStatus(StrictEnum): + OK = "OK" + VIOLATION = "VIOLATION" diff --git a/src/dodal/devices/beamlines/i06_1/magnet/movement.py b/src/dodal/devices/beamlines/i06_1/magnet/movement.py new file mode 100644 index 00000000000..050701c161e --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnet/movement.py @@ -0,0 +1,253 @@ +from abc import abstractmethod +from dataclasses import dataclass + +from dodal.devices.beamlines.i06_1.magnet.coordinates import ( + MagnetPosition, + MagnetRequest, +) +from dodal.devices.beamlines.i06_1.magnet.enums import MagnetMode + + +class MagnetPositionError(Exception): + @classmethod + def total_field_mag_outside_limit( + cls, mode: MagnetMode, limit: float, pos: MagnetPosition + ): + return cls( + f"Target field magnitude of {pos.field_magnitude}T exceeds " + f"limit of {limit} for mode {mode}. Requested position: {pos}." + ) + + @classmethod + def axis_below_limit(cls, mode: MagnetMode, limit: float, pos: float, axis: str): + return cls( + f"Axis {axis} with value {pos}T is below minimum {limit}T for mode {mode}." + ) + + @classmethod + def epics_hardware_limit_exceeded(cls, axis: str, pos: float, limit: float): + return cls( + f"Requested {axis} axis position {pos}T exceeds EPICS hardware limit " + f"({limit}T)." + ) + + +class MovementStrategy: + """Define the movement and safety constraints for a magnet operating mode. + + A movement strategy validates requested positions against the constraints of a + particular magnet mode and generates a sequence of safe movement steps to reach + the requested position. Each generated requested step is validated against the current + readback before it is applied. + """ + + def _check_epics_hardware_limits( + self, target_pos: MagnetPosition, limits: MagnetPosition + ) -> None: + """Utility method to validate target position against live EPICS PV limits.""" + if abs(target_pos.x) > limits.x: + raise MagnetPositionError.epics_hardware_limit_exceeded( + "x", target_pos.x, limits.x + ) + if abs(target_pos.y) > limits.y: + raise MagnetPositionError.epics_hardware_limit_exceeded( + "y", target_pos.y, limits.y + ) + if abs(target_pos.z) > limits.z: + raise MagnetPositionError.epics_hardware_limit_exceeded( + "z", target_pos.z, limits.z + ) + + @abstractmethod + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + """Validate that a requested position is within the mode's limits. + + Raises: + MagnetPositionError: If the requested position is invalid or outside + the limits of this movement strategy. + """ + + @abstractmethod + def move_steps( + self, current_readback: MagnetPosition, target: MagnetRequest + ) -> list[MagnetRequest]: + """Generate a sequence of safe movement steps to reach the target. + + The returned steps may be used to move through intermediate positions + required to safely reach the target from the current readback position. + """ + + +class SphericalMovement(MovementStrategy): + LIMIT = 1.75 + MODE = MagnetMode.SPHERICAL + + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + + mag_pos_after_move = target.resolve_pos(current_readback) + self._check_epics_hardware_limits(mag_pos_after_move, limits) + if mag_pos_after_move.field_magnitude > self.LIMIT: + raise MagnetPositionError.total_field_mag_outside_limit( + self.MODE, self.LIMIT, mag_pos_after_move + ) + + def move_steps( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + ) -> list[MagnetRequest]: + + steps: list[MagnetRequest] = [] + decreases = [ + ("z", current_readback.z, target.z), + ("x", current_readback.x, target.x), + ("y", current_readback.y, target.y), + ] + # Do all decreasing of axes first as one step. This ensures we don't + # quench the magnet. + decrease_kwargs = {} + increase_kwargs = {} + for axis, old, new in decreases: + if new is not None and old != new: + if abs(new) < abs(old): + decrease_kwargs[axis] = new + else: + increase_kwargs[axis] = new + + if decrease_kwargs: + steps.append(MagnetRequest(**decrease_kwargs)) + if increase_kwargs: + steps.append(MagnetRequest(**increase_kwargs)) + return steps + + +class CubicMovement(MovementStrategy): + MODE = MagnetMode.CUBIC + + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + target_pos = target.resolve_pos(current_readback) + self._check_epics_hardware_limits(target_pos, limits) + + def move_steps( + self, current_readback: MagnetPosition, target: MagnetRequest + ) -> list[MagnetRequest]: + return [target] + + +class PlanarXZMovement(MovementStrategy): + LIMIT = 2.0 + MODE = MagnetMode.PLANAR_XZ + + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + target_pos = target.resolve_pos(current_readback) + self._check_epics_hardware_limits(target_pos, limits) + + if target_pos.field_magnitude > self.LIMIT: + raise MagnetPositionError.total_field_mag_outside_limit( + self.MODE, self.LIMIT, target_pos + ) + + def move_steps( + self, current_readback: MagnetPosition, target: MagnetRequest + ) -> list[MagnetRequest]: + return [MagnetRequest(x=target.x, z=target.z)] + + +@dataclass +class UniaxialMovement(MovementStrategy): + mode: MagnetMode + + def _target_axes_map(self, target: MagnetRequest) -> dict[str, float | None]: + return { + "x": target.x, + "y": target.y, + "z": target.z, + } + + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + target_pos = target.resolve_pos(current_readback) + self._check_epics_hardware_limits(target_pos, limits) + + def move_steps( + self, current_readback: MagnetPosition, target: MagnetRequest + ) -> list[MagnetRequest]: + axis = self.mode.axis_alias + return [MagnetRequest(**{axis: self._target_axes_map(target)[axis]})] + + +class QuadrantXYMovement(MovementStrategy): + """Start from zero field in all axes. Appy 2.0 T in + y direction. Move field + direction through positive quadrant to 2.0 T in the +x direction. At all times + | B | <= 2.0 T. Ramp field back to zero in x direction. + """ + + LIMIT = 2.0 + MODE = MagnetMode.QUADRANT_XY + + def check_within_limits( + self, + current_readback: MagnetPosition, + target: MagnetRequest, + limits: MagnetPosition, + ) -> None: + target_pos = target.resolve_pos(current_readback) + self._check_epics_hardware_limits(target_pos, limits) + # Positive quadrant only as manual state`Move field direction through positive quadrant.` + if target_pos.x < 0: + raise MagnetPositionError.axis_below_limit( + self.MODE, 0.0, target_pos.x, "x" + ) + + if target_pos.y < 0: + raise MagnetPositionError.axis_below_limit( + self.MODE, 0.0, target_pos.y, "y" + ) + + if target_pos.field_magnitude > self.LIMIT: + raise MagnetPositionError.total_field_mag_outside_limit( + self.MODE, self.LIMIT, target_pos + ) + + def move_steps( + self, current_readback: MagnetPosition, target: MagnetRequest + ) -> list[MagnetRequest]: + steps: list[MagnetRequest] = [] + + # If increasing Y while X is non-zero could exceed the field limit, + # reduce X before increasing Y. + if ( + target.y is not None + and target.y > current_readback.y + and current_readback.x > 0 + and MagnetPosition(x=current_readback.x, y=target.y, z=0).field_magnitude + >= self.LIMIT + ): + steps.append(MagnetRequest(x=0)) + + steps.append(MagnetRequest(x=target.x, y=target.y, z=target.z)) + return steps diff --git a/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py b/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py new file mode 100644 index 00000000000..3a15ff31210 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnet/superconducting_magnet.py @@ -0,0 +1,596 @@ +import asyncio +import dataclasses +import math +from typing import Protocol + +from bluesky.protocols import Flyable, Movable, Preparable +from ophyd_async.core import ( + DEFAULT_TIMEOUT, + AsyncStatus, + ConfinedModel, + DeviceMock, + Reference, + SignalRW, + StandardReadable, + StandardReadableFormat, + callback_on_mock_execute, + callback_on_mock_put, + default_mock_class, + derived_signal_rw, + error_if_none, + set_mock_value, + wait_for_value, +) +from ophyd_async.epics.core import ( + epics_signal_r, + epics_signal_rw, + epics_triggerable_command, +) +from pydantic import Field + +from dodal.devices.beamlines.i06_1.magnet.coordinates import ( + MagnetPosition, + MagnetRequest, + MagnetSphericalPosition, + MagnetSphericalRequest, +) +from dodal.devices.beamlines.i06_1.magnet.enums import ( + MagnetLimitStatus, + MagnetMode, + MagnetRampStatus, +) +from dodal.devices.beamlines.i06_1.magnet.movement import ( + CubicMovement, + MagnetPositionError, + MovementStrategy, + PlanarXZMovement, + QuadrantXYMovement, + SphericalMovement, + UniaxialMovement, +) + + +class FlyVectorMagnetInfo(ConfinedModel): + fly_axis: str = Field(frozen=True) + """Axis to fly along, one of 'x', 'y' or 'z'. Todo theta, phi or rho.""" + start_position: float = Field(frozen=True) + """Start position of the magnet move. in Tesla""" + + end_position: float = Field(frozen=True) + """End position of the magnet move, in Tesla.""" + + ramp_rate: float = Field(frozen=True, gt=0) + """Ramp rate of the magnet move, in Tesla/s.""" + + +class MagnetMoveWithinBoundary(Protocol): + async def __call__(self, value: MagnetRequest): ... + + +@default_mock_class(DeviceMock) +# Don't want to sync readback and setpoint. IOC behaviour will move readback to demand +# once start_ramp is triggered. This logic lives in the SuperConductingMagnetController. +class MagnetAxis(StandardReadable): + """Represents one cartesian axis of a superconducting vector magnet. + + Although each axis can be moved independently, all moves are delegated to + the parent :class:`SuperConductingMagnetController` so that coordinated motion can + enforce safe field transitions before updating the underlying demand PVs. + + Each axis is associated with a :class:`MagnetAxisRampRateController`, which is used + to configure the axis ramp rate during preparation for fly scans. + """ + + def __init__( + self, + prefix: str, + axis: str, + axis_number: str, + name: str = "", + ): + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.readback = epics_signal_r(float, f"{prefix}-MAG-01:{axis}:RBV") + self.demand = epics_signal_rw(float, f"{prefix}-MAG-01:{axis}:DMD") + self.limit = epics_signal_r(float, f"{prefix}-SMC-0{axis_number}:LIM:FIELD:NOW") + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.ramp_rate = epics_signal_rw( + float, + read_pv=f"{prefix}-SMC-0{axis_number}:STS:RAMPRATE:TPM", + write_pv=f"{prefix}-SMC-0{axis_number}:SET:DMD:RAMPRATE:TPM", + ) + self.ramp_limit = epics_signal_r( + float, f"{prefix}-SMC-0{axis_number}:LIM:RAMPRATE:TPM" + ) + super().__init__(name) + + +class MockSuperConductingMagnetController( + DeviceMock["SuperConductingMagnetController"] +): + """Add additional callback logic to our device to get the mock behaviour to simulate + the hardware as best we can. + """ + + async def connect(self, device: "SuperConductingMagnetController"): + + async def _trigger_start_ramp(): + # Whenever ramp is triggered for the ioc, readback values move to the + # demand values. Simulate this behaviour here. + x_d, y_d, z_d = await asyncio.gather( + device.x.demand.get_value(), + device.y.demand.get_value(), + device.z.demand.get_value(), + ) + set_mock_value(device.ramp_status, MagnetRampStatus.RAMPING) + set_mock_value(device.x.readback, x_d) + set_mock_value(device.y.readback, y_d) + set_mock_value(device.z.readback, z_d) + set_mock_value(device.ramp_status, MagnetRampStatus.RAMP_MADE) + + callback_on_mock_execute(device._start_ramp, _trigger_start_ramp) # noqa: SLF001 + + async def _set_mode(value): + # Whenever mode is set, ioc automatically sets everything to zero and + # triggers a ramp. + set_mock_value(device.x.demand, 0.0) + set_mock_value(device.y.demand, 0.0) + set_mock_value(device.z.demand, 0.0) + await _trigger_start_ramp() + + callback_on_mock_put(device.mode, _set_mode) + set_mock_value(device.limit_status, MagnetLimitStatus.OK) + set_mock_value(device.ramp_status, MagnetRampStatus.RAMP_MADE) + set_mock_value(device.x.limit, 2.0) + set_mock_value(device.y.limit, 2.0) + set_mock_value(device.z.limit, 6.0) + set_mock_value(device.x.ramp_limit, 0.5) + set_mock_value(device.y.ramp_limit, 0.25) + set_mock_value(device.z.ramp_limit, 1.0) + set_mock_value(device.x.ramp_rate, 0.5) + set_mock_value(device.y.ramp_rate, 0.25) + set_mock_value(device.z.ramp_rate, 1.0) + + +@default_mock_class(MockSuperConductingMagnetController) +class SuperConductingMagnetController(StandardReadable, Flyable, Preparable): + """A three-axis superconducting vector magnet. + + The magnet provides both cartesian and spherical coordinate interfaces for + reading and setting the magnetic field. Cartesian coordinates expose the + physical X, Y and Z magnet axes, while spherical coordinates provide the + equivalent field magnitude and orientation. + + The permitted operating region and the sequence of axis movements are + determined by the current magnet operating mode. Each operating mode uses a + dedicated movement strategy to validate requested positions and generate a + safe sequence of cartesian magnet moves before ramping the field. + """ + + _MODE_MOVEMENT_STRATEGY: dict[MagnetMode, MovementStrategy] = { + MagnetMode.SPHERICAL: SphericalMovement(), + MagnetMode.CUBIC: CubicMovement(), + MagnetMode.PLANAR_XZ: PlanarXZMovement(), + MagnetMode.QUADRANT_XY: QuadrantXYMovement(), + MagnetMode.UNIAXIAL_X: UniaxialMovement(MagnetMode.UNIAXIAL_X), + MagnetMode.UNIAXIAL_Y: UniaxialMovement(MagnetMode.UNIAXIAL_Y), + MagnetMode.UNIAXIAL_Z: UniaxialMovement(MagnetMode.UNIAXIAL_Z), + } + + def __init__( + self, + prefix: str, + name: str = "", + ): + with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL): + self.mode = epics_signal_rw(MagnetMode, prefix + "-MAG-01:MODE") + self.x = MagnetAxis(prefix, axis="X", axis_number="1") + self.y = MagnetAxis(prefix, axis="Y", axis_number="2") + self.z = MagnetAxis(prefix, axis="Z", axis_number="3") + + self.ramp_status = epics_signal_rw( + MagnetRampStatus, prefix + "-MAG-01:RAMPSTATUS" + ) + self.limit_status = epics_signal_rw( + MagnetLimitStatus, prefix + "-MAG-01:LIMITSTATUS" + ) + # Make private so cannot trigger without going through magnet API. + self._start_ramp = epics_triggerable_command(prefix + "-MAG-01:STARTRAMP.PROC") + # Used to block parallel moves that are not submitted together. Allows us to + # always be sure we safely move the magnet using coordinated logic. + self._move_lock = asyncio.Lock() + self._fly_info: FlyVectorMagnetInfo | None = None + self._fly_status: AsyncStatus | None = None + + super().__init__(name) + + async def _trigger_ramp(self, timeout: float | None = None): + # Setting invalid demand values should put the limit status in violation state. + # Block ramp if in violation state. + limit_status = await self.limit_status.get_value() + if limit_status == MagnetLimitStatus.VIOLATION: + raise MagnetPositionError(f"{self.limit_status.name} is at {limit_status}") + self.log.info("About to start ramping the magnet.") + await self._start_ramp.trigger() + await wait_for_value( + self.ramp_status, MagnetRampStatus.RAMP_MADE, timeout=timeout + ) + self.log.info( + f"Ramping complete. Ramp status is now {MagnetRampStatus.RAMP_MADE}" + ) + + async def _apply_step( + self, step: MagnetRequest, timeout: float | None = None + ) -> None: + """Apply a single movement step and wait for the magnet ramp to complete. + + A movement step may update one or more cartesian axes simultaneously. Once + the requested demand values have been written, the magnet is instructed to + ramp and this method waits for the ramp to complete before returning. + """ + tasks = [] + if step.x is not None: + tasks.append(self.x.demand.set(step.x)) + if step.y is not None: + tasks.append(self.y.demand.set(step.y)) + if step.z is not None: + tasks.append(self.z.demand.set(step.z)) + self.log.info(f"About to set demand values of the magnet to {step}.") + await asyncio.gather(*tasks) + await self._trigger_ramp(timeout=timeout) + + @AsyncStatus.wrap + async def set_within_boundary(self, value: MagnetRequest) -> None: + """Move the magnet to a new cartesian position while respecting mode limits. + + Any coordinates left as ``None`` retain their current values. The requested + target position is first validated against the limits of the current magnet + operating mode. The corresponding movement strategy then generates a sequence + of intermediate movement steps. + + Each movement step is validated against the latest magnet readback before it + is applied. The readback position is updated after each step, ensuring that + subsequent steps are validated against the magnet's actual current position. + """ + if self._move_lock.locked(): + raise RuntimeError(f"Cannot do move for {self.name}. It is already moving.") + + async with self._move_lock: + current_readback, mode, field_limit = await asyncio.gather( + self.get_readback_position(), + self.mode.get_value(), + self.get_field_limit(), + ) + movement_strategy = self._MODE_MOVEMENT_STRATEGY.get(mode) + if movement_strategy is None: + raise ValueError( + f"No movement strategy has been configured for device {self.name} for mode {mode}." + ) + self.log.debug( + f"Attempting move in mode {mode} with parameters {value}. " + f"Current readback is {current_readback}." + ) + # Check final requested position is within limits. + movement_strategy.check_within_limits(current_readback, value, field_limit) + for step in movement_strategy.move_steps(current_readback, value): + # Check each generated step with the current readback is within limits. + movement_strategy.check_within_limits( + current_readback, step, field_limit + ) + timeout = await self.calculate_timeout_per_step(current_readback, step) + await self._apply_step(step, timeout=timeout) + self.log.debug( + f"Applying movement step {step}. Current readback is " + f"{current_readback} with timeout {timeout}." + ) + current_readback = await self.get_readback_position() + + async def calculate_timeout_per_step( + self, current_readback: MagnetPosition, value: MagnetRequest + ) -> float: + """Calculate the timeout for a move based on the distance to move and the ramp rate.""" + ramp_rate = await self.get_ramp_rate() + dx = (value.x - current_readback.x) if value.x is not None else 0.0 + dy = (value.y - current_readback.y) if value.y is not None else 0.0 + dz = (value.z - current_readback.z) if value.z is not None else 0.0 + distance = math.sqrt(dx**2 + dy**2 + dz**2) + combined_ramp_rate = max(ramp_rate.x, ramp_rate.y, ramp_rate.z) + return distance / combined_ramp_rate + DEFAULT_TIMEOUT + + async def get_readback_position(self) -> MagnetPosition: + """Get the current magnet readback position as a :class:`MagnetPosition`.""" + x, y, z = await asyncio.gather( + self.x.readback.get_value(), + self.y.readback.get_value(), + self.z.readback.get_value(), + ) + return MagnetPosition(x=x, y=y, z=z) + + async def get_field_limit(self) -> MagnetPosition: + """Get the current field limit of the magnet. + + The field limit is a live readback from the IOC and may change depending on + the current magnet operating mode and the current demand values. + """ + x, y, z = await asyncio.gather( + self.x.limit.get_value(), + self.y.limit.get_value(), + self.z.limit.get_value(), + ) + return MagnetPosition(x=x, y=y, z=z) + + async def get_ramp_rate_limit(self) -> MagnetPosition: + """Get the current ramp rate limit of the magnet. + + The ramp rate limit is a live readback from the IOC and may change depending on + the current magnet operating mode and the current demand values. + """ + x, y, z = await asyncio.gather( + self.x.ramp_limit.get_value(), + self.y.ramp_limit.get_value(), + self.z.ramp_limit.get_value(), + ) + return MagnetPosition(x=x, y=y, z=z) + + async def get_ramp_rate(self) -> MagnetPosition: + """Get the current ramp rate of the magnet. + + The ramp rate is a live readback from the IOC and may change depending on + the current magnet operating mode and the current demand values. + """ + x, y, z = await asyncio.gather( + self.x.ramp_rate.get_value(), + self.y.ramp_rate.get_value(), + self.z.ramp_rate.get_value(), + ) + return MagnetPosition(x=x, y=y, z=z) + + @AsyncStatus.wrap + async def prepare(self, value: FlyVectorMagnetInfo) -> None: + """Prepare the magnet for a fly scan. + The requested target position is validated against the limits of the current + magnet operating mode. The corresponding movement strategy then generates a + sequence of intermediate movement steps. + Each movement step is validated against the latest magnet readback before it + is applied. The readback position is updated after each step, ensuring that + subsequent steps are validated against the magnet's actual current position. + """ + self._fly_info = value + current_readback, mode, field_limit, ramp_rate_limit = await asyncio.gather( + self.get_readback_position(), + self.mode.get_value(), + self.get_field_limit(), + self.get_ramp_rate_limit(), + ) + if mode in { + MagnetMode.UNIAXIAL_X, + MagnetMode.UNIAXIAL_Y, + MagnetMode.UNIAXIAL_Z, + MagnetMode.CUBIC, + }: + movement_strategy = self._MODE_MOVEMENT_STRATEGY.get(mode) + max_ramp_rate: float = getattr(ramp_rate_limit, value.fly_axis) + if movement_strategy is None: + raise ValueError( + f"No movement strategy has been configured for device {self.name} for mode {mode}." + ) + if max_ramp_rate < value.ramp_rate or value.ramp_rate <= 0: + raise ValueError( + f"Requested ramp rate {value.ramp_rate} exceeds the ramp rate limit" + f" of {max_ramp_rate} for axis {value.fly_axis}." + ) + + start_req = MagnetRequest(**{value.fly_axis: value.start_position}) + end_req = MagnetRequest(**{value.fly_axis: value.end_position}) + + movement_strategy.check_within_limits( + current_readback, start_req, field_limit + ) + + start_readback = dataclasses.replace( + current_readback, **{value.fly_axis: value.start_position} + ) + + movement_strategy.check_within_limits(start_readback, end_req, field_limit) + + else: + raise ValueError( + f"Cannot prepare fly scan in mode {mode}. Only uniaxial and cubic modes are supported for now." + ) + start_position = MagnetRequest(**{value.fly_axis: value.start_position}) + await self.set_within_boundary(start_position) + fly_axis = getattr(self, value.fly_axis) + await fly_axis.ramp_rate.set(value.ramp_rate) + + @AsyncStatus.wrap + async def kickoff(self): + fly_info = error_if_none( + self._fly_info, + f"{self.name} must be prepared before attempting to kickoff.", + ) + + self.log.info( + f"Starting fly scan along axis {fly_info.fly_axis} from " + f"{fly_info.start_position} to {fly_info.end_position} " + f"at ramp rate {fly_info.ramp_rate}." + ) + + self._fly_info = None + self._fly_status = self.set_within_boundary( + MagnetRequest(**{fly_info.fly_axis: fly_info.end_position}) + ) + + def complete(self) -> AsyncStatus: + fly_status = error_if_none( + self._fly_status, f"kickoff for magnet {self.name} not called." + ) + self._fly_status = None + return fly_status + + +class MagnetCartesianCoordinates(StandardReadable, Movable[MagnetPosition]): + """Cartesian interface to the superconducting magnet. + + Exposes the physical X, Y and Z magnet axes as individual movable signals, + together with a grouped ``MagnetPosition`` interface for moving all three axes + simultaneously. + + Individual axis moves and grouped cartesian moves are delegated to the parent + ``SuperConductingMagnetController``, which applies the active movement strategy for the + current operating mode before commanding the hardware. + """ + + def __init__( + self, + controller: SuperConductingMagnetController, + name: str = "", + ) -> None: + self._controller_ref = Reference(controller) + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.x = derived_signal_rw( + raw_to_derived=self._read_position, + set_derived=self._set_x, + val=self._controller_ref().x.readback, + ) + self.y = derived_signal_rw( + raw_to_derived=self._read_position, + set_derived=self._set_y, + val=self._controller_ref().y.readback, + ) + self.z = derived_signal_rw( + raw_to_derived=self._read_position, + set_derived=self._set_z, + val=self._controller_ref().z.readback, + ) + + super().__init__(name) + + def _read_position(self, val: float) -> float: + """Read the current magnet position as a :class:`MagnetPosition`.""" + return val + + @AsyncStatus.wrap + async def set(self, value: MagnetPosition | MagnetRequest): + if isinstance(value, MagnetPosition): + value = MagnetRequest(x=value.x, y=value.y, z=value.z) + await self._controller_ref().set_within_boundary(value) + + async def _set_x(self, x: float): + await self.set(MagnetRequest(x=x)) + + async def _set_y(self, y: float): + await self.set(MagnetRequest(y=y)) + + async def _set_z(self, z: float): + await self.set(MagnetRequest(z=z)) + + +class MagnetSphericalCoordinates(StandardReadable, Movable[MagnetSphericalPosition]): + """Spherical coordinate interface to the superconducting magnet. + + Exposes the magnet field in spherical coordinates using the derived signals + ``rho``, ``theta`` and ``phi``. These signals are calculated from the + underlying cartesian magnet axes and may also be written individually or as a + complete ``MagnetSphericalRequest``. + + Writes are converted to cartesian coordinates before being delegated to the + parent ``SuperConductingMagnetController``, allowing the active movement strategy to + determine a safe sequence of cartesian moves. + """ + + def __init__( + self, + controller: SuperConductingMagnetController, + name: str = "", + ): + self._controller_ref = Reference(controller) + + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.rho = derived_signal_rw( + raw_to_derived=self._get_spherical_component, + set_derived=self._set_rho, + component="rho", + x=self._controller_ref().x.readback, + y=self._controller_ref().y.readback, + z=self._controller_ref().z.readback, + ) + self.theta = derived_signal_rw( + raw_to_derived=self._get_spherical_component, + set_derived=self._set_theta, + component="theta", + x=self._controller_ref().x.readback, + y=self._controller_ref().y.readback, + z=self._controller_ref().z.readback, + ) + self.phi = derived_signal_rw( + raw_to_derived=self._get_spherical_component, + set_derived=self._set_phi, + component="phi", + x=self._controller_ref().x.readback, + y=self._controller_ref().y.readback, + z=self._controller_ref().z.readback, + ) + super().__init__(name) + + def _get_spherical_component( + self, component: str, x: float, y: float, z: float + ) -> float: + cart_pos = MagnetPosition( + x=x, + y=y, + z=z, + ) + return getattr(cart_pos.to_spherical(), component) + + @AsyncStatus.wrap + async def set(self, value: MagnetSphericalRequest): + """Set the requested spherical coordinates. + + Any unspecified spherical coordinates are taken from the current magnet + readback position. The resulting complete spherical position is converted + to Cartesian coordinates and passed to the movement controller, which + determines a safe sequence of Cartesian moves for the active magnet mode. + """ + current_readback = await self._controller_ref().get_readback_position() + target = value.resolve_pos(current_readback.to_spherical()) + await self._controller_ref().set_within_boundary( + target.to_cartesian().to_request_pos() + ) + + async def _set_rho(self, rho: float): + await self.set(MagnetSphericalRequest(rho=rho)) + + async def _set_theta(self, theta: float): + await self.set(MagnetSphericalRequest(theta=theta)) + + async def _set_phi(self, phi: float): + await self.set(MagnetSphericalRequest(phi=phi)) + + +class SuperConductingMagnet(StandardReadable, Preparable, Flyable): + """Factory function to create a superconducting magnet device with both cartesian and + spherical coordinate interfaces. + """ + + def __init__(self, controller: SuperConductingMagnetController, name: str = ""): + + with self.add_children_as_readables(): + self.cart = MagnetCartesianCoordinates(controller) + self.sph = MagnetSphericalCoordinates(controller) + self.controller = controller + + super().__init__(name) + + @property + def mode(self) -> SignalRW[MagnetMode]: + return self.controller.mode + + @AsyncStatus.wrap + async def prepare(self, value: FlyVectorMagnetInfo) -> None: + await self.controller.prepare(value=value) + + @AsyncStatus.wrap + async def kickoff(self) -> None: + await self.controller.kickoff() + + def complete(self) -> AsyncStatus: + return self.controller.complete() diff --git a/tests/devices/beamlines/i06_1/magnet/__init__.py b/tests/devices/beamlines/i06_1/magnet/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/devices/beamlines/i06_1/magnet/test_coordinates.py b/tests/devices/beamlines/i06_1/magnet/test_coordinates.py new file mode 100644 index 00000000000..2e923df54d5 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnet/test_coordinates.py @@ -0,0 +1,51 @@ +import pytest + +from dodal.devices.beamlines.i06_1.magnet.coordinates import ( + MagnetPosition, + MagnetRequest, + MagnetSphericalPosition, + MagnetSphericalRequest, +) +from tests.devices.beamlines.i06_1.magnet.utils import ( + EXPECTED_CARTESIAN_SPHERICAL_CONVERSION, +) + + +def test_cartesian_and_spherical_are_inverse() -> None: + cartesian = MagnetPosition(x=10, y=20, z=30) + result = cartesian.to_spherical().to_cartesian() + assert result.x == pytest.approx(cartesian.x) + assert result.y == pytest.approx(cartesian.y) + assert result.z == pytest.approx(cartesian.z) + + +def test_spherical_and_cartesian_are_inverse() -> None: + spherical = MagnetSphericalPosition(rho=10, theta=20, phi=30) + result = spherical.to_cartesian().to_spherical() + assert result.rho == pytest.approx(spherical.rho) + assert result.theta == pytest.approx(spherical.theta) + assert result.phi == pytest.approx(spherical.phi) + + +@pytest.mark.parametrize( + "cartesian, spherical", EXPECTED_CARTESIAN_SPHERICAL_CONVERSION +) +def test_cartesian_and_spherical_conversion_is_correct( + cartesian: MagnetPosition, spherical: MagnetSphericalPosition +) -> None: + result = cartesian.to_spherical() + assert result.rho == pytest.approx(spherical.rho) + assert result.theta == pytest.approx(spherical.theta) + assert result.phi == pytest.approx(spherical.phi) + + result = spherical.to_cartesian() + assert result.x == pytest.approx(cartesian.x) + assert result.y == pytest.approx(cartesian.y) + assert result.z == pytest.approx(cartesian.z) + + +def test_all_none_values_for_magnet_position_request_raises_error(): + with pytest.raises(ValueError): + MagnetRequest() + with pytest.raises(ValueError): + MagnetSphericalRequest() diff --git a/tests/devices/beamlines/i06_1/magnet/test_enums.py b/tests/devices/beamlines/i06_1/magnet/test_enums.py new file mode 100644 index 00000000000..2a0b591e7fd --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnet/test_enums.py @@ -0,0 +1,30 @@ +import pytest + +from dodal.devices.beamlines.i06_1.magnet import MagnetMode + + +@pytest.mark.parametrize( + "mode, expected_alias", + [ + (MagnetMode.UNIAXIAL_X, "x"), + (MagnetMode.UNIAXIAL_Y, "y"), + (MagnetMode.UNIAXIAL_Z, "z"), + ], +) +def test_magnet_mode_axis_alias(mode: MagnetMode, expected_alias: str) -> None: + assert mode.axis_alias == expected_alias + + +@pytest.mark.parametrize( + "mode", + [ + MagnetMode.CUBIC, + MagnetMode.PLANAR_XZ, + MagnetMode.QUADRANT_XY, + MagnetMode.SPHERICAL, + MagnetMode.UNDEFINED, + ], +) +def test_magnet_mode_axis_alias_raises_error_when_not_axis(mode: MagnetMode) -> None: + with pytest.raises(ValueError): + mode.axis_alias # noqa: B018 diff --git a/tests/devices/beamlines/i06_1/magnet/test_movements.py b/tests/devices/beamlines/i06_1/magnet/test_movements.py new file mode 100644 index 00000000000..62abc67c4b2 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnet/test_movements.py @@ -0,0 +1,319 @@ +import pytest + +from dodal.devices.beamlines.i06_1.magnet.coordinates import ( + MagnetPosition, + MagnetRequest, + MagnetSphericalPosition, + MagnetSphericalRequest, +) +from dodal.devices.beamlines.i06_1.magnet.enums import MagnetMode +from dodal.devices.beamlines.i06_1.magnet.movement import ( + CubicMovement, + MagnetPositionError, + PlanarXZMovement, + QuadrantXYMovement, + SphericalMovement, + UniaxialMovement, +) +from tests.devices.beamlines.i06_1.magnet.utils import ( + EXPECTED_CARTESIAN_SPHERICAL_CONVERSION, +) + + +def test_magnet_position_error_total_field_mag_outside_limit(): + position = MagnetPosition(x=1.0, y=2.0, z=3.0) + error = MagnetPositionError.total_field_mag_outside_limit( + MagnetMode.SPHERICAL, 1.75, position + ) + assert isinstance(error, MagnetPositionError) + assert str(error) == ( + f"Target field magnitude of {position.field_magnitude}T exceeds " + f"limit of 1.75 for mode {MagnetMode.SPHERICAL}. Requested position: {position}." + ) + + +def test_axis_below_limit() -> None: + error = MagnetPositionError.axis_below_limit(MagnetMode.QUADRANT_XY, 0.0, -1.0, "x") + assert str(error) == ( + "Axis x with value -1.0T is below minimum 0.0T for mode QUADRANT_XY." + ) + + +def test_cartesian_and_spherical_are_inverse() -> None: + cartesian = MagnetPosition(x=10, y=20, z=30) + result = cartesian.to_spherical().to_cartesian() + assert result.x == pytest.approx(cartesian.x) + assert result.y == pytest.approx(cartesian.y) + assert result.z == pytest.approx(cartesian.z) + + +def test_spherical_and_cartesian_are_inverse() -> None: + spherical = MagnetSphericalPosition(rho=10, theta=20, phi=30) + result = spherical.to_cartesian().to_spherical() + assert result.rho == pytest.approx(spherical.rho) + assert result.theta == pytest.approx(spherical.theta) + assert result.phi == pytest.approx(spherical.phi) + + +@pytest.mark.parametrize( + "cartesian, spherical", EXPECTED_CARTESIAN_SPHERICAL_CONVERSION +) +def test_cartesian_and_spherical_conversion_is_correct( + cartesian: MagnetPosition, spherical: MagnetSphericalPosition +) -> None: + result = cartesian.to_spherical() + assert result.rho == pytest.approx(spherical.rho) + assert result.theta == pytest.approx(spherical.theta) + assert result.phi == pytest.approx(spherical.phi) + + result = spherical.to_cartesian() + assert result.x == pytest.approx(cartesian.x) + assert result.y == pytest.approx(cartesian.y) + assert result.z == pytest.approx(cartesian.z) + + +def test_all_none_values_for_magnet_position_request_raises_error(): + with pytest.raises(ValueError): + MagnetRequest() + with pytest.raises(ValueError): + MagnetSphericalRequest() + + +@pytest.mark.parametrize( + "current, target, expected", + [ + # All decreases done as one move. + ( + MagnetPosition(x=3, y=3, z=3), + MagnetRequest(x=1, y=1, z=1), + [MagnetRequest(x=1, y=1, z=1)], + ), + # Check decreases done as one move and done before the increase + ( + MagnetPosition(x=3, y=3, z=3), + MagnetRequest(x=4, y=2, z=1), + [ + MagnetRequest(y=2, z=1), + MagnetRequest(x=4), + ], + ), + # Check if all an increase, all done together + ( + MagnetPosition(x=1, y=1, z=1), + MagnetRequest(x=2, y=2, z=2), + [ + MagnetRequest(x=2, y=2, z=2), + ], + ), + # Check requesting same position will do no movement + ( + MagnetPosition(x=1, y=2, z=3), + MagnetRequest(x=1, y=2, z=3), + [], + ), + ], +) +def test_spherical_movement( + current: MagnetPosition, + target: MagnetRequest, + expected: list[MagnetRequest], +) -> None: + assert SphericalMovement().move_steps(current, target) == expected + + +def test_spherical_movement_within_limit(): + current = MagnetPosition(x=0.0, y=0.0, z=0.0) + target = MagnetRequest(x=1.0, y=0.0, z=0.0) + hard_limits = MagnetPosition(x=2.0, y=2.0, z=2.0) + SphericalMovement().check_within_limits(current, target, hard_limits) + + +def test_spherical_movement_outside_limit_raises_error(): + current = MagnetPosition(x=0.0, y=0.0, z=0.0) + target = MagnetRequest(x=2.0, y=0.0, z=0.0) + hard_limits = MagnetPosition(x=2.0, y=2.0, z=2.0) + with pytest.raises(MagnetPositionError): + SphericalMovement().check_within_limits(current, target, hard_limits) + + +def test_spherical_movement_at_limit_is_allowed(): + current = MagnetPosition(x=0.0, y=0.0, z=0.0) + target = MagnetRequest(x=1.75, y=0.0, z=0.0) + hard_limits = MagnetPosition(x=2.0, y=2.0, z=2.0) + SphericalMovement().check_within_limits(current, target, hard_limits) + + +@pytest.mark.parametrize( + "target", + [ + MagnetRequest(x=1.6, y=0, z=0), + MagnetRequest(x=0, y=-1.6, z=0), + MagnetRequest(x=0, y=0, z=2), + ], +) +def test_cubic_rejects_outside_limits(target: MagnetRequest) -> None: + with pytest.raises(MagnetPositionError): + CubicMovement().check_within_limits( + MagnetPosition(x=0, y=0, z=0), target, MagnetPosition(x=1.5, y=1.5, z=1.5) + ) + + +def test_cubic_returns_single_step() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetRequest(x=1, y=1.2, z=-0.5) + assert CubicMovement().move_steps(current, target) == [ + MagnetRequest(x=1, y=1.2, z=-0.5) + ] + + +def test_planar_xz_returns_single_step() -> None: + target = MagnetRequest(x=1.2, y=0, z=-0.4) + assert PlanarXZMovement().move_steps(MagnetPosition(x=0, y=0, z=0), target) == [ + MagnetRequest(x=1.2, z=-0.4), + ] + + +def test_planar_xz_rejects_nonzero_y() -> None: + hard_limits = MagnetPosition(x=2.0, y=0.0, z=2.0) + with pytest.raises(MagnetPositionError): + PlanarXZMovement().check_within_limits( + MagnetPosition(x=0, y=0, z=0), MagnetRequest(x=1, y=1, z=0), hard_limits + ) + + +def test_planar_xz_rejects_outside_radius() -> None: + hard_limits = MagnetPosition(x=2.0, y=0.0, z=2.0) + with pytest.raises(MagnetPositionError): + PlanarXZMovement().check_within_limits( + MagnetPosition(x=0, y=0, z=0), MagnetRequest(x=2, y=0, z=2), hard_limits + ) + + +@pytest.mark.parametrize( + "mode, target", + [ + ( + MagnetMode.UNIAXIAL_X, + MagnetRequest(x=1, y=0, z=0), + ), + ( + MagnetMode.UNIAXIAL_Y, + MagnetRequest(x=0, y=-1, z=0), + ), + ( + MagnetMode.UNIAXIAL_Z, + MagnetRequest(x=0, y=0, z=3), + ), + ], +) +def test_uniaxial_returns_single_step(mode: MagnetMode, target: MagnetRequest) -> None: + assert UniaxialMovement(mode).move_steps(MagnetPosition(x=0, y=0, z=0), target) == [ + MagnetRequest(**{mode.axis_alias: getattr(target, mode.axis_alias)}) + ] + + +@pytest.mark.parametrize( + "mode, hard_limits, target", + [ + ( + MagnetMode.UNIAXIAL_X, + MagnetPosition(x=6, y=0, z=0), + MagnetRequest(x=1, y=1, z=0), + ), + ( + MagnetMode.UNIAXIAL_X, + MagnetPosition(x=6, y=0, z=0), + MagnetRequest(x=1, y=0, z=1), + ), + ( + MagnetMode.UNIAXIAL_Y, + MagnetPosition(x=1, y=1, z=0), + MagnetRequest(x=0, y=2, z=0), + ), + ( + MagnetMode.UNIAXIAL_Y, + MagnetPosition(x=0, y=2, z=0), + MagnetRequest(x=0, y=1, z=1), + ), + ( + MagnetMode.UNIAXIAL_Z, + MagnetPosition(x=0, y=0, z=2), + MagnetRequest(x=1, y=0, z=1), + ), + ( + MagnetMode.UNIAXIAL_Z, + MagnetPosition(x=0, y=0, z=2), + MagnetRequest(x=0, y=1, z=1), + ), + ], +) +def test_uniaxial_rejects_other_axes( + mode: MagnetMode, hard_limits: MagnetPosition, target: MagnetRequest +) -> None: + with pytest.raises(MagnetPositionError): + UniaxialMovement(mode).check_within_limits( + MagnetPosition(x=0, y=0, z=0), target, hard_limits + ) + + +def test_uniaxial_raise_error_when_above_limit() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetRequest(x=10, y=0, z=0) + hard_limits = MagnetPosition(x=5, y=0, z=0) + with pytest.raises(MagnetPositionError): + UniaxialMovement(MagnetMode.UNIAXIAL_X).check_within_limits( + current, target, hard_limits + ) + + +def test_quadrant_xy_sequence() -> None: + current = MagnetPosition(x=1, y=0.5, z=0) + target = MagnetRequest(x=1.5, y=1, z=0) + assert QuadrantXYMovement().move_steps(current, target) == [ + MagnetRequest(x=1.5, y=1, z=0) + ] + + +def test_quadrant_xy_reduces_x_before_increasing_y() -> None: + current = MagnetPosition(x=1.8, y=0, z=0) + target = MagnetRequest(x=1.8, y=1) + assert QuadrantXYMovement().move_steps(current, target) == [ + MagnetRequest(x=0), + MagnetRequest(x=1.8, y=1), + ] + + +def test_quadrant_xy_moves_x_and_y_together_from_zero() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetRequest(x=1, y=1) + assert QuadrantXYMovement().move_steps(current, target) == [MagnetRequest(x=1, y=1)] + + +def test_quadrant_xy_moves_x_when_y_is_unchanged() -> None: + current = MagnetPosition(x=0.5, y=1.5, z=0) + target = MagnetRequest(x=1) + assert QuadrantXYMovement().move_steps(current, target) == [MagnetRequest(x=1)] + + +def test_quadrant_xy_increases_y_without_zeroing_x_when_safe() -> None: + current = MagnetPosition(x=1, y=0.5, z=0) + target = MagnetRequest(x=1, y=1) + assert QuadrantXYMovement().move_steps(current, target) == [MagnetRequest(x=1, y=1)] + + +@pytest.mark.parametrize( + "target", [MagnetRequest(x=-1), MagnetRequest(y=-1), MagnetRequest(z=1)] +) +def test_quadrant_xy_rejects_invalid_target(target: MagnetRequest) -> None: + current = MagnetPosition(x=0, y=0, z=0) + hard_limits = MagnetPosition(x=1, y=1, z=0) + with pytest.raises(MagnetPositionError): + QuadrantXYMovement().check_within_limits(current, target, hard_limits) + + +def test_quadrant_xy_rejects_field_above_limit() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetRequest(x=1.5, y=1.5) + hard_limits = MagnetPosition(x=2, y=2, z=0) + with pytest.raises(MagnetPositionError): + QuadrantXYMovement().check_within_limits(current, target, hard_limits) diff --git a/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py b/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py new file mode 100644 index 00000000000..f54207a7f41 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnet/test_superconducting_magnet.py @@ -0,0 +1,522 @@ +import asyncio +import math +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch + +import pytest +from bluesky import FailedStatus, RunEngine +from bluesky.plan_stubs import mv +from bluesky.protocols import Reading +from ophyd_async.core import init_devices, set_mock_value +from ophyd_async.testing import assert_configuration, assert_reading, partial_reading + +from dodal.devices.beamlines.i06_1.magnet import ( + FlyVectorMagnetInfo, + MagnetAxis, + MagnetLimitStatus, + MagnetMode, + MagnetPosition, + MagnetPositionError, + MagnetRampStatus, + MagnetRequest, + MagnetSphericalPosition, + SuperConductingMagnet, + SuperConductingMagnetController, + movement, +) +from dodal.devices.beamlines.i06_1.magnet.movement import MovementStrategy +from tests.devices.beamlines.i06_1.magnet.utils import ( + EXPECTED_CARTESIAN_SPHERICAL_CONVERSION, +) + + +@pytest.fixture +def scmc() -> SuperConductingMagnetController: + with init_devices(mock=True): + scmc = SuperConductingMagnetController("TEST:") + return scmc + + +@pytest.fixture +def scm(scmc: SuperConductingMagnetController) -> SuperConductingMagnet: + with init_devices(mock=True): + scm = SuperConductingMagnet(scmc) + return scm + + +async def test_scm_read(scm: SuperConductingMagnet) -> None: + + await assert_reading( + scm, + { + "scm-cart-x": partial_reading(0), + "scm-cart-y": partial_reading(0), + "scm-cart-z": partial_reading(0), + "scm-sph-theta": partial_reading(0), + "scm-sph-rho": partial_reading(0), + "scm-sph-phi": partial_reading(0), + }, + ) + + +async def test_scmc_configuration(scm: SuperConductingMagnet) -> None: + await assert_configuration( + scm, {"scm-controller-mode": partial_reading(MagnetMode.UNIAXIAL_X)} + ) + + +@pytest.mark.parametrize( + "axis, value, expected_x, expected_y, expected_z", + [ + ("x", 0.1, 0.1, 0, 0), + ("y", 0.5, 0, 0.5, 0), + ("z", 0.2, 0, 0, 0.2), + ], +) +async def test_scmc_axis_set_uses_correct_axis( + scm: SuperConductingMagnet, + axis: str, + value: float, + expected_x: float, + expected_y: float, + expected_z: float, +) -> None: + await scm.controller.mode.set(getattr(MagnetMode, "UNIAXIAL_" + axis.capitalize())) + await getattr(scm.cart, axis).set(value) + x, y, z = await asyncio.gather( + scm.cart.x.get_value(), + scm.cart.y.get_value(), + scm.cart.z.get_value(), + ) + assert x == expected_x + assert y == expected_y + assert z == expected_z + + +@pytest.mark.parametrize( + "cartesian, spherical", EXPECTED_CARTESIAN_SPHERICAL_CONVERSION +) +async def test_scmc_sph_set_using_spherical( + scm: SuperConductingMagnet, + cartesian: MagnetPosition, + spherical: MagnetSphericalPosition, +) -> None: + await scm.controller.mode.set(MagnetMode.SPHERICAL) + await scm.sph.set(spherical.to_request_pos()) + rho, theta, phi = await asyncio.gather( + scm.sph.rho.get_value(), + scm.sph.theta.get_value(), + scm.sph.phi.get_value(), + ) + assert rho == pytest.approx(spherical.rho) + assert theta == pytest.approx(spherical.theta) + assert phi == pytest.approx(spherical.phi) + + x, y, z = await asyncio.gather( + scm.cart.x.get_value(), + scm.cart.y.get_value(), + scm.cart.z.get_value(), + ) + assert x == pytest.approx(cartesian.x) + assert y == pytest.approx(cartesian.y) + assert z == pytest.approx(cartesian.z) + + +@pytest.mark.parametrize( + "cartesian, spherical", EXPECTED_CARTESIAN_SPHERICAL_CONVERSION +) +async def test_scmc_cart_set_using_cartesian( + scm: SuperConductingMagnet, + cartesian: MagnetPosition, + spherical: MagnetSphericalPosition, +) -> None: + await scm.controller.mode.set(MagnetMode.CUBIC) + await scm.cart.set(cartesian.to_request_pos()) + x, y, z = await asyncio.gather( + scm.cart.x.get_value(), + scm.cart.y.get_value(), + scm.cart.z.get_value(), + ) + assert x == pytest.approx(cartesian.x) + assert y == pytest.approx(cartesian.y) + assert z == pytest.approx(cartesian.z) + + rho, theta, phi = await asyncio.gather( + scm.sph.rho.get_value(), + scm.sph.theta.get_value(), + scm.sph.phi.get_value(), + ) + assert rho == pytest.approx(spherical.rho) + assert theta == pytest.approx(spherical.theta) + assert phi == pytest.approx(spherical.phi) + + +@pytest.mark.parametrize( + "axes, values, expected_rho, expected_phi, expected_theta, expected_x, expected_y, expected_z", + [ + (("rho",), (1,), 1, 0, 0, 0, 1, 0), + (("rho", "phi"), (1, 90), 1, 90, 0, 0, 0, 1), + (("rho", "phi", "theta"), (1, 45, 45), 1, 45, 45, -0.5, math.sqrt(2) / 2, 0.5), + ], +) +async def test_scmc_spherical_individual_axis_set( + scm: SuperConductingMagnet, + axes: tuple[str], + values: tuple[float], + expected_rho: float, + expected_phi: float, + expected_theta: float, + expected_x: float, + expected_y: float, + expected_z: float, +) -> None: + await scm.controller.mode.set(MagnetMode.SPHERICAL) + for axis, value in zip(axes, values, strict=True): + await getattr(scm.sph, axis).set(value) + + await assert_reading( + scm, + { + "scm-cart-x": partial_reading(expected_x), + "scm-cart-y": partial_reading(expected_y), + "scm-cart-z": partial_reading(expected_z), + "scm-sph-rho": partial_reading(expected_rho), + "scm-sph-phi": partial_reading(expected_phi), + "scm-sph-theta": partial_reading(expected_theta), + }, + ) + + +async def test_scmc_raises_error_if_limit_status_is_violation( + scm: SuperConductingMagnet, +) -> None: + await scm.mode.set(MagnetMode.UNIAXIAL_X) + set_mock_value(scm.controller.limit_status, MagnetLimitStatus.VIOLATION) + with pytest.raises(MagnetPositionError): + await scm.cart.x.set(1) + + +@pytest.mark.parametrize( + "mode", [MagnetMode.UNIAXIAL_X, MagnetMode.UNIAXIAL_Y, MagnetMode.UNIAXIAL_Z] +) +async def test_scmc_axis_with_ramp_rate_wired_correctly_with_prepare( + scm: SuperConductingMagnet, + mode: MagnetMode, +) -> None: + await scm.controller.mode.set(mode) + magnet_axis: MagnetAxis = getattr(scm.controller, mode.axis_alias) + + fly_info = FlyVectorMagnetInfo( + fly_axis=mode.axis_alias, start_position=0.2, end_position=1.2, ramp_rate=0.25 + ) + await scm.prepare(fly_info) + assert await magnet_axis.demand.get_value() == fly_info.start_position + assert await magnet_axis.ramp_rate.get_value() == fly_info.ramp_rate + + +async def test_scmc_axis_kickoff_and_complete( + scm: SuperConductingMagnet, +) -> None: + await scm.mode.set(MagnetMode.UNIAXIAL_X) + fly_info = FlyVectorMagnetInfo( + fly_axis=MagnetMode.UNIAXIAL_X.axis_alias, + start_position=1, + end_position=2, + ramp_rate=0.5, + ) + assert scm.controller._fly_info is None + await scm.prepare(fly_info) + assert scm.controller._fly_info is fly_info + assert scm.controller._fly_status is None + await scm.kickoff() + assert scm.controller._fly_info is None + await scm.complete() + assert await scm.cart.x.get_value() == fly_info.end_position + + +async def test_scmc_axis_kickoff_and_complete_raises_error_without_prepare( + scm: SuperConductingMagnet, +): + await scm.mode.set(MagnetMode.UNIAXIAL_X) + # Do multiple times to make sure you cannot kickoff or complete without preparing + # each time + for _ in range(3): + with pytest.raises(RuntimeError): + await scm.kickoff() + with pytest.raises(RuntimeError): + await scm.complete() + + fly_info = FlyVectorMagnetInfo( + fly_axis=MagnetMode.UNIAXIAL_X.axis_alias, + start_position=1, + end_position=2, + ramp_rate=0.5, + ) + await scm.prepare(fly_info) + await scm.kickoff() + await scm.complete() + + +async def test_scmc_mock_device_behaviour( + scm: SuperConductingMagnet, +) -> None: + assert await scm.controller.limit_status.get_value() == MagnetLimitStatus.OK + + set_mock_value(scm.controller.x.demand, 1) + set_mock_value(scm.controller.y.demand, 1) + set_mock_value(scm.controller.z.demand, 1) + + ramp_states: list[MagnetRampStatus] = [] + + def _ramp_state_callback(value: dict[str, Reading[MagnetRampStatus]]): + ramp_states.append(value[scm.controller.ramp_status.name]["value"]) + + scm.controller.ramp_status.subscribe_reading(_ramp_state_callback) + + # Set a mode value and check to see if all values are set back to zero and ramp + # status change happened + await scm.mode.set(MagnetMode.UNIAXIAL_X) + + x_d, y_d, z_d, x_rbv, y_rbv, z_rbv = await asyncio.gather( + scm.controller.x.demand.get_value(), + scm.controller.y.demand.get_value(), + scm.controller.z.demand.get_value(), + scm.cart.x.get_value(), + scm.cart.y.get_value(), + scm.cart.z.get_value(), + ) + assert x_d == y_d == z_d == x_rbv == y_rbv == z_rbv == 0 + assert ramp_states == [ + MagnetRampStatus.RAMP_MADE, + MagnetRampStatus.RAMPING, + MagnetRampStatus.RAMP_MADE, + ] + + +async def test_scmc_no_movement_strategy_for_mode( + scm: SuperConductingMagnet, +) -> None: + mode = await scm.controller.mode.get_value() + with patch.dict(scm.controller._MODE_MOVEMENT_STRATEGY, {}, clear=True): + with pytest.raises( + ValueError, + match=f"No movement strategy has been configured for device scm-controller for mode {mode}", + ): + await scm.cart.x.set(5) + + +@pytest.mark.parametrize( + "mode, expected", + [ + (MagnetMode.SPHERICAL, movement.SphericalMovement), + (MagnetMode.CUBIC, movement.CubicMovement), + (MagnetMode.PLANAR_XZ, movement.PlanarXZMovement), + (MagnetMode.QUADRANT_XY, movement.QuadrantXYMovement), + (MagnetMode.UNIAXIAL_X, movement.UniaxialMovement), + (MagnetMode.UNIAXIAL_Y, movement.UniaxialMovement), + (MagnetMode.UNIAXIAL_Z, movement.UniaxialMovement), + ], +) +async def test_scmc_magnet_mode_to_movement_strategy_configuration( + scm: SuperConductingMagnet, + mode: MagnetMode, + expected: type[movement.MovementStrategy], +) -> None: + assert isinstance(scm.controller._MODE_MOVEMENT_STRATEGY[mode], expected) + + +async def test_scmc_magnet_mode_to_uniaxial_movement_strategy_configuration( + scm: SuperConductingMagnet, +) -> None: + assert ( + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.UNIAXIAL_X].mode # type:ignore + == MagnetMode.UNIAXIAL_X + ) + assert ( + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.UNIAXIAL_Y].mode # type:ignore + == MagnetMode.UNIAXIAL_Y + ) + assert ( + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.UNIAXIAL_Z].mode # type:ignore + == MagnetMode.UNIAXIAL_Z + ) + + +async def test_scmc_executes_movement_stragey_and_ramp_at_each_step( + scm: SuperConductingMagnet, +) -> None: + + movement_strategy = MagicMock() + + mov_str_return_values = [ + movement.MagnetRequest(z=2), + movement.MagnetRequest(x=1, y=2, z=3), + ] + expected_apply_step_calls = [ + call(mov_str_return_values[0], timeout=ANY), + call(mov_str_return_values[1], timeout=ANY), + ] + movement_strategy.move_steps.return_value = mov_str_return_values + + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.UNIAXIAL_X] = movement_strategy + + scm.controller._trigger_ramp = AsyncMock() + + with patch( + "dodal.devices.beamlines.i06_1.magnet.superconducting_magnet.SuperConductingMagnetController._apply_step", + wraps=scm.controller._apply_step, + ) as mock_apply_step: + await scm.mode.set(MagnetMode.UNIAXIAL_X) + await scm.cart.x.set(4) + + assert mock_apply_step.call_args_list == expected_apply_step_calls + assert scm.controller._trigger_ramp.call_count == len(mov_str_return_values) + + +async def test_external_parallel_moves_for_scmc_raise_error( + scm: SuperConductingMagnet, run_engine: RunEngine +) -> None: + run_engine(mv(scm.mode, MagnetMode.CUBIC)) + # Coordinated parallel move of axes submitted together is okay. + run_engine(mv(scm.cart, MagnetPosition(x=1, y=1, z=1))) + + # Coordinated parallel move on axes done separately fails. + with pytest.raises(FailedStatus): + run_engine(mv(scm.cart.x, 0.5, scm.cart.y, 0.5)) + + # Check after a blocked move and the first move finished, we can still do another move. + run_engine(mv(scm.cart, MagnetPosition(x=0.1, y=0.2, z=0.3))) + + +async def test_scmc_set_within_boundary_checks_each_movement_step( + scm: SuperConductingMagnet, +) -> None: + movement_strategy = MagicMock(spec=MovementStrategy) + initial_position = MagnetPosition(x=0, y=0, z=0) + position_after_step_1 = MagnetPosition(x=0, y=1, z=0) + position_after_step_2 = MagnetPosition(x=1, y=1, z=0) + + target = MagnetRequest(x=1, y=1) + + steps = [MagnetRequest(y=1), MagnetRequest(x=1)] + movement_strategy.move_steps.return_value = steps + + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.QUADRANT_XY] = movement_strategy + + scm.controller.get_readback_position = AsyncMock( + side_effect=[ + initial_position, + position_after_step_1, + position_after_step_2, + ] + ) + filed_limit = await scm.controller.get_field_limit() + scm.controller.mode.get_value = AsyncMock(return_value=MagnetMode.QUADRANT_XY) + scm.controller._apply_step = AsyncMock() + + await scm.controller.set_within_boundary(target) + + assert movement_strategy.check_within_limits.call_args_list == [ + call(initial_position, target, filed_limit), + call(initial_position, steps[0], filed_limit), + call(position_after_step_1, steps[1], filed_limit), + ] + assert scm.controller._apply_step.call_args_list == [ + call(steps[0], timeout=ANY), + call(steps[1], timeout=ANY), + ] + assert scm.controller.get_readback_position.call_count == 3 + + +@pytest.mark.asyncio +async def test_scmc_set_within_boundary_stops_when_step_is_outside_limits( + scm: SuperConductingMagnet, +) -> None: + movement_strategy = MagicMock(spec=MovementStrategy) + initial_position = MagnetPosition(x=0, y=0, z=0) + target = MagnetRequest(x=1, y=1) + steps = [MagnetRequest(y=1), MagnetRequest(x=1)] + + movement_strategy.move_steps.return_value = steps + movement_strategy.check_within_limits.side_effect = [ + None, # Final target is valid. + None, # First step is valid. + MagnetPositionError.axis_below_limit(MagnetMode.QUADRANT_XY, 0, -1, "x"), + ] + scm.controller._MODE_MOVEMENT_STRATEGY[MagnetMode.QUADRANT_XY] = movement_strategy + scm.controller.get_readback_position = AsyncMock( + side_effect=[initial_position, MagnetPosition(x=0, y=1, z=0)] + ) + scm.controller.mode.get_value = AsyncMock(return_value=MagnetMode.QUADRANT_XY) + scm.controller._apply_step = AsyncMock() + + with pytest.raises(MagnetPositionError): + await scm.controller.set_within_boundary(target) + # The invalid second step must never be applied. + scm.controller._apply_step.assert_awaited_once_with(steps[0], timeout=ANY) + + +@pytest.mark.asyncio +async def test_prepare_raises_when_no_movement_strategy( + scm: SuperConductingMagnet, +): + + original_strategies = dict(scm.controller._MODE_MOVEMENT_STRATEGY) + scm.controller._MODE_MOVEMENT_STRATEGY.clear() + + fly_info = FlyVectorMagnetInfo( + fly_axis=MagnetMode.UNIAXIAL_X.axis_alias, + start_position=1, + end_position=2, + ramp_rate=0.5, + ) + with pytest.raises( + ValueError, + match="No movement strategy has been configured for device", + ): + await scm.prepare(fly_info) + scm.controller._MODE_MOVEMENT_STRATEGY.update(original_strategies) + + +async def test_prepare_raises_on_invalid_ramp_rate( + scm: SuperConductingMagnet, +): + invalid_ramp_rate = 15.0 + await scm.mode.set(MagnetMode.UNIAXIAL_X) + fly_info = FlyVectorMagnetInfo( + fly_axis="x", + start_position=1, + end_position=2, + ramp_rate=invalid_ramp_rate, + ) + + with pytest.raises( + ValueError, + match="Requested ramp rate", + ): + await scm.prepare(fly_info) + + +@pytest.mark.parametrize( + "mode", + [ + (MagnetMode.SPHERICAL), + (MagnetMode.PLANAR_XZ), + (MagnetMode.QUADRANT_XY), + ], +) +async def test_prepare_raises_on_invalid_mode( + scm: SuperConductingMagnet, mode: MagnetMode +): + await scm.mode.set(mode) + fly_info = FlyVectorMagnetInfo( + fly_axis="Y", + start_position=1, + end_position=2, + ramp_rate=0.2, + ) + + with pytest.raises( + ValueError, + match="Only uniaxial and cubic modes are supported", + ): + await scm.prepare(fly_info) diff --git a/tests/devices/beamlines/i06_1/magnet/utils.py b/tests/devices/beamlines/i06_1/magnet/utils.py new file mode 100644 index 00000000000..b6fd21b5d22 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnet/utils.py @@ -0,0 +1,32 @@ +import math + +from dodal.devices.beamlines.i06_1.magnet.coordinates import ( + MagnetPosition, + MagnetSphericalPosition, +) + +EXPECTED_CARTESIAN_SPHERICAL_CONVERSION = [ + (MagnetPosition(x=0, y=1, z=0), MagnetSphericalPosition(rho=1, theta=0, phi=0)), + (MagnetPosition(x=0, y=0, z=1), MagnetSphericalPosition(rho=1, theta=0, phi=90)), + (MagnetPosition(x=0, y=-1, z=0), MagnetSphericalPosition(rho=1, theta=0, phi=180)), + (MagnetPosition(x=-1, y=0, z=0), MagnetSphericalPosition(rho=1, theta=90, phi=90)), + (MagnetPosition(x=0, y=0, z=-1), MagnetSphericalPosition(rho=1, theta=180, phi=90)), + (MagnetPosition(x=1, y=0, z=0), MagnetSphericalPosition(rho=1, theta=270, phi=90)), + # 45 degree increments + ( + MagnetPosition(x=-math.sqrt(2) / 2, y=0, z=math.sqrt(2) / 2), + MagnetSphericalPosition(rho=1, theta=45, phi=90), + ), + ( + MagnetPosition(x=-math.sqrt(2) / 2, y=0, z=-math.sqrt(2) / 2), + MagnetSphericalPosition(rho=1, theta=135, phi=90), + ), + ( + MagnetPosition(x=math.sqrt(2) / 2, y=0, z=-math.sqrt(2) / 2), + MagnetSphericalPosition(rho=1, theta=225, phi=90), + ), + ( + MagnetPosition(x=math.sqrt(2) / 2, y=0, z=math.sqrt(2) / 2), + MagnetSphericalPosition(rho=1, theta=315, phi=90), + ), +]