feat: migrate to v4 scans

xtreme_bec/scans/hyst_scan.py

@@ -1,112 +1,202 @@
+"""
+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": ["mono", "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)],
+        mono: Annotated[DeviceBase, ScanArgument(display_name="Mono", description="Mono.")],
+        e1: Annotated[float, ScanArgument(display_name="E1", description="Energy 1", units=Units.eV)],
+        e2: Annotated[float, ScanArgument(display_name="E2", description="Energy 2", units=Units.eV)],
+        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.
+            mono (DeviceBase): Mono.
+            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 = mono
+        self.e1 = e1
+        self.e2 = e2
         self.ramp_rate = ramp_rate
+        self.motors = [field_motor, 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.

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):
-    scan_name = "otf_scan"
-    required_kwargs = ["e1", "e2", "time"]
-    arg_input = {}
-    arg_bundle_size = {"bundle": len(arg_input), "min": None, "max": None}
+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
 
-    def __init__(self, *args, parameter: dict = None, **kwargs):
-        """Scans the energy from e1 to e2 in <time> minutes.
+    # 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_v4"
 
-        Examples:
-            >>> scans.otf_scan(e1=700, e2=740, time=4)
+    gui_config = {"Scan Parameters": ["e1", "e2", "time"]}
 
+    def __init__(
+        self,
+        #fmt: off
+        e1: Annotated[float, ScanArgument(display_name="E1", description="Start energy", units=Units.eV, ge=250, le=2000)],
+        e2: Annotated[float, ScanArgument(display_name="E2", description="End energy", units=Units.eV, ge=250, le=2000)],
+        time: Annotated[float, ScanArgument(display_name="Time", description="Travel time", units=Units.min, ge=0)],
+        **kwargs,
+        #fmt: on
+    ):
         """
-        super().__init__(parameter=parameter, **kwargs)
-        self.axis = []
-        self.scan_motors = []
-        self.num_pos = 0
-        self.mono = self.caller_kwargs.get("mono", "mono")
-        self.otf_device = self.caller_kwargs.get("otf", "otf")
+        OTF scan using the PGM
 
+        Args:
+            e1 (float): Start energy
+            e2 (float): End energy
+            time (float): Travel time
+
+        Returns:
+            ScanReport
+        """
+        super().__init__(**kwargs)
+        self.components = XtremeBecScanComponents(self)
+        self._baseline_readout_status = None
+        self.e1 = e1
+        self.e2 = e2
+        self.time = time
+        self.mono = self.dev["mono"]
+        self.otf = self.dev["otf"]
+
+        self.update_scan_info()
+
+    @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.add_scan_report_instruction_device_progress(device=self.otf)
+        self._baseline_readout_status = self.actions.read_baseline_devices(wait=False)
+
+    @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()
+
+    @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):
-        yield None
+        """
+        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()
 
-    def scan_report_instructions(self):
-        yield from self.stubs.scan_report_instruction({"device_progress": [self.otf_device]})
-
-    @property
-    def monitor_sync(self):
-        return self.otf_device
+        # Move to the first energy
+        self.mono.set(self.e1).wait()
 
+    @scan_hook
     def scan_core(self):
-        yield from self.stubs.set(device=self.mono, value=self.caller_kwargs["e1"], wait=True)
-        yield from self.stubs.kickoff(
-            device=self.otf_device,
-            parameter={
-                key: val for key, val in self.caller_kwargs.items() if key in ["e1", "e2", "time"]
-            },
-            wait=True,
-        )
-        status = yield from self.stubs.complete(device=self.otf_device)
+        """
+        Core scan logic to be executed during the scan.
+        This is where the main scan logic should be implemented.
+        """
+
+        e1_status = self.otf.e1.set(self.e1)
+        e2_status = self.otf.e2.set(self.e2)
+        time_status = self.otf.time.set(self.time)
+
+        e1_status.wait()
+        e2_status.wait()
+        time_status.wait()
+
+        self.actions.kickoff(device=self.otf)
+
+        status = self.actions.complete(device=self.otf, wait=False)
 
         while not status.done:
-            yield from self.stubs.read(group="monitored", wait=True)
-            progress = self.stubs.get_device_progress(
-                device=self.otf_device, RID=self.metadata["RID"]
-            )
-            if progress:
-                self.num_pos = progress
-            time.sleep(1)
+            self.at_each_point()
+
+    @scan_hook
+    def at_each_point(self):
+        """
+        Logic to be executed at each acquisition point during the scan.
+        """
+        self.actions.read_monitored_devices()
+
+    @scan_hook
+    def post_scan(self):
+        """
+        Post-scan steps to be executed after the main scan logic.
+        """
+        self.actions.complete_all_devices()
+
+    @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.