From 555cbedbad249f055a2b977b69958f2fd5e8b45c Mon Sep 17 00:00:00 2001 From: wakonig_k Date: Wed, 29 Jul 2026 17:54:54 +0200 Subject: [PATCH] feat: add FlyScanV4 and HklScanV4 implementations --- addams_bec/scans/__init__.py | 2 + addams_bec/scans/fly_scan_v4.py | 280 ++++++++++++++++++++++++++++++++ addams_bec/scans/hkl_scan_v4.py | 204 +++++++++++++++++++++++ 3 files changed, 486 insertions(+) create mode 100644 addams_bec/scans/fly_scan_v4.py create mode 100644 addams_bec/scans/hkl_scan_v4.py diff --git a/addams_bec/scans/__init__.py b/addams_bec/scans/__init__.py index 1b876c6..2144c71 100644 --- a/addams_bec/scans/__init__.py +++ b/addams_bec/scans/__init__.py @@ -1,2 +1,4 @@ from .fly_scan import HklFlyScan +from .fly_scan_v4 import FlyScanV4 from .hkl_scan import HklScan +from .hkl_scan_v4 import HklScanV4 diff --git a/addams_bec/scans/fly_scan_v4.py b/addams_bec/scans/fly_scan_v4.py new file mode 100644 index 0000000..5d6304b --- /dev/null +++ b/addams_bec/scans/fly_scan_v4.py @@ -0,0 +1,280 @@ +""" +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_lib.scan_args import DefaultArgType, ScanArgument +from bec_server.scan_server.scans import ScanAbortion +from bec_server.scan_server.scans.scan_base import ScanBase, ScanType +from bec_server.scan_server.scans.scan_modifier import scan_hook + +from addams_bec.scans.scan_customization.scan_components import AddamsBecScanComponents + +logger = bec_logger.logger + + +class FlyScanV4(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 = "fly_scan_v4" + + gui_config = { + "Scan Parameters": ["diffract", "controller", "start", "stop", "optimize_profile"], + "Acquisition Parameters": ["exp_time"], + } + + def __init__( + #fmt: off + self, + + start: Annotated[list[float], ScanArgument(display_name="Start", description="Start.")], + stop: Annotated[list[float], ScanArgument(display_name="Stop", description="Stop.")], + points: Annotated[int, ScanArgument(display_name="Points", description="Number of points.")], + optimize_profile: Annotated[bool, ScanArgument(display_name="Optimize Profile", description="Optimize profile.")], + exp_time: DefaultArgType.ExposureTime, + diffract: Annotated[DeviceBase | None, ScanArgument(display_name="Diffract", description="Diffractometer")] = None, + controller: Annotated[DeviceBase | None, ScanArgument(display_name="Controller", description="Controller")] = None, + **kwargs, + #fmt: on + ): + """ + Scan implementation. + + Args: + diffract (DeviceBase): Diffractometer + controller (DeviceBase): Controller + start (list[float]): Start. + stop (list[float]): Stop. + optimize_profile (bool): Optimize profile. + exp_time (float): Exposure time in seconds + + Returns: + ScanReport + """ + super().__init__(**kwargs) + self.components = AddamsBecScanComponents(self) + self._baseline_readout_status = None + self.diffract, self.controller = self._get_endstation_devices(diffract, controller) + self.start = start + self.stop = stop + self.optimize_profile = optimize_profile + self.exp_time = exp_time + self.points = points + + self.update_scan_info(exp_time=exp_time, scan_report_devices=["h", "k", "l"]) + self.real_motors = [] + + @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. + """ + + # TODO: not sure if this needs to be fetched or can be cached? If cached, let's move it to init + self.real_motors = self.diffract.real_axes.get() + + hkls = np.linspace(self.start, self.stop, self.points).tolist() + + positions = self.diffract.angles_from_hkls(hkls) + + _positions = [] + for pos in positions: + _positions.append(pos[:-2]) + + self.positions = np.array(_positions) + self.components.check_limits(self.real_motors, self.positions) + + num_points = len(self.positions) + total_time = num_points * self.exp_time + + self.update_scan_info( + positions=self.positions, + num_points=len(self.positions), + num_monitored_readouts=len(self.positions), + ) + + self.actions.add_scan_report_instruction_scan_progress( + points=self.scan_info.num_monitored_readouts, show_table=False + ) + self._baseline_readout_status = self.actions.read_baseline_devices(wait=False) + + if self.optimize_profile: + self.controller.num_points.put(2) + self.controller.num_pulses.put(num_points) + self.controller.start_pulses.put(1) + self.controller.end_pulses.put(2) + self.controller.time_mode.put(1) + self.controller.times.put([total_time, total_time]) + + for index, axis_name in enumerate(self.real_motors): + getattr(self.controller, f"{axis_name}.positions").put( + (self.positions[0, index], self.positions[-1, index]) + ) + getattr(self.controller, f"{axis_name}.use_axis").put(1) + else: + self.controller.num_points.put(num_points) + self.controller.num_pulses.put(num_points) + self.controller.start_pulses.put(1) + self.controller.end_pulses.put(num_points) + self.controller.time_mode.put(0) + self.controller.fixed_time.put(self.exp_time) + + for index, axis_name in enumerate(self.real_motors): + getattr(self.controller, f"{axis_name}.positions").put(self.positions[:, index]) + getattr(self.controller, f"{axis_name}.use_axis").put(1) + + self.controller.build_profile().wait() + build_status = self.controller.build_status.get() + if build_status != 1: + raise ScanAbortion("Profile build failed") + + @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): + """ + 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() + + @scan_hook + def scan_core(self): + """ + Core scan logic to be executed during the scan. + This is where the main scan logic should be implemented. + """ + + profile_status = self.controller.execute_profile() + while not profile_status.done: + self.at_each_point() + + exit_status = self.controller.execute_status.get() + if exit_status != 1: + raise ScanAbortion("Profile execution failed") + + self.controller.readback_profile().wait() + readback_status = self.controller.readback_status.get() + if readback_status != 1: + raise ScanAbortion("Profile readback failed") + + # FIXME: This is where the old scan started to publish data. We should + # implement it on the device instead. + + @scan_hook + def at_each_point(self): + """ + Logic to be executed at each acquisition point during the scan. + """ + + @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 ########### + ####################################################### + + def _get_endstation_devices( + self, diffract: DeviceBase | None, controller: DeviceBase | None + ) -> list[DeviceBase | None]: + """ + Helper method to retrieve the endstation devices (diffractometer and controller) from the device manager. + This method allows for flexibility in how the devices are provided to the scan, either directly as arguments + or by looking them up in the device manager based on some criteria (e.g. device type, name, etc.). + """ + out = [diffract, controller] + if diffract is not None and controller is not None: + return out + + if diffract is None: + if "x04h" in self.device_manager.devices: + out[0] = self.device_manager.devices["x04h"] + elif "x04v" in self.device_manager.devices: + out[0] = self.device_manager.devices["x04v"] + else: + raise ValueError( + "Diffractometer device not provided and could not be found in the device manager." + ) + if controller is None: + if "profileMove" in self.device_manager.devices: + out[1] = self.device_manager.devices["profileMove"] + else: + raise ValueError( + "Controller device not provided and could not be found in the device manager." + ) + return out diff --git a/addams_bec/scans/hkl_scan_v4.py b/addams_bec/scans/hkl_scan_v4.py new file mode 100644 index 0000000..92570f3 --- /dev/null +++ b/addams_bec/scans/hkl_scan_v4.py @@ -0,0 +1,204 @@ +""" +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.scan_args import DefaultArgType, ScanArgument +from bec_server.scan_server.scans.scan_base import ScanBase, ScanType +from bec_server.scan_server.scans.scan_modifier import scan_hook + + +class HklScanV4(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 = "hkl_scan_v4" + + gui_config = { + "Scan Parameters": ["diffract", "start", "stop", "points", "relative"], + "Acquisition Parameters": ["exp_time"], + } + + def __init__( + self, + #fmt: off + diffract: Annotated[DeviceBase, ScanArgument(display_name="Diffract", description="Diffractometer")], + start: Annotated[list[float], ScanArgument(display_name="Start", description="Start position")], + stop: Annotated[list[float], ScanArgument(display_name="Stop", description="Stop position")], + points: Annotated[int, ScanArgument(display_name="Points", description="Number of points")], + exp_time: DefaultArgType.ExposureTime = 0, + **kwargs, + #fmt: on + ): + """ + Scan implementation. + + Args: + diffract (DeviceBase): Diffractometer + start (list[float]): Start position + stop (list[float]): Stop position + points (int): Number of points + exp_time (float): Exposure time in seconds + + Returns: + ScanReport + """ + super().__init__(**kwargs) + self._baseline_readout_status = None + self.diffract = diffract + self.start = start + self.stop = stop + self.points = points + self.exp_time = exp_time + self.motors = [diffract, self.dev.h, self.dev.k, self.dev.l] + + self.update_scan_info(exp_time=exp_time, scan_report_devices=["h", "k", "l"]) + self.real_motors = [] + + self.actions.set_device_readout_priority(self.motors, priority="monitored") + + @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. + """ + + # TODO: not sure if this needs to be fetched or can be cached? If cached, let's move it to init + self.real_motors = self.diffract.real_axes.get() + + hkls = np.linspace(self.start, self.stop, self.points).tolist() + + positions = self.diffract.angles_from_hkls(hkls) + + _positions = [] + for pos in positions: + _positions.append(pos[:-2]) + + self.positions = np.array(_positions) + + self.components.check_limits(self.real_motors, self.positions) + + self.update_scan_info( + positions=self.positions, + num_points=len(self.positions), + num_monitored_readouts=len(self.positions), + ) + + self.actions.add_scan_report_instruction_scan_progress( + points=self.scan_info.num_monitored_readouts, show_table=False + ) + 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): + """ + 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() + + @scan_hook + def scan_core(self): + """ + Core scan logic to be executed during the scan. + This is where the main scan logic should be implemented. + """ + self.components.step_scan( + motors=self.real_motors, positions=self.positions, at_each_point=self.at_each_point + ) + + @scan_hook + def at_each_point( + self, + motors: list[str | DeviceBase], + positions: np.ndarray, + last_positions: np.ndarray | None, + ): + """ + Logic to be executed at each acquisition point during the scan. + """ + self.components.step_scan_at_each_point(motors, positions, last_positions=last_positions) + + @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.