From fa1e092c9c527ace22d97052c9b2e3841dc7b0fa Mon Sep 17 00:00:00 2001 From: gac-x04sa Date: Sun, 27 Apr 2025 20:42:59 +0200 Subject: [PATCH] generic line-/gridscan with profileMove --- addams_bec/scans/__init__.py | 3 +- addams_bec/scans/profilemove_scan.py | 233 +++++++++++++++++++++++++++ 2 files changed, 235 insertions(+), 1 deletion(-) create mode 100644 addams_bec/scans/profilemove_scan.py diff --git a/addams_bec/scans/__init__.py b/addams_bec/scans/__init__.py index e82e226..7eb92e6 100644 --- a/addams_bec/scans/__init__.py +++ b/addams_bec/scans/__init__.py @@ -1,3 +1,4 @@ from .hkl_scan import HklScan from .fly_scan import HklFlyScan -from .mesh_scan import HklMeshScan, HklMeshFlyScan \ No newline at end of file +from .mesh_scan import HklMeshScan, HklMeshFlyScan +from .profilemove_scan import LineProfileMove, GridProfileMove diff --git a/addams_bec/scans/profilemove_scan.py b/addams_bec/scans/profilemove_scan.py new file mode 100644 index 0000000..1d81fe5 --- /dev/null +++ b/addams_bec/scans/profilemove_scan.py @@ -0,0 +1,233 @@ +import numpy + +from bec_lib import messages +from bec_lib.endpoints import MessageEndpoints +from bec_lib.devicemanager import DeviceManagerBase +from bec_lib.logger import bec_logger +from bec_server.scan_server.scans import ScanAbortion, ScanArgType, ScanBase, AsyncFlyScanBase + +logger = bec_logger.logger + +class LineProfileMove(AsyncFlyScanBase): + scan_name = 'line_profilemove' + scan_type = 'fly' + arg_input = { + 'motor': ScanArgType.DEVICE, + 'start': ScanArgType.FLOAT, + 'stop': ScanArgType.FLOAT, + } + arg_bundle_size = {"bundle": len(arg_input), "min": 1, "max": None} + required_kwargs = ['controller', 'points'] + use_scan_progress_report = False + + def __init__(self, *args, controller: str, points: int, **kwargs): + self.controller = controller + self.points = self.num_pos = points + super().__init__(**kwargs) + + self.all_motors = self.device_manager.devices[self.controller].axes.get() + if not all(m in self.all_motors for m in self.scan_motors): + raise ValueError(f"Scan motors {self.scan_motors} not in {self.all_motors}") + + def update_readout_priority(self): + self.readout_priority['async'].extend(self.scan_motors) + + def update_scan_motors(self): + ScanBase.update_scan_motors(self) + + def prepare_positions(self): + """prepare the positions for the scan""" + self._calculate_positions() + self._optimize_trajectory() + self.num_pos = len(self.positions) * self.burst_at_each_point + yield from self._set_position_offset() + self._check_limits() + + def _calculate_positions(self): + axis = [] + for _, val in self.caller_args.items(): + ax_pos = numpy.linspace(val[0], val[1], self.points, dtype=float) + axis.append(ax_pos) + self.positions = numpy.array(list(zip(*axis)), dtype=float) + + def scan_core(self): + + total_time = self.exp_time * self.points + yield from self.stubs.send_rpc_and_wait(self.controller, 'num_points.put', 2) + yield from self.stubs.send_rpc_and_wait(self.controller, 'num_pulses.put', self.points) + yield from self.stubs.send_rpc_and_wait(self.controller, 'start_pulses.put', 1) + yield from self.stubs.send_rpc_and_wait(self.controller, 'end_pulses.put', 2) + yield from self.stubs.send_rpc_and_wait(self.controller, 'time_mode.put', 1) + yield from self.stubs.send_rpc_and_wait(self.controller, 'times.put', [total_time, total_time]) + + for axis_name in self.all_motors: + if axis_name not in self.scan_motors: + yield from self.stubs.send_rpc_and_wait(self.controller, f'{axis_name}.use_axis.put', 0) + else: + index = self.scan_motors.index(axis_name) + yield from self.stubs.send_rpc_and_wait(self.controller, + f'{axis_name}.positions.put', + (self.caller_args[axis_name][0], self.caller_args[axis_name][1])) + yield from self.stubs.send_rpc_and_wait(self.controller, f'{axis_name}.use_axis.put', 1) + + yield from self.stubs.send_rpc_and_wait(self.controller, 'build_profile') + build_status = yield from self.stubs.send_rpc_and_wait(self.controller, 'build_status.get') + if build_status != 1: + raise ScanAbortion('Profile build failed') + + yield from self.stubs.send_rpc_and_wait(self.controller, 'execute_profile') + execute_status = yield from self.stubs.send_rpc_and_wait(self.controller, 'execute_status.get') + if execute_status != 1: + raise ScanAbortion('Profile execute failed') + + yield from self.stubs.send_rpc_and_wait(self.controller, 'readback_profile') + readback_status = yield from self.stubs.send_rpc_and_wait(self.controller,'readback_status.get') + if readback_status != 1: + raise ScanAbortion('Profile readback failed') + + for index, axis_name in enumerate(self.scan_motors): + readbacks = yield from self.stubs.send_rpc_and_wait(self.controller, f'{axis_name}.readbacks.get') + self._publish_readbacks(axis_name, readbacks) + + logger.success(f'{self.scan_name} finished') + + def _publish_readbacks(self, device, readbacks): + metadata = {"async_update": "append", "max_shape": [None, None]} + msg = messages.DeviceMessage( + signals={device: {'value': readbacks} }, metadata=metadata + ) + + self.stubs.connector.xadd( + topic=MessageEndpoints.device_async_readback( + scan_id=self.scan_id, device=device + ), + msg_dict={"data": msg}, + expire=1800, + ) + + +class GridProfileMove(AsyncFlyScanBase): + scan_name = 'grid_profilemove' + scan_type = 'fly' + arg_input = { + 'motor': ScanArgType.DEVICE, + 'start': ScanArgType.FLOAT, + 'stop': ScanArgType.FLOAT, + 'points': ScanArgType.INT, + } + arg_bundle_size = {"bundle": len(arg_input), "min": 2, "max": 2} + required_kwargs = ['controller'] + use_scan_progress_report = False + + def __init__(self, *args, controller: str, **kwargs): + self.controller = controller + super().__init__(**kwargs) + + self.all_motors = self.device_manager.devices[self.controller].axes.get() + if not all(m in self.all_motors for m in self.scan_motors): + raise ValueError(f"Scan motors {self.scan_motors} not in {self.all_motors}") + + def update_readout_priority(self): + self.readout_priority['async'].extend(self.scan_motors) + + def update_scan_motors(self): + ScanBase.update_scan_motors(self) + + def prepare_positions(self): + """prepare the positions for the scan""" + self._calculate_positions() + self._optimize_trajectory() + self.num_pos = len(self.positions) * self.burst_at_each_point + yield from self._set_position_offset() + self._check_limits() + + def _calculate_positions(self): + inner_motor, outer_motor = self.scan_motors + outer_start, outer_stop, outer_points = self.caller_args[outer_motor] + inner_start, inner_stop, inner_points = self.caller_args[inner_motor] + + self.positions = [] + outer_vect = numpy.linspace(outer_start, outer_stop, outer_points, dtype=float) + for i, outer in enumerate(outer_vect): + if i % 2 == 0: + inner_vect = numpy.linspace(inner_start, inner_stop, inner_points, dtype=float) + else: + inner_vect = numpy.linspace(inner_stop, inner_start, inner_points, dtype=float) + + for inner in inner_vect: + self.positions.append((inner, outer)) + + def scan_core(self): + inner_motor, outer_motor = self.scan_motors + outer_start, outer_stop, outer_points = self.caller_args[outer_motor] + inner_start, inner_stop, inner_points = self.caller_args[inner_motor] + + num_pos = 0 + outer_vect = numpy.linspace(outer_start, outer_stop, outer_points, dtype=float) + for i, outer in enumerate(outer_vect): + if i % 2 == 0: + start, stop = inner_start, inner_stop + else: + start, stop = inner_stop, inner_start + + # Move outer motor to start position + yield from self.stubs.set(device=outer_motor, value=outer) + outer_readback = yield from self.stubs.send_rpc_and_wait(outer_motor, "position") + + inner_time = self.exp_time * inner_points + yield from self.stubs.send_rpc_and_wait(self.controller, 'num_points.put', 2) + yield from self.stubs.send_rpc_and_wait(self.controller, 'num_pulses.put', inner_points) + yield from self.stubs.send_rpc_and_wait(self.controller, 'start_pulses.put', 1) + yield from self.stubs.send_rpc_and_wait(self.controller, 'end_pulses.put', 2) + yield from self.stubs.send_rpc_and_wait(self.controller, 'time_mode.put', 1) + yield from self.stubs.send_rpc_and_wait(self.controller, 'times.put', [inner_time, inner_time]) + + for axis_name in self.all_motors: + if axis_name == inner_motor: + yield from self.stubs.send_rpc_and_wait(self.controller, + f'{axis_name}.positions.put', + (start, stop)) + yield from self.stubs.send_rpc_and_wait(self.controller, f'{axis_name}.use_axis.put', 1) + else: + yield from self.stubs.send_rpc_and_wait(self.controller, f'{axis_name}.use_axis.put', 0) + + yield from self.stubs.send_rpc_and_wait(self.controller, 'build_profile') + build_status = yield from self.stubs.send_rpc_and_wait(self.controller, 'build_status.get') + if build_status != 1: + raise ScanAbortion('Profile build failed') + + yield from self.stubs.send_rpc_and_wait(self.controller, 'execute_profile') + execute_status = yield from self.stubs.send_rpc_and_wait(self.controller, 'execute_status.get') + if execute_status != 1: + raise ScanAbortion('Profile execute failed') + + yield from self.stubs.send_rpc_and_wait(self.controller, 'readback_profile') + readback_status = yield from self.stubs.send_rpc_and_wait(self.controller,'readback_status.get') + if readback_status != 1: + raise ScanAbortion('Profile readback failed') + + readbacks = yield from self.stubs.send_rpc_and_wait(self.controller, f'{inner_motor}.readbacks.get') + self._publish_readbacks(inner_motor, readbacks) + + readbacks = [outer_readback] * len(readbacks) + self._publish_readbacks(outer_motor, readbacks) + + # motor readbacks have more points than generated pulses + num_pos += len(readbacks) + + self.num_pos = num_pos + logger.success(f'{self.scan_name} finished') + + def _publish_readbacks(self, device, readbacks): + metadata = {"async_update": "append", "max_shape": [None]} + msg = messages.DeviceMessage( + signals={device: {'value': readbacks} }, metadata=metadata + ) + + self.stubs.connector.xadd( + topic=MessageEndpoints.device_async_readback( + scan_id=self.scan_id, device=device + ), + msg_dict={"data": msg}, + expire=1800, + )