diff --git a/src/dodal/beamlines/i06_1.py b/src/dodal/beamlines/i06_1.py index b98a1d6a9e9..1bbb5e8cb5c 100644 --- a/src/dodal/beamlines/i06_1.py +++ b/src/dodal/beamlines/i06_1.py @@ -2,6 +2,10 @@ 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.magnets import ( + MagnetThreeAxesRampRateController, + SuperConductingMagnetController, +) 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 +21,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 +38,17 @@ def xabs() -> XYThetaStage: @devices.factory() def dd() -> DiffractionDichroism: return DiffractionDichroism(f"{PREFIX.beamline_prefix}-EA-DDIFF-01:") + + +@devices.factory() +def mag_ramp_rate() -> MagnetThreeAxesRampRateController: + return MagnetThreeAxesRampRateController(f"{PREFIX.beamline_prefix}-EA-SMC") + + +@devices.factory() +def scmc( + mag_ramp_rate: MagnetThreeAxesRampRateController, +) -> SuperConductingMagnetController: + return SuperConductingMagnetController( + f"{PREFIX.beamline_prefix}-EA-MAG-01:", mag_ramp_rate + ) diff --git a/src/dodal/devices/beamlines/i06_1/magnets/__init__.py b/src/dodal/devices/beamlines/i06_1/magnets/__init__.py new file mode 100644 index 00000000000..c5e558e5a84 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnets/__init__.py @@ -0,0 +1,27 @@ +from .movement import MagnetPosition, MagnetPositionError, MagnetSphericalPosition +from .ramp_controller import ( + MagnetAxisRampRateController, + MagnetThreeAxesRampRateController, +) +from .superconducting_magnet import ( + FlyMagnetInfo, + MagnetAxis, + MagnetLimitStatus, + MagnetModes, + MagnetRampStatus, + SuperConductingMagnetController, +) + +__all__ = [ + "MagnetPosition", + "MagnetPositionError", + "MagnetSphericalPosition", + "MagnetAxisRampRateController", + "MagnetThreeAxesRampRateController", + "FlyMagnetInfo", + "MagnetAxis", + "MagnetLimitStatus", + "MagnetModes", + "MagnetRampStatus", + "SuperConductingMagnetController", +] diff --git a/src/dodal/devices/beamlines/i06_1/magnets/enums.py b/src/dodal/devices/beamlines/i06_1/magnets/enums.py new file mode 100644 index 00000000000..0de08351f74 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnets/enums.py @@ -0,0 +1,22 @@ +from ophyd_async.core import StrictEnum + + +class MagnetModes(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" + + +class MagnetRampStatus(StrictEnum): + RAMP_MADE = "RAMP MADE" + RAMPING = "RAMPING" + + +class MagnetLimitStatus(StrictEnum): + OK = "OK" + VIOLTATION = "VIOLATION" diff --git a/src/dodal/devices/beamlines/i06_1/magnets/movement.py b/src/dodal/devices/beamlines/i06_1/magnets/movement.py new file mode 100644 index 00000000000..99a97e09040 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnets/movement.py @@ -0,0 +1,187 @@ +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +def read_theta(x: float, z: float) -> float: + theta = math.degrees(math.atan2(-x, z)) + return theta % 360 + + +def read_rho(x: float, y: float, z: float) -> float: + return math.hypot(x, y, z) + + +def read_phi(x: float, y: float, z: float) -> float: + return math.degrees(math.atan2(math.hypot(x, z), y)) + + +@dataclass(kw_only=True) +class MagnetPosition: + x: float + y: float + z: float + + def to_spherical(self) -> "MagnetSphericalPosition": + return MagnetSphericalPosition( + rho=read_rho(self.x, self.y, self.z), + theta=read_theta(self.x, self.z), + phi=read_phi(self.x, self.y, 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), + ) + + +@dataclass(frozen=True, kw_only=True) +class MagnetStep: + x: float | None = None + y: float | None = None + z: float | None = None + + +class MagnetPositionError(Exception): + pass + + +class MovementStrategy(ABC): + @abstractmethod + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + """Return a sequence of safe moves.""" + + +class SphericalMovement(MovementStrategy): + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + steps: list[MagnetStep] = [] + + decreases = [ + ("z", current.z, target.z), + ("x", current.x, target.x), + ("y", current.y, target.y), + ] + + for axis, old, new in decreases: + if abs(new) < abs(old): + kwargs = {axis: new} + steps.append(MagnetStep(**kwargs)) + + if any(abs(new) > abs(old) for _, old, new in decreases): + steps.append(MagnetStep(x=target.x, y=target.y, z=target.z)) + + return steps + + +class CubicMovement(MovementStrategy): + LIMIT = 1.5 + + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + for axis in (target.x, target.y, target.z): + if abs(axis) > self.LIMIT: + raise MagnetPositionError( + f"Target {target} outside cubic operating region limit of {self.LIMIT}" + ) + + return [MagnetStep(x=target.x, y=target.y, z=target.z)] + + +class PlanarXZMovement(MovementStrategy): + LIMIT = 2.0 + + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + if target.y != 0: + raise MagnetPositionError( + f"Y must remain zero in XZ mode. Requested value was {target.y}" + ) + + hypot = math.hypot(target.x, target.z) + if hypot > self.LIMIT: + raise MagnetPositionError( + f"Outside XZ operating region. math.hypot(x={target.x}, y={target.y}) = {hypot}. Limit is {self.LIMIT}" + ) + + return [MagnetStep(x=target.x, z=target.z)] + + +@dataclass +class UniaxialMovement(MovementStrategy): + axis: str + limit: float + + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + coords = {"x": target.x, "y": target.y, "z": target.z} + for axis, value in coords.items(): + if axis == self.axis: + abs_value = abs(value) + if abs_value > self.limit: + raise MagnetPositionError( + f"{axis} value of {abs_value} exceeds limit of {self.limit}" + ) + elif value != 0: + raise MagnetPositionError( + f"{axis} must remain zero in {self.axis.upper()} mode. Requested value was {value} " + ) + + return [MagnetStep(**{self.axis: coords[self.axis]})] + + +class QuadrantXYMovement(MovementStrategy): + """Start from zero field in all axes. Appy 2.o 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 + + def moves( + self, current: MagnetPosition, target: MagnetPosition + ) -> list[MagnetStep]: + if target.z != 0: + raise MagnetPositionError( + f"Z must remain zero in QUADRANT_XY mode. You provided {target.z}" + ) + + hypot = math.hypot(target.x, target.y) + if hypot > self.LIMIT: + raise MagnetPositionError( + f"Target math.hypot(x={target.x}, y={target.y}) {hypot} is outside the XY operating region limit of {self.LIMIT}." + ) + + steps: list[MagnetStep] = [] + + # Step 1: remove X component first + if current.x != 0: + steps.append(MagnetStep(x=0)) + + # Step 2: move Y + if current.y != target.y: + steps.append(MagnetStep(y=target.y)) + + # Step 3: move X to its final value + if target.x != 0: + steps.append(MagnetStep(x=target.x)) + + return steps diff --git a/src/dodal/devices/beamlines/i06_1/magnets/ramp_controller.py b/src/dodal/devices/beamlines/i06_1/magnets/ramp_controller.py new file mode 100644 index 00000000000..77413ed012b --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnets/ramp_controller.py @@ -0,0 +1,76 @@ +from dataclasses import dataclass +from functools import cached_property + +from ophyd_async.core import ( + InstantMovableMock, + MovableLogic, + SignalR, + StandardMovable, + StandardReadable, + StandardReadableFormat, + TimeoutCalculator, + default_mock_class, + set_mock_value, +) +from ophyd_async.epics.core import epics_signal_rw + + +@dataclass +class RampRateMovableLogic(MovableLogic[float]): + limit: SignalR[float] + + async def check_move(self, new_position: float) -> None: + limit = await self.limit.get_value() + if new_position > limit: + raise ValueError( + f"Requested ramp rate {new_position} exceeds the maximum limit of {limit} for device {self.readback.name}." + ) + + async def move(self, new_position: float, timeout: TimeoutCalculator): + await self.setpoint.set(new_position, timeout=timeout()) + + +class MockMagnetAxisRampRateController(InstantMovableMock): + async def connect(self, device: StandardMovable): + await super().connect(device) + # Extend to set a sensible default value for the limit. + set_mock_value(device.limit, 2) # type: ignore + + +@default_mock_class(MockMagnetAxisRampRateController) +class MagnetAxisRampRateController(StandardMovable[float], StandardReadable): + """Controls the ramp rate of a single superconducting magnet axis. + + Exposes the readback ramp rate, demand ramp rate and the maximum + permitted ramp rate for one magnet axis. + """ + + def __init__(self, prefix: str, name: str = ""): + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.readback = epics_signal_rw(float, prefix + "STS:RAMPRATE:TPM") + self.demand = epics_signal_rw(float, prefix + "SET:DMD:RAMPRATE:TPM") + self.limit = epics_signal_rw(float, prefix + "LIM:RAMPRATE:TPM") + super().__init__(name) + + @cached_property + def movable_logic(self) -> RampRateMovableLogic: + return RampRateMovableLogic( + readback=self.readback, setpoint=self.demand, limit=self.limit + ) + + +class MagnetThreeAxesRampRateController(StandardReadable): + """Groups the ramp rate controllers for the x, y and z magnet axes. + + This device is passed to :class:`SuperConductingMagnetController` so that each + :class:`MagnetAxis` can configure its own ramp rate during preparation for + fly scans. + """ + + def __init__(self, prefix: str, name: str = ""): + with self.add_children_as_readables(): + self.x = MagnetAxisRampRateController(prefix + "-01:") + self.y = MagnetAxisRampRateController(prefix + "-02:") + self.z = MagnetAxisRampRateController(prefix + "-03:") + + super().__init__(name) diff --git a/src/dodal/devices/beamlines/i06_1/magnets/superconducting_magnet.py b/src/dodal/devices/beamlines/i06_1/magnets/superconducting_magnet.py new file mode 100644 index 00000000000..e3ae44ede17 --- /dev/null +++ b/src/dodal/devices/beamlines/i06_1/magnets/superconducting_magnet.py @@ -0,0 +1,396 @@ +import asyncio +from dataclasses import dataclass +from functools import cached_property +from typing import Protocol + +from bluesky.protocols import Flyable, Movable, Preparable +from ophyd_async.core import ( + AsyncStatus, + DeviceMock, + MovableLogic, + Reference, + StandardMovable, + StandardReadable, + StandardReadableFormat, + TimeoutCalculator, + WatchableAsyncStatus, + 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_signal_x +from pydantic import BaseModel, Field + +from dodal.devices.beamlines.i06_1.magnets.enums import ( + MagnetLimitStatus, + MagnetModes, + MagnetRampStatus, +) +from dodal.devices.beamlines.i06_1.magnets.movement import ( + CubicMovement, + MagnetPosition, + MagnetPositionError, + MagnetSphericalPosition, + MagnetStep, + MovementStrategy, + PlanarXZMovement, + QuadrantXYMovement, + SphericalMovement, + UniaxialMovement, + read_phi, + read_rho, + read_theta, +) +from dodal.devices.beamlines.i06_1.magnets.ramp_controller import ( + MagnetAxisRampRateController, + MagnetThreeAxesRampRateController, +) + + +class MagnetMoveWithinBoundary(Protocol): + async def __call__(self, x: float | None, y: float | None, z: float | None): ... + + +class FlyMagnetInfo(BaseModel): + start_position: float = Field(frozen=True) + end_position: float = Field(frozen=True) + ramp_rate: float = Field(frozen=True, gt=0) + + +@dataclass +class MagnetAxisMovableLogic(MovableLogic[float]): + axis: str + magnet_set_within_boundary: MagnetMoveWithinBoundary + + async def move(self, new_position: float, timeout: TimeoutCalculator) -> None: + values = {self.axis: new_position} + await self.magnet_set_within_boundary(**values) + + +@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, StandardMovable[float], Flyable, Preparable): + """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, + ramp_rate: MagnetAxisRampRateController, + magnet_set_within_boundary: MagnetMoveWithinBoundary, + name: str = "", + ): + self.ramp_rate_ref = Reference(ramp_rate) + + with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL): + self.readback = epics_signal_r(float, prefix + "RBV") + self.demand = epics_signal_rw(float, prefix + "DMD") + + self._movable_logic = MagnetAxisMovableLogic( + readback=self.readback, + setpoint=self.demand, + axis=axis, + magnet_set_within_boundary=magnet_set_within_boundary, + ) + self._fly_info: FlyMagnetInfo | None = None + self._fly_status: WatchableAsyncStatus | None = None + super().__init__(name) + + @cached_property + def movable_logic(self) -> MagnetAxisMovableLogic: + return self._movable_logic + + @AsyncStatus.wrap + async def prepare(self, value: FlyMagnetInfo): + self._fly_info = value + await self.ramp_rate_ref().set(value.ramp_rate) + await self.set(value.start_position) + + @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._fly_status = self.set(fly_info.end_position) + self._fly_info = None + + def complete(self) -> WatchableAsyncStatus: + fly_status = error_if_none(self._fly_status, f"{self.name} kickoff not called.") + self._fly_status = None + return fly_status + + +class MagnetCartesianCoorindates(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, + prefix: str, + ramp_rate: MagnetThreeAxesRampRateController, + set_mag_within_boundary: MagnetMoveWithinBoundary, + name: str = "", + ) -> None: + self._set_mag_within_boundary = set_mag_within_boundary + with self.add_children_as_readables(): + self.x = MagnetAxis( + prefix + "X:", "x", ramp_rate.x, set_mag_within_boundary + ) + self.y = MagnetAxis( + prefix + "Y:", "y", ramp_rate.y, set_mag_within_boundary + ) + self.z = MagnetAxis( + prefix + "Z:", "z", ramp_rate.z, set_mag_within_boundary + ) + super().__init__(name) + + @AsyncStatus.wrap + async def set(self, value: MagnetPosition): + await self._set_mag_within_boundary(x=value.x, y=value.y, z=value.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 ``MagnetSphericalPosition``. + + 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, + cart: MagnetCartesianCoorindates, + set_mag_within_boundary: MagnetMoveWithinBoundary, + name: str = "", + ): + with self.add_children_as_readables(): + self.theta = derived_signal_rw( + read_theta, self._set_theta, x=cart.x, z=cart.z + ) + self.rho = derived_signal_rw( + read_rho, self._set_rho, x=cart.x, y=cart.y, z=cart.z + ) + self.phi = derived_signal_rw( + read_phi, + self._set_phi, + x=cart.x, + y=cart.y, + z=cart.z, + ) + self._set_mag_within_boundary = set_mag_within_boundary + self._cart_ref = Reference(cart) + super().__init__(name) + + @AsyncStatus.wrap + async def set(self, value: MagnetSphericalPosition): + cart = value.to_cartesian() + await self._set_mag_within_boundary(x=cart.x, y=cart.y, z=cart.z) + + async def _set_spherical( + self, + rho: float | None = None, + theta: float | None = None, + phi: float | None = None, + ): + x, y, z = await asyncio.gather( + self._cart_ref().x.readback.get_value(), + self._cart_ref().y.readback.get_value(), + self._cart_ref().z.readback.get_value(), + ) + current = MagnetPosition(x=x, y=y, z=z).to_spherical() + spherical_pos = MagnetSphericalPosition( + rho=current.rho if rho is None else rho, + theta=current.theta if theta is None else theta, + phi=current.phi if phi is None else phi, + ) + await self.set(spherical_pos) + + async def _set_rho(self, rho: float): + await self._set_spherical(rho=rho) + + async def _set_theta(self, theta: float): + await self._set_spherical(theta=theta) + + async def _set_phi(self, phi: float): + await self._set_spherical(phi=phi) + + +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(value): + # 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.cart.x.demand.get_value(), + device.cart.y.demand.get_value(), + device.cart.z.demand.get_value(), + ) + set_mock_value(device.ramp_status, MagnetRampStatus.RAMPING) + set_mock_value(device.cart.x.readback, x_d) + set_mock_value(device.cart.y.readback, y_d) + set_mock_value(device.cart.z.readback, z_d) + set_mock_value(device.ramp_status, MagnetRampStatus.RAMP_MADE) + + callback_on_mock_put(device.start_ramp, _trigger_start_ramp) + + async def _set_mode(value): + # Whenever mode is set, ioc automatically sets everything to zero and + # triggers a ramp. + set_mock_value(device.cart.x.demand, 0.0) + set_mock_value(device.cart.y.demand, 0.0) + set_mock_value(device.cart.z.demand, 0.0) + await _trigger_start_ramp(value) + + 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) + + +@default_mock_class(MockSuperConductingMagnetController) +class SuperConductingMagnetController(StandardReadable): + """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[MagnetModes, MovementStrategy] = { + MagnetModes.SPHERICAL: SphericalMovement(), + MagnetModes.CUBIC: CubicMovement(), + MagnetModes.PLANAR_XZ: PlanarXZMovement(), + MagnetModes.QUADRANT_XY: QuadrantXYMovement(), + MagnetModes.UNIAXIAL_X: UniaxialMovement("x", limit=2.0), + MagnetModes.UNIAXIAL_Y: UniaxialMovement("y", limit=2.0), + MagnetModes.UNIAXIAL_Z: UniaxialMovement("z", limit=6.0), + } + + def __init__( + self, + prefix: str, + ramp_controllers: MagnetThreeAxesRampRateController, + name: str = "", + ): + with self.add_children_as_readables(): + self.cart = MagnetCartesianCoorindates( + prefix, ramp_controllers, self.set_within_boundary + ) + self.sph = MagnetSphericalCoordinates(self.cart, self.set_within_boundary) + + with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL): + self.mode = epics_signal_rw(MagnetModes, prefix + "MODE") + + self.ramp_status = epics_signal_rw(MagnetRampStatus, prefix + "RAMPSTATUS") + self.limit_status = epics_signal_rw(MagnetLimitStatus, prefix + "LIMITSTATUS") + self.start_ramp = epics_signal_x(prefix + "STARTRAMP.PROC") + self.timeout = 600 + + super().__init__(name) + + async def _ramp(self): + # 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.VIOLTATION: + 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(timeout=self.timeout) + await wait_for_value(self.ramp_status, MagnetRampStatus.RAMP_MADE, self.timeout) + self.log.info( + f"Ramping complete. Ramp status is now {MagnetRampStatus.RAMP_MADE}" + ) + + async def _apply_step(self, step: MagnetStep) -> 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.cart.x.demand.set(step.x)) + if step.y is not None: + tasks.append(self.cart.y.demand.set(step.y)) + if step.z is not None: + tasks.append(self.cart.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._ramp() + + async def set_within_boundary( + self, + x: float | None = None, + y: float | None = None, + z: float | None = None, + ): + """Move the magnet to a new cartesian position. + + Any coordinates left as ``None`` retain their current values. The requested + target position is validated against the current magnet operating mode and + converted into a sequence of movement steps by the corresponding movement + strategy. Each step is then applied sequentially until the requested target + position is reached. + """ + if x is None and y is None and z is None: + raise MagnetPositionError("x, y, and z cannot all be None.") + current = MagnetPosition( + x=await self.cart.x.readback.get_value(), + y=await self.cart.y.readback.get_value(), + z=await self.cart.z.readback.get_value(), + ) + target = MagnetPosition( + x=current.x if x is None else x, + y=current.y if y is None else y, + z=current.z if z is None else z, + ) + mode = await self.mode.get_value() + 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 {target}. Current position is {current}" + ) + for step in movement_strategy.moves(current, target): + await self._apply_step(step) diff --git a/tests/devices/beamlines/i06_1/magnets/__init__.py b/tests/devices/beamlines/i06_1/magnets/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/devices/beamlines/i06_1/magnets/test_movements.py b/tests/devices/beamlines/i06_1/magnets/test_movements.py new file mode 100644 index 00000000000..5d2ac5d90de --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnets/test_movements.py @@ -0,0 +1,228 @@ +import pytest + +from dodal.devices.beamlines.i06_1.magnets.movement import ( + CubicMovement, + MagnetPosition, + MagnetPositionError, + MagnetSphericalPosition, + MagnetStep, + PlanarXZMovement, + QuadrantXYMovement, + SphericalMovement, + UniaxialMovement, +) +from tests.devices.beamlines.i06_1.magnets.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_spherical_movement_decreases_z_then_x_then_y() -> None: + move_stragegy = SphericalMovement() + steps = move_stragegy.moves( + current=MagnetPosition(x=10, y=10, z=10), + target=MagnetPosition(x=5, y=5, z=2), + ) + assert steps == [MagnetStep(z=2), MagnetStep(x=5), MagnetStep(y=5)] + + +def test_spherical_movement_adds_final_combined_increase() -> None: + move_stragegy = SphericalMovement() + steps = move_stragegy.moves( + current=MagnetPosition(x=10, y=10, z=10), + target=MagnetPosition(x=20, y=5, z=2), + ) + assert steps == [MagnetStep(z=2), MagnetStep(y=5), MagnetStep(x=20, y=5, z=2)] + + +@pytest.mark.parametrize( + "current, target, expected", + [ + ( + MagnetPosition(x=3, y=3, z=3), + MagnetPosition(x=1, y=1, z=1), + [ + MagnetStep(z=1), + MagnetStep(x=1), + MagnetStep(y=1), + ], + ), + ( + MagnetPosition(x=3, y=3, z=3), + MagnetPosition(x=4, y=2, z=1), + [ + MagnetStep(z=1), + MagnetStep(y=2), + MagnetStep(x=4, y=2, z=1), + ], + ), + ( + MagnetPosition(x=1, y=1, z=1), + MagnetPosition(x=2, y=2, z=2), + [ + MagnetStep(x=2, y=2, z=2), + ], + ), + ( + MagnetPosition(x=1, y=2, z=3), + MagnetPosition(x=1, y=2, z=3), + [], + ), + ], +) +def test_spherical_movement( + current: MagnetPosition, target: MagnetPosition, expected: list[MagnetStep] +) -> None: + assert SphericalMovement().moves(current, target) == expected + + +@pytest.mark.parametrize( + "target", + [ + MagnetPosition(x=1.6, y=0, z=0), + MagnetPosition(x=0, y=-1.6, z=0), + MagnetPosition(x=0, y=0, z=2), + ], +) +def test_cubic_rejects_outside_limits(target: MagnetPosition) -> None: + with pytest.raises(MagnetPositionError): + CubicMovement().moves(MagnetPosition(x=0, y=0, z=0), target) + + +def test_cubic_returns_single_step() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetPosition(x=1, y=1.2, z=-0.5) + + assert CubicMovement().moves(current, target) == [MagnetStep(x=1, y=1.2, z=-0.5)] + + +def test_planar_xz_returns_single_step() -> None: + target = MagnetPosition(x=1.2, y=0, z=-0.4) + + assert PlanarXZMovement().moves(MagnetPosition(x=0, y=0, z=0), target) == [ + MagnetStep(x=1.2, z=-0.4), + ] + + +def test_planar_xz_rejects_nonzero_y() -> None: + with pytest.raises(MagnetPositionError): + PlanarXZMovement().moves( + MagnetPosition(x=0, y=0, z=0), MagnetPosition(x=1, y=1, z=0) + ) + + +def test_planar_xz_rejects_outside_radius() -> None: + with pytest.raises(MagnetPositionError): + PlanarXZMovement().moves( + MagnetPosition(x=0, y=0, z=0), MagnetPosition(x=2, y=0, z=2) + ) + + +@pytest.mark.parametrize( + "axis, limit, target", + [ + ("x", 2, MagnetPosition(x=1, y=0, z=0)), + ("y", 2, MagnetPosition(x=0, y=-1, z=0)), + ("z", 5, MagnetPosition(x=0, y=0, z=3)), + ], +) +def test_uniaxial_returns_single_step( + axis: str, limit: float, target: MagnetPosition +) -> None: + assert UniaxialMovement(axis, limit).moves( + MagnetPosition(x=0, y=0, z=0), target + ) == [MagnetStep(**{axis: getattr(target, axis)})] + + +@pytest.mark.parametrize( + "axis,target", + [ + ("x", MagnetPosition(x=1, y=1, z=0)), + ("x", MagnetPosition(x=1, y=0, z=1)), + ("y", MagnetPosition(x=1, y=1, z=0)), + ("y", MagnetPosition(x=0, y=1, z=1)), + ("z", MagnetPosition(x=1, y=0, z=1)), + ("z", MagnetPosition(x=0, y=1, z=1)), + ], +) +def test_uniaxial_rejects_other_axes(axis: str, target: MagnetPosition) -> None: + with pytest.raises(MagnetPositionError): + UniaxialMovement(axis, limit=5).moves(MagnetPosition(x=0, y=0, z=0), target) + + +def test_uniaxial_raise_error_when_above_limit() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetPosition(x=10, y=0, z=0) + with pytest.raises(MagnetPositionError): + UniaxialMovement("x", limit=5).moves(current, target) + + +def test_quadrant_xy_sequence() -> None: + current = MagnetPosition(x=1, y=0.5, z=0) + target = MagnetPosition(x=1.5, y=1, z=0) + strategy = QuadrantXYMovement() + assert strategy.moves(current, target) == [ + MagnetStep(x=0), + MagnetStep(y=1), + MagnetStep(x=1.5), + ] + + +def test_quadrant_xy_skips_initial_x_zero() -> None: + current = MagnetPosition(x=0, y=0, z=0) + target = MagnetPosition(x=1, y=1, z=0) + strategy = QuadrantXYMovement() + assert strategy.moves(current, target) == [MagnetStep(y=1), MagnetStep(x=1)] + + +def test_quadrant_xy_skips_y_move() -> None: + target = MagnetPosition(x=1, y=1.5, z=0) + current = MagnetPosition(x=0.5, y=1.5, z=0) + strategy = QuadrantXYMovement() + assert strategy.moves(target, current) == [MagnetStep(x=0), MagnetStep(x=0.5)] + + +def test_quadrant_xy_rejects_nonzero_z() -> None: + with pytest.raises(MagnetPositionError): + QuadrantXYMovement().moves( + MagnetPosition(x=0, y=0, z=0), MagnetPosition(x=1, y=1, z=1) + ) + + +def test_quadrant_xy_rejects_outside_radius() -> None: + with pytest.raises(MagnetPositionError): + QuadrantXYMovement().moves( + MagnetPosition(x=0, y=0, z=0), MagnetPosition(x=2, y=2, z=0) + ) diff --git a/tests/devices/beamlines/i06_1/magnets/test_ramp_controller.py b/tests/devices/beamlines/i06_1/magnets/test_ramp_controller.py new file mode 100644 index 00000000000..7d7604fd8a7 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnets/test_ramp_controller.py @@ -0,0 +1,60 @@ +import pytest +from ophyd_async.core import init_devices +from ophyd_async.testing import assert_reading, partial_reading + +from dodal.devices.beamlines.i06_1.magnets import ( + MagnetAxisRampRateController, + MagnetThreeAxesRampRateController, +) + + +@pytest.fixture +def magx_ramp_rate() -> MagnetAxisRampRateController: + with init_devices(mock=True): + magx_ramp_rate = MagnetAxisRampRateController("TEST:") + return magx_ramp_rate + + +async def test_magx_ramp_rate_read( + magx_ramp_rate: MagnetAxisRampRateController, +) -> None: + await assert_reading(magx_ramp_rate, {"magx_ramp_rate": partial_reading(0)}) + + +async def test_magx_ramp_rate_set( + magx_ramp_rate: MagnetAxisRampRateController, +) -> None: + ramp_rate = 1 + await magx_ramp_rate.set(ramp_rate) + assert await magx_ramp_rate.readback.get_value() == ramp_rate + + +async def test_magx_ramp_rate_set_above_limit_throws_error( + magx_ramp_rate: MagnetAxisRampRateController, +) -> None: + ramp_rate = 10 + with pytest.raises( + ValueError, + match=f"Requested ramp rate {ramp_rate} exceeds the maximum limit of 2.0 for device {magx_ramp_rate.name}.", + ): + await magx_ramp_rate.set(ramp_rate) + + +@pytest.fixture +def mag_three_axis_ramp_rate() -> MagnetThreeAxesRampRateController: + with init_devices(mock=True): + mag_three_axis_ramp_rate = MagnetThreeAxesRampRateController("TEST:") + return mag_three_axis_ramp_rate + + +async def test_mag_three_axis_ramp_rate_read( + mag_three_axis_ramp_rate: MagnetThreeAxesRampRateController, +) -> None: + await assert_reading( + mag_three_axis_ramp_rate, + { + "mag_three_axis_ramp_rate-x": partial_reading(0), + "mag_three_axis_ramp_rate-y": partial_reading(0), + "mag_three_axis_ramp_rate-z": partial_reading(0), + }, + ) diff --git a/tests/devices/beamlines/i06_1/magnets/test_superconducting_magnet.py b/tests/devices/beamlines/i06_1/magnets/test_superconducting_magnet.py new file mode 100644 index 00000000000..9b6bc81a237 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnets/test_superconducting_magnet.py @@ -0,0 +1,365 @@ +import asyncio +import math +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest +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.magnets import ( + FlyMagnetInfo, + MagnetAxis, + MagnetAxisRampRateController, + MagnetLimitStatus, + MagnetModes, + MagnetPosition, + MagnetPositionError, + MagnetRampStatus, + MagnetSphericalPosition, + MagnetThreeAxesRampRateController, + SuperConductingMagnetController, + movement, +) +from tests.devices.beamlines.i06_1.magnets.utils import ( + EXPECTED_CARTESIAN_SPHERICAL_CONVERSION, +) + + +@pytest.fixture +def ramp_rate() -> MagnetThreeAxesRampRateController: + with init_devices(mock=True): + ramp_rate = MagnetThreeAxesRampRateController("TEST:") + return ramp_rate + + +@pytest.fixture +def scmc( + ramp_rate: MagnetThreeAxesRampRateController, +) -> SuperConductingMagnetController: + with init_devices(mock=True): + scmc = SuperConductingMagnetController("TEST:", ramp_rate) + return scmc + + +async def test_scmc_read(scmc: SuperConductingMagnetController) -> None: + await assert_reading( + scmc, + { + "scmc-cart-x": partial_reading(0), + "scmc-cart-y": partial_reading(0), + "scmc-cart-z": partial_reading(0), + "scmc-sph-theta": partial_reading(0), + "scmc-sph-rho": partial_reading(0), + "scmc-sph-phi": partial_reading(0), + }, + ) + + +async def test_scmc_configuration(scmc: SuperConductingMagnetController) -> None: + await assert_configuration( + scmc, {"scmc-mode": partial_reading(MagnetModes.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( + scmc: SuperConductingMagnetController, + axis: str, + value: float, + expected_x: float, + expected_y: float, + expected_z: float, +) -> None: + await scmc.mode.set(getattr(MagnetModes, "UNIAXIAL_" + axis.capitalize())) + await getattr(scmc.cart, axis).set(value) + x, y, z = await asyncio.gather( + scmc.cart.x.readback.get_value(), + scmc.cart.y.readback.get_value(), + scmc.cart.z.readback.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( + scmc: SuperConductingMagnetController, + cartesian: MagnetPosition, + spherical: MagnetSphericalPosition, +) -> None: + await scmc.mode.set(MagnetModes.SPHERICAL) + await scmc.sph.set(spherical) + rho, theta, phi = await asyncio.gather( + scmc.sph.rho.get_value(), + scmc.sph.theta.get_value(), + scmc.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( + scmc.cart.x.readback.get_value(), + scmc.cart.y.readback.get_value(), + scmc.cart.z.readback.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( + scmc: SuperConductingMagnetController, + cartesian: MagnetPosition, + spherical: MagnetSphericalPosition, +) -> None: + await scmc.mode.set(MagnetModes.CUBIC) + await scmc.cart.set(cartesian) + x, y, z = await asyncio.gather( + scmc.cart.x.readback.get_value(), + scmc.cart.y.readback.get_value(), + scmc.cart.z.readback.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( + scmc.sph.rho.get_value(), + scmc.sph.theta.get_value(), + scmc.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( + scmc: SuperConductingMagnetController, + 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 scmc.mode.set(MagnetModes.SPHERICAL) + for axis, value in zip(axes, values, strict=True): + await getattr(scmc.sph, axis).set(value) + + await assert_reading( + scmc, + { + "scmc-cart-x": partial_reading(expected_x), + "scmc-cart-y": partial_reading(expected_y), + "scmc-cart-z": partial_reading(expected_z), + "scmc-sph-rho": partial_reading(expected_rho), + "scmc-sph-phi": partial_reading(expected_phi), + "scmc-sph-theta": partial_reading(expected_theta), + }, + ) + + +async def test_scmc_raises_error_if_limit_status_is_violation( + scmc: SuperConductingMagnetController, +) -> None: + await scmc.mode.set(MagnetModes.UNIAXIAL_X) + set_mock_value(scmc.limit_status, MagnetLimitStatus.VIOLTATION) + with pytest.raises(MagnetPositionError): + await scmc.cart.x.set(1) + + +async def test_scmc_set_within_boundary_raises_error_if_all_values_none( + scmc: SuperConductingMagnetController, +) -> None: + with pytest.raises(MagnetPositionError): + await scmc.set_within_boundary() + + +@pytest.mark.parametrize( + "axis, mode", + [ + ("x", MagnetModes.UNIAXIAL_X), + ("y", MagnetModes.UNIAXIAL_Y), + ("z", MagnetModes.UNIAXIAL_Z), + ], +) +async def test_scmc_axis_with_ramp_rate_wired_correctly_with_prepare( + scmc: SuperConductingMagnetController, + ramp_rate: MagnetThreeAxesRampRateController, + axis: str, + mode: MagnetModes, +) -> None: + await scmc.mode.set(mode) + magnet_axis: MagnetAxis = getattr(scmc.cart, axis) + ramp_axis: MagnetAxisRampRateController = getattr(ramp_rate, axis) + + fly_info = FlyMagnetInfo(start_position=1, end_position=6, ramp_rate=1.5) + await magnet_axis.prepare(fly_info) + assert await magnet_axis.demand.get_value() == fly_info.start_position + assert await ramp_axis.demand.get_value() == fly_info.ramp_rate + + +async def test_scmc_axis_kickoff_and_complete( + scmc: SuperConductingMagnetController, +) -> None: + await scmc.mode.set(MagnetModes.UNIAXIAL_X) + fly_info = FlyMagnetInfo(start_position=1, end_position=2, ramp_rate=2) + assert scmc.cart.x._fly_info is None + await scmc.cart.x.prepare(fly_info) + assert scmc.cart.x._fly_info is fly_info + assert scmc.cart.x._fly_status is None + await scmc.cart.x.kickoff() + assert scmc.cart.x._fly_info is None + await scmc.cart.x.complete() + assert await scmc.cart.x.readback.get_value() == fly_info.end_position + + +async def test_scmc_axis_kickoff_and_complete_raises_error_without_prepare( + scmc: SuperConductingMagnetController, +): + await scmc.mode.set(MagnetModes.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 scmc.cart.x.kickoff() + + with pytest.raises(RuntimeError): + await scmc.cart.x.complete() + + fly_info = FlyMagnetInfo(start_position=1, end_position=2, ramp_rate=1) + await scmc.cart.x.prepare(fly_info) + await scmc.cart.x.kickoff() + await scmc.cart.x.complete() + + +async def test_scmc_mock_device_behaviour( + scmc: SuperConductingMagnetController, +) -> None: + assert await scmc.limit_status.get_value() == MagnetLimitStatus.OK + + set_mock_value(scmc.cart.x.demand, 1) + set_mock_value(scmc.cart.y.demand, 1) + set_mock_value(scmc.cart.z.demand, 1) + + ramp_states: list[MagnetRampStatus] = [] + + def _ramp_state_callback(value: dict[str, Reading[MagnetRampStatus]]): + ramp_states.append(value[scmc.ramp_status.name]["value"]) + + scmc.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 scmc.mode.set(MagnetModes.UNIAXIAL_X) + + x_d, y_d, z_d, x_rbv, y_rbv, z_rbv = await asyncio.gather( + scmc.cart.x.demand.get_value(), + scmc.cart.y.demand.get_value(), + scmc.cart.z.demand.get_value(), + scmc.cart.x.readback.get_value(), + scmc.cart.y.readback.get_value(), + scmc.cart.z.readback.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( + scmc: SuperConductingMagnetController, +) -> None: + mode = await scmc.mode.get_value() + with patch.dict(scmc._MODE_MOVEMENT_STRATEGY, {}, clear=True): + with pytest.raises( + ValueError, + match=f"No movement strategy has been configured for device scmc for mode {mode}", + ): + await scmc.cart.x.set(5) + + +@pytest.mark.parametrize( + "mode,expected", + [ + (MagnetModes.SPHERICAL, movement.SphericalMovement), + (MagnetModes.CUBIC, movement.CubicMovement), + (MagnetModes.PLANAR_XZ, movement.PlanarXZMovement), + (MagnetModes.QUADRANT_XY, movement.QuadrantXYMovement), + (MagnetModes.UNIAXIAL_X, movement.UniaxialMovement), + (MagnetModes.UNIAXIAL_Y, movement.UniaxialMovement), + (MagnetModes.UNIAXIAL_Z, movement.UniaxialMovement), + ], +) +async def test_scmc_magnet_mode_to_movement_strategy_configuration( + scmc: SuperConductingMagnetController, + mode: MagnetModes, + expected: type[movement.MovementStrategy], +) -> None: + assert isinstance(scmc._MODE_MOVEMENT_STRATEGY[mode], expected) + + +async def test_scmc_magnet_mode_to_uniaxial_movement_strategy_configuration( + scmc: SuperConductingMagnetController, +) -> None: + assert scmc._MODE_MOVEMENT_STRATEGY[MagnetModes.UNIAXIAL_X].axis == "x" # type:ignore + assert scmc._MODE_MOVEMENT_STRATEGY[MagnetModes.UNIAXIAL_Y].axis == "y" # type:ignore + assert scmc._MODE_MOVEMENT_STRATEGY[MagnetModes.UNIAXIAL_Z].axis == "z" # type:ignore + + +async def test_scmc_executes_movement_stragey_and_ramp_at_each_step( + scmc: SuperConductingMagnetController, +) -> None: + + movement_strategy = MagicMock() + + mov_str_return_values = [ + movement.MagnetStep(z=2), + movement.MagnetStep(x=1, y=2, z=3), + ] + expected_apply_step_calls = [ + call(movement.MagnetStep(z=2)), + call(movement.MagnetStep(x=1, y=2, z=3)), + ] + movement_strategy.moves.return_value = mov_str_return_values + + scmc._MODE_MOVEMENT_STRATEGY[MagnetModes.UNIAXIAL_X] = movement_strategy + + scmc._ramp = AsyncMock() + + with patch( + "dodal.devices.beamlines.i06_1.magnets.superconducting_magnet.SuperConductingMagnetController._apply_step", + wraps=scmc._apply_step, + ) as mock_apply_step: + await scmc.mode.set(MagnetModes.UNIAXIAL_X) + await scmc.cart.x.set(4) + + assert mock_apply_step.call_args_list == expected_apply_step_calls + assert scmc._ramp.call_count == len(mov_str_return_values) diff --git a/tests/devices/beamlines/i06_1/magnets/utils.py b/tests/devices/beamlines/i06_1/magnets/utils.py new file mode 100644 index 00000000000..f55e6863044 --- /dev/null +++ b/tests/devices/beamlines/i06_1/magnets/utils.py @@ -0,0 +1,32 @@ +import math + +from dodal.devices.beamlines.i06_1.magnets.movement 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), + ), +]