From 44bdecc8a20fb5c8fbe93a760713b2cf749beb03 Mon Sep 17 00:00:00 2001 From: olliesilvester <122091460+olliesilvester@users.noreply.github.com> Date: Tue, 11 Jun 2024 08:33:13 +0100 Subject: [PATCH] Refactor and create ophyd-async FGS devices (#422) * Convert FGS device to ophyd async * Combine panda and zebra FGS devices and params * Param positions now a function instead of attribute --------- Co-authored-by: Dominic Oram --- src/dodal/beamlines/i03.py | 17 +- src/dodal/beamlines/i04.py | 12 +- src/dodal/devices/fast_grid_scan.py | 377 +++++++++--------- src/dodal/devices/panda_fast_grid_scan.py | 162 -------- .../system_tests/test_gridscan_system.py | 68 +--- tests/devices/unit_tests/test_gridscan.py | 297 ++++++++------ .../devices/unit_tests/test_panda_gridscan.py | 112 ------ 7 files changed, 378 insertions(+), 667 deletions(-) delete mode 100644 src/dodal/devices/panda_fast_grid_scan.py delete mode 100644 tests/devices/unit_tests/test_panda_gridscan.py diff --git a/src/dodal/beamlines/i03.py b/src/dodal/beamlines/i03.py index 76f0c5184a..439197f899 100644 --- a/src/dodal/beamlines/i03.py +++ b/src/dodal/beamlines/i03.py @@ -14,12 +14,11 @@ from dodal.devices.detector import DetectorParams from dodal.devices.detector.detector_motion import DetectorMotion from dodal.devices.eiger import EigerDetector -from dodal.devices.fast_grid_scan import FastGridScan +from dodal.devices.fast_grid_scan import PandAFastGridScan, ZebraFastGridScan from dodal.devices.flux import Flux from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, VFMMirrorVoltages from dodal.devices.oav.oav_detector import OAV, OAVConfigParams from dodal.devices.oav.pin_image_recognition import PinTipDetection -from dodal.devices.panda_fast_grid_scan import PandAFastGridScan from dodal.devices.qbpm1 import QBPM1 from dodal.devices.robot import BartRobot from dodal.devices.s4_slit_gaps import S4SlitGaps @@ -200,16 +199,16 @@ def set_params(eiger: EigerDetector): ) -def fast_grid_scan( +def zebra_fast_grid_scan( wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False -) -> FastGridScan: - """Get the i03 fast_grid_scan device, instantiate it if it hasn't already been. +) -> ZebraFastGridScan: + """Get the i03 zebra_fast_grid_scan device, instantiate it if it hasn't already been. If this is called when already instantiated in i03, it will return the existing object. """ return device_instantiation( - device_factory=FastGridScan, - name="fast_grid_scan", - prefix="-MO-SGON-01:", + device_factory=ZebraFastGridScan, + name="zebra_fast_grid_scan", + prefix="-MO-SGON-01:FGS:", wait=wait_for_connection, fake=fake_with_ophyd_sim, ) @@ -220,7 +219,7 @@ def panda_fast_grid_scan( ) -> PandAFastGridScan: """Get the i03 panda_fast_grid_scan device, instantiate it if it hasn't already been. If this is called when already instantiated in i03, it will return the existing object. - This is used instead of the fast_grid_scan device when using the PandA. + This is used instead of the zebra_fast_grid_scan device when using the PandA. """ return device_instantiation( device_factory=PandAFastGridScan, diff --git a/src/dodal/beamlines/i04.py b/src/dodal/beamlines/i04.py index b37ebf2da3..fcdf08d9d5 100644 --- a/src/dodal/beamlines/i04.py +++ b/src/dodal/beamlines/i04.py @@ -8,7 +8,7 @@ from dodal.devices.detector import DetectorParams from dodal.devices.detector.detector_motion import DetectorMotion from dodal.devices.eiger import EigerDetector -from dodal.devices.fast_grid_scan import FastGridScan +from dodal.devices.fast_grid_scan import ZebraFastGridScan from dodal.devices.flux import Flux from dodal.devices.i04.transfocator import Transfocator from dodal.devices.ipin import IPin @@ -260,15 +260,15 @@ def set_params(eiger: EigerDetector): ) -def fast_grid_scan( +def zebra_fast_grid_scan( wait_for_connection: bool = True, fake_with_ophyd_sim: bool = False -) -> FastGridScan: - """Get the i04 fast_grid_scan device, instantiate it if it hasn't already been. +) -> ZebraFastGridScan: + """Get the i04 zebra_fast_grid_scan device, instantiate it if it hasn't already been. If this is called when already instantiated in i04, it will return the existing object. """ return device_instantiation( - device_factory=FastGridScan, - name="fast_grid_scan", + device_factory=ZebraFastGridScan, + name="zebra_fast_grid_scan", prefix="-MO-SGON-01:", wait=wait_for_connection, fake=fake_with_ophyd_sim, diff --git a/src/dodal/devices/fast_grid_scan.py b/src/dodal/devices/fast_grid_scan.py index 7d58f1768f..0a87d48a53 100644 --- a/src/dodal/devices/fast_grid_scan.py +++ b/src/dodal/devices/fast_grid_scan.py @@ -1,24 +1,27 @@ -import threading -import time -from typing import Any +from abc import ABC +from typing import Any, Generic, TypeVar import numpy as np from bluesky.plan_stubs import mv from numpy import ndarray -from ophyd import ( - Component, +from ophyd_async.core import ( + AsyncStatus, Device, - EpicsSignal, - EpicsSignalRO, - EpicsSignalWithRBV, Signal, + SignalR, + SoftSignalBackend, + StandardReadable, + wait_for_value, ) -from ophyd.status import DeviceStatus, StatusBase +from ophyd_async.epics.signal import ( + epics_signal_r, + epics_signal_rw, +) +from ophyd_async.epics.signal.signal import epics_signal_rw_rbv from pydantic import validator from pydantic.dataclasses import dataclass from dodal.devices.motors import XYZLimitBundle -from dodal.devices.status import await_value from dodal.log import LOGGER from dodal.parameters.experiment_parameter_base import AbstractExperimentWithBeamParams @@ -50,9 +53,6 @@ class GridScanParamsCommon(AbstractExperimentWithBeamParams): layout to EPICS. The parameters and functions of this class are common to both the zebra and panda triggered fast grid scans. - Motion program will do a grid in x-y then rotate omega +90 and perform - a grid in x-z. - The grid specified is where data is taken e.g. it can be assumed the first frame is at x_start, y1_start, z1_start and subsequent frames are N*step_size away. """ @@ -75,6 +75,21 @@ class GridScanParamsCommon(AbstractExperimentWithBeamParams): # Whether to set the stub offsets after centering set_stub_offsets: bool = False + def get_param_positions(self) -> dict: + return { + "x_steps": self.x_steps, + "y_steps": self.y_steps, + "z_steps": self.z_steps, + "x_step_size": self.x_step_size, + "y_step_size": self.y_step_size, + "z_step_size": self.z_step_size, + "x_start": self.x_start, + "y1_start": self.y1_start, + "y2_start": self.y2_start, + "z1_start": self.z1_start, + "z2_start": self.z2_start, + } + class Config: arbitrary_types_allowed = True fields = { @@ -153,21 +168,21 @@ def grid_position_to_motor_position(self, grid_position: ndarray) -> ndarray: ) -class GridScanParams(GridScanParamsCommon): - """ - Holder class for the parameters of a grid scan in a similar - layout to EPICS. These params are used for the zebra-triggered - fast grid scan +ParamType = TypeVar("ParamType", bound=GridScanParamsCommon) - Motion program will do a grid in x-y then rotate omega +90 and perform - a grid in x-z. - The grid specified is where data is taken e.g. it can be assumed the first frame is - at x_start, y1_start, z1_start and subsequent frames are N*step_size away. +class ZebraGridScanParams(GridScanParamsCommon): + """ + Params for standard Zebra FGS. Adds on the dwell time """ dwell_time_ms: float = 10 + def get_param_positions(self): + param_positions = super().get_param_positions() + param_positions["dwell_time_ms"] = self.dwell_time_ms + return param_positions + @validator("dwell_time_ms", always=True, check_fields=True) def non_integer_dwell_time(cls, dwell_time_ms: float) -> float: dwell_time_floor_rounded = np.floor(dwell_time_ms) @@ -181,181 +196,159 @@ def non_integer_dwell_time(cls, dwell_time_ms: float) -> float: return dwell_time_ms -class GridScanCompleteStatus(DeviceStatus): +class PandAGridScanParams(GridScanParamsCommon): + """ + Params for panda constant-motion scan. Adds on the goniometer run-up distance + """ + + run_up_distance_mm: float = 0.17 + + def get_param_positions(self): + param_positions = super().get_param_positions() + param_positions["run_up_distance_mm"] = self.run_up_distance_mm + return param_positions + + +class MotionProgram(Device): + def __init__(self, prefix: str, name: str = "") -> None: + super().__init__(name) + self.running = epics_signal_r(float, prefix + "PROGBITS") + self.program_number = epics_signal_r(float, prefix + "CS1:PROG_NUM") + + +class ExpectedImages(SignalR[int]): + def __init__(self, parent: "FastGridScanCommon") -> None: + super().__init__(SoftSignalBackend(int)) + self.parent: "FastGridScanCommon" = parent + + async def get_value(self): + x = await self.parent.x_steps.get_value() + y = await self.parent.y_steps.get_value() + z = await self.parent.z_steps.get_value() + first_grid = x * y + second_grid = x * z + return first_grid + second_grid + + +class FastGridScanCommon(StandardReadable, ABC, Generic[ParamType]): + """Device for a general fast grid scan + + When the motion program is started, the goniometer will move in a snake-like grid trajectory, + with X as the fast axis and Y as the slow axis. If Z steps isn't 0, the goniometer will + then rotate in the omega direction such that it moves from the X-Y, to the X-Z plane then + do a second grid scan. The detector is triggered after every x step. + See https://github.com/DiamondLightSource/hyperion/wiki/Coordinate-Systems for more """ - A Status for the grid scan completion - A special status object that notifies watchers (progress bars) - based on comparing device.expected_images to device.position_counter. + + def __init__(self, prefix: str, name: str = "") -> None: + super().__init__(name) + self.x_steps = epics_signal_rw_rbv(int, "X_NUM_STEPS") + self.y_steps = epics_signal_rw_rbv( + int, "X_NUM_STEPS" + ) # Number of vertical steps during the first grid scan + self.z_steps = epics_signal_rw_rbv( + int, "X_NUM_STEPS" + ) # Number of vertical steps during the second grid scan, after the rotation in omega + self.x_step_size = epics_signal_rw_rbv(float, "X_STEP_SIZE") + self.y_step_size = epics_signal_rw_rbv(float, "Y_STEP_SIZE") + self.z_step_size = epics_signal_rw_rbv(float, "Z_STEP_SIZE") + self.x_start = epics_signal_rw_rbv(float, "X_START") + self.y1_start = epics_signal_rw_rbv(float, "Y_START") + self.y2_start = epics_signal_rw_rbv(float, "Y2_START") + self.z1_start = epics_signal_rw_rbv(float, "Z_START") + self.z2_start = epics_signal_rw_rbv(float, "Z2_START") + + self.position_counter = epics_signal_rw( + int, "POS_COUNTER", write_pv="POS_COUNTER_WRITE" + ) + self.x_counter = epics_signal_r(int, "X_COUNTER") + self.y_counter = epics_signal_r(int, "Y_COUNTER") + self.scan_invalid = epics_signal_r(float, "SCAN_INVALID") + + self.run_cmd = epics_signal_rw(int, "RUN.PROC") + self.stop_cmd = epics_signal_rw(int, "STOP.PROC") + self.status = epics_signal_r(float, "SCAN_STATUS") + + self.expected_images = ExpectedImages(parent=self) + + self.motion_program = MotionProgram(prefix) + + # Kickoff timeout in seconds + self.KICKOFF_TIMEOUT: float = 5.0 + + self.COMPLETE_STATUS: float = 60.0 + + self.movable_params: dict[str, Signal] = { + "x_steps": self.x_steps, + "y_steps": self.y_steps, + "z_steps": self.z_steps, + "x_step_size": self.x_step_size, + "y_step_size": self.y_step_size, + "z_step_size": self.z_step_size, + "x_start": self.x_start, + "y1_start": self.y1_start, + "y2_start": self.y2_start, + "z1_start": self.z1_start, + "z2_start": self.z2_start, + } + + @AsyncStatus.wrap + async def kickoff(self): + curr_prog = await self.motion_program.program_number.get_value() + running = await self.motion_program.running.get_value() + if running: + LOGGER.info(f"Motion program {curr_prog} still running, waiting...") + await wait_for_value(self.motion_program.running, 0, self.KICKOFF_TIMEOUT) + + LOGGER.debug("Running scan") + await self.run_cmd.set(1) + LOGGER.info("Waiting for FGS to start") + await wait_for_value(self.status, 1, self.KICKOFF_TIMEOUT) + LOGGER.debug("FGS kicked off") + + @AsyncStatus.wrap + async def complete(self): + await wait_for_value(self.status, 0, self.COMPLETE_STATUS) + + +class ZebraFastGridScan(FastGridScanCommon[ZebraGridScanParams]): + """Device for standard Zebra FGS. In this scan, the goniometer's velocity profile follows a parabolic shape between X steps, + with the slowest points occuring at each X step. """ - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.start_ts = time.time() - - self.device.position_counter.subscribe(self._notify_watchers) - self.device.status.subscribe(self._running_changed) - - self._name = self.device.name - self._target_count = self.device.expected_images.get() - - def _notify_watchers(self, value, *args, **kwargs): - if not self._watchers: - return - time_elapsed = time.time() - self.start_ts - try: - fraction = 1 - value / self._target_count - except ZeroDivisionError: - fraction = 0 - time_remaining = 0 - except Exception as e: - fraction = None - time_remaining = None - self.set_exception(e) - self.clean_up() - else: - time_remaining = time_elapsed / fraction - for watcher in self._watchers: - watcher( - name=self._name, - current=value, - initial=0, - target=self._target_count, - unit="images", - precision=0, - fraction=fraction, - time_elapsed=time_elapsed, - time_remaining=time_remaining, + def __init__(self, prefix: str, name: str = "") -> None: + super().__init__(prefix, name) + + # Time taken to travel between X steps + self.dwell_time_ms = epics_signal_rw_rbv(float, "DWELL_TIME") + self.movable_params["dwell_time_ms"] = self.dwell_time_ms + + +class PandAFastGridScan(FastGridScanCommon[PandAGridScanParams]): + """Device for panda constant-motion scan""" + + def __init__(self, prefix: str, name: str = "") -> None: + super().__init__(prefix, name) + self.time_between_x_steps_ms = ( + epics_signal_rw_rbv( # Used by motion controller to set goniometer velocity + float, "TIME_BETWEEN_X_STEPS" ) + ) - def _running_changed(self, value=None, old_value=None, **kwargs): - if (old_value == 1) and (value == 0): - self.set_finished() - self.clean_up() + # Distance before and after the grid given to allow goniometer to reach desired speed while it is within the + # grid + self.run_up_distance_mm = epics_signal_rw_rbv(float, "RUNUP_DISTANCE") + self.movable_params["run_up_distance_mm"] = self.run_up_distance_mm - def clean_up(self): - self.device.position_counter.clear_sub(self._notify_watchers) - self.device.status.clear_sub(self._running_changed) +def set_fast_grid_scan_params(scan: FastGridScanCommon[ParamType], params: ParamType): + to_move = [] -class MotionProgram(Device): - running = Component(EpicsSignalRO, "PROGBITS") - program_number = Component(EpicsSignalRO, "CS1:PROG_NUM") - - -class FastGridScan(Device): - x_steps = Component(EpicsSignalWithRBV, "FGS:X_NUM_STEPS") - y_steps = Component(EpicsSignalWithRBV, "FGS:Y_NUM_STEPS") - z_steps = Component(EpicsSignalWithRBV, "FGS:Z_NUM_STEPS") - - x_step_size = Component(EpicsSignalWithRBV, "FGS:X_STEP_SIZE") - y_step_size = Component(EpicsSignalWithRBV, "FGS:Y_STEP_SIZE") - z_step_size = Component(EpicsSignalWithRBV, "FGS:Z_STEP_SIZE") - - dwell_time_ms = Component(EpicsSignalWithRBV, "FGS:DWELL_TIME") - - x_start = Component(EpicsSignalWithRBV, "FGS:X_START") - y1_start = Component(EpicsSignalWithRBV, "FGS:Y_START") - y2_start = Component(EpicsSignalWithRBV, "FGS:Y2_START") - z1_start = Component(EpicsSignalWithRBV, "FGS:Z_START") - z2_start = Component(EpicsSignalWithRBV, "FGS:Z2_START") - - position_counter = Component( - EpicsSignal, "FGS:POS_COUNTER", write_pv="FGS:POS_COUNTER_WRITE" - ) - x_counter = Component(EpicsSignalRO, "FGS:X_COUNTER") - y_counter = Component(EpicsSignalRO, "FGS:Y_COUNTER") - scan_invalid = Component(EpicsSignalRO, "FGS:SCAN_INVALID") - - run_cmd = Component(EpicsSignal, "FGS:RUN.PROC") - stop_cmd = Component(EpicsSignal, "FGS:STOP.PROC") - status = Component(EpicsSignalRO, "FGS:SCAN_STATUS") - - expected_images = Component(Signal) - - motion_program = Component(MotionProgram, "") - - # Kickoff timeout in seconds - KICKOFF_TIMEOUT: float = 5.0 - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def set_expected_images(*_, **__): - x = int(self.x_steps.get()) - y = int(self.y_steps.get()) - z = int(self.z_steps.get()) - first_grid = x * y - second_grid = x * z - self.expected_images.put(first_grid + second_grid) - - self.x_steps.subscribe(set_expected_images) - self.y_steps.subscribe(set_expected_images) - self.z_steps.subscribe(set_expected_images) - - def is_invalid(self) -> bool: - if "GONP" in self.scan_invalid.pvname: - return False - return bool(self.scan_invalid.get()) - - def kickoff(self) -> StatusBase: - st = DeviceStatus(device=self, timeout=self.KICKOFF_TIMEOUT) - - def scan(): - try: - curr_prog = self.motion_program.program_number.get() - running = self.motion_program.running.get() - if running: - LOGGER.info(f"Motion program {curr_prog} still running, waiting...") - await_value(self.motion_program.running, 0).wait() - LOGGER.debug("Running scan") - self.run_cmd.put(1) - LOGGER.info("Waiting for FGS to start") - await_value(self.status, 1).wait() - st.set_finished() - LOGGER.debug(f"{st} finished, exiting FGS kickoff thread") - except Exception as e: - st.set_exception(e) - - threading.Thread(target=scan, daemon=True).start() - LOGGER.info("Returning FGS kickoff status") - return st - - def complete(self) -> DeviceStatus: - return GridScanCompleteStatus(self) - - def collect(self): - return {} - - def describe_collect(self): - return {} - - -def set_fast_grid_scan_params(scan: FastGridScan, params: GridScanParams): - yield from mv( - scan.x_steps, - params.x_steps, - scan.y_steps, - params.y_steps, - scan.z_steps, - params.z_steps, - scan.x_step_size, - params.x_step_size, - scan.y_step_size, - params.y_step_size, - scan.z_step_size, - params.z_step_size, - scan.dwell_time_ms, - params.dwell_time_ms, - scan.x_start, - params.x_start, - scan.y1_start, - params.y1_start, - scan.y2_start, - params.y2_start, - scan.z1_start, - params.z1_start, - scan.z2_start, - params.z2_start, - scan.position_counter, - 0, - ) + # Create arguments for bps.mv + for key in scan.movable_params.keys(): + to_move.extend([scan.movable_params[key], params.__dict__[key]]) + + # Counter should always start at 0 + to_move.extend([scan.position_counter, 0]) + + yield from mv(*to_move) diff --git a/src/dodal/devices/panda_fast_grid_scan.py b/src/dodal/devices/panda_fast_grid_scan.py deleted file mode 100644 index c4dfeb61a7..0000000000 --- a/src/dodal/devices/panda_fast_grid_scan.py +++ /dev/null @@ -1,162 +0,0 @@ -import threading - -from bluesky.plan_stubs import mv -from ophyd import ( - Component, - Device, - EpicsSignal, - EpicsSignalRO, - EpicsSignalWithRBV, - Signal, -) -from ophyd.status import DeviceStatus, StatusBase - -from dodal.devices.fast_grid_scan import GridScanParamsCommon -from dodal.devices.status import await_value - - -class GridScanCompleteStatus(DeviceStatus): - """ - A Status for the grid scan completion - Progress bar functionality has been removed for now in the panda fast grid scan - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.device.status.subscribe(self._running_changed) - - def _running_changed(self, value=None, old_value=None, **kwargs): - if (old_value == 1) and (value == 0): - self.set_finished() - self.clean_up() - - def clean_up(self): - self.device.status.clear_sub(self._running_changed) - - -class PandAGridScanParams(GridScanParamsCommon): - """ - Holder class for the parameters of a grid scan in a similar - layout to EPICS. These params are used for the panda-triggered - constant motion grid scan - - Motion program will do a grid in x-y then rotate omega +90 and perform - a grid in x-z. - - The grid specified is where data is taken e.g. it can be assumed the first frame is - at x_start, y1_start, z1_start and subsequent frames are N*step_size away. - """ - - run_up_distance_mm: float = 0.17 - - -class PandAFastGridScan(Device): - """This is similar to the regular FastGridScan device. It has two extra PVs: runup distance and time between x steps. - Dwell time is not moved in this scan. - """ - - x_steps: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "X_NUM_STEPS") - y_steps: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Y_NUM_STEPS") - z_steps: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Z_NUM_STEPS") - - x_step_size: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "X_STEP_SIZE") - y_step_size: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Y_STEP_SIZE") - z_step_size: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Z_STEP_SIZE") - - # This value is fixed by the time between X steps detector deadtime. The only reason it is a PV - # Is so the value can be read by the motion program in the PPMAC - time_between_x_steps_ms = Component(EpicsSignalWithRBV, "TIME_BETWEEN_X_STEPS") - - run_up_distance: EpicsSignalWithRBV = Component(EpicsSignal, "RUNUP_DISTANCE") - - x_start: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "X_START") - y1_start: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Y_START") - y2_start: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Y2_START") - z1_start: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Z_START") - z2_start: EpicsSignalWithRBV = Component(EpicsSignalWithRBV, "Z2_START") - - position_counter: EpicsSignalRO = Component(EpicsSignalRO, "Y_COUNTER") - scan_invalid: EpicsSignalRO = Component(EpicsSignalRO, "SCAN_INVALID") - - run_cmd: EpicsSignal = Component(EpicsSignal, "RUN.PROC") - stop_cmd: EpicsSignal = Component(EpicsSignal, "STOP.PROC") - status: EpicsSignalRO = Component(EpicsSignalRO, "SCAN_STATUS") - - expected_images: Signal = Component(Signal) - - # Kickoff timeout in seconds - KICKOFF_TIMEOUT: float = 5.0 - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def set_expected_images(*_, **__): - x, y, z = self.x_steps.get(), self.y_steps.get(), self.z_steps.get() - first_grid = x * y - second_grid = x * z - self.expected_images.put(first_grid + second_grid) - - self.x_steps.subscribe(set_expected_images) - self.y_steps.subscribe(set_expected_images) - self.z_steps.subscribe(set_expected_images) - - def is_invalid(self) -> bool: - if "GONP" in self.scan_invalid.pvname: - return False - return self.scan_invalid.get() - - def kickoff(self) -> StatusBase: - # Check running already here? - st = DeviceStatus(device=self, timeout=self.KICKOFF_TIMEOUT) - - def scan(): - try: - self.log.debug("Running scan") - self.run_cmd.put(1) - self.log.debug("Waiting for scan to start") - await_value(self.status, 1).wait() - st.set_finished() - except Exception as e: - st.set_exception(e) - - threading.Thread(target=scan, daemon=True).start() - return st - - def complete(self) -> DeviceStatus: - return GridScanCompleteStatus(self) - - def collect(self): - return {} - - def describe_collect(self): - return {} - - -def set_fast_grid_scan_params(scan: PandAFastGridScan, params: PandAGridScanParams): - yield from mv( - scan.x_steps, - params.x_steps, - scan.y_steps, - params.y_steps, - scan.z_steps, - params.z_steps, - scan.x_step_size, - params.x_step_size, - scan.y_step_size, - params.y_step_size, - scan.z_step_size, - params.z_step_size, - scan.x_start, - params.x_start, - scan.y1_start, - params.y1_start, - scan.y2_start, - params.y2_start, - scan.z1_start, - params.z1_start, - scan.z2_start, - params.z2_start, - scan.run_up_distance, - params.run_up_distance_mm, - ) diff --git a/tests/devices/system_tests/test_gridscan_system.py b/tests/devices/system_tests/test_gridscan_system.py index d2bbb1013e..2da4e6ab75 100644 --- a/tests/devices/system_tests/test_gridscan_system.py +++ b/tests/devices/system_tests/test_gridscan_system.py @@ -1,15 +1,14 @@ import bluesky.plan_stubs as bps import pytest -from bluesky.run_engine import RunEngine from dodal.devices.fast_grid_scan import ( - FastGridScan, - GridScanParams, + ZebraFastGridScan, + ZebraGridScanParams, set_fast_grid_scan_params, ) -def wait_for_fgs_valid(fgs_motors: FastGridScan, timeout=0.5): +def wait_for_fgs_valid(fgs_motors: ZebraFastGridScan, timeout=0.5): SLEEP_PER_CHECK = 0.1 times_to_check = int(timeout / SLEEP_PER_CHECK) for _ in range(times_to_check): @@ -22,60 +21,23 @@ def wait_for_fgs_valid(fgs_motors: FastGridScan, timeout=0.5): @pytest.fixture() -def fast_grid_scan(): - fast_grid_scan = FastGridScan(name="fast_grid_scan", prefix="BL03S-MO-SGON-01:FGS:") - yield fast_grid_scan +def zebra_fast_grid_scan(): + zebra_fast_grid_scan = ZebraFastGridScan( + name="zebra_fast_grid_scan", prefix="BL03S-MO-SGON-01:FGS:" + ) + yield zebra_fast_grid_scan @pytest.mark.s03 -def test_when_program_data_set_and_staged_then_expected_images_correct( - fast_grid_scan: FastGridScan, RE: RunEngine +async def test_when_program_data_set_and_staged_then_expected_images_correct( + zebra_fast_grid_scan: ZebraFastGridScan, RE, RunEngine ): RE( set_fast_grid_scan_params( - fast_grid_scan, - GridScanParams(transmission_fraction=0.01, x_steps=2, y_steps=2), + zebra_fast_grid_scan, + ZebraGridScanParams(transmission_fraction=0.01, x_steps=2, y_steps=2), ) ) - assert fast_grid_scan.expected_images.get() == 2 * 2 - fast_grid_scan.stage() - assert fast_grid_scan.position_counter.get() == 0 - - -@pytest.mark.s03 -def test_given_valid_params_when_kickoff_then_completion_status_increases_and_finishes( - fast_grid_scan: FastGridScan, RE: RunEngine -): - def set_and_wait_plan(fast_grid_scan: FastGridScan): - yield from set_fast_grid_scan_params( - fast_grid_scan, - GridScanParams(transmission_fraction=0.01, x_steps=3, y_steps=3), - ) - yield from wait_for_fgs_valid(fast_grid_scan) - - prev_current, prev_fraction = None, None - - def progress_watcher(*args, **kwargs): - nonlocal prev_current, prev_fraction - if "current" in kwargs.keys() and "fraction" in kwargs.keys(): - current, fraction = kwargs["current"], kwargs["fraction"] - if not prev_current: - prev_current, prev_fraction = current, fraction - else: - assert current > prev_current - assert fraction > prev_fraction - assert 0 < fraction < 1 - assert 0 < prev_fraction < 1 - - RE(set_and_wait_plan(fast_grid_scan)) - assert fast_grid_scan.position_counter.get() == 0 - - # S03 currently is giving 2* the number of expected images (see #13) - fast_grid_scan.expected_images.put(3 * 3 * 2) - - fast_grid_scan.kickoff() - complete_status = fast_grid_scan.complete() - complete_status.watch(progress_watcher) - complete_status.wait() - assert prev_current is not None - assert prev_fraction is not None + assert await zebra_fast_grid_scan.expected_images.get_value() == 2 * 2 + zebra_fast_grid_scan.stage() + assert await zebra_fast_grid_scan.position_counter.get_value() == 0 diff --git a/tests/devices/unit_tests/test_gridscan.py b/tests/devices/unit_tests/test_gridscan.py index 9c6c2c4ed9..1610a9c30b 100644 --- a/tests/devices/unit_tests/test_gridscan.py +++ b/tests/devices/unit_tests/test_gridscan.py @@ -1,17 +1,19 @@ +from asyncio import wait_for + import numpy as np import pytest from bluesky import plan_stubs as bps from bluesky import preprocessors as bpp from bluesky.run_engine import RunEngine -from mockito import mock, verify, when -from mockito.matchers import ANY, ARGS, KWARGS from ophyd.sim import make_fake_device from ophyd.status import DeviceStatus, Status -from ophyd.utils.errors import StatusTimeoutError +from ophyd_async.core import DeviceCollector, set_mock_value from dodal.devices.fast_grid_scan import ( - FastGridScan, - GridScanParams, + PandAFastGridScan, + PandAGridScanParams, + ZebraFastGridScan, + ZebraGridScanParams, set_fast_grid_scan_params, ) from dodal.devices.smargon import Smargon @@ -25,97 +27,67 @@ def discard_status(st: Status | DeviceStatus): @pytest.fixture -def fast_grid_scan(request): - FakeFastGridScan = make_fake_device(FastGridScan) - fast_grid_scan: FastGridScan = FakeFastGridScan( - name=f"test fake FGS: {request.node.name}" - ) - fast_grid_scan.scan_invalid.pvname = "" - yield fast_grid_scan +async def zebra_fast_grid_scan(): + async with DeviceCollector(mock=True): + zebra_fast_grid_scan = ZebraFastGridScan(name="fake_FGS", prefix="FGS") + return zebra_fast_grid_scan -def test_given_settings_valid_when_kickoff_then_run_started( - fast_grid_scan: FastGridScan, -): - when(fast_grid_scan.scan_invalid).get().thenReturn(False) - when(fast_grid_scan.position_counter).get().thenReturn(0) - - mock_run_set_status = mock() - when(fast_grid_scan.run_cmd).put(ANY).thenReturn(mock_run_set_status) - fast_grid_scan.status.subscribe = lambda func, **_: func(1) # type: ignore - status = fast_grid_scan.kickoff() - - status.wait() - assert status.exception() is None +@pytest.fixture +async def panda_fast_grid_scan(): + async with DeviceCollector(mock=True): + panda_fast_grid_scan = PandAFastGridScan(name="fake_PGS", prefix="PGS") - verify(fast_grid_scan.run_cmd).put(1) + return panda_fast_grid_scan -def test_waits_for_running_motion( - fast_grid_scan: FastGridScan, +@pytest.mark.parametrize( + "use_pgs", + [(False), (True)], +) +async def test_given_settings_valid_when_kickoff_then_run_started( + use_pgs, + zebra_fast_grid_scan: ZebraFastGridScan, + panda_fast_grid_scan: PandAFastGridScan, ): - when(fast_grid_scan.motion_program.running).get().thenReturn(1) - - fast_grid_scan.KICKOFF_TIMEOUT = 0.01 - - with pytest.raises(StatusTimeoutError): - status = fast_grid_scan.kickoff() - status.wait() - - fast_grid_scan.KICKOFF_TIMEOUT = 1 + grid_scan: ZebraFastGridScan | PandAFastGridScan = ( + panda_fast_grid_scan if use_pgs else zebra_fast_grid_scan + ) + set_mock_value(grid_scan.scan_invalid, False) + set_mock_value(grid_scan.position_counter, 0) + set_mock_value(grid_scan.status, 1) - mock_run_set_status = mock() - when(fast_grid_scan.run_cmd).put(ANY).thenReturn(mock_run_set_status) - fast_grid_scan.status.subscribe = lambda func, **_: func(1) # type: ignore + await grid_scan.kickoff() - when(fast_grid_scan.motion_program.running).get().thenReturn(0) - status = fast_grid_scan.kickoff() - status.wait() - verify(fast_grid_scan.run_cmd).put(1) + assert await grid_scan.run_cmd.get_value() == 1 -def run_test_on_complete_watcher( - RE: RunEngine, fast_grid_scan: FastGridScan, num_pos_1d, put_value, expected_frac +@pytest.mark.parametrize( + "use_pgs", + [(False), (True)], +) +async def test_waits_for_running_motion( + use_pgs, + zebra_fast_grid_scan: ZebraFastGridScan, + panda_fast_grid_scan: PandAFastGridScan, ): - RE( - set_fast_grid_scan_params( - fast_grid_scan, - GridScanParams( - x_steps=num_pos_1d, - y_steps=num_pos_1d, - transmission_fraction=0.01, - ), - ) + grid_scan: ZebraFastGridScan | PandAFastGridScan = ( + panda_fast_grid_scan if use_pgs else zebra_fast_grid_scan ) + set_mock_value(grid_scan.motion_program.running, 1) - complete_status = fast_grid_scan.complete() - watcher = mock() - complete_status.watch(watcher) - - fast_grid_scan.position_counter.sim_put(put_value) # type: ignore - verify(watcher).__call__( - *ARGS, - current=put_value, - target=num_pos_1d**2, - fraction=expected_frac, - **KWARGS, - ) - return complete_status + grid_scan.KICKOFF_TIMEOUT = 0.01 + with pytest.raises(TimeoutError): + await grid_scan.kickoff() -def test_when_new_image_then_complete_watcher_notified( - fast_grid_scan: FastGridScan, RE: RunEngine -): - status = run_test_on_complete_watcher(RE, fast_grid_scan, 2, 1, 3 / 4) - discard_status(status) + grid_scan.KICKOFF_TIMEOUT = 1 - -def test_given_0_expected_images_then_complete_watcher_correct( - fast_grid_scan: FastGridScan, RE: RunEngine -): - status = run_test_on_complete_watcher(RE, fast_grid_scan, 0, 1, 0) - discard_status(status) + set_mock_value(grid_scan.motion_program.running, 0) + set_mock_value(grid_scan.status, 1) + await grid_scan.kickoff() + assert await grid_scan.run_cmd.get_value() == 1 @pytest.mark.parametrize( @@ -126,44 +98,55 @@ def test_given_0_expected_images_then_complete_watcher_correct( ((7, 0, 5), 35), ], ) -def test_given_different_step_numbers_then_expected_images_correct( - fast_grid_scan: FastGridScan, steps, expected_images +async def test_given_different_step_numbers_then_expected_images_correct( + zebra_fast_grid_scan: ZebraFastGridScan, steps, expected_images ): - fast_grid_scan.x_steps.sim_put(steps[0]) # type: ignore - fast_grid_scan.y_steps.sim_put(steps[1]) # type: ignore - fast_grid_scan.z_steps.sim_put(steps[2]) # type: ignore - - assert fast_grid_scan.expected_images.get() == expected_images + set_mock_value(zebra_fast_grid_scan.x_steps, steps[0]) + set_mock_value(zebra_fast_grid_scan.y_steps, steps[1]) + set_mock_value(zebra_fast_grid_scan.z_steps, steps[2]) - -def test_given_invalid_image_number_then_complete_watcher_correct( - fast_grid_scan: FastGridScan, RE: RunEngine -): - complete_status = run_test_on_complete_watcher(RE, fast_grid_scan, 1, "BAD", None) - assert complete_status.exception() + assert await zebra_fast_grid_scan.expected_images.get_value() == expected_images -def test_running_finished_with_all_images_done_then_complete_status_finishes_not_in_error( - fast_grid_scan: FastGridScan, RE: RunEngine +@pytest.mark.parametrize( + "use_pgs", + [(False), (True)], +) +async def test_running_finished_with_all_images_done_then_complete_status_finishes_not_in_error( + use_pgs, + zebra_fast_grid_scan: ZebraFastGridScan, + panda_fast_grid_scan: PandAFastGridScan, + RE: RunEngine, ): num_pos_1d = 2 - RE( - set_fast_grid_scan_params( - fast_grid_scan, - GridScanParams( - transmission_fraction=0.01, x_steps=num_pos_1d, y_steps=num_pos_1d - ), + if use_pgs: + grid_scan = panda_fast_grid_scan + RE( + set_fast_grid_scan_params( + grid_scan, + PandAGridScanParams( + transmission_fraction=0.01, x_steps=num_pos_1d, y_steps=num_pos_1d + ), + ) ) - ) - - fast_grid_scan.status.sim_put(1) # type: ignore + else: + grid_scan = zebra_fast_grid_scan + RE( + set_fast_grid_scan_params( + grid_scan, + ZebraGridScanParams( + transmission_fraction=0.01, x_steps=num_pos_1d, y_steps=num_pos_1d + ), + ) + ) + set_mock_value(grid_scan.status, 1) - complete_status = fast_grid_scan.complete() + complete_status = grid_scan.complete() assert not complete_status.done - fast_grid_scan.position_counter.sim_put(num_pos_1d**2) # type: ignore - fast_grid_scan.status.sim_put(0) # type: ignore + set_mock_value(grid_scan.position_counter, num_pos_1d**2) + set_mock_value(grid_scan.status, 0) - complete_status.wait() + await wait_for(complete_status, 0.1) assert complete_status.done assert complete_status.exception() is None @@ -218,7 +201,7 @@ def test_within_limits_check(position, expected_in_limit): ) def test_scan_within_limits_1d(start, steps, size, expected_in_limits): motor_bundle = create_motor_bundle_with_limits(0.0, 10.0) - grid_params = GridScanParams( + grid_params = ZebraGridScanParams( transmission_fraction=0.01, x_start=start, x_steps=steps, x_step_size=size ) assert grid_params.is_valid(motor_bundle.get_xyz_limits()) == expected_in_limits @@ -236,7 +219,7 @@ def test_scan_within_limits_2d( x_start, x_steps, x_size, y1_start, y_steps, y_size, z1_start, expected_in_limits ): motor_bundle = create_motor_bundle_with_limits(0.0, 10.0) - grid_params = GridScanParams( + grid_params = ZebraGridScanParams( transmission_fraction=0.01, x_start=x_start, x_steps=x_steps, @@ -293,7 +276,7 @@ def test_scan_within_limits_3d( expected_in_limits, ): motor_bundle = create_motor_bundle_with_limits(0.0, 10.0) - grid_params = GridScanParams( + grid_params = ZebraGridScanParams( transmission_fraction=0.01, x_start=x_start, x_steps=x_steps, @@ -311,8 +294,26 @@ def test_scan_within_limits_3d( @pytest.fixture -def grid_scan_params(): - yield GridScanParams( +def zebra_grid_scan_params(): + yield ZebraGridScanParams( + transmission_fraction=0.01, + x_steps=10, + y_steps=15, + z_steps=20, + x_step_size=0.3, + y_step_size=0.2, + z_step_size=0.1, + x_start=0, + y1_start=1, + y2_start=2, + z1_start=3, + z2_start=4, + ) + + +@pytest.fixture +def panda_grid_scan_params(): + yield PandAGridScanParams( transmission_fraction=0.01, x_steps=10, y_steps=15, @@ -338,19 +339,25 @@ def grid_scan_params(): ], ) def test_given_x_y_z_out_of_range_then_converting_to_motor_coords_raises( - grid_scan_params: GridScanParams, grid_position + zebra_grid_scan_params: ZebraGridScanParams, + panda_grid_scan_params: PandAGridScanParams, + grid_position, ): with pytest.raises(IndexError): - grid_scan_params.grid_position_to_motor_position(grid_position) + zebra_grid_scan_params.grid_position_to_motor_position(grid_position) + with pytest.raises(IndexError): + panda_grid_scan_params.grid_position_to_motor_position(grid_position) def test_given_x_y_z_of_origin_when_get_motor_positions_then_initial_positions_returned( - grid_scan_params: GridScanParams, + zebra_grid_scan_params: ZebraGridScanParams, + panda_grid_scan_params: PandAGridScanParams, ): - motor_positions = grid_scan_params.grid_position_to_motor_position( - np.array([0, 0, 0]) - ) - assert np.allclose(motor_positions, np.array([0, 1, 4])) + motor_positions = [ + zebra_grid_scan_params.grid_position_to_motor_position(np.array([0, 0, 0])), + panda_grid_scan_params.grid_position_to_motor_position(np.array([0, 0, 0])), + ] + assert [np.allclose(position, np.array([0, 1, 4])) for position in motor_positions] @pytest.mark.parametrize( @@ -362,33 +369,57 @@ def test_given_x_y_z_of_origin_when_get_motor_positions_then_initial_positions_r ], ) def test_given_various_x_y_z_when_get_motor_positions_then_expected_positions_returned( - grid_scan_params: GridScanParams, grid_position, expected_x, expected_y, expected_z + zebra_grid_scan_params: ZebraGridScanParams, + panda_grid_scan_params: PandAGridScanParams, + grid_position, + expected_x, + expected_y, + expected_z, ): - motor_positions = grid_scan_params.grid_position_to_motor_position(grid_position) - np.testing.assert_allclose( - motor_positions, np.array([expected_x, expected_y, expected_z]) - ) + motor_positions = [ + zebra_grid_scan_params.grid_position_to_motor_position(grid_position), + panda_grid_scan_params.grid_position_to_motor_position(grid_position), + ] + [ + np.testing.assert_allclose( + position, np.array([expected_x, expected_y, expected_z]) + ) + for position in motor_positions + ] -def test_can_run_fast_grid_scan_in_run_engine(fast_grid_scan, RE: RunEngine): +@pytest.mark.parametrize( + "pgs", + [(False), (True)], +) +def test_can_run_fast_grid_scan_in_run_engine( + pgs, + zebra_fast_grid_scan: ZebraFastGridScan, + panda_fast_grid_scan: PandAFastGridScan, + RE: RunEngine, +): @bpp.run_decorator() def kickoff_and_complete(device): yield from bps.kickoff(device, group="kickoff") - device.status.sim_put(1) + set_mock_value(device.status, 1) yield from bps.wait("kickoff") yield from bps.complete(device, group="complete") - device.position_counter.sim_put(device.expected_images) - device.status.sim_put(0) + set_mock_value(device.position_counter, device.expected_images) + set_mock_value(device.status, 0) yield from bps.wait("complete") - RE(kickoff_and_complete(fast_grid_scan)) + RE(kickoff_and_complete(panda_fast_grid_scan)) if pgs else RE( + kickoff_and_complete(zebra_fast_grid_scan) + ) assert RE.state == "idle" def test_given_x_y_z_steps_when_full_number_calculated_then_answer_is_as_expected( - grid_scan_params: GridScanParams, + zebra_grid_scan_params: ZebraGridScanParams, + panda_grid_scan_params: PandAGridScanParams, ): - assert grid_scan_params.get_num_images() == 350 + assert zebra_grid_scan_params.get_num_images() == 350 + assert panda_grid_scan_params.get_num_images() == 350 @pytest.mark.parametrize( @@ -414,14 +445,14 @@ def test_given_x_y_z_steps_when_full_number_calculated_then_answer_is_as_expecte ) def test_non_test_integer_dwell_time(test_dwell_times, expected_dwell_time_is_integer): if expected_dwell_time_is_integer: - params = GridScanParams( + params = ZebraGridScanParams( dwell_time_ms=test_dwell_times, transmission_fraction=0.01, ) assert params.dwell_time_ms == test_dwell_times else: with pytest.raises(ValueError): - GridScanParams( + ZebraGridScanParams( dwell_time_ms=test_dwell_times, transmission_fraction=0.01, ) diff --git a/tests/devices/unit_tests/test_panda_gridscan.py b/tests/devices/unit_tests/test_panda_gridscan.py deleted file mode 100644 index 88456a2a3b..0000000000 --- a/tests/devices/unit_tests/test_panda_gridscan.py +++ /dev/null @@ -1,112 +0,0 @@ -import pytest -from bluesky import plan_stubs as bps -from bluesky import preprocessors as bpp -from bluesky.run_engine import RunEngine -from mockito import mock, verify, when -from mockito.matchers import ANY -from ophyd.sim import instantiate_fake_device, make_fake_device - -from dodal.devices.panda_fast_grid_scan import ( - PandAFastGridScan, - PandAGridScanParams, - set_fast_grid_scan_params, -) -from dodal.devices.smargon import Smargon - - -@pytest.fixture -def fast_grid_scan(request): - fast_grid_scan: PandAFastGridScan = instantiate_fake_device( - PandAFastGridScan, name=f"test fake FGS: {request.node.name}" - ) - fast_grid_scan.scan_invalid.pvname = "" - - yield fast_grid_scan - - -def test_given_settings_valid_when_kickoff_then_run_started( - fast_grid_scan: PandAFastGridScan, -): - when(fast_grid_scan.scan_invalid).get().thenReturn(False) - when(fast_grid_scan.position_counter).get().thenReturn(0) - - mock_run_set_status = mock() - when(fast_grid_scan.run_cmd).put(ANY).thenReturn(mock_run_set_status) - fast_grid_scan.status.subscribe = lambda func, **_: func(1) - - status = fast_grid_scan.kickoff() - - status.wait() - assert status.exception() is None - - verify(fast_grid_scan.run_cmd).put(1) - - -@pytest.mark.parametrize( - "steps, expected_images", - [ - ((10, 10, 0), 100), - ((30, 5, 10), 450), - ((7, 0, 5), 35), - ], -) -def test_given_different_step_numbers_then_expected_images_correct( - fast_grid_scan: PandAFastGridScan, steps, expected_images -): - fast_grid_scan.x_steps.sim_put(steps[0]) # type: ignore - fast_grid_scan.y_steps.sim_put(steps[1]) # type: ignore - fast_grid_scan.z_steps.sim_put(steps[2]) # type: ignore - - assert fast_grid_scan.expected_images.get() == expected_images - - -def test_running_finished_with_all_images_done_then_complete_status_finishes_not_in_error( - fast_grid_scan: PandAFastGridScan, RE: RunEngine -): - num_pos_1d = 2 - RE( - set_fast_grid_scan_params( - fast_grid_scan, - PandAGridScanParams( - x_steps=num_pos_1d, - y_steps=num_pos_1d, - transmission_fraction=0.01, - ), - ) - ) - - fast_grid_scan.status.sim_put(1) # type: ignore - - complete_status = fast_grid_scan.complete() - assert not complete_status.done - fast_grid_scan.position_counter.sim_put(num_pos_1d**2) # type: ignore - fast_grid_scan.status.sim_put(0) # type: ignore - - complete_status.wait() - - assert complete_status.done - assert complete_status.exception() is None - - -def create_motor_bundle_with_limits(low_limit, high_limit) -> Smargon: - FakeSmargon = make_fake_device(Smargon) - grid_scan_motor_bundle: Smargon = FakeSmargon(name="test fake Smargon") - grid_scan_motor_bundle.wait_for_connection() - for axis in [ - grid_scan_motor_bundle.x, - grid_scan_motor_bundle.y, - grid_scan_motor_bundle.z, - ]: - axis.low_limit_travel.sim_put(low_limit) # type: ignore - axis.high_limit_travel.sim_put(high_limit) # type: ignore - return grid_scan_motor_bundle - - -def test_can_run_fast_grid_scan_in_run_engine(fast_grid_scan, RE: RunEngine): - @bpp.run_decorator() - def kickoff_and_complete(device): - yield from bps.kickoff(device) - yield from bps.complete(device) - - RE(kickoff_and_complete(fast_grid_scan)) - assert RE.state == "idle"