diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer.py b/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer.py index ac7379c..670706f 100644 --- a/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer.py +++ b/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer.py @@ -6,8 +6,7 @@ import os import subprocess import sys from datetime import datetime -from functools import partial -from typing import Literal, Optional, cast +from typing import Literal from bec_lib import bec_logger from bec_lib.endpoints import MessageEndpoints @@ -15,27 +14,12 @@ from bec_widgets.utils.bec_dispatcher import BECDispatcher from bec_widgets.utils.bec_widget import BECWidget from bec_widgets.utils.colors import apply_theme, get_accent_colors from bec_widgets.utils.error_popups import SafeSlot -from qtpy.QtCore import Qt # pylint: disable=E0611 -from qtpy.QtGui import QFont +from qtpy.QtWidgets import QApplication, QVBoxLayout, QWidget -# pylint: disable=E0611 -from qtpy.QtWidgets import ( - QApplication, - QComboBox, - QDoubleSpinBox, - QGroupBox, - QHBoxLayout, - QLabel, - QLayout, - QPushButton, - QVBoxLayout, - QWidget, -) - -from superxas_bec.bec_widgets.widgets.data_viewer.qt_widgets import TaggedListWidget -from superxas_bec.bec_widgets.widgets.data_viewer.viewer import HDF5Viewer +from .panels.input_panel import InputPanel +from .panels.scan_view import ScanViewer logger = bec_logger.logger @@ -54,20 +38,15 @@ class DataViewer(BECWidget, QWidget): super().__init__(parent=parent, theme_update=True, *arg, **kwargs) self.get_bec_shortcuts() - logger.info(f"Type of self.client: {type(self.client)}") - central = QWidget() self.root_layout = QVBoxLayout(central) - self.input = InputPanel() - self.viewer = HDF5Viewer() - + self.viewer = ScanViewer() self.root_layout.addWidget(self.input, 0) self.root_layout.addWidget(self.viewer, 1) self.setLayout(self.root_layout) self.setWindowTitle("Data Viewer") - # self.resize(1800, 800) self.history = [] self.bec_dispatcher.connect_slot(self.on_history_update, MessageEndpoints.scan_history()) @@ -75,14 +54,9 @@ class DataViewer(BECWidget, QWidget): self.current_row = 0 - logger.info(self.client.acl) - logger.info(self.client.active_account) - logger.info(self.client.proc) - logger.info(self.client.username) - self.input.scan_sel.currentItemChanged_connect(self.scan_sel_changed) - self.input.load_button.clicked_connect(self.load_dataset) - self.input.unload_button.clicked_connect(self.unload_all_datasets) + self.input.load_button.clicked_connect(self.load_scan) + self.input.unload_button.clicked_connect(self.unload_all_scans) self.input.open_fm_button.clicked_connect(self.open_in_file_manager) def apply_theme(self, theme: Literal["dark", "light"]): @@ -97,10 +71,12 @@ class DataViewer(BECWidget, QWidget): @SafeSlot() def scan_sel_changed(self, *_, **kwargs): + """Updates the current row value of the scan selection list""" self.current_row = kwargs["value"]().row() @SafeSlot() def open_in_file_manager(self, *_): + """Open the scan folder in the systems default file manager""" if len(self.history) > 0: scan = self.history[self.current_row] filepath = scan["file_components"][0].decode().rsplit("/", 1)[0] @@ -113,54 +89,93 @@ class DataViewer(BECWidget, QWidget): ) @SafeSlot() - def load_dataset( - self, *_ - ): # TODO: Check scan file components for combined xas/xrd scans. Is the Pilatus file in there as well? + def load_scan(self, *_): + """ + Loads a scan. Find all files within the scan folder, sort them and + then load the files in the scan view + """ if len(self.history) > 0: scan = self.history[self.current_row] - file = scan["file_components"][0] + b"_master." + scan["file_components"][1] - logger.info(file.decode()) - self.viewer.load_files([file.decode()]) + base_filepath = scan["file_components"][0].decode().rsplit("/", 1)[0] + filenames = [ + f + for f in os.listdir(base_filepath) + if os.path.isfile(os.path.join(base_filepath, f)) + ] + + def sort_priority(name): + if "master" in name: + return 0 + if name.endswith(".h5"): + return 1 + return 2 + + sorted_files = [ + f"{base_filepath}/{name}" for name in sorted(filenames, key=sort_priority) + ] + self.viewer.load_files(sorted_files) @SafeSlot() - def unload_all_datasets(self, *_): + def unload_all_scans(self, *_): + """Removes all scans from the scan view""" self.viewer.clear_files() - def duration_string(self, start: str, end: str) -> str: + def duration_formatted(self, start: str, end: str) -> str: + """ + Calculates the duration of a scan based on start end end time and + formats it as an easy readable string. + + Args: + start(str): start time in iso-format + end(str): end time in iso-format + + Returns: + str: Formatted duration, e.g. '1min 10s' or '1h 13min' + """ start_dt = datetime.fromisoformat(start) end_dt = datetime.fromisoformat(end) - seconds = abs(int((end_dt - start_dt).total_seconds())) days, remainder = divmod(seconds, 86400) hours, remainder = divmod(remainder, 3600) - minutes, _ = divmod(remainder, 60) + minutes, seconds = divmod(remainder, 60) parts = [] - if days: parts.append(f"{days}d") if hours: parts.append(f"{hours}h") if minutes: parts.append(f"{minutes}min") - + if not days and not hours and minutes < 10: + parts.append(f"{seconds}s") return " ".join(parts) if parts else "<1min" + def time_formatted(self, iso_time: str) -> str: + """ + Formates a time as an easy readable string. + + Args: + iso_time(str): Time in iso-format + + Returns: + str: Time formatted with format '%d.%m.%Y %H:%M', e.g. '14.01.1995 08:12' + """ + dt = datetime.fromisoformat(iso_time) + return dt.strftime("%d.%m.%Y %H:%M") + @SafeSlot() def on_history_update(self, *_): + """Updates the scan list based on the bec scan history.""" self.history = [] self.input.scan_sel.clear() - # Get the length of the scan history, which is 0 when the bec server was started - # and no scan has finished yet. Limit the history to the latest 20 scans. + if self.client.history is None: + return max_scans = min(len(self.client.history), MAX_HIST_LEN) for n in range(1, max_scans): # last scans, limited by MAX_HIST_LEN - # logger.info(self.client.history[-n].metadata["bec"]["status"]) - start_time = self.client.history[-n].metadata["start_time"] - end_time = self.client.history[-n].metadata["end_time"] - # logger.info(type(start_time)) - scan_data = self.client.history[-n].metadata["bec"] - # logger.info(scan_data) + start_time = self.client.history[-n].metadata["start_time"] # type: ignore + end_time = self.client.history[-n].metadata["end_time"] # type: ignore + scan_data = self.client.history[-n].metadata["bec"] # type: ignore scan_number = scan_data["scan_number"] scan_name = scan_data["scan_name"] comment = scan_data["metadata"]["user_metadata"]["comment"] @@ -186,142 +201,13 @@ class DataViewer(BECWidget, QWidget): tags.append((comment, get_accent_colors().warning.name())) if status == "closed": tags.append((status, get_accent_colors().success.name())) - elif status == "halted": + elif status == "halted" or status == "aborted": tags.append((status, get_accent_colors().emergency.name())) else: tags.append((status, "#656365")) - tags.append((self.duration_string(start_time, end_time), "#656365")) + tags.append((self.duration_formatted(start_time, end_time), "#656365")) + tags.append((self.time_formatted(start_time), "#656365")) self.input.scan_sel.addTaggedItem(label=str(scan_number), tags=tags) - # logger.info(f"Scan history: {self.history}") - - -class InputPanel(QWidget): - """Panel for scan selection of the data viewer widget""" - - def __init__(self, parent=None): - super().__init__(parent) - self._layout = QHBoxLayout(self) - # self._layout.setSizeConstraint(QLayout.SetFixedSize) # type: ignore - - # Scan selection - self.scan_sel = ListWidget("scan_sel", "Scan", ["Si", "Rh", "Pt"]) - self.load_button = Button(label_button="Load Dataset", enabled=True) - self.unload_button = Button(label_button="Unload all", enabled=True) - self.open_fm_button = Button(label_button="Open in File Manager", enabled=True) - - self._button_layout = QVBoxLayout() - self._button_layout.addWidget(self.load_button) - self._button_layout.addWidget(self.unload_button) - self._button_layout.addWidget(self.open_fm_button) - self._button_layout.addStretch() - - # Assemble complete scan selection group - self.input_group = Group( - "Scan selection", [self._button_layout, self.scan_sel], orientation="horizontal" - ) - - self._layout.addWidget(self.input_group) - # self._layout.addStretch() - - -class Group(QGroupBox): - def __init__(self, label, objs, orientation="vertical"): - super().__init__(label) - if orientation == "vertical": - self._layout = QVBoxLayout(self) # type: ignore - elif orientation == "horizontal": # assume horizontal - self._layout = QHBoxLayout(self) # type: ignore - else: - raise ValueError(f"Orientation {orientation} is not supported!") - for obj in objs: - if isinstance(obj, QWidget): - self._layout.addWidget(obj) # type: ignore - elif isinstance(obj, QLayout): - self._layout.addLayout(obj) - - -class ListWidget(QWidget): - def __init__(self, identifier="", label="", enums=[]): - super().__init__() - layout = QHBoxLayout(self) - layout.setContentsMargins(10, 0, 0, 0) - layout.setSpacing(0) - self.identifier = identifier - # self.label = QLabel(label) - # self.label.setFixedWidth(140) - # self.label.setContentsMargins(0, 0, 10, 0) - # self.label.setWordWrap(True) - # layout.addWidget(self.label) - self.value = TaggedListWidget() - # self.value.setFixedWidth(400) - # for entry in enums: - # self.value.addItem(entry) - layout.addWidget(self.value) - - def clear(self): - self.value.clear() - - def addTaggedItem(self, label, tags): - self.value.addTaggedItem(label, tags) - - def setCurrentIndex(self, text): - self.value.setCurrentIndex(text) - - # def currentIndex(self) -> int: - # return self.value.currentIndex() - - # def has_focus(self) -> bool: - # return QApplication.focusWidget() is self.value.view() - - def currentItemChanged_connect(self, func): - """Connect a function to the Enter/Return key press.""" - self.value.currentItemChanged.connect( - partial( - func, - identifier=self.identifier, - value_obj=self.value, - value=lambda: self.value.currentIndex(), - ) - ) - - def setDisabled(self, disable): - self.value.setDisabled(disable) - - -class Button(QWidget): - def __init__(self, label=None, label_button: str = "", enabled=False): - super().__init__() - layout = QHBoxLayout(self) - layout.setContentsMargins(10, 0, 0, 0) - layout.setSpacing(0) - if label is not None: - self.label = QLabel(label) - self.label.setFixedWidth(140) - layout.addWidget(self.label) - self.button = QPushButton(label_button) - if label is not None: - self.button.setFixedWidth(160) - self.enable_button(enabled) - layout.addWidget(self.button) - - def clicked_connect(self, func): - """Connect a function to the button press.""" - self.button.clicked.connect(func) - - def enable_button(self, enable: bool = False): - if enable: - self.button.setStyleSheet( - f"QPushButton {{background-color: {get_accent_colors().default.name()}; color: white;}}" - ) - self.button.setEnabled(True) - else: # disabled - self.button.setStyleSheet( - "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}" - ) - self.button.setDisabled(True) - - def setText(self, text): - self.button.setText(text) if __name__ == "__main__": diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer_plugin.py b/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer_plugin.py index f797743..0846e80 100644 --- a/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer_plugin.py +++ b/superxas_bec/bec_widgets/widgets/data_viewer/data_viewer_plugin.py @@ -5,7 +5,7 @@ from bec_widgets.utils.bec_designer import designer_material_icon from qtpy.QtDesigner import QDesignerCustomWidgetInterface from qtpy.QtWidgets import QWidget -from superxas_bec.bec_widgets.widgets.data_viewer.data_viewer import DataViewer +from .data_viewer import DataViewer DOM_XML = """ diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/loaders.py b/superxas_bec/bec_widgets/widgets/data_viewer/loaders.py new file mode 100644 index 0000000..5b260ab --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/data_viewer/loaders.py @@ -0,0 +1,229 @@ +""" +File-format loader (HDF5, images). +""" + +import os +from abc import ABC, abstractmethod +from typing import Iterator, Literal, Optional + +import h5py +import numpy as np + + +class NodeInfo: + """Describes a single node (group or dataset) inside a loaded file.""" + + __slots__ = ("name", "path", "kind", "dtype", "shape") + + def __init__( + self, + name: str, + path: str, + kind: Literal["group", "dataset"], + dtype: str = "", + shape: tuple[int, ...] = (), + ): + self.name = name + self.path = path + self.kind = kind + self.dtype = dtype # e.g. "float", "int", "str", … + self.shape = shape # empty tuple for scalars / groups + + +class BaseFileLoader(ABC): + """ + Abstract base class for file-format loaders. + + Subclass this to add support for a new format. Three things are required: + + 1. ``EXTENSIONS`` — tuple of lowercase extensions this loader handles, + e.g. ``(".h5", ".hdf5")``. + + 2. ``open(filepath)`` — open the file and keep any handles alive. + + 3. ``iter_nodes(path)`` — yield ``NodeInfo`` objects for the direct + children of *path* (depth-1 walk; the tree widget calls this + recursively as the user expands nodes). + + 4. ``read_dataset(path)`` — return the dataset at *path* as a + ``numpy.ndarray``. + + 5. ``close()`` — release any open file handles. + + 6. ``child_count(path)`` — return the number of direct children of a + group node (used for the status label; override if cheap to compute). + """ + + EXTENSIONS: tuple[str, ...] = () + + @abstractmethod + def open(self, filepath: str) -> None: ... + + @abstractmethod + def iter_nodes(self, path: str) -> Iterator[NodeInfo]: ... + + @abstractmethod + def read_dataset(self, path: str) -> np.ndarray: ... + + @abstractmethod + def close(self) -> None: ... + + def child_count(self, path: str) -> int: + return sum(1 for _ in self.iter_nodes(path)) + + +class HDF5Loader(BaseFileLoader): + """Loader for HDF5 / NeXus files (.h5, .hdf5, .nxs, .nx).""" + + EXTENSIONS = (".h5", ".hdf5", ".hdf", ".nxs", ".nx") + + def __init__(self): + self._file: Optional[h5py.File] = None + + def open(self, filepath: str) -> None: + self._file = h5py.File(filepath, "r") + + def close(self) -> None: + if self._file is not None: + self._file.close() + self._file = None + + def iter_nodes(self, path: str) -> Iterator[NodeInfo]: + assert self._file is not None, "File not open" + obj = self._file[path] if path != "/" else self._file + + if not isinstance(obj, h5py.Group): + return + + for key in obj.keys(): + try: + child = obj[key] + except Exception: + continue + + child_path = child.name # h5py always gives the absolute path + + if isinstance(child, h5py.Group): + yield NodeInfo(name=key, path=child_path, kind="group") + + elif isinstance(child, h5py.Dataset): + shape_tuple = child.shape + d = child.dtype + dtype = "unknown" + if np.issubdtype(d, np.integer): + dtype = "int" + if np.issubdtype(d, np.floating): + dtype = "float" + if np.issubdtype(d, np.complexfloating): + dtype = "complex" + if d.kind in ("S", "U", "O"): + dtype = "str" + + yield NodeInfo( + name=key, path=child_path, kind="dataset", dtype=dtype, shape=shape_tuple + ) + + def read_dataset(self, path: str) -> np.ndarray: + assert self._file is not None, "File not open" + return self._file[path][()] + + def child_count(self, path: str) -> int: + assert self._file is not None, "File not open" + obj = self._file[path] if path != "/" else self._file + return len(obj) if isinstance(obj, h5py.Group) else 0 + + +class ImageLoader(BaseFileLoader): + """ + Loader for raster image files. + + The file is treated as a single, flat dataset. ``iter_nodes`` yields one + leaf node. + + Requires: Pillow (``pip install Pillow``) + """ + + EXTENSIONS = ( + ".jpg", + ".jpeg", + ".png", + ".gif", + ".tiff", + ".tif", + ".bmp", + ".webp", + ".ico", + ".ppm", + ".pgm", + ".pbm", + ) + + def __init__(self): + self._filepath: Optional[str] = None + + def open(self, filepath: str) -> None: + # Validate that Pillow can open it; keep only the path. + try: + from PIL import Image as _PILImage # noqa: F401 — existence check + + _PILImage.open(filepath).verify() + except Exception as exc: + raise OSError(f"Cannot open image {filepath!r}: {exc}") from exc + self._filepath = filepath + + def close(self) -> None: + self._filepath = None + + def iter_nodes(self, path: str) -> Iterator[NodeInfo]: + """Images have no internal hierarchy — yield a single leaf node.""" + if path != "/" or self._filepath is None: + return + from PIL import Image as _PILImage + + with _PILImage.open(self._filepath) as img: + w, h = img.size + mode = img.mode # e.g. "RGB", "RGBA", "L", … + + name = os.path.basename(self._filepath) + yield NodeInfo(name=name, path="/image", kind="dataset", dtype=mode, shape=(h, w)) + + def read_dataset(self, path: str) -> np.ndarray: + """Return the image as a uint8 numpy array (H x W x C or H x W).""" + from PIL import Image as _PILImage + + with _PILImage.open(self._filepath) as img: # type: ignore[arg-type] + # Animated GIF → first frame only + if hasattr(img, "n_frames") and img.n_frames > 1: + img.seek(0) + return np.asarray(img) + + def child_count(self, path: str) -> int: + return 1 if path == "/" else 0 + + +class LoaderRegistry: + """Maps file extensions to loader classes.""" + + def __init__(self): + self._registry: dict[str, type[BaseFileLoader]] = {} + + def register(self, loader_cls: type[BaseFileLoader]) -> None: + """Register a loader class for all extensions it declares.""" + for ext in loader_cls.EXTENSIONS: + self._registry[ext.lower()] = loader_cls + + def get_loader(self, filepath: str) -> Optional[BaseFileLoader]: + """Return a fresh loader instance for *filepath*, or None if unsupported.""" + ext = os.path.splitext(filepath)[1].lower() + cls = self._registry.get(ext) + return cls() if cls is not None else None + + @property + def supported_extensions(self) -> list[str]: + return sorted(self._registry) + + +# Default global registry — pre-populated with built-in loaders. +registry = LoaderRegistry() +registry.register(HDF5Loader) +registry.register(ImageLoader) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/panels/__init__,py b/superxas_bec/bec_widgets/widgets/data_viewer/panels/__init__,py new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/panels/data_view.py b/superxas_bec/bec_widgets/widgets/data_viewer/panels/data_view.py new file mode 100644 index 0000000..e9b368c --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/data_viewer/panels/data_view.py @@ -0,0 +1,393 @@ +""" +Data viewer displaying the data +""" + +from typing import Literal, Optional + +import numpy as np +import pyqtgraph as pg +from bec_lib import bec_logger +from bec_widgets.utils.colors import Colors +from qtpy.QtCore import Qt + +# pylint: disable=E0611 +from qtpy.QtGui import QFont, QPixmap + +# pylint: disable=E0611 +from qtpy.QtWidgets import ( + QAbstractItemView, + QApplication, + QGroupBox, + QHBoxLayout, + QHeaderView, + QLabel, + QRadioButton, + QSizePolicy, + QSlider, + QTableWidget, + QTableWidgetItem, + QVBoxLayout, + QWidget, +) + +logger = bec_logger.logger + +MAX_ROWS = 2000 +MAX_COLS = 500 + + +class DataView(QWidget): + def __init__(self): + super().__init__() + self._layout = QVBoxLayout(self) + + header = QHBoxLayout() + self.path_label = QLabel("") + self.path_label.setObjectName("path_label") + self.info_label = QLabel("") + self.info_label.setObjectName("info_label") + header.addWidget(self.path_label, 1) + header.addWidget(self.info_label) + self._layout.addLayout(header) + + mode_box = QGroupBox("View mode") + mode_layout = QHBoxLayout(mode_box) + mode_layout.setContentsMargins(8, 4, 8, 4) + self.rb_auto = QRadioButton("Auto") + self.rb_plot = QRadioButton("Plot") + self.rb_image = QRadioButton("Image") + self.rb_table = QRadioButton("Table") + self.rb_auto.setChecked(True) + for rb in (self.rb_auto, self.rb_image, self.rb_plot, self.rb_table): + mode_layout.addWidget(rb) + rb.toggled.connect(self._on_mode_change) + mode_layout.addStretch() + self._layout.addWidget(mode_box) + + self.content = QWidget() + self.content_layout = QVBoxLayout(self.content) + self.content_layout.setContentsMargins(0, 0, 0, 0) + self._layout.addWidget(self.content, 1) + + self.plot_widget = None + self.image_widget = None + self._current_data = None + self.show_empty() + + 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 + + bg_color = pg.getConfigOption("background") + fg_color = pg.getConfigOption("foreground") + if self.plot_widget is not None: + n_curves = len(self.plot_widget.listDataItems()) + colors = Colors.golden_angle_color( + colormap="plasma", num=max(10, n_curves + 1), format="HEX" + ) + for idx, curve in enumerate(self.plot_widget.listDataItems()): + curve.setPen(pg.mkPen(color=colors[idx])) + # Background + self.plot_widget.setBackground(bg_color) + # Axes (tick marks, tick labels, axis line) + for axis in ["left", "bottom", "right", "top"]: + ax = self.plot_widget.getAxis(axis) + ax.setPen(pg.mkPen(color=fg_color)) + ax.setTextPen(pg.mkPen(color=fg_color)) + + if self.image_widget is not None: + self.image_widget.getView().setBackgroundColor(bg_color) + self.image_widget.ui.histogram.setBackground(bg_color) + + def _clear_stack(self): + while self.content_layout.count(): + item = self.content_layout.takeAt(0) + if item.widget(): + item.widget().deleteLater() + self.plot_widget = None + self.image_widget = None + + def show_empty(self): + """Empties the content area.""" + self._clear_stack() + empty_label = QLabel("No data selected") + empty_label.setObjectName("info_label") + empty_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.path_label.setText("") + self.info_label.setText("") + self.content_layout.addWidget(empty_label) + + def show_unsupported(self, path: str = "") -> None: + """Display a friendly 'not implemented' message for unknown file types.""" + self._clear_stack() + self._current_data = None + self.path_label.setText(path) + self.info_label.setText("") + lbl = QLabel("File type not supported") + lbl.setObjectName("info_label") + lbl.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.content_layout.addWidget(lbl) + + def _on_mode_change(self): + if self._current_data is not None: + self.display(self._current_data, self.path_label.text()) + + def _active_mode(self): + if self.rb_plot.isChecked(): + return "plot" + if self.rb_image.isChecked(): + return "image" + if self.rb_table.isChecked(): + return "table" + return "auto" + + def display(self, data, path: str = "") -> None: + """ + Render *data* in the panel. + + Parameters + ---------- + data: + A ``numpy.ndarray`` (or anything convertible to one). + path: + Human-readable label shown in the header. + """ + self._current_data = data + self.path_label.setText(path) + + if not isinstance(data, np.ndarray): + data = np.array(data) + + self.info_label.setText(f"shape {data.shape}, dtype {data.dtype}, {data.size} elements") + + mode = self._active_mode() + if mode == "auto": + if data.ndim <= 1 and data.size > 1: + mode = "plot" + elif data.ndim <= 4 and min(data.shape, default=0) > 1: + mode = "image" + else: + mode = "table" + + if mode == "plot": + self._show_plot_1d(data) + elif mode == "image": + self._show_image_2d(data) + else: + self._show_table(data) + + def _show_plot_1d(self, data): + self._clear_stack() + + is_2d = data.ndim == 2 + + if is_2d: + n_rows, _ = data.shape + row_data = data[0].astype(np.float32) + else: + row_data = data.reshape(-1).astype(np.float32) + + x = np.arange(row_data.size, dtype=np.float32) + + self.plot_widget = pg.PlotWidget() + plot_item = self.plot_widget.getPlotItem() + assert plot_item is not None, "PlotWidget has no PlotItem" + plot_item.showGrid(x=True, y=True, alpha=0.25) + self.plot_widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + + plot_item.setAutoVisible(y=False) # type: ignore[attr-defined] + + curve = pg.PlotDataItem( + x, row_data, pen=pg.mkPen(color="#2980b9", width=1.6), antialias=False + ) + self.plot_widget.addItem(curve) + + curve.setDownsampling(auto=True, method="peak") + curve.setClipToView(True) + curve.setSkipFiniteCheck(True) + + plot_item.enableAutoRange() # type: ignore[attr-defined] + + if is_2d: + slider = QSlider(Qt.Orientation.Vertical) + slider.setMinimum(0) + slider.setMaximum(n_rows - 1) + slider.setValue(0) + slider.setFixedWidth(32) + slider.setPageStep(1) + + current_row_label = QLabel("0") + current_row_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + current_row_label.setFixedWidth(32) + + max_row_label = QLabel(f"0:{n_rows-1}") + max_row_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + max_row_label.setFixedWidth(32) + + def on_row_changed(row): + current_row_label.setText(str(row)) + new_data = data[row].astype(np.float32) + new_x = np.arange(new_data.size, dtype=np.float32) + curve.setData(new_x, new_data) + plot_item.enableAutoRange() # type: ignore[attr-defined] + + slider.valueChanged.connect(on_row_changed) + + slider_col = QWidget() + slider_col.setFixedWidth(36) + col_layout = QVBoxLayout(slider_col) + col_layout.setContentsMargins(0, 0, 0, 0) + col_layout.setSpacing(2) + col_layout.addWidget(max_row_label) + col_layout.addWidget(slider) + col_layout.addWidget(current_row_label) + + container = QWidget() + h_layout = QHBoxLayout(container) + h_layout.setContentsMargins(0, 0, 0, 0) + h_layout.setSpacing(4) + h_layout.addWidget(slider_col) + h_layout.addWidget(self.plot_widget) + + self.content_layout.addWidget(container) + else: + self.content_layout.addWidget(self.plot_widget) + + self.apply_theme() + + def _show_image_2d(self, data): + self._clear_stack() + + stacked = False + n_images = 0 + rgb = False + img = data + if data.ndim == 3: + if data.shape[-1] in (3, 4): + rgb = True + else: + stacked = True + n_images = data.shape[0] + img = data[0, :] + elif data.ndim == 4: + if data.shape[-1] in (3, 4): + rgb = True + stacked = True + n_images = data.shape[0] + img = data[0, :] + + self.image_widget = pg.ImageView() + + self.image_widget.ui.roiBtn.hide() + self.image_widget.ui.menuBtn.hide() + + if not rgb: + self.image_widget.setColorMap(pg.colormap.get("inferno", source="matplotlib")) + + def set_image(img, autoLevels=True, autoHistogramRange=True): + if rgb: + self.image_widget.imageItem.setOpts(axisOrder="row-major") + self.image_widget.setImage(img) + else: + self.image_widget.setImage( + img.T, autoLevels=autoLevels, autoHistogramRange=autoHistogramRange + ) + + set_image(img) + self.image_widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + + if stacked: + slider = QSlider(Qt.Orientation.Vertical) + slider.setMinimum(0) + slider.setMaximum(n_images - 1) + slider.setValue(0) + slider.setFixedWidth(32) + slider.setPageStep(1) + + current_image_label = QLabel("0") + current_image_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + current_image_label.setFixedWidth(32) + + max_image_label = QLabel(f"0:{n_images-1}") + max_image_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + max_image_label.setFixedWidth(32) + + def on_image_changed(row): + current_image_label.setText(str(row)) + set_image(data[row, :], autoLevels=False, autoHistogramRange=False) + + slider.valueChanged.connect(on_image_changed) + + slider_col = QWidget() + slider_col.setFixedWidth(36) + col_layout = QVBoxLayout(slider_col) + col_layout.setContentsMargins(0, 0, 0, 0) + col_layout.setSpacing(2) + col_layout.addWidget(max_image_label) + col_layout.addWidget(slider) + col_layout.addWidget(current_image_label) + + container = QWidget() + h_layout = QHBoxLayout(container) + h_layout.setContentsMargins(0, 0, 0, 0) + h_layout.setSpacing(4) + h_layout.addWidget(slider_col) + h_layout.addWidget(self.image_widget) + + self.content_layout.addWidget(container) + else: + self.content_layout.addWidget(self.image_widget) + + self.apply_theme() + + def _show_table(self, data): + self._clear_stack() + + if data.ndim == 0: + flat = data.reshape(1, 1) + elif data.ndim == 1: + flat = data.reshape(-1, 1) + elif data.ndim == 2: + flat = data + else: + flat = data.reshape(-1, data.shape[-1]) + + rows, cols = flat.shape + show_rows = min(rows, MAX_ROWS) + show_cols = min(cols, MAX_COLS) + + if rows > MAX_ROWS or cols > MAX_COLS: + note = QLabel(f"⚠ Showing {show_rows}/{rows} rows x {show_cols}/{cols} columns") + note.setObjectName("info_label") + note.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.content_layout.addWidget(note) + + table = QTableWidget(show_rows, show_cols) + table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers) + table.setSelectionMode(QAbstractItemView.SelectionMode.ContiguousSelection) + table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Interactive) + + is_float = np.issubdtype(flat.dtype, np.floating) + is_complex = np.iscomplexobj(flat) + is_bytes = flat.dtype.kind == "S" + + for r in range(show_rows): + for c in range(show_cols): + val = flat[r, c] + txt = ( + f"{val:.6g}" + if is_float + else f"{val:.4g}" if is_complex else str(val.decode()) if is_bytes else str(val) + ) + cell = QTableWidgetItem(txt) + cell.setTextAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter) + table.setItem(r, c, cell) + + self.content_layout.addWidget(table) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/panels/input_panel.py b/superxas_bec/bec_widgets/widgets/data_viewer/panels/input_panel.py new file mode 100644 index 0000000..f249b94 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/data_viewer/panels/input_panel.py @@ -0,0 +1,32 @@ +# pylint: disable=E0611 +from qtpy.QtWidgets import QHBoxLayout, QVBoxLayout, QWidget + +# pylint: disable=E0402 +from ..widgets.qt_widgets import Button, Group, ListWidget + + +class InputPanel(QWidget): + """Panel for scan selection of the data viewer widget""" + + def __init__(self, parent=None): + super().__init__(parent) + self._layout = QHBoxLayout(self) + + # Scan selection + self.scan_sel = ListWidget("scan_sel") + self.load_button = Button(label_button="Load Dataset", enabled=True) + self.unload_button = Button(label_button="Unload all", enabled=True) + self.open_fm_button = Button(label_button="Open in File Manager", enabled=True) + + self._button_layout = QVBoxLayout() + self._button_layout.addWidget(self.load_button) + self._button_layout.addWidget(self.unload_button) + self._button_layout.addWidget(self.open_fm_button) + self._button_layout.addStretch() + + # Assemble complete scan selection group + self.input_group = Group( + "Scan selection", [self._button_layout, self.scan_sel], orientation="horizontal" + ) + + self._layout.addWidget(self.input_group) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/panels/scan_view.py b/superxas_bec/bec_widgets/widgets/data_viewer/panels/scan_view.py new file mode 100644 index 0000000..aa3ebe7 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/data_viewer/panels/scan_view.py @@ -0,0 +1,213 @@ +""" +Scan viewer. Displays files of one or more scans in a tree view +""" + +import os +from typing import Literal, Optional + +from bec_lib import bec_logger +from bec_qthemes import material_icon +from bec_widgets.utils.colors import get_accent_colors +from qtpy.QtCore import Qt + +# pylint: disable=E0611 +from qtpy.QtGui import QBrush, QColor + +# pylint: disable=E0611 +from qtpy.QtWidgets import ( + QHBoxLayout, + QHeaderView, + QMainWindow, + QSplitter, + QTreeWidget, + QTreeWidgetItem, + QVBoxLayout, + QWidget, +) + +# pylint: disable=E0402 +from ..loaders import BaseFileLoader, registry +from ..widgets.qt_widgets import Group +from .data_view import DataView + +logger = bec_logger.logger + +ICON_SIZE = 20 + + +class ScanViewer(QMainWindow): + """ + Generic scan viewer. Supports any format registered in *registry*. + + Args: + filepath(str): Optional path to open on startup. + """ + + def __init__(self, filepath: Optional[str] = None): + super().__init__() + self.registry = registry + + self._open_files: dict[str, tuple[BaseFileLoader, str]] = {} + + central = QWidget() + self.setCentralWidget(central) + root_layout = QHBoxLayout(central) + + splitter = QSplitter(Qt.Orientation.Horizontal) + splitter.setChildrenCollapsible(False) + + left_pane = QWidget() + left_layout = QVBoxLayout(left_pane) + + self.tree = QTreeWidget() + self.tree.setMinimumWidth(250) + self.tree.setHeaderLabels(["Name", "Type", "Shape"]) + self.tree.header().setStretchLastSection(False) + self.tree.header().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch) + self.tree.header().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) + self.tree.header().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents) + self.tree.itemClicked.connect(self._on_item_clicked) + + left_layout.addWidget(self.tree, 1) + + self.data_panel = DataView() + + splitter.addWidget(left_pane) + splitter.addWidget(self.data_panel) + + splitter.setStretchFactor(0, 1) + splitter.setStretchFactor(1, 3) + splitter.setSizes([300, 900]) + splitter.setHandleWidth(6) + splitter.setChildrenCollapsible(False) + + self.scan_view_group = Group("Scan view", [splitter]) + + root_layout.addWidget(self.scan_view_group) + + if filepath: + self.load_files([filepath]) + + def apply_theme(self, theme: Literal["dark", "light"]): + """ + Apply the theme + + Args: + theme (str): Theme, either "dark" or "light" + """ + self.data_panel.apply_theme(theme) + + def load_files(self, filepaths: list[str]) -> None: + """Open one or more files and add each as a top-level tree node.""" + for fp in filepaths: + if fp in self._open_files: + continue # already loaded + + loader = self.registry.get_loader(fp) + if loader is None: + supported = ", ".join(self.registry.supported_extensions) + logger.warning("No loader found for %r. Supported extensions: %s", fp, supported) + continue + + try: + loader.open(fp) + except Exception as exc: + logger.error("Failed to open %r: %s", fp, exc) + continue + + display_name = os.path.basename(fp) + self._open_files[fp] = (loader, display_name) + self._add_file_to_tree(fp, loader, display_name) + + def clear_files(self) -> None: + """Close all open files and reset the tree.""" + for loader, _ in self._open_files.values(): + loader.close() + self._open_files.clear() + self.tree.clear() + self.data_panel.show_empty() + + def closeEvent(self, event): + """Close all""" + for loader, _ in self._open_files.values(): + loader.close() + self._open_files.clear() + super().closeEvent(event) + + def _add_file_to_tree(self, filepath: str, loader: BaseFileLoader, display_name: str) -> None: + """Add a single file as a new top-level node in the tree.""" + dataset_icon = material_icon( + "dataset", size=(ICON_SIZE, ICON_SIZE), color=get_accent_colors().default.name() + ) + root_item = QTreeWidgetItem(self.tree, [display_name, "Group", ""]) + root_item.setIcon(0, dataset_icon) + root_item.setData(0, Qt.ItemDataRole.UserRole, "/") # path + root_item.setData(0, Qt.ItemDataRole.UserRole + 1, "group") # kind + root_item.setData(0, Qt.ItemDataRole.UserRole + 2, filepath) # file key + + self._populate_tree(root_item, loader, "/") + self.tree.addTopLevelItem(root_item) + + # Expand first 2 levels by default + self.tree.expandItem(root_item) + for i in range(root_item.childCount()): + self.tree.expandItem(root_item.child(i)) + + self.tree.setCurrentItem(root_item) + + def _populate_tree( + self, parent_item: QTreeWidgetItem, loader: BaseFileLoader, path: str + ) -> None: + folder_icon = material_icon("folder", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") + vector_icon = material_icon("show_chart", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") + array_icon = material_icon( + "stacked_line_chart", size=(ICON_SIZE, ICON_SIZE), color="#2980b9" + ) + scalar_icon = material_icon("point_scan", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") + str_icon = material_icon("text_snippet", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") + + for node in loader.iter_nodes(path): + shape_str = "x".join(str(s) for s in node.shape) if node.shape else "scalar" + + if node.kind == "group": + item = QTreeWidgetItem(parent_item, [node.name, "Group", ""]) + item.setIcon(0, folder_icon) + item.setData(0, Qt.ItemDataRole.UserRole, node.path) + item.setData(0, Qt.ItemDataRole.UserRole + 1, "group") + self._populate_tree(item, loader, node.path) + + else: # dataset + if shape_str == "scalar": + icon = str_icon if node.dtype == "str" else scalar_icon + elif "x" in shape_str: + icon = array_icon + else: + icon = vector_icon + + item = QTreeWidgetItem(parent_item, [node.name, node.dtype, shape_str]) + item.setIcon(0, icon) + item.setData(0, Qt.ItemDataRole.UserRole, node.path) + item.setData(0, Qt.ItemDataRole.UserRole + 1, "dataset") + item.setForeground(1, QBrush(QColor(get_accent_colors().success.name()))) + item.setForeground(2, QBrush(QColor("#656365"))) + + def _get_filepath_for_item(self, item: QTreeWidgetItem) -> str: + """Walk up the tree to find the filepath stored on the root node.""" + node = item + while node.parent(): + node = node.parent() + return node.data(0, Qt.ItemDataRole.UserRole + 2) + + def _on_item_clicked(self, item: QTreeWidgetItem, _col: int) -> None: + path = item.data(0, Qt.ItemDataRole.UserRole) + kind = item.data(0, Qt.ItemDataRole.UserRole + 1) + filepath = self._get_filepath_for_item(item) + + if not path or not filepath: + return + + loader, _ = self._open_files[filepath] + + if kind == "dataset": + data = loader.read_dataset(path) + self.data_panel.display(data, path) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/register_data_viewer.py b/superxas_bec/bec_widgets/widgets/data_viewer/register_data_viewer.py index ba6a4a3..10c5154 100644 --- a/superxas_bec/bec_widgets/widgets/data_viewer/register_data_viewer.py +++ b/superxas_bec/bec_widgets/widgets/data_viewer/register_data_viewer.py @@ -6,7 +6,7 @@ def main(): # pragma: no cover return from PySide6.QtDesigner import QPyDesignerCustomWidgetCollection - from superxas_bec.bec_widgets.widgets.data_viewer.data_viewer_plugin import DataViewerPlugin + from .data_viewer_plugin import DataViewerPlugin QPyDesignerCustomWidgetCollection.addCustomWidget(DataViewerPlugin()) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/viewer.py b/superxas_bec/bec_widgets/widgets/data_viewer/viewer.py deleted file mode 100644 index 96ce3d5..0000000 --- a/superxas_bec/bec_widgets/widgets/data_viewer/viewer.py +++ /dev/null @@ -1,514 +0,0 @@ -""" -HDF5 Viewer — qtpy + pyqtgraph + h5py -""" - -from typing import Literal, Optional, cast - -import h5py -import numpy as np -import pyqtgraph as pg -from bec_lib import bec_logger -from bec_qthemes import material_icon -from bec_widgets.utils.colors import Colors, get_accent_colors -from qtpy.QtCore import Qt - -# pylint: disable=E0611 -from qtpy.QtGui import QBrush, QColor, QFont - -# pylint: disable=E0611 -from qtpy.QtWidgets import ( - QAbstractItemView, - QApplication, - QGroupBox, - QHBoxLayout, - QHeaderView, - QLabel, - QMainWindow, - QRadioButton, - QSizePolicy, - QSlider, - QSplitter, - QTableWidget, - QTableWidgetItem, - QTreeWidget, - QTreeWidgetItem, - QVBoxLayout, - QWidget, -) - -logger = bec_logger.logger - -ICON_SIZE = 20 - - -# ── Data / Plot panel ────────────────────────────────────────────────────── -class DataPanel(QWidget): - def __init__(self): - super().__init__() - self._layout = QVBoxLayout(self) - - # Header - hdr = QHBoxLayout() - self.path_label = QLabel("Select a dataset from the tree") - self.path_label.setObjectName("path_label") - self.info_label = QLabel("") - self.info_label.setObjectName("info_label") - hdr.addWidget(self.path_label, 1) - hdr.addWidget(self.info_label) - self._layout.addLayout(hdr) - - # View-mode selector - mode_box = QGroupBox("View mode") - mode_layout = QHBoxLayout(mode_box) - mode_layout.setContentsMargins(8, 4, 8, 4) - self.rb_auto = QRadioButton("Auto") - self.rb_plot = QRadioButton("Plot") - self.rb_image = QRadioButton("Image") - self.rb_table = QRadioButton("Table") - self.rb_auto.setChecked(True) - for rb in (self.rb_auto, self.rb_plot, self.rb_image, self.rb_table): - mode_layout.addWidget(rb) - rb.toggled.connect(self._on_mode_change) - mode_layout.addStretch() - self._layout.addWidget(mode_box) - - # Content stack - self.stack = QWidget() - self.stack_layout = QVBoxLayout(self.stack) - self.stack_layout.setContentsMargins(0, 0, 0, 0) - self._layout.addWidget(self.stack, 1) - - self.plot_widget = None - self.image_widget = None - - self._current_data = None - self.show_empty() - - # ── helpers ─────────────────────────────────────────────────────────── - def _clear_stack(self): - while self.stack_layout.count(): - item = self.stack_layout.takeAt(0) - if item.widget(): - item.widget().deleteLater() - self.plot_widget = None - self.image_widget = None - - def show_empty(self): - self._clear_stack() - lbl = QLabel("No data selected") - lbl.setObjectName("info_label") - lbl.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.stack_layout.addWidget(lbl) - - def _on_mode_change(self): - if self._current_data is not None: - self.display(self._current_data, self.path_label.text()) - - def _active_mode(self): - if self.rb_plot.isChecked(): - return "plot" - if self.rb_image.isChecked(): - return "image" - if self.rb_table.isChecked(): - return "table" - return "auto" - - # ── public ──────────────────────────────────────────────────────────── - def display(self, data, path=""): - self._current_data = data - self.path_label.setText(path) - - if not isinstance(data, np.ndarray): - data = np.array(data) - - self.info_label.setText( - f"shape {data.shape} · dtype {data.dtype} · {data.size:,} elements" - ) - - mode = self._active_mode() - if mode == "auto": - if data.ndim <= 1 and data.size > 1: - mode = "plot" - elif data.ndim == 2 and min(data.shape) > 1: - mode = "image" - else: - mode = "table" - - if mode == "plot": - self._show_plot_1d(data) - elif mode == "image": - self._show_image_2d(data) - else: - self._show_table(data) - - # ── 1-D line plot ────────────────────────────────────────────────────── - def _show_plot_1d(self, data): - self._clear_stack() - - is_2d = data.ndim == 2 - - if is_2d: - n_rows, n_cols = data.shape - row_data = data[0].astype(np.float32) - else: - row_data = data.reshape(-1).astype(np.float32) - - x = np.arange(row_data.size, dtype=np.float32) - - self.plot_widget = pg.PlotWidget() - plot_item = self.plot_widget.getPlotItem() - assert plot_item is not None, "PlotWidget has no PlotItem" - plot_item.showGrid(x=True, y=True, alpha=0.25) - self.plot_widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - - plot_item.setAutoVisible(y=False) # type: ignore[attr-defined] - - curve = pg.PlotDataItem( - x, row_data, pen=pg.mkPen(color="#2980b9", width=1.6), antialias=False - ) - self.plot_widget.addItem(curve) - - curve.setDownsampling(auto=True, method="peak") - curve.setClipToView(True) - curve.setSkipFiniteCheck(True) - - plot_item.enableAutoRange() # type: ignore[attr-defined] - - if is_2d: - # --- Slider --- - slider = QSlider(Qt.Orientation.Vertical) - slider.setMinimum(0) - slider.setMaximum(n_rows - 1) - slider.setValue(0) - slider.setFixedWidth(32) - slider.setPageStep(1) - - row_label = QLabel("0") - row_label.setAlignment(Qt.AlignmentFlag.AlignCenter) - row_label.setFixedWidth(32) - - def on_row_changed(row): - row_label.setText(str(row)) - new_data = data[row].astype(np.float32) - new_x = np.arange(new_data.size, dtype=np.float32) - curve.setData(new_x, new_data) - plot_item.enableAutoRange() # type: ignore[attr-defined] - - slider.valueChanged.connect(on_row_changed) - - slider_col = QWidget() - slider_col.setFixedWidth(36) - col_layout = QVBoxLayout(slider_col) - col_layout.setContentsMargins(0, 0, 0, 0) - col_layout.setSpacing(2) - col_layout.addWidget(row_label) - col_layout.addWidget(slider) - - container = QWidget() - h_layout = QHBoxLayout(container) - h_layout.setContentsMargins(0, 0, 0, 0) - h_layout.setSpacing(4) - h_layout.addWidget(slider_col) - h_layout.addWidget(self.plot_widget) - - self.stack_layout.addWidget(container) - else: - self.stack_layout.addWidget(self.plot_widget) - - 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 - - bg_color = pg.getConfigOption("background") - fg_color = pg.getConfigOption("foreground") - if self.plot_widget is not None: - n_curves = len(self.plot_widget.listDataItems()) - colors = Colors.golden_angle_color( - colormap="plasma", num=max(10, n_curves + 1), format="HEX" - ) - for idx, curve in enumerate(self.plot_widget.listDataItems()): - curve.setPen(pg.mkPen(color=colors[idx])) - # Background - self.plot_widget.setBackground(bg_color) - # Axes (tick marks, tick labels, axis line) - for axis in ["left", "bottom", "right", "top"]: - ax = self.plot_widget.getAxis(axis) - ax.setPen(pg.mkPen(color=fg_color)) - ax.setTextPen(pg.mkPen(color=fg_color)) - - if self.image_widget is not None: - self.image_widget.getView().setBackgroundColor(bg_color) - self.image_widget.ui.histogram.setBackground(bg_color) - - # ── 2-D image ────────────────────────────────────────────────────────── - def _show_image_2d(self, data): - self._clear_stack() - - squeezed = np.squeeze(data) - if squeezed.ndim > 2: - squeezed = squeezed.reshape(-1, squeezed.shape[-1]) - - # complex → magnitude - img_data = np.abs(squeezed) if np.iscomplexobj(squeezed) else squeezed.astype(float) - - # ImageView gives us colorbar + histogram + zoom for free - self.image_widget = pg.ImageView() - self.image_widget.ui.roiBtn.hide() - self.image_widget.ui.menuBtn.hide() - - # Use 'inferno'-like LUT - self.image_widget.setColorMap(pg.colormap.get("inferno", source="matplotlib")) - - # pyqtgraph ImageView expects (cols, rows) — transpose so row 0 is at top - self.image_widget.setImage(img_data.T, autoLevels=True, autoHistogramRange=True) - self.image_widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) - - self.stack_layout.addWidget(self.image_widget) - - self.apply_theme() - - # ── Table ────────────────────────────────────────────────────────────── - def _show_table(self, data): - self._clear_stack() - - MAX_ROWS, MAX_COLS = 2000, 500 - - if data.ndim == 0: - flat = data.reshape(1, 1) - elif data.ndim == 1: - flat = data.reshape(-1, 1) - elif data.ndim == 2: - flat = data - else: - flat = data.reshape(-1, data.shape[-1]) - - rows, cols = flat.shape - show_rows = min(rows, MAX_ROWS) - show_cols = min(cols, MAX_COLS) - - if rows > MAX_ROWS or cols > MAX_COLS: - note = QLabel(f"⚠ Showing {show_rows}/{rows} rows × {show_cols}/{cols} cols") - note.setObjectName("info_label") - note.setAlignment(Qt.AlignmentFlag.AlignCenter) - self.stack_layout.addWidget(note) - - table = QTableWidget(show_rows, show_cols) - table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers) - table.setSelectionMode(QAbstractItemView.SelectionMode.ContiguousSelection) - table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Interactive) - table.horizontalHeader().setDefaultSectionSize(90) - table.verticalHeader().setDefaultSectionSize(22) - table.setFont(QFont("JetBrains Mono, Consolas, monospace", 10)) - - is_float = np.issubdtype(flat.dtype, np.floating) - is_complex = np.iscomplexobj(flat) - is_bytes = flat.dtype.kind == "S" - - for r in range(show_rows): - for c in range(show_cols): - val = flat[r, c] - txt = ( - f"{val:.6g}" - if is_float - else f"{val:.4g}" if is_complex else str(val.decode()) if is_bytes else str(val) - ) - cell = QTableWidgetItem(txt) - cell.setTextAlignment(Qt.AlignmentFlag.AlignRight | Qt.AlignmentFlag.AlignVCenter) - table.setItem(r, c, cell) - - self.stack_layout.addWidget(table) - - -# ── Main window ──────────────────────────────────────────────────────────── -class HDF5Viewer(QMainWindow): - def __init__(self, filepath=None): - super().__init__() - self.h5files = {} # filepath -> h5py.File - self._build_ui() - if filepath: - self.load_files([filepath]) - - def apply_theme(self, theme: Literal["dark", "light"]): - """ - Apply the theme - - Args: - theme (str): Theme, either "dark" or "light" - """ - self.data_panel.apply_theme(theme) - - def load_files(self, filepaths: list[str]): - """Open one or more HDF5 files and add each as a top-level tree node.""" - for f in filepaths: - if f in self.h5files: - continue # already loaded - self.h5files[f] = h5py.File(f, "r") - self._add_file_to_tree(f) - - def _build_ui(self): - central = QWidget() - self.setCentralWidget(central) - - root_layout = QHBoxLayout(central) - - splitter = QSplitter(Qt.Orientation.Horizontal) - splitter.setChildrenCollapsible(False) - - # ── Left pane ── - left_pane = QWidget() - left_layout = QVBoxLayout(left_pane) - - self.tree = QTreeWidget() - self.tree.setHeaderLabels(["Name", "Type", "Shape"]) - self.tree.header().setStretchLastSection(False) - self.tree.header().setSectionResizeMode(0, QHeaderView.ResizeMode.Stretch) - self.tree.header().setSectionResizeMode(1, QHeaderView.ResizeMode.ResizeToContents) - self.tree.header().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents) - self.tree.itemClicked.connect(self._on_item_clicked) - - left_layout.addWidget(self.tree, 1) - - # ── Right pane ── - self.data_panel = DataPanel() - - splitter.addWidget(left_pane) - splitter.addWidget(self.data_panel) - - splitter.setStretchFactor(0, 1) - splitter.setStretchFactor(1, 3) - splitter.setSizes([300, 900]) - splitter.setHandleWidth(6) - splitter.setChildrenCollapsible(False) - - root_layout.addWidget(splitter) - - def _add_file_to_tree(self, filepath: str): - """Add a single file as a new top-level node in the tree.""" - h5file = self.h5files[filepath] - filename = filepath.split("/")[-1] - - dataset_icon = material_icon( - "dataset", size=(ICON_SIZE, ICON_SIZE), color=get_accent_colors().default.name() - ) - root_item = QTreeWidgetItem(self.tree, [filename, "Group", ""]) - root_item.setIcon(0, dataset_icon) - root_item.setData(0, Qt.ItemDataRole.UserRole, "/") - root_item.setData(0, Qt.ItemDataRole.UserRole + 1, "group") - root_item.setData(0, Qt.ItemDataRole.UserRole + 2, filepath) # so clicks know which file - - self.populate_tree(root_item, h5file) - self.tree.addTopLevelItem(root_item) - - # Expand first 2 levels - self.tree.expandItem(root_item) - for i in range(root_item.childCount()): - self.tree.expandItem(root_item.child(i)) - - self.tree.setCurrentItem(root_item) - - def populate_tree(self, parent_item, h5_obj): - folder_icon = material_icon("folder", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") - vector_icon = material_icon("show_chart", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") - array_icon = material_icon( - "stacked_line_chart", size=(ICON_SIZE, ICON_SIZE), color="#2980b9" - ) - scalar_icon = material_icon("point_scan", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") - str_icon = material_icon("text_snippet", size=(ICON_SIZE, ICON_SIZE), color="#2980b9") - - if not isinstance(h5_obj, h5py.Group): - return - - for key in h5_obj.keys(): - try: - child = h5_obj[key] - except Exception: - continue - - if isinstance(child, h5py.Group): - item = QTreeWidgetItem(parent_item, [key, "Group", ""]) - item.setIcon(0, folder_icon) - - item.setData(0, Qt.ItemDataRole.UserRole, child.name) - item.setData(0, Qt.ItemDataRole.UserRole + 1, "group") - - self.populate_tree(item, child) - - elif isinstance(child, h5py.Dataset): - shape_str = "x".join(str(s) for s in child.shape) or "scalar" - - d = child.dtype - dtype = "unknown" - if np.issubdtype(d, np.integer): - dtype = "int" - if np.issubdtype(d, np.floating): - dtype = "float" - if np.issubdtype(d, np.complexfloating): - dtype = "complex" - if d.kind in ("S", "U", "O"): - dtype = "str" - - if shape_str == "scalar": - if dtype == "str": - icon = str_icon - else: - icon = scalar_icon - elif "x" in shape_str: - icon = array_icon - else: - icon = vector_icon - - item = QTreeWidgetItem(parent_item, [key, dtype, shape_str]) - item.setIcon(0, icon) - - item.setData(0, Qt.ItemDataRole.UserRole, child.name) - item.setData(0, Qt.ItemDataRole.UserRole + 1, "dataset") - - item.setForeground(1, QBrush(QColor(get_accent_colors().success.name()))) - item.setForeground(2, QBrush(QColor("#656365"))) - - def _get_filepath_for_item(self, item: QTreeWidgetItem) -> str: - """Walk up the tree to find the filepath stored on the root node.""" - node = item - while node.parent(): - node = node.parent() - return node.data(0, Qt.ItemDataRole.UserRole + 2) - - def _on_item_clicked(self, item, _col): - path = item.data(0, Qt.ItemDataRole.UserRole) - kind = item.data(0, Qt.ItemDataRole.UserRole + 1) - filepath = self._get_filepath_for_item(item) - - if not path or not filepath: - return - - h5file = self.h5files[filepath] - obj = h5file[path] if path != "/" else h5file - - if kind == "dataset": - data = h5file[path][()] - self.data_panel.display(data, path) - else: - n = len(obj) if isinstance(obj, h5py.Group) else 0 - - def clear_files(self): - """Close all open HDF5 files and reset the tree.""" - for h5file in self.h5files.values(): - h5file.close() - self.h5files.clear() - self.tree.clear() - self.data_panel.show_empty() - - def closeEvent(self, event): - for h5file in self.h5files.values(): - h5file.close() - self.h5files.clear() - super().closeEvent(event) diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/widgets/__init__,py b/superxas_bec/bec_widgets/widgets/data_viewer/widgets/__init__,py new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/data_viewer/qt_widgets.py b/superxas_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py similarity index 50% rename from superxas_bec/bec_widgets/widgets/data_viewer/qt_widgets.py rename to superxas_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py index 1e1d7e9..187157e 100644 --- a/superxas_bec/bec_widgets/widgets/data_viewer/qt_widgets.py +++ b/superxas_bec/bec_widgets/widgets/data_viewer/widgets/qt_widgets.py @@ -1,151 +1,235 @@ -""" -TaggedListWidget — a QListWidget where each item shows a label + styled tag pills. -Selection works natively; no popup, no paintEvent hacks. -""" - -import sys - -from qtpy.QtCore import QPoint, QRect, QSize, Qt -from qtpy.QtGui import QColor, QFont, QFontMetrics, QPainter, QPen -from qtpy.QtWidgets import ( - QApplication, - QLabel, - QListWidget, - QListWidgetItem, - QStyle, - QStyledItemDelegate, - QStyleOptionViewItem, - QVBoxLayout, - QWidget, -) - -# ── Design tokens ────────────────────────────────────────────────────────────── -ITEM_HEIGHT = 30 -H_PAD = 12 -TAG_H_PAD = 7 -TAG_V_PAD = 3 -TAG_GAP = 5 -LABEL_TAG_GAP = 12 -CORNER_RADIUS = 4 -TAG_TEXT_COLOR = "#FFFFFF" - -# ── Enum compat (PySide6 nested vs PyQt5 flat) ──────────────────────────────── -try: - _DEMIBOLD = QFont.Weight.DemiBold - _MEDIUM = QFont.Weight.Medium -except AttributeError: - _DEMIBOLD = QFont.DemiBold - _MEDIUM = QFont.Medium - -_AlignVCenter = Qt.AlignmentFlag.AlignVCenter if hasattr(Qt, "AlignmentFlag") else Qt.AlignVCenter -_AlignCenter = Qt.AlignmentFlag.AlignCenter if hasattr(Qt, "AlignmentFlag") else Qt.AlignCenter -_UserRole = Qt.ItemDataRole.UserRole if hasattr(Qt, "ItemDataRole") else Qt.UserRole -_NoPen = Qt.PenStyle.NoPen if hasattr(Qt, "PenStyle") else Qt.NoPen -_AA = QPainter.RenderHint.Antialiasing if hasattr(QPainter, "RenderHint") else QPainter.Antialiasing - -try: - _State_Selected = QStyle.StateFlag.State_Selected - _State_MouseOver = QStyle.StateFlag.State_MouseOver -except AttributeError: - _State_Selected = QStyle.State_Selected - _State_MouseOver = QStyle.State_MouseOver - - -# ── Delegate ────────────────────────────────────────────────────────────────── -class TaggedDelegate(QStyledItemDelegate): - - def sizeHint(self, option: QStyleOptionViewItem, index) -> QSize: - return QSize(option.rect.width() or 300, ITEM_HEIGHT) - - def paint(self, painter: QPainter, option: QStyleOptionViewItem, index) -> None: - painter.save() - - is_selected = bool(option.state & _State_Selected) - is_hover = bool(option.state & _State_MouseOver) - - # Background - if is_selected: - painter.fillRect(option.rect, option.palette.highlight()) - elif is_hover: - painter.fillRect(option.rect, QColor("#F1F5F9")) - else: - painter.fillRect(option.rect, option.palette.base()) - - label_text = ( - index.data( - Qt.ItemDataRole.DisplayRole if hasattr(Qt, "ItemDataRole") else Qt.DisplayRole - ) - or "" - ) - tags: list = index.data(_UserRole) or [] - - # Label - label_font = painter.font() # inherit default font - label_font.setWeight(_DEMIBOLD) - painter.setFont(label_font) - label_color = ( - option.palette.highlightedText().color() - if is_selected - else option.palette.text().color() - ) - painter.setPen(QPen(label_color)) - - fm = QFontMetrics(label_font) - label_w = fm.horizontalAdvance(label_text) - label_y = option.rect.top() + (ITEM_HEIGHT - fm.height()) // 2 + fm.ascent() - painter.drawText(QPoint(option.rect.left() + H_PAD, label_y), label_text) - - # Tag pills - tag_font = painter.font() # inherit default font - tag_font.setWeight(_MEDIUM) - fm_tag = QFontMetrics(tag_font) - - x = option.rect.left() + H_PAD + label_w + LABEL_TAG_GAP - for tag_text, hex_color in tags: - tag_text = str(tag_text) - tw = fm_tag.horizontalAdvance(tag_text) + 2 * TAG_H_PAD - th = fm_tag.height() + 2 * TAG_V_PAD - ty = option.rect.top() + (ITEM_HEIGHT - th) // 2 - pill = QRect(x, ty, tw, th) - - fill = QColor(hex_color) - if is_selected: - fill = fill.lighter(140) - - painter.setRenderHint(_AA) - painter.setPen(_NoPen) - painter.setBrush(fill) - painter.drawRoundedRect(pill, CORNER_RADIUS, CORNER_RADIUS) - - painter.setFont(tag_font) - painter.setPen(QPen(QColor(TAG_TEXT_COLOR))) - painter.drawText(pill, _AlignCenter, tag_text) - - x += tw + TAG_GAP - - painter.restore() - - -# ── Widget ──────────────────────────────────────────────────────────────────── -class TaggedListWidget(QListWidget): - """QListWidget with label + coloured tag pills per row.""" - - def __init__(self, parent=None) -> None: - super().__init__(parent) - self.setItemDelegate(TaggedDelegate(self)) - self.setMouseTracking(True) # enables hover highlight - - def addTaggedItem(self, label: str, tags: list | None = None) -> QListWidgetItem: - """ - Add a row. tags is a list of (text, hex_color) pairs, e.g. - [("v1.26", "#2563EB"), ("stable", "#16A34A")] - """ - item = QListWidgetItem(str(label)) - if tags: - item.setData(_UserRole, list(tags)) - self.addItem(item) - return item - - def currentTags(self) -> list: - item = self.currentItem() - return item.data(_UserRole) or [] if item else [] +""" +Universal Qt widgets +""" + +from functools import partial + +from bec_widgets.utils.colors import get_accent_colors + +# pylint: disable=E0611 +from qtpy.QtCore import QPoint, QRect, QSize, Qt +from qtpy.QtGui import QColor, QFont, QFontMetrics, QPainter, QPen +from qtpy.QtWidgets import ( + QGroupBox, + QHBoxLayout, + QLabel, + QLayout, + QListWidget, + QListWidgetItem, + QPushButton, + QStyle, + QStyledItemDelegate, + QStyleOptionViewItem, + QVBoxLayout, + QWidget, +) + + +class Group(QGroupBox): + def __init__(self, label, objs, orientation="vertical"): + super().__init__(label) + if orientation == "vertical": + self._layout = QVBoxLayout(self) + elif orientation == "horizontal": + self._layout = QHBoxLayout(self) + else: + raise ValueError(f"Orientation {orientation} is not supported!") + for obj in objs: + if isinstance(obj, QWidget): + self._layout.addWidget(obj) + elif isinstance(obj, QLayout): + self._layout.addLayout(obj) + + +class Button(QWidget): + def __init__(self, label=None, label_button: str = "", enabled=False): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(10, 0, 0, 0) + layout.setSpacing(0) + if label is not None: + self.label = QLabel(label) + self.label.setFixedWidth(140) + layout.addWidget(self.label) + self.button = QPushButton(label_button) + if label is not None: + self.button.setFixedWidth(160) + self.enable_button(enabled) + layout.addWidget(self.button) + + def clicked_connect(self, func): + """Connect a function to the button press.""" + self.button.clicked.connect(func) + + def enable_button(self, enable: bool = False): + if enable: + self.button.setStyleSheet( + f"QPushButton {{background-color: {get_accent_colors().default.name()}; color: white;}}" + ) + self.button.setEnabled(True) + else: # disabled + self.button.setStyleSheet( + "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}" + ) + self.button.setDisabled(True) + + def setText(self, text): + self.button.setText(text) + + +class ListWidget(QWidget): + def __init__(self, identifier=""): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(10, 0, 0, 0) + layout.setSpacing(0) + self.identifier = identifier + self.value = TaggedListWidget() + layout.addWidget(self.value) + + def clear(self): + self.value.clear() + + def addTaggedItem(self, label, tags): + self.value.addTaggedItem(label, tags) + + def setCurrentIndex(self, text): + self.value.setCurrentIndex(text) + + def currentItemChanged_connect(self, func): + """Connect a function to the Enter/Return key press.""" + self.value.currentItemChanged.connect( + partial( + func, + identifier=self.identifier, + value_obj=self.value, + value=lambda: self.value.currentIndex(), + ) + ) + + def setDisabled(self, disable): + self.value.setDisabled(disable) + + +class TaggedListWidget(QListWidget): + """QListWidget with label + coloured tag pills per row.""" + + def __init__(self, parent=None) -> None: + super().__init__(parent) + self.setItemDelegate(TaggedDelegate(self)) + self.setMouseTracking(True) # enables hover highlight + + def addTaggedItem(self, label: str, tags: list | None = None) -> QListWidgetItem: + """ + Add a row. tags is a list of (text, hex_color) pairs, e.g. + [("v1.26", "#2563EB"), ("stable", "#16A34A")] + """ + item = QListWidgetItem(str(label)) + if tags: + item.setData(_UserRole, list(tags)) + self.addItem(item) + return item + + def currentTags(self) -> list: + item = self.currentItem() + return item.data(_UserRole) or [] if item else [] + + +ITEM_HEIGHT = 30 +H_PAD = 12 +TAG_H_PAD = 7 +TAG_V_PAD = 3 +TAG_GAP = 5 +LABEL_TAG_GAP = 12 +CORNER_RADIUS = 4 +TAG_TEXT_COLOR = "#FFFFFF" +_DEMIBOLD = QFont.Weight.DemiBold +_MEDIUM = QFont.Weight.Medium +_AlignCenter = ( + Qt.AlignmentFlag.AlignCenter if hasattr(Qt, "AlignmentFlag") else Qt.AlignmentFlag.AlignCenter +) +_UserRole = Qt.ItemDataRole.UserRole if hasattr(Qt, "ItemDataRole") else Qt.ItemDataRole.UserRole +_NoPen = Qt.PenStyle.NoPen if hasattr(Qt, "PenStyle") else Qt.PenStyle.NoPen +_AA = ( + QPainter.RenderHint.Antialiasing + if hasattr(QPainter, "RenderHint") + else QPainter.RenderHint.Antialiasing +) +_State_Selected = QStyle.StateFlag.State_Selected +_State_MouseOver = QStyle.StateFlag.State_MouseOver + + +class TaggedDelegate(QStyledItemDelegate): + + def sizeHint(self, option: QStyleOptionViewItem, index) -> QSize: + return QSize(option.rect.width() or 300, ITEM_HEIGHT) + + def paint(self, painter: QPainter, option: QStyleOptionViewItem, index) -> None: + painter.save() + + is_selected = bool(option.state & _State_Selected) + is_hover = bool(option.state & _State_MouseOver) + + # Background + if is_selected: + painter.fillRect(option.rect, option.palette.highlight()) + elif is_hover: + painter.fillRect(option.rect, QColor("#F1F5F9")) + else: + painter.fillRect(option.rect, option.palette.base()) + + label_text = ( + index.data( + Qt.ItemDataRole.DisplayRole + if hasattr(Qt, "ItemDataRole") + else Qt.ItemDataRole.DisplayRole + ) + or "" + ) + tags: list = index.data(_UserRole) or [] + + # Label + label_font = painter.font() # inherit default font + label_font.setWeight(_DEMIBOLD) + painter.setFont(label_font) + label_color = ( + option.palette.highlightedText().color() + if is_selected + else option.palette.text().color() + ) + painter.setPen(QPen(label_color)) + + fm = QFontMetrics(label_font) + label_w = fm.horizontalAdvance(label_text) + label_y = option.rect.top() + (ITEM_HEIGHT - fm.height()) // 2 + fm.ascent() + painter.drawText(QPoint(option.rect.left() + H_PAD, label_y), label_text) + + # Tag pills + tag_font = painter.font() # inherit default font + tag_font.setWeight(_MEDIUM) + fm_tag = QFontMetrics(tag_font) + + x = option.rect.left() + H_PAD + label_w + LABEL_TAG_GAP + for tag_text, hex_color in tags: + tag_text = str(tag_text) + tw = fm_tag.horizontalAdvance(tag_text) + 2 * TAG_H_PAD + th = fm_tag.height() + 2 * TAG_V_PAD + ty = option.rect.top() + (ITEM_HEIGHT - th) // 2 + pill = QRect(x, ty, tw, th) + + fill = QColor(hex_color) + if is_selected: + fill = fill.lighter(140) + + painter.setRenderHint(_AA) + painter.setPen(_NoPen) + painter.setBrush(fill) + painter.drawRoundedRect(pill, CORNER_RADIUS, CORNER_RADIUS) + + painter.setFont(tag_font) + painter.setPen(QPen(QColor(TAG_TEXT_COLOR))) + painter.drawText(pill, _AlignCenter, tag_text) + + x += tw + TAG_GAP + + painter.restore() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/__init__.py b/superxas_bec/bec_widgets/widgets/digital_twin/__init__.py new file mode 100644 index 0000000..a42cb09 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/__init__.py @@ -0,0 +1,3 @@ +from .beamline import get_parameters + +parameters = get_parameters() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/beamline.py b/superxas_bec/bec_widgets/widgets/digital_twin/beamline.py new file mode 100644 index 0000000..3653284 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/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/superxas_bec/bec_widgets/widgets/digital_twin/calculations/__init__.py b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py new file mode 100644 index 0000000..6347d4b --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_positions.py @@ -0,0 +1,259 @@ +""" +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 ConfigDict + +logger = bec_logger.logger + + +def calc_positions(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 + 1} # Add 0.5 mm space on both sides of the beam + + ## 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 + 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 + 1} # Add 0.5 mm space on both sides of the beam + + ## 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 + + # 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) + + return pos diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py new file mode 100644 index 0000000..135c76d --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_sideview.py @@ -0,0 +1,70 @@ +""" +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/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py new file mode 100644 index 0000000..9411ded --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_surfaces.py @@ -0,0 +1,157 @@ +""" +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"] = [] + + # 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/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py new file mode 100644 index 0000000..33ad7cf --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/calculations/calc_varia.py @@ -0,0 +1,444 @@ +""" +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 + 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 + """ + 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] + """ + 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 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 + """ + 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"pitch: {pitch}") + return pitch, qy + + +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 = [] + 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 = [] + 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 in bl.tables: + return bl.tables[table]["smpl"] + raise ValueError(f"Table {table} not found in beamline parameter file") diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.py b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.py new file mode 100644 index 0000000..180e9fb --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.py @@ -0,0 +1,937 @@ +""" +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 +import yaml +from bec_lib import bec_logger +from bec_lib.endpoints import MessageEndpoints +from bec_widgets.utils.bec_dispatcher import BECDispatcher +from bec_widgets.utils.bec_widget import BECWidget +from bec_widgets.utils.colors import apply_theme, get_accent_colors +from bec_widgets.utils.error_popups import SafeSlot + +# pylint: disable=E0611 +from qtpy.QtCore import Qt, QTimer +from qtpy.QtGui import QFont +from qtpy.QtWidgets import ( + QApplication, + QComboBox, + QDialog, + QDialogButtonBox, + QFrame, + QHBoxLayout, + QLabel, + QPlainTextEdit, + QPushButton, + QScrollArea, + QSizePolicy, + QStyle, + QVBoxLayout, + QWidget, +) + +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 ( + 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") + + +class DigitalTwin(BECWidget, QWidget): + """ + Main widget of Digital Twin + """ + + PLUGIN = True + ICON_NAME = "lightbulb" + + def __init__(self, *arg, parent=None, **kwargs): + super().__init__(parent=parent, theme_update=True, *arg, **kwargs) + self.get_bec_shortcuts() + + 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( + self.check_bec_config, MessageEndpoints.device_config_update() + ) + + logger.info(f"Start Digital Twin with beamline {self.beamline} and all devices available") + + self.content_widget = QWidget(self) + self.root_layout = QHBoxLayout(self.content_widget) + self.root_layout.setContentsMargins(6, 6, 6, 6) + self.root_layout.setSpacing(6) + + self.input_widget = QWidget() + self.input_layout = QVBoxLayout(self.input_widget) + self.input_layout.setContentsMargins(4, 4, 4, 4) + self.input_layout.setSpacing(6) + self.input = InputPanel(self.beamline) + self.settings = SettingsPanel() + self.input_layout.addWidget(self.input) + self.input_layout.addWidget(self.settings) + self.input_layout.addStretch() + + self.plot_widget = 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.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) + + self.mover = MoverPanel(self.dev) + + self.input_scroll = self._scroll_area(self.input_widget, min_width=320, max_width=360) + self.mover_scroll = self._scroll_area(self.mover, min_width=380, max_width=460) + + self.root_layout.addWidget(self.input_scroll) + self.root_layout.addWidget(self.plot_widget, stretch=1) + self.root_layout.addWidget(self.mover_scroll) + widget_layout = self.layout() + if widget_layout is None: + widget_layout = QVBoxLayout(self) + widget_layout.setContentsMargins(0, 0, 0, 0) + widget_layout.setSpacing(0) + widget_layout.addWidget(self.content_widget) + self.setWindowTitle("Digital Twin") + self.resize(1450, 950) + + self.input.energy.value_changed_connect(self.calc_assistant) + self.input.sldi_hacc.value_changed_connect(self.calc_assistant) + self.input.sldi_vacc.value_changed_connect(self.calc_assistant) + self.input.cm_stripe.activated_connect(self.calc_assistant) + self.input.cm_pitch.value_changed_connect(self.calc_assistant) + self.input.mo1_mode.activated_connect(self.calc_assistant) + self.input.mo1_xtal.activated_connect(self.calc_assistant) + self.input.fm_stripe.activated_connect(self.calc_assistant) + self.input.fm_focus.activated_connect(self.calc_assistant) + self.input.fm_rotx.value_changed_connect(self.calc_assistant) + self.input.fm_focx.value_changed_connect(self.calc_assistant) + self.input.fm_focy.value_changed_connect(self.calc_assistant) + match self.input.smpl: + case InputNumberField(): + self.input.smpl.value_changed_connect(self.calc_assistant) + case ComboBox(): + self.input.smpl.activated_connect(self.calc_assistant) + + self.input.adapt_reality.clicked_connect(self.adapt_reality) + self.settings.load_offsets.clicked_connect(self.load_offsets) + self.settings.show_offsets.clicked_connect(self.show_offsets) + + self.bragg_angle = 0.0 + self.qy = 0.0 + self.offsets = {} + + # Initialize all values + self.load_offsets(recalculate=False) + self.calc_assistant(identifier="init") + + # Timer: update reality plots every 1 second + self._timer = QTimer(self) + self._timer.setInterval(1000) + self._timer.timeout.connect(self.calc_reality) + self._timer.start() + + @staticmethod + def _scroll_area(widget: QWidget, min_width: int, max_width: int) -> QScrollArea: + """Wrap a side panel in a compact vertical scroll area.""" + scroll = QScrollArea() + scroll.setWidgetResizable(True) + widget.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.MinimumExpanding) + scroll.setWidget(widget) + scroll.setFrameShape(QFrame.Shape.NoFrame) + scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff) + scroll.setMinimumWidth(min_width) + scroll.setMaximumWidth(max_width) + scroll.setSizePolicy(QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Expanding) + return scroll + + def apply_theme(self, theme: Literal["dark", "light"]): + """ + Apply the theme + + Args: + theme (str): Theme, either "dark" or "light" + """ + self.sideview_plot.apply_theme(theme) + self.surface_plots.apply_theme(theme) + self.mover.apply_theme(theme) + + @SafeSlot() + def check_bec_config(self, *args): + """ + Checks the BEC config and opens a window if not all necessary + devices are loaded in the config. If called from a slot from + BEC dispatcher whenever there is a config update, stop the timer + that updates the plot in the background. + """ + reload_config = (args[0] if args else {}).get("action") == "reload" + if reload_config: + self._timer.stop() + devices = [ + "abs", + "sldi_gapx", + "sldi_gapy", + "cm_trx", + "cm_try", + "cm_bnd", + "cm_rotx", + "mo1_bragg", + "mo1_trx", + "mo1_try", + "sl1_centery", + "sl1_gapy", + "bm1_try", + "fm_trx", + "fm_try", + "fm_bnd", + "fm_rotx", + "fm_roty", + "fm_rotz", + "bm2_try", + "ot_try", + "es0wi_try", + ] + if self.beamline == "x01da": # X01DA specific devices + devices.extend( + [ + "cm_bnd_radius", + "fm_bnd_radius", + "sl2_centery", + "sl2_gapy", + "ot_try", + "ot_es1_trz", + ] + ) + if self.beamline == "x10da": # X10DA specific devices + devices.extend(["mo1_rotx", "es1_try", "es1ic0_try", "es1ic1_try", "es1ic2_try"]) + + while True: + missing = [d for d in devices if d not in self.dev] + if not missing: + break + dialog = QDialog(self) + dialog.setWindowTitle("Digital Twin - Config Check") + dialog.setFixedWidth(400) + layout = QVBoxLayout() + + top = QHBoxLayout() + icon = QLabel() + icon_pixmap = ( + QApplication.style() + .standardIcon(QStyle.StandardPixmap.SP_MessageBoxWarning) + .pixmap(48, 48) + ) + icon.setPixmap(icon_pixmap) + icon.setAlignment(Qt.AlignmentFlag.AlignTop) + top.addWidget(icon) + + text = QLabel( + "The current config does not include all required devices to run Digital Twin." + + "Reload the config with the correct devices." + ) + text.setWordWrap(True) + text.setAlignment(Qt.AlignmentFlag.AlignTop) + top.addWidget(text, stretch=1) + layout.addLayout(top) + + info = QLabel("Missing devices:\n" + ", ".join(missing)) + info.setWordWrap(True) + info.setAlignment(Qt.AlignmentFlag.AlignTop) + layout.addWidget(info) + layout.addStretch() + + buttons = QHBoxLayout() + check_again = QPushButton("Check Again") + close_app = QPushButton("Close Application") + check_again.clicked.connect(dialog.accept) + close_app.clicked.connect(dialog.reject) + buttons.addWidget(check_again) + buttons.addWidget(close_app) + layout.addLayout(buttons) + + dialog.setLayout(layout) + dialog.show() + info.setMinimumHeight(info.heightForWidth(info.width())) + if dialog.exec_() == QDialog.DialogCode.Rejected: + running_app = QApplication.instance() + if running_app is not None: + running_app.exit(0) + if reload_config: + self._timer.start() + + @SafeSlot() + def calc_assistant(self, *_, **kwargs): + """ + Calculates various values for the assistant. + If called from a qt slot, the identifier represents + the button pressed / value changed. Based on the identifier, + calculate different values. + Note: identifier=init calculates all values + """ + identifier = kwargs["identifier"] + match identifier: + case "init": + self.update_mo1_mode() + self.calc_mo1_bragg_angle() + self.calc_cm_crit_pitch() + self.calc_cm_reflectivity() + self.update_fm_mode() + self.calc_fm_reflectivity() + self.calc_cm_fm_harm_suppr() + self.calc_fm_ideal_pitch() + self.calc_mo1_energy_resolution() + case "energy": + self.calc_mo1_bragg_angle() + self.calc_cm_crit_pitch() + self.calc_cm_reflectivity() + self.calc_fm_reflectivity() + self.calc_cm_fm_harm_suppr() + self.calc_mo1_energy_resolution() + case "hacc": + self.calc_fm_ideal_pitch() + case "vacc": + self.calc_fm_ideal_pitch() + case "cm_stripe": + self.calc_cm_crit_pitch() + self.calc_cm_reflectivity() + self.calc_cm_fm_harm_suppr() + case "cm_pitch": + self.calc_cm_reflectivity() + self.calc_cm_fm_harm_suppr() + case "mo1_mode": + self.update_mo1_mode() + case "mo1_xtal": + self.calc_mo1_bragg_angle() + self.calc_mo1_energy_resolution() + case "fm_focus": + self.update_fm_mode() + self.calc_fm_ideal_pitch() + case "fm_focx": + self.calc_fm_ideal_pitch() + case "fm_focy": + self.calc_fm_ideal_pitch() + case "fm_rotx": + self.calc_fm_reflectivity() + self.calc_cm_fm_harm_suppr() + case "fm_stripe": + self.calc_fm_reflectivity() + self.calc_cm_fm_harm_suppr() + self.calc_fm_ideal_pitch() + case "smpl": + self.calc_fm_ideal_pitch() + self.calc_positions() + self.calc_assistant_sideview() + self.calc_assistant_surfaces() + + def get_assistant_config(self, apply_offset: bool = False) -> ConfigDict: + """ + Assembles the digital twin config from the assistants input. + + Args: + apply_offset(bool): Applies the offset values to the config. + Defaults to False + + Returns: + ConfigDict: config of the assistant + """ + fm_focus = self.input.fm_focus.currentText() + if fm_focus == "Manual": + fm_rotx = self.input.fm_rotx.value() + fm_qy = None + elif fm_focus == "Focused": + fm_rotx = self.input.fm_rotx_ideal.value() + fm_qy = None + else: # Focused + fm_rotx = self.input.fm_rotx_ideal.value() + fm_qy = self.qy + + cm_stripe = self.input.cm_stripe.currentText() + cm_trx = cm_stripe_to_trx(cm_stripe) + fm_stripe = self.input.fm_stripe.currentText() + fm_trx = 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}!" + + match self.input.smpl: + case InputNumberField(): + smpl = self.input.smpl.value() + case ComboBox(): + table = self.input.smpl.currentText() + smpl = table_to_smpl_pos(table) + + config: ConfigDict = { + "energy": self.input.energy.value(), + "h_acc": self.input.sldi_hacc.value(), + "v_acc": self.input.sldi_vacc.value(), + "cm_pitch": -self.input.cm_pitch.value(), + "cm_stripe": cm_stripe, + "cm_trx": cm_trx, + "mo1_mode": self.input.mo1_mode.currentText(), + "mo1_xtal": self.input.mo1_xtal.currentText(), + "mo1_bragg": self.bragg_angle, + "fm_rotx": -fm_rotx, + "fm_stripe": fm_stripe, + "fm_trx": fm_trx, + "fm_qy": fm_qy, + "fm_gain_height": 1, + "smpl": smpl, + } + + # 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"] + + # Convert to SI units! + config["h_acc"] *= 1e-3 + config["v_acc"] *= 1e-3 + config["cm_pitch"] *= 1e-3 + config["fm_rotx"] *= 1e-3 + + # logger.info(f'Config created: {config}') + return config + + def get_reality_config(self) -> ConfigDict: + """ + Assembles the digital twin config based on the real axis positions. + + Returns: + ConfigDict: config of the reality + """ + mo1_trx = self.dev.mo1_trx.read(cached=True)["mo1_trx"]["value"] + if abs(mo1_trx) > 5: + mo1_mode = "Monochromatic" + else: + mo1_mode = "Pinkbeam" + 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) + cm_trx = self.dev.cm_trx.read(cached=True)["cm_trx"]["value"] + cm_stripe = 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_rotx = self.dev.fm_rotx.read(cached=True)["fm_rotx"]["value"] + fm_rotx_real = 2 * cm_pitch - fm_rotx + smpl = self.dev.ot_es1_trz.read(cached=True)["ot_es1_trz"]["value"] + raw = { # Config in SI units! + "energy": mo1_bragg["mo1_bragg"]["value"], + "h_acc": h_acc, + "v_acc": v_acc, + "cm_pitch": -cm_pitch * 1e-3, + "cm_stripe": cm_stripe, + "cm_trx": cm_trx, + "mo1_mode": mo1_mode, + "mo1_xtal": mo1_bragg["mo1_bragg_crystal_current_xtal_string"]["value"], + "mo1_bragg": mo1_bragg["mo1_bragg_angle"]["value"] / 180 * np.pi, + "fm_rotx": -fm_rotx_real * 1e-3, + "fm_stripe": fm_stripe, + "fm_trx": fm_trx, + "fm_qy": None, + "fm_gain_height": 1, + "smpl": smpl, + } + config = cast(ConfigDict, raw) + # logger.info(f'Config created: {config}') + + abs_open = self.dev.abs.read(cached=True)["abs_status_string"]["value"] == "OPEN" + if not abs_open: + ready = True + for mover in self.mover.mover_widgets: + if mover.status in ("moving", "error"): + ready = False + if ready: + self.mover.abs.enable_open(True) # Enable open button + else: + self.mover.abs.enable_open(False) # Disable open button + else: + self.mover.abs.enable_open(False) # Disable open button + + self.mover.sldi_gapx.set_feedback(sldi_gapx) + self.mover.sldi_gapy.set_feedback(sldi_gapy) + self.mover.cm_trx.set_feedback(cm_trx) + self.mover.cm_try.set_feedback(self.dev.cm_try.read(cached=True)["cm_try"]["value"]) + self.mover.cm_bnd.set_feedback( + self.dev.cm_bnd_radius.read(cached=True)["cm_bnd_radius"]["value"] + ) + self.mover.cm_rotx.set_feedback(cm_pitch) + self.mover.mo1_bragg_angle.set_feedback(mo1_bragg["mo1_bragg_angle"]["value"]) + self.mover.mo1_trx.set_feedback(mo1_trx) + self.mover.mo1_try.set_feedback(self.dev.mo1_try.read(cached=True)["mo1_try"]["value"]) + self.mover.sl1_centery.set_feedback( + self.dev.sl1_centery.read(cached=True)["sl1_centery"]["value"] + ) + self.mover.sl1_gapy.set_feedback(self.dev.sl1_gapy.read(cached=True)["sl1_gapy"]["value"]) + self.mover.bm1_try.set_feedback(self.dev.bm1_try.read(cached=True)["bm1_try"]["value"]) + self.mover.fm_trx.set_feedback(fm_trx) + self.mover.fm_try.set_feedback(self.dev.fm_try.read(cached=True)["fm_try"]["value"]) + self.mover.fm_bnd.set_feedback( + self.dev.fm_bnd_radius.read(cached=True)["fm_bnd_radius"]["value"] + ) + self.mover.fm_rotx.set_feedback(fm_rotx) + self.mover.fm_roty.set_feedback(self.dev.fm_roty.read(cached=True)["fm_roty"]["value"]) + self.mover.fm_rotz.set_feedback(self.dev.fm_rotz.read(cached=True)["fm_rotz"]["value"]) + self.mover.sl2_centery.set_feedback( + self.dev.sl2_centery.read(cached=True)["sl2_centery"]["value"] + ) + self.mover.sl2_gapy.set_feedback(self.dev.sl2_gapy.read(cached=True)["sl2_gapy"]["value"]) + self.mover.bm2_try.set_feedback(self.dev.bm2_try.read(cached=True)["bm2_try"]["value"]) + self.mover.ot_try.set_feedback(self.dev.ot_try.read(cached=True)["ot_try"]["value"]) + self.mover.ot_rotx.set_feedback(self.dev.ot_rotx.read(cached=True)["ot_rotx"]["value"]) + self.mover.ot_es1_trz.set_feedback(smpl) + self.mover.es0wi_try.set_feedback( + self.dev.es0wi_try.read(cached=True)["es0wi_try"]["value"] + ) + self.mover.abs.set_feedback(abs_open) + return config + + @SafeSlot() + def adapt_reality(self, *_): + """ + Based on the real axis positions, adjust the assistant to reflect + the reality. + """ + pos = {} + pos["sldi_gapx"] = self.dev.sldi_gapx.read(cached=True)["sldi_gapx"]["value"] + pos["sldi_gapy"] = self.dev.sldi_gapy.read(cached=True)["sldi_gapy"]["value"] + pos["cm_trx"] = self.dev.cm_trx.read(cached=True)["cm_trx"]["value"] + pos["cm_rotx"] = self.dev.cm_rotx.read(cached=True)["cm_rotx"]["value"] + pos["mo1_trx"] = self.dev.mo1_trx.read(cached=True)["mo1_trx"]["value"] + pos["fm_trx"] = self.dev.fm_trx.read(cached=True)["fm_trx"]["value"] + pos["fm_rotx"] = self.dev.fm_rotx.read(cached=True)["fm_rotx"]["value"] + 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"] + + 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"]) + 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_pitch.set_number(pos["cm_rotx"]) + if abs(pos["mo1_trx"]) > 5: + mo1_mode = "Monochromatic" + else: + mo1_mode = "Pinkbeam" + self.input.mo1_mode.set_current_text(mo1_mode) + self.input.mo1_xtal.set_current_text( + self.dev.mo1_bragg.read(cached=True)["mo1_bragg_crystal_current_xtal_string"]["value"] + ) + self.input.fm_stripe.set_current_text(fm_trx_to_stripe(-pos["fm_trx"])) + self.input.fm_focus.set_current_text("Manual") + fm_rotx_real = 2 * pos["cm_rotx"] - pos["fm_rotx"] + self.input.fm_rotx.set_number(fm_rotx_real) + match self.input.smpl: + case InputNumberField(): + self.input.smpl.set_number(pos["ot_es1_trz"]) + case ComboBox(): + table = self.ask_table_selection(self.input.smpl.currentText()) + self.input.smpl.set_current_text(table) + + self.calc_assistant(identifier="init") + + def ask_table_selection(self, preset=None) -> str | None: + """ + Opens a dialog asking the user to select a table (ES1 or ES2). + + Args: + preset (str): Preset text for the table, either 'ES1' or 'ES2'. + + Returns: + The selected table ('ES1' or 'ES2'), or None if the user cancelled. + """ + dialog = QDialog(self) + dialog.setWindowTitle("Select Table") + layout = QVBoxLayout(dialog) + + text = QLabel("Select the current table in use.") + + combo = QComboBox() + choice = ["ES1", "ES2"] + combo.addItems(choice) + if preset is not None: + if preset in choice: + combo.setCurrentText(preset) + + buttons = QDialogButtonBox( + QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel + ) + buttons.accepted.connect(dialog.accept) + buttons.rejected.connect(dialog.reject) + + layout.addWidget(text) + layout.addWidget(combo) + layout.addWidget(buttons) + + if dialog.exec() == QDialog.DialogCode.Accepted: + return combo.currentText() + return None + + @SafeSlot() + def load_offsets(self, *_, recalculate: bool = True): + """ + Loads or unloads the offsets from the file + + Args: + recalculate(bool): Recalculates the assistant values after loading. + 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 + + if recalculate: + self.calc_assistant(identifier="init") + + self.settings.load_offsets.setText("Unload") + self.settings.offsets_status.setText("Loaded and applied") + self.settings.offsets_status.setColor(get_accent_colors().success.name()) + self.settings.show_offsets.enable_button(True) + else: + # Unload offsets + self.offsets = {} + self.calc_assistant(identifier="init") + + self.settings.load_offsets.setText("Load") + self.settings.offsets_status.setText("No offsets") + self.settings.offsets_status.setColor(get_accent_colors().default.name()) + self.settings.show_offsets.enable_button(False) + + @SafeSlot() + def show_offsets(self, *_): + """ + Shows the offsets in a popup window + """ + dialog = QDialog() + dialog.setWindowTitle("Digital Twin - Offsets") + dialog.setFixedWidth(500) + layout = QVBoxLayout(dialog) + layout.setSpacing(12) + layout.setContentsMargins(20, 20, 20, 20) + + intro_label = QLabel("The offsets are saved in the digital twin BEC widget folder:") + intro_label.setWordWrap(True) + layout.addWidget(intro_label) + + file = QLabel(str(self.offset_file)) + file.setWordWrap(True) + font = QFont() + font.setItalic(True) + file.setFont(font) + layout.addWidget(file) + + text_edit = QPlainTextEdit() + text_edit.setReadOnly(True) + text_edit.setFont(QFont("Consolas", 9)) + + class InlineListDumper(yaml.Dumper): + """YAML dumper that renders all sequences on a single line.""" + + 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)) + layout.addWidget(text_edit) + + buttons = QDialogButtonBox(QDialogButtonBox.StandardButton.Close) + buttons.rejected.connect(dialog.reject) + layout.addWidget(buttons) + + dialog.exec() + + def update_fm_mode(self): + """ + Updates the focusing mirror input group based on the + selection of the focus strategy. + """ + fm_focus = self.input.fm_focus.currentText() + if fm_focus == "Manual": + self.input.fm_rotx.setVisible(True) + self.input.fm_rotx_ideal.setVisible(True) + self.input.fm_focx.setVisible(False) + self.input.fm_focy.setVisible(False) + self.input.fm_rotx_ideal.setLabel("Incidence Angle for focused beam") + elif fm_focus == "Focused": + self.input.fm_rotx.setVisible(False) + self.input.fm_rotx_ideal.setVisible(True) + self.input.fm_focx.setVisible(False) + self.input.fm_focy.setVisible(False) + self.input.fm_rotx_ideal.setLabel("Incidence Angle for focused beam") + else: # Defocused + self.input.fm_rotx.setVisible(False) + self.input.fm_rotx_ideal.setVisible(True) + self.input.fm_focx.setVisible(True) + self.input.fm_focy.setVisible(True) + self.input.fm_rotx_ideal.setLabel("Incidence Angle for defocused beam") + + @SafeSlot() + def calc_reality(self): + """ + Updates the plots for the reality scene + """ + config = self.get_reality_config() + data = calc_sideview(config) + self.sideview_plot.update_curves("reality", data=data) + surfaces = calc_surfaces(config) + self.surface_plots.update_surfaces(scene="reality", data=surfaces) + + def calc_mo1_energy_resolution(self): + """ + Calculates the energy resolution of the monochromator + """ + xtal = self.input.mo1_xtal.currentText().translate( + str.maketrans("", "", "()") + ) # 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)) + + def calc_cm_reflectivity(self): + """ + Calculates the collimating mirror reflectivity + """ + 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.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.setLabel(f"Reflectivity at \n{3*energy:.0f} eV") + + def calc_fm_reflectivity(self): + """ + Calculates the focusing mirror reflectivity + """ + fm_stripe = self.input.fm_stripe.currentText() + fm_focus = self.input.fm_focus.currentText() + if fm_focus == "Manual": + fm_rotx = -self.input.fm_rotx.value() * 1e-3 + 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.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.setLabel(f"Reflectivity at \n{3*energy:.0f} eV") + + def calc_cm_fm_harm_suppr(self): + """ + Calculates the combined harmonics suppression of both mirrors + """ + harm_suppr = (self.input.cm_refl.value() * self.input.fm_refl.value()) / ( + self.input.cm_refl_harm.value() * self.input.fm_refl_harm.value() + ) + self.input.cm_fm_harm_suppr.setValue(harm_suppr) + self.input.cm_fm_harm_suppr.setLabel( + f"Total Suppression Factor at {3 * self.input.energy.value():.0f} eV" + ) + + def calc_assistant_sideview(self): + """ + Updates the sideview plot based on the assistant values + """ + config = self.get_assistant_config(apply_offset=True) + data = 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()) + 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.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"] + + self.mover.sldi_gapx.set_target(out["sldi_gapx"]["value"]) + self.mover.sldi_gapy.set_target(out["sldi_gapy"]["value"]) + self.mover.cm_trx.set_target(out["cm_trx"]["value"]) + self.mover.cm_try.set_target(out["cm_try"]["value"]) + self.mover.cm_bnd.set_target(out["cm_bnd_radius"]["value"]) + self.mover.cm_rotx.set_target(out["cm_rotx"]["value"]) + self.mover.mo1_bragg_angle.set_target(out["mo1_bragg_angle"]["value"]) + self.mover.mo1_trx.set_target(out["mo1_trx"]["value"]) + self.mover.mo1_try.set_target(out["mo1_try"]["value"]) + self.mover.sl1_centery.set_target(out["sl1_centery"]["value"]) + self.mover.sl1_gapy.set_target(out["sl1_gapy"]["value"]) + self.mover.bm1_try.set_target(out["bm1_try"]["value"]) + self.mover.fm_trx.set_target(out["fm_trx"]["value"]) + self.mover.fm_try.set_target(out["fm_try"]["value"]) + self.mover.fm_bnd.set_target(out["fm_bnd_radius"]["value"]) + self.mover.fm_rotx.set_target(out["fm_rotx"]["value"]) + self.mover.fm_roty.set_target(out["fm_roty"]["value"]) + self.mover.fm_rotz.set_target(out["fm_rotz"]["value"]) + self.mover.sl2_centery.set_target(out["sl2_centery"]["value"]) + self.mover.sl2_gapy.set_target(out["sl2_gapy"]["value"]) + self.mover.bm2_try.set_target(out["bm2_try"]["value"]) + self.mover.ot_try.set_target(out["ot_try"]["value"]) + self.mover.ot_rotx.set_target(out["ot_rotx"]["value"]) + self.mover.ot_es1_trz.set_target(out["ot_es1_trz"]["value"]) + self.mover.es0wi_try.set_target(out["es0wi_try"]["value"]) + + def calc_mo1_bragg_angle(self): + """ + Calculates bragg angle in rad + """ + xtal = self.input.mo1_xtal.currentText() + if xtal == "Si(111)": + d_spacing = self.dev.mo1_bragg.crystal.d_spacing_si111.read(cached=True)[ + "mo1_bragg_crystal_d_spacing_si111" + ]["value"] + elif xtal == "Si(311)": + d_spacing = self.dev.mo1_bragg.crystal.d_spacing_si311.read(cached=True)[ + "mo1_bragg_crystal_d_spacing_si311" + ]["value"] + else: + raise ValueError(f"Invalid xtal selection: {xtal}") + 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) + self.bragg_angle = theta + self.input.mo1_bragg_angle.setValue(theta / np.pi * 180) + + def update_mo1_mode(self): + """ + Updates the monochromator input group based on the + selection of the mode. + """ + if self.input.mo1_mode.currentText() == "Monochromatic": + self.input.mo1_xtal.setVisible(True) + self.input.mo1_bragg_angle.setVisible(True) + self.input.mo1_eres.setVisible(True) + else: + self.input.mo1_xtal.setVisible(False) + self.input.mo1_bragg_angle.setVisible(False) + self.input.mo1_eres.setVisible(False) + + def calc_fm_ideal_pitch(self): + """ + Calculate the ideal pitch for the focusing mirror. + """ + fm_focus = cast( + Literal["Defocused", "Focused", "Manual"], self.input.fm_focus.currentText() + ) + fm_stripe = self.input.fm_stripe.currentText() + match self.input.smpl: + case InputNumberField(): + smpl = self.input.smpl.value() + case ComboBox(): + table = self.input.smpl.currentText() + smpl = 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_focus, fm_stripe, smpl, sldi_hacc, sldi_vacc, fm_focx, fm_focy + ) + self.qy = qy + self.input.fm_rotx_ideal.setValue(-fm_rotx * 1e3) + + def calc_cm_crit_pitch(self): + """ + Calculate the critical pitch for the collimating mirror + """ + 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) + + +if __name__ == "__main__": + app = QApplication(sys.argv) + apply_theme("light") + dispatcher = BECDispatcher(gui_id="digital_twin") + win = DigitalTwin() + win.show() + sys.exit(app.exec_()) diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.pyproject b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.pyproject new file mode 100644 index 0000000..226dc5f --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin.pyproject @@ -0,0 +1 @@ +{'files': ['digital_twin.py']} \ No newline at end of file diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin_plugin.py b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin_plugin.py new file mode 100644 index 0000000..921fd6b --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/digital_twin_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 .digital_twin import DigitalTwin + +DOM_XML = """ + + + + +""" + + +class DigitalTwinPlugin(QDesignerCustomWidgetInterface): # pragma: no cover + def __init__(self): + super().__init__() + self._form_editor = None + + def createWidget(self, parent): + if parent is None: + return QWidget() + t = DigitalTwin(parent) + return t + + def domXml(self): + return DOM_XML + + def group(self): + return "" + + def icon(self): + return designer_material_icon(DigitalTwin.ICON_NAME) + + def includeFile(self): + return "digital_twin" + + 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 "DigitalTwin" + + def toolTip(self): + return "DigitalTwin" + + def whatsThis(self): + return self.toolTip() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/panels/__init__.py b/superxas_bec/bec_widgets/widgets/digital_twin/panels/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py b/superxas_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py new file mode 100644 index 0000000..c0ae458 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/panels/input_panel.py @@ -0,0 +1,186 @@ +""" +Panel for user inputs of the digital twin widget +""" + +from typing import Union + +# pylint: disable=E0611 +from qtpy.QtWidgets import QVBoxLayout, QWidget + +from ..types import BeamlineId +from ..widgets.qt_widgets import Button, ComboBox, Group, InputNumberField, NumberIndicator + + +class InputPanel(QWidget): + """ + Panel for user inputs of the digital twin widget + + Args: + beamline (BeamlineId): Beamline id type + """ + + def __init__(self, beamline: BeamlineId, parent=None): + super().__init__(parent) + self._layout = QVBoxLayout(self) + self._layout.setContentsMargins(4, 4, 4, 4) + self._layout.setSpacing(4) + + # Adapt to reality + self.adapt_reality = Button(label_button="Adapt to reality", enabled=True) + + # Energy + self.energy = InputNumberField( + "energy", "Energy", unit="eV", init=8979, decimals=0, single_step=100, ll=4000, hl=65000 + ) + + # FE Slits Acceptance + self.sldi_hacc = InputNumberField( + "h_acc", + "Horizontal", + unit="mrad", + prefix="±", + init=0.25, + decimals=3, + single_step=0.01, + ll=-0.1, + hl=0.9, + ) + self.sldi_vacc = InputNumberField( + "v_acc", + "Vertical", + unit="mrad", + prefix="±", + init=0.1, + decimals=3, + single_step=0.01, + ll=-0.1, + hl=0.5, + ) + self.sldi_ass_group = Group("FE Slits Acceptance", [self.sldi_hacc, self.sldi_vacc]) + + # Collimating mirror + self.cm_stripe = ComboBox("cm_stripe", "Stripe", ["Si", "Rh", "Pt"]) + self.cm_pitch = InputNumberField( + "cm_pitch", + "Pitch", + unit="mrad", + init=-2.391, + decimals=3, + single_step=0.01, + ll=-4.6, + hl=-1.2, + ) + self.cm_pitch_critical = NumberIndicator("Critical Pitch", "mrad", decimals=3) + self.cm_refl = NumberIndicator("Reflectivity at x eV", "%", decimals=0) + self.cm_refl_harm = NumberIndicator("Reflectivity at x eV", "%", decimals=0) + self.cm_ass_group = Group( + "Collimating Mirror", + [ + self.cm_stripe, + self.cm_pitch, + self.cm_pitch_critical, + self.cm_refl, + self.cm_refl_harm, + ], + ) + + # Monochromator + self.mo1_mode = ComboBox("mo1_mode", "Mode", ["Monochromatic", "Pinkbeam"]) + self.mo1_xtal = ComboBox("mo1_xtal", "Crystal", ["Si(111)", "Si(311)"]) + self.mo1_bragg_angle = NumberIndicator("Bragg Angle", "deg", decimals=1) + self.mo1_eres = NumberIndicator("Energy Resolution", "eV", decimals=2) + self.mo1_ass_group = Group( + "Monochromator", [self.mo1_mode, self.mo1_xtal, self.mo1_bragg_angle, self.mo1_eres] + ) + + # Focusing Mirror + stripes: dict[BeamlineId, list[str]] = { + BeamlineId.X01DA: ["Rh (toroid)", "Rh (flat)", "Pt (toroid)", "Pt (flat)"], + BeamlineId.X10DA: ["Rh (toroid)", "Pt (toroid)"], + } + self.fm_stripe = ComboBox("fm_stripe", "Stripe", stripes[beamline]) + self.fm_focus = ComboBox("fm_focus", "Focus Type", ["Manual", "Focused", "Defocused"]) + self.fm_rotx = InputNumberField( + "fm_rotx", + "Incidence Angle", + unit="mrad", + init=-2.391, + decimals=3, + single_step=0.01, + ll=-10, + hl=2, + ) + self.fm_focx = InputNumberField( + "fm_focx", + "Beam Size Horizontal", + unit="mm", + init=1, + decimals=1, + single_step=0.1, + ll=0, + hl=30, + ) + self.fm_focy = InputNumberField( + "fm_focy", + "Beam Size Vertical", + unit="mm", + init=1, + decimals=1, + single_step=0.1, + ll=0, + hl=10, + ) + self.fm_rotx_ideal = NumberIndicator("Incidence Angle for focused beam", "mrad", decimals=3) + self.fm_refl = NumberIndicator("Reflectivity at x eV", "%", decimals=0) + self.fm_refl_harm = NumberIndicator("Reflectivity at x eV", "%", decimals=0) + self.fm_ass_group = Group( + "Focusing Mirror", + [ + self.fm_stripe, + self.fm_focus, + self.fm_rotx, + self.fm_focx, + self.fm_focy, + self.fm_rotx_ideal, + self.fm_refl, + self.fm_refl_harm, + ], + ) + + # Sample + self.cm_fm_harm_suppr = NumberIndicator("Total Suppression Factor at x eV", "", decimals=0) + self.smpl = self._create_smpl(beamline) + + # Assemble complete assistant group + self.input_group = Group( + "User Input", + [ + self.adapt_reality, + self.energy, + self.sldi_ass_group, + self.cm_ass_group, + self.mo1_ass_group, + self.fm_ass_group, + self.cm_fm_harm_suppr, + self.smpl, + ], + ) + + self._layout.addWidget(self.input_group) + self._layout.addStretch() + + def _create_smpl(self, beamline: BeamlineId) -> Union[InputNumberField, ComboBox]: + match beamline: + case BeamlineId.X01DA: + return InputNumberField( + "smpl", + "Sample Position", + unit="mm", + init=23511, + decimals=0, + single_step=100, + ll=23000, + hl=30000, + ) + case BeamlineId.X10DA: + return ComboBox("smpl", "Sample Position", ["ES1", "ES2"]) diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py b/superxas_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py new file mode 100644 index 0000000..1e1ba99 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/panels/mover_panel.py @@ -0,0 +1,230 @@ +""" +Panel to move an axis to a certain position +""" + +from typing import Literal + +# pylint: disable=E0611 +from qtpy.QtWidgets import QVBoxLayout, QWidget + +from ..widgets.move_widget import AbsorberWidget, MoveWidget +from ..widgets.qt_widgets import Group + + +class MoverPanel(QWidget): + """ "Panel to move an axis to a certain position""" + + def __init__(self, dev, parent=None): + super().__init__(parent) + self._layout = QVBoxLayout(self) + self._layout.setContentsMargins(4, 4, 4, 4) + self._layout.setSpacing(4) + + self.mover_widgets = [] + + # FE Slits + self.sldi_gapx = MoveWidget( + dev=dev, motor="sldi_gapx", label="GAPX", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.sldi_gapx) + + self.sldi_gapy = MoveWidget( + dev=dev, motor="sldi_gapy", label="GAPY", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.sldi_gapy) + + self.sldi_mov_group = Group("FE Slits", [self.sldi_gapx, self.sldi_gapy]) + + # Absorber + self.abs = AbsorberWidget(absorber=dev.abs, label="") + + self.abs_group = Group("Absorber", [self.abs]) + + # Collimating mirror + self.cm_trx = MoveWidget( + dev=dev, motor="cm_trx", label="TRX", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.cm_trx) + + self.cm_try = MoveWidget( + dev=dev, motor="cm_try", label="TRY", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.cm_try) + + self.cm_bnd = MoveWidget( + dev=dev, motor="cm_bnd", label="BENDER", unit="km", decimals=2, deadband=0.2 + ) + self.mover_widgets.append(self.cm_bnd) + + self.cm_rotx = MoveWidget( + dev=dev, motor="cm_rotx", label="PITCH", unit="mrad", decimals=3, deadband=0.01 + ) + self.mover_widgets.append(self.cm_rotx) + + self.cm_mov_group = Group( + "Collimating Mirror", [self.cm_trx, self.cm_try, self.cm_bnd, self.cm_rotx] + ) + + # Monochromator + self.mo1_bragg_angle = MoveWidget( + dev=dev, + motor="mo1_bragg_angle", + label="Bragg Angle", + unit="deg", + decimals=3, + deadband=0.01, + ) + self.mover_widgets.append(self.mo1_bragg_angle) + + self.mo1_trx = MoveWidget( + dev=dev, motor="mo1_trx", label="TRX", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.mo1_trx) + + self.mo1_try = MoveWidget( + dev=dev, motor="mo1_try", label="TRY", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.mo1_try) + + self.mo1_mov_group = Group( + "Monochromator", [self.mo1_bragg_angle, self.mo1_trx, self.mo1_try] + ) + + # OP Slits 1 + self.sl1_centery = MoveWidget( + dev=dev, motor="sl1_centery", label="CENTERY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.sl1_centery) + + self.sl1_gapy = MoveWidget( + dev=dev, motor="sl1_gapy", label="GAPY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.sl1_gapy) + + self.sl1_mov_group = Group("OP Slits 1", [self.sl1_centery, self.sl1_gapy]) + + # OP Beam Monitor 1 + self.bm1_try = MoveWidget( + dev=dev, motor="bm1_try", label="TRY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.bm1_try) + + self.bm1_mov_group = Group("OP Beam Monitor 1", [self.bm1_try]) + + # Focusing Mirror + self.fm_trx = MoveWidget( + dev=dev, motor="fm_trx", label="TRX", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.fm_trx) + + self.fm_try = MoveWidget( + dev=dev, motor="fm_try", label="TRY", unit="mm", decimals=2, deadband=0.01 + ) + self.mover_widgets.append(self.fm_try) + + self.fm_bnd = MoveWidget( + dev=dev, motor="fm_bnd", label="BENDER", unit="km", decimals=2, deadband=0.2 + ) + self.mover_widgets.append(self.fm_bnd) + + self.fm_rotx = MoveWidget( + dev=dev, motor="fm_rotx", label="PITCH", unit="mrad", decimals=3, deadband=0.01 + ) + self.mover_widgets.append(self.fm_rotx) + + self.fm_roty = MoveWidget( + dev=dev, motor="fm_roty", label="YAW", unit="mrad", decimals=3, deadband=0.01 + ) + self.mover_widgets.append(self.fm_roty) + + self.fm_rotz = MoveWidget( + dev=dev, motor="fm_rotz", label="ROLL", unit="mrad", decimals=3, deadband=0.01 + ) + self.mover_widgets.append(self.fm_rotz) + + self.fm_mov_group = Group( + "Focusing Mirror", + [self.fm_trx, self.fm_try, self.fm_bnd, self.fm_rotx, self.fm_roty, self.fm_rotz], + ) + + # OP Slits 2 + self.sl2_centery = MoveWidget( + dev=dev, motor="sl2_centery", label="CENTERY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.sl2_centery) + + self.sl2_gapy = MoveWidget( + dev=dev, motor="sl2_gapy", label="GAPY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.sl2_gapy) + + self.sl2_mov_group = Group("OP Slits 2", [self.sl2_centery, self.sl2_gapy]) + + # OP Beam Monitor 2 + self.bm2_try = MoveWidget( + dev=dev, motor="bm2_try", label="TRY", unit="mm", decimals=2, deadband=0.1 + ) + self.mover_widgets.append(self.bm2_try) + + self.bm2_mov_group = Group("OP Beam Monitor 2", [self.bm2_try]) + + # Optical Table + self.ot_try = MoveWidget( + dev=dev, motor="ot_try", label="TRY", unit="mm", decimals=2, deadband=0.2 + ) + self.mover_widgets.append(self.ot_try) + + self.ot_rotx = MoveWidget( + dev=dev, motor="ot_rotx", label="ROTX", unit="mrad", decimals=3, deadband=0.05 + ) + self.mover_widgets.append(self.ot_rotx) + + self.ot_mov_group = Group("Optical Table", [self.ot_try, self.ot_rotx]) + + # Experimental Station 0 + self.es0wi_try = MoveWidget( + dev=dev, motor="es0wi_try", label="ES0 WI", unit="mm", decimals=0, deadband=0.1 + ) + self.mover_widgets.append(self.es0wi_try) + + self.es0_mov_group = Group("Experimental Station 0", [self.es0wi_try]) + + # Experimental Station 1 + self.ot_es1_trz = MoveWidget( + dev=dev, motor="ot_es1_trz", label="ES1 TRZ", unit="mm", decimals=0, deadband=5 + ) + self.mover_widgets.append(self.ot_es1_trz) + + self.es1_mov_group = Group("Experimental Station 1", [self.ot_es1_trz]) + + # Assemble complete mover group + self.mover_group = Group( + "Mover", + [ + self.sldi_mov_group, + self.abs_group, + self.cm_mov_group, + self.mo1_mov_group, + self.sl1_mov_group, + self.bm1_mov_group, + self.fm_mov_group, + self.sl2_mov_group, + self.bm2_mov_group, + self.ot_mov_group, + self.es0_mov_group, + self.es1_mov_group, + ], + ) + + self._layout.addWidget(self.mover_group) + self._layout.addStretch() + + def apply_theme(self, theme: Literal["dark", "light"]): + """ + Apply the theme + + Args: + theme (str): Theme, either "dark" or "light" + """ + for widget in self.mover_widgets: + widget.apply_theme(theme) diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/panels/plots.py b/superxas_bec/bec_widgets/widgets/digital_twin/panels/plots.py new file mode 100644 index 0000000..0080ef0 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/panels/plots.py @@ -0,0 +1,332 @@ +""" +Two plot classes to plot side-view and surface-view +""" + +from typing import Literal, Optional, cast + +import numpy as np +import pyqtgraph as pg +from bec_lib import bec_logger + +# pylint: disable=E0611 +from qtpy.QtCore import Qt +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 ..types import DataDict, SurfaceDict +from ..widgets.qt_widgets import Group + +logger = bec_logger.logger + + +class SurfacePlots(QWidget): + """Plot widget with two curves and legend.""" + + def __init__(self, parent=None): + super().__init__(parent=parent) + self._layout = QHBoxLayout(self) + self._layout.setContentsMargins(4, 4, 4, 4) + self._layout.setSpacing(6) + + self.surfaces: dict[str, SurfaceDict] = { + "assistant": { + "cm": {"x": [], "y": []}, + "mo1_1": {"x": [], "y": []}, + "mo1_2": {"x": [], "y": []}, + "fm": {"x": [], "y": []}, + }, + "reality": { + "cm": {"x": [], "y": []}, + "mo1_1": {"x": [], "y": []}, + "mo1_2": {"x": [], "y": []}, + "fm": {"x": [], "y": []}, + }, + } + + self.plots = {"fm": {}, "mo1_2": {}, "mo1_1": {}, "cm": {}} + + self.color_impenetrable = (0, 0, 0) + self.colors = [(255, 255, 0), (255, 0, 255)] + self.text_color = (255, 255, 255) + + # Create plot widgets + for name, widget in self.plots.items(): + plot_widget = pg.PlotWidget() + plot_widget.getAxis("bottom").enableAutoSIPrefix(False) + + plot_group = Group("Surface " + name, [plot_widget]) + + plot_widget.setLabel("left", "Z [mm]") + plot_widget.setLabel("bottom", "X [mm]") + plot_widget.setMouseEnabled(x=False, y=False) + plot_widget.setMenuEnabled(False) + plot_widget.hideButtons() + + widget["widget"] = plot_widget + self._layout.addWidget(plot_group) + + # Create surfaces + for idx, scene in enumerate(self.surfaces): + for name, _ in self.surfaces[scene].items(): + if scene == "assistant": + brush = QBrush(QColor(*self.colors[idx], 255), Qt.BrushStyle.DiagCrossPattern) + pen = pg.mkPen( + QColor(*self.colors[idx], 255), width=1, style=Qt.PenStyle.DashLine + ) + z_value = 2 + else: + brush = QBrush(QColor(*self.colors[idx], 255)) + pen = pg.mkPen(QColor(*self.colors[idx], 255), width=1) + z_value = 1 + widget = self.plots[name] + self.plots[name][scene] = widget["widget"].plot( + [], [], pen=pen, name=scene, brush=brush, fillLevel=0 + ) + self.plots[name][scene].setZValue(z_value) + + self.walls = [] + self.texts = [] + + self.plot_walls() + + 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 + + bg_color = pg.getConfigOption("background") + fg_color = pg.getConfigOption("foreground") + for _, plot in self.plots.items(): + # Background + plot["widget"].setBackground(bg_color) + # Axes (tick marks, tick labels, axis line) + for axis in ["left", "bottom", "right", "top"]: + ax = plot["widget"].getAxis(axis) + ax.setPen(pg.mkPen(color=fg_color)) + ax.setTextPen(pg.mkPen(color=fg_color)) + + if theme == "light": + self.color_impenetrable = (30, 30, 30) + self.colors = [(79, 163, 224), (240, 128, 60)] + self.text_color = (255, 255, 255) + else: # dark theme + self.color_impenetrable = (180, 180, 180) + self.colors = [(26, 111, 173), (212, 83, 10)] + self.text_color = (0, 0, 0) + + for idx, scene in enumerate(self.surfaces): + for name, _ in self.surfaces[scene].items(): + if scene == "assistant": + brush = QBrush(QColor(*self.colors[idx], 255), Qt.BrushStyle.DiagCrossPattern) + pen = pg.mkPen( + QColor(*self.colors[idx], 255), width=1, style=Qt.PenStyle.DashLine + ) + else: + brush = QBrush(QColor(*self.colors[idx], 255)) + pen = pg.mkPen(QColor(*self.colors[idx], 255), width=0) + self.plots[name][scene].setPen(pen) + self.plots[name][scene].setBrush(brush) + + for wall in self.walls: + wall.setPen(pg.mkPen(color=self.color_impenetrable, width=2)) + wall.setBrush(QBrush(QColor(*self.color_impenetrable))) + + for text in self.texts: + text.setColor(self.text_color) + + def plot_walls(self): + """Plot walls""" + + def plot_surface(widget, surfaces): + for name, surface in surfaces.items(): + rect = QGraphicsRectItem(*surface) + rect.setBrush(QBrush(QColor(*self.color_impenetrable))) + rect.setPen(pg.mkPen(color=self.color_impenetrable, width=2)) + widget.addItem(rect) + text = pg.TextItem(name, color=self.text_color, anchor=(0.5, 0.5)) + widget.addItem(text) + text.setPos(surface[0] + surface[2] / 2, surface[1] + surface[3] / 2) + text.setZValue(10) + self.walls.append(rect) + self.texts.append(text) + + for name, plot in self.plots.items(): + if name == "cm": + plot_surface(plot["widget"], mirror_surface_geometries("cm")) + elif name == "mo1_1": + plot_surface(plot["widget"], mo_surface_geometries("mo1", 0)) + elif name == "mo1_2": + plot_surface(plot["widget"], mo_surface_geometries("mo1", 1)) + elif name == "fm": + plot_surface(plot["widget"], mirror_surface_geometries("fm_flat")) + plot_surface(plot["widget"], mirror_surface_geometries("fm_toroid")) + else: + raise ValueError(f"Plot {name} not found!") + for name, plot in self.plots.items(): + plot["widget"].disableAutoRange() + + def update_surfaces(self, scene: Literal["assistant", "reality"], data: SurfaceDict): + """Update the curves of the plot + + Args: + scene (str): The scene to update, either "assistant" or "reality". + data (DataDict): The new data to plot, with keys "x" and "y", + each containing a list of values. + """ + self.surfaces[scene] = data + for name, device in self.surfaces[scene].items(): + device = cast(DataDict, device) + plot = self.plots[name][scene] + x = np.array(device["x"] + [device["x"][0]]) if len(device["x"]) != 0 else np.array([]) + y = np.array(device["y"] + [device["y"][0]]) if len(device["y"]) != 0 else np.array([]) + plot.setData(x=x, y=y) + + +class SideviewPlot(QWidget): + """Plot widget with two curves and legend.""" + + def __init__(self, parent=None): + super().__init__(parent=parent) + self._layout = QVBoxLayout(self) + self._layout.setContentsMargins(4, 4, 4, 4) + self._layout.setSpacing(0) + + self.plot_widget = pg.PlotWidget() + self.plot_widget.getAxis("bottom").enableAutoSIPrefix(False) + self.plot_widget.invertX(True) + self.plot_widget.addLegend() + + self.color_impenetrable = (0, 0, 0) + self.colors = [(255, 255, 0), (255, 0, 255)] + + self.data: dict[str, DataDict] = { + "assistant": {"x": [0, 1000, 2000], "y": [0, 20, 30]}, + "reality": {"x": [0, 1000, 2000], "y": [0, 15, 50]}, + } + + self.plots = {} + + self.pipes = [] + self.walls = [] + + for idx, scene in enumerate(self.data.keys()): + if scene == "assistant": + pen = pg.mkPen(color=self.colors[idx], width=2, style=Qt.PenStyle.DotLine) + z_value = 2 + else: + pen = pg.mkPen(color=self.colors[idx], width=2) + z_value = 1 + self.plots[scene] = self.plot_widget.plot([], [], pen=pen, name=scene) + self.plots[scene].setZValue(z_value) + + self.plot_group = Group("Side View", [self.plot_widget]) + + self.plot_widget.setLabel("left", "Height [mm]") + self.plot_widget.setLabel("bottom", "Distance [mm]") + self.plot_widget.setMouseEnabled(x=False, y=False) + self.plot_widget.setXRange(0, 25000, 0.1) # pylint: disable=E1121 # type: ignore + self.plot_widget.setYRange(-20, 120, 0.1) # pylint: disable=E1121 # type: ignore + self.plot_widget.setMenuEnabled(False) + self.plot_widget.hideButtons() + + self._layout.addWidget(self.plot_group) + + self.plot_vacuum_pipes() + self.plot_walls() + + 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 + + bg_color = pg.getConfigOption("background") + fg_color = pg.getConfigOption("foreground") + # Background + self.plot_widget.setBackground(bg_color) + # Axes (tick marks, tick labels, axis line) + for axis in ["left", "bottom", "right", "top"]: + ax = self.plot_widget.getAxis(axis) + ax.setPen(pg.mkPen(color=fg_color)) + ax.setTextPen(pg.mkPen(color=fg_color)) + + if theme == "light": + self.color_impenetrable = (30, 30, 30) + self.colors = [(79, 163, 224), (240, 128, 60)] + self.text_color = (255, 255, 255) + else: # dark theme + self.color_impenetrable = (180, 180, 180) + self.colors = [(26, 111, 173), (212, 83, 10)] + self.text_color = (0, 0, 0) + + for idx, scene in enumerate(self.data): + if scene == "assistant": + brush = QBrush(QColor(*self.colors[idx], 255), Qt.BrushStyle.DiagCrossPattern) + pen = pg.mkPen(QColor(*self.colors[idx], 255), width=3, style=Qt.PenStyle.DashLine) + else: + brush = QBrush(QColor(*self.colors[idx], 255)) + pen = pg.mkPen(QColor(*self.colors[idx], 255), width=3) + self.plots[scene].setPen(pen) + self.plots[scene].setBrush(brush) + + for wall in self.walls: + wall.setPen(pg.mkPen(color=self.color_impenetrable, width=3)) + wall.setBrush(QBrush(QColor(*self.color_impenetrable))) + + for pipe in self.pipes: + pipe.setPen(pg.mkPen(color=self.color_impenetrable, width=3)) + + def plot_vacuum_pipes(self): + """Plot vacuum pipes""" + pipes = pipe_geometries() + for pipe in pipes: + self.pipes.append( + self.plot_widget.plot( + x=pipe["x"], y=pipe["y"], pen=pg.mkPen(color=self.color_impenetrable, width=2) + ) + ) + + def plot_walls(self): + """Plot walls""" + walls = wall_geometries() + for wall in walls: + rect = QGraphicsRectItem(wall[0], wall[1], wall[2], wall[3]) + rect.setBrush(QBrush(QColor(*self.color_impenetrable))) + rect.setPen(pg.mkPen(color=self.color_impenetrable, width=2)) + self.plot_widget.addItem(rect) + self.walls.append(rect) + + def update_curves(self, scene: Literal["assistant", "reality"], data: DataDict): + """Update the curves of the plot + + Args: + scene (str): The scene to update, either "assistant" or "reality". + data (DataDict): The new data to plot, with keys "x" and "y", + each containing a list of values. + """ + self.data[scene] = data + plot = self.plots[scene] + plot.setData(x=self.data[scene]["x"], y=self.data[scene]["y"]) diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/panels/settings_panel.py b/superxas_bec/bec_widgets/widgets/digital_twin/panels/settings_panel.py new file mode 100644 index 0000000..2bf74d2 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/panels/settings_panel.py @@ -0,0 +1,31 @@ +""" +Settings panel for the digital twin widget +""" + +# pylint: disable=E0611 +from qtpy.QtWidgets import QVBoxLayout, QWidget + +from ..widgets.qt_widgets import Button, Group, TextIndicator + + +class SettingsPanel(QWidget): + """Settings panel for the digital twin widget""" + + def __init__(self, parent=None): + super().__init__(parent) + self._layout = QVBoxLayout(self) + self._layout.setContentsMargins(4, 4, 4, 4) + self._layout.setSpacing(4) + + # Reload offsets + self.load_offsets = Button(label="Load Offsets", label_button="Load", enabled=True) + self.offsets_status = TextIndicator(label="Offsets") + self.show_offsets = Button(label="Show Offsets", label_button="Show", enabled=True) + + # Assemble complete offset group + self.offset_group = Group( + "Axes Offsets", [self.load_offsets, self.offsets_status, self.show_offsets] + ) + + self._layout.addWidget(self.offset_group) + self._layout.addStretch() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/register_digital_twin.py b/superxas_bec/bec_widgets/widgets/digital_twin/register_digital_twin.py new file mode 100644 index 0000000..024c888 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/register_digital_twin.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 .digital_twin_plugin import DigitalTwinPlugin + + QPyDesignerCustomWidgetCollection.addCustomWidget(DigitalTwinPlugin()) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/types.py b/superxas_bec/bec_widgets/widgets/digital_twin/types.py new file mode 100644 index 0000000..d5c0393 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/types.py @@ -0,0 +1,83 @@ +"""Types used for the beamline config and for plotting data""" + +from enum import Enum +from typing import TypedDict + + +class BeamlineId(str, Enum): + """ + Identifier for supported beamlines. + """ + + X01DA = "x01da" + X10DA = "x10da" + + +class ConfigDict(TypedDict): + """ + Typed dictionary representing the beamline configuration. + + Attributes: + energy (float): Beam energy. + h_acc (float): Horizontal acceptance. + v_acc (float): Vertical acceptance. + cm_pitch (float): CM pitch angle. + cm_stripe (str): CM stripe name. + cm_trx (float): CM translation x. + mo1_mode (str): MO1 mode. + mo1_xtal (str): MO1 crystal. + mo1_bragg (float): MO1 Bragg angle. + fm_rotx (float): FM rotation x. + fm_stripe (str): FM stripe name. + fm_trx (float): FM translation x. + fm_qy (float): FM qy value. + fm_gain_height (int): FM gain height. + smpl (float): Sample value. + """ + + energy: float + h_acc: float + v_acc: float + cm_pitch: float + cm_stripe: str + cm_trx: float + mo1_mode: str + mo1_xtal: str + mo1_bragg: float + fm_rotx: float + fm_stripe: str + fm_trx: float + fm_qy: None | float + fm_gain_height: int + smpl: float + + +class DataDict(TypedDict): + """ + Typed dictionary representing plot data. + + Attributes: + x (list[float]): List of x-axis values. + y (list[float]): List of y-axis values. + """ + + x: list + y: list + + +class SurfaceDict(TypedDict): + """ + Typed dictionary representing the surfaces of a scene, + grouping plot data by surface type. + + Attributes: + cm (DataDict): Data for the cm surface. + mo1_1 (DataDict): Data for the mo1_1 surface. + mo1_2 (DataDict): Data for the mo1_2 surface. + fm (DataDict): Data for the fm surface. + """ + + cm: DataDict + mo1_1: DataDict + mo1_2: DataDict + fm: DataDict diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/widgets/__init__.py b/superxas_bec/bec_widgets/widgets/digital_twin/widgets/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py b/superxas_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py new file mode 100644 index 0000000..4bd529c --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/widgets/move_widget.py @@ -0,0 +1,591 @@ +"""Move widget to display an axis and also move it through BEC""" + +import threading +import time +from typing import Literal, Optional + +from bec_lib import bec_logger +from bec_qthemes import material_icon +from bec_widgets.utils.colors import get_accent_colors + +# pylint: disable=E0611 +from qtpy.QtCore import Property # type: ignore[attr-defined] +from qtpy.QtCore import Signal # type: ignore[attr-defined] +from qtpy.QtCore import QObject, QPropertyAnimation, Qt, QThread +from qtpy.QtGui import QTransform +from qtpy.QtWidgets import QApplication, QHBoxLayout, QLabel, QPushButton, QWidget + +# pylint: disable=E0402 +from .....devices.absorber import STATUS as ABS_STATUS + +logger = bec_logger.logger + + +class Status: + """Status class for the axis""" + + IN_POSITION = "in_position" # green mdi.check-circle + NOT_IN_POSITION = "not_in_position" # orange mdi.close-circle + MOVING = "moving" # blue mdi.loading (spinning) + ERROR = "error" # red mdi.alert-circle + + +class StatusIcon(QWidget): + """ + Displays a status icon using bec_qthemes Material Design Icons. + Handles its own spin animation for the MOVING state via QPropertyAnimation. + """ + + ICON_SIZE = 20 + + _ICON_MAP = { + Status.IN_POSITION: ("check_circle", "#27ae60"), + Status.NOT_IN_POSITION: ("cancel", "#e6d922"), + Status.ERROR: ("warning", "#e74c3c"), + Status.MOVING: ("cycle", "#2980b9"), + } + + def __init__(self, parent=None): + super().__init__(parent=parent) + self._status = None + self._rotation = 0.0 + + self._label = QLabel(self) + self._label.setFixedSize(self.ICON_SIZE, self.ICON_SIZE) + self._label.setAlignment(Qt.AlignmentFlag.AlignCenter) + self.setFixedSize(self.ICON_SIZE, self.ICON_SIZE) + + self._spin_anim = QPropertyAnimation(self, b"rotation") # type: ignore[call-arg] + self._spin_anim.setStartValue(0) + self._spin_anim.setEndValue(360) + self._spin_anim.setDuration(1000) + self._spin_anim.setLoopCount(-1) # Loop indefinitely + + self.set_status(Status.NOT_IN_POSITION) + + def get_rotation(self) -> float: + """ + Return the current rotation angle in degrees. + + Returns: + float: Rotation angle in deg + """ + return self._rotation + + def set_rotation(self, angle: float): + """ + Set the rotation angle and update the displayed pixmap. + + Rotates the current base pixmap around its center point using a smooth + transformation. Has no effect on the display if no base pixmap is set. + + Args: + angle (float): Rotation angle in degrees, clockwise. + """ + self._rotation = angle + if self._current_pixmap_base is not None: + cx = self._current_pixmap_base.width() / 2 + cy = self._current_pixmap_base.height() / 2 + t = QTransform().translate(cx, cy).rotate(angle).translate(-cx, -cy) + self._label.setPixmap( + self._current_pixmap_base.transformed(t, Qt.TransformationMode.SmoothTransformation) + ) + + rotation = Property(float, get_rotation, set_rotation) # type: ignore[call-arg] + + def set_status(self, status: str): + """ + Update the widget's status and refresh the displayed icon accordingly. + + Looks up the icon name and color associated with the given status from + ``_ICON_MAP``, renders a new pixmap, and starts or stops the spin + animation depending on whether the status is ``Status.MOVING``. Returns + early without any updates if the status has not changed. + + Args: + status (str): The new status value. Must be a key in ``_ICON_MAP``. + """ + if status == self._status: + return + self._status = status + + icon_name, color = self._ICON_MAP[status] + icon = material_icon( + icon_name, size=(self.ICON_SIZE, self.ICON_SIZE), color=color, convert_to_pixmap=True + ) + self._current_pixmap_base = icon + + if status == Status.MOVING: + self._spin_anim.start() + else: + self._spin_anim.stop() + self._label.setPixmap(icon) + + +class MotionWorker(QObject): + """ + Executes motion on the specified motor and includes some safety during + motion for certain motors. + """ + + position_changed = Signal(float) + error = Signal() + finished = Signal() + + def __init__(self, dev, motor, target_pos: float): + super().__init__() + self.dev = dev + self.motor = motor + self._target = target_pos + self._stop_flag = threading.Event() + + def stop(self): + """Sets the stop flag""" + self._stop_flag.set() + + def run(self): + """Prepares the movement based on the axis (motor)""" + match self.motor: + case "sldi_gapx" | "sldi_gapy" | "sldi_centerx" | "sldi_centery": + self.motion() + case "cm_trx": + self.motion( + abs_closed=True, + surveyed_axes=[{"device": self.dev["cm_roty"], "abs_tol": 0.05}], + ) + case "cm_roty": + self.motion( + abs_closed=True, surveyed_axes=[{"device": self.dev["cm_trx"], "abs_tol": 0.05}] + ) + case "cm_try": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["cm_rotx"], "abs_tol": 0.05}, + {"device": self.dev["cm_rotz"], "abs_tol": 0.05}, + ], + ) + case "cm_rotx": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["cm_try"], "abs_tol": 0.05}, + {"device": self.dev["cm_rotz"], "abs_tol": 0.05}, + ], + ) + case "cm_rotz": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["cm_try"], "abs_tol": 0.05}, + {"device": self.dev["cm_rotx"], "abs_tol": 0.05}, + ], + ) + case "cm_bnd": + p1 = ( + 1 / (self.dev.cm_bnd_radius.read()["cm_bnd_radius"]["value"] * 1e3) + 0.0284 + ) / 2e-6 + p2 = (1 / (self._target * 1e3) + 0.0284) / 2e-6 + self._target = p2 - p1 + self.motion(relative=True, rb={"device": self.dev["cm_bnd_radius"]}) + case "mo1_try" | "mo1_trx" | "mo1_roty": + self.motion(abs_closed=True) + case "mo1_bragg_angle": + self.motion() + case "sl1_centery" | "sl1_gapy" | "bm1_try": + self.motion() + case "fm_trx": + self.motion( + abs_closed=True, + surveyed_axes=[{"device": self.dev["fm_roty"], "abs_tol": 0.05}], + ) + case "fm_roty": + self.motion( + abs_closed=True, surveyed_axes=[{"device": self.dev["fm_trx"], "abs_tol": 0.05}] + ) + case "fm_try": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["fm_rotx"], "abs_tol": 0.05}, + {"device": self.dev["fm_rotz"], "abs_tol": 0.05}, + ], + ) + case "fm_rotx": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["fm_try"], "abs_tol": 0.05}, + {"device": self.dev["fm_rotz"], "abs_tol": 0.05}, + ], + ) + case "fm_rotz": + self.motion( + abs_closed=True, + surveyed_axes=[ + {"device": self.dev["fm_try"], "abs_tol": 0.05}, + {"device": self.dev["fm_rotx"], "abs_tol": 0.05}, + ], + ) + case "fm_bnd": + p1 = ( + 1 / (self.dev.fm_bnd_radius.read()["fm_bnd_radius"]["value"] * 1e3) + 4.28e-5 + ) / 1.84e-9 + p2 = (1 / (self._target * 1e3) + 4.28e-5) / 1.84e-9 + self._target = p2 - p1 + self.motion(relative=True, rb={"device": self.dev["fm_bnd_radius"]}) + case "sl2_centery" | "sl2_gapy" | "bm2_try": + self.motion() + case "ot_try" | "ot_rotx" | "ot_es1_trz": + self.motion() + case _: + logger.warning(f"Motor {self.motor} not integrated in digital twin!") + + def motion(self, abs_closed: bool = False, relative: bool = False, rb=None, surveyed_axes=None): + """ + Moves an axis while surverying a set of axes (if set). + Example surveyed_axes: + [{'device': bec_device_object, 'abs_tol': 0.1},] + + Args: + surveyed_axes (list): List of dictionaries of devices + """ + 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: + surv_ax["name"] = surv_ax["device"].dotted_name + surv_ax["old_value"] = surv_ax["device"].read(cached=True)[surv_ax["name"]]["value"] + if rb is not None: + rb["name"] = rb["device"].dotted_name + status = self.dev[self.motor].move(self._target, relative=relative) + last_check = time.time() + update_interval = 0.1 + while status.status == "RUNNING": + now = time.time() + if time.time() - last_check < update_interval: + time.sleep(0.01) + last_check = now + if self._stop_flag.is_set(): + self.dev[self.motor].stop() + self._stop_flag.clear() + if rb is not None: + self.position_changed.emit(rb["device"].read(cached=True)[rb["name"]]["value"]) + else: + self.position_changed.emit( + self.dev[self.motor].read(cached=True)[self.motor]["value"] + ) + if surveyed_axes is not None: + for surv_ax in surveyed_axes: + fb = surv_ax["device"].read(cached=True)[surv_ax["name"]]["value"] + if abs(fb - surv_ax["old_value"]) > surv_ax["abs_tol"]: + self.dev[self.motor].stop() + self.error.emit() + break + self.finished.emit() + + +class MoveWidget(QWidget): + """ + One motor stage control group containing: + - Target label (target position) + - Feedback label (current position) + - Status icon (bec_qthemes) + - Start / Stop button + """ + + def __init__(self, dev, motor, label: str = "", unit=None, decimals=3, deadband=0.0): + super().__init__() + self.fb = 0.0 + self.target = 0 + self.dev = dev + self.motor = motor + self.deadband = deadband + self.status = Status.IN_POSITION + self._thread: QThread | None = None + self._worker: MotionWorker | None = None + + self.text_color = (0, 0, 0) + + self.unit = unit + self.decimals = decimals + + layout = QHBoxLayout(self) + layout.setContentsMargins(4, 0, 4, 0) + layout.setSpacing(4) + + # Name + self.label = QLabel(label) + self.label.setFixedWidth(76) + self.label.setWordWrap(True) + layout.addWidget(self.label) + + # Target + self.target_label = QLabel("-") + self.target_label.setFixedWidth(84) + layout.addWidget(self.target_label) + + # Feedback + self.fb_label = QLabel("-") + self.fb_label.setFixedWidth(84) + layout.addWidget(self.fb_label) + + # Status icon + self.status_icon = StatusIcon() + self.status_icon.setFixedWidth(24) + layout.addWidget(self.status_icon) + + # Start / Stop button + self.btn_action = QPushButton("Move") + self.btn_action.setFixedWidth(64) + self.btn_action.setFixedHeight(20) + self.btn_action.clicked.connect(self._on_button_clicked) + layout.addWidget(self.btn_action) + self.btn_mode = "start" + + self._apply_button_style("start") + + 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 + + if theme == "light": + self.text_color = {"target": (79, 163, 224), "fb": (240, 128, 60)} + else: # dark theme + self.text_color = {"target": (26, 111, 173), "fb": (212, 83, 10)} + r, g, b = self.text_color["target"] + self.target_label.setStyleSheet(f"QLabel {{color: rgb({r}, {g}, {b})}}") + r, g, b = self.text_color["fb"] + self.fb_label.setStyleSheet(f"QLabel {{color: rgb({r}, {g}, {b})}}") + + if self.btn_mode == "start": + self.btn_action.setStyleSheet( + "QPushButton " + + f"{{background-color: {get_accent_colors().success.name()}; color: white;}}" + ) + else: + self.btn_action.setStyleSheet( + "QPushButton " + + f"{{background-color: {get_accent_colors().emergency.name()}; color: white;}}" + ) + + def set_target(self, target): + """Change the target value in the ui""" + self.target = target + text = f"{target:.{int(self.decimals)}f}" + if self.unit is not None: + text = text + " " + self.unit + self.target_label.setText(text) + self._on_target_or_fb_changed() + + def set_feedback(self, fb): + """Change the feedback value in the ui""" + if self.status != Status.MOVING: + self.fb = fb + text = f"{fb:.{int(self.decimals)}f}" + if self.unit is not None: + text = text + " " + self.unit + self.fb_label.setText(text) + self._on_target_or_fb_changed() + + def _apply_button_style(self, mode: str): + """Apply a button style depending on if the button shows start or stop""" + self.btn_mode = mode + if mode == "start": + self.btn_action.setText("Move") + self.btn_action.setStyleSheet( + "QPushButton " + + f"{{background-color: {get_accent_colors().success.name()}; color: white;}}" + ) + else: # stop + self.btn_action.setText("Stop") + self.btn_action.setStyleSheet( + "QPushButton " + + f"{{background-color: {get_accent_colors().emergency.name()}; color: white;}}" + ) + + def _set_status(self, status: str): + """Set the current status icon in the ui""" + self.status = status + self.status_icon.set_status(status) + + def _on_target_or_fb_changed(self): + """Re-evaluate in-position status whenever the target value changes.""" + if self.status in (Status.ERROR, Status.MOVING): + return + if abs(self.fb - self.target) <= self.deadband: + self._set_status(Status.IN_POSITION) + else: + self._set_status(Status.NOT_IN_POSITION) + + def _on_button_clicked(self): + """Starts or stops motion depending on current situation""" + if self._thread and self._thread.isRunning(): + self._stop_motion() + else: + self._start_motion() + + def _start_motion(self): + """Start a motion""" + target = self.target + if abs(target - self.fb) <= self.deadband: + self._set_status(Status.IN_POSITION) + return + + self._set_status(Status.MOVING) + self._apply_button_style("stop") + + self._worker = MotionWorker(self.dev, self.motor, target) + self._thread = QThread() + self._worker.moveToThread(self._thread) + + self._thread.started.connect(self._worker.run) + self._worker.position_changed.connect(self._on_position_changed) + self._worker.error.connect(self._on_error) + self._worker.error.connect(self._thread.quit) + self._worker.finished.connect(self._on_motion_finished) + self._worker.finished.connect(self._thread.quit) + self._thread.finished.connect(self._cleanup_thread) + + self._thread.start() + + def _on_error(self): + """Called when an error occurs""" + self._set_status(Status.ERROR) + self._apply_button_style("start") + + def _stop_motion(self): + """Attempts to stop the motion""" + if self._worker: + self._worker.stop() + + def _on_position_changed(self, pos: float): + """Change the feedback value in the ui""" + self.fb = pos + text = f"{pos:.{int(self.decimals)}f}" + if self.unit is not None: + text = text + " " + self.unit + self.fb_label.setText(text) + + def _on_motion_finished(self): + """Finished a movement""" + target = self.target + if self.status != Status.ERROR: + if abs(self.fb - target) <= self.deadband: + self._set_status(Status.IN_POSITION) + else: + self._set_status(Status.NOT_IN_POSITION) + self._apply_button_style("start") + + def _cleanup_thread(self): + """Cleaning up of the mover thread""" + if self._thread: + self._thread.deleteLater() + self._thread = None + if self._worker: + self._worker.deleteLater() + self._worker = None + + def shutdown(self): + """Cleaning up of the mover when shutting down the application""" + if self._worker: + self._worker.stop() + if self._thread: + self._thread.quit() + self._thread.wait(2000) # max 2 s grace period + + +class AbsorberWidget(QWidget): + """ + Control of the frontend absorber (only open) + """ + + def __init__(self, absorber, label: str = "Absorber"): + super().__init__() + self.absorber = absorber + self.fb = False + self.text_color = (0, 0, 0) + + layout = QHBoxLayout(self) + layout.setContentsMargins(4, 0, 4, 0) + layout.setSpacing(4) + + # Name + self.label = QLabel(label) + self.label.setFixedWidth(76) + self.label.setWordWrap(True) + layout.addWidget(self.label) + + # Blank + self.blank_label = QLabel("") + self.blank_label.setFixedWidth(84) + layout.addWidget(self.blank_label) + + # Feedback + self.fb_label = QLabel("-") + self.fb_label.setFixedWidth(84) + layout.addWidget(self.fb_label) + + # Blank icon + self.blank_icon = QLabel("") + self.blank_icon.setFixedWidth(24) + layout.addWidget(self.blank_icon) + + # Open + self.btn_action = QPushButton("Open") + self.btn_action.setFixedWidth(64) + self.btn_action.setFixedHeight(20) + self.btn_action.clicked.connect(self._on_button_clicked) + layout.addWidget(self.btn_action) + + def set_feedback(self, fb: bool): + """ + Displays the status of the absober in the ui + + Args: + fb (bool): True will set the button to Open, False to Closed + """ + self.fb = fb + if fb: + self.fb_label.setText("Open") + self.fb_label.setStyleSheet(f"QLabel {{color: {get_accent_colors().success.name()}}}") + else: + self.fb_label.setText("Closed") + self.fb_label.setStyleSheet(f"QLabel {{color: {get_accent_colors().emergency.name()}}}") + + def enable_open(self, enable: bool = False): + """ + Enable or disable the open/close button + + Args: + enable (bool): Enables and disables the button + """ + if enable: + self.btn_action.setStyleSheet( + "QPushButton " + + f"{{background-color: {get_accent_colors().success.name()}; color: white;}}" + ) + self.btn_action.setEnabled(True) + else: # disabled + self.btn_action.setStyleSheet( + "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}" + ) + self.btn_action.setDisabled(True) + + def _on_button_clicked(self): + """Open absorber""" + self.absorber.open() diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/widgets/qt_widgets.py b/superxas_bec/bec_widgets/widgets/digital_twin/widgets/qt_widgets.py new file mode 100644 index 0000000..b127418 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/widgets/qt_widgets.py @@ -0,0 +1,223 @@ +""" +Universal Qt widgets +""" + +from functools import partial + +from bec_widgets.utils.colors import get_accent_colors +from qtpy.QtCore import Qt + +# pylint: disable=E0611 +from qtpy.QtGui import QFont +from qtpy.QtWidgets import ( + QApplication, + QComboBox, + QDoubleSpinBox, + QGroupBox, + QHBoxLayout, + QLabel, + QPushButton, + QVBoxLayout, + QWidget, +) + +LABEL_WIDTH = 118 +ROW_MARGINS = (4, 0, 4, 0) +ROW_SPACING = 6 + + +class Group(QGroupBox): + def __init__(self, label, widgets): + super().__init__(label) + self.layout = QVBoxLayout(self) # type: ignore + self.layout.setContentsMargins(6, 6, 6, 6) + self.layout.setSpacing(4) + for widget in widgets: + self.layout.addWidget(widget) # type: ignore + + +class NumberIndicator(QWidget): + def __init__(self, label="", unit=None, highlight=False, decimals=3): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(*ROW_MARGINS) + layout.setSpacing(ROW_SPACING) + self.label = QLabel(label) + self.label.setFixedWidth(LABEL_WIDTH) + self.label.setWordWrap(True) + layout.addWidget(self.label) + self.val = QLabel("-") + self.val.setAlignment(Qt.AlignTop) # type: ignore + layout.addWidget(self.val) + self.unit = unit + self.highlight = highlight + self.decimals = decimals + self.number = 0 + if highlight: + font = QFont() + font.setBold(True) + font.setPointSize(14) + self.label.setFont(font) + self.val.setFont(font) + + def value(self) -> float: + return self.number + + def setLabel(self, label) -> None: + self.label.setText(label) + + def setValue(self, number): + self.number = number + text = f"{number:.{int(self.decimals)}f}" + if self.unit is not None: + text = text + " " + self.unit + self.val.setText(text) + + +class InputNumberField(QWidget): + def __init__( + self, + identifier="", + label="", + unit=None, + prefix=None, + init=0.0, + decimals=1, + single_step=0.1, + ll=-1e6, + hl=1e6, + ): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(*ROW_MARGINS) + layout.setSpacing(ROW_SPACING) + self.identifier = identifier + self.label = QLabel(label) + self.label.setFixedWidth(LABEL_WIDTH) + self.label.setWordWrap(True) + layout.addWidget(self.label) + self.val = QDoubleSpinBox() + self.val.setRange(ll, hl) + self.val.setDecimals(decimals) + self.val.setSingleStep(single_step) + self.val.setValue(init) + if unit is not None: + self.val.setSuffix(" " + unit) + if prefix is not None: + self.val.setPrefix(prefix + " ") + layout.addWidget(self.val) + + def set_number(self, number): + self.val.setValue(number) + + def has_focus(self) -> bool: + return self.val.hasFocus() + + def value(self) -> float: + return self.val.value() + + def value_changed_connect(self, func): + """Connect a function to the Enter/Return key press.""" + self.val.valueChanged.connect( + partial( + func, identifier=self.identifier, value_obj=self.val, value=lambda: self.val.value() + ) + ) + + +class ComboBox(QWidget): + def __init__(self, identifier="", label="", enums=None): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(*ROW_MARGINS) + layout.setSpacing(ROW_SPACING) + self.identifier = identifier + self.label = QLabel(label) + self.label.setFixedWidth(LABEL_WIDTH) + self.label.setWordWrap(True) + layout.addWidget(self.label) + self.value = QComboBox() + for entry in enums or []: + self.value.addItem(entry) + layout.addWidget(self.value) + + def set_current_text(self, text): + self.value.setCurrentText(text) + + def currentText(self) -> str: + return self.value.currentText() + + def has_focus(self) -> bool: + return QApplication.focusWidget() is self.value.view() + + def activated_connect(self, func): + """Connect a function to the Enter/Return key press.""" + self.value.activated.connect( + partial( + func, + identifier=self.identifier, + value_obj=self.value, + value=lambda: self.value.currentText(), + ) + ) + + def setDisabled(self, disable): + self.value.setDisabled(disable) + + +class Button(QWidget): + def __init__(self, label=None, label_button: str = "", enabled=False): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(*ROW_MARGINS) + layout.setSpacing(ROW_SPACING) + if label is not None: + self.label = QLabel(label) + self.label.setFixedWidth(LABEL_WIDTH) + layout.addWidget(self.label) + self.button = QPushButton(label_button) + self.enable_button(enabled) + layout.addWidget(self.button) + + def clicked_connect(self, func): + """Connect a function to the button press.""" + self.button.clicked.connect(func) + + def enable_button(self, enable: bool = False): + if enable: + self.button.setStyleSheet( + f"QPushButton {{background-color: {get_accent_colors().default.name()}; color: white;}}" + ) + self.button.setEnabled(True) + else: # disabled + self.button.setStyleSheet( + "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}" + ) + self.button.setDisabled(True) + + def setText(self, text): + self.button.setText(text) + + +class TextIndicator(QWidget): + def __init__(self, label): + super().__init__() + layout = QHBoxLayout(self) + layout.setContentsMargins(*ROW_MARGINS) + layout.setSpacing(ROW_SPACING) + self.label = QLabel(label) + self.label.setFixedWidth(LABEL_WIDTH) + self.label.setWordWrap(True) + layout.addWidget(self.label) + self.text = QLabel("-") + self.text.setAlignment(Qt.AlignTop) # type: ignore + layout.addWidget(self.text) + + def setLabel(self, label) -> None: + self.label.setText(label) + + def setText(self, text): + self.text.setText(text) + + def setColor(self, color: str): + self.text.setStyleSheet(f"QLabel {{color:{color}}}") diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/x01da_offsets.yaml b/superxas_bec/bec_widgets/widgets/digital_twin/x01da_offsets.yaml new file mode 100644 index 0000000..f475d96 --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/x01da_offsets.yaml @@ -0,0 +1,50 @@ +cm_try: + offset: 0.15 + +mo1_trx: + modifier: + axis: mo1_trx + range: [[-30, -0.1], [0.1, 30]] + offset: [-2.3, 1.31] + +mo1_try: + modifier: + axis: mo1_trx + range: [[-30, -0.1], [0.1, 30]] + offset: [-1.78, -1.78] + +sl1_centery: + offset: -1.2 + +fm_trx: + modifier: + axis: fm_trx + range: [[-66, -31], [-24, 7], [11, 31], [38, 66]] + offset: [-0.61, 0, 0, -0.16] + +fm_try: + modifier: + axis: fm_trx + range: [[-66, -31], [-24, 7], [11, 31], [38, 66]] + offset: [0.028, 0, 0, -0.45] + +fm_rotx: + modifier: + axis: fm_trx + range: [[-66, -31], [-24, 7], [11, 31], [38, 66]] + offset: [0.027, 0, 0, 0.045] + +fm_roty: + modifier: + axis: fm_trx + range: [[-66, -31], [-24, 7], [11, 31], [38, 66]] + offset: [-0.038, 0, 0, -0.053] + +sl2_centery: + offset: -0.7 + +ot_try: + offset: -0.49 + +ot_rotx: + offset: 0 \ No newline at end of file diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/x01da_parameters.py b/superxas_bec/bec_widgets/widgets/digital_twin/x01da_parameters.py new file mode 100644 index 0000000..214cb7d --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/x01da_parameters.py @@ -0,0 +1,323 @@ +""" +X01DA / Debye Beamline Parameters. +This file describes the parameter of each component of the Debye 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, 5045, sourceHeight), + center2=(0, 5289.5, sourceHeight), + maxDivH=1.8e-3, + maxDivV=0.8e-3, +) + +# FE Window +filt = namedtuple( + "filt", ["name", "center", "pitch", "limPhysX", "limPhysY", "surface", "material", "thickness"] +) + +feWindow = filt( + name="FE-WINDOW", + center=(0.0, 7020, sourceHeight), + pitch=np.pi / 2, + limPhysX=(-6, 6), + limPhysY=(-3.0, 3.0), + surface="None", + material=filterDiamond, + thickness=0.1, +) +feWindow = feWindow._replace(surface=f"CVD Diamond window {feWindow.thickness*1e3:0.0f} $\\mu$m") + +# 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, 6890, sourceHeight], + surface=("Si", "Pt", "Rh"), + material=(stripeSi, stripePt, stripeRh), + limPhysX=(-34, 34), + limPhysY=(-600, 600), + limOptX=((-21, -7, 14), (-11, 11, 23)), + limOptY=((-500, -500, -500), (500, 500, 500)), + R=[3e6, 15e6], + pitch=[-5.0e-3, -0.0e-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, 8815, sourceHeight], opening=[-20.0, 20.0, -20.0 + 12.5, 20.0 + 12.5] +) # left, right, bottom, top + +opWbBsBlock = apertures( + name="OP-WB-BS-BLOCK", center=[0.0, 13860, sourceHeight], opening=[-18.0, 18.0, 25, 85.5] +) # left, right, bottom, top +# opening=[-18., 18., 42, 76], # X10DA + +# 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-MO1", + center=[0.0, 11750, sourceHeight], + xtal=("Si311", "Si111"), + material1=(si311_1, si111_1), + material2=(si311_2, si111_2), + xtalWidth=(24, 24), + xtalOffsetX=(-21.2, 21.2), + xtalLength1=(55, 55), + xtalLength2=(105, 105), + xtalGap=(8, 8), + rotOffset=6, + heightOffset=8.5, + braggLim=[3.6, 33], + jack1=[0.0, 11350.0, 0.0], # Tripod maybe not available! + jack2=[-400.0, 12350.0, 0.0], + jack3=[400.0, 12350.0, 0.0], + tx=0.0, +) # X-Stage [x] + +mo2 = monochromator( + name="OP-CCM2", + center=[0.0, 13250, sourceHeight], + xtal=("Si311", "Si111"), + material1=(si311_1, si111_1), + material2=(si311_2, si111_2), + xtalWidth=(24, 24), + xtalOffsetX=(-21, 21), + xtalLength1=(55, 55), + xtalLength2=(105, 105), + xtalGap=(8, 8), + rotOffset=6, + heightOffset=8.5, + braggLim=[3.6, 33], + jack1=[0.0, 13350.0, 0.0], # Tripod maybe not available! + jack2=[-400.0, 14350.0, 0.0], + jack3=[400.0, 14350.0, 0.0], + tx=0.0, +) # X-Stage [x] + +# OP Slits +op_slits = namedtuple("op_slits", ["name", "center"]) + +opSlits1 = op_slits(name="OP-SLITS 1", center=(0, 14349.6, sourceHeight)) + +opSlits2 = op_slits(name="OP-SLITS 2", center=(0, 18134.8, sourceHeight)) + +# OP Beam Monitors +op_bm = namedtuple("op_bm", ["name", "center"]) + +opBM1 = op_bm(name="OP Beam Monitor 1", center=(0, 14599.6, sourceHeight)) + +opBM2 = op_bm(name="OP Beam Monitor 2", center=(0, 18384.8, sourceHeight)) + +# Focusing mirror +focusingMirror = namedtuple( + "focusingMirror", + [ + "name", + "center", + "surfaceToroid", + "materialToroid", + "surfaceFlat", + "materialFlat", + "limPhysXToroid", + "limPhysYToroid", + "limPhysXFlat", + "limPhysYFlat", + "limOptXToroid", + "limOptYToroid", + "limOptXFlat", + "limOptYFlat", + "R", + "pitch", + "r", + "xToroid", + "xFlat", + "hToroid", + "jack1", + "jack2", + "jack3", + "tx1", + "tx2", + ], +) + +fm = focusingMirror( + name="OP-FM", + center=[0.0, 15670, sourceHeight], # nominal height 58 mm above ring, SLS1! + surfaceToroid=("Rh", "Pt"), + materialToroid=(stripeRh, stripePt), + surfaceFlat=("Rh", "Pt"), + materialFlat=(stripeRh, stripePt), + limPhysXToroid=(-79.0, 79.0), + limPhysYToroid=(-575.0, 575.0), + limPhysXFlat=(-79.0, 79.0), + limPhysYFlat=(-575.0, 575.0), + limOptXToroid=((-38, 66), (-66, 31)), + limOptYToroid=((-500.0, -500.0), (500.0, 500.0)), + limOptXFlat=((-11.45, 23.55), (-30.45, -6.45)), + limOptYFlat=((-500.0, -500.0), (500.0, 500.0)), + R=[3e6, 15e6], + pitch=[-5.0e-3, 0e-3], + r=[35.510, 24.986], + xToroid=[-52, 48.5], # offset in local x + xFlat=[-20.95, 8.55], + hToroid=[2.88, 7.15], # depth of the cylinder at x = xCylinder1 and x = xCylinder2. + jack1=[-130.0, 15535 - 538.0, 0.0], + jack2=[130.0, 15535 + 538.0, 0.0], + jack3=[0.0, 15535 + 538.0, 0.0], + tx1=[0.0, -575.0], # X-Stage 1 [x, y] + tx2=[0.0, 575.0], +) # X-Stage 2 [x, y] + +# EH Window +ehWindow = filt( + name="EH-WINDOW", + center=(0.0, 19998.3, sourceHeight), + pitch=np.pi / 2, + limPhysX=(-20.0, 20.0), + limPhysY=(-4, 4), + surface="None", + material=filterSi3N4, + thickness=0.002, +) +ehWindow = ehWindow._replace(surface=f"Beryllium window {ehWindow.thickness*1e3:0.0f} $\\mu$m") + +# Sample +sample = namedtuple("sample", ["name", "center"]) + +smpl = sample(name="EH-SMPL", center=[0, 23365, sourceHeight]) + +smpl2 = sample(name="EH-SMPL2", center=[0, 27500, sourceHeight]) + +tables = {} + +# Vacuum pipes +# DN40CF ID = 35 mm oder 37 mm +# DN50CF ID = 47.5 mm +# DN63CF ID = 60.2 mm oder 66 mm +# DN100CF ID = 97.4 mm oder 104 mm +pipe = namedtuple("pipes", ["center", "diameter", "start", "end"]) +vacuum_pipes = pipe( + center=[27.5, (37.5 + 27.5) / 2, 37.5, 62.5, 72.5], + diameter=[97.4, 97.4, 97.4, 97.4, 97.4], + start=[10952.88, 11750 + 250, mo2.center[1] + 250, 14000, fm.center[1]], + end=[11750 - 250, mo2.center[1] - 250, 14000, fm.center[1], ehWindow.center[1]], +) + +Walls = namedtuple("walls", ["start", "end", "height"]) +walls = Walls(start=[13999.30], end=[13999 + 75.5 + 30], height=[[-20, 25]]) diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/x10da_offsets.yaml b/superxas_bec/bec_widgets/widgets/digital_twin/x10da_offsets.yaml new file mode 100644 index 0000000..e69de29 diff --git a/superxas_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py b/superxas_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py new file mode 100644 index 0000000..a933b3e --- /dev/null +++ b/superxas_bec/bec_widgets/widgets/digital_twin/x10da_parameters.py @@ -0,0 +1,269 @@ +""" +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 +slits = namedtuple("slits", ["name", "center", "maxDivH", "maxDivV"]) + +feSlits = slits(name="FE-SLITS", center=(0, 5290, 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, 7618, sourceHeight], + surface=("Si", "Pt", "Rh"), + material=(stripeSi, stripePt, stripeRh), + limPhysX=(-30, 30), + limPhysY=(-600, 600), + limOptX=((-21, -8, 5), (-11, 2, 21)), + 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 + +opSlits = apertures( + name="OP-SLITS", center=[0, 14145 - 135, sourceHeight], opening=[-35 / 2, 35 / 2, 47.5, 82.5] +) + +# 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", + ], +) + +fm = focusingMirror( + name="OP-FM", + center=[0.0, 15580 - 135, sourceHeight], + surfaceToroid=("Rh (toroid)", "Pt (toroid)"), + materialToroid=(stripeRh, stripePt), + limPhysXToroid=(-54.0, 54.0), + limPhysYToroid=(-565.0, 565.0), + limOptXToroid=((4.865, -40.882), (43.388, -4.865)), + 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, -22, 874], # 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] + +ehWindow = filt( + name="EH-WINDOW", + center=(0.0, 22225 - 135, 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"]) + +smpl = sample(name="OP-SMPL", center=[0, 24000 - 135, sourceHeight]) + +ES1 = sample(name="EH-ES1", center=[0, 24000, sourceHeight]) +ES2 = sample(name="EH-ES2", center=[0, 25000, sourceHeight])