diff --git a/tests/tests_scans/test_hyst_scan.py b/tests/tests_scans/test_hyst_scan.py deleted file mode 100644 index aa35057..0000000 --- a/tests/tests_scans/test_hyst_scan.py +++ /dev/null @@ -1,132 +0,0 @@ -from unittest import mock - -import pytest - -from bec_server.scan_server.tests.fixtures import * # noqa: F401, F403 - -from xtreme_bec.scans.hyst_scan import HystScan - -# pylint: disable=missing-function-docstring -# pylint: disable=protected-access - - -@pytest.fixture -def hyst_scan(scan_assembler, device_manager_mock): - """HystScan instance with samx as field motor (flyer) and samy as mono.""" - return scan_assembler( - HystScan, - device_manager_mock.devices["samx"], # field_motor / flyer - 0.0, # start_field - 0.5, # end_field - device_manager_mock.devices["samy"], # mono / energy_motor - 600.0, # energy1 - 640.0, # energy2 - ramp_rate=3.0, # distinct from default_ramp_rate=2 for clearer assertions - parameter={"args": {}, "kwargs": {}}, - ) - - -@pytest.mark.parametrize( - "in_args, num_steps, reference_positions", - [ - ([600], 5, [600, 600, 600, 600, 600]), - ([600, 640], 6, [600, 640, 600, 640, 600, 640]), - ([600, 620, 640], 9, [600, 620, 640, 620, 600, 620, 640, 620, 600]), - ([1, 2, 3, 4], 10, [1, 2, 3, 4, 3, 2, 1, 2, 3, 4]), - ], -) -def test_get_next_scan_motor_position(in_args, num_steps, reference_positions, hyst_scan): - hyst_scan.energy_positions = in_args - gen = hyst_scan._get_next_scan_motor_position() - assert [next(gen) for _ in range(num_steps)] == reference_positions - - -def test_hyst_scan_init(scan_assembler, device_manager_mock): - field_motor = device_manager_mock.devices["samx"] - mono = device_manager_mock.devices["samy"] - scan = scan_assembler( - HystScan, - field_motor, 0.0, 0.5, - mono, 600.0, 640.0, - parameter={"args": {}, "kwargs": {}}, - ) - assert scan.flyer is field_motor - assert scan.energy_motor is mono - assert scan.flyer_positions == [0.0, 0.5] - assert scan.energy_positions == [600.0, 640.0] - assert scan.ramp_rate == HystScan.default_ramp_rate - assert scan.scan_motors == [mono, field_motor] - - -def test_hyst_scan_prepare_positions(hyst_scan): - with mock.patch.object(hyst_scan, "_check_limits"): - list(hyst_scan.prepare_positions()) - assert hyst_scan.positions == [[600.0], [640.0]] - assert hyst_scan.num_pos == 0 - - -def test_hyst_scan_scan_core(hyst_scan, device_manager_mock, ScanStubStatusMock): - """scan_core ramps the field while stepping the energy motor once per loop iteration.""" - field_motor = device_manager_mock.devices["samx"] - mono = device_manager_mock.devices["samy"] - - def fake_done(): - yield False # first while-check: enter loop - yield True # second while-check: exit loop - - def fake_set(*args, **kwargs): - yield "fake_set" - return ScanStubStatusMock(done_func=fake_done) - - def fake_rpc(*args, **kwargs): - yield "fake_rpc" - - def fake_read(*args, **kwargs): - yield "fake_read" - - # Replace connector with a MagicMock so send_client_info is freely callable - hyst_scan.connector = mock.MagicMock() - - devices_cls = type(device_manager_mock.devices) - with mock.patch.object(hyst_scan.stubs, "set", side_effect=fake_set) as mock_set: - with mock.patch.object( - hyst_scan.stubs, "send_rpc_and_wait", side_effect=fake_rpc - ) as mock_rpc: - with mock.patch.object(hyst_scan.stubs, "read", side_effect=fake_read) as mock_read: - with mock.patch.object(devices_cls, "monitored_devices", return_value=[]): - output = list(hyst_scan.scan_core()) - - assert output == [ - "fake_rpc", # send_rpc_and_wait(flyer, "ramprate.set", default_ramp_rate) - "fake_set", # set(flyer, start_field=0.0) - "fake_rpc", # send_rpc_and_wait(flyer, "ramprate.set", ramp_rate=3.0) - "fake_set", # set(flyer, end_field=0.5) ← status captured here - "fake_set", # set(energy_motor, energy1=600.0) — inside loop - "fake_read", # read(flyer, energy_motor) — inside loop - "fake_rpc", # send_rpc_and_wait(flyer, "ramprate.set", default_ramp_rate) — after loop - ] - assert hyst_scan.num_pos == 1 - assert hyst_scan.point_id == 1 - - # Ramp-rate sequence: brake to default → set scan rate → restore default - assert mock_rpc.call_args_list[0] == mock.call( - field_motor, "ramprate.set", HystScan.default_ramp_rate - ) - assert mock_rpc.call_args_list[1] == mock.call(field_motor, "ramprate.set", 3.0) - assert mock_rpc.call_args_list[2] == mock.call( - field_motor, "ramprate.set", HystScan.default_ramp_rate - ) - - # Field motor: move to start position, then begin flying to end (non-blocking) - assert mock_set.call_args_list[0] == mock.call(device=field_motor, value=0.0) - assert mock_set.call_args_list[1] == mock.call(device=field_motor, value=0.5, wait=False) - # Energy motor steps to first energy position inside the loop - assert mock_set.call_args_list[2] == mock.call(device=mono, value=600.0) - - # Read issued once (one loop iteration), with flyer and energy motor - mock_read.assert_called_once_with(device=[field_motor, mono], point_id=0) - - # Progress message sent for the energy step - hyst_scan.connector.send_client_info.assert_called_once_with( - "Moving mono to 600.0.", source="scan_server", expire=10 - ) diff --git a/xtreme_bec/devices/pgm_otf.py b/xtreme_bec/devices/pgm_otf.py index 56d2419..1bbfbeb 100644 --- a/xtreme_bec/devices/pgm_otf.py +++ b/xtreme_bec/devices/pgm_otf.py @@ -2,7 +2,14 @@ from __future__ import annotations from ophyd import Component as Cpt from ophyd import Device, EpicsSignal, EpicsSignalRO, Kind -from ophyd_devices import AsyncSignal, DeviceStatus, PSIDeviceBase, StatusBase, TransitionStatus +from ophyd_devices import ( + AsyncSignal, + DeviceStatus, + ProgressSignal, + PSIDeviceBase, + StatusBase, + TransitionStatus, +) class PGMOtfControl(Device): @@ -62,10 +69,8 @@ class PGMOtf(PSIDeviceBase, PGMOtfControl): docs="Aggregated fdata signal from the PGM OTF scan.", ) - # FIXME: Once we have the proper energy readback, we should set up the progress update - # to use the energy progress between e1 and e2. - # energy = Cpt(EpicsSignalRO, "ENERGY_RBV", kind=Kind.omitted, auto_monitor=True) - # progress = Cpt(ProgressSignal) + energy = Cpt(EpicsSignalRO, "X07MA-PGM:rbkenergy", kind=Kind.omitted, auto_monitor=True) + progress = Cpt(ProgressSignal) ######################################## # Beamline Specific Implementations # @@ -118,10 +123,14 @@ class PGMOtf(PSIDeviceBase, PGMOtfControl): def on_complete(self) -> DeviceStatus | StatusBase | None: """Called to inquire if a device has completed a scans.""" + self.complete_status.add_callback( + lambda status: self.progress.put(value=100, max_value=100, done=True) + ) return self.complete_status def on_kickoff(self) -> DeviceStatus | StatusBase | None: """Called to kickoff a device for a fly scan. Has to be called explicitly.""" + self.progress.put(value=0, max_value=100, done=False) self.acquire.put(1, use_complete=True) def on_stop(self) -> None: @@ -131,6 +140,21 @@ class PGMOtf(PSIDeviceBase, PGMOtfControl): def on_destroy(self) -> None: """Called when the device is destroyed. Cleanup resources here.""" + def _get_progress(self) -> float: + """ + Get the current progress of the OTF scan. The progress is calculated + based on the energy range (e2 - e1) and the current energy readback. + + Returns: + float: The current progress as a percentage (0 to 100). + """ + e1 = self.e1.get() + e2 = self.e2.get() + energy_range = abs(e2 - e1) + current_energy = self.current_energy.get() + progress = (current_energy - e1) / energy_range * 100 + return progress + def _update_data(self, value, obj, **kwargs): """ Callback function to update the aggregated data signals when the raw signals change. @@ -144,3 +168,4 @@ class PGMOtf(PSIDeviceBase, PGMOtfControl): self.idata.put(value) case "raw_fdata": self.fdata.put(value) + self.progress.put(value=self._get_progress(), max_value=100, done=False) diff --git a/xtreme_bec/scans/hyst_scan.py b/xtreme_bec/scans/hyst_scan.py index 0a2dee7..ea87188 100644 --- a/xtreme_bec/scans/hyst_scan.py +++ b/xtreme_bec/scans/hyst_scan.py @@ -1,112 +1,200 @@ +""" +Scan implementation. + +Scan procedure: + - prepare_scan + - open_scan + - stage + - pre_scan + - scan_core + - at_each_point (optionally called by scan_core) + - post_scan + - unstage + - close_scan + - on_exception (called if any exception is raised during the scan) +""" + from __future__ import annotations +from typing import Annotated + +import numpy as np from bec_lib.device import DeviceBase from bec_lib.logger import bec_logger -from bec_server.scan_server.scans import ScanBase +from bec_lib.scan_args import ScanArgument, Units +from bec_server.scan_server.scans import position_generators +from bec_server.scan_server.scans.scan_base import ScanBase, ScanType +from bec_server.scan_server.scans.scan_modifier import scan_hook logger = bec_logger.logger class HystScan(ScanBase): + # Scan Type: Hardware triggered or software triggered? + # If the main trigger and readout logic is done within the at_each_point method in scan_core, choose SOFTWARE_TRIGGERED. + # If the main trigger and readout logic is implemented on a device that is simply kicked off in this scan, choose HARDWARE_TRIGGERED. + # This primarily serves as information for devices: The device may need to react differently if a software trigger is expected + # for every point. + scan_type = ScanType.SOFTWARE_TRIGGERED + + # Scan name: This is the name of the scan, e.g. "line_scan". This is used for display purposes and to identify the scan type in user interfaces. + # Choose a descriptive name that does not conflict with existing scan names. + # It must be a valid Python identifier, that is, it can only contain letters, numbers, and underscores, and must not start with a number. scan_name = "hyst_scan" - required_kwargs = [] - scan_type = "step" + + gui_config = { + "Field Parameters": ["field_motor", "start_field", "end_field", "ramp_rate"], + "Energy Parameters": ["e1", "e2"], + } default_ramp_rate = 2 def __init__( + # fmt: off self, - field_motor: DeviceBase, - start_field: float, - end_field: float, - mono: DeviceBase, - energy1: float, - energy2: float, - ramp_rate: float = default_ramp_rate, + field_motor: Annotated[DeviceBase, ScanArgument(display_name="Field Motor", description="Field motor.")], + start_field: Annotated[float, ScanArgument(display_name="Start Field", description="Start field.", units=Units.T)], + end_field: Annotated[float, ScanArgument(display_name="End Field", description="End field.", units=Units.T)], + e1: Annotated[float, ScanArgument(display_name="Energy 1", description="Energy 1", units=Units.eV, ge=250, le=2000)], + e2: Annotated[float, ScanArgument(display_name="Energy 2", description="Energy 2", units=Units.eV, ge=250, le=2000)], + ramp_rate: Annotated[float, ScanArgument(display_name="Ramp Rate", description="Ramp rate.", units=Units.T / Units.min)] = default_ramp_rate, **kwargs, + # fmt: on ): """ - A hysteresis scan that steps the energy between energy1 and energy2 while - ramping the field from start_field to end_field and back to start_field. - The field is ramped at ramp_rate T/min. + Scan implementation. - scans.hyst_scan(field_motor, start_field, end_field, mono, energy1, energy2) - - Examples: - >>> scans.hyst_scan(dev.field_x, 0, 0.5, dev.mono, 600, 640, ramp_rate=2) + Args: + field_motor (DeviceBase): Field motor. + start_field (float): Start field. + end_field (float): End field. + e1 (float): Energy 1 + e2 (float): Energy 2 + ramp_rate (float): Ramp rate. + Returns: + ScanReport """ super().__init__(**kwargs) - self.axis = [] - self.flyer = field_motor - self.energy_motor = mono - self.scan_motors = [self.energy_motor, self.flyer] - self.flyer_positions = [start_field, end_field] - self.energy_positions = [energy1, energy2] - self._current_scan_motor_index = 0 - self._scan_motor_direction = 1 + self._baseline_readout_status = None + self.field_motor = field_motor + self.start_field = start_field + self.end_field = end_field + self.mono = self.dev["mono"] + self.e1 = e1 + self.e2 = e2 self.ramp_rate = ramp_rate + self.motors = [field_motor, self.mono] - def prepare_positions(self): - self.positions = [[pos] for pos in self.energy_positions] - self.num_pos = 0 - yield None - self._check_limits() + self.update_scan_info() - def _get_next_scan_motor_position(self): - positions = self.energy_positions - n = len(positions) + @scan_hook + def prepare_scan(self): + """ + Prepare the scan. This can include any steps that need to be executed + before the scan is opened, such as preparing the positions (if not done already) + or setting up the devices. + """ + self.actions.set_device_readout_priority(self.motors, priority="monitored") - if n == 1: - # Only one position, so just yield it indefinitely - while True: - yield positions[0] + # We don't know the number of points in advance, so we set it to 0 + self.actions.add_scan_report_instruction_scan_progress(points=0, show_table=False) - # Multiple positions, so yield them in a back-and-forth pattern - # For example, if positions = [600, 620, 640], yield: - # 600, 620, 640, 620, 600, 620, 640, 620, 600, ... - idx = 0 - direction = 1 + self._baseline_readout_status = self.actions.read_baseline_devices(wait=False) - while True: - yield positions[idx] - if idx == n - 1: - direction = -1 - elif idx == 0: - direction = 1 + @scan_hook + def open_scan(self): + """ + Open the scan. + This step must call self.actions.open_scan() to ensure that a new scan is + opened. Make sure to prepare the scan metadata before, either in + prepare_scan() or in open_scan() itself and call self.update_scan_info(...) + to update the scan metadata if needed. + """ + self.actions.open_scan() - idx += direction + @scan_hook + def stage(self): + """ + Stage the devices for the upcoming scan. The stage logic is typically + implemented on the device itself (i.e. by the device's stage method). + However, if there are any additional steps that need to be executed before + staging the devices, they can be implemented here. + """ + self.actions.stage_all_devices() + @scan_hook + def pre_scan(self): + """ + Pre-scan steps to be executed before the main scan logic. + This is typically the last chance to prepare the devices before the core scan + logic is executed. For example, this is a good place to initialize time-criticial + devices, e.g. devices that have a short timeout. + The pre-scan logic is typically implemented on the device itself. + """ + self.actions.pre_scan_all_devices() + + # Set the ramp rate to the default ramp rate to be a bit faster. + self.field_motor.ramprate.set(self.default_ramp_rate).wait() + + # move to the start position + self.field_motor.set(self.start_field).wait() + + # Now we set the proper ramp rate for the scan + self.field_motor.ramprate.set(self.ramp_rate).wait() + + @scan_hook def scan_core(self): - # yield from self._move_scan_motors_and_wait(self.positions[0]) - yield from self.stubs.send_rpc_and_wait(self.flyer, "ramprate.set", self.default_ramp_rate) - yield from self.stubs.set(device=self.flyer, value=self.flyer_positions[0]) + """ + Core scan logic to be executed during the scan. + This is where the main scan logic should be implemented. + """ - yield from self.stubs.send_rpc_and_wait(self.flyer, "ramprate.set", self.ramp_rate) + pos_generator = position_generators.oscillating_positions(values=[self.e1, self.e2]) - # send the slow motor on its way - status = yield from self.stubs.set( - device=self.flyer, value=self.flyer_positions[1], wait=False - ) - - pos_generator = self._get_next_scan_motor_position() - - dev = self.device_manager.devices + status = self.field_motor.set(self.end_field) while not status.done: val = next(pos_generator) - self.connector.send_client_info( - f"Moving mono to {val}.", source="scan_server", expire=10 - ) - yield from self.stubs.set(device=self.energy_motor, value=val) + self.mono.set(val).wait() + self.at_each_point() - monitored_devices = [_dev.name for _dev in dev.monitored_devices([])] - yield from self.stubs.read( - device=[self.flyer, self.scan_motors[0], *monitored_devices], point_id=self.point_id - ) - # time.sleep(1) - self.point_id += 1 - self.num_pos += 1 + @scan_hook + def at_each_point(self): + """ + Logic to be executed at each acquisition point during the scan. + """ + self.actions.read_monitored_devices() - yield from self.stubs.send_rpc_and_wait(self.flyer, "ramprate.set", self.default_ramp_rate) + @scan_hook + def post_scan(self): + """ + Post-scan steps to be executed after the main scan logic. + """ + self.actions.complete_all_devices() + self.field_motor.ramprate.set(self.default_ramp_rate).wait() - def move_to_start(self): - yield None + @scan_hook + def unstage(self): + """Unstage the scan by executing post-scan steps.""" + self.actions.unstage_all_devices() + + @scan_hook + def close_scan(self): + """Close the scan.""" + if self._baseline_readout_status is not None: + self._baseline_readout_status.wait() + self.actions.close_scan() + self.actions.check_for_unchecked_statuses() + + @scan_hook + def on_exception(self, exception: Exception): + """ + Handle exceptions that occur during the scan. + This is a good place to implement any cleanup logic that needs to be executed in case of an exception, + such as returning the devices to a safe state or moving the motors back to their starting position. + """ + + ####################################################### + ######### Helper methods for the scan logic ########### + ####################################################### + + # Implement scan-specific helper methods below. diff --git a/xtreme_bec/scans/otf_scan.py b/xtreme_bec/scans/otf_scan.py index eb72f91..815d9bf 100644 --- a/xtreme_bec/scans/otf_scan.py +++ b/xtreme_bec/scans/otf_scan.py @@ -1,59 +1,184 @@ +""" +OTF scan using the PGM + +Scan procedure: + - prepare_scan + - open_scan + - stage + - pre_scan + - scan_core + - at_each_point (optionally called by scan_core) + - post_scan + - unstage + - close_scan + - on_exception (called if any exception is raised during the scan) +""" + from __future__ import annotations -import time +from typing import Annotated from bec_lib.logger import bec_logger -from bec_server.scan_server.scans import SyncFlyScanBase +from bec_lib.scan_args import ScanArgument, Units +from bec_server.scan_server.scans.scan_base import ScanBase, ScanType +from bec_server.scan_server.scans.scan_modifier import scan_hook + +from xtreme_bec.scans.scan_customization.scan_components import XtremeBecScanComponents logger = bec_logger.logger -class OTFScan(SyncFlyScanBase): +class OTFScan(ScanBase): + # Scan Type: Hardware triggered or software triggered? + # If the main trigger and readout logic is done within the at_each_point method in scan_core, choose SOFTWARE_TRIGGERED. + # If the main trigger and readout logic is implemented on a device that is simply kicked off in this scan, choose HARDWARE_TRIGGERED. + # This primarily serves as information for devices: The device may need to react differently if a software trigger is expected + # for every point. + scan_type = ScanType.HARDWARE_TRIGGERED + + # Scan name: This is the name of the scan, e.g. "line_scan". This is used for display purposes and to identify the scan type in user interfaces. + # Choose a descriptive name that does not conflict with existing scan names. + # It must be a valid Python identifier, that is, it can only contain letters, numbers, and underscores, and must not start with a number. scan_name = "otf_scan" - required_kwargs = ["e1", "e2", "time"] - arg_input = {} - arg_bundle_size = {"bundle": len(arg_input), "min": None, "max": None} - def __init__(self, *args, parameter: dict = None, **kwargs): - """Scans the energy from e1 to e2 in