diff --git a/csaxs_bec/scans/omny_fermat_scan_v4.py b/csaxs_bec/scans/omny_fermat_scan_v4.py new file mode 100644 index 0000000..7b3678f --- /dev/null +++ b/csaxs_bec/scans/omny_fermat_scan_v4.py @@ -0,0 +1,328 @@ +""" +OMNY Fermat's spiral scan + +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 import messages +from bec_lib.logger import bec_logger +from bec_lib.scan_args import DefaultArgType, ScanArgument, Units +from bec_server.scan_server.scans import MessageEndpoints, 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 OmnyFermatScanV4(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 = "omny_fermat_scan_v4" + + gui_config = { + "Scan Parameters": [ + "fovx", + "fovy", + "cenx", + "ceny", + "step", + "zshift", + "angle", + "corridor_size", + ], + "Acquisition Parameters": ["exp_time", "frames_per_trigger", "readout_time"], + } + + def __init__( + # fmt: off + self, + fovx: Annotated[float, ScanArgument(display_name="Fovx", description="FOV in the piezo plane (i.e. piezo range).", units=Units.µm, gt=0, lt=200)], + fovy: Annotated[float, ScanArgument(display_name="Fovy", description="FOV in the piezo plane (i.e. piezo range).", units=Units.µm, gt=0, lt=100)], + cenx: Annotated[float, ScanArgument(display_name="Cenx", description="Center position in x.", units=Units.µm)], + ceny: Annotated[float, ScanArgument(display_name="Ceny", description="Center position in y.", units=Units.µm)], + step: Annotated[float, ScanArgument(display_name="Step", description="Step size.", units=Units.µm)], + zshift: Annotated[float, ScanArgument(display_name="Zshift", description="Shift in z.", units=Units.µm)], + angle: Annotated[float | None, ScanArgument(display_name="Angle", description="Rotation angle (will rotate first)", units=Units.deg)] = None, + corridor_size: Annotated[float, ScanArgument(display_name="Corridor Size", description="Corridor size for the corridor optimization. ", units=Units.µm)] = 3, + exp_time: DefaultArgType.ExposureTime = 0, + frames_per_trigger: DefaultArgType.FramesPerTrigger = 1, + readout_time: DefaultArgType.ReadoutTime = 0, + **kwargs, + # fmt: on + ): + """ + OMNY Fermat's spiral scan + + Args: + fovx (float): FOV in the piezo plane (i.e. piezo range). + fovy (float): FOV in the piezo plane (i.e. piezo range). + cenx (float): Center position in x. + ceny (float): Center position in y. + step (float): Step size. + zshift (float): Shift in z. + angle (float | None): Rotation angle (will rotate first) + corridor_size (float): Corridor size for the corridor optimization. + exp_time (float): Exposure time in seconds + frames_per_trigger (int): Number of frames per trigger for devices that support configurable frame counts per trigger. + readout_time (float): Configuration for devices that support configurable readout times. + + Returns: + ScanReport + """ + super().__init__(**kwargs) + self._baseline_readout_status = None + self.fovx = fovx + self.fovy = fovy + self.cenx = cenx + self.ceny = ceny + self.step = step + self.zshift = zshift + self.angle = angle + self.corridor_size = corridor_size + self.exp_time = exp_time + self.frames_per_trigger = frames_per_trigger + self.readout_time = readout_time + + if self.zshift > 100: + logger.warning("The zshift is larger than 100 um. It will be limited to 100 um.") + self.zshift = 100 + elif self.zshift < -100: + logger.warning("The zshift is smaller than -100 um. It will be limited to -100 um.") + self.zshift = -100 + + self.update_scan_info( + exp_time=exp_time, frames_per_trigger=frames_per_trigger, readout_time=readout_time + ) + + self.omny_rotation_status = None + + @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.positions = self.get_omny_fermat_spiral_pos( + -np.abs(self.fovx / 2), + np.abs(self.fovx / 2), + -np.abs(self.fovy / 2), + np.abs(self.fovy / 2), + step=self.step, + spiral_type=0, + center=False, + ) + + # TODO: If necessary, prepare the positions for the scan here and update the scan + # info with the updated num_points and num_monitored_readouts using self.update_scan_info(). + # Take the following example of a hexagonal scan as a reference: + + # self.positions = position_generators.hex_grid_2d( + # [ + # (self.start_motor1, self.stop_motor1, self.step_motor1), + # (self.start_motor2, self.stop_motor2, self.step_motor2), + # ], + # snaked=self.snaked, + # ) + + # if self.relative: + # self.start_positions = self.components.get_start_positions(self.motors) + # self.positions += self.start_positions + + # self.components.check_limits(self.motors, self.positions) + + # self.update_scan_info( + # positions=self.positions, + # num_points=len(self.positions), + # num_monitored_readouts=len(self.positions) * self.burst_at_each_point, + # ) + + # Delete the comment if not needed. + + 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. + """ + + # TODO: Implement the core scan logic here. If applicable, we recommend to call + # self.at_each_point() at each acquisition point in the scan to allow scan modifiers + # to extend the behavior of the scan at each point. + # Delete the comment if not needed. + + @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 reverse_trajectory(self): + """ + Reverse the trajectory. Every other scan should be reversed to + shorten the movement time. In order to keep the last state, even if the + server is restarted, the state is stored in a global variable in redis. + """ + msg = self.redis_connector.get(MessageEndpoints.global_vars("reverse_omny_trajectory")) + if msg: + val = msg.content.get("value", False) + else: + val = False + self.redis_connector.set( + MessageEndpoints.global_vars("reverse_omny_trajectory"), + messages.VariableMessage(value=(not val)), + ) + return val + + def get_omny_fermat_spiral_pos( + self, + m1_start: float, + m1_stop: float, + m2_start: float, + m2_stop: float, + step: float = 1, + spiral_type: int = 0, + center: bool = False, + ): + """ + Calculate positions for a Fermat spiral scan. + + Args: + m1_start(float): start position in m1 + m1_stop(float): stop position in m1 + m2_start(float): start position in m2 + m2_stop(float): stop position in m2 + step(float): stepsize + spiral_type(int): 0 for traditional Fermat spiral + center(bool): whether to include the center position + + Returns: + positions(array): positions + """ + positions = [] + phi = 2 * np.pi * ((1 + np.sqrt(5)) / 2.0) + spiral_type * np.pi + + start = int(not center) + + length_axis1 = np.abs(m1_stop - m1_start) + length_axis2 = np.abs(m2_stop - m2_start) + n_max = int(length_axis1 * length_axis2 * 3.2 / step / step) + + z_pos = self.zshift + + for ii in range(start, n_max): + radius = step * 0.57 * np.sqrt(ii) + # FOV is restructed below at check pos in range + if abs(radius * np.sin(ii * phi)) > length_axis1 / 2: + continue + if abs(radius * np.cos(ii * phi)) > length_axis2 / 2: + continue + x = radius * np.sin(ii * phi) + y = radius * np.cos(ii * phi) + positions.append([x + self.cenx, y + self.ceny, z_pos]) + left_lower_corner = [ + min(m1_start, m1_stop) + self.cenx, + min(m2_start, m2_stop) + self.ceny, + z_pos, + ] + right_upper_corner = [ + max(m1_start, m1_stop) + self.cenx, + max(m2_start, m2_stop) + self.ceny, + z_pos, + ] + positions.append(left_lower_corner) + positions.append(right_upper_corner) + return np.array(positions)