diff --git a/debye_bec/bec_ipython_client/plugins/auto_gain.py b/debye_bec/bec_ipython_client/plugins/auto_gain.py new file mode 100644 index 0000000..a5ccd5c --- /dev/null +++ b/debye_bec/bec_ipython_client/plugins/auto_gain.py @@ -0,0 +1,259 @@ +"""Module to automatically set the gains for the selected amplifiers""" + +import builtins +import time +from bisect import bisect_right + +import xraydb +from bec_lib import bec_logger + +from ...devices.absorber import STATUS as ABS_STATUS +from ...devices.eh_shutter import STATUS as EH_PH_STATUS +from ...devices.ionization_chambers.ionization_chamber_enums import AmplifierEnable +from ...devices.nidaq.nidaq_enums import NidaqState +from ...devices.op_shutter import STATUS as OP_PH_STATUS + +logger = bec_logger.logger + +EMIN = -100 +EMAX = 200 + +MIN_RING_CURRENT = 5 # Minimum ring current to use auto-gain +NOMINAL_RING_CURRENT = 400 # Nominal ring current of SLS2 + +MONO_VELOCITY = 20 # Move velocity in deg/s +TIMEOUT_MONO_PV = 5 # Timeout to set a PV on the mono +TIMEOUT_MONO_MOVE = 30 # Timeout to finish a movement on the mono + +AVAILABLE_GAINS = [1e6, 1e7, 5e7, 1e8, 1e9] # ascending order +MIN_SIGNAL = 0.05 # Minimum signal to count as valid signal +FULL_SCALE_V = 10.0 # NIDAQ AI full-scale range +SAFETY_MARGIN = 0.9 # keep max signal under 90% of full scale + + +class AutoGainError(Exception): + """AutoGain specific error""" + + +class AutoGain: + """Module to automatically set the gains for the selected amplifiers""" + + def __init__(self): + dev = builtins.__dict__.get("dev") + bec = builtins.__dict__.get("bec") + if dev is None: + raise AutoGainError("Did not get dev") + if bec is None: + raise AutoGainError("Did not get bec") + self.dev = dev + self.bec = bec + + def start( + self, + element: str, + edge: str, + amplifier: list[str] | None = None, + comp_ring_current: bool = True, + ) -> None: + """Start the auto-gain sequence. Measure the signals of the specified + amplifiers and set the gains accordingly. Makes sure there is actually beam available. + + Args: + element(str): Element which defines the energy at which the gain will be set, e.g. 'Cu' + edge(str): Corresponding edge, e.g. 'L1' + amplifier(list[str]): Amplifiers where auto-gain should be applied to + Defaults to all amplifiers -> ['ic0', 'ic1', 'ic2', 'pips'] + comp_ring_current(bool): Respects the current ring current and calculates the gain(s) + for a nominal ring current of 400 mA. Defaults to True + + Raises: + If NIDAQ is not in measurement mode + If Ring current is below 5 mA + If Absorber, OP Photon Shutter or EH Photon Shutter is closed + If a bec scan is running + If the energy for the supplied element/edge cannot be found + If the energy is outside the movement range of the monochromator + If a selected amplifier is switched off + If the high voltage of a selected ionization chamber is not enabled or < 1000 V + If the gas filling of a selected ionization chamber is not OK + If the ring current drops to 0 mA during the measurement (beamdump) + """ + + if amplifier is None: + amplifier = ["ic0", "ic1", "ic2", "pips"] + + # Make sure NIDAQ is in standby mode + if self.dev.nidaq.state.get() != NidaqState.STANDBY: + raise AutoGainError("NIDAQ was not in Standby mode, cannot proceed.") + + # Check for beam availability + if self._get_ring_current() < MIN_RING_CURRENT: + raise AutoGainError(f"Ring current is below {MIN_RING_CURRENT} mA") + if self.dev.abs.status.get() != ABS_STATUS.OPEN: + raise AutoGainError("Absorber is closed, no beam") + if self.dev.op_sh.status.get() != OP_PH_STATUS.NOT_CLOSED: + raise AutoGainError("OP Photon Shutter is closed, no beam") + if self.dev.eh_sh.status.get() != EH_PH_STATUS.NOT_CLOSED: + raise AutoGainError("EH Photon Shutter is closed, no beam") + + # Check if no scan is running + scan_id = self.bec.queue.scan_storage.current_scan_id + if len(scan_id) > 0: + raise AutoGainError(f"Scan with ID {scan_id} is currently running, cannot continue") + + # Get edge energy + energy = xraydb.xray_edge(element, edge, True) + if energy is None: + raise ValueError(f"Could not find edge energy for element/edge {element}/{edge}") + emin = energy + EMIN + emax = energy + EMAX + + # Check range of mono + low_limit = self.dev.mo1_bragg.low_lim.get() + high_limit = self.dev.mo1_bragg.high_lim.get() + if emin < low_limit or emax > high_limit: + raise ValueError( + f"Chosen element/edge {element}/{edge} with edge energy of {energy}" + + " is outside of accessible range of monochromator " + + f"{low_limit:.1f} eV - {high_limit:.1f} eV" + ) + + # Map amplifier names to their NIDAQ channels + channel_map = { + "ic0": {"signal": self.dev.nidaq.ai0, "self.dev": self.dev.ic0}, + "ic1": {"signal": self.dev.nidaq.ai2, "self.dev": self.dev.ic1}, + "ic2": {"signal": self.dev.nidaq.ai4, "self.dev": self.dev.ic2}, + "pips": {"signal": self.dev.nidaq.ai6, "self.dev": self.dev.pips}, + } + active_channels = {name: ch for name, ch in channel_map.items() if name in amplifier} + + # Check if amplifieres are switched on + for name, ch in active_channels.items(): + if ch["self.dev"].amp.cOnOff.get() != AmplifierEnable.ON: + raise AutoGainError(f"Amplifier of self.device {name} is not enabled") + + # Check high voltage on ionization chambers + for name, ch in active_channels.items(): + if name != "pips": + if ch["self.dev"].hv_en.ena.get() != 1: + raise AutoGainError(f"High voltage of ionization chamber {name} is not enabled") + if ch["self.dev"].hv.hv_v.get() < 1000: + raise AutoGainError(f"HV voltage of ionization chamber {name} is < 1000") + if ch["self.dev"].hv.grid_v.get() < 1000: + raise AutoGainError(f"Grid voltage of ionization chamber {name} is < 1000") + + # Check gas filling of ionization chambers + for name, ch in active_channels.items(): + if name != "pips": + if ch["self.dev"].gmes.status.get() != 1: + raise AutoGainError(f"Gas filling of ionization chamber {name} is not OK") + + logger.info("All checks done, start preparing for measurement") + + # Get initial monochromator position and velocity + init_pos = self.dev.mo1_bragg.position.get() + init_vel = self.dev.mo1_bragg.velocity.get() + + logger.info(f"Move mono to start of {emin} eV") + status = self.dev.mo1_bragg.move(emin) + status.wait(TIMEOUT_MONO_MOVE) + + # Set NIDAQ to max mode + # TODO implement + + # Set gains to lowest gain + for name, ch in active_channels.items(): + lowest_gain = AVAILABLE_GAINS[0] + ch["self.dev"].set_gain(lowest_gain) + ch["gain"] = lowest_gain + + remeasure = True + logger.info(f"Start measurement from {emin} eV to {emax} eV") + while remeasure: + # Create temporary storage for max signal per channel + data = {name: 0 for name in active_channels} + + # Measure current ring current + ring_current_1 = self._get_ring_current() + logger.info(f"Ring current right before measurement: {ring_current_1} mA") + if ring_current_1 == 0: + raise AutoGainError("Ring current dropped to 0 mA right before measurement") + + # Scan range, recording the peak NIDAQ signal per channel + status = self.dev.mo1_bragg.velocity.put(MONO_VELOCITY) + status.wait(TIMEOUT_MONO_PV) + self.dev.mo1_bragg.move(emax).wait(timeout=TIMEOUT_MONO_MOVE) + status.wait(TIMEOUT_MONO_MOVE) + for name, ch in active_channels.items(): + data[name] = max(data[name], ch["signal"].get()) + + # Rest max values of NIDAQ signals + # TODO implement + + # Measure current ring current again + ring_current_2 = self._get_ring_current() + logger.info(f"Ring current right after measurement: {ring_current_2} mA") + if ring_current_2 == 0: + raise AutoGainError("Ring current dropped to 0 mA during measurement") + ring_current = (ring_current_1 + ring_current_2) / 2 + + # Move back to first monochromator position + status = self.dev.mo1_bragg.move(emin) + + # Choose gain per channel based on the max signal recorded during the scan + remeasure = False + for name, ch in active_channels.items(): + raw_signal = data[name] + logger.info(f"Raw signal for self.device {name} is {raw_signal} V") + if comp_ring_current: + raw_signal = raw_signal * NOMINAL_RING_CURRENT / ring_current + logger.info(f"Compensate for ring current, new raw signal is {raw_signal} V") + if raw_signal < MIN_SIGNAL: + logger.info(f"Raw signal for self.device {name} is below {MIN_SIGNAL}") + # Choose next gain to be 100x the current gain, or if this gain does not exist, + # choose the next smaller one + if ch["gain"] == AVAILABLE_GAINS[-1]: + logger.warning( + f"Amplifier of {name} at highest gain {ch['gain']} and still not" + + f" measured signal above {MIN_SIGNAL}" + ) + else: + next_gain = AVAILABLE_GAINS[ + bisect_right(AVAILABLE_GAINS, ch["gain"] * 100) - 1 + ] + ch["self.dev"].set_gain(next_gain) + logger.info( + f"Setting gain of self.device {name} to {next_gain:.0e} and remeasure" + ) + remeasure = True + else: + gain = max( + ( + g + for g in AVAILABLE_GAINS + if raw_signal / ch["gain"] * g <= FULL_SCALE_V * SAFETY_MARGIN + ), + default=min(AVAILABLE_GAINS), + ) + ch["self.dev"].set_gain(gain) + logger.info(f"Calculated final gain for {name} of {gain:.0e}") + + # Wait for mono to return to start position + status.wait(TIMEOUT_MONO_MOVE) + + # Reset NIDAQ to mean mode + # TODO implement + + # Wait for mono to move to initial position and reset velocity + status = self.dev.mo1_bragg.move(init_pos) + status.wait(TIMEOUT_MONO_MOVE) + self.dev.mo1_bragg.velocity.put(init_vel) + + def _get_ring_current(self) -> float: + ring_current = 0 + retries = 0 + while ring_current == 0 and retries < 10: + ring_current = self.dev.curr.get() + retries += 1 + time.sleep(0.01) + return ring_current diff --git a/debye_bec/bec_ipython_client/plugins/digital_twin_core/__init__.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/__init__.py new file mode 100644 index 0000000..a42cb09 --- /dev/null +++ b/debye_bec/bec_ipython_client/plugins/digital_twin_core/__init__.py @@ -0,0 +1,3 @@ +from .beamline import get_parameters + +parameters = get_parameters() diff --git a/debye_bec/bec_ipython_client/plugins/digital_twin_core/beamline.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/beamline.py new file mode 100644 index 0000000..3653284 --- /dev/null +++ b/debye_bec/bec_ipython_client/plugins/digital_twin_core/beamline.py @@ -0,0 +1,51 @@ +import socket + +from bec_lib import bec_logger + +from .types import BeamlineId + +logger = bec_logger.logger + + +def get_beamline_id() -> BeamlineId: + """ + Based on the bec servers hostname, tries to extract the beamline + identifier (e.g. x01da, x10da, etc). + + Raises: + ValueError if beamline cannot be extracted from hostname or beamline not implemented. + """ + bec_hostname = socket.gethostname() + start = bec_hostname.find("x") + if start != -1: + beamline = bec_hostname[start : start + 5] + match beamline: + case "x01da": + return BeamlineId.X01DA + case "x10da": + return BeamlineId.X10DA + case _: + raise ValueError(f"Not implemented beamline {beamline}") + else: + logger.warning(f"Failed to extract beamline from bec server hostname {bec_hostname}") + choice = input("Do you want to manually select a beamline? (yes/no): ").strip().lower() + if choice in ["yes", "y"]: + bl = input(f"Choose from: {[bl.value for bl in BeamlineId]}") + if bl in BeamlineId: + logger.info(f"Manually selected beamline {bl}") + return BeamlineId(bl) + else: + raise ValueError(f"Wrong selection {bl}") + else: + raise ValueError("Cannot open digital twin without a beamline") + + +def get_parameters(): + beamline = get_beamline_id() + if beamline == "x01da": + from . import x01da_parameters as parameters + elif beamline == "x10da": + from . import x10da_parameters as parameters + else: + raise ValueError(f"Unknown beamline: {beamline}") + return parameters diff --git a/debye_bec/bec_ipython_client/plugins/digital_twin_core/digital_twin_core.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/digital_twin_core.py new file mode 100644 index 0000000..8acc1b5 --- /dev/null +++ b/debye_bec/bec_ipython_client/plugins/digital_twin_core/digital_twin_core.py @@ -0,0 +1,1119 @@ +import re +from pathlib import Path +from typing import Literal, cast + +import numpy as np +import yaml +from bec_lib import bec_logger +from scipy.interpolate import UnivariateSpline +from xrt.backends.raycing.physconsts import AVOGADRO, CHeVcm + +from . import parameters as bl +from .beamline import get_beamline_id +from .types import BeamlineId, ConfigDict, DataDict, SurfaceDict + +H = 6.62606957e-34 +E = 1.602176634e-19 +C = 299792458 +RE = 2.8179e-15 + +OFFSET_FILE_X01DA = Path(__file__).with_name("x01da_offsets.yaml") +OFFSET_FILE_X10DA = Path(__file__).with_name("x10da_offsets.yaml") + +logger = bec_logger.logger + +""" +The idea is to move the core logic of digital_twin to this file. Only keep the gui elements in the widget. +This way, the scheduler widget can access digital twin without loading the GUI (GUI is still needed for the item creation, but not the item execution) + Scheduler will extract assistant inputs (get_assistant_config) during item creation + Scheduler will use digital_twin.calculate_positons to calculate positions and digital_twin.move_all to move the motors +""" + + +class DigitalTwinCore: + + def __init__(self): + logger.info("This is the digital twin from the ipython client!") + self.beamline = get_beamline_id() + self.offset_file = Path() + match self.beamline: + case "x01da": + self.offset_file = OFFSET_FILE_X01DA + case "x10da": + self.offset_file = OFFSET_FILE_X10DA + self.offsets = {} + self.load_offsets() + + def move_with_config(self, config): + positions = self.calc_positions(self.beamline, config) + positions = self.apply_offsets(positions, nested_config=True) + logger.info(f"Would now move to these positions: {positions}") + + def load_offsets(self): + if self.offsets == {}: + logger.info("Load beamline offsets") + if not self.offset_file.exists(): + raise FileNotFoundError(f"Offset file not found: {self.offset_file}") + + with self.offset_file.open("r", encoding="utf-8") as f: + data = yaml.safe_load(f) + + if not isinstance(data, dict): + raise ValueError(f"Expected a YAML mapping, got {type(data).__name__}") + + self.offsets = data + else: + logger.info("Unload beamline offsets") + self.offsets = {} + + def apply_offsets(self, config, nested_config=False): + for axis, axis_data in config.items(): + if axis in self.offsets: + axis_offsets = self.offsets[axis] + if "modifier" in axis_offsets and "offset" in axis_offsets: + for idx, rng in enumerate(axis_offsets["modifier"]["range"]): + modifier_axis = axis_offsets["modifier"]["axis"] + modifier_value = ( + config[modifier_axis]["value"] + if nested_config + else config[modifier_axis] + ) + if rng[0] < modifier_value < rng[1]: + if nested_config: + axis_data["value"] += axis_offsets["offset"][idx] + else: + config[axis] += axis_offsets["offset"][idx] + break + elif "offset" in axis_offsets: + if nested_config: + axis_data["value"] += axis_offsets["offset"] + else: + config[axis] += axis_offsets["offset"] + return config + + def remove_offsets(self, config): + for axis, _ in config.items(): + if axis in self.offsets: + axis_offsets = self.offsets[axis] + if "modifier" in axis_offsets and "offset" in axis_offsets: + for idx, rng in enumerate(axis_offsets["modifier"]["range"]): + if rng[0] < config[axis_offsets["modifier"]["axis"]] < rng[1]: + config[axis] -= axis_offsets["offset"][idx] + break + elif "offset" in axis_offsets: + config[axis] -= axis_offsets["offset"] + return config + + @staticmethod + def calc_positions(beamline: BeamlineId, cfg: ConfigDict) -> dict[str, dict[str, float]]: + """ + Calculates the positions of axes based on a beamline config. + + Args: + cfg(ConfigDict): Dictionary with beamline config + + Returns: + dict[str, dict[str, float]]: Dictionary mapping device names to dictionaries + containing a "value" key with the corresponding float value (position). + """ + + pos = {} + + ## FE slits + trxr = -np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1] + trxw = ( + (np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1]) + / bl.feSlits.center1[1] + * bl.feSlits.center2[1] + ) + tryb = -np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1] + tryt = ( + (np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1]) + / bl.feSlits.center1[1] + * bl.feSlits.center2[1] + ) + + xgap = trxw - trxr + ygap = tryt - tryb + + pos["sldi_gapx"] = {"value": xgap} + pos["sldi_gapy"] = {"value": ygap} + + ## Collimating Mirror + obj_dist = bl.cm.center[1] # object distance + beam_vs = 2 * obj_dist * np.tan(cfg["v_acc"]) # vertical size of beam after CM + + # TRX + if cfg["cm_stripe"] in bl.cm.surface: + index = bl.cm.surface.index(cfg["cm_stripe"]) + else: + raise ValueError(f"Requested stripe {cfg['cm_stripe']} not found in parameters!") + cm_trx = -(bl.cm.limOptX[0][index] + bl.cm.limOptX[1][index]) / 2 + pos["cm_trx"] = {"value": cm_trx} + + # TRY + height = obj_dist * np.tan(cfg["v_acc"]) ** 2 * 1 / np.tan(cfg["cm_pitch"]) + pos["cm_try"] = {"value": height} + + # Pitch + pos["cm_rotx"] = { + "value": -cfg["cm_pitch"] * 1e3 + } # invert and convert to mrad (same as EGU of rotx axis) + + # Bending Radius + radius = ( + 2.0 * obj_dist / np.sin(cfg["cm_pitch"]) + ) # Elements of modern X-ray Physics, page 108 ff. + pos["cm_bnd_radius"] = {"value": radius * 1e-6} # Convert to km + + ## Monochromator + if cfg["mo1_mode"] == "Monochromatic": + # Add 2x CM pitch to the bragg angle + bragg = cfg["mo1_bragg"] + elif cfg["mo1_mode"] == "Pinkbeam": + # Align xtal surfaces parallel to beam + bragg = 0 + else: + raise ValueError("Monochromator mode not supported") + pos["mo1_bragg_angle"] = {"value": bragg / np.pi * 180} # Bragg angle in deg + + # TRY, Height + l = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) + yhor = l * np.cos(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"])) + yver = yhor * np.tan(2.0 * cfg["cm_pitch"]) + + if cfg["mo1_mode"] == "Monochromatic": + beam_offset_mo1 = ( + l * np.sin(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"])) - yver + ) # Resultat ist korrekt! + elif cfg["mo1_mode"] == "Pinkbeam": + beam_offset_mo1 = 0 + else: + raise ValueError("Monochromator mode not supported") + + def csc(a): + return 1 / np.sin(a) + + def cot(a): + return 1 / np.tan(a) + + # calculate height of center of first crystal surface + f = bl.mo1.rotOffset # rotation offset, mm + d = bl.mo1.heightOffset # xtal height offset, mm + c = d * csc(cfg["mo1_bragg"]) - f * cot(cfg["mo1_bragg"]) + + # Calculate height of center of rotation + b = np.sqrt( + d**2 * csc(cfg["mo1_bragg"]) ** 2 + - 2 * d * f * cot(cfg["mo1_bragg"]) * csc(cfg["mo1_bragg"]) + + f**2 * cot(cfg["mo1_bragg"]) ** 2 + + f**2 + ) + h = np.cos(np.pi / 2 - np.arctan(f / c) - cfg["mo1_bragg"] - 2 * cfg["cm_pitch"]) * b + h2 = ((bl.mo1.center[1] - bl.cm.center[1]) - np.sqrt(b**2 - h**2)) * np.tan( + 2 * cfg["cm_pitch"] + ) + height_mo1_real = ( + h + h2 + ) # per design, the height should not change if the pitch of the CM is not changed! + if cfg["mo1_mode"] == "Monochromatic": + pass + elif cfg["mo1_mode"] == "Pinkbeam": + height_mo1_real = ( + height_mo1_real - 13 + ) # Move down to let beam pass between both crystal without touching copper cooler + else: + raise ValueError("Monochromator mode not supported") + pos["mo1_try"] = {"value": height_mo1_real} + + # TRX, Crystal selection + if cfg["mo1_mode"] == "Monochromatic": + xtal = cfg["mo1_xtal"].translate( + str.maketrans("", "", "()") + ) # Remove brackets from xtal name to conform with parameters + if xtal in bl.mo1.xtal: + index = bl.mo1.xtal.index(xtal) + else: + raise ValueError(f"Requested xtal {xtal} not found in parameters!") + pos["mo1_trx"] = {"value": bl.mo1.xtalOffsetX[index]} + else: + pos["mo1_trx"] = {"value": 0} + + diag = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) # Calculations for Mono + dz = diag * np.cos(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) + + ## Slits 1 + d = bl.opSlits1.center[1] - bl.cm.center[1] - dz + sl1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 + pos["sl1_centery"] = {"value": sl1_beam_height} + pos["sl1_gapy"] = {"value": beam_vs} + + ## Beam Monitor 1 + d = bl.opBM1.center[1] - bl.cm.center[1] - dz + bm1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 + pos["bm1_try"] = {"value": bm1_beam_height} + + ## Focusing Mirror + p = bl.fm.center[1] + q = cfg["smpl"] - bl.fm.center[1] + f = (p * q) / (p + q) # focal length + + # Bender radius + if cfg["fm_qy"] is None: + radius = 2 * q / np.sin(cfg["fm_rotx"]) # ideal bending radius for focused beam + else: + radius = ( + 2 * cfg["fm_qy"] / np.sin(cfg["fm_rotx"]) + ) # ideal bending radius for unfocused beam + pos["fm_bnd_radius"] = {"value": radius * 1e-6} # Convert to km + + # Pitch + d = bl.fm.center[1] - bl.cm.center[1] - dz + fm_rotx = ( + 2 * cfg["cm_pitch"] - cfg["fm_rotx"] + ) # calculate pitch in absolute values (according to horizontal plane) + pos["fm_rotx"] = { + "value": -fm_rotx * 1e3 + } # invert and convert to mrad (same as EGU of rotx axis) + + if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): + + # TRY + if cfg["fm_stripe"] == "Rh (toroid)": + r = bl.fm.r[0] + h_cyl = bl.fm.hToroid[0] + else: # PT toroid + r = bl.fm.r[1] + h_cyl = bl.fm.hToroid[1] + width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"] * 1e-3) + alpha = np.arccos(1 - width_beam**2 / (2 * r**2)) + h = r - (r * np.cos(alpha / 2)) + fm_beam_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg[ + "fm_gain_height" + ] + fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 - h_cyl + h / 2) * cfg[ + "fm_gain_height" + ] + pos["fm_try"] = {"value": fm_height} + + # TRX + if cfg["fm_stripe"] == "Rh (toroid)": + x_cyl = -bl.fm.xToroid[0] + else: + x_cyl = -bl.fm.xToroid[1] + pos["fm_trx"] = {"value": x_cyl} + + elif cfg["fm_stripe"] in ("Rh (flat)", "Pt (flat)"): + + # TRY + fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg["fm_gain_height"] + fm_beam_height = fm_height + pos["fm_try"] = {"value": fm_height} + + # TRX + if cfg["fm_stripe"] == "Rh (flat)": + x_flat = -bl.fm.xFlat[0] + else: + x_flat = -bl.fm.xFlat[1] + pos["fm_trx"] = {"value": x_flat} + + else: + raise ValueError("FM Stripe selection not valid") + + pos["fm_roty"] = {"value": 0} + pos["fm_rotz"] = {"value": 0} + + ## Slits 2 + if hasattr(bl, "opSlits2"): + d = bl.opSlits2.center[1] - bl.fm.center[1] + sl2_beam_height = fm_beam_height - d * np.tan( + -(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]) + ) + pos["sl2_centery"] = {"value": sl2_beam_height} + pos["sl2_gapy"] = {"value": beam_vs} + + ## Beam Monitor 2 + d = bl.opBM2.center[1] - bl.fm.center[1] + bm2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) + pos["bm2_try"] = {"value": bm2_beam_height} + + ## Optical Table + + if beamline == "x01da": + # TRY + d = bl.ehWindow.center[1] - bl.fm.center[1] + ot_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) + pos["ot_try"] = {"value": ot_height} + + # Pitch + ot_pitch = -(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]) + pos["ot_rotx"] = {"value": ot_pitch * 1e3} + + # TRZ ES1 + ot_es1_trz = cfg["smpl"] + pos["ot_es1_trz"] = {"value": ot_es1_trz} + + # ES0 exit window + pos["es0wi_try"] = { + "value": 5 + } # At 5mm, the middle of the window is 500 mm from the table (neutral position) + else: + # Exit window height + d = bl.ehWindow.center[1] - bl.fm.center[1] + es0wi_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) + pos["es0wi_try"] = {"value": es0wi_try} + + # ES1 table height + d = bl.es1.center[1] - bl.fm.center[1] + es1_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) + pos["es1_try"] = {"value": es1_try} + + # IC0 height + d = bl.es1ic0.center[1] - bl.fm.center[1] + es1ic0_try = ( + fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try + ) + pos["es1ic0_try"] = {"value": es1ic0_try} + + # IC1 height + d = bl.es1ic1.center[1] - bl.fm.center[1] + es1ic1_try = ( + fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try + ) + pos["es1ic1_try"] = {"value": es1ic1_try} + + # IC2 height + d = bl.es1ic2.center[1] - bl.fm.center[1] + es1ic2_try = ( + fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try + ) + pos["es1ic2_try"] = {"value": es1ic2_try} + + # ES2 table height + d = bl.es2.center[1] - bl.fm.center[1] + es2_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) + pos["es2_try"] = {"value": es2_try} + + return pos + + @staticmethod + def sldi_gap_to_acc(sldi_gapx: float, sldi_gapy: float) -> tuple[float, float]: + """ + Calculate the slits acceptance based on the gap values + + Args: + sldi_gapx(float): GAPX value of the slits in mm + sldi_gapy(float): GAPY value of the slits in mm + + Returns: + tuple[float, float]: Horizontal and vertical acceptance in rad + """ + d1 = bl.feSlits.center1[1] + d2 = bl.feSlits.center2[1] + h_acc = np.tan(sldi_gapx / (d2 + d1)) + v_acc = np.tan(sldi_gapy / (d2 + d1)) + return h_acc, v_acc + + @staticmethod + def cm_trx_to_stripe(cm_trx: float) -> str | None: + """ + Based on the trx value of the collimating mirror, return + the correct stripe + + Args: + cm_trx(float): Collimating mirror trx value + + Returns + str | None: Stripe of the mirror, None if not found + """ + cm_stripe = None + for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]): + if low <= cm_trx <= high: + cm_stripe = name + return cm_stripe + + @staticmethod + def cm_stripe_to_trx(cm_stripe: str) -> float | None: + """ + Based on the stripe of the collimating mirror, return + the trx value + + Args: + cm_stripe(str): Stripe of the collimating mirror + + Returns: + float | None: TRX value of the stripe. None if not found + """ + for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]): + if cm_stripe == name: + return -(low + high) / 2 + return None + + @staticmethod + def fm_trx_to_stripe(fm_trx: float) -> str | None: + """ + Based on the trx value of the focusing mirror, return + the correct stripe + + Args: + fm_trx(float): focusing mirror trx value + + Returns + str | None: Stripe of the mirror, None if not found + """ + fm_stripe = None + if hasattr(bl.fm, "surfaceFlat"): + for name, low, high in zip( + bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0] + ): + if low <= fm_trx <= high: + fm_stripe = name + " (flat)" + for name, low, high in zip( + bl.fm.surfaceToroid, bl.fm.limOptXToroid[1], bl.fm.limOptXToroid[0] + ): + if low <= fm_trx <= high: + fm_stripe = name + " (toroid)" + return fm_stripe + + @staticmethod + def fm_stripe_to_trx(fm_stripe: str) -> float | None: + """ + Based on the stripe of the focusing mirror, return + the trx value + + Args: + fm_stripe(str): Stripe of the focusing mirror + + Returns: + float | None: TRX value of the stripe. None if not found + """ + if hasattr(bl.fm, "surfaceFlat"): + for name, low, high in zip( + bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0] + ): + if fm_stripe == name + " (flat)": + return (low + high) / 2 + for name, low, high in zip( + bl.fm.surfaceToroid, bl.fm.limOptXToroid[1], bl.fm.limOptXToroid[0] + ): + if fm_stripe == name + " (toroid)": + return -(low + high) / 2 + return None + + @staticmethod + def mo1_energy_resolution(xtal: Literal["Si111", "Si311"], energy: float) -> float: + """ + Calculate the energy resolution of the monochromator + + Args: + xtal(str): Xtal name. "Si111" or "Si311" + energy(float): Energy in eV + + Returns: + float: Energy resolution in eV + """ + index = bl.mo1.xtal.index(xtal) + crystal = bl.mo1.material1[index] + + dtheta = np.linspace(-30, 90, 601) + theta = crystal.get_Bragg_angle(energy) + dtheta * 1e-6 + refl = np.abs(crystal.get_amplitude(energy, np.sin(theta))[0]) ** 2 # single crystal + + refl2 = refl**2 # DCM with parallel crystals + + # FWHM of the DCM curve + spline = UnivariateSpline(dtheta, refl2 - refl2.max() / 2, s=0) + roots = cast(np.ndarray, spline.roots()) + r1, r2 = float(roots[0]), float(roots[1]) + fwhm_rad = (r2 - r1) * 1e-6 # µrad → rad + + # Energy resolution + theta_b = crystal.get_Bragg_angle(energy) + de_over_e = fwhm_rad / np.tan(theta_b) + de = de_over_e * energy + + # logger.info(f"DCM FWHM : {r2-r1:.2f} µrad") + # logger.info(f"ΔE/E : {dE_over_E:.2e}") + # logger.info(f"ΔE : {dE:.3f} eV at {E} eV") + + return de + + @staticmethod + def cm_reflectivity(cm_stripe: str, cm_pitch: float, energy: float) -> float: + """ + Calculate the reflectivity of the mirror stripe based + on the pitch and energy. + + Args: + cm_stripe(str): Mirror stripe + cm_pitch(float): Pitch of the mirror (beam incidence angle) + energy(float): Energy of the beam in eV + + Returns: + float: Reflectivity [0-1] + """ + if cm_stripe is None: + return np.nan + index = bl.cm.surface.index(cm_stripe) + rs, _ = bl.cm.material[index].get_amplitude(energy, np.sin(cm_pitch))[0:2] + refl = abs(rs) ** 2 + return refl + + @staticmethod + def fm_reflectivity(fm_stripe: str, fm_pitch: float, energy: float) -> float: + """ + Calculate the reflectivity of the mirror stripe based + on the pitch and energy. + + Args: + cm_stripe(str): Mirror stripe + cm_pitch(float): Pitch of the mirror (beam incidence angle) + energy(float): Energy of the beam in eV + + Returns: + float: Reflectivity [0-1] + """ + if fm_stripe is None: + return np.nan + if fm_stripe in ("Rh (toroid)", "Pt (toroid)"): + surface = bl.fm.surfaceToroid + material = bl.fm.materialToroid + stripe = re.sub(r"\s*\(.*?\)", "", fm_stripe).strip() + index = surface.index(stripe) + else: + surface = bl.fm.surfaceFlat + material = bl.fm.materialFlat + stripe = re.sub(r"\s*\(.*?\)", "", fm_stripe).strip() + index = surface.index(stripe) + rs, _ = material[index].get_amplitude(energy, np.sin(fm_pitch))[0:2] + refl = abs(rs) ** 2 + return refl + + @staticmethod + def mo1_bragg_angle( + mo_mode: Literal["Monochromatic", "Pinkbeam"], + d_spacing: float, + energy: float, + cm_pitch: float, + ) -> tuple[float, float]: + """ + Calculate the bragg angle of the monochromator. + Corrects for the collimating mirror pitch. + + Args: + mo_mode(str): Monochromator mode. "Monochromatic" or "Pinkbeam" + d_spacing(float): D-spacing of the crystal in Angstrom + energy(float): Energy of the beam in eV + cm_pitch(float): Pitch of collimating mirror in rad + + Returns: + tuple[float, float]: Bragg angle and corrected bragg angle + """ + wl = C * H / (E * energy) + val = wl / (2 * d_spacing * 1e-10) + bragg_angle = 0 + if val > -1 and val < 1: + bragg_angle = np.asin(val) + if mo_mode == "Monochromatic": + # Add 2x CM pitch to the bragg angle + bragg_angle_cor = (2 * cm_pitch) + bragg_angle + else: + # Align xtal surfaces parallel to beam + bragg_angle_cor = 2 * cm_pitch + return bragg_angle, bragg_angle_cor + + @staticmethod + def fm_ideal_pitch( + fm_focus: Literal["Defocused", "Focused", "Manual"], + fm_stripe: str, + smpl: float, + sldi_hacc: float | None = None, + sldi_vacc: float | None = None, + fm_focx: float | None = None, + fm_focy: float | None = None, + ) -> tuple[float, float | None]: + """ + Calculates the ideal pitch for the focusing mirror depending on the + focusing strategy. + If "Defocused" is chosed, sldi_hacc, sldi_vacc, fm_focx and fm_focy + must be provided. + + Args: + fm_focus(str): Focus strategy. "Defocused", "Focused" or "Manual + fm_stripe(str): Mirror stripe + smpl(float): Sample position in mm from source + sldi_hacc(float): Horizontal acceptance of frontend slits. Defaults to None + sldi_vacc(float): Vertical acceptance of frontend slits. Defaults to None + fm_focx(float): Requested horizontal spot size in mm. Defaults to None + fm_focy(float): Requested vertical spot size in mm. Defaults to None + + Returns: + tuple[float, float | None]: Pitch of mirror in rad, qy in mm + """ + + # logger.info("Calculate pitch and qy now...") + # logger.info(f"sldi_hacc: {sldi_hacc}") + # logger.info(f"sldi_vacc: {sldi_vacc}") + # logger.info(f"fm_stripe: {fm_stripe}") + # logger.info(f"smpl: {smpl}") + p_cm = bl.cm.center[1] # posCM + p = bl.fm.center[1] # posFM + q = smpl - bl.fm.center[1] # dist posFM to posEX + if fm_focus == "Defocused": + assert sldi_hacc is not None, "sldi_hacc must be provided for Defocused mode" + assert sldi_vacc is not None, "sldi_vacc must be provided for Defocused mode" + assert fm_focx is not None, "fm_focx must be provided for Defocused mode" + assert fm_focy is not None, "fm_focy must be provided for Defocused mode" + a = 2 * np.tan(sldi_hacc) * bl.fm.center[1] # Beam width at focusing mirror + # logger.info(f"a: {a}") + # logger.info(f"sldi_hacc: {sldi_hacc}") + # logger.info(f"bl.fm.center[1]: {bl.fm.center[1]}") + # logger.info(f"p: {p}") + # logger.info(f"q: {q}") + b = ( + 2 * np.tan(sldi_vacc) * bl.cm.center[1] + ) # Beam height at focusing mirror (collimated beam) + x = fm_focx + # logger.info(f"x: {x}") + x = 0.098821 * x**2 + 0.512344 * x # polynom to correct for spot size + # logger.info(f"x (corrected): {x}") + y = fm_focy + y = 3.183562 * y**2 + 1.258364 * y # polynom to correct for spot size + qx = q + x * p / a + qy = q + y * p_cm / b + f = (p * qx) / (p + qx) # focal length + # logger.info(f"qx: {qx}") + # logger.info(f"f: {f}") + else: # Calculate for focused beam on sample in "manual" and "focused" mode + qy = None + f = (p * q) / (p + q) # focal length + pitch = 0 + if "Rh" in fm_stripe: + pitch = np.arcsin(bl.fm.r[0] / (2 * f)) # ideal pitch for FM + if "Pt" in fm_stripe: + pitch = np.arcsin(bl.fm.r[1] / (2 * f)) # ideal pitch for FM + # logger.info(f"fm_pitch: {pitch}") + # logger.info(f"qy: {qy}") + return pitch, qy + + @staticmethod + def calc_beamsize( + sldi_hacc: float, + sldi_vacc: float, + fm_stripe: str, + fm_pitch: float, + fm_radius: float, + smpl: float, + ) -> tuple[float, float | None]: + """ + Calculate the resulting beamsize according to the input parameters + + Args: + sldi_hacc(float): Horizontal acceptance of frontend slits + sldi_vacc(float): Vertical acceptance of frontend slits + fm_stripe(str): Mirror stripe + fm_pitch(float): Focusing mirror pitch in rad + fm_radius(float): Focusing mirror bender radius in m + smpl(float): Sample position in mm from source + + Returns: + tuple[float, float | None]: horizontal spot size, vertical spot size, both in mm + """ + + # logger.info("Calculate beamsize now...") + # logger.info(f"sldi_hacc: {sldi_hacc}") + # logger.info(f"sldi_vacc: {sldi_vacc}") + # logger.info(f"fm_stripe: {fm_stripe}") + # logger.info(f"fm_pitch: {fm_pitch}") + # logger.info(f"fm_radius: {fm_radius}") + # logger.info(f"smpl: {smpl}") + p_cm = bl.cm.center[1] # posCM + p = bl.fm.center[1] # posFM + q = smpl - bl.fm.center[1] # dist posFM to posEX + qy = fm_radius * np.sin(fm_pitch) / 2 + a = 2 * np.tan(sldi_hacc) * bl.fm.center[1] # Beam width at focusing mirror + b = ( + 2 * np.tan(sldi_vacc) * bl.cm.center[1] + ) # Beam height at focusing mirror (collimated beam) + f = 0 + if "Rh" in fm_stripe: + f = bl.fm.r[0] / (2 * np.sin(fm_pitch)) + if "Pt" in fm_stripe: + f = bl.fm.r[1] / (2 * np.sin(fm_pitch)) + qx = p * f / (p - f) + x = a * (qx - q) / p + y = b * (qy - q) / p_cm + # Change this | to a plus if calculation is not correct + fm_focx = -4 * (64043 - 125000 * np.sqrt(0.26249637 + 0.395284 * x)) / 98821 + # Change this | to a plus if calculation is not correct + fm_focy = -1 * (314591 - 250000 * np.sqrt(1.58347995 + 12.734248 * y)) / 1591781 + # logger.info(f"f: {f}") + # logger.info(f"qx: {qx}") + # logger.info(f"qy: {qy}") + # logger.info(f"fm_focx: {fm_focx}") + # logger.info(f"fm_focy: {fm_focy}") + return fm_focx, fm_focy + + @staticmethod + def cm_critical_angle(cm_stripe: Literal["Si", "Pt", "Rh"], energy) -> float: + """ + Calculate the critical angle of the mirror stripe + + Args: + cm_stripe(str): Mirror stripe. "Si", "Pt" or "Rh" + energy(float): Energy in eV + + Returns: + float: Critical angle in rad + """ + if cm_stripe == "Si": + stripe = bl.stripeSi + elif cm_stripe == "Pt": + stripe = bl.stripePt + else: + stripe = bl.stripeRh + w = CHeVcm / 100 / energy # convert energy [eV] to wavelength [m] + f1 = stripe.elements[0].Z + np.real(stripe.elements[0].get_f1f2(energy)) + number_density = stripe.rho * 1e3 * AVOGADRO / (stripe.elements[0].mass / 1e3) + critical_angle = np.sqrt(number_density * RE * w**2 * f1 / np.pi) + return critical_angle + + @staticmethod + def mirror_surface_geometries( + mirror: Literal["cm", "fm_toroid", "fm_flat"], + ) -> dict[str, tuple[float, float, float, float]]: + """ + Return the mirror stripe geometries + + Args: + mirror(str): Mirror. "cm", "fm_toroid" or "fm_flat" + + Returns: + dict[str, tuple[float, float, float, float]]: Dictionary mapping surface + names to tuples of (x, y, width, height). + """ + if mirror == "cm": + surface = bl.cm.surface + lim_opt_x = bl.cm.limOptX + lim_opt_y = bl.cm.limOptY + elif mirror == "fm_toroid": + surface = bl.fm.surfaceToroid + lim_opt_x = bl.fm.limOptXToroid + lim_opt_y = bl.fm.limOptYToroid + elif mirror == "fm_flat": + surface = bl.fm.surfaceFlat + lim_opt_x = bl.fm.limOptXFlat + lim_opt_y = bl.fm.limOptYFlat + else: + raise ValueError(f"Requested mirror {mirror} not available!") + geom = {} + for sf, lx, hx, ly, hy in zip( + surface, lim_opt_x[0], lim_opt_x[1], lim_opt_y[0], lim_opt_y[1] + ): + geom[sf] = (lx, ly, hx - lx, hy - ly) + return geom + + @staticmethod + def mo_surface_geometries( + mo: Literal["mo1"], plane: Literal[0, 1] + ) -> dict[str, tuple[float, float, float, float]]: + """ + Return the monochromator xtal geometries + + Args: + mo(str): Monochromator. Only "mo1" implemented + plane(int): Surface of xtal. 0 and 1 (First and second) + + Returns: + dict[str, tuple[float, float, float, float]]: Dictionary mapping surface + names to tuples of (x, y, width, height). + """ + if mo == "mo1": + xtal = bl.mo1.xtal + xtal_width = bl.mo1.xtalWidth + xtal_offset_x = bl.mo1.xtalOffsetX + if plane == 0: + xtal_length = bl.mo1.xtalLength1 + else: + xtal_length = bl.mo1.xtalLength2 + else: + return {} + geom = {} + for sf, w, offx, length in zip(xtal, xtal_width, xtal_offset_x, xtal_length): + geom[sf] = (offx - w / 2, -length / 2, w, length) + return geom + + @staticmethod + def wall_geometries() -> list[list[float]]: + """ + Return the wall geometries + + Returns: + list[list[float]]: List of [x, y, width, height] geometry values for each wall. + """ + geom = [] + if not hasattr(bl, "walls"): + return geom + for i, _ in enumerate(bl.walls.start): + geom.append( + [ + bl.walls.start[i], + bl.walls.height[i][0], + bl.walls.end[i] - bl.walls.start[i], + bl.walls.height[i][1] - bl.walls.height[i][0], + ] + ) + return geom + + @staticmethod + def pipe_geometries() -> list[dict[str, np.ndarray]]: + """ + Return the wall geometries + + Returns: + list[dict[str, np.ndarray]]: List of dictionaries with keys "x" and "y", + each containing a numpy array of two float values representing + the start and end coordinates of the pipe top and bottom edges. + """ + pipes = [] + if not hasattr(bl, "vacuum_pipes"): + return pipes + for i, _ in enumerate(bl.vacuum_pipes.center): + top = bl.vacuum_pipes.center[i] + bl.vacuum_pipes.diameter[i] / 2 + bl.sourceHeight + bottom = bl.vacuum_pipes.center[i] - bl.vacuum_pipes.diameter[i] / 2 + bl.sourceHeight + pipes.append( + { + "x": np.array([bl.vacuum_pipes.start[i], bl.vacuum_pipes.end[i]]), + "y": np.array([top, top]), + } + ) + pipes.append( + { + "x": np.array([bl.vacuum_pipes.start[i], bl.vacuum_pipes.end[i]]), + "y": np.array([bottom, bottom]), + } + ) + return pipes + + @staticmethod + def table_to_smpl_pos(table: str) -> float: + """ + Return the sample position based on the table name. + + Args: + table (str): Table name, e.g. ES1 or ES2 + """ + + if table == bl.es1.name: + return bl.es1.center[1] + if table == bl.es2.name: + return bl.es2.center[1] + raise ValueError(f"Table {table} not found in beamline parameter file") + + @staticmethod + def calc_sideview(cfg: ConfigDict) -> DataDict: + """ + Calculates the sideview coordinates based on a beamline config. + + Args: + cfg(ConfigDict): Dictionary with beamline config + + Returns: + DataDict: Sideview data + """ + + beam: DataDict = {"x": [], "y": []} + + beam["x"] = [] + beam["y"] = [] + beam["x"].append(0) # Source + beam["y"].append(bl.sourceHeight) + beam["x"].append(bl.cm.center[1]) # CM + beam["y"].append(bl.sourceHeight) + if cfg["mo1_mode"] == "Monochromatic": + diag = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) # Calculations for Mono + dy = diag * np.sin(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) + dz = diag * np.cos(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) + beam["x"].append(bl.mo1.center[1] - dz / 2) # Mono 1.1 + beam["y"].append( + bl.sourceHeight + + np.tan(2 * cfg["cm_pitch"]) * (bl.mo1.center[1] - dz / 2 - bl.cm.center[1]) + ) + beam["x"].append(bl.mo1.center[1] + dz / 2) # Mono 1.2 + beam["y"].append( + bl.sourceHeight + + np.tan(2 * cfg["cm_pitch"]) * (bl.mo1.center[1] - dz / 2 - bl.cm.center[1]) + + dy + ) + beam["x"].append(bl.fm.center[1]) # FM + beam["y"].append( + bl.sourceHeight + + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1] - dz) + + dy + ) + beam["x"].append(cfg["smpl"]) # Experiment + beam["y"].append( + bl.sourceHeight + + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1] - dz) + + dy + + np.tan(2 * (cfg["cm_pitch"] - cfg["fm_rotx"])) * (cfg["smpl"] - bl.fm.center[1]) + ) + elif cfg["mo1_mode"] == "Pinkbeam": + beam["x"].append(bl.fm.center[1]) # FM + beam["y"].append( + bl.sourceHeight + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1]) + ) + beam["x"].append(cfg["smpl"]) # Experiment + beam["y"].append( + bl.sourceHeight + + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1]) + + np.tan(2 * (cfg["cm_pitch"] - cfg["fm_rotx"])) * (cfg["smpl"] - bl.fm.center[1]) + ) + + return beam + + @staticmethod + def calc_surfaces(cfg: ConfigDict) -> SurfaceDict: + """ + Calculates the surface coordinates based on a beamline config. + + Args: + cfg(ConfigDict): Dictionary with beamline config + + Returns: + SurfaceDict: Surface data + """ + + out: SurfaceDict = { + "cm": {"x": [], "y": []}, + "mo1_1": {"x": [], "y": []}, + "mo1_2": {"x": [], "y": []}, + "fm": {"x": [], "y": []}, + } + + # Collimating mirror + l = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"]) / np.sin(cfg["cm_pitch"]) + + w1 = 2 * (bl.cm.center[1] - l / 2) * np.tan(cfg["h_acc"]) + w2 = 2 * (bl.cm.center[1] + l / 2) * np.tan(cfg["h_acc"]) + + # index = bl.cm.surface.index(cfg["cm_stripe"]) + + cen = -cfg["cm_trx"] + + out["cm"]["x"] = [cen - w1 / 2, cen - w2 / 2, cen + w2 / 2, cen + w1 / 2] + out["cm"]["y"] = [-l / 2, l / 2, l / 2, -l / 2] + + # Monochromator + # calculate height of center of first crystal surface + c = bl.mo1.heightOffset * 1 / np.sin(cfg["mo1_bragg"]) - bl.mo1.rotOffset * 1 / np.tan( + cfg["mo1_bragg"] + ) + e = bl.mo1.xtalGap[0] / np.tan(cfg["mo1_bragg"]) - c + + xtal = cfg["mo1_xtal"].translate( + str.maketrans("", "", "()") + ) # Remove brackets from xtal name to conform with parameters + index = bl.mo1.xtal.index(xtal) + + xtal_pos = bl.mo1.xtalOffsetX[index] + xtal_length_1 = bl.mo1.xtalLength1[index] + xtal_length_2 = bl.mo1.xtalLength2[index] + + width_beam = 2 * bl.mo1.center[1] * np.tan(cfg["h_acc"]) + + height_beam = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"]) + w = height_beam / np.sin(cfg["mo1_bragg"]) + + if cfg["mo1_mode"] == "Monochromatic": + out["mo1_1"]["x"] = [ + xtal_pos - width_beam / 2, + xtal_pos + width_beam / 2, + xtal_pos + width_beam / 2, + xtal_pos - width_beam / 2, + ] + out["mo1_1"]["y"] = [ + xtal_length_1 / 2 - c - w / 2, + xtal_length_1 / 2 - c - w / 2, + xtal_length_1 / 2 - c + w / 2, + xtal_length_1 / 2 - c + w / 2, + ] + out["mo1_2"]["x"] = [ + xtal_pos - width_beam / 2, + xtal_pos + width_beam / 2, + xtal_pos + width_beam / 2, + xtal_pos - width_beam / 2, + ] + out["mo1_2"]["y"] = [ + -xtal_length_2 / 2 + e - w / 2, + -xtal_length_2 / 2 + e - w / 2, + -xtal_length_2 / 2 + e + w / 2, + -xtal_length_2 / 2 + e + w / 2, + ] + else: # Pinkbeam + out["mo1_1"]["x"] = [] + out["mo1_1"]["y"] = [] + out["mo1_2"]["x"] = [] + out["mo1_2"]["y"] = [] + + if cfg["fm_stripe"] is None: + return out + # Focusing mirror + if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): + surface = bl.fm.surfaceToroid + stripe = re.sub(r"\s*\(.*?\)", "", cfg["fm_stripe"]).strip() + index = surface.index(stripe) + r = bl.fm.r[index] + else: + surface = bl.fm.surfaceFlat + stripe = re.sub(r"\s*\(.*?\)", "", cfg["fm_stripe"]).strip() + index = surface.index(stripe) + r = bl.fm.r[index] + off = -cfg["fm_trx"] + + width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"]) + + if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): + + l = height_beam / np.sin(cfg["fm_rotx"]) + alpha = np.arccos(1 - width_beam**2 / (2 * r**2)) + h = r - (r * np.cos(alpha / 2)) + z = h / np.tan(cfg["fm_rotx"]) + + x = [off - width_beam / 2, off - width_beam / 2] + y = [l / 2 - z / 2, -l / 2 - z / 2] + + res = 20 + x_elipse = np.linspace(0, np.pi, res) + y_elipse = np.linspace(0, np.pi, res) + x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse] + y_elipse = [width_beam * np.sin(i) * z / width_beam - l / 2 - z / 2 for i in y_elipse] + + x.extend(x_elipse) + y.extend(y_elipse) + + x.extend([off + width_beam / 2, off + width_beam / 2]) + y.extend([-l / 2 - z / 2, l / 2 - z / 2]) + + res = 50 + x_elipse = np.linspace(np.pi, 0, res) + y_elipse = np.linspace(np.pi, 0, res) + x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse] + y_elipse = [width_beam * np.sin(i) * z / width_beam + l / 2 - z / 2 for i in y_elipse] + + x.extend(x_elipse) + y.extend(y_elipse) + + out["fm"]["x"] = x + out["fm"]["y"] = y + + else: # flat surface, no toroid + l = height_beam / np.sin(cfg["fm_rotx"]) + + w1 = 2 * (bl.fm.center[1] - l / 2) * np.tan(cfg["h_acc"]) + w2 = 2 * (bl.fm.center[1] + l / 2) * np.tan(cfg["h_acc"]) + + out["fm"]["x"] = [off - w1 / 2, off + w1 / 2, off + w2 / 2, off - w2 / 2] + out["fm"]["y"] = [-l / 2, -l / 2, l / 2, l / 2] + + return out diff --git a/debye_bec/bec_widgets/widgets/digital_twin/types.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/types.py similarity index 100% rename from debye_bec/bec_widgets/widgets/digital_twin/types.py rename to debye_bec/bec_ipython_client/plugins/digital_twin_core/types.py diff --git a/debye_bec/bec_widgets/widgets/digital_twin/x01da_offsets.yaml b/debye_bec/bec_ipython_client/plugins/digital_twin_core/x01da_offsets.yaml similarity index 100% rename from debye_bec/bec_widgets/widgets/digital_twin/x01da_offsets.yaml rename to debye_bec/bec_ipython_client/plugins/digital_twin_core/x01da_offsets.yaml diff --git a/debye_bec/bec_widgets/widgets/digital_twin/x01da_parameters.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/x01da_parameters.py similarity index 100% rename from debye_bec/bec_widgets/widgets/digital_twin/x01da_parameters.py rename to debye_bec/bec_ipython_client/plugins/digital_twin_core/x01da_parameters.py diff --git a/debye_bec/bec_widgets/widgets/digital_twin/x10da_offsets.yaml b/debye_bec/bec_ipython_client/plugins/digital_twin_core/x10da_offsets.yaml similarity index 100% rename from debye_bec/bec_widgets/widgets/digital_twin/x10da_offsets.yaml rename to debye_bec/bec_ipython_client/plugins/digital_twin_core/x10da_offsets.yaml diff --git a/debye_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py b/debye_bec/bec_ipython_client/plugins/digital_twin_core/x10da_parameters.py similarity index 96% rename from debye_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py rename to debye_bec/bec_ipython_client/plugins/digital_twin_core/x10da_parameters.py index b7b40a5..c0bd4e0 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py +++ b/debye_bec/bec_ipython_client/plugins/digital_twin_core/x10da_parameters.py @@ -1,296 +1,296 @@ -""" -X10DA / SuperXAS Beamline Parameters. -This file describes the parameter of each component of the SuperXAS beamline -to be used for raytracing and geometrical calculations. -""" - -from collections import namedtuple - -import numpy as np -import xrt.backends.raycing.materials as rm - -# XRT definitions -filterBeryl = rm.Material("Be", rho=1.85, kind="plate") # pyright: ignore[reportArgumentType] -filterDiamond = rm.Material("C", rho=3.52, kind="plate") # pyright: ignore[reportArgumentType] -filterGraphite = rm.Material("C", rho=2.266, kind="plate") # pyright: ignore[reportArgumentType] - -stripeSi = rm.Material("Si", rho=2.33) # pyright: ignore[reportArgumentType] -stripePt = rm.Material("Pt", rho=21.45) # pyright: ignore[reportArgumentType] -stripeRh = rm.Material("Rh", rho=12.41) # pyright: ignore[reportArgumentType] -stripeCr = rm.Material("Cr", rho=7.14) # pyright: ignore[reportArgumentType] -stripePyrex = rm.Material( - "Si", rho=2.20 -) # Use Si as bare element and the density of SiO2 # pyright: ignore[reportArgumentType] - -si111_1 = rm.CrystalSi(hkl=(1, 1, 1), tK=77) # first xtal surface -si311_1 = rm.CrystalSi(hkl=(3, 1, 1), tK=77) # first xtal surface -si333_1 = rm.CrystalSi(hkl=(3, 3, 3), tK=77) # first xtal surface -si511_1 = rm.CrystalSi(hkl=(5, 1, 1), tK=77) # first xtal surface -si111_2 = rm.CrystalSi(hkl=(1, 1, 1), tK=77) # second xtal surface -si311_2 = rm.CrystalSi(hkl=(3, 1, 1), tK=77) # second xtal surface -si333_2 = rm.CrystalSi(hkl=(3, 3, 3), tK=77) # second xtal surface -si511_2 = rm.CrystalSi(hkl=(5, 1, 1), tK=77) # second xtal surface - -filterDiamond = rm.Material("C", rho=3.52, kind="plate") # pyright: ignore[reportArgumentType] -filterBe = rm.Material("Be", rho=1.85, kind="plate") # pyright: ignore[reportArgumentType] -filterSi3N4 = rm.Material( - ["Si", "N"], quantities=[3, 4], rho=3.44, kind="plate" -) # pyright: ignore[reportArgumentType] -filterAl = rm.Material("Al", rho=2.69, kind="plate") # pyright: ignore[reportArgumentType] -filterGraphite = rm.Material("C", rho=2.266, kind="plate") # pyright: ignore[reportArgumentType] - -# General parameters -sourceHeight = 0 - -# Synchrotron -synchrotron = namedtuple( - "synchrotron", ["eE", "eI", "eEspread", "eEpsilonX", "eEpsilonZ", "betaX", "betaZ"] -) - -sls1 = synchrotron( - eE=2.4, eI=0.4, eEspread=0.878e-3, eEpsilonX=5.63, eEpsilonZ=0.007, betaX=0.45, betaZ=14.4 -) - -sls2 = synchrotron( - eE=2.7, eI=0.4, eEspread=1.147e-3, eEpsilonX=0.156, eEpsilonZ=0.01, betaX=0.18, betaZ=4.6 -) - -# Source -bendingMagnet = namedtuple("bendingMagnet", ["name", "center", "sync", "B0"]) - -sls1_14t = bendingMagnet(name="FE-BM-SLS1-1.4T", center=(0, 0, 0), sync=sls1, B0=1.4) - -sls2_21t = bendingMagnet(name="FE-BM-SLS2-2.1T", center=(0, 0, 0), sync=sls2, B0=2.1) - -sls2_35t = bendingMagnet(name="FE-BM-SLS2-3.5T", center=(0, 0, 0), sync=sls2, B0=3.5) - -sls2_50t = bendingMagnet(name="FE-BM-SLS2-5.0T", center=(0, 0, 0), sync=sls2, B0=5.0) - -# FE slits -fe_slits = namedtuple("slits", ["name", "center", "center1", "center2", "maxDivH", "maxDivV"]) - -feSlits = fe_slits( - name="FE-SLITS", - center=(0, 6117, sourceHeight), - center1=(0, 5038.4, sourceHeight), - center2=(0, 5282.9, sourceHeight), - maxDivH=1.8e-3, - maxDivV=0.8e-3, -) - -# Filters -filt = namedtuple( - "filt", ["name", "center", "pitch", "limPhysX", "limPhysY", "surface", "material", "thickness"] -) - -feWindow = filt( - name="FE-WINDOW", - center=(0.0, 6158, sourceHeight), - pitch=np.pi / 2, - limPhysX=(-6, 6), - limPhysY=(-3.0, 3.0), - surface="None", - material=filterDiamond, - thickness=0.1, -) -feWindow = feWindow._replace( - surface="CVD Diamond window {0:0.0f} $\\mu$m".format(feWindow.thickness * 1e3) -) - -feFilt = filt( - name="FE-FI", - center=(0.0, 6590, sourceHeight), - pitch=np.pi / 2, - limPhysX=(-15, 15), - limPhysY=(-10, 10), - surface="None", - material=filterGraphite, - thickness=0.25, -) -feFilt = feFilt._replace(surface="Graphite filter {0:0.0f} $\\mu$m".format(feFilt.thickness * 1e3)) - -# Collimating mirror -collimatingMirror = namedtuple( - "collimatingMirror", - [ - "name", - "center", - "surface", - "material", - "limPhysX", - "limPhysY", - "limOptX", - "limOptY", - "R", - "pitch", - "jack1", - "jack2", - "jack3", - "tx1", - "tx2", - ], -) - -cm = collimatingMirror( - name="FE-CM", - center=[0, 7560.8, sourceHeight], - surface=("Pt", "Si", "Rh"), - material=(stripePt, stripeSi, stripeRh), - limPhysX=(-30, 30), - limPhysY=(-600, 600), - limOptX=((-21, -0.5, 11), (-4, 9.5, 23)), - limOptY=((-500, -500, -500), (500, 500, 500)), - R=[3e6, 15e6], - pitch=[1.4e-3, 4.5e-3], - jack1=[0.0, 7210.0, 0.0], # Tripod X, Y, Z (global) - jack2=[-210.0, 8310.0, 0.0], - jack3=[210.0, 8310.0, 0.0], - tx1=[0.0, -575.5], # X-Stage 1 [x, y] (local) - tx2=[0.0, 575], -) # X-Stage 2 - -apertures = namedtuple("apertures", ["name", "center", "opening"]) - -fePS = apertures( - name="FE-PS", center=[0, 8760, sourceHeight], opening=[-39 / 2, 39 / 2, -10, 29] -) # left, right, bottom, top - -opWbBsBlock = apertures( - name="OP-WB-BS-BLOCK", center=[0.0, 13606 - 135, sourceHeight], opening=[-18.0, 18.0, 42, 76] -) # left, right, bottom, top - -opSlits1 = apertures( - name="OP-SLITS 1", center=[0, 14145 - 135, sourceHeight], opening=[-35 / 2, 35 / 2, 47.5, 82.5] -) - -# OP Beam Monitors -op_bm = namedtuple("op_bm", ["name", "center"]) - -opBM1 = op_bm(name="OP Beam Monitor 1", center=(0, 14525 - 135, sourceHeight)) - -opBM2 = op_bm(name="OP Beam Monitor 2", center=(0, 17161.6 - 135, sourceHeight)) - -# Monochromator -monochromator = namedtuple( - "monochromator", - [ - "name", - "center", - "xtal", - "material1", - "material2", - "xtalWidth", - "xtalOffsetX", - "xtalLength1", - "xtalLength2", - "xtalGap", - "rotOffset", - "heightOffset", - "braggLim", - "jack1", - "jack2", - "jack3", - "tx", - ], -) - -mo1 = monochromator( - name="OP-CCM1", - center=[0.0, 11670 - 135, sourceHeight], - xtal=("Si311", "Si111"), - material1=(si311_1, si111_1), - material2=(si311_2, si111_2), - xtalWidth=(20, 20), - xtalOffsetX=(19.2, -19.2), - xtalLength1=(60, 60), - xtalLength2=(60, 60), - xtalGap=(8, 8), - rotOffset=6, # not sure what it is - heightOffset=8.5, # not sure what it is - braggLim=[4, 35], - jack1=[0.0, 11350.0, 0.0], # Tripod not available! - jack2=[-400.0, 12350.0, 0.0], - jack3=[400.0, 12350.0, 0.0], - tx=0.0, -) # X-Stage [x] - -# Focusing mirror -focusingMirror = namedtuple( - "focusingMirror", - [ - "name", - "center", - "surfaceToroid", - "materialToroid", - "limPhysXToroid", - "limPhysYToroid", - "limOptXToroid", - "limOptYToroid", - "R", - "pitch", - "r", - "xToroid", - "hToroid", - "jack1", - "jack2", - "jack3", - "tx1", - "tx2", - ], -) - -OFFSET_TRX = 46.8735 - -fm = focusingMirror( - name="OP-FM", - center=[0.0, 15580 - 135, sourceHeight], - surfaceToroid=("Rh", "Pt"), - materialToroid=(stripeRh, stripePt), - limPhysXToroid=(-54.0, 54.0), - limPhysYToroid=(-565.0, 565.0), - limOptXToroid=( - (43.388 + OFFSET_TRX, -4.865 + OFFSET_TRX), - (4.865 + OFFSET_TRX, -40.882 + OFFSET_TRX), - ), - limOptYToroid=((-500.0, -500.0), (500.0, 500.0)), - R=[3e6, 15e6], - pitch=[1.4e-3, 4.5e-3], - r=[30, 20], - xToroid=[24.126 + OFFSET_TRX, -22 + OFFSET_TRX], # offset in local x - hToroid=[7.0, 11.3], # depth of the cylinder at x = xCylinder1 and x = xCylinder2. - jack1=[0.0, 14980.0, 0.0], - jack2=[-75.0, 16180.0, 0.0], - jack3=[75.0, 16180.0, 0.0], - tx1=[0.0, -575.0], # X-Stage 1 [x, y] - tx2=[0.0, 575.0], -) # X-Stage 2 [x, y] - -# Entry wall experimental hutch: 21593 mm from source (SLS2) - -# Exit window -ehWindow = filt( - name="EH-WINDOW", - center=(0.0, 22063, sourceHeight), - pitch=np.pi / 2, - limPhysX=(-10.0, 10.0), - limPhysY=(17.5, 92.5), - surface="None", - material=filterBe, - thickness=0.25, -) -ehWindow = ehWindow._replace( - surface="Beryllium window {0:0.0f} $\\mu$m".format(ehWindow.thickness * 1e3) -) - -# Sample -sample = namedtuple("sample", ["name", "center"]) - -es1 = sample(name="ES1", center=[0, 23823, sourceHeight]) -es2 = sample(name="ES2", center=[0, 25843, sourceHeight]) - -# Ionization chambers -ic = namedtuple("sample", ["name", "center"]) - -es1ic0 = ic(name="ES1 IC0", center=[0, 23633, sourceHeight]) -es1ic1 = ic(name="ES1 IC1", center=[0, 24383, sourceHeight]) -es1ic2 = ic(name="ES1 IC2", center=[0, 24723, sourceHeight]) +""" +X10DA / SuperXAS Beamline Parameters. +This file describes the parameter of each component of the SuperXAS beamline +to be used for raytracing and geometrical calculations. +""" + +from collections import namedtuple + +import numpy as np +import xrt.backends.raycing.materials as rm + +# XRT definitions +filterBeryl = rm.Material("Be", rho=1.85, kind="plate") # pyright: ignore[reportArgumentType] +filterDiamond = rm.Material("C", rho=3.52, kind="plate") # pyright: ignore[reportArgumentType] +filterGraphite = rm.Material("C", rho=2.266, kind="plate") # pyright: ignore[reportArgumentType] + +stripeSi = rm.Material("Si", rho=2.33) # pyright: ignore[reportArgumentType] +stripePt = rm.Material("Pt", rho=21.45) # pyright: ignore[reportArgumentType] +stripeRh = rm.Material("Rh", rho=12.41) # pyright: ignore[reportArgumentType] +stripeCr = rm.Material("Cr", rho=7.14) # pyright: ignore[reportArgumentType] +stripePyrex = rm.Material( + "Si", rho=2.20 +) # Use Si as bare element and the density of SiO2 # pyright: ignore[reportArgumentType] + +si111_1 = rm.CrystalSi(hkl=(1, 1, 1), tK=77) # first xtal surface +si311_1 = rm.CrystalSi(hkl=(3, 1, 1), tK=77) # first xtal surface +si333_1 = rm.CrystalSi(hkl=(3, 3, 3), tK=77) # first xtal surface +si511_1 = rm.CrystalSi(hkl=(5, 1, 1), tK=77) # first xtal surface +si111_2 = rm.CrystalSi(hkl=(1, 1, 1), tK=77) # second xtal surface +si311_2 = rm.CrystalSi(hkl=(3, 1, 1), tK=77) # second xtal surface +si333_2 = rm.CrystalSi(hkl=(3, 3, 3), tK=77) # second xtal surface +si511_2 = rm.CrystalSi(hkl=(5, 1, 1), tK=77) # second xtal surface + +filterDiamond = rm.Material("C", rho=3.52, kind="plate") # pyright: ignore[reportArgumentType] +filterBe = rm.Material("Be", rho=1.85, kind="plate") # pyright: ignore[reportArgumentType] +filterSi3N4 = rm.Material( + ["Si", "N"], quantities=[3, 4], rho=3.44, kind="plate" +) # pyright: ignore[reportArgumentType] +filterAl = rm.Material("Al", rho=2.69, kind="plate") # pyright: ignore[reportArgumentType] +filterGraphite = rm.Material("C", rho=2.266, kind="plate") # pyright: ignore[reportArgumentType] + +# General parameters +sourceHeight = 0 + +# Synchrotron +synchrotron = namedtuple( + "synchrotron", ["eE", "eI", "eEspread", "eEpsilonX", "eEpsilonZ", "betaX", "betaZ"] +) + +sls1 = synchrotron( + eE=2.4, eI=0.4, eEspread=0.878e-3, eEpsilonX=5.63, eEpsilonZ=0.007, betaX=0.45, betaZ=14.4 +) + +sls2 = synchrotron( + eE=2.7, eI=0.4, eEspread=1.147e-3, eEpsilonX=0.156, eEpsilonZ=0.01, betaX=0.18, betaZ=4.6 +) + +# Source +bendingMagnet = namedtuple("bendingMagnet", ["name", "center", "sync", "B0"]) + +sls1_14t = bendingMagnet(name="FE-BM-SLS1-1.4T", center=(0, 0, 0), sync=sls1, B0=1.4) + +sls2_21t = bendingMagnet(name="FE-BM-SLS2-2.1T", center=(0, 0, 0), sync=sls2, B0=2.1) + +sls2_35t = bendingMagnet(name="FE-BM-SLS2-3.5T", center=(0, 0, 0), sync=sls2, B0=3.5) + +sls2_50t = bendingMagnet(name="FE-BM-SLS2-5.0T", center=(0, 0, 0), sync=sls2, B0=5.0) + +# FE slits +fe_slits = namedtuple("slits", ["name", "center", "center1", "center2", "maxDivH", "maxDivV"]) + +feSlits = fe_slits( + name="FE-SLITS", + center=(0, 6117, sourceHeight), + center1=(0, 5038.4, sourceHeight), + center2=(0, 5282.9, sourceHeight), + maxDivH=1.8e-3, + maxDivV=0.8e-3, +) + +# Filters +filt = namedtuple( + "filt", ["name", "center", "pitch", "limPhysX", "limPhysY", "surface", "material", "thickness"] +) + +feWindow = filt( + name="FE-WINDOW", + center=(0.0, 6158, sourceHeight), + pitch=np.pi / 2, + limPhysX=(-6, 6), + limPhysY=(-3.0, 3.0), + surface="None", + material=filterDiamond, + thickness=0.1, +) +feWindow = feWindow._replace( + surface="CVD Diamond window {0:0.0f} $\\mu$m".format(feWindow.thickness * 1e3) +) + +feFilt = filt( + name="FE-FI", + center=(0.0, 6590, sourceHeight), + pitch=np.pi / 2, + limPhysX=(-15, 15), + limPhysY=(-10, 10), + surface="None", + material=filterGraphite, + thickness=0.25, +) +feFilt = feFilt._replace(surface="Graphite filter {0:0.0f} $\\mu$m".format(feFilt.thickness * 1e3)) + +# Collimating mirror +collimatingMirror = namedtuple( + "collimatingMirror", + [ + "name", + "center", + "surface", + "material", + "limPhysX", + "limPhysY", + "limOptX", + "limOptY", + "R", + "pitch", + "jack1", + "jack2", + "jack3", + "tx1", + "tx2", + ], +) + +cm = collimatingMirror( + name="FE-CM", + center=[0, 7560.8, sourceHeight], + surface=("Pt", "Si", "Rh"), + material=(stripePt, stripeSi, stripeRh), + limPhysX=(-30, 30), + limPhysY=(-600, 600), + limOptX=((-21, -0.5, 11), (-4, 9.5, 23)), + limOptY=((-500, -500, -500), (500, 500, 500)), + R=[3e6, 15e6], + pitch=[1.4e-3, 4.5e-3], + jack1=[0.0, 7210.0, 0.0], # Tripod X, Y, Z (global) + jack2=[-210.0, 8310.0, 0.0], + jack3=[210.0, 8310.0, 0.0], + tx1=[0.0, -575.5], # X-Stage 1 [x, y] (local) + tx2=[0.0, 575], +) # X-Stage 2 + +apertures = namedtuple("apertures", ["name", "center", "opening"]) + +fePS = apertures( + name="FE-PS", center=[0, 8760, sourceHeight], opening=[-39 / 2, 39 / 2, -10, 29] +) # left, right, bottom, top + +opWbBsBlock = apertures( + name="OP-WB-BS-BLOCK", center=[0.0, 13606 - 135, sourceHeight], opening=[-18.0, 18.0, 42, 76] +) # left, right, bottom, top + +opSlits1 = apertures( + name="OP-SLITS 1", center=[0, 14145 - 135, sourceHeight], opening=[-35 / 2, 35 / 2, 47.5, 82.5] +) + +# OP Beam Monitors +op_bm = namedtuple("op_bm", ["name", "center"]) + +opBM1 = op_bm(name="OP Beam Monitor 1", center=(0, 14525 - 135, sourceHeight)) + +opBM2 = op_bm(name="OP Beam Monitor 2", center=(0, 17161.6 - 135, sourceHeight)) + +# Monochromator +monochromator = namedtuple( + "monochromator", + [ + "name", + "center", + "xtal", + "material1", + "material2", + "xtalWidth", + "xtalOffsetX", + "xtalLength1", + "xtalLength2", + "xtalGap", + "rotOffset", + "heightOffset", + "braggLim", + "jack1", + "jack2", + "jack3", + "tx", + ], +) + +mo1 = monochromator( + name="OP-CCM1", + center=[0.0, 11670 - 135, sourceHeight], + xtal=("Si311", "Si111"), + material1=(si311_1, si111_1), + material2=(si311_2, si111_2), + xtalWidth=(20, 20), + xtalOffsetX=(19.2, -19.2), + xtalLength1=(60, 60), + xtalLength2=(60, 60), + xtalGap=(8, 8), + rotOffset=6, # not sure what it is + heightOffset=8.5, # not sure what it is + braggLim=[4, 35], + jack1=[0.0, 11350.0, 0.0], # Tripod not available! + jack2=[-400.0, 12350.0, 0.0], + jack3=[400.0, 12350.0, 0.0], + tx=0.0, +) # X-Stage [x] + +# Focusing mirror +focusingMirror = namedtuple( + "focusingMirror", + [ + "name", + "center", + "surfaceToroid", + "materialToroid", + "limPhysXToroid", + "limPhysYToroid", + "limOptXToroid", + "limOptYToroid", + "R", + "pitch", + "r", + "xToroid", + "hToroid", + "jack1", + "jack2", + "jack3", + "tx1", + "tx2", + ], +) + +OFFSET_TRX = 46.8735 + +fm = focusingMirror( + name="OP-FM", + center=[0.0, 15580 - 135, sourceHeight], + surfaceToroid=("Rh", "Pt"), + materialToroid=(stripeRh, stripePt), + limPhysXToroid=(-54.0, 54.0), + limPhysYToroid=(-565.0, 565.0), + limOptXToroid=( + (43.388 + OFFSET_TRX, -4.865 + OFFSET_TRX), + (4.865 + OFFSET_TRX, -40.882 + OFFSET_TRX), + ), + limOptYToroid=((-500.0, -500.0), (500.0, 500.0)), + R=[3e6, 15e6], + pitch=[1.4e-3, 4.5e-3], + r=[30, 20], + xToroid=[24.126 + OFFSET_TRX, -22 + OFFSET_TRX], # offset in local x + hToroid=[7.0, 11.3], # depth of the cylinder at x = xCylinder1 and x = xCylinder2. + jack1=[0.0, 14980.0, 0.0], + jack2=[-75.0, 16180.0, 0.0], + jack3=[75.0, 16180.0, 0.0], + tx1=[0.0, -575.0], # X-Stage 1 [x, y] + tx2=[0.0, 575.0], +) # X-Stage 2 [x, y] + +# Entry wall experimental hutch: 21593 mm from source (SLS2) + +# Exit window +ehWindow = filt( + name="EH-WINDOW", + center=(0.0, 22063, sourceHeight), + pitch=np.pi / 2, + limPhysX=(-10.0, 10.0), + limPhysY=(17.5, 92.5), + surface="None", + material=filterBe, + thickness=0.25, +) +ehWindow = ehWindow._replace( + surface="Beryllium window {0:0.0f} $\\mu$m".format(ehWindow.thickness * 1e3) +) + +# Sample +sample = namedtuple("sample", ["name", "center"]) + +es1 = sample(name="ES1", center=[0, 23823, sourceHeight]) +es2 = sample(name="ES2", center=[0, 25843, sourceHeight]) + +# Ionization chambers +ic = namedtuple("sample", ["name", "center"]) + +es1ic0 = ic(name="ES1 IC0", center=[0, 23633, sourceHeight]) +es1ic1 = ic(name="ES1 IC1", center=[0, 24383, sourceHeight]) +es1ic2 = ic(name="ES1 IC2", center=[0, 24723, sourceHeight]) diff --git a/debye_bec/bec_ipython_client/plugins/move_to_label.py b/debye_bec/bec_ipython_client/plugins/move_to_label.py deleted file mode 100644 index dd14970..0000000 --- a/debye_bec/bec_ipython_client/plugins/move_to_label.py +++ /dev/null @@ -1,82 +0,0 @@ -from __future__ import annotations - -import builtins -from typing import TYPE_CHECKING - -from bec_lib import bec_logger -from debye_bec.devices.absorber import STATUS as ABS_STATUS - -logger = bec_logger.logger -# import builtins to avoid linter errors -dev = builtins.__dict__.get("dev") - -class MoveToLabelError(Exception): - """Exception for the MoveToLabel function""" - -def move_to_label(): - """ - Function to move several motors to a specific position defined in the label dict. - """ - - label = get_device_conditions(label="digitalTwin") - - # Get absorber status and close if open - logger.info("Check Frontend Absorber Status") - abs_was_open = dev.abs.status.get() == ABS_STATUS.OPEN - if abs_was_open: - logger.info(" Close Frontend Absorber") - status = dev.abs.close() - status.wait() - - # Move Frontend Slits - logger.info("Move Frontend Slits into position") - devices = ["sldi_centerx", "sldi_centery", "sldi_gapx", "sldi_gapy"] - matches = {key: label[key] for key in devices if key in label} - statuses = [] - for device in matches.values(): - statuses.append(device['device'].move(device['value'])) - for status in statuses: - status.wait(timeout=30) - - # Move Collimating mirror - logger.info("Move Collimating Mirror into position") - if "cm_rotx" in label: # pitch - logger.info(" Move pitch into position") - surveyed_movement( - axis=label['cm_rotx'], - surveyed_axes= [ - {'device': dev.cm_rotz, 'abs_tol': 0.1}, - ] - ) - - # Restore absorber position - logger.info("Restore Frontend Absorber Status") - if abs_was_open: - status = dev.abs.open() - status.wait() - - -def surveyed_movement(axis, surveyed_axes): - """ - Moves an axis while surverying a set of axes. - - Args: - axis (DeviceCondition): Device condition - surveyed_axes (list): List of dicts (same format as DeviceCondition) - - Raises: - If during movement of axis, one of the surveyed axes moves out of tolerance. - """ - - for surv_ax in surveyed_axes: - surv_ax['old_value'] = surv_ax['device'].read() - status = axis['device'].move(axis['value']) - while status.status == 'RUNNING': - for surv_ax in surveyed_axes: - if abs(surv_ax['device'].read() - surv_ax['old_value']) > surv_ax['abs_tol']: - axis['device'].stop() - raise MoveToLabelError( - f"During movement of {axis['device'].name}, {surv_ax['device'].name} " + - f"started to move unexpectedly (old pos: {surv_ax['old_value']}, " + - f"current pos: {surv_ax['device'].read()})" - ) diff --git a/debye_bec/bec_ipython_client/startup/post_startup.py b/debye_bec/bec_ipython_client/startup/post_startup.py index 07d6da4..0fd9d82 100644 --- a/debye_bec/bec_ipython_client/startup/post_startup.py +++ b/debye_bec/bec_ipython_client/startup/post_startup.py @@ -34,3 +34,18 @@ to setup the prompts. """ # pylint: disable=invalid-name, unused-import, import-error, undefined-variable, unused-variable, unused-argument, no-name-in-module + +from bec_lib import bec_logger + +logger = bec_logger.logger + +logger.info("Using the Debye startup script.") + +from debye_bec.bec_ipython_client.plugins.digital_twin_core.digital_twin_core import DigitalTwinCore +from debye_bec.bec_ipython_client.plugins.auto_gain import AutoGain + +digital_twin = DigitalTwinCore() +logger.success("Digital Twin Core loaded. Use 'digital_twin' to access it.") + +auto_gain = AutoGain() +logger.success("Auto-Gain module loaded. Use 'auto_gain' to access it.") diff --git a/debye_bec/bec_widgets/widgets/client.py b/debye_bec/bec_widgets/widgets/client.py index c974339..94aa2c0 100644 --- a/debye_bec/bec_widgets/widgets/client.py +++ b/debye_bec/bec_widgets/widgets/client.py @@ -17,6 +17,7 @@ _Widgets = { "DigitalTwin": "DigitalTwin", "RestartServer": "RestartServer", "ScanControlXAS": "ScanControlXAS", + "Scheduler": "Scheduler", } @@ -69,7 +70,7 @@ class DigitalTwin(RPCBase): class RestartServer(RPCBase): - """Main widget of server restart widget""" + """Main widget of server restart widget.""" _IMPORT_MODULE = "debye_bec.bec_widgets.widgets.restart_server.restart_server" @@ -115,3 +116,102 @@ class ScanControlXAS(RPCBase): """ Take a screenshot of the dock area and save it to a file. """ + + +class Scheduler(RPCBase): + """Schedule, persist and execute a sequence of BEC scan/device commands.""" + + _IMPORT_MODULE = "debye_bec.bec_widgets.widgets.scheduler.scheduler" + + @rpc_call + def run_schedule(self): + """ + Run the schedule, continuing from wherever it last left off. + """ + + @rpc_call + def abort_schedule(self): + """ + Request that the schedule stop after the current item. + """ + + @rpc_call + def reset_schedule(self): + """ + Clear all execution state and start the schedule over from item 1. + """ + + @rpc_call + def get_status(self) -> "dict": + """ + RPC-exposed: current schedule state, e.g. for another widget or a script. + """ + + @rpc_call + def add_item( + self, + command: "str", + index: "int | None" = None, + kind: "str" = "custom", + form_state: "dict | None" = None, + ) -> "str": + """ + RPC-exposed: insert a new, PENDING command into the schedule. + Safe to call while the schedule is running. + + Args: + command: command text, evaluated the same way as the existing + items (against `scans`/`dev`) once the schedule runs. + index: position to insert at (0 = first). Clamped so the item + can never land before something already running or + finished. Defaults to appending at the end. + kind/form_state: optional structured description of how + `command` was built (see `schedule_item.ScheduleItem`), + used to reopen the Edit dialog pre-filled. Leave as + defaults for a plain, hand-typed command. + + Returns: + The new item's item_id. + """ + + @rpc_call + def edit_item( + self, + item_id: "str", + command: "str", + kind: "str" = "custom", + form_state: "dict | None" = None, + ): + """ + RPC-exposed: change the command of an item that has not started + yet. Raises `RuntimeError` for an item that is already + running/finished. + """ + + @rpc_call + def delete_item(self, item_id: "str"): + """ + RPC-exposed: remove an item that has not started yet. Raises + `RuntimeError` for an item that is already running/finished. + """ + + @rpc_call + def move_item(self, item_id: "str", new_index: "int"): + """ + RPC-exposed: move an item that has not started yet to a new + position (0 = first, but never before something already + running/finished). Raises `RuntimeError` for an item that is + already running/finished. + """ + + @rpc_call + def move_item_up(self, item_id: "str"): + """ + RPC-exposed: swap an item with the one directly before it. + """ + + @rpc_call + def move_item_down(self, item_id: "str"): + """ + RPC-exposed: swap an item with the one directly after it. + """ diff --git a/debye_bec/bec_widgets/widgets/data_viewer/data_viewer.py b/debye_bec/bec_widgets/widgets/data_viewer/data_viewer.py index d8addb3..f4ab865 100644 --- a/debye_bec/bec_widgets/widgets/data_viewer/data_viewer.py +++ b/debye_bec/bec_widgets/widgets/data_viewer/data_viewer.py @@ -55,6 +55,7 @@ class DataViewer(BECWidget, QWidget): self.current_row = 0 self.input.scan_sel.currentItemChanged_connect(self.scan_sel_changed) + self.input.scan_sel.itemDoubleClicked_connect(self.scan_sel_double_click) self.input.load_button.clicked_connect(self.load_scan_from_history) self.input.load_from_folder_button.clicked_connect(self.load_scan_from_folder) self.viewer.unload_button.clicked_connect(self.unload_all_scans) @@ -77,6 +78,12 @@ class DataViewer(BECWidget, QWidget): """Updates the current row value of the scan selection list""" self.current_row = kwargs["value"]().row() + @SafeSlot() + def scan_sel_double_click(self, *_, **kwargs): + """Updates the current row value of the scan selection list and loads the scan""" + self.current_row = kwargs["value"]().row() + self.load_scan_from_history() + @SafeSlot() def open_in_file_manager(self, *_): """Open the scan folder in the systems default file manager""" diff --git a/debye_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py b/debye_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py index 7242c1e..4f9add1 100644 --- a/debye_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py +++ b/debye_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py @@ -105,6 +105,17 @@ class ListWidget(QWidget): ) ) + def itemDoubleClicked_connect(self, func): + """Connect a function to Double Click event.""" + self.value.itemDoubleClicked.connect( + partial( + func, + identifier=self.identifier, + value_obj=self.value, + value=lambda: self.value.currentIndex(), + ) + ) + def setDisabled(self, disable): self.value.setDisabled(disable) diff --git a/debye_bec/bec_widgets/widgets/designer_plugins.py b/debye_bec/bec_widgets/widgets/designer_plugins.py index 77ddd10..aaf7251 100644 --- a/debye_bec/bec_widgets/widgets/designer_plugins.py +++ b/debye_bec/bec_widgets/widgets/designer_plugins.py @@ -15,6 +15,7 @@ designer_plugins = { "debye_bec.bec_widgets.widgets.scan_control_xas.scan_control_xas", "ScanControlXAS", ), + "Scheduler": ("debye_bec.bec_widgets.widgets.scheduler.scheduler", "Scheduler"), } widget_icons = { @@ -22,4 +23,5 @@ widget_icons = { "DigitalTwin": "lightbulb", "RestartServer": "restart_alt", "ScanControlXAS": "tune", + "Scheduler": "assignment_add", } diff --git a/debye_bec/bec_widgets/widgets/digital_twin/__init__.py b/debye_bec/bec_widgets/widgets/digital_twin/__init__.py index a42cb09..e69de29 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/__init__.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/__init__.py @@ -1,3 +0,0 @@ -from .beamline import get_parameters - -parameters = get_parameters() diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py b/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py deleted file mode 100644 index 3047bdf..0000000 --- a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py +++ /dev/null @@ -1,297 +0,0 @@ -""" -Calculates the positions of axes based on a beamline config -""" - -import numpy as np -from bec_lib import bec_logger - -from .. import parameters as bl -from ..types import BeamlineId, ConfigDict - -logger = bec_logger.logger - - -def calc_positions(beamline: BeamlineId, cfg: ConfigDict) -> dict[str, dict[str, float]]: - """ - Calculates the positions of axes based on a beamline config. - - Args: - cfg(ConfigDict): Dictionary with beamline config - - Returns: - dict[str, dict[str, float]]: Dictionary mapping device names to dictionaries - containing a "value" key with the corresponding float value (position). - """ - - pos = {} - - ## FE slits - trxr = -np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1] - trxw = ( - (np.arctan(cfg["h_acc"]) * bl.feSlits.center1[1]) - / bl.feSlits.center1[1] - * bl.feSlits.center2[1] - ) - tryb = -np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1] - tryt = ( - (np.arctan(cfg["v_acc"]) * bl.feSlits.center1[1]) - / bl.feSlits.center1[1] - * bl.feSlits.center2[1] - ) - - xgap = trxw - trxr - ygap = tryt - tryb - - pos["sldi_gapx"] = {"value": xgap} - pos["sldi_gapy"] = {"value": ygap} - - ## Collimating Mirror - obj_dist = bl.cm.center[1] # object distance - beam_vs = 2 * obj_dist * np.tan(cfg["v_acc"]) # vertical size of beam after CM - - # TRX - if cfg["cm_stripe"] in bl.cm.surface: - index = bl.cm.surface.index(cfg["cm_stripe"]) - else: - raise ValueError(f"Requested stripe {cfg['cm_stripe']} not found in parameters!") - cm_trx = -(bl.cm.limOptX[0][index] + bl.cm.limOptX[1][index]) / 2 - pos["cm_trx"] = {"value": cm_trx} - - # TRY - height = obj_dist * np.tan(cfg["v_acc"]) ** 2 * 1 / np.tan(cfg["cm_pitch"]) - pos["cm_try"] = {"value": height} - - # Pitch - pos["cm_rotx"] = { - "value": -cfg["cm_pitch"] * 1e3 - } # invert and convert to mrad (same as EGU of rotx axis) - - # Bending Radius - radius = ( - 2.0 * obj_dist / np.sin(cfg["cm_pitch"]) - ) # Elements of modern X-ray Physics, page 108 ff. - pos["cm_bnd_radius"] = {"value": radius * 1e-6} # Convert to km - - ## Monochromator - if cfg["mo1_mode"] == "Monochromatic": - # Add 2x CM pitch to the bragg angle - bragg = cfg["mo1_bragg"] - elif cfg["mo1_mode"] == "Pinkbeam": - # Align xtal surfaces parallel to beam - bragg = 0 - else: - raise ValueError("Monochromator mode not supported") - pos["mo1_bragg_angle"] = {"value": bragg / np.pi * 180} # Bragg angle in deg - - # TRY, Height - l = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) - yhor = l * np.cos(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"])) - yver = yhor * np.tan(2.0 * cfg["cm_pitch"]) - - if cfg["mo1_mode"] == "Monochromatic": - beam_offset_mo1 = ( - l * np.sin(2.0 * (cfg["mo1_bragg"] + cfg["cm_pitch"])) - yver - ) # Resultat ist korrekt! - elif cfg["mo1_mode"] == "Pinkbeam": - beam_offset_mo1 = 0 - else: - raise ValueError("Monochromator mode not supported") - - def csc(a): - return 1 / np.sin(a) - - def cot(a): - return 1 / np.tan(a) - - # calculate height of center of first crystal surface - f = bl.mo1.rotOffset # rotation offset, mm - d = bl.mo1.heightOffset # xtal height offset, mm - c = d * csc(cfg["mo1_bragg"]) - f * cot(cfg["mo1_bragg"]) - - # Calculate height of center of rotation - b = np.sqrt( - d**2 * csc(cfg["mo1_bragg"]) ** 2 - - 2 * d * f * cot(cfg["mo1_bragg"]) * csc(cfg["mo1_bragg"]) - + f**2 * cot(cfg["mo1_bragg"]) ** 2 - + f**2 - ) - h = np.cos(np.pi / 2 - np.arctan(f / c) - cfg["mo1_bragg"] - 2 * cfg["cm_pitch"]) * b - h2 = ((bl.mo1.center[1] - bl.cm.center[1]) - np.sqrt(b**2 - h**2)) * np.tan(2 * cfg["cm_pitch"]) - height_mo1_real = ( - h + h2 - ) # per design, the height should not change if the pitch of the CM is not changed! - if cfg["mo1_mode"] == "Monochromatic": - pass - elif cfg["mo1_mode"] == "Pinkbeam": - height_mo1_real = ( - height_mo1_real - 13 - ) # Move down to let beam pass between both crystal without touching copper cooler - else: - raise ValueError("Monochromator mode not supported") - pos["mo1_try"] = {"value": height_mo1_real} - - # TRX, Crystal selection - if cfg["mo1_mode"] == "Monochromatic": - xtal = cfg["mo1_xtal"].translate( - str.maketrans("", "", "()") - ) # Remove brackets from xtal name to conform with parameters - if xtal in bl.mo1.xtal: - index = bl.mo1.xtal.index(xtal) - else: - raise ValueError(f"Requested xtal {xtal} not found in parameters!") - pos["mo1_trx"] = {"value": bl.mo1.xtalOffsetX[index]} - else: - pos["mo1_trx"] = {"value": 0} - - diag = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) # Calculations for Mono - dz = diag * np.cos(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) - - ## Slits 1 - d = bl.opSlits1.center[1] - bl.cm.center[1] - dz - sl1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 - pos["sl1_centery"] = {"value": sl1_beam_height} - pos["sl1_gapy"] = {"value": beam_vs} - - ## Beam Monitor 1 - d = bl.opBM1.center[1] - bl.cm.center[1] - dz - bm1_beam_height = d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 - pos["bm1_try"] = {"value": bm1_beam_height} - - ## Focusing Mirror - p = bl.fm.center[1] - q = cfg["smpl"] - bl.fm.center[1] - f = (p * q) / (p + q) # focal length - - # Bender radius - if cfg["fm_qy"] is None: - radius = 2 * q / np.sin(cfg["fm_rotx"]) # ideal bending radius for focused beam - else: - radius = ( - 2 * cfg["fm_qy"] / np.sin(cfg["fm_rotx"]) - ) # ideal bending radius for unfocused beam - pos["fm_bnd_radius"] = {"value": radius * 1e-6} # Convert to km - - # Pitch - d = bl.fm.center[1] - bl.cm.center[1] - dz - fm_rotx = ( - 2 * cfg["cm_pitch"] - cfg["fm_rotx"] - ) # calculate pitch in absolute values (according to horizontal plane) - pos["fm_rotx"] = { - "value": -fm_rotx * 1e3 - } # invert and convert to mrad (same as EGU of rotx axis) - - if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): - - # TRY - if cfg["fm_stripe"] == "Rh (toroid)": - r = bl.fm.r[0] - h_cyl = bl.fm.hToroid[0] - else: # PT toroid - r = bl.fm.r[1] - h_cyl = bl.fm.hToroid[1] - width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"] * 1e-3) - alpha = np.arccos(1 - width_beam**2 / (2 * r**2)) - h = r - (r * np.cos(alpha / 2)) - fm_beam_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg["fm_gain_height"] - fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1 - h_cyl + h / 2) * cfg[ - "fm_gain_height" - ] - pos["fm_try"] = {"value": fm_height} - - # TRX - if cfg["fm_stripe"] == "Rh (toroid)": - x_cyl = -bl.fm.xToroid[0] - else: - x_cyl = -bl.fm.xToroid[1] - pos["fm_trx"] = {"value": x_cyl} - - elif cfg["fm_stripe"] in ("Rh (flat)", "Pt (flat)"): - - # TRY - fm_height = (d * np.tan(2 * cfg["cm_pitch"]) + beam_offset_mo1) * cfg["fm_gain_height"] - fm_beam_height = fm_height - pos["fm_try"] = {"value": fm_height} - - # TRX - if cfg["fm_stripe"] == "Rh (flat)": - x_flat = -bl.fm.xFlat[0] - else: - x_flat = -bl.fm.xFlat[1] - pos["fm_trx"] = {"value": x_flat} - - else: - raise ValueError("FM Stripe selection not valid") - - pos["fm_roty"] = {"value": 0} - pos["fm_rotz"] = {"value": 0} - - ## Slits 2 - if hasattr(bl, "opSlits2"): - d = bl.opSlits2.center[1] - bl.fm.center[1] - sl2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["sl2_centery"] = {"value": sl2_beam_height} - pos["sl2_gapy"] = {"value": beam_vs} - - ## Beam Monitor 2 - d = bl.opBM2.center[1] - bl.fm.center[1] - bm2_beam_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["bm2_try"] = {"value": bm2_beam_height} - - ## Optical Table - - if beamline == "x01da": - # TRY - d = bl.ehWindow.center[1] - bl.fm.center[1] - ot_height = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["ot_try"] = {"value": ot_height} - - # Pitch - ot_pitch = -(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"]) - pos["ot_rotx"] = {"value": ot_pitch * 1e3} - - # TRZ ES1 - ot_es1_trz = cfg["smpl"] - pos["ot_es1_trz"] = {"value": ot_es1_trz} - - # ES0 exit window - pos["es0wi_try"] = { - "value": 5 - } # At 5mm, the middle of the window is 500 mm from the table (neutral position) - else: - # Exit window height - d = bl.ehWindow.center[1] - bl.fm.center[1] - es0wi_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["es0wi_try"] = {"value": es0wi_try} - - # ES1 table height - d = bl.es1.center[1] - bl.fm.center[1] - es1_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["es1_try"] = {"value": es1_try} - - # IC0 height - d = bl.es1ic0.center[1] - bl.fm.center[1] - es1ic0_try = ( - fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try - ) - pos["es1ic0_try"] = {"value": es1ic0_try} - - # IC1 height - d = bl.es1ic1.center[1] - bl.fm.center[1] - es1ic1_try = ( - fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try - ) - pos["es1ic1_try"] = {"value": es1ic1_try} - - # IC2 height - d = bl.es1ic2.center[1] - bl.fm.center[1] - es1ic2_try = ( - fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - es1_try - ) - pos["es1ic2_try"] = {"value": es1ic2_try} - - # ES2 table height - d = bl.es2.center[1] - bl.fm.center[1] - es2_try = fm_beam_height - d * np.tan(-(2 * cfg["cm_pitch"] - 2 * cfg["fm_rotx"])) - pos["es2_try"] = {"value": es2_try} - - return pos diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py b/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py deleted file mode 100644 index 135c76d..0000000 --- a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py +++ /dev/null @@ -1,70 +0,0 @@ -""" -Calculates the sideview coordinates based on a beamline config. -""" - -import numpy as np - -from .. import parameters as bl -from ..types import ConfigDict, DataDict - - -def calc_sideview(cfg: ConfigDict) -> DataDict: - """ - Calculates the sideview coordinates based on a beamline config. - - Args: - cfg(ConfigDict): Dictionary with beamline config - - Returns: - DataDict: Sideview data - """ - - beam: DataDict = {"x": [], "y": []} - - beam["x"] = [] - beam["y"] = [] - beam["x"].append(0) # Source - beam["y"].append(bl.sourceHeight) - beam["x"].append(bl.cm.center[1]) # CM - beam["y"].append(bl.sourceHeight) - if cfg["mo1_mode"] == "Monochromatic": - diag = bl.mo1.xtalGap[0] / np.sin(cfg["mo1_bragg"]) # Calculations for Mono - dy = diag * np.sin(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) - dz = diag * np.cos(2 * (cfg["cm_pitch"] + cfg["mo1_bragg"])) - beam["x"].append(bl.mo1.center[1] - dz / 2) # Mono 1.1 - beam["y"].append( - bl.sourceHeight - + np.tan(2 * cfg["cm_pitch"]) * (bl.mo1.center[1] - dz / 2 - bl.cm.center[1]) - ) - beam["x"].append(bl.mo1.center[1] + dz / 2) # Mono 1.2 - beam["y"].append( - bl.sourceHeight - + np.tan(2 * cfg["cm_pitch"]) * (bl.mo1.center[1] - dz / 2 - bl.cm.center[1]) - + dy - ) - beam["x"].append(bl.fm.center[1]) # FM - beam["y"].append( - bl.sourceHeight - + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1] - dz) - + dy - ) - beam["x"].append(cfg["smpl"]) # Experiment - beam["y"].append( - bl.sourceHeight - + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1] - dz) - + dy - + np.tan(2 * (cfg["cm_pitch"] - cfg["fm_rotx"])) * (cfg["smpl"] - bl.fm.center[1]) - ) - elif cfg["mo1_mode"] == "Pinkbeam": - beam["x"].append(bl.fm.center[1]) # FM - beam["y"].append( - bl.sourceHeight + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1]) - ) - beam["x"].append(cfg["smpl"]) # Experiment - beam["y"].append( - bl.sourceHeight - + np.tan(2 * cfg["cm_pitch"]) * (bl.fm.center[1] - bl.cm.center[1]) - + np.tan(2 * (cfg["cm_pitch"] - cfg["fm_rotx"])) * (cfg["smpl"] - bl.fm.center[1]) - ) - - return beam diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py b/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py deleted file mode 100644 index 2ef4a59..0000000 --- a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py +++ /dev/null @@ -1,159 +0,0 @@ -""" -Calculates the surface coordinates based on a beamline config. -""" - -import re - -import numpy as np -from bec_lib import bec_logger - -from .. import parameters as bl -from ..types import ConfigDict, SurfaceDict - -logger = bec_logger.logger - - -def calc_surfaces(cfg: ConfigDict) -> SurfaceDict: - """ - Calculates the surface coordinates based on a beamline config. - - Args: - cfg(ConfigDict): Dictionary with beamline config - - Returns: - SurfaceDict: Surface data - """ - - out: SurfaceDict = { - "cm": {"x": [], "y": []}, - "mo1_1": {"x": [], "y": []}, - "mo1_2": {"x": [], "y": []}, - "fm": {"x": [], "y": []}, - } - - # Collimating mirror - l = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"]) / np.sin(cfg["cm_pitch"]) - - w1 = 2 * (bl.cm.center[1] - l / 2) * np.tan(cfg["h_acc"]) - w2 = 2 * (bl.cm.center[1] + l / 2) * np.tan(cfg["h_acc"]) - - # index = bl.cm.surface.index(cfg["cm_stripe"]) - - cen = -cfg["cm_trx"] - - out["cm"]["x"] = [cen - w1 / 2, cen - w2 / 2, cen + w2 / 2, cen + w1 / 2] - out["cm"]["y"] = [-l / 2, l / 2, l / 2, -l / 2] - - # Monochromator - # calculate height of center of first crystal surface - c = bl.mo1.heightOffset * 1 / np.sin(cfg["mo1_bragg"]) - bl.mo1.rotOffset * 1 / np.tan( - cfg["mo1_bragg"] - ) - e = bl.mo1.xtalGap[0] / np.tan(cfg["mo1_bragg"]) - c - - xtal = cfg["mo1_xtal"].translate( - str.maketrans("", "", "()") - ) # Remove brackets from xtal name to conform with parameters - index = bl.mo1.xtal.index(xtal) - - xtal_pos = bl.mo1.xtalOffsetX[index] - xtal_length_1 = bl.mo1.xtalLength1[index] - xtal_length_2 = bl.mo1.xtalLength2[index] - - width_beam = 2 * bl.mo1.center[1] * np.tan(cfg["h_acc"]) - - height_beam = 2 * bl.cm.center[1] * np.tan(cfg["v_acc"]) - w = height_beam / np.sin(cfg["mo1_bragg"]) - - if cfg["mo1_mode"] == "Monochromatic": - out["mo1_1"]["x"] = [ - xtal_pos - width_beam / 2, - xtal_pos + width_beam / 2, - xtal_pos + width_beam / 2, - xtal_pos - width_beam / 2, - ] - out["mo1_1"]["y"] = [ - xtal_length_1 / 2 - c - w / 2, - xtal_length_1 / 2 - c - w / 2, - xtal_length_1 / 2 - c + w / 2, - xtal_length_1 / 2 - c + w / 2, - ] - out["mo1_2"]["x"] = [ - xtal_pos - width_beam / 2, - xtal_pos + width_beam / 2, - xtal_pos + width_beam / 2, - xtal_pos - width_beam / 2, - ] - out["mo1_2"]["y"] = [ - -xtal_length_2 / 2 + e - w / 2, - -xtal_length_2 / 2 + e - w / 2, - -xtal_length_2 / 2 + e + w / 2, - -xtal_length_2 / 2 + e + w / 2, - ] - else: # Pinkbeam - out["mo1_1"]["x"] = [] - out["mo1_1"]["y"] = [] - out["mo1_2"]["x"] = [] - out["mo1_2"]["y"] = [] - - if cfg["fm_stripe"] is None: - return out - # Focusing mirror - if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): - surface = bl.fm.surfaceToroid - stripe = re.sub(r"\s*\(.*?\)", "", cfg["fm_stripe"]).strip() - index = surface.index(stripe) - r = bl.fm.r[index] - else: - surface = bl.fm.surfaceFlat - stripe = re.sub(r"\s*\(.*?\)", "", cfg["fm_stripe"]).strip() - index = surface.index(stripe) - r = bl.fm.r[index] - off = -cfg["fm_trx"] - - width_beam = 2 * bl.fm.center[1] * np.tan(cfg["h_acc"]) - - if cfg["fm_stripe"] in ("Rh (toroid)", "Pt (toroid)"): - - l = height_beam / np.sin(cfg["fm_rotx"]) - alpha = np.arccos(1 - width_beam**2 / (2 * r**2)) - h = r - (r * np.cos(alpha / 2)) - z = h / np.tan(cfg["fm_rotx"]) - - x = [off - width_beam / 2, off - width_beam / 2] - y = [l / 2 - z / 2, -l / 2 - z / 2] - - res = 20 - x_elipse = np.linspace(0, np.pi, res) - y_elipse = np.linspace(0, np.pi, res) - x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse] - y_elipse = [width_beam * np.sin(i) * z / width_beam - l / 2 - z / 2 for i in y_elipse] - - x.extend(x_elipse) - y.extend(y_elipse) - - x.extend([off + width_beam / 2, off + width_beam / 2]) - y.extend([-l / 2 - z / 2, l / 2 - z / 2]) - - res = 50 - x_elipse = np.linspace(np.pi, 0, res) - y_elipse = np.linspace(np.pi, 0, res) - x_elipse = [-width_beam / 2 * np.cos(i) + off for i in x_elipse] - y_elipse = [width_beam * np.sin(i) * z / width_beam + l / 2 - z / 2 for i in y_elipse] - - x.extend(x_elipse) - y.extend(y_elipse) - - out["fm"]["x"] = x - out["fm"]["y"] = y - - else: # flat surface, no toroid - l = height_beam / np.sin(cfg["fm_rotx"]) - - w1 = 2 * (bl.fm.center[1] - l / 2) * np.tan(cfg["h_acc"]) - w2 = 2 * (bl.fm.center[1] + l / 2) * np.tan(cfg["h_acc"]) - - out["fm"]["x"] = [off - w1 / 2, off + w1 / 2, off + w2 / 2, off - w2 / 2] - out["fm"]["y"] = [-l / 2, -l / 2, l / 2, l / 2] - - return out diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py b/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py deleted file mode 100644 index 5c80afb..0000000 --- a/debye_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py +++ /dev/null @@ -1,519 +0,0 @@ -""" -Various calculations for the digital twin -""" - -import re -from typing import Literal, cast - -import numpy as np -from bec_lib import bec_logger -from scipy.interpolate import UnivariateSpline -from xrt.backends.raycing.physconsts import AVOGADRO, CHeVcm - -from .. import parameters as bl - -logger = bec_logger.logger - -H = 6.62606957e-34 -E = 1.602176634e-19 -C = 299792458 -RE = 2.8179e-15 - - -def sldi_gap_to_acc(sldi_gapx: float, sldi_gapy: float) -> tuple[float, float]: - """ - Calculate the slits acceptance based on the gap values - - Args: - sldi_gapx(float): GAPX value of the slits in mm - sldi_gapy(float): GAPY value of the slits in mm - - Returns: - tuple[float, float]: Horizontal and vertical acceptance in rad - """ - d1 = bl.feSlits.center1[1] - d2 = bl.feSlits.center2[1] - h_acc = np.tan(sldi_gapx / (d2 + d1)) - v_acc = np.tan(sldi_gapy / (d2 + d1)) - return h_acc, v_acc - - -def cm_trx_to_stripe(cm_trx: float) -> str | None: - """ - Based on the trx value of the collimating mirror, return - the correct stripe - - Args: - cm_trx(float): Collimating mirror trx value - - Returns - str | None: Stripe of the mirror, None if not found - """ - cm_stripe = None - for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]): - if low <= cm_trx <= high: - cm_stripe = name - return cm_stripe - - -def cm_stripe_to_trx(cm_stripe: str) -> float | None: - """ - Based on the stripe of the collimating mirror, return - the trx value - - Args: - cm_stripe(str): Stripe of the collimating mirror - - Returns: - float | None: TRX value of the stripe. None if not found - """ - for name, low, high in zip(bl.cm.surface, bl.cm.limOptX[0], bl.cm.limOptX[1]): - if cm_stripe == name: - return -(low + high) / 2 - return None - - -def fm_trx_to_stripe(fm_trx: float) -> str | None: - """ - Based on the trx value of the focusing mirror, return - the correct stripe - - Args: - fm_trx(float): focusing mirror trx value - - Returns - str | None: Stripe of the mirror, None if not found - """ - fm_stripe = None - if hasattr(bl.fm, "surfaceFlat"): - for name, low, high in zip(bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0]): - if low <= fm_trx <= high: - fm_stripe = name + " (flat)" - for name, low, high in zip(bl.fm.surfaceToroid, bl.fm.limOptXToroid[1], bl.fm.limOptXToroid[0]): - if low <= fm_trx <= high: - fm_stripe = name + " (toroid)" - return fm_stripe - - -def fm_stripe_to_trx(fm_stripe: str) -> float | None: - """ - Based on the stripe of the focusing mirror, return - the trx value - - Args: - fm_stripe(str): Stripe of the focusing mirror - - Returns: - float | None: TRX value of the stripe. None if not found - """ - if hasattr(bl.fm, "surfaceFlat"): - for name, low, high in zip(bl.fm.surfaceFlat, bl.fm.limOptXFlat[1], bl.fm.limOptXFlat[0]): - if fm_stripe == name + " (flat)": - return (low + high) / 2 - for name, low, high in zip(bl.fm.surfaceToroid, bl.fm.limOptXToroid[1], bl.fm.limOptXToroid[0]): - if fm_stripe == name + " (toroid)": - return -(low + high) / 2 - return None - - -def mo1_energy_resolution(xtal: Literal["Si111", "Si311"], energy: float) -> float: - """ - Calculate the energy resolution of the monochromator - - Args: - xtal(str): Xtal name. "Si111" or "Si311" - energy(float): Energy in eV - - Returns: - float: Energy resolution in eV - """ - index = bl.mo1.xtal.index(xtal) - crystal = bl.mo1.material1[index] - - dtheta = np.linspace(-30, 90, 601) - theta = crystal.get_Bragg_angle(energy) + dtheta * 1e-6 - refl = np.abs(crystal.get_amplitude(energy, np.sin(theta))[0]) ** 2 # single crystal - - refl2 = refl**2 # DCM with parallel crystals - - # FWHM of the DCM curve - spline = UnivariateSpline(dtheta, refl2 - refl2.max() / 2, s=0) - roots = cast(np.ndarray, spline.roots()) - r1, r2 = float(roots[0]), float(roots[1]) - fwhm_rad = (r2 - r1) * 1e-6 # µrad → rad - - # Energy resolution - theta_b = crystal.get_Bragg_angle(energy) - de_over_e = fwhm_rad / np.tan(theta_b) - de = de_over_e * energy - - # logger.info(f"DCM FWHM : {r2-r1:.2f} µrad") - # logger.info(f"ΔE/E : {dE_over_E:.2e}") - # logger.info(f"ΔE : {dE:.3f} eV at {E} eV") - - return de - - -def cm_reflectivity(cm_stripe: str, cm_pitch: float, energy: float) -> float: - """ - Calculate the reflectivity of the mirror stripe based - on the pitch and energy. - - Args: - cm_stripe(str): Mirror stripe - cm_pitch(float): Pitch of the mirror (beam incidence angle) - energy(float): Energy of the beam in eV - - Returns: - float: Reflectivity [0-1] - """ - if cm_stripe is None: - return np.nan - index = bl.cm.surface.index(cm_stripe) - rs, _ = bl.cm.material[index].get_amplitude(energy, np.sin(cm_pitch))[0:2] - refl = abs(rs) ** 2 - return refl - - -def fm_reflectivity(fm_stripe: str, fm_pitch: float, energy: float) -> float: - """ - Calculate the reflectivity of the mirror stripe based - on the pitch and energy. - - Args: - cm_stripe(str): Mirror stripe - cm_pitch(float): Pitch of the mirror (beam incidence angle) - energy(float): Energy of the beam in eV - - Returns: - float: Reflectivity [0-1] - """ - if fm_stripe is None: - return np.nan - if fm_stripe in ("Rh (toroid)", "Pt (toroid)"): - surface = bl.fm.surfaceToroid - material = bl.fm.materialToroid - stripe = re.sub(r"\s*\(.*?\)", "", fm_stripe).strip() - index = surface.index(stripe) - else: - surface = bl.fm.surfaceFlat - material = bl.fm.materialFlat - stripe = re.sub(r"\s*\(.*?\)", "", fm_stripe).strip() - index = surface.index(stripe) - rs, _ = material[index].get_amplitude(energy, np.sin(fm_pitch))[0:2] - refl = abs(rs) ** 2 - return refl - - -def mo1_bragg_angle( - mo_mode: Literal["Monochromatic", "Pinkbeam"], d_spacing: float, energy: float, cm_pitch: float -) -> tuple[float, float]: - """ - Calculate the bragg angle of the monochromator. - Corrects for the collimating mirror pitch. - - Args: - mo_mode(str): Monochromator mode. "Monochromatic" or "Pinkbeam" - d_spacing(float): D-spacing of the crystal in Angstrom - energy(float): Energy of the beam in eV - cm_pitch(float): Pitch of collimating mirror in rad - - Returns: - tuple[float, float]: Bragg angle and corrected bragg angle - """ - wl = C * H / (E * energy) - val = wl / (2 * d_spacing * 1e-10) - bragg_angle = 0 - if val > -1 and val < 1: - bragg_angle = np.asin(val) - if mo_mode == "Monochromatic": - # Add 2x CM pitch to the bragg angle - bragg_angle_cor = (2 * cm_pitch) + bragg_angle - else: - # Align xtal surfaces parallel to beam - bragg_angle_cor = 2 * cm_pitch - return bragg_angle, bragg_angle_cor - - -def fm_ideal_pitch( - fm_focus: Literal["Defocused", "Focused", "Manual"], - fm_stripe: str, - smpl: float, - sldi_hacc: float | None = None, - sldi_vacc: float | None = None, - fm_focx: float | None = None, - fm_focy: float | None = None, -) -> tuple[float, float | None]: - """ - Calculates the ideal pitch for the focusing mirror depending on the - focusing strategy. - If "Defocused" is chosed, sldi_hacc, sldi_vacc, fm_focx and fm_focy - must be provided. - - Args: - fm_focus(str): Focus strategy. "Defocused", "Focused" or "Manual - fm_stripe(str): Mirror stripe - smpl(float): Sample position in mm from source - sldi_hacc(float): Horizontal acceptance of frontend slits. Defaults to None - sldi_vacc(float): Vertical acceptance of frontend slits. Defaults to None - fm_focx(float): Requested horizontal spot size in mm. Defaults to None - fm_focy(float): Requested vertical spot size in mm. Defaults to None - - Returns: - tuple[float, float | None]: Pitch of mirror in rad, qy in mm - """ - - # logger.info("Calculate pitch and qy now...") - # logger.info(f"sldi_hacc: {sldi_hacc}") - # logger.info(f"sldi_vacc: {sldi_vacc}") - # logger.info(f"fm_stripe: {fm_stripe}") - # logger.info(f"smpl: {smpl}") - p_cm = bl.cm.center[1] # posCM - p = bl.fm.center[1] # posFM - q = smpl - bl.fm.center[1] # dist posFM to posEX - if fm_focus == "Defocused": - assert sldi_hacc is not None, "sldi_hacc must be provided for Defocused mode" - assert sldi_vacc is not None, "sldi_vacc must be provided for Defocused mode" - assert fm_focx is not None, "fm_focx must be provided for Defocused mode" - assert fm_focy is not None, "fm_focy must be provided for Defocused mode" - a = 2 * np.tan(sldi_hacc) * bl.fm.center[1] # Beam width at focusing mirror - # logger.info(f"a: {a}") - # logger.info(f"sldi_hacc: {sldi_hacc}") - # logger.info(f"bl.fm.center[1]: {bl.fm.center[1]}") - # logger.info(f"p: {p}") - # logger.info(f"q: {q}") - b = ( - 2 * np.tan(sldi_vacc) * bl.cm.center[1] - ) # Beam height at focusing mirror (collimated beam) - x = fm_focx - # logger.info(f"x: {x}") - x = 0.098821 * x**2 + 0.512344 * x # polynom to correct for spot size - # logger.info(f"x (corrected): {x}") - y = fm_focy - y = 3.183562 * y**2 + 1.258364 * y # polynom to correct for spot size - qx = q + x * p / a - qy = q + y * p_cm / b - f = (p * qx) / (p + qx) # focal length - # logger.info(f"qx: {qx}") - # logger.info(f"f: {f}") - else: # Calculate for focused beam on sample in "manual" and "focused" mode - qy = None - f = (p * q) / (p + q) # focal length - pitch = 0 - if "Rh" in fm_stripe: - pitch = np.arcsin(bl.fm.r[0] / (2 * f)) # ideal pitch for FM - if "Pt" in fm_stripe: - pitch = np.arcsin(bl.fm.r[1] / (2 * f)) # ideal pitch for FM - # logger.info(f"fm_pitch: {pitch}") - # logger.info(f"qy: {qy}") - return pitch, qy - - -def calc_beamsize( - sldi_hacc: float, - sldi_vacc: float, - fm_stripe: str, - fm_pitch: float, - fm_radius: float, - smpl: float, -) -> tuple[float, float | None]: - """ - Calculate the resulting beamsize according to the input parameters - - Args: - sldi_hacc(float): Horizontal acceptance of frontend slits - sldi_vacc(float): Vertical acceptance of frontend slits - fm_stripe(str): Mirror stripe - fm_pitch(float): Focusing mirror pitch in rad - fm_radius(float): Focusing mirror bender radius in m - smpl(float): Sample position in mm from source - - Returns: - tuple[float, float | None]: horizontal spot size, vertical spot size, both in mm - """ - - # logger.info("Calculate beamsize now...") - # logger.info(f"sldi_hacc: {sldi_hacc}") - # logger.info(f"sldi_vacc: {sldi_vacc}") - # logger.info(f"fm_stripe: {fm_stripe}") - # logger.info(f"fm_pitch: {fm_pitch}") - # logger.info(f"fm_radius: {fm_radius}") - # logger.info(f"smpl: {smpl}") - p_cm = bl.cm.center[1] # posCM - p = bl.fm.center[1] # posFM - q = smpl - bl.fm.center[1] # dist posFM to posEX - qy = fm_radius * np.sin(fm_pitch) / 2 - a = 2 * np.tan(sldi_hacc) * bl.fm.center[1] # Beam width at focusing mirror - b = 2 * np.tan(sldi_vacc) * bl.cm.center[1] # Beam height at focusing mirror (collimated beam) - f = 0 - if "Rh" in fm_stripe: - f = bl.fm.r[0] / (2 * np.sin(fm_pitch)) - if "Pt" in fm_stripe: - f = bl.fm.r[1] / (2 * np.sin(fm_pitch)) - qx = p * f / (p - f) - x = a * (qx - q) / p - y = b * (qy - q) / p_cm - # Change this | to a plus if calculation is not correct - fm_focx = -4 * (64043 - 125000 * np.sqrt(0.26249637 + 0.395284 * x)) / 98821 - # Change this | to a plus if calculation is not correct - fm_focy = -1 * (314591 - 250000 * np.sqrt(1.58347995 + 12.734248 * y)) / 1591781 - # logger.info(f"f: {f}") - # logger.info(f"qx: {qx}") - # logger.info(f"qy: {qy}") - # logger.info(f"fm_focx: {fm_focx}") - # logger.info(f"fm_focy: {fm_focy}") - return fm_focx, fm_focy - - -def cm_critical_angle(cm_stripe: Literal["Si", "Pt", "Rh"], energy) -> float: - """ - Calculate the critical angle of the mirror stripe - - Args: - cm_stripe(str): Mirror stripe. "Si", "Pt" or "Rh" - energy(float): Energy in eV - - Returns: - float: Critical angle in rad - """ - if cm_stripe == "Si": - stripe = bl.stripeSi - elif cm_stripe == "Pt": - stripe = bl.stripePt - else: - stripe = bl.stripeRh - w = CHeVcm / 100 / energy # convert energy [eV] to wavelength [m] - f1 = stripe.elements[0].Z + np.real(stripe.elements[0].get_f1f2(energy)) - number_density = stripe.rho * 1e3 * AVOGADRO / (stripe.elements[0].mass / 1e3) - critical_angle = np.sqrt(number_density * RE * w**2 * f1 / np.pi) - return critical_angle - - -def mirror_surface_geometries( - mirror: Literal["cm", "fm_toroid", "fm_flat"], -) -> dict[str, tuple[float, float, float, float]]: - """ - Return the mirror stripe geometries - - Args: - mirror(str): Mirror. "cm", "fm_toroid" or "fm_flat" - - Returns: - dict[str, tuple[float, float, float, float]]: Dictionary mapping surface - names to tuples of (x, y, width, height). - """ - if mirror == "cm": - surface = bl.cm.surface - lim_opt_x = bl.cm.limOptX - lim_opt_y = bl.cm.limOptY - elif mirror == "fm_toroid": - surface = bl.fm.surfaceToroid - lim_opt_x = bl.fm.limOptXToroid - lim_opt_y = bl.fm.limOptYToroid - elif mirror == "fm_flat": - surface = bl.fm.surfaceFlat - lim_opt_x = bl.fm.limOptXFlat - lim_opt_y = bl.fm.limOptYFlat - else: - raise ValueError(f"Requested mirror {mirror} not available!") - geom = {} - for sf, lx, hx, ly, hy in zip(surface, lim_opt_x[0], lim_opt_x[1], lim_opt_y[0], lim_opt_y[1]): - geom[sf] = (lx, ly, hx - lx, hy - ly) - return geom - - -def mo_surface_geometries( - mo: Literal["mo1"], plane: Literal[0, 1] -) -> dict[str, tuple[float, float, float, float]]: - """ - Return the monochromator xtal geometries - - Args: - mo(str): Monochromator. Only "mo1" implemented - plane(int): Surface of xtal. 0 and 1 (First and second) - - Returns: - dict[str, tuple[float, float, float, float]]: Dictionary mapping surface - names to tuples of (x, y, width, height). - """ - if mo == "mo1": - xtal = bl.mo1.xtal - xtal_width = bl.mo1.xtalWidth - xtal_offset_x = bl.mo1.xtalOffsetX - if plane == 0: - xtal_length = bl.mo1.xtalLength1 - else: - xtal_length = bl.mo1.xtalLength2 - else: - return {} - geom = {} - for sf, w, offx, length in zip(xtal, xtal_width, xtal_offset_x, xtal_length): - geom[sf] = (offx - w / 2, -length / 2, w, length) - return geom - - -def wall_geometries() -> list[list[float]]: - """ - Return the wall geometries - - Returns: - list[list[float]]: List of [x, y, width, height] geometry values for each wall. - """ - geom = [] - if not hasattr(bl, "walls"): - return geom - for i, _ in enumerate(bl.walls.start): - geom.append( - [ - bl.walls.start[i], - bl.walls.height[i][0], - bl.walls.end[i] - bl.walls.start[i], - bl.walls.height[i][1] - bl.walls.height[i][0], - ] - ) - return geom - - -def pipe_geometries() -> list[dict[str, np.ndarray]]: - """ - Return the wall geometries - - Returns: - list[dict[str, np.ndarray]]: List of dictionaries with keys "x" and "y", - each containing a numpy array of two float values representing - the start and end coordinates of the pipe top and bottom edges. - """ - pipes = [] - if not hasattr(bl, "vacuum_pipes"): - return pipes - for i, _ in enumerate(bl.vacuum_pipes.center): - top = bl.vacuum_pipes.center[i] + bl.vacuum_pipes.diameter[i] / 2 + bl.sourceHeight - bottom = bl.vacuum_pipes.center[i] - bl.vacuum_pipes.diameter[i] / 2 + bl.sourceHeight - pipes.append( - { - "x": np.array([bl.vacuum_pipes.start[i], bl.vacuum_pipes.end[i]]), - "y": np.array([top, top]), - } - ) - pipes.append( - { - "x": np.array([bl.vacuum_pipes.start[i], bl.vacuum_pipes.end[i]]), - "y": np.array([bottom, bottom]), - } - ) - return pipes - - -def table_to_smpl_pos(table: str) -> float: - """ - Return the sample position based on the table name. - - Args: - table (str): Table name, e.g. ES1 or ES2 - """ - - if table == bl.es1.name: - return bl.es1.center[1] - if table == bl.es2.name: - return bl.es2.center[1] - raise ValueError(f"Table {table} not found in beamline parameter file") diff --git a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py index f5fcbfe..2a01e02 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py @@ -3,7 +3,6 @@ Digital Twin: Custom BEC widget to support the beamline alignment. """ import sys -from pathlib import Path from typing import Literal, cast import numpy as np @@ -35,38 +34,18 @@ from qtpy.QtWidgets import ( QWidget, ) +from ....bec_ipython_client.plugins.digital_twin_core.beamline import get_beamline_id +from ....bec_ipython_client.plugins.digital_twin_core.digital_twin_core import DigitalTwinCore +from ....bec_ipython_client.plugins.digital_twin_core.types import ConfigDict from ..edge_selector import EdgeSelector -from .beamline import get_beamline_id -from .calculations.calc_positions import calc_positions -from .calculations.calc_sideview import calc_sideview -from .calculations.calc_surfaces import calc_surfaces -from .calculations.calc_varia import ( - calc_beamsize, - cm_critical_angle, - cm_reflectivity, - cm_stripe_to_trx, - cm_trx_to_stripe, - fm_ideal_pitch, - fm_reflectivity, - fm_stripe_to_trx, - fm_trx_to_stripe, - mo1_bragg_angle, - mo1_energy_resolution, - sldi_gap_to_acc, - table_to_smpl_pos, -) from .panels.input_panel import InputPanel from .panels.mover_panel import MoverPanel from .panels.plots import SideviewPlot, SurfacePlots from .panels.settings_panel import SettingsPanel -from .types import ConfigDict from .widgets.qt_widgets import ComboBox, InputNumberField logger = bec_logger.logger -OFFSET_FILE_X01DA = Path(__file__).with_name("x01da_offsets.yaml") -OFFSET_FILE_X10DA = Path(__file__).with_name("x10da_offsets.yaml") - X01DA_E_MIN = 4500 X01DA_E_MAX = 60000 X10DA_E_MIN = 4500 @@ -85,17 +64,12 @@ class DigitalTwin(BECWidget, QWidget): super().__init__(parent=parent, *arg, **kwargs) self.get_bec_shortcuts() + self.core = DigitalTwinCore() + self.beamline = get_beamline_id() # Debugging, override beamline! # self.beamline = BeamlineId.X10DA - self.offset_file = Path() - match self.beamline: - case "x01da": - self.offset_file = OFFSET_FILE_X01DA - case "x10da": - self.offset_file = OFFSET_FILE_X10DA - # Check if devices are all in config self.check_bec_config() self.bec_dispatcher.connect_slot( @@ -123,8 +97,8 @@ class DigitalTwin(BECWidget, QWidget): self.plot_layout = QVBoxLayout(self.plot_widget) self.plot_layout.setContentsMargins(4, 4, 4, 4) self.plot_layout.setSpacing(6) - self.sideview_plot = SideviewPlot() - self.surface_plots = SurfacePlots(self.beamline) + self.sideview_plot = SideviewPlot(self.core) + self.surface_plots = SurfacePlots(self.beamline, self.core) self.plot_layout.addWidget(self.sideview_plot, stretch=1) self.plot_layout.addWidget(self.surface_plots, stretch=1) self.plot_widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) @@ -172,7 +146,6 @@ class DigitalTwin(BECWidget, QWidget): self.edge_selector_energy = 0.0 self.bragg_angle = 0.0 self.qy = 0.0 - self.offsets = {} # Initialize all values self.load_offsets(recalculate=False) @@ -400,9 +373,9 @@ class DigitalTwin(BECWidget, QWidget): fm_qy = self.qy cm_stripe = self.input.cm_stripe.currentText() - cm_trx = cm_stripe_to_trx(cm_stripe) + cm_trx = self.core.cm_stripe_to_trx(cm_stripe) fm_stripe = self.input.fm_stripe.currentText() - fm_trx = fm_stripe_to_trx(fm_stripe) + fm_trx = self.core.fm_stripe_to_trx(fm_stripe) assert cm_trx is not None, f"No cm_trx found for given stripe {cm_stripe}!" assert fm_trx is not None, f"No fm_trx found for given stripe {fm_stripe}!" @@ -412,7 +385,7 @@ class DigitalTwin(BECWidget, QWidget): smpl = self.input.smpl.value() case ComboBox(): table = self.input.smpl.currentText() - smpl = table_to_smpl_pos(table) + smpl = self.core.table_to_smpl_pos(table) config: ConfigDict = { "energy": self.input.energy.value(), @@ -434,16 +407,7 @@ class DigitalTwin(BECWidget, QWidget): # Apply offsets if apply_offset: - for axis, _ in config.items(): - if axis in self.offsets: - axis_offsets = self.offsets[axis] - if "modifier" in axis_offsets and "offset" in axis_offsets: - for idx, rng in enumerate(axis_offsets["modifier"]["range"]): - if rng[0] < config[axis_offsets["modifier"]["axis"]] < rng[1]: - config[axis] += axis_offsets["offset"][idx] - break - elif "offset" in axis_offsets: - config[axis] += axis_offsets["offset"] + config = self.core.apply_offsets(config, nested_config=False) # Convert to SI units! config["h_acc"] *= 1e-3 @@ -469,12 +433,12 @@ class DigitalTwin(BECWidget, QWidget): mo1_bragg = self.dev.mo1_bragg.read(cached=True) sldi_gapx = self.dev.sldi_gapx.read(cached=True)["sldi_gapx"]["value"] sldi_gapy = self.dev.sldi_gapy.read(cached=True)["sldi_gapy"]["value"] - h_acc, v_acc = sldi_gap_to_acc(sldi_gapx, sldi_gapy) + h_acc, v_acc = self.core.sldi_gap_to_acc(sldi_gapx, sldi_gapy) cm_trx = self.dev.cm_trx.read(cached=True)["cm_trx"]["value"] - cm_stripe = cm_trx_to_stripe(-cm_trx) + cm_stripe = self.core.cm_trx_to_stripe(-cm_trx) cm_pitch = self.dev.cm_rotx.read(cached=True)["cm_rotx"]["value"] fm_trx = self.dev.fm_trx.read(cached=True)["fm_trx"]["value"] - fm_stripe = fm_trx_to_stripe(-fm_trx) + fm_stripe = self.core.fm_trx_to_stripe(-fm_trx) fm_rotx = self.dev.fm_rotx.read(cached=True)["fm_rotx"]["value"] fm_rotx_real = 2 * cm_pitch - fm_rotx @@ -483,7 +447,7 @@ class DigitalTwin(BECWidget, QWidget): smpl = self.dev.ot_es1_trz.read(cached=True)["ot_es1_trz"]["value"] case ComboBox(): table = self.input.smpl.currentText() - smpl = table_to_smpl_pos(table) + smpl = self.core.table_to_smpl_pos(table) raw = { # Config in SI units! "energy": mo1_bragg["mo1_bragg"]["value"], @@ -592,22 +556,13 @@ class DigitalTwin(BECWidget, QWidget): pos["ot_es1_trz"] = self.dev.ot_es1_trz.read(cached=True)["ot_es1_trz"]["value"] # Removing offsets - for axis, _ in pos.items(): - if axis in self.offsets: - axis_offsets = self.offsets[axis] - if "modifier" in axis_offsets and "offset" in axis_offsets: - for idx, rng in enumerate(axis_offsets["modifier"]["range"]): - if rng[0] < pos[axis_offsets["modifier"]["axis"]] < rng[1]: - pos[axis] -= axis_offsets["offset"][idx] - break - elif "offset" in axis_offsets: - pos[axis] -= axis_offsets["offset"] + pos = self.core.remove_offsets(pos) self.input.energy.set_number(self.dev.mo1_bragg.read(cached=True)["mo1_bragg"]["value"]) - h_acc, v_acc = sldi_gap_to_acc(pos["sldi_gapx"], pos["sldi_gapy"]) + h_acc, v_acc = self.core.sldi_gap_to_acc(pos["sldi_gapx"], pos["sldi_gapy"]) self.input.sldi_hacc.set_number(h_acc * 1e3) self.input.sldi_vacc.set_number(v_acc * 1e3) - self.input.cm_stripe.set_current_text(cm_trx_to_stripe(-pos["cm_trx"])) + self.input.cm_stripe.set_current_text(self.core.cm_trx_to_stripe(-pos["cm_trx"])) self.input.cm_pitch.set_number(pos["cm_rotx"]) if abs(pos["mo1_trx"]) > 5: mo1_mode = "Monochromatic" @@ -617,7 +572,7 @@ class DigitalTwin(BECWidget, QWidget): self.input.mo1_xtal.set_current_text( self.dev.mo1_bragg.read(cached=True)["mo1_bragg_crystal_current_xtal_string"]["value"] ) - fm_stripe = fm_trx_to_stripe(-pos["fm_trx"]) + fm_stripe = self.core.fm_trx_to_stripe(-pos["fm_trx"]) self.input.fm_stripe.set_current_text(fm_stripe) fm_rotx_real = 2 * pos["cm_rotx"] - pos["fm_rotx"] self.input.fm_rotx.set_number(fm_rotx_real) @@ -628,10 +583,10 @@ class DigitalTwin(BECWidget, QWidget): self.input.smpl.set_number(pos["ot_es1_trz"]) case ComboBox(): table = self.ask_table_selection(self.input.smpl.currentText()) - smpl = table_to_smpl_pos(table) + smpl = self.core.table_to_smpl_pos(table) self.input.smpl.set_current_text(table) - fm_focx, fm_focy = calc_beamsize( + fm_focx, fm_focy = self.core.calc_beamsize( h_acc, v_acc, fm_stripe, -fm_rotx_real * 1e-3, pos["fm_bnd_radius"] * 1e6, smpl ) if fm_focx < 0.08 and fm_focy < 0.08: @@ -690,19 +645,9 @@ class DigitalTwin(BECWidget, QWidget): Defaults to True """ - if self.offsets == {}: - # Load offsets - if not self.offset_file.exists(): - raise FileNotFoundError(f"Offset file not found: {self.offset_file}") - - with self.offset_file.open("r", encoding="utf-8") as f: - data = yaml.safe_load(f) - - if not isinstance(data, dict): - raise ValueError(f"Expected a YAML mapping, got {type(data).__name__}") - - self.offsets = data - + self.core.load_offsets() + if self.core.offsets != {}: + # Offsets were loaded if recalculate: self.calc_assistant(identifier="init") @@ -711,8 +656,7 @@ class DigitalTwin(BECWidget, QWidget): self.settings.offsets_status.setColor(get_accent_colors().success.name()) self.settings.show_offsets.enable_button(True) else: - # Unload offsets - self.offsets = {} + # Offsets were unloaded self.calc_assistant(identifier="init") self.settings.load_offsets.setText("Load") @@ -736,7 +680,7 @@ class DigitalTwin(BECWidget, QWidget): intro_label.setWordWrap(True) layout.addWidget(intro_label) - file = QLabel(str(self.offset_file)) + file = QLabel(str(self.core.offset_file)) file.setWordWrap(True) font = QFont() font.setItalic(True) @@ -753,7 +697,9 @@ class DigitalTwin(BECWidget, QWidget): def represent_sequence(self, tag, sequence, *_): return super().represent_sequence(tag, sequence, flow_style=True) - text_edit.setPlainText(yaml.dump(self.offsets, Dumper=InlineListDumper, sort_keys=False)) + text_edit.setPlainText( + yaml.dump(self.core.offsets, Dumper=InlineListDumper, sort_keys=False) + ) layout.addWidget(text_edit) buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Close) @@ -793,9 +739,9 @@ class DigitalTwin(BECWidget, QWidget): Updates the plots for the reality scene """ config = self.get_reality_config() - data = calc_sideview(config) + data = self.core.calc_sideview(config) self.sideview_plot.update_curves("reality", data=data) - surfaces = calc_surfaces(config) + surfaces = self.core.calc_surfaces(config) self.surface_plots.update_surfaces(scene="reality", data=surfaces) @SafeSlot() @@ -825,7 +771,7 @@ class DigitalTwin(BECWidget, QWidget): ) # Remove brackets from xtal name to conform with parameters xtal = cast(Literal["Si111", "Si311"], xtal) energy = self.input.energy.value() - self.input.mo1_eres.setValue(mo1_energy_resolution(xtal, energy)) + self.input.mo1_eres.setValue(self.core.mo1_energy_resolution(xtal, energy)) def calc_cm_reflectivity(self): """ @@ -834,9 +780,11 @@ class DigitalTwin(BECWidget, QWidget): cm_stripe = self.input.cm_stripe.currentText() cm_pitch = -self.input.cm_pitch.value() * 1e-3 energy = self.input.energy.value() - self.input.cm_refl.setValue(100 * cm_reflectivity(cm_stripe, cm_pitch, energy)) + self.input.cm_refl.setValue(100 * self.core.cm_reflectivity(cm_stripe, cm_pitch, energy)) self.input.cm_refl.setLabel(f"Reflectivity at \n{energy:.0f} eV") - self.input.cm_refl_harm.setValue(100 * cm_reflectivity(cm_stripe, cm_pitch, 3 * energy)) + self.input.cm_refl_harm.setValue( + 100 * self.core.cm_reflectivity(cm_stripe, cm_pitch, 3 * energy) + ) self.input.cm_refl_harm.setLabel(f"Reflectivity at \n{3*energy:.0f} eV") def calc_fm_reflectivity(self): @@ -850,9 +798,11 @@ class DigitalTwin(BECWidget, QWidget): else: fm_rotx = -self.input.fm_rotx_ideal.value() * 1e-3 energy = self.input.energy.value() - self.input.fm_refl.setValue(100 * fm_reflectivity(fm_stripe, fm_rotx, energy)) + self.input.fm_refl.setValue(100 * self.core.fm_reflectivity(fm_stripe, fm_rotx, energy)) self.input.fm_refl.setLabel(f"Reflectivity at \n{energy:.0f} eV") - self.input.fm_refl_harm.setValue(100 * fm_reflectivity(fm_stripe, fm_rotx, 3 * energy)) + self.input.fm_refl_harm.setValue( + 100 * self.core.fm_reflectivity(fm_stripe, fm_rotx, 3 * energy) + ) self.input.fm_refl_harm.setLabel(f"Reflectivity at \n{3*energy:.0f} eV") def calc_cm_fm_harm_suppr(self): @@ -872,33 +822,23 @@ class DigitalTwin(BECWidget, QWidget): Updates the sideview plot based on the assistant values """ config = self.get_assistant_config(apply_offset=True) - data = calc_sideview(config) + data = self.core.calc_sideview(config) self.sideview_plot.update_curves("assistant", data) def calc_assistant_surfaces(self): """ Updates the surface plot based on the assistant values """ - surfaces = calc_surfaces(self.get_assistant_config()) + surfaces = self.core.calc_surfaces(self.get_assistant_config()) self.surface_plots.update_surfaces(scene="assistant", data=surfaces) def calc_positions(self): """ Calculates the positions for the axes based on the assistant values """ - out = calc_positions(self.beamline, self.get_assistant_config()) - - # Apply offsets - for axis, axis_data in out.items(): - if axis in self.offsets: - axis_offsets = self.offsets[axis] - if "modifier" in axis_offsets and "offset" in axis_offsets: - for idx, rng in enumerate(axis_offsets["modifier"]["range"]): - if rng[0] < out[axis_offsets["modifier"]["axis"]]["value"] < rng[1]: - axis_data["value"] += axis_offsets["offset"][idx] - break - elif "offset" in axis_offsets: - axis_data["value"] += axis_offsets["offset"] + config = self.get_assistant_config() + out = self.core.calc_positions(self.beamline, config) + out = self.core.apply_offsets(out, nested_config=True) self.mover.sldi_gapx.set_target(out["sldi_gapx"]["value"]) self.mover.sldi_gapy.set_target(out["sldi_gapy"]["value"]) @@ -952,7 +892,7 @@ class DigitalTwin(BECWidget, QWidget): cm_pitch = -self.dev.cm_rotx.read(cached=True)["cm_rotx"]["value"] * 1e-3 mo1_mode = cast(Literal["Monochromatic", "Pinkbeam"], self.input.mo1_mode.currentText()) energy = self.input.energy.value() - theta, _ = mo1_bragg_angle(mo1_mode, d_spacing, energy, cm_pitch) + theta, _ = self.core.mo1_bragg_angle(mo1_mode, d_spacing, energy, cm_pitch) self.bragg_angle = theta self.input.mo1_bragg_angle.setValue(theta / np.pi * 180) @@ -983,12 +923,12 @@ class DigitalTwin(BECWidget, QWidget): smpl = self.input.smpl.value() case ComboBox(): table = self.input.smpl.currentText() - smpl = table_to_smpl_pos(table) + smpl = self.core.table_to_smpl_pos(table) sldi_hacc = self.input.sldi_hacc.value() * 1e-3 sldi_vacc = self.input.sldi_vacc.value() * 1e-3 fm_focx = self.input.fm_focx.value() fm_focy = self.input.fm_focy.value() - fm_rotx, qy = fm_ideal_pitch( + fm_rotx, qy = self.core.fm_ideal_pitch( fm_focus, fm_stripe, smpl, sldi_hacc, sldi_vacc, fm_focx, fm_focy ) self.qy = qy @@ -1000,7 +940,7 @@ class DigitalTwin(BECWidget, QWidget): """ cm_stripe = cast(Literal["Si", "Pt", "Rh"], self.input.cm_stripe.currentText()) energy = self.input.energy.value() - self.input.cm_pitch_critical.setValue(-cm_critical_angle(cm_stripe, energy) * 1e3) + self.input.cm_pitch_critical.setValue(-self.core.cm_critical_angle(cm_stripe, energy) * 1e3) if __name__ == "__main__": diff --git a/debye_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py b/debye_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py index 8bebec6..5c4cd1b 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py @@ -7,7 +7,7 @@ from typing import Union # pylint: disable=E0611 from qtpy.QtWidgets import QHBoxLayout, QLabel, QPushButton, QVBoxLayout, QWidget -from ..types import BeamlineId +from .....bec_ipython_client.plugins.digital_twin_core.types import BeamlineId from ..widgets.qt_widgets import ( Button, ComboBox, diff --git a/debye_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py b/debye_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py index 4b43a3d..148d0ef 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py @@ -7,7 +7,7 @@ from typing import Literal # pylint: disable=E0611 from qtpy.QtWidgets import QVBoxLayout, QWidget -from ..types import BeamlineId +from .....bec_ipython_client.plugins.digital_twin_core.types import BeamlineId from ..widgets.move_widget import AbsorberWidget, MoveWidget from ..widgets.qt_widgets import Group diff --git a/debye_bec/bec_widgets/widgets/digital_twin/panels/plots.py b/debye_bec/bec_widgets/widgets/digital_twin/panels/plots.py index b7d7131..0d8f383 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/panels/plots.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/panels/plots.py @@ -15,13 +15,11 @@ from qtpy.QtGui import QBrush, QColor # pylint: disable=E0611 from qtpy.QtWidgets import QApplication, QGraphicsRectItem, QHBoxLayout, QVBoxLayout, QWidget -from ..calculations.calc_varia import ( - mirror_surface_geometries, - mo_surface_geometries, - pipe_geometries, - wall_geometries, +from .....bec_ipython_client.plugins.digital_twin_core.types import ( + BeamlineId, + DataDict, + SurfaceDict, ) -from ..types import BeamlineId, DataDict, SurfaceDict from ..widgets.qt_widgets import Group logger = bec_logger.logger @@ -30,9 +28,10 @@ logger = bec_logger.logger class SurfacePlots(QWidget): """Plot widget with two curves and legend.""" - def __init__(self, beamline: BeamlineId, parent=None): + def __init__(self, beamline: BeamlineId, core, parent=None): super().__init__(parent=parent) self.beamline = beamline + self.core = core self._layout = QHBoxLayout(self) self._layout.setContentsMargins(4, 4, 4, 4) self._layout.setSpacing(6) @@ -169,15 +168,15 @@ class SurfacePlots(QWidget): for name, plot in self.plots.items(): if name == "cm": - plot_surface(plot["widget"], mirror_surface_geometries("cm")) + plot_surface(plot["widget"], self.core.mirror_surface_geometries("cm")) elif name == "mo1_1": - plot_surface(plot["widget"], mo_surface_geometries("mo1", 0)) + plot_surface(plot["widget"], self.core.mo_surface_geometries("mo1", 0)) elif name == "mo1_2": - plot_surface(plot["widget"], mo_surface_geometries("mo1", 1)) + plot_surface(plot["widget"], self.core.mo_surface_geometries("mo1", 1)) elif name == "fm": if self.beamline == "x01da": - plot_surface(plot["widget"], mirror_surface_geometries("fm_flat")) - plot_surface(plot["widget"], mirror_surface_geometries("fm_toroid")) + plot_surface(plot["widget"], self.core.mirror_surface_geometries("fm_flat")) + plot_surface(plot["widget"], self.core.mirror_surface_geometries("fm_toroid")) else: raise ValueError(f"Plot {name} not found!") for name, plot in self.plots.items(): @@ -203,8 +202,9 @@ class SurfacePlots(QWidget): class SideviewPlot(QWidget): """Plot widget with two curves and legend.""" - def __init__(self, parent=None): + def __init__(self, core, parent=None): super().__init__(parent=parent) + self.core = core self._layout = QVBoxLayout(self) self._layout.setContentsMargins(4, 4, 4, 4) self._layout.setSpacing(0) @@ -303,7 +303,7 @@ class SideviewPlot(QWidget): def plot_vacuum_pipes(self): """Plot vacuum pipes""" - pipes = pipe_geometries() + pipes = self.core.pipe_geometries() for pipe in pipes: self.pipes.append( self.plot_widget.plot( @@ -313,7 +313,7 @@ class SideviewPlot(QWidget): def plot_walls(self): """Plot walls""" - walls = wall_geometries() + walls = self.core.wall_geometries() for wall in walls: rect = QGraphicsRectItem(wall[0], wall[1], wall[2], wall[3]) rect.setBrush(QBrush(QColor(*self.color_impenetrable))) diff --git a/debye_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py b/debye_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py index aa57582..6df4065 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py @@ -15,9 +15,10 @@ from qtpy.QtCore import QObject, QPropertyAnimation, Qt, QThread from qtpy.QtGui import QTransform from qtpy.QtWidgets import QApplication, QHBoxLayout, QLabel, QPushButton, QWidget +from .....bec_ipython_client.plugins.digital_twin_core.types import BeamlineId + # pylint: disable=E0402 from .....devices.absorber import STATUS as ABS_STATUS -from ..types import BeamlineId logger = bec_logger.logger @@ -292,15 +293,13 @@ class MotionWorker(QObject): Args: surveyed_axes (list): List of dictionaries of devices """ + logger.info(f"Move axis {self.motor} to target {self._target}, move_relative={relative}") try: if alias: self.motor = alias if abs_closed: if self.dev.abs.status.get() == ABS_STATUS.OPEN: status = self.dev.abs.close() - # TODO Set timeout to 0.001 and check if it actually raises - # (it should not start motion). - # Check of behavior of digital twin afterwards. status.wait(timeout=5) if surveyed_axes is not None: for surv_ax in surveyed_axes: @@ -336,7 +335,8 @@ class MotionWorker(QObject): self.finished.emit() break self.finished.emit() - except: + except Exception as e: + logger.error(f"Error during movement of {self.motor}: {e}") self.error.emit() self.finished.emit() diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calculations/__init__.py b/debye_bec/bec_widgets/widgets/scheduler/__init__.py similarity index 100% rename from debye_bec/bec_widgets/widgets/digital_twin/calculations/__init__.py rename to debye_bec/bec_widgets/widgets/scheduler/__init__.py diff --git a/debye_bec/bec_widgets/widgets/scheduler/endpoints.py b/debye_bec/bec_widgets/widgets/scheduler/endpoints.py new file mode 100644 index 0000000..06fb111 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/endpoints.py @@ -0,0 +1,61 @@ +""" +Redis endpoint(s) for the schedule widget plugin. + +Design notes (why this looks the way it does) +----------------------------------------------- +BEC never touches Redis with raw keys/commands. Every channel is described +by an `EndpointInfo` (endpoint string + message type + allowed operations, +see `bec_lib.endpoints`), and `RedisConnector` enforces both: calling an +operation that isn't in the endpoint's `MessageOp` raises +`IncompatibleRedisOperation`, and passing a message that isn't an instance +of the endpoint's declared `message_type` raises +`IncompatibleMessageForEndpoint`. Plain string topics still work but are +deprecated. So a "conforming" custom endpoint means building a real +`EndpointInfo`, exactly like `bec_lib.endpoints.MessageEndpoints` does +internally. + +Namespace: we use `EndpointType.USER` ("user/...”), the same prefix BEC's +own `MessageEndpoints.scan_queue_schedule()` uses for user-writable, +persisted data (as opposed to `internal/`, `public/`, etc.). + +Message type: we deliberately do NOT declare our own `BECMessage` +subclass. BEC's msgpack codec (`bec_lib.codecs.BECMessageEncoder.decode`) +resolves an incoming message's class by name via +`getattr(bec_lib.messages, type_name)` - i.e. only classes physically +defined inside `bec_lib.messages` are resolvable this way. A message class +declared in plugin code would fail to deserialize unless you monkey-patch +it into that module, which is exactly the kind of side-door this plugin +is meant to avoid. Instead we reuse `bec_lib.messages.VariableMessage`, +a first-class, exported message type built for carrying an arbitrary +(msgpack-serializable) payload under `.value`. +`bec_lib.script_executor.upload_script` persists a script's text in Redis +the same way, for the same reason - it's the established BEC pattern for +"a plugin needs to stash its own structured data in Redis". +""" + +from __future__ import annotations + +from bec_lib.endpoints import EndpointInfo, EndpointType, MessageOp +from bec_lib.messages import VariableMessage + + +def schedule(schedule_name: str) -> EndpointInfo: + """ + Endpoint for one named widget schedule (an ordered list of commands + plus their execution status). The whole schedule is stored as a single + `VariableMessage` document that gets overwritten and republished on + every change (`MessageOp.SET_PUBLISH`), so any other subscriber - + another instance of this widget, a monitoring script, ... - stays in + sync live. + + Args: + schedule_name: a stable, user-chosen name for the schedule. Unlike + a widget's `gui_id` (regenerated every time the widget is + constructed), this name is what lets a widget that was closed + and reopened find its own previously persisted schedule again. + """ + return EndpointInfo( + endpoint=f"{EndpointType.USER.value}/schedule_widget/schedule/{schedule_name}", + message_type=VariableMessage, + message_op=MessageOp.SET_PUBLISH, + ) diff --git a/debye_bec/bec_widgets/widgets/scheduler/enums.py b/debye_bec/bec_widgets/widgets/scheduler/enums.py new file mode 100644 index 0000000..71bc706 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/enums.py @@ -0,0 +1,11 @@ +from __future__ import annotations + +from enum import Enum + + +class ScheduleItemStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + ABORTED = "aborted" diff --git a/debye_bec/bec_widgets/widgets/scheduler/guard.py b/debye_bec/bec_widgets/widgets/scheduler/guard.py new file mode 100644 index 0000000..2f0dd11 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/guard.py @@ -0,0 +1,188 @@ +""" +A small, self-contained utility for "if signal X drops below A, pause; once +it's back above B, resume" (hysteresis) behavior - e.g. auto-pausing scans +while the beam current is too low. + +This module knows nothing about schedules, scans, or the rest of this +plugin. It only monitors one BEC device's live readback value +(`MessageEndpoints.device_readback`, the same endpoint the device server +publishes for every monitored device) and emits Qt signals when the value +crosses one of the two configured thresholds. `schedule_widget.py` is the +only place that connects those signals to schedule-specific behavior +(aborting/deferring a scan item) - see its module docstring for how. +""" + +from __future__ import annotations + +import threading +import time + +from bec_lib.endpoints import MessageEndpoints +from bec_lib.logger import bec_logger +from pydantic import BaseModel +from qtpy.QtCore import QObject, Signal + +logger = bec_logger.logger + +MIN_UPDATE_INTERVAL = 1 + + +class GuardSettings(BaseModel): + """Persisted configuration for one `SignalGuard`.""" + + enabled: bool = False + device_name: str | None = None + pause_below: float | None = None + resume_above: float | None = None + + +class SignalGuard(QObject): + """ + Monitors one device's readback value and reports hysteresis-based + pause/resume crossings. + + - `paused` fires the first time the value drops below `pause_below` + (not again while it stays low). + - `resumed` fires the first time it then climbs back above + `resume_above` (not again while it stays high). + A value sitting between the two thresholds never re-triggers either + signal - that gap is the point of using two thresholds instead of one, + so a value oscillating right at a single cutoff wouldn't cause rapid + pause/resume flapping. + + Callbacks from `RedisConnector.register()` run on a background + (Redis-listener) thread, not the Qt GUI thread; `paused`/`resumed` are + Qt signals, so connecting to them with the default (auto) connection + type safely marshals delivery onto whatever thread the receiver lives + on - no extra locking needed on the receiving end. + """ + + value_update = Signal(float) + paused = Signal(float) + resumed = Signal(float) + + def __init__(self, connector, parent=None, dev=None): + super().__init__(parent) + self.dev = dev + self._connector = connector + self._lock = threading.Lock() + self._clear_event = threading.Event() + self._clear_event.set() # not blocking until configured/proven otherwise + + self.enabled = False + self.device_name: str | None = None + self.pause_below: float | None = None + self.resume_above: float | None = None + self.current_value: float | None = None + self.units: str = "" + self.prec: int = 3 + + self._subscribed_endpoint = None + + self.last_val_update = time.time() + + def configure(self, settings: GuardSettings): + """(Re)configure and (re)subscribe. Safe to call repeatedly, e.g. after editing settings.""" + self._unsubscribe() + + self.enabled = settings.enabled + self.device_name = settings.device_name + self.pause_below = settings.pause_below + self.resume_above = settings.resume_above + self.current_value = None + self._clear_event.set() + + if ( + self.enabled + and self.device_name + and self.pause_below is not None + and self.resume_above is not None + ): + self._subscribed_endpoint = MessageEndpoints.device_readback(self.device_name) + self._connector.register(topics=self._subscribed_endpoint, cb=self._on_readback) + description = self.dev[self.device_name].describe()[self.device_name] + self.units = description["units"] + self.prec = description["precision"] + + def _unsubscribe(self): + if self._subscribed_endpoint is not None: + try: + self._connector.unregister(topics=self._subscribed_endpoint, cb=self._on_readback) + except Exception: # pylint: disable=broad-except + logger.exception( + f"Failed to unsubscribe SignalGuard from {self._subscribed_endpoint}" + ) + self._subscribed_endpoint = None + + def _on_readback(self, msg): + device_msg = getattr(msg, "value", None) + if device_msg is None: + return + value = _extract_value(device_msg, self.device_name) + if value is None: + return + + crossed_pause = crossed_resume = False + with self._lock: + self.current_value = value + if self._clear_event.is_set(): + if value < self.pause_below: + self._clear_event.clear() + crossed_pause = True + elif value > self.resume_above: + self._clear_event.set() + crossed_resume = True + + if crossed_pause: + logger.info( + f"SignalGuard: {self.device_name} dropped to {value} (below {self.pause_below}) - pausing." + ) + self.paused.emit(value) + elif crossed_resume: + logger.info( + f"SignalGuard: {self.device_name} recovered to {value} (above {self.resume_above}) - resuming." + ) + self.resumed.emit(value) + + current_time = time.time() + if current_time - self.last_val_update > MIN_UPDATE_INTERVAL: + self.last_val_update = current_time + self.value_update.emit(value) + + def is_clear(self) -> bool: + """True if not currently blocking (disabled, or value above the resume threshold).""" + return (not self.enabled) or self._clear_event.is_set() + + def wait_until_clear(self, should_abort, poll_interval: float = 0.5) -> bool: + """ + Blocks the calling (non-GUI) thread until `is_clear()` becomes + True, checking `should_abort()` between polls so an operator abort + can still interrupt the wait. Returns True if it cleared, False if + `should_abort()` returned True first. + """ + if self.is_clear(): + return True + while not should_abort(): + if self._clear_event.wait(timeout=poll_interval): + return True + return False + + def cleanup(self): + self._unsubscribe() + + +def _extract_value(device_msg, device_name: str | None): + """ + `DeviceMessage.signals` is keyed by signal name, e.g. {"beam_current": + {"value": ..., "timestamp": ...}, ...}. For a simple scalar device the + primary signal is usually named after the device itself; fall back to + the first signal found if not (e.g. a differently-named primary signal + on a compound device). + """ + signals = getattr(device_msg, "signals", None) or {} + if device_name in signals: + return signals[device_name].get("value") + for signal in signals.values(): + if "value" in signal: + return signal["value"] + return None diff --git a/debye_bec/bec_widgets/widgets/scheduler/guard_dialog.py b/debye_bec/bec_widgets/widgets/scheduler/guard_dialog.py new file mode 100644 index 0000000..3a7654f --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/guard_dialog.py @@ -0,0 +1,94 @@ +"""Small settings dialog for `guard.GuardSettings` - kept separate from the guard logic itself.""" + +from __future__ import annotations + +from qtpy.QtWidgets import ( + QCheckBox, + QComboBox, + QDialog, + QDialogButtonBox, + QDoubleSpinBox, + QFormLayout, + QLabel, + QMessageBox, + QVBoxLayout, +) + +from .guard import GuardSettings + +_DSPIN_RANGE = (-1e12, 1e12) + + +class GuardSettingsDialog(QDialog): + """Configure (or disable) the auto-pause/resume guard for scan items.""" + + def __init__(self, settings: GuardSettings, device_names: list[str], parent=None): + super().__init__(parent) + self.setWindowTitle("Auto-pause on signal") + + layout = QVBoxLayout(self) + info_label = QLabel( + "If enabled, the currently running scan is aborted as soon as the chosen " + "signal drops below the pause value, and automatically restarted once it " + "rises back above the resume value. Only scan items are affected - device " + "moves and custom/RPC commands are never interrupted by this." + ) + info_label.setWordWrap(True) + layout.addWidget(info_label) + + form = QFormLayout() + layout.addLayout(form) + + self.enabled_check = QCheckBox("Enable auto-pause") + self.enabled_check.setChecked(settings.enabled) + form.addRow("", self.enabled_check) + + self.device_combo = QComboBox() + self.device_combo.addItems(device_names) + if settings.device_name: + idx = self.device_combo.findText(settings.device_name) + if idx >= 0: + self.device_combo.setCurrentIndex(idx) + form.addRow("Signal (device)", self.device_combo) + + self.pause_below_spin = QDoubleSpinBox() + self.pause_below_spin.setDecimals(4) + self.pause_below_spin.setRange(*_DSPIN_RANGE) + if settings.pause_below is not None: + self.pause_below_spin.setValue(settings.pause_below) + form.addRow("Pause below", self.pause_below_spin) + + self.resume_above_spin = QDoubleSpinBox() + self.resume_above_spin.setDecimals(4) + self.resume_above_spin.setRange(*_DSPIN_RANGE) + if settings.resume_above is not None: + self.resume_above_spin.setValue(settings.resume_above) + form.addRow("Resume above", self.resume_above_spin) + + buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) + buttons.accepted.connect(self._on_accept) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + + def _on_accept(self): + if self.enabled_check.isChecked(): + if not self.device_combo.currentText(): + QMessageBox.warning(self, "Missing input", "Select a device to monitor.") + return + if self.resume_above_spin.value() <= self.pause_below_spin.value(): + QMessageBox.warning( + self, + "Invalid thresholds", + "'Resume above' must be greater than 'Pause below' (hysteresis gap).", + ) + return + self.accept() + + def result_settings(self) -> GuardSettings: + """Valid after `exec_()` returns `QDialog.Accepted`.""" + return GuardSettings( + enabled=self.enabled_check.isChecked(), + device_name=self.device_combo.currentText() or None, + pause_below=self.pause_below_spin.value(), + resume_above=self.resume_above_spin.value(), + ) diff --git a/debye_bec/bec_widgets/widgets/scheduler/item_dialog.py b/debye_bec/bec_widgets/widgets/scheduler/item_dialog.py new file mode 100644 index 0000000..02066bf --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/item_dialog.py @@ -0,0 +1,623 @@ +""" +The dialog behind the schedule widget's "Add..."/"Edit..." buttons. + +Rather than asking the operator to remember and type +`scans.xas_simple_scan(12000, 14000, 2, 10)`-style commands, this presents: + +- a "Scan" tab: BEC's own `bec_widgets` `ScanControl` widget, embedded + as-is - scan selection, its live-generated per-scan argument form, + docs tooltips, metadata, "recall last scan parameters", all of it. + Reusing it instead of a plugin-owned reimplementation means this stays + in sync with BEC's scan capabilities for free, and looks/behaves exactly + like the scan controls an operator already knows from elsewhere in the + GUI. `ScanControl.button_run_scan` ("Start") is hidden here: this dialog + only ever wants the configured scan name/args/kwargs, never an + immediate submission - see `_collect_scan_result`. +- a "Move" tab: pick a device and a target value/relative flag; +- a "Digital Twin" tab: the beamline-alignment `DigitalTwin` widget, + embedded the same way as ScanControl. Instead of submitting anything + itself, OK captures a *snapshot* of `DigitalTwin.get_assistant_config()` + and stores it - execution later calls `move_all_axes(...)` with that + frozen config, so editing the (possibly separately open) Digital Twin + widget afterwards never affects an already-added schedule item, exactly + like a Scan item's captured args/kwargs aren't affected by reopening + ScanControl elsewhere. See `_collect_digital_twin_result`. +- an "Other" tab: a free-text field for anything else (including RPC + calls to other widgets), plus a couple of beamline-specific quick-fill + forms (ionization chamber gas mix, reference foil). + +Whichever tab is used, the dialog's only output is the same kind of plain +command string the executor already knows how to run - this dialog adds a +friendlier way to *build* that string, it doesn't change what happens with +it afterwards. `kind`/`form_state` are carried along purely so "Edit..." +can reopen the dialog pre-filled instead of asking the user to start over +- except for the Digital Twin tab, where "Edit..." currently falls back to +showing the generated command as read/write text on the "Other" tab rather +than reloading the captured config back into DigitalTwin's input fields; +see the note on `_apply_initial`. + +The Digital Twin item reuses `kind="move"` (not a new kind): it submits +through `scans.mv(...)`, exactly like the plain Move tab, so it should be +treated the same way everywhere else in the plugin that branches on kind - +not guard-protected, and counted under the "Movements" notification +toggle. A `form_state["source"] = "digital_twin"` marker is only used +locally, by this dialog, to tell the two apart when reopening for Edit. +""" + +from __future__ import annotations + +import math + +from bec_lib.logger import bec_logger +from bec_widgets.widgets.control.device_input.device_combobox.device_combobox import ( + BECDeviceFilter, + DeviceComboBox, +) +from bec_widgets.widgets.control.scan_control.scan_control import ScanControl, ScanParameterConfig + +# pylint: disable=E0611 +from qtpy.QtWidgets import ( + QCheckBox, + QComboBox, + QDialog, + QDialogButtonBox, + QDoubleSpinBox, + QFormLayout, + QGroupBox, + QHBoxLayout, + QLabel, + QLineEdit, + QMessageBox, + QScrollArea, + QTabWidget, + QVBoxLayout, + QWidget, +) + +from .qt_widgets import MyButton + +logger = bec_logger.logger + +# Tab indices, named instead of magic numbers now that there are four - +# see _collect_result()/_apply_initial(). +_TAB_SCAN = 0 +_TAB_MOVE = 1 +_TAB_DIGITAL_TWIN = 2 +_TAB_OTHER = 3 + + +class ScheduleItemDialog(QDialog): + """Add or edit one schedule item, via ScanControl, a move form, Digital Twin, or free text.""" + + def __init__( + self, + scans, + dev, + parent=None, + initial: dict | None = None, + client=None, + beamline: str | None = None, + ): + super().__init__(parent) + self.setWindowTitle("Schedule item") + self.setMinimumSize(520, 480) + + self._scans = scans + self._dev = dev + self._client = client + + self.beamline = beamline + if self.beamline in ["x01da", "x10da"]: + logger.info(f"Loading bl-specific modules for beamline {self.beamline}") + from ....bec_ipython_client.plugins.digital_twin_core.digital_twin_core import DigitalTwinCore + from ..digital_twin.digital_twin import DigitalTwin + from ..edge_selector import EdgeSelector + from ..scan_control_xas.scan_control_xas import ScanControlXAS + + self.DigitalTwinCore = DigitalTwinCore + self.DigitalTwin = DigitalTwin + self.EdgeSelector = EdgeSelector + self.ScanControlXAS = ScanControlXAS + + layout = QVBoxLayout(self) + self.tabs = QTabWidget() + layout.addWidget(self.tabs) + + self._build_scan_tab() + self._build_move_tab() + if self.beamline in ["x01da", "x10da"]: + self._build_digital_twin_tab() + self._build_custom_tab() + + buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) + buttons.button(QDialogButtonBox.Ok).setText("Add") + buttons.setStyleSheet("QPushButton {qproperty-icon: none;}") + buttons.accepted.connect(self._on_accept) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + + # The embedded DigitalTwin instance runs a 1s polling timer and a + # bec_dispatcher subscription for its whole lifetime - stop both + # when this dialog closes, however it closes (OK, Cancel, or the + # window's own close button), not just on accept. + self.finished.connect(self._cleanup_digital_twin) + + self._apply_initial(initial or {}) + + # ------------------------------------------------------------------ # + # Scan tab - embeds BEC's own ScanControl widget + # ------------------------------------------------------------------ # + def _build_scan_tab(self): + tab = QWidget() + layout = QVBoxLayout(tab) + + # client=None resolves to the same process-wide BEC client + # (bec_dispatcher.client) our own widget uses - no second Redis + # connection is opened. + if self.beamline in ["x01da", "x10da"]: + self.scan_control = self.ScanControlXAS(parent=tab, client=self._client) + else: + self.scan_control = ScanControl(parent=tab, client=self._client) + self.scan_control.button_run_scan.hide() + layout.addWidget(self.scan_control) + + self.tabs.addTab(tab, "Scan") + + def _collect_scan_result(self) -> dict: + # Same call ScanControl.run_scan() makes before actually + # submitting, to resolve a typed-but-unconfirmed scan name. + self.scan_control.validate_scan_selection() + scan_name = self.scan_control.current_scan + if not scan_name: + raise ValueError("No scan selected.") + + # bec_object=False: plain, repr-able values (e.g. a device name + # string rather than the live DeviceBase instance) - needed since + # the result has to survive a round-trip through Redis as text and + # be re-evaluated later, not just used in-process immediately. + args, kwargs = self.scan_control.get_scan_parameters(bec_object=False) + command = _format_scan_call(scan_name, args, kwargs) + return { + "command": command, + "kind": "scan", + "form_state": {"scan_name": scan_name, "args": args, "kwargs": kwargs}, + } + + def _prefill_scan_tab(self, scan_name: str, args: list, kwargs: dict): + # ScanControl restores parameters for a scan from its own config + # cache (see `ScanControl.restore_scan_parameters`); pre-loading + # that cache before switching to the scan reuses that mechanism + # instead of poking at its internal argument widgets directly. + self.scan_control.config.scans[scan_name] = ScanParameterConfig( + name=scan_name, args=args, kwargs=kwargs + ) + self.scan_control.current_scan = scan_name + self.scan_control.restore_scan_parameters(scan_name) + + # ------------------------------------------------------------------ # + # Move tab + # ------------------------------------------------------------------ # + def _build_move_tab(self): + tab = QWidget() + form = QFormLayout(tab) + form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint) + + self.move_device_combo = DeviceComboBox(self, device_filter=[BECDeviceFilter.POSITIONER]) + + form.addRow("Device", self.move_device_combo) + + self.move_value_spin = QDoubleSpinBox() + form.addRow("Value", self.move_value_spin) + + self.move_relative_check = QCheckBox("") + form.addRow("Relative move", self.move_relative_check) + + self.move_device_combo.currentIndexChanged.connect(self._adjust_spinbox) + + self.tabs.addTab(tab, "Move") + + def _adjust_spinbox(self, _) -> None: + prec = self._dev[self.move_device_combo.currentText()].precision + units = self._dev[self.move_device_combo.currentText()].egu() + ll = self._dev[self.move_device_combo.currentText()].low_limit + hl = self._dev[self.move_device_combo.currentText()].high_limit + + self.move_value_spin.setDecimals(prec) + self.move_value_spin.setSuffix(f" {units}") + if (hl - ll) > 0: + self.move_value_spin.setMinimum(ll) + self.move_value_spin.setMaximum(hl) + self.move_value_spin.setSingleStep(10 ** round(math.log10((hl - ll) / 100))) + else: + self.move_value_spin.setMinimum(1e6) + self.move_value_spin.setMaximum(-1e6) + self.move_value_spin.setSingleStep(1) + + def _collect_move_result(self) -> dict: + device_name = self.move_device_combo.currentText() + if not device_name: + raise ValueError("No movable device available/selected.") + value = self.move_value_spin.value() + relative = self.move_relative_check.isChecked() + return { + "command": f"dev.{device_name}.move({value!r}, relative={relative!r})", + "kind": "move", + "form_state": {"device_name": device_name, "value": value, "relative": relative}, + } + + # ------------------------------------------------------------------ # + # Digital Twin tab - embeds the beamline-alignment DigitalTwin widget + # ------------------------------------------------------------------ # + def _build_digital_twin_tab(self): + tab = QWidget() + outer = QVBoxLayout(tab) + + # DigitalTwin defaults to a large fixed size (it's normally its + # own top-level window) - wrap it in a scroll area so this dialog + # doesn't have to grow to match it. + scroll = QScrollArea() + scroll.setWidgetResizable(True) + self.digital_twin = self.DigitalTwin(parent=scroll, client=self._client) + # Hide move and abs open buttons + for mover in self.digital_twin.mover.mover_widgets: + mover.btn_action.hide() + self.digital_twin.mover.abs.btn_action.hide() + + scroll.setWidget(self.digital_twin) + outer.addWidget(scroll) + + hint = QLabel( + "Configure the beamline alignment above, then confirm with Add below - the " + "computed motor targets are captured now and moved together (in one combined " + "move) when this schedule item runs, not immediately." + ) + hint.setWordWrap(True) + hint.setStyleSheet("color: gray;") + outer.addWidget(hint) + + self.tabs.addTab(tab, "Digital Twin") + + def _collect_digital_twin_result(self) -> dict: + config = self.digital_twin.get_assistant_config() + # beamline = self.digital_twin.beamline + + # # Init the class when the scheduler is opened + # digital_twin = DigitalTwinCore() + # # The command below would then execute the movement + # digital_twin.move_with_config(config) + + cmd = f"digital_twin.move_with_config({config})" + + return {"command": f"{cmd}", "kind": "custom", "form_state": {"text": cmd}} + + def _cleanup_digital_twin(self, *_): + digital_twin = getattr(self, "digital_twin", None) + if digital_twin is None: + return + try: + digital_twin._timer.stop() # pylint: disable=protected-access + except Exception: # pylint: disable=broad-except + logger.exception("Failed to stop the Digital Twin's reality-update timer.") + try: + digital_twin.cleanup() + except Exception: # pylint: disable=broad-except + logger.exception("Failed to clean up the embedded Digital Twin widget.") + + # ------------------------------------------------------------------ # + # Custom tab + # ------------------------------------------------------------------ # + def _build_custom_tab(self): + tab = QWidget() + layout = QVBoxLayout(tab) + layout.addWidget( + QLabel( + "Free-form command, evaluated against `scans` and `dev` - use this for " + "anything the other tabs don't cover, e.g. an RPC call to another widget." + ) + ) + self.custom_edit = QLineEdit() + # TODO Change to a different placeholder text + self.custom_edit.setPlaceholderText("scans.xas_simple_scan(12000, 14000, 2, 10)") + layout.addWidget(self.custom_edit) + + if self.beamline in ["x01da", "x10da"]: + abs_form = self._create_abs_form() + if abs_form is not None: + layout.addWidget(abs_form) + + ic_form = self._create_ionization_chamber_form() + if ic_form is not None: + layout.addWidget(ic_form) + + reffoil_form = self._create_reffoil_form() + if reffoil_form is not None: + layout.addWidget(reffoil_form) + + auto_gain_form = self._create_auto_gain_form() + if auto_gain_form is not None: + layout.addWidget(auto_gain_form) + + layout.addStretch(1) + self.tabs.addTab(tab, "Other") + + def _create_abs_form(self): + if "abs" in self._dev: + abs_group = QGroupBox("Frontend Absorber") + layout = QVBoxLayout(abs_group) + form = QFormLayout() + layout.addLayout(form) + form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint) + self.abs_selector = QComboBox() + self.abs_selector.addItems(["Open", "Force open", "Close"]) + form.addRow("Action", self.abs_selector) + + button_layout = QHBoxLayout() + generate_cmd = MyButton("Generate command", "default") + button_layout.addWidget(generate_cmd) + button_layout.addStretch(1) + layout.addLayout(button_layout) + + generate_cmd.clicked.connect(self._generate_abs_command) + + return abs_group + return None + + def _generate_abs_command(self): + match self.abs_selector.currentText(): + case "Open": + suffix = "open()" + case "Force open": + suffix = "open(force=True)" + case "Close": + suffix = "close()" + cmd = f"dev.abs.{suffix}" + self.custom_edit.setText(cmd) + + def _create_ionization_chamber_form(self): + if all(key in self._dev for key in ("ic0", "ic1", "ic2")): + ic_group = QGroupBox("Ionization chamber filling") + layout = QVBoxLayout(ic_group) + form = QFormLayout() + layout.addLayout(form) + form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint) + self.ic_selector = QComboBox() + self.ic_selector.addItems(["IC0", "IC1", "IC2"]) + gases = ["He", "N2", "Ar", "Kr"] + self.gas1 = QComboBox() + self.gas2 = QComboBox() + self.gas1.addItems(gases) + self.gas2.addItems(gases) + self.conc1 = QDoubleSpinBox() + self.conc2 = QDoubleSpinBox() + for conc in [self.conc1, self.conc2]: + conc.setDecimals(0) + conc.setSuffix(" %") + conc.setMinimum(0) + conc.setMaximum(100) + conc.setSingleStep(1) + self.pressure = QDoubleSpinBox() + self.pressure.setDecimals(3) + self.pressure.setSuffix(" bar abs") + self.pressure.setMinimum(1) + self.pressure.setMaximum(3) + self.pressure.setSingleStep(0.1) + + form.addRow("Ionization chamber", self.ic_selector) + form.addRow("Gas 1", self.gas1) + form.addRow("Concentration 1", self.conc1) + form.addRow("Gas 2", self.gas2) + form.addRow("Concentration 2", self.conc2) + form.addRow("Pressure", self.pressure) + + button_layout = QHBoxLayout() + generate_cmd = MyButton("Generate command", "default") + button_layout.addWidget(generate_cmd) + button_layout.addStretch(1) + layout.addLayout(button_layout) + + self.conc1.valueChanged.connect(self._equalize_ic_conc) + self.conc2.valueChanged.connect(self._equalize_ic_conc) + generate_cmd.clicked.connect(self._generate_ic_command) + + return ic_group + return None + + def _equalize_ic_conc(self, new_val): + if self.conc1.value() == new_val: # conc1 was changed + self.conc2.setValue(100 - new_val) + else: + self.conc1.setValue(100 - new_val) + + def _generate_ic_command(self): + if self.conc1.value() + self.conc2.value() != 100: + return + match self.ic_selector.currentText(): + case "IC0": + ic = "ic0" + case "IC1": + ic = "ic1" + case "IC2": + ic = "ic2" + cmd = ( + f"dev.{ic}.fill(" + + f"gas1='{self.gas1.currentText()}', conc1={self.conc1.value()}, " + + f"gas2='{self.gas2.currentText()}', conc2={self.conc2.value()}, " + + f"pressure={self.pressure.value()}, wait=True)" + ) + self.custom_edit.setText(cmd) + + def _create_reffoil_form(self): + if "reffoilchanger" in self._dev: + reffoil_group = QGroupBox("Reference foil changer") + layout = QVBoxLayout(reffoil_group) + form = QFormLayout() + layout.addLayout(form) + form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint) + self.reffoil_selector = QComboBox() + available_foils = self._dev.reffoilchanger.get_all_foils() + self.reffoil_selector.addItems(available_foils) + + form.addRow("Reference foil", self.reffoil_selector) + + button_layout = QHBoxLayout() + generate_cmd = MyButton("Generate command", "default") + button_layout.addWidget(generate_cmd) + button_layout.addStretch(1) + layout.addLayout(button_layout) + + generate_cmd.clicked.connect(self._generate_reffoil_command) + + return reffoil_group + return None + + def _generate_reffoil_command(self): + cmd = f"dev.reffoilchanger.insert(ref='{self.reffoil_selector.currentText()}', wait=True)" + self.custom_edit.setText(cmd) + + def _create_auto_gain_form(self): + auto_gain_group = QGroupBox("Auto Gain") + layout = QVBoxLayout(auto_gain_group) + + edge_selector_layout = QHBoxLayout() + edge_selector_label = QLabel("Absorption edge:") + self.edge_selector_button = MyButton("Choose", "default") + self.edge_label = QLabel("No edge selected") + edge_selector_layout.addWidget(edge_selector_label) + edge_selector_layout.addWidget(self.edge_selector_button) + edge_selector_layout.addWidget(self.edge_label) + edge_selector_layout.addStretch() + self.edge_element = None + self.edge_edge = None + + layout.addLayout(edge_selector_layout) + + layout.addWidget(QLabel("Auto gain on")) + + form = QFormLayout() + layout.addLayout(form) + form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint) + + self.ic0_check = QCheckBox("") + form.addRow("IC0", self.ic0_check) + self.ic1_check = QCheckBox("") + form.addRow("IC1", self.ic1_check) + self.ic2_check = QCheckBox("") + form.addRow("IC2", self.ic2_check) + self.pips_check = QCheckBox("") + form.addRow("PIPS", self.pips_check) + + button_layout = QHBoxLayout() + generate_cmd = MyButton("Generate command", "default") + button_layout.addWidget(generate_cmd) + button_layout.addStretch(1) + layout.addLayout(button_layout) + + self.edge_selector_button.clicked.connect(self._update_edge) + generate_cmd.clicked.connect(self._generate_auto_gain_command) + + return auto_gain_group + + def _update_edge(self, *_): + match self.beamline: + case "x01da": + dlg = self.EdgeSelector(self) + case "x10da": + dlg = self.EdgeSelector(self) + case _: + dlg = self.EdgeSelector(self) + if dlg.exec_(): + self.edge_energy = dlg.selected_energy + self.edge_label.setText( + f"{dlg.selected_element}, {dlg.selected_edge}-edge, {dlg.selected_energy:0.1f} eV" + ) + self.edge_element = dlg.selected_element + self.edge_edge = dlg.selected_edge + + def _generate_auto_gain_command(self): + if self.edge_edge is None or self.edge_element is None: + return + amplifiers = [] + for amp, name in [ + (self.ic0_check, "ic0"), + (self.ic1_check, "ic1"), + (self.ic2_check, "ic2"), + (self.pips_check, "pips"), + ]: + if amp.isChecked(): + amplifiers.append(name) + if amplifiers == []: + return + cmd = f"auto_gain.start(element={self.edge_element}, edge={self.edge_edge}, {amplifiers}, comp_ring_current=True)" + self.custom_edit.setText(cmd) + + def _collect_custom_result(self) -> dict: + text = self.custom_edit.text().strip() + if not text: + raise ValueError("Command must not be empty.") + return {"command": text, "kind": "custom", "form_state": {"text": text}} + + # ------------------------------------------------------------------ # + # pre-fill (edit mode) / result extraction + # ------------------------------------------------------------------ # + def _apply_initial(self, initial: dict): + kind = initial.get("kind") + state = initial.get("form_state") or {} + + if kind == "scan" and state.get("scan_name"): + self._prefill_scan_tab( + state["scan_name"], state.get("args") or [], state.get("kwargs") or {} + ) + self.tabs.setCurrentIndex(_TAB_SCAN) + elif kind == "move" and state.get("source") == "digital_twin": + # Reloading a captured config back into DigitalTwin's own input + # fields would need inverting get_assistant_config()'s unit + # conversions and mode branching (fm_focus, mo1_mode, ...) + # field-by-field - not implemented yet. Fall back to showing + # the generated command as read/write text instead of silently + # dropping the captured config; the "Add" (=Ok) button below + # will just resubmit that text unchanged unless it's edited. + self.custom_edit.setText(initial.get("command", "")) + self.tabs.setCurrentIndex(_TAB_OTHER) + elif kind == "move" and state.get("device_name"): + idx = self.move_device_combo.findText(state["device_name"]) + if idx >= 0: + self.move_device_combo.setCurrentIndex(idx) + self.move_value_spin.setValue(float(state.get("value", 0.0))) + self.move_relative_check.setChecked(bool(state.get("relative", False))) + self.tabs.setCurrentIndex(_TAB_MOVE) + elif initial.get("command"): + # "custom" kind, or a legacy/unrecognized item - fall back to + # showing the raw command text as-is. + self.custom_edit.setText(state.get("text", initial["command"])) + self.tabs.setCurrentIndex(_TAB_OTHER) + else: + logger.warning(f"Unknown kind: {kind}") + + def _on_accept(self): + try: + result = self._collect_result() + except ValueError as exc: + QMessageBox.warning(self, "Missing input", str(exc)) + return + self._result = result + self.accept() + + def _collect_result(self) -> dict: + current = self.tabs.currentIndex() + if current == _TAB_SCAN: + return self._collect_scan_result() + if current == _TAB_MOVE: + return self._collect_move_result() + if current == _TAB_DIGITAL_TWIN: + return self._collect_digital_twin_result() + return self._collect_custom_result() + + def result(self) -> dict: + """Valid after `exec_()` returns `QDialog.Accepted`.""" + return self._result + + +def _format_scan_call(scan_name: str, args: list, kwargs: dict) -> str: + parts = [repr(a) for a in args] + parts += [f"{name}={value!r}" for name, value in kwargs.items()] + return f"scans.{scan_name}({', '.join(parts)})" diff --git a/debye_bec/bec_widgets/widgets/scheduler/notification_dialog.py b/debye_bec/bec_widgets/widgets/scheduler/notification_dialog.py new file mode 100644 index 0000000..cc7dd81 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/notification_dialog.py @@ -0,0 +1,231 @@ +""" +Settings dialog for `notifications.NotificationSettings` - kept separate +from the send logic itself (`notifications.py`). +""" + +from __future__ import annotations + +from bec_lib.logger import bec_logger +from qtpy.QtWidgets import ( + QCheckBox, + QComboBox, + QDialog, + QDialogButtonBox, + QGroupBox, + QLabel, + QMessageBox, + QPushButton, + QVBoxLayout, +) + +from .enums import ScheduleItemStatus +from .notifications import NotificationSettings, NotificationTarget, send_notification + +logger = bec_logger.logger + + +def _service_state(client, name: str) -> tuple[bool, list[str]]: + """ + Best-effort read of a messaging service's enabled/scopes state (see + the module docstring in notifications.py for why this is the only + introspection available). Never raises - a client without `messaging` + set up yet just looks like "not enabled, no scopes". + """ + messaging = getattr(client, "messaging", None) + service = getattr(messaging, name, None) if messaging is not None else None + if service is None: + return False, [] + return bool(getattr(service, "_enabled", False)), sorted(getattr(service, "_scopes", set())) + + +class NotificationSettingsDialog(QDialog): + """Choose whether/where/for-which-item-types finish/fail notifications are sent.""" + + def __init__(self, settings: NotificationSettings, client, parent=None): + super().__init__(parent) + self._client = client + self.setWindowTitle("Notifications") + self.setMinimumWidth(380) + + layout = QVBoxLayout(self) + info_label = QLabel( + "Sends a message through the selected BEC messaging service whenever a " + "schedule item finishes or fails." + ) + info_label.setWordWrap(True) + layout.addWidget(info_label) + + self.enabled_check = QCheckBox("Send notifications") + self.enabled_check.setChecked(settings.enabled) + layout.addWidget(self.enabled_check) + + kind_box = QGroupBox("Notify for") + kind_layout = QVBoxLayout(kind_box) + self.scan_check = QCheckBox("Scans") + self.scan_check.setChecked(settings.notify_scan) + self.move_check = QCheckBox("Movements") + self.move_check.setChecked(settings.notify_move) + self.rpc_check = QCheckBox("RPC / custom commands") + self.rpc_check.setChecked(settings.notify_rpc) + for check in (self.scan_check, self.move_check, self.rpc_check): + kind_layout.addWidget(check) + layout.addWidget(kind_box) + + target_box = QGroupBox("Send via") + target_layout = QVBoxLayout(target_box) + self.scilog_check = QCheckBox("SciLog") + self.teams_check = QCheckBox("Microsoft Teams") + self.signal_check = QCheckBox("Signal Messenger") + for check in (self.scilog_check, self.teams_check, self.signal_check): + target_layout.addWidget(check) + layout.addWidget(target_box) + + event_box = QGroupBox("Send on") + event_layout = QVBoxLayout(event_box) + self.completed_check = QCheckBox("Completed") + self.aborted_check = QCheckBox("Aborted") + self.failed_check = QCheckBox("Failed") + for check in (self.completed_check, self.aborted_check, self.failed_check): + event_layout.addWidget(check) + layout.addWidget(event_box) + + # Teams: contact/channel can only be a scope an admin has already + # registered for the deployment - there's no way to address an + # arbitrary Teams user from the client, so this is a dropdown, not + # free text (see notifications.py's module docstring). + self.teams_box = QGroupBox("Teams settings") + teams_layout = QVBoxLayout(self.teams_box) + teams_enabled, teams_scopes = _service_state(client, "teams") + self.teams_combo = QComboBox() + self.teams_combo.addItems(teams_scopes) + self.teams_combo.setEnabled(teams_enabled) + teams_layout.addWidget(QLabel("Contact or channel")) + teams_layout.addWidget(self.teams_combo) + if not teams_enabled: + teams_layout.addWidget(_dim_label("Teams messaging is not enabled for this session.")) + elif not teams_scopes: + teams_layout.addWidget(_dim_label("No Teams contacts/channels are registered.")) + layout.addWidget(self.teams_box) + + # Signal: a raw phone number works directly (BEC normalizes it - + # see the docs), so this is an editable combo box: type a number, + # or pick a pre-registered scope if any exist. + self.signal_box = QGroupBox("Signal settings") + signal_layout = QVBoxLayout(self.signal_box) + signal_enabled, signal_scopes = _service_state(client, "signal") + self.signal_combo = QComboBox() + self.signal_combo.setEditable(True) + self.signal_combo.addItems(signal_scopes) + self.signal_combo.setEnabled(signal_enabled) + signal_layout.addWidget(QLabel("Phone number or contact")) + self.signal_combo.setEditText("") + self.signal_combo.lineEdit().setPlaceholderText("e.g. +41791234567") + signal_layout.addWidget(self.signal_combo) + if not signal_enabled: + signal_layout.addWidget(_dim_label("Signal messaging is not enabled for this session.")) + layout.addWidget(self.signal_box) + + # pre-fill from the persisted settings + { + NotificationTarget.SCILOG: self.scilog_check, + NotificationTarget.TEAMS: self.teams_check, + NotificationTarget.SIGNAL: self.signal_check, + }.get(settings.target, self.scilog_check).setChecked(True) + for event in settings.events: + { + ScheduleItemStatus.COMPLETED: self.completed_check, + ScheduleItemStatus.ABORTED: self.aborted_check, + ScheduleItemStatus.FAILED: self.failed_check, + }.get(event).setChecked(True) + if settings.teams_scope: + idx = self.teams_combo.findText(settings.teams_scope) + if idx >= 0: + self.teams_combo.setCurrentIndex(idx) + if settings.signal_number: + self.signal_combo.setEditText(settings.signal_number) + + self.scilog_check.toggled.connect(lambda *_: self._sync_target_widgets(self.scilog_check)) + self.teams_check.toggled.connect(lambda *_: self._sync_target_widgets(self.teams_check)) + self.signal_check.toggled.connect(lambda *_: self._sync_target_widgets(self.signal_check)) + + self.test_btn = QPushButton("Send test message") + self.test_btn.clicked.connect(self._on_test_clicked) + layout.addWidget(self.test_btn) + + buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) + buttons.accepted.connect(self.accept) + buttons.rejected.connect(self.reject) + layout.addWidget(buttons) + + self._sync_target_widgets() + + def _sync_target_widgets(self, current_button: QCheckBox = None): + if current_button is None: + for button in [self.scilog_check, self.teams_check, self.signal_check]: + if button.isChecked(): + current_button = button + break + if current_button.isChecked(): + for button in [self.scilog_check, self.teams_check, self.signal_check]: + if button != current_button: + button.setChecked(False) + if current_button == self.teams_check: + self.teams_box.setVisible(True) + self.signal_box.setVisible(False) + elif current_button == self.signal_check: + self.teams_box.setVisible(False) + self.signal_box.setVisible(True) + else: + self.teams_box.setVisible(False) + self.signal_box.setVisible(False) + self.layout().activate() + self.adjustSize() + + def _current_settings(self) -> NotificationSettings: + if self.scilog_check.isChecked(): + target = NotificationTarget.SCILOG + elif self.teams_check.isChecked(): + target = NotificationTarget.TEAMS + elif self.signal_check.isChecked(): + target = NotificationTarget.SIGNAL + else: + raise ValueError(f"Target {target} not supported") + + events = [] + if self.completed_check.isChecked(): + events.append(ScheduleItemStatus.COMPLETED) + if self.aborted_check.isChecked(): + events.append(ScheduleItemStatus.ABORTED) + if self.failed_check.isChecked(): + events.append(ScheduleItemStatus.FAILED) + + return NotificationSettings( + enabled=self.enabled_check.isChecked(), + notify_scan=self.scan_check.isChecked(), + notify_move=self.move_check.isChecked(), + notify_rpc=self.rpc_check.isChecked(), + target=target, + events=events, + teams_scope=self.teams_combo.currentText().strip() or None, + signal_number=self.signal_combo.currentText().strip() or None, + ) + + def _on_test_clicked(self): + settings = self._current_settings() + try: + send_notification(self._client, settings, "Test message from the BEC schedule widget.") + except Exception as exc: # pylint: disable=broad-except + QMessageBox.critical(self, "Test message failed", str(exc)) + return + QMessageBox.information(self, "Test message sent", "The test message was sent.") + + def result_settings(self) -> NotificationSettings: + """Valid after `exec_()` returns `QDialog.Accepted`.""" + return self._current_settings() + + +def _dim_label(text: str) -> QLabel: + label = QLabel(text) + label.setStyleSheet("color: gray;") + label.setWordWrap(True) + return label diff --git a/debye_bec/bec_widgets/widgets/scheduler/notifications.py b/debye_bec/bec_widgets/widgets/scheduler/notifications.py new file mode 100644 index 0000000..bd8fbf3 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/notifications.py @@ -0,0 +1,159 @@ +""" +A small, self-contained "tell the user when something finished/failed" +feature, independent of scheduling/execution logic. + +Sends through BEC's own `client.messaging` container +(`bec_lib.messaging_services.MessagingContainer`, exposed as `bec.messaging` +in the IPython client) - the same mechanism documented at +https://bec.readthedocs.io/latest/how-to/general/send-messages-to-signal.html +and +https://bec.readthedocs.io/latest/how-to/general/send-messages-to-scilog.html: + + bec.messaging.signal.new("Beamline checks completed.").send(scope="+41791234567") + bec.messaging.scilog.new("Beamline checks completed.").send() + bec.messaging.teams.new("Beamline checks completed.").send(scope=) + +Notes on each service, read directly from `bec_lib.messaging_services` +rather than assumed: + +- **SciLog** needs no `scope` - it posts to the logbook of the currently + active pgroup automatically. +- **Signal** accepts a raw phone number as `scope` (BEC normalizes it, + defaulting to the Swiss country code for numbers without one - see the + docs above) *or* a pre-registered scope name for a group. +- **Teams** is a plain `MessagingService` with no Teams-specific methods: + a contact/channel can only be addressed by a `scope` an admin has + already registered for the deployment (`service._scopes`) - there is no + way to address an arbitrary Teams user/email directly from the client. + This is why the Teams UI in `notification_dialog.py` is a dropdown of + known scopes rather than free text, while Signal's is free text (with + any registered scopes offered as suggestions). + +`client.messaging.._enabled` / `._scopes` are the only +introspection available (no public accessor exists at the time of +writing) - used here and in the dialog to grey out a target that isn't +configured for the current session instead of only failing at send time. +""" + +from __future__ import annotations + +from enum import Enum + +from bec_lib.logger import bec_logger +from pydantic import BaseModel + +from .enums import ScheduleItemStatus + +logger = bec_logger.logger + + +class NotificationTarget(str, Enum): + SCILOG = "scilog" + TEAMS = "teams" + SIGNAL = "signal" + + +class NotificationSettings(BaseModel): + """Persisted per-schedule notification preferences.""" + + enabled: bool = False + notify_scan: bool = True + notify_move: bool = True + # "custom" is what this plugin calls the free-text tab (see + # item_dialog.py) - it covers RPC calls to other widgets as well as + # any hand-typed command, so it's exposed to the user as "RPC". + notify_rpc: bool = True + + target: NotificationTarget = NotificationTarget.SCILOG + events: list[ScheduleItemStatus] = [] + teams_scope: str | None = None # a pre-registered Teams scope (channel/contact) + signal_number: str | None = None # a phone number, or a pre-registered Signal scope + + +def _enabled_for_kind(settings: NotificationSettings, kind: str) -> bool: + if not settings.enabled: + return False + return { + "scan": settings.notify_scan, + "move": settings.notify_move, + "custom": settings.notify_rpc, + }.get(kind, False) + + +def send_notification(client, settings: NotificationSettings, text: str) -> None: + """ + Send `text` through whichever service `settings.target` selects. + + Raises on failure (no target selected, missing Teams scope/Signal + number, the target service not enabled for this session, ...) - + callers that want a fire-and-forget send should catch around this. + `notify_item_finished` does; the dialog's "Send test message" button + deliberately does not, so the operator sees exactly what went wrong. + """ + if settings.target == NotificationTarget.SCILOG: + client.messaging.scilog.new(text).send() + elif settings.target == NotificationTarget.TEAMS: + if not settings.teams_scope: + raise ValueError("No Teams contact/channel selected.") + client.messaging.teams.new(text).send(scope=settings.teams_scope) + elif settings.target == NotificationTarget.SIGNAL: + if not settings.signal_number: + raise ValueError("No Signal phone number/contact set.") + client.messaging.signal.new(text).send(scope=settings.signal_number) + else: + raise ValueError("No notification target selected.") + + +def notify_item_finished( + client, + kind: str, + command: str, + scan_number: str, + final_status: ScheduleItemStatus, + settings: NotificationSettings, + error: str | None = None, +): + """ + Sends a message if notifications are enabled for `kind` and a target + is configured. Never raises - a notification failing to send should + not affect schedule execution. + """ + if not _enabled_for_kind(settings, kind): + return + if final_status not in settings.events: + return + + status_word = final_status.name + text = f"**BEC Scheduler Widget**\nStatus: {kind} {status_word}" + if scan_number is not None: + text = text + f"\nScan number: {scan_number}" + text = text + f"\n\nCommand: {command}" + if error: + text += f"\n\n{error.strip().splitlines()[-1]}" + + try: + send_notification(client, settings, text) + except Exception: # pylint: disable=broad-except + logger.exception(f"Failed to send schedule notification via {settings.target.value}") + + +def notify_schedule_state(client, settings, schedule_name, state): + """ + Sends a message if notifications are enabled and a schedule is started or stopped. + Never raises - a notification failing to send should not affect schedule execution. + """ + text = f'**BEC Scheduler Widget**\nSchedule "{schedule_name}" ' + if state == "started": + text = text + "has started" + elif state == "aborted": + text = text + "was aborted" + elif state == "widget_closed": + text = text + "has finished prematurely because widget was closed" + elif state == "finished": + text = text + "has finished" + else: + logger.warning(f"Unknown schedule state {state}") + try: + send_notification(client, settings, text) + except Exception: # pylint: disable=broad-except + logger.exception(f"Failed to send schedule notification via {settings.target.value}") diff --git a/debye_bec/bec_widgets/widgets/scheduler/qt_widgets.py b/debye_bec/bec_widgets/widgets/scheduler/qt_widgets.py new file mode 100644 index 0000000..4f47e49 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/qt_widgets.py @@ -0,0 +1,61 @@ +from bec_widgets.utils.colors import get_accent_colors +from qtpy.QtCore import Qt, Signal + +# pylint: disable=E0611 +from qtpy.QtGui import QKeySequence +from qtpy.QtWidgets import QListWidget, QPushButton + + +class MyListWidget(QListWidget): + deletePressed = Signal() + emptySpaceClicked = Signal() + copyPressed = Signal() + pastePressed = Signal() + + def keyPressEvent(self, event): + if event.key() == Qt.Key_Delete and self.currentItem() is not None: + self.deletePressed.emit() + event.accept() + return + + if event.matches(QKeySequence.StandardKey.Copy) and self.currentItem() is not None: + self.copyPressed.emit() + event.accept() + return + + if event.matches(QKeySequence.StandardKey.Paste): + self.pastePressed.emit() + event.accept() + return + + super().keyPressEvent(event) + + def mousePressEvent(self, event): + if event.button() == Qt.LeftButton and self.itemAt(event.position().toPoint()) is None: + self.clearSelection() + self.setCurrentRow(-1) + self.emptySpaceClicked.emit() + return + + super().mousePressEvent(event) + + +class MyButton(QPushButton): + def __init__(self, text="", color="default", parent=None): + self.color = color + super().__init__(text, parent) + self.apply_theme() + + def apply_theme(self): + if self.isEnabled(): + colors = get_accent_colors() + color = getattr(colors, self.color).name() + self.setStyleSheet(f"QPushButton {{ background-color: {color}; color: white; }}") + else: + self.setStyleSheet( + "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}" + ) + + def setEnabled(self, enable: bool = True): + super().setEnabled(enable) + self.apply_theme() diff --git a/debye_bec/bec_widgets/widgets/scheduler/register_scheduler.py b/debye_bec/bec_widgets/widgets/scheduler/register_scheduler.py new file mode 100644 index 0000000..6b6f1a1 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/register_scheduler.py @@ -0,0 +1,15 @@ +def main(): # pragma: no cover + from qtpy import PYSIDE6 + + if not PYSIDE6: + print("PYSIDE6 is not available in the environment. Cannot patch designer.") + return + from PySide6.QtDesigner import QPyDesignerCustomWidgetCollection + + from .scheduler_plugin import SchedulerPlugin + + QPyDesignerCustomWidgetCollection.addCustomWidget(SchedulerPlugin()) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/debye_bec/bec_widgets/widgets/scheduler/schedule_item.py b/debye_bec/bec_widgets/widgets/scheduler/schedule_item.py new file mode 100644 index 0000000..b675094 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/schedule_item.py @@ -0,0 +1,69 @@ +""" +Plain (non-BECMessage) data model for one schedule entry and the schedule +as a whole. + +These are ordinary pydantic models used only client-side, for validation +and convenience. They are never sent over Redis as their own type - only +ever as the `.value` payload of a `bec_lib.messages.VariableMessage` (see +`endpoints.schedule` for why). +""" + +from __future__ import annotations + +from enum import Enum +from typing import Literal + +from pydantic import BaseModel, Field + +from .enums import ScheduleItemStatus +from .guard import GuardSettings +from .notifications import NotificationSettings + +# Which tab of the Add/Edit dialog built `command`, and therefore how to +# re-open that dialog pre-filled with the same values for editing. "custom" +# covers hand-typed text (including RPC calls to other widgets), which has +# no structured `form_state` to restore beyond the raw text itself. +ScheduleItemKind = Literal["scan", "move", "custom"] + + +class ScheduleItem(BaseModel): + """A single command in the schedule, plus its execution bookkeeping.""" + + item_id: str + command: str + status: ScheduleItemStatus = ScheduleItemStatus.PENDING + + # How `command` was produced, and (for "scan"/"move") the structured + # inputs behind it, so the Add/Edit dialog can be reopened pre-filled + # instead of asking the user to re-type everything. `command` itself + # remains the single source of truth actually evaluated at execution + # time - `kind`/`form_state` only drive the UI. + kind: ScheduleItemKind = "custom" + form_state: dict | None = None + + # Bookkeeping used to reconnect to a submission that is still (or was) + # in flight on the BEC scan/device server, after this widget has been + # closed and reopened. + request_id: str | None = None + scan_id: str | None = None + scan_number: int | None = None + + error: str | None = None + started_at: float | None = None + finished_at: float | None = None + + +class Schedule(BaseModel): + """The full, persisted state of one schedule-widget instance.""" + + schedule_name: str + items: list[ScheduleItem] = Field(default_factory=list) + is_running: bool = False + notes: str = "" + + # Settings for two independent, optional features (see notifications.py + # and guard.py) - persisted here alongside the schedule itself so they + # survive a widget restart too, but their *logic* lives entirely in + # those separate modules; this is just where their settings are stored. + notifications: NotificationSettings = Field(default_factory=NotificationSettings) + guard: GuardSettings = Field(default_factory=GuardSettings) diff --git a/debye_bec/bec_widgets/widgets/scheduler/schedule_logic.py b/debye_bec/bec_widgets/widgets/scheduler/schedule_logic.py new file mode 100644 index 0000000..4073793 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/schedule_logic.py @@ -0,0 +1,71 @@ +""" +Pure bookkeeping over a list of `ScheduleItem`s - no Qt, no locking, no I/O. +Kept separate from `schedule_widget.py` for the same reason `guard.py` and +`notifications.py` are: it's testable on its own, and the widget should +only be responsible for Qt/orchestration, not figuring out which item runs +next. + +The one invariant everything here assumes and preserves: at any moment, a +schedule's items form a (possibly empty) prefix that is no longer PENDING +(RUNNING - at most one, the one currently executing - COMPLETED, FAILED or +ABORTED), followed by a suffix that is entirely PENDING. `ScheduleWidget` +is responsible for only ever mutating the PENDING suffix (see its module +docstring) and for calling these functions while holding its lock; nothing +here does its own locking. +""" + +from __future__ import annotations + +from typing import Literal + +from bec_lib.logger import bec_logger + +from .schedule_item import ScheduleItem, ScheduleItemStatus + +logger = bec_logger.logger + + +def protected_prefix_length(items: list[ScheduleItem]) -> int: + """How many items, from the start, are no longer PENDING.""" + count = 0 + for item in items: + if item.status != ScheduleItemStatus.PENDING: + count += 1 + else: + break + return count + + +def index_of(items: list[ScheduleItem], item_id: str | None) -> int | None: + if item_id is None: + return None + for i, item in enumerate(items): + if item.item_id == item_id: + return i + return None + + +def pick_next_runnable( + items: list[ScheduleItem], repeat_aborted_item +) -> ScheduleItem | Literal["stop"] | None: + """ + What the execution loop should do next: always re-derived from scratch + (never a remembered index/object) so edits made to the PENDING suffix + between calls are picked up correctly. + + Returns the next item to run/attach to, the string `"stop"` if an + earlier item failed/was aborted (execution stays parked there until the + operator intervenes), or `None` if every item is COMPLETED. + """ + logger.info(f"pick next runnable, repeat abort item is {repeat_aborted_item}") + for item in items: + logger.info(f"item: {item}") + if item.status == ScheduleItemStatus.COMPLETED: + continue + if item.status == ScheduleItemStatus.ABORTED: + if not repeat_aborted_item: + continue + if item.status == ScheduleItemStatus.FAILED: + return "stop" + return item # PENDING, RUNNING or ABORTED if repeat_aborted_item is set (reconciled as still active) + return None diff --git a/debye_bec/bec_widgets/widgets/scheduler/scheduler.py b/debye_bec/bec_widgets/widgets/scheduler/scheduler.py new file mode 100644 index 0000000..97e25a2 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/scheduler.py @@ -0,0 +1,1305 @@ +""" +ScheduleWidget - persists an ordered schedule of BEC commands (scans, +device moves, ...) to Redis, executes them one by one, and survives being +closed and reopened - including while an item is still running +server-side. Items that have not started yet can be added, reordered or +deleted while the schedule is executing; items that have already started +or finished cannot. +""" + +from __future__ import annotations + +import json +import sys +import threading +import time +import traceback +import uuid +from typing import Literal, Optional + +import numpy as np +from bec_lib.endpoints import MessageEndpoints +from bec_lib.logger import bec_logger +from bec_lib.messages import VariableMessage +from bec_qthemes import material_icon +from bec_widgets.utils.bec_connector import ConnectionConfig +from bec_widgets.utils.bec_dispatcher import BECDispatcher +from bec_widgets.utils.bec_widget import BECWidget +from bec_widgets.utils.error_popups import SafeSlot +from pydantic import ValidationError + +# pylint: disable=E0611 +from qtpy.QtCore import Qt, QTimer, Signal +from qtpy.QtWidgets import ( + QApplication, + QDialog, + QFileDialog, + QGroupBox, + QHBoxLayout, + QLabel, + QLineEdit, + QListWidgetItem, + QMessageBox, + QPlainTextEdit, + QVBoxLayout, + QWidget, +) + +from .endpoints import schedule as schedule_endpoint +from .enums import ScheduleItemStatus +from .guard import SignalGuard +from .guard_dialog import GuardSettingsDialog +from .item_dialog import ScheduleItemDialog +from .notification_dialog import NotificationSettingsDialog +from .notifications import ( + NotificationSettings, + NotificationTarget, + notify_item_finished, + notify_schedule_state, +) +from .qt_widgets import MyButton, MyListWidget +from .schedule_item import Schedule, ScheduleItem +from .schedule_logic import index_of, pick_next_runnable, protected_prefix_length + +logger = bec_logger.logger + +_ACTIVE_QUEUE_STATES = ("PENDING", "RUNNING") +_ITEM_ID_ROLE = Qt.UserRole + 1 # QListWidgetItem data role used to map a row back to an item_id + +ICON_SIZE = 20 +_ICON_MAP = { + ScheduleItemStatus.PENDING: ("hourglass", "#919090"), + ScheduleItemStatus.RUNNING: ("cycle", "#2980b9"), + ScheduleItemStatus.COMPLETED: ("check", "#27ae60"), + ScheduleItemStatus.FAILED: ("warning", "#e74c3c"), + ScheduleItemStatus.ABORTED: ("cancel", "#e6d922"), +} + + +class ScheduleWidgetConfig(ConnectionConfig): + """ + Per-instance identity/settings for `ScheduleWidget`, following the same + pattern `bec_widgets` itself uses (e.g. `ScanControlConfig`): a + `ConnectionConfig` subclass constructed in `__init__` when none is + passed, carrying whatever small settings distinguish one instance from + another. + + `schedule_name` lives here rather than as a bare constructor kwarg for + that consistency - but note this is *not* where the schedule's actual + contents (items, guard/notification settings) live: `self.config` is + used for RPC identity and optional explicit GUI-layout save/restore, + not auto-persisted across an ordinary widget close/reopen the way this + plugin's Redis-backed `Schedule` document is. Conflating the two would + quietly reintroduce the "closed widget loses its schedule" problem this + plugin exists to solve. + """ + + schedule_name: str = "default_schedule" + + +class Scheduler(BECWidget, QWidget): + """Schedule, persist and execute a sequence of BEC scan/device commands.""" + + PLUGIN = True + ICON_NAME = "assignment_add" + USER_ACCESS = [ + "run_schedule", + "abort_schedule", + "reset_schedule", + "get_status", + "add_item", + "edit_item", + "delete_item", + "move_item", + "move_item_up", + "move_item_down", + ] + + # emitted (thread-safe, may come from the background execution task) + # whenever the schedule state changes and the UI should redraw + schedule_changed = Signal() + + def __init__( + self, + parent=None, + client=None, + config: ScheduleWidgetConfig | dict | None = None, + gui_id: str | None = None, + schedule_name: str | None = None, + **kwargs, + ): + if config is None: + config = ScheduleWidgetConfig( + widget_class=self.__class__.__name__, + schedule_name=schedule_name or "default_schedule", + ) + elif schedule_name is not None: + # Allow overriding even when a config object/dict was supplied, + # mirroring how ScanControl treats its own optional constructor + # args as overrides onto whatever config ends up in self.config. + if isinstance(config, dict): + config = dict(config, schedule_name=schedule_name) + else: + config.schedule_name = schedule_name + super().__init__(parent=parent, client=client, gui_id=gui_id, config=config, **kwargs) + + # A stable name, NOT self.gui_id: gui_id is regenerated every time + # the widget is constructed, so it could never be used to find a + # previously persisted schedule again after the widget is reopened. + self.schedule_name = self.config.schedule_name + + self.get_bec_shortcuts() # -> self.client, self.dev, self.scans, self.queue + self.beamline = self.get_beamline() + self.connector = self.client.connector + + if self.beamline in ['x01da', 'x10da']: + logger.info( + 'Scheduler running at X01DA or X10DA, import and load digital twin and auto-gain' + ) + from ....bec_ipython_client.plugins.digital_twin_core.digital_twin_core import DigitalTwinCore + from ....bec_ipython_client.plugins.auto_gain import AutoGain + self.digital_twin = DigitalTwinCore() + self.auto_gain = AutoGain() + else: + self.digital_twin = None + self.auto_gain = None + + self._endpoint = schedule_endpoint(self.schedule_name) + self._lock = threading.Lock() # guards self.schedule (see module docstring) + self._abort_requested = False + self._guard_interrupt_requested = False # see "Guard interaction" in module docstring + self._current_report = None # ScanReport of the item currently executing, if any + self._current_item_kind: str | None = None # kind of the item _current_report belongs to + self._selected_item_id: str | None = None # survives list repopulation on refresh + self._clipboard_item: dict | None = None # copy/paste buffer, see _on_copy_clicked + + # Notes are auto-saved on a debounce timer rather than on every + # keystroke - see _on_notes_changed/_persist_notes. + self._notes_save_timer = QTimer(self) + self._notes_save_timer.setSingleShot(True) + self._notes_save_timer.setInterval(200) + self._notes_save_timer.timeout.connect(self._persist_notes) + + self.schedule_changed.connect(self._refresh_ui) + + self.schedule: Schedule = self._load_or_seed_schedule() + self._reconcile_with_live_state() + + self._guard = SignalGuard(self.connector, parent=self, dev=self.dev) + self._guard.paused.connect(self._on_guard_paused) + self._guard.resumed.connect(self._on_guard_resumed) + self._guard.value_update.connect(self._update_guard_label) + self._guard.configure(self.schedule.guard) + + self._build_ui() + self._refresh_ui() + + # Stay in sync with any other subscriber touching the same + # schedule (another instance of this widget, a script, ...). + self.bec_dispatcher.connect_slot(self._on_remote_update, self._endpoint) + + self._closing = threading.Event() + + def get_beamline(self): + hostname = self.client._hostname + start = hostname.find("x") + if start != -1: + return hostname[start : start + 5] + + # ------------------------------------------------------------------ # + # UI + # ------------------------------------------------------------------ # + def _build_ui(self): + layout = QVBoxLayout(self) + + schedule_group = QGroupBox("Schedule") + schedule_layout = QVBoxLayout(schedule_group) + + schedule_name_layout = QHBoxLayout() + name_label = QLabel("Name") + self.schedule_name_input = QLineEdit() + schedule_name_layout.addWidget(name_label) + schedule_name_layout.addWidget(self.schedule_name_input) + schedule_name_layout.addStretch() + schedule_layout.addLayout(schedule_name_layout) + + self.schedule_name_input.editingFinished.connect(self._on_schedule_name_change) + + edit_row = QHBoxLayout() + self.add_btn = MyButton("Add", "default") + self.edit_btn = MyButton("Edit", "default") + self.delete_btn = MyButton("Delete", "default") + self.move_up_btn = MyButton("Move Up", "default") + self.move_down_btn = MyButton("Move Down", "default") + self.copy_btn = MyButton("Copy", "default") + self.paste_btn = MyButton("Paste", "default") + for button in ( + self.add_btn, + self.edit_btn, + self.delete_btn, + self.move_up_btn, + self.move_down_btn, + self.copy_btn, + self.paste_btn, + ): + edit_row.addWidget(button) + edit_row.addStretch() + schedule_layout.addLayout(edit_row) + + self.list_widget = MyListWidget() + self.list_widget.currentRowChanged.connect(self._on_selection_changed) + self.list_widget.deletePressed.connect(self._on_delete_clicked) + self.list_widget.emptySpaceClicked.connect(self._on_empty_space_clicked) + self.list_widget.itemDoubleClicked.connect(self._on_edit_clicked) + self.list_widget.copyPressed.connect(self._on_copy_clicked) + self.list_widget.pastePressed.connect(self._on_paste_clicked) + + schedule_layout.addWidget(self.list_widget) + + schedule_layout.addWidget(QLabel("Notes")) + self.notes_edit = QPlainTextEdit() + self.notes_edit.setPlaceholderText("Notes about this schedule...") + self.notes_edit.setPlainText(self.schedule.notes) + self.notes_edit.setMaximumHeight(100) + self.notes_edit.textChanged.connect(self._on_notes_changed) + schedule_layout.addWidget(self.notes_edit) + + self.add_btn.clicked.connect(self._on_add_clicked) + self.edit_btn.clicked.connect(self._on_edit_clicked) + self.delete_btn.clicked.connect(self._on_delete_clicked) + self.move_up_btn.clicked.connect(self._on_move_up_clicked) + self.move_down_btn.clicked.connect(self._on_move_down_clicked) + self.copy_btn.clicked.connect(self._on_copy_clicked) + self.paste_btn.clicked.connect(self._on_paste_clicked) + + layout.addWidget(schedule_group) + + control_group = QGroupBox("Control") + control_layout = QVBoxLayout(control_group) + + run_row = QHBoxLayout() + self.run_btn = MyButton("Run / Continue", "success") + self.abort_btn = MyButton("Abort", "emergency") + self.reset_btn = MyButton("Reset", "warning") + for button in (self.run_btn, self.abort_btn, self.reset_btn): + run_row.addWidget(button) + run_row.addStretch() + control_layout.addLayout(run_row) + + self.run_btn.clicked.connect(self.run_schedule) + self.abort_btn.clicked.connect(self.abort_schedule) + self.reset_btn.clicked.connect(self.reset_schedule) + + layout.addWidget(control_group) + + settings_group = QGroupBox("Settings") + settings_layout = QVBoxLayout(settings_group) + + guard_row = QHBoxLayout() + self.guard_btn = MyButton("Auto-pause", "default") + self.guard_status_label = QLabel() + guard_row.addWidget(self.guard_btn) + guard_row.addWidget(self.guard_status_label) + guard_row.addStretch() + notifications_row = QHBoxLayout() + self.notifications_btn = MyButton("Notifications", "default") + self.notification_status_label = QLabel() + notifications_row.addWidget(self.notifications_btn) + notifications_row.addWidget(self.notification_status_label) + notifications_row.addStretch() + settings_layout.addLayout(guard_row) + settings_layout.addLayout(notifications_row) + + self.notifications_btn.clicked.connect(self._on_notifications_clicked) + self.guard_btn.clicked.connect(self._on_guard_clicked) + + layout.addWidget(settings_group) + + file_group = QGroupBox("File") + file_layout = QHBoxLayout(file_group) + self.save_file_btn = MyButton("Save to file", "default") + self.load_file_btn = MyButton("Load from file", "default") + file_layout.addWidget(self.save_file_btn) + file_layout.addWidget(self.load_file_btn) + file_layout.addStretch() + + self.save_file_btn.clicked.connect(self._on_save_to_file_clicked) + self.load_file_btn.clicked.connect(self._on_load_from_file_clicked) + + layout.addWidget(file_group) + + self._update_guard_label() + self._update_notifications_label() + + self.apply_theme() + + def apply_theme(self, theme: Optional[Literal["dark", "light"]] = None): + """ + Apply the theme + + Args: + theme (Optional[str]): Theme, either "dark", "light", or None. Defaults to None. + """ + if theme is None: + app = QApplication.instance() + theme = app.theme.theme # type: ignore + + for button in [ + self.add_btn, + self.edit_btn, + self.delete_btn, + self.move_down_btn, + self.move_up_btn, + self.copy_btn, + self.paste_btn, + self.run_btn, + self.abort_btn, + self.reset_btn, + self.notifications_btn, + self.guard_btn, + self.save_file_btn, + self.load_file_btn, + ]: + button.apply_theme() + + # ------------------------------------------------------------------ # + # persistence + # ------------------------------------------------------------------ # + def _load_or_seed_schedule(self) -> Schedule: + logger.info("load or seed schedule") + msg: VariableMessage | None = self.connector.get(self._endpoint) + if msg is not None: + logger.info(f"Got msg from endpoint: {msg.value}") + return Schedule.model_validate(msg.value) + + logger.info("No schedule found, create one now") + schedule = Schedule(schedule_name=self.schedule_name, items=[]) + self.connector.set_and_publish( + self._endpoint, VariableMessage(value=schedule.model_dump(mode="json")) + ) + return schedule + + def _persist_locked(self): + """Write `self.schedule` to Redis. Caller must hold `self._lock`.""" + self.connector.set_and_publish( + self._endpoint, VariableMessage(value=self.schedule.model_dump(mode="json")) + ) + + def _persist(self): + with self._lock: + self._persist_locked() + + @SafeSlot(dict, dict) + def _on_remote_update(self, msg_content: dict, _): + value = msg_content.get("value") if msg_content else None + if not value or value.get("schedule_name") != self.schedule_name: + return + with self._lock: + new_schedule = Schedule.model_validate(value) + if new_schedule.is_running: + # This instance's own execution task owns the current + # ScheduleItem objects (and will re-persist its own state + # again shortly); accepting a wholesale replacement here + # would orphan those objects. See module docstring. + logger.info( + f"Ignoring remote schedule update for {self.schedule_name} while a local run is in progress." + ) + return + self.schedule = new_schedule + self._apply_guard_settings() + self._refresh_ui() + + # ------------------------------------------------------------------ # + # reconciliation: figure out real state after being closed/reopened + # ------------------------------------------------------------------ # + def _reconcile_with_live_state(self): + """ + For any item last seen as RUNNING, ask BEC's own queue state - + not our local memory, which may be stale or from a previous + process - what actually happened to it while we were gone. + """ + queue_storage = self.client.queue.queue_storage + + # Force a fresh pull of the current queue snapshot from Redis + # rather than relying on pubsub messages we may have missed while + # this widget instance did not exist. + status_msg = self.connector.get(MessageEndpoints.scan_queue_status()) + if status_msg is not None: + queue_storage.update_with_status(status_msg) + + with self._lock: + self.schedule.is_running = False + for item in self.schedule.items: + if item.status != ScheduleItemStatus.RUNNING or not item.request_id: + continue + queue_item = queue_storage.find_queue_item_by_requestID(item.request_id) + if queue_item is not None and self._is_queue_item_active(queue_item): + # genuinely still in flight server-side; leave it + # RUNNING - run_schedule() will re-attach to it + # instead of resubmitting. + continue + item.status = ScheduleItemStatus.COMPLETED + item.finished_at = time.time() + self._persist_locked() + + @staticmethod + def _is_queue_item_active(queue_item) -> bool: + return str(getattr(queue_item, "status", "")).upper() in _ACTIVE_QUEUE_STATES + + # ------------------------------------------------------------------ # + # the "protected prefix" invariant (see module docstring) - the actual + # logic lives in schedule_logic.py, pure and independently testable; + # these thin wrappers just apply it to the current, lock-held state. + # ------------------------------------------------------------------ # + def _protected_prefix_length_locked(self) -> int: + return protected_prefix_length(self.schedule.items) + + def _index_of_locked(self, item_id: str | None) -> int | None: + return index_of(self.schedule.items, item_id) + + @SafeSlot() + def _refresh_ui(self): + with self._lock: + schedule_name = self.schedule.schedule_name + items = list(self.schedule.items) + notes = self.schedule.notes + + self.schedule_name_input.setText(schedule_name) + + self.list_widget.blockSignals(True) + self.list_widget.clear() + selected_row = None + for row, item in enumerate(items): + text = f"{item.command}" + if item.error: + text += f" ({item.error.strip().splitlines()[-1]})" + list_item = QListWidgetItem(text, self.list_widget) + list_item.setData(_ITEM_ID_ROLE, item.item_id) + if item.item_id == self._selected_item_id: + selected_row = row + self._update_icons() + if selected_row is not None: + self.list_widget.setCurrentRow(selected_row) + else: + self._selected_item_id = None + self.list_widget.blockSignals(False) + + # Don't stomp on notes the operator is actively typing (e.g. a + # remote update, or our own debounce timer firing right after + # local edits) - only resync when the field isn't focused and the + # text actually differs from what we already have. + if not self.notes_edit.hasFocus() and self.notes_edit.toPlainText() != notes: + self.notes_edit.blockSignals(True) + self.notes_edit.setPlainText(notes) + self.notes_edit.blockSignals(False) + + def _update_guard_label(self, *_): + if self._guard.enabled: + state = "OK" if self._guard.is_clear() else "PAUSED - waiting to resume" + value = self._guard.current_value + units = self._guard.units + value_text = f"{value:.{self._guard.prec}f}" if value is not None else "?" + self.guard_status_label.setText( + f"Enabled: {self._guard.device_name} = {value_text} {units} " + f"(pause < {self._guard.pause_below:g}, resume > {self._guard.resume_above:g}) [{state}]" + ) + else: + self.guard_status_label.setText("Disabled") + + def _update_buttons(self): + with self._lock: + items = list(self.schedule.items) + is_running = self.schedule.is_running + + idx = None + selected_item = None + + if self._selected_item_id is not None: + for i, item in enumerate(items): + if item.item_id == self._selected_item_id: + idx, selected_item = i, item + break + + can_edit_selected = ( + selected_item is not None and selected_item.status == ScheduleItemStatus.PENDING + ) + + self.run_btn.setEnabled(not is_running) + self.add_btn.setEnabled(True) + self.edit_btn.setEnabled(can_edit_selected) + self.delete_btn.setEnabled(can_edit_selected) + self.move_up_btn.setEnabled(can_edit_selected and idx not in (None, 0)) + self.move_down_btn.setEnabled( + can_edit_selected and idx is not None and idx < len(items) - 1 + ) + # Copy just captures the command text, so it's fine on any item + # (e.g. copying a COMPLETED item's command to reuse it later). + # Paste always goes through add_item(), which already clamps the + # insert position past the protected prefix - see _on_paste_clicked. + self.copy_btn.setEnabled(selected_item is not None) + self.paste_btn.setEnabled(self._clipboard_item is not None) + + def _update_icons(self): + items = list(self.schedule.items) + for row, item in enumerate(items): + icon_name, color = _ICON_MAP[item.status] + if self._selected_item_id == item.item_id: + color = "#FFFFFF" + icon = material_icon( + icon_name, size=(ICON_SIZE, ICON_SIZE), color=color, convert_to_pixmap=True + ) + list_item = self.list_widget.item(row) + list_item.setIcon(icon) + + @SafeSlot() + def _on_schedule_name_change(self): + name = self.schedule_name_input.text() + with self._lock: + self.schedule.schedule_name = name + self._persist_locked() + + @SafeSlot() + def _on_selection_changed(self, row: int): + item = self.list_widget.item(row) if row >= 0 else None + self._selected_item_id = item.data(_ITEM_ID_ROLE) if item is not None else None + self._update_icons() + self._update_buttons() + + @SafeSlot() + def _on_empty_space_clicked(self): + self.list_widget.setCurrentRow(-1) + self._refresh_ui() + + # ---- notes ---- # + def _on_notes_changed(self): + # Debounced: restart the timer on every keystroke, only persist + # once typing pauses, so we're not writing to Redis on every + # character. + self._notes_save_timer.start() + + def _persist_notes(self): + text = self.notes_edit.toPlainText() + with self._lock: + if self.schedule.notes == text: + return + self.schedule.notes = text + self._persist_locked() + + # ---- save / restore to file ---- # + @SafeSlot() + def _on_save_to_file_clicked(self): + with self._lock: + data = self.schedule.model_dump(mode="json") + default_name = "schedule.json" + + active_account = self.client.active_account + default_name = f"/sls/{self.beamline}/data/{active_account}/raw/{default_name}" + + path, _ = QFileDialog.getSaveFileName( + self, "Save schedule to file", default_name, "JSON files (*.json);;All files (*)" + ) + if not path: + return + try: + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + except OSError as exc: + QMessageBox.critical(self, "Save failed", f"Could not write to {path}:\n{exc}") + return + logger.info(f"Saved schedule '{self.schedule_name}' to {path}") + + @SafeSlot() + def _on_load_from_file_clicked(self): + with self._lock: + if self.schedule.is_running: + QMessageBox.warning( + self, + "Cannot load", + "Stop the running schedule (Abort) before loading a new one from file.", + ) + return + + start_folder = "" + + active_account = self.client.active_account + start_folder = f"/sls/{self.beamline}/data/{active_account}/raw" + + path, _ = QFileDialog.getOpenFileName( + self, "Load schedule from file", start_folder, "JSON files (*.json);;All files (*)" + ) + if not path: + return + + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + loaded = Schedule.model_validate(data) + except (OSError, json.JSONDecodeError, ValidationError) as exc: + QMessageBox.critical( + self, "Load failed", f"Could not load a schedule from {path}:\n{exc}" + ) + return + + with self._lock: + current_count = len(self.schedule.items) + confirm = QMessageBox.question( + self, + "Load schedule", + f"Replace the current {current_count} item(s) with {len(loaded.items)} " + f"item(s) from the file?\n\nAll items will be loaded as PENDING, regardless " + f"of their status when the file was saved.", + ) + if confirm != QMessageBox.StandardButton.Yes: + return + + # Loading a file restores a *definition* (what to run), not a past + # execution: every item comes back PENDING with a fresh item_id and + # its execution bookkeeping cleared, and is_running is always + # forced False here - regardless of what the file says - so we + # never end up "reconciling" against request/scan ids from a + # different session. schedule_name is deliberately left as this + # widget's own name, not overwritten by the file's. + new_items = [ + item.model_copy( + update={ + "item_id": str(uuid.uuid4()), + "status": ScheduleItemStatus.PENDING, + "request_id": None, + "scan_id": None, + "scan_number": None, + "error": None, + "started_at": None, + "finished_at": None, + } + ) + for item in loaded.items + ] + + with self._lock: + if self.schedule.is_running: + QMessageBox.warning( + self, "Cannot load", "The schedule started running while the dialog was open." + ) + return + self.schedule.schedule_name = loaded.schedule_name + self.schedule.items = new_items + self.schedule.notes = loaded.notes + self.schedule.guard = loaded.guard + self.schedule.notifications = loaded.notifications + self._persist_locked() + + self._guard.configure(self.schedule.guard) + self._selected_item_id = None + logger.info( + f"Loaded schedule from {path} ({len(new_items)} item(s)) into '{self.schedule_name}'" + ) + self._refresh_ui() + self._update_buttons() + + # ---- notifications / auto-pause settings ---- # + @SafeSlot() + def _on_notifications_clicked(self): + dialog = NotificationSettingsDialog(self.schedule.notifications, self.client, parent=self) + if dialog.exec_() != QDialog.Accepted: + return + with self._lock: + self.schedule.notifications = dialog.result_settings() + self._update_notifications_label() + self._persist_locked() + self._refresh_ui() + + def _update_notifications_label(self): + if self.schedule.notifications.enabled: + types = [] + if self.schedule.notifications.notify_scan: + types.append("scan") + if self.schedule.notifications.notify_move: + types.append("move") + if self.schedule.notifications.notify_rpc: + types.append("rpc") + match self.schedule.notifications.target: + case NotificationTarget.SIGNAL: + target = "Signal Messenger" + case NotificationTarget.TEAMS: + target = "Microsoft Teams" + case NotificationTarget.SCILOG: + target = "Scilog" + events = [] + for event in self.schedule.notifications.events: + events.append(event.value) + text = f"Enabled: Notify through {target}, for item type {types}, and events {events}" + else: + text = "Disabled" + self.notification_status_label.setText(text) + + @SafeSlot() + def _on_guard_clicked(self): + device_names = sorted(self.dev.keys()) + dialog = GuardSettingsDialog(self.schedule.guard, device_names, parent=self) + if dialog.exec_() != QDialog.Accepted: + return + with self._lock: + self.schedule.guard = dialog.result_settings() + self._persist_locked() + self._guard.configure(self.schedule.guard) + self._refresh_ui() + self._update_guard_label() + + def _apply_guard_settings(self): + """Re-subscribe the guard after `self.schedule.guard` changes (e.g. from a remote update).""" + self._guard.configure(self.schedule.guard) + self._update_guard_label() + + @SafeSlot(float) + def _on_guard_paused(self, value: float): + logger.info( + f"Auto-pause guard tripped for schedule {self.schedule_name} (value = {value}); interrupting the running scan, if any." + ) + if self._current_report is not None and self._current_item_kind == "scan": + self._guard_interrupt_requested = True + try: + self._current_report.cancel() + except Exception: # pylint: disable=broad-except + logger.exception("Failed to cancel the running scan for the auto-pause guard") + self.schedule_changed.emit() + self._update_guard_label() + + @SafeSlot(float) + def _on_guard_resumed(self, value: float): + logger.info( + f"Auto-pause guard cleared for schedule {self.schedule_name} (value = {value}); the execution loop will " + "resume the next scan item automatically." + ) + self.schedule_changed.emit() + + # ---- UI-triggered edit actions ---- # + @SafeSlot() + def _on_add_clicked(self): + dialog = ScheduleItemDialog(self.scans, self.dev, parent=self, client=self.client, beamline=self.beamline) + if dialog.exec_() != QDialog.Accepted: + return + result = dialog.result() + with self._lock: + idx = self._index_of_locked(self._selected_item_id) + insert_at = None if idx is None else idx + 1 + new_id = self.add_item( + result["command"], index=insert_at, kind=result["kind"], form_state=result["form_state"] + ) + self._selected_item_id = new_id + self._refresh_ui() + self._update_buttons() + + @SafeSlot() + def _on_edit_clicked(self, _): + if self._selected_item_id is None: + return + with self._lock: + idx = self._index_of_locked(self._selected_item_id) + item = self.schedule.items[idx] if idx is not None else None + initial = ( + {"kind": item.kind, "command": item.command, "form_state": item.form_state} + if item is not None + else None + ) + if initial is None: + return + + logger.info(f"Initial: {initial}") + + dialog = ScheduleItemDialog( + self.scans, self.dev, parent=self, initial=initial, client=self.client + ) + if dialog.exec_() != QDialog.Accepted: + return + result = dialog.result() + try: + self.edit_item( + self._selected_item_id, + result["command"], + kind=result["kind"], + form_state=result["form_state"], + ) + except RuntimeError as exc: + QMessageBox.warning(self, "Cannot edit item", str(exc)) + return + self._refresh_ui() + + @SafeSlot() + def _on_delete_clicked(self): + if self._selected_item_id is None: + return + row = self.list_widget.currentRow() + try: + self.delete_item(self._selected_item_id) + except RuntimeError as exc: + QMessageBox.warning(self, "Cannot delete item", str(exc)) + return + self._selected_item_id = None + self._on_selection_changed(max(0, row - 1)) + self._refresh_ui() + + @SafeSlot() + def _on_move_up_clicked(self): + if self._selected_item_id is None: + return + try: + self.move_item_up(self._selected_item_id) + except RuntimeError as exc: + QMessageBox.warning(self, "Cannot move item", str(exc)) + + @SafeSlot() + def _on_move_down_clicked(self): + if self._selected_item_id is None: + return + try: + self.move_item_down(self._selected_item_id) + except RuntimeError as exc: + QMessageBox.warning(self, "Cannot move item", str(exc)) + + @SafeSlot() + def _on_copy_clicked(self): + """ + Copy the selected item's command onto an internal clipboard. Only + the command/kind/form_state are captured (not status, timestamps, + request/scan ids, ...) since pasting always creates a fresh, + PENDING item - this is "copy the command", not "duplicate the + history". + """ + if self._selected_item_id is None: + return + with self._lock: + idx = self._index_of_locked(self._selected_item_id) + if idx is None: + return + item = self.schedule.items[idx] + self._clipboard_item = { + "command": item.command, + "kind": item.kind, + "form_state": dict(item.form_state) if item.form_state else None, + } + logger.info(f"Copied schedule item: {self._clipboard_item['command']}") + self._update_buttons() + + @SafeSlot() + def _on_paste_clicked(self): + """ + Paste the copied command as a new item, right after whatever is + currently selected (or at the end, if nothing is selected) - same + placement `_on_add_clicked` uses. This goes through `add_item()` + unchanged, so the existing protected-prefix clamping applies here + too: pasting can never land before an item that's already + running or finished. + """ + if self._clipboard_item is None: + return + with self._lock: + idx = self._index_of_locked(self._selected_item_id) + insert_at = None if idx is None else idx + 1 + new_id = self.add_item( + self._clipboard_item["command"], + index=insert_at, + kind=self._clipboard_item["kind"], + form_state=self._clipboard_item["form_state"], + ) + self._selected_item_id = new_id + self._refresh_ui() + self._update_buttons() + + # ------------------------------------------------------------------ # + # RPC-exposed actions (USER_ACCESS) + # ------------------------------------------------------------------ # + @SafeSlot() + def run_schedule(self): + """Run the schedule, continuing from wherever it last left off.""" + with self._lock: + if self.schedule.is_running: + return + repeat_aborted_item = False + if any(item.status == ScheduleItemStatus.ABORTED for item in self.schedule.items): + repeat_aborted_item = ( + QMessageBox.question( + self, + "Repeat aborted scan", + "Would you like to repeat the aborted scan?", + QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, + ) + == QMessageBox.StandardButton.Yes + ) + self._abort_requested = False + self.schedule.is_running = True + self._persist_locked() + self.schedule_changed.emit() + self.submit_task( + self._run_all, + repeat_aborted_item, + on_complete=self._on_run_finished, + on_failed=self._on_run_failed, + ) + + @SafeSlot() + def abort_schedule(self): + """Request that the schedule stop after the current item.""" + self._abort_requested = True + if self._current_report is not None: + try: + self._current_report.cancel() + except Exception: # pylint: disable=broad-except + logger.exception("Failed to cancel the currently running schedule item") + + @SafeSlot() + def reset_schedule(self): + """Clear all execution state and start the schedule over from item 1.""" + with self._lock: + if self.schedule.is_running: + logger.info("Cannot reset schedule, as schedule is running!") + return + for item in self.schedule.items: + logger.info(f"Reset schedule item {item}") + item.status = ScheduleItemStatus.PENDING + item.request_id = None + item.scan_id = None + item.scan_number = None + item.error = None + item.started_at = None + item.finished_at = None + self._persist_locked() + self.schedule_changed.emit() + + def get_status(self) -> dict: + """RPC-exposed: current schedule state, e.g. for another widget or a script.""" + with self._lock: + return self.schedule.model_dump(mode="json") + + # ------------------------------------------------------------------ # + # RPC-exposed schedule editing (add / edit / move / delete) + # + # All of these only ever touch the PENDING suffix of the schedule - + # see the "protected prefix" section of the module docstring. They can + # be called at any time, including while the schedule is running. + # ------------------------------------------------------------------ # + def add_item( + self, + command: str, + index: int | None = None, + kind: str = "custom", + form_state: dict | None = None, + ) -> str: + """ + RPC-exposed: insert a new, PENDING command into the schedule. + Safe to call while the schedule is running. + + Args: + command: command text, evaluated the same way as the existing + items (against `scans`/`dev`) once the schedule runs. + index: position to insert at (0 = first). Clamped so the item + can never land before something already running or + finished. Defaults to appending at the end. + kind/form_state: optional structured description of how + `command` was built (see `schedule_item.ScheduleItem`), + used to reopen the Edit dialog pre-filled. Leave as + defaults for a plain, hand-typed command. + + Returns: + The new item's item_id. + """ + command = command.strip() + if not command: + raise ValueError("command must not be empty") + + with self._lock: + protected = self._protected_prefix_length_locked() + item = ScheduleItem( + item_id=str(uuid.uuid4()), command=command, kind=kind, form_state=form_state + ) + if index is None or index >= len(self.schedule.items): + self.schedule.items.append(item) + else: + self.schedule.items.insert(max(index, protected), item) + self._persist_locked() + self.schedule_changed.emit() + return item.item_id + + def edit_item( + self, item_id: str, command: str, kind: str = "custom", form_state: dict | None = None + ): + """ + RPC-exposed: change the command of an item that has not started + yet. Raises `RuntimeError` for an item that is already + running/finished. + """ + command = command.strip() + if not command: + raise ValueError("command must not be empty") + + with self._lock: + idx = self._index_of_locked(item_id) + if idx is None: + return + item = self.schedule.items[idx] + if item.status != ScheduleItemStatus.PENDING: + raise RuntimeError("Only items that have not started yet (PENDING) can be edited.") + item.command = command + item.kind = kind + item.form_state = form_state + self._persist_locked() + self.schedule_changed.emit() + + def delete_item(self, item_id: str): + """ + RPC-exposed: remove an item that has not started yet. Raises + `RuntimeError` for an item that is already running/finished. + """ + with self._lock: + idx = self._index_of_locked(item_id) + if idx is None: + return + if self.schedule.items[idx].status != ScheduleItemStatus.PENDING: + raise RuntimeError("Only items that have not started yet (PENDING) can be deleted.") + del self.schedule.items[idx] + self._persist_locked() + self.schedule_changed.emit() + + def move_item(self, item_id: str, new_index: int): + """ + RPC-exposed: move an item that has not started yet to a new + position (0 = first, but never before something already + running/finished). Raises `RuntimeError` for an item that is + already running/finished. + """ + with self._lock: + idx = self._index_of_locked(item_id) + if idx is None: + return + if self.schedule.items[idx].status != ScheduleItemStatus.PENDING: + raise RuntimeError("Only items that have not started yet (PENDING) can be moved.") + protected = self._protected_prefix_length_locked() + item = self.schedule.items.pop(idx) + new_index = max(protected, min(new_index, len(self.schedule.items))) + self.schedule.items.insert(new_index, item) + self._persist_locked() + self.schedule_changed.emit() + + def move_item_up(self, item_id: str): + """RPC-exposed: swap an item with the one directly before it.""" + with self._lock: + idx = self._index_of_locked(item_id) + target = idx - 1 if idx is not None and idx > 0 else None + if target is not None: + self.move_item(item_id, target) + self._refresh_ui() + self._update_buttons() + + def move_item_down(self, item_id: str): + """RPC-exposed: swap an item with the one directly after it.""" + with self._lock: + idx = self._index_of_locked(item_id) + n = len(self.schedule.items) + target = idx + 1 if idx is not None and idx < n - 1 else None + if target is not None: + self.move_item(item_id, target) + self._refresh_ui() + self._update_buttons() + + # ------------------------------------------------------------------ # + # execution - runs in a background task (BECConnector.submit_task, + # backed by a QThreadPool worker thread; UI updates below go through + # the `schedule_changed` Qt signal so they are marshalled back onto + # the GUI thread instead of touching widgets directly) + # ------------------------------------------------------------------ # + def _pick_next_runnable_locked(self, repeat_aborted_item): + """Must be called while holding `self._lock`; see `schedule_logic.pick_next_runnable`.""" + return pick_next_runnable(self.schedule.items, repeat_aborted_item) + + @SafeSlot() + def _on_run_finished(self): + self._refresh_ui() + + @SafeSlot(str) + def _on_run_failed(self, error: str): + logger.error(f"Schedule execution task failed unexpectedly: {error}") + with self._lock: + self.schedule.is_running = False + self._persist_locked() + self._refresh_ui() + + def _emit_schedule_changed(self): + if self._closing.is_set(): + return + + self.schedule_changed.emit() + + def closeEvent(self, event): + """Ensure cleanup is invoked when the widget is closed via the window system.""" + self.cleanup() + super().closeEvent(event) + + def cleanup(self): + """Stop background execution loops without cancelling server-side scans.""" + self._closing.set() + + # Flush any not-yet-saved notes edit instead of losing it - the + # debounce timer won't get a chance to fire once we're closing. + if self._notes_save_timer.isActive(): + self._notes_save_timer.stop() + self._persist_notes() + + # Interrupt guard wait loops if any thread is blocked in self._guard.wait_until_clear() + self._guard.cleanup() + + super().cleanup() + + def _run_all(self, repeat_aborted_item): + notify_schedule_state( + self.client, self.schedule.notifications, self.schedule.schedule_name, "started" + ) + namespace = { + "scans": self.scans, + "dev": self.dev, + "digital_twin": self.digital_twin, + "auto_gain": self.auto_gain, + "np": np, + } + while not self._abort_requested and not self._closing.is_set(): + with self._lock: + next_item = self._pick_next_runnable_locked(repeat_aborted_item) + if next_item is None or next_item == "stop": + logger.info("Next item is None or stop") + break + if next_item.status == ScheduleItemStatus.RUNNING: + self._await_running_item(next_item) + else: + self._execute_item(next_item, namespace) + if self._abort_requested: + notify_schedule_state( + self.client, self.schedule.notifications, self.schedule.schedule_name, "aborted" + ) + elif self._closing.is_set(): + notify_schedule_state( + self.client, + self.schedule.notifications, + self.schedule.schedule_name, + "widget_closed", + ) + else: + notify_schedule_state( + self.client, self.schedule.notifications, self.schedule.schedule_name, "finished" + ) + with self._lock: + logger.info("in _run_all, set is_running to false") + self.schedule.is_running = False + self._persist_locked() + self._emit_schedule_changed() + + def _execute_item(self, item: ScheduleItem, namespace: dict): + if item.kind == "scan": + # Block here until the guard is clear, operator aborts, or widget closes + cleared = self._guard.wait_until_clear( + should_abort=lambda: self._abort_requested or self._closing.is_set() + ) + self._emit_schedule_changed() + if not cleared or self._closing.is_set(): + return + + with self._lock: + item.status = ScheduleItemStatus.RUNNING + item.error = None + item.started_at = time.time() + self._persist_locked() + self._emit_schedule_changed() + + try: + report = eval( + item.command, {"__builtins__": {}}, namespace + ) # noqa: S307 pylint: disable=eval-used + self._current_report = report + self._current_item_kind = item.kind + + request = getattr(report, "request", None) + with self._lock: + item.request_id = getattr(request, "requestID", None) + self._persist_locked() + + # RPC commands may not return a status, thus have no wait() + if item.kind != "custom": + report.wait() + + if self._closing.is_set(): + # The widget was closed mid-item; exit without changing status in Redis. + # Reconciliation logic will handle state on reopen. + return + + scan = getattr(report, "scan", None) + with self._lock: + item.scan_id = getattr(scan, "scan_id", None) if scan else None + item.scan_number = getattr(scan, "scan_number", None) if scan else None + item.status = ( + ScheduleItemStatus.ABORTED + if self._abort_requested + else ScheduleItemStatus.COMPLETED + ) + except Exception: # pylint: disable=broad-except + if self._closing.is_set(): + return + with self._lock: + if self._guard_interrupt_requested: + item.status = ScheduleItemStatus.PENDING + item.request_id = None + item.scan_id = None + item.scan_number = None + item.error = None + item.started_at = None + self._guard_interrupt_requested = False + elif self._abort_requested: + item.status = ScheduleItemStatus.ABORTED + else: + item.status = ScheduleItemStatus.FAILED + item.error = traceback.format_exc() + logger.error(f"Schedule item failed: {item.command}\n{item.error}") + finally: + if not self._closing.is_set(): + with self._lock: + if item.status != ScheduleItemStatus.PENDING: + item.finished_at = time.time() + self._persist_locked() + final_status = item.status + self._current_report = None + self._current_item_kind = None + self._emit_schedule_changed() + else: + final_status = None + + if final_status in ( + ScheduleItemStatus.COMPLETED, + ScheduleItemStatus.FAILED, + ScheduleItemStatus.ABORTED, + ): + notify_item_finished( + self.client, + item.kind, + item.command, + item.scan_number, + final_status=final_status, + settings=self.schedule.notifications, + error=item.error, + ) + + def _await_running_item(self, item: ScheduleItem): + queue_storage = self.client.queue.queue_storage + while not self._abort_requested and not self._closing.is_set(): + queue_item = queue_storage.find_queue_item_by_requestID(item.request_id) + if queue_item is None or not self._is_queue_item_active(queue_item): + with self._lock: + item.status = ScheduleItemStatus.COMPLETED + item.finished_at = time.time() + self._persist_locked() + self._emit_schedule_changed() + return + time.sleep(0.5) + + +if __name__ == "__main__": + app = QApplication(sys.argv) + dispatcher = BECDispatcher(gui_id="Scheduler") + win = Scheduler() + win.show() + sys.exit(app.exec_()) diff --git a/debye_bec/bec_widgets/widgets/scheduler/scheduler.pyproject b/debye_bec/bec_widgets/widgets/scheduler/scheduler.pyproject new file mode 100644 index 0000000..0b99313 --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/scheduler.pyproject @@ -0,0 +1 @@ +{'files': ['scheduler.py']} \ No newline at end of file diff --git a/debye_bec/bec_widgets/widgets/scheduler/scheduler_plugin.py b/debye_bec/bec_widgets/widgets/scheduler/scheduler_plugin.py new file mode 100644 index 0000000..50ce14a --- /dev/null +++ b/debye_bec/bec_widgets/widgets/scheduler/scheduler_plugin.py @@ -0,0 +1,57 @@ +# Copyright (C) 2022 The Qt Company Ltd. +# SPDX-License-Identifier: LicenseRef-Qt-Commercial OR BSD-3-Clause + +from bec_widgets.utils.bec_designer import designer_material_icon +from qtpy.QtDesigner import QDesignerCustomWidgetInterface +from qtpy.QtWidgets import QWidget + +from .scheduler import Scheduler + +DOM_XML = """ + + + + +""" + + +class SchedulerPlugin(QDesignerCustomWidgetInterface): # pragma: no cover + def __init__(self): + super().__init__() + self._form_editor = None + + def createWidget(self, parent): + if parent is None: + return QWidget() + t = Scheduler(parent) + return t + + def domXml(self): + return DOM_XML + + def group(self): + return "" + + def icon(self): + return designer_material_icon(Scheduler.ICON_NAME) + + def includeFile(self): + return "scheduler" + + def initialize(self, form_editor): + self._form_editor = form_editor + + def isContainer(self): + return False + + def isInitialized(self): + return self._form_editor is not None + + def name(self): + return "Scheduler" + + def toolTip(self): + return "Scheduler" + + def whatsThis(self): + return self.toolTip() diff --git a/debye_bec/device_configs/x01da_experimental_hutch.yaml b/debye_bec/device_configs/x01da_experimental_hutch.yaml index f2c241e..cec7109 100644 --- a/debye_bec/device_configs/x01da_experimental_hutch.yaml +++ b/debye_bec/device_configs/x01da_experimental_hutch.yaml @@ -1,3 +1,18 @@ + +####################################### +## Experimental Hutch Photon Shutter ## +####################################### + +eh_sh: + readoutPriority: baseline + description: Experimental Hutch Photon Shutter + deviceClass: debye_bec.devices.eh_shutter.EHPhotonShutter + deviceConfig: + prefix: "X01DA-" + onFailure: retry + enabled: true + softwareTrigger: false + ################################### ## Optical Table ## ################################### diff --git a/debye_bec/device_configs/x01da_frontend.yaml b/debye_bec/device_configs/x01da_frontend.yaml index 3a9edb7..77a4236 100644 --- a/debye_bec/device_configs/x01da_frontend.yaml +++ b/debye_bec/device_configs/x01da_frontend.yaml @@ -240,4 +240,4 @@ cm_xstripe: prefix: X01DA-FE-CM:XSTRIPE onFailure: retry enabled: true - softwareTrigger: false \ No newline at end of file + softwareTrigger: false diff --git a/debye_bec/device_configs/x01da_machine.yaml b/debye_bec/device_configs/x01da_machine.yaml index cd957dd..cbade1e 100644 --- a/debye_bec/device_configs/x01da_machine.yaml +++ b/debye_bec/device_configs/x01da_machine.yaml @@ -15,4 +15,18 @@ curr: onFailure: buffer enabled: true readOnly: true + softwareTrigger: false + +bl_status: + readoutPriority: baseline + description: BL status for machine + deviceClass: ophyd.EpicsSignal + deviceConfig: + auto_monitor: false + read_pv: AGEOP-BL:STATUS-X01DA + deviceTags: + - machine + onFailure: buffer + enabled: true + readOnly: false softwareTrigger: false \ No newline at end of file diff --git a/debye_bec/device_configs/x01da_optics.yaml b/debye_bec/device_configs/x01da_optics.yaml index a168e92..c96564d 100644 --- a/debye_bec/device_configs/x01da_optics.yaml +++ b/debye_bec/device_configs/x01da_optics.yaml @@ -1,4 +1,18 @@ +################################### +## Optics Photon Shutter ## +################################### + +op_sh: + readoutPriority: baseline + description: Optics Hutch Photon Shutter + deviceClass: debye_bec.devices.op_shutter.OPPhotonShutter + deviceConfig: + prefix: "X01DA-" + onFailure: retry + enabled: true + softwareTrigger: false + ################################### ## Monochromator ## ################################### @@ -408,4 +422,4 @@ sl2_gapy: softwareTrigger: false deviceTags: - optics - - slits \ No newline at end of file + - slits diff --git a/debye_bec/devices/absorber.py b/debye_bec/devices/absorber.py index 73da60a..29c3c9f 100644 --- a/debye_bec/devices/absorber.py +++ b/debye_bec/devices/absorber.py @@ -10,9 +10,13 @@ from ophyd import EpicsSignal, EpicsSignalRO from ophyd_devices import CompareStatus, DeviceStatus from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase +from .utils.bl_status_enum import BlStatus + if TYPE_CHECKING: from bec_lib.devicemanager import ScanInfo +TIMEOUT_FOR_PV = 5 + class AbsorberError(Exception): """Absorber specific exception""" @@ -38,6 +42,13 @@ class STATUS(int, enum.Enum): UNDEFINED = 14 +class BL_ENABLE(int, enum.Enum): + """Beamline enable""" + + DISABLE = 0 + ENABLE = 1 + + class Absorber(PSIDeviceBase): """Class for the Frontend Absorber""" @@ -55,6 +66,7 @@ class Absorber(PSIDeviceBase): string=True, doc="Absorber Status", ) + close4bl = Cpt(EpicsSignal, suffix="CLOSE4BL", kind="config", doc="Beamline enable") def __init__(self, *, name: str, prefix: str = "", scan_info: ScanInfo | None = None, **kwargs): super().__init__(name=name, prefix=prefix, scan_info=scan_info, **kwargs) @@ -63,12 +75,25 @@ class Absorber(PSIDeviceBase): # Wait for connection on all components, ensure IOC is connected self.wait_for_connection(all_signals=True, timeout=5) - def open(self) -> DeviceStatus | None: - """Open the Absorber""" + def open(self, force: bool = False) -> DeviceStatus | None: + """Open the Absorber + + Args: + force(bool): If needed, set bl status to enable and bl enable to ENABLE, defaults to False + + """ + if force and self.device_manager.devices.get("bl_status", None) is None: + raise AbsorberError("bl_status is not in device config, thus cannot use force = True") if self.status.get() == STATUS.CLOSED: + if force: + if self.device_manager.bl_status.get() == BlStatus.OFFLINE: + status = self.device_manager.bl_status.put(BlStatus.ATTENDED) + status.wait(timeout=TIMEOUT_FOR_PV) + if self.close4bl.get() == BL_ENABLE.DISABLE: + status = self.close4bl.set(BL_ENABLE.ENABLE) + status.wait(timeout=TIMEOUT_FOR_PV) self.request.put(1) - status_open = CompareStatus(self.status, STATUS.OPEN, timeout=self.timeout_for_move) - status = status_open + status = CompareStatus(self.status, STATUS.OPEN, timeout=self.timeout_for_move) return status else: return None @@ -77,8 +102,7 @@ class Absorber(PSIDeviceBase): """Close the Absorber""" if self.status.get() == STATUS.OPEN: self.request.put(1) - status_close = CompareStatus(self.status, STATUS.CLOSED, timeout=self.timeout_for_move) - status = status_close + status = CompareStatus(self.status, STATUS.CLOSED, timeout=self.timeout_for_move) return status else: return None diff --git a/debye_bec/devices/eh_shutter.py b/debye_bec/devices/eh_shutter.py new file mode 100644 index 0000000..27b141e --- /dev/null +++ b/debye_bec/devices/eh_shutter.py @@ -0,0 +1,85 @@ +"""Experimental Hutch Photon Shutter""" + +from __future__ import annotations + +import enum +from typing import TYPE_CHECKING + +from ophyd import Component as Cpt +from ophyd import EpicsSignal, EpicsSignalRO +from ophyd_devices import CompareStatus, DeviceStatus +from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase + +from .utils.bl_status_enum import BlStatus + +if TYPE_CHECKING: + from bec_lib.devicemanager import ScanInfo + +TIMEOUT_FOR_PV = 5 + + +class STATUS(int, enum.Enum): + """Shutter States""" + + NOT_CLOSED = 0 + CLOSED = 1 + + +class BL_ENABLE(int, enum.Enum): + """Beamline enable""" + + DISABLE = 0 + ENABLE = 1 + + +class EHPhotonShutter(PSIDeviceBase): + """Class for the Experimental Hutch Photon Shutter""" + + USER_ACCESS = ["open", "close"] + + request_open = Cpt( + EpicsSignal, suffix="EH1-PSYS:SH-A-OPEN-SET", kind="config", doc="Open Shutter" + ) + request_close = Cpt( + EpicsSignal, suffix="EH1-PSYS:SH-A-CLOSE-SET", kind="config", doc="Close Shutter" + ) + status = Cpt( + EpicsSignalRO, + suffix="EH1-PSYS:SH-A-CLOSE", + kind="normal", + auto_monitor=True, + doc="Shutter Status", + ) + status_string = Cpt( + EpicsSignalRO, + suffix="EH1-PSYS:SH-A-CLOSE", + kind="normal", + auto_monitor=True, + string=True, + doc="Shutter Status", + ) + + def __init__(self, *, name: str, prefix: str = "", scan_info: ScanInfo | None = None, **kwargs): + super().__init__(name=name, prefix=prefix, scan_info=scan_info, **kwargs) + + self.timeout_for_move = 10 + # Wait for connection on all components, ensure IOC is connected + self.wait_for_connection(all_signals=True, timeout=5) + + def open(self) -> DeviceStatus | None: + """Open the Shutter""" + if self.status.get() == STATUS.CLOSED: + self.request_open.put(1) + status = CompareStatus(self.status, STATUS.NOT_CLOSED, timeout=self.timeout_for_move) + return status + else: + return None + + def close(self) -> DeviceStatus | None: + """Close the Shutter""" + if self.status.get() == STATUS.NOT_CLOSED: + self.request_close.put(1) + status = CompareStatus(self.status, STATUS.CLOSED, timeout=self.timeout_for_move) + return status + else: + return None diff --git a/debye_bec/devices/ionization_chambers/ionization_chamber.py b/debye_bec/devices/ionization_chambers/ionization_chamber.py index 25b81d8..9d2bd2c 100644 --- a/debye_bec/devices/ionization_chambers/ionization_chamber.py +++ b/debye_bec/devices/ionization_chambers/ionization_chamber.py @@ -263,7 +263,7 @@ class IonizationChamber0(PSIDeviceBase): self.gmes.gas2_req.set(gas2).wait(timeout=3) self.gmes.conc2_req.set(conc2).wait(timeout=3) - status = TransitionStatus(self.gmes.status.get(), [0, 1]) + status = TransitionStatus(self.gmes.status, [0, 1]) self.cancel_on_stop(status) self.gmes.fill.put(1) if wait: @@ -373,6 +373,7 @@ class IonizationChamber2(IonizationChamber0): } hv_en = Dcpt(hv_en_signals) + class Pips(IonizationChamber0): """Pips, prefix should be 'X01DA-'.""" diff --git a/debye_bec/devices/nidaq/nidaq.py b/debye_bec/devices/nidaq/nidaq.py index be52d92..d3543bc 100644 --- a/debye_bec/devices/nidaq/nidaq.py +++ b/debye_bec/devices/nidaq/nidaq.py @@ -178,6 +178,9 @@ class NidaqControl(Device): heartbeat = Cpt(EpicsSignal, suffix="NIDAQ-Heartbeat", kind=Kind.config, auto_monitor=True) time_left = Cpt(EpicsSignalRO, suffix="NIDAQ-TimeLeft", kind=Kind.config, auto_monitor=True) + epics_mode = Cpt(EpicsSignal, suffix="NIDAQ-EpicsMode", kind=Kind.config, auto_monitor=True) + epics_max_reset = Cpt(EpicsSignal, suffix="NIDAQ-EpicsMaxReset", kind=Kind.config, auto_monitor=True) + ai_chans = Cpt(EpicsSignal, suffix="NIDAQ-AIChans", kind=Kind.config, auto_monitor=True) ci_chans = Cpt(EpicsSignal, suffix="NIDAQ-CIChans", kind=Kind.config, auto_monitor=True) di_chans = Cpt(EpicsSignal, suffix="NIDAQ-DIChans", kind=Kind.config, auto_monitor=True) diff --git a/debye_bec/devices/nidaq/nidaq_enums.py b/debye_bec/devices/nidaq/nidaq_enums.py index 14e4e5c..9b659ee 100644 --- a/debye_bec/devices/nidaq/nidaq_enums.py +++ b/debye_bec/devices/nidaq/nidaq_enums.py @@ -58,3 +58,9 @@ class EncoderFactors(int, enum.Enum): X1 = 4 X2 = 5 X4 = 6 + +class EpicsMode(int, enum.Enum): + """Mode when sending through EPICS""" + + MEAN = 0 + MAX = 0 diff --git a/debye_bec/devices/op_shutter.py b/debye_bec/devices/op_shutter.py new file mode 100644 index 0000000..4f963ad --- /dev/null +++ b/debye_bec/devices/op_shutter.py @@ -0,0 +1,106 @@ +"""Optics Photon Shutter""" + +from __future__ import annotations + +import enum +from typing import TYPE_CHECKING + +from ophyd import Component as Cpt +from ophyd import EpicsSignal, EpicsSignalRO +from ophyd_devices import CompareStatus, DeviceStatus +from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase + +from .utils.bl_status_enum import BlStatus + +if TYPE_CHECKING: + from bec_lib.devicemanager import ScanInfo + +TIMEOUT_FOR_PV = 5 + + +class OPPhotonShutterError(Exception): + """Shutter specific exception""" + + +class STATUS(int, enum.Enum): + """Shutter States""" + + NOT_CLOSED = 0 + CLOSED = 1 + + +class BL_ENABLE(int, enum.Enum): + """Beamline enable""" + + DISABLE = 0 + ENABLE = 1 + + +class OPPhotonShutter(PSIDeviceBase): + """Class for the Optics Photon Shutter""" + + USER_ACCESS = ["open", "close"] + + request_open = Cpt( + EpicsSignal, suffix="OP-PSYS:SH-A-OPEN-SET", kind="config", doc="Open Shutter" + ) + request_close = Cpt( + EpicsSignal, suffix="OP-PSYS:SH-A-CLOSE-SET", kind="config", doc="Close Shutter" + ) + status = Cpt( + EpicsSignalRO, + suffix="OP-PSYS:SH-A-CLOSE", + kind="normal", + auto_monitor=True, + doc="Shutter Status", + ) + status_string = Cpt( + EpicsSignalRO, + suffix="OP-PSYS:SH-A-CLOSE", + kind="normal", + auto_monitor=True, + string=True, + doc="Shutter Status", + ) + close4bl = Cpt(EpicsSignal, suffix="FE-BST1:CLOSE4BL", kind="config", doc="Beamline enable") + + def __init__(self, *, name: str, prefix: str = "", scan_info: ScanInfo | None = None, **kwargs): + super().__init__(name=name, prefix=prefix, scan_info=scan_info, **kwargs) + + self.timeout_for_move = 10 + # Wait for connection on all components, ensure IOC is connected + self.wait_for_connection(all_signals=True, timeout=5) + + def open(self, force: bool = False) -> DeviceStatus | None: + """Open the Shutter + + Args: + force(bool): If needed, set bl status to enable and bl enable to ENABLE, defaults to False + + """ + if force and self.device_manager.devices.get("bl_status", None) is None: + raise OPPhotonShutterError( + "bl_status is not in device config, thus cannot use force = True" + ) + if self.status.get() == STATUS.CLOSED: + if force: + if self.device_manager.bl_status.get() == BlStatus.OFFLINE: + status = self.device_manager.bl_status.put(BlStatus.ATTENDED) + status.wait(timeout=TIMEOUT_FOR_PV) + if self.close4bl.get() == BL_ENABLE.DISABLE: + status = self.close4bl.set(BL_ENABLE.ENABLE) + status.wait(timeout=TIMEOUT_FOR_PV) + self.request_open.put(1) + status = CompareStatus(self.status, STATUS.NOT_CLOSED, timeout=self.timeout_for_move) + return status + else: + return None + + def close(self) -> DeviceStatus | None: + """Close the Shutter""" + if self.status.get() == STATUS.NOT_CLOSED: + self.request_close.put(1) + status = CompareStatus(self.status, STATUS.CLOSED, timeout=self.timeout_for_move) + return status + else: + return None diff --git a/debye_bec/devices/reffoilchanger.py b/debye_bec/devices/reffoilchanger.py index b59970d..9bd5a9c 100644 --- a/debye_bec/devices/reffoilchanger.py +++ b/debye_bec/devices/reffoilchanger.py @@ -37,7 +37,7 @@ class OpMode(int, enum.Enum): class Reffoilchanger(PSIDeviceBase): """Class for the ES2 Reference Foil Changer""" - USER_ACCESS = ["insert"] + USER_ACCESS = ["get_all_foils", "insert"] inserted = Cpt( EpicsSignalRO, suffix="ES2-REF:TRY-FilterInserted", kind="config", doc="Inserted indicator" @@ -53,13 +53,21 @@ class Reffoilchanger(PSIDeviceBase): EpicsSignal, suffix="ES2-REF:SELN-FilterState-ENUM_RBV", kind="config", doc="Status" ) status_string = Cpt( - EpicsSignal, suffix="ES2-REF:SELN-FilterState-ENUM_RBV", kind="config", doc="Status", string=True + EpicsSignal, + suffix="ES2-REF:SELN-FilterState-ENUM_RBV", + kind="config", + doc="Status", + string=True, ) op_mode = Cpt( EpicsSignalWithRBV, suffix="ES2-REF:SELN-OpMode-ENUM", kind="config", doc="Status" ) op_mode_string = Cpt( - EpicsSignalWithRBV, suffix="ES2-REF:SELN-OpMode-ENUM", kind="config", doc="Status", string=True + EpicsSignalWithRBV, + suffix="ES2-REF:SELN-OpMode-ENUM", + kind="config", + doc="Status", + string=True, ) ref_set = Cpt(EpicsSignal, suffix="ES2-REF:SELN-SET", kind="config", doc="Requested reference") ref_rb = Cpt( @@ -149,6 +157,13 @@ class Reffoilchanger(PSIDeviceBase): self.foil38, ] + def get_all_foils(self) -> list[str]: + """Returns a list of strings of all available foils""" + foils_list = [] + for foil in self.foils: + foils_list.append(foil.get()) + return foils_list + def insert(self, ref: str, wait: bool = False) -> DeviceStatus: """Insert a reference diff --git a/debye_bec/devices/utils/bl_status_enum.py b/debye_bec/devices/utils/bl_status_enum.py new file mode 100644 index 0000000..2913745 --- /dev/null +++ b/debye_bec/devices/utils/bl_status_enum.py @@ -0,0 +1,8 @@ +import enum + +class BlStatus(str, enum.Enum): + """Beamline status enum""" + + OFFLINE = 0 + ATTENDED = 1 + REMOTE = 2