wip widget

This commit is contained in:
x01da
2026-08-10 21:27:06 +02:00
parent 76cde8903a
commit 688a2fa89d
@@ -1,17 +1,34 @@
import sys
from functools import partial
from typing import Literal, Optional
import numpy as np
import pyqtgraph as pg
import xraydb
from bec_lib import bec_logger
from bec_lib.endpoints import MessageEndpoints
from bec_widgets.utils.bec_dispatcher import BECDispatcher
from bec_widgets.utils.colors import apply_theme
from bec_widgets.utils.colors import Colors, apply_theme, get_accent_colors
from bec_widgets.utils.error_popups import SafeSlot
from bec_widgets.widgets.control.scan_control.scan_control import ScanControl
from qtpy.QtCore import Qt, QTimer, Signal
from qtpy.QtCore import Qt
# pylint: disable=E0611
from qtpy.QtWidgets import QApplication, QDoubleSpinBox, QPushButton, QVBoxLayout, QWidget
from qtpy.QtWidgets import (
QApplication,
QComboBox,
QDialog,
QFrame,
QGridLayout,
QGroupBox,
QHBoxLayout,
QLabel,
QPushButton,
QScrollArea,
QSizePolicy,
QVBoxLayout,
QWidget,
)
from ....devices.mo1_bragg.mo1_bragg_utils import compute_spline
@@ -24,11 +41,16 @@ SHOW_MOTION_PROFILE = [
"xas_advanced_scan_with_xrd",
]
# TODO: Changing edge trigger change self.edge_energy. Theme change edge marker. Name of edge marker label should be element and edge
class ScanControlAdvanced(ScanControl):
def __init__(self, *args, **kwargs):
self._motion_profile_group = QGroupBox("Motion Profile Preview")
motion_profile_layout = QVBoxLayout(self._motion_profile_group)
self.motion_profile_widget = MotionProfileWidget(self)
motion_profile_layout.addWidget(self.motion_profile_widget)
super().__init__(*args, **kwargs)
self.show_motion_profile(self._selected_scan)
@@ -45,7 +67,7 @@ class ScanControlAdvanced(ScanControl):
self.motion_profile_widget.update_plot(d_spacing=d_spacing)
def _add_metadata_form(self):
self.layout.addWidget(self.motion_profile_widget)
self.layout.addWidget(self._motion_profile_group, stretch=1)
super()._add_metadata_form()
self.connect_scan_parameter_signals()
@@ -81,48 +103,222 @@ class ScanControlAdvanced(ScanControl):
pass
def scan_parameter_changed(self, *_):
# logger.info(f"Fe K edge: {xraydb.xray_edge("Fe", "K", energy_only=True)}")
params = self.get_scan_parameters()[1]
logger.info(f"Scan parameters: {params}")
# logger.info(f"Scan parameters: {params}")
self.motion_profile_widget.update_plot(
scan_name=self._selected_scan, scan_parameters=params
)
def apply_theme(self, theme: Literal["dark", "light"]):
"""
Apply the theme
Args:
theme (str): Theme, either "dark" or "light"
"""
self.motion_profile_widget.apply_theme(theme)
PLOT_RESOLUTION = 100
H = 6.62606957e-34
E = 1.602176634e-19
C = 299792458
ELEMENTS = {
# period 4
"Sc": (4, 3),
"Ti": (4, 4),
"V": (4, 5),
"Cr": (4, 6),
"Mn": (4, 7),
"Fe": (4, 8),
"Co": (4, 9),
"Ni": (4, 10),
"Cu": (4, 11),
"Zn": (4, 12),
"Ga": (4, 13),
"Ge": (4, 14),
"As": (4, 15),
"Se": (4, 16),
"Br": (4, 17),
"Kr": (4, 18),
# period 5
"Rb": (5, 1),
"Sr": (5, 2),
"Y": (5, 3),
"Zr": (5, 4),
"Nb": (5, 5),
"Mo": (5, 6),
"Tc": (5, 7),
"Ru": (5, 8),
"Rh": (5, 9),
"Pd": (5, 10),
"Ag": (5, 11),
"Cd": (5, 12),
"In": (5, 13),
"Sn": (5, 14),
"Sb": (5, 15),
"Te": (5, 16),
"I": (5, 17),
"Xe": (5, 18),
# period 6
"Cs": (6, 1),
"Ba": (6, 2),
"La": (6, 3),
"Ce": (8, 4),
"Pr": (8, 5),
"Nd": (8, 6),
"Pm": (8, 7),
"Sm": (8, 8),
"Eu": (8, 9),
"Gd": (8, 10),
"Tb": (8, 11),
"Dy": (8, 12),
"Ho": (8, 13),
"Er": (8, 14),
"Tm": (8, 15),
"Pb": (6, 14),
"Bi": (6, 15),
"Po": (6, 16),
"At": (6, 17),
"Rn": (6, 18),
# period 7
"Fr": (7, 1),
"Ra": (7, 2),
"Ac": (9, 3),
"Th": (9, 4),
"Pa": (9, 5),
"U": (9, 6),
}
class PeriodicTableDialog(QDialog):
def __init__(self, parent=None):
super().__init__(parent)
self.selected = None
layout = QGridLayout(self)
for element, (row, col) in ELEMENTS.items():
button = QPushButton(element)
button.setFixedSize(42, 32)
button.clicked.connect(lambda checked=False, e=element: self.select(e))
layout.addWidget(button, row - 1, col - 1)
def select(self, element):
self.selected = element
self.accept()
class ElementSelector(QWidget):
def __init__(self):
super().__init__()
self.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
layout = QVBoxLayout(self)
layout.setContentsMargins(*ROW_MARGINS)
layout.setSpacing(ROW_SPACING)
self.label = QLabel("Element")
layout_selection = QHBoxLayout(self)
self.button = QPushButton("Select")
self.button.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
self.button.setStyleSheet(
f"QPushButton {{background-color: {get_accent_colors().default.name()}; color: white;}}"
)
self.selection = QLabel("")
layout_selection.addWidget(self.button)
layout_selection.addWidget(self.selection)
layout.addWidget(self.label)
layout.addLayout(layout_selection)
self.button.clicked.connect(self.open_table)
def open_table(self):
dialog = PeriodicTableDialog(self)
if dialog.exec_():
self.selection.setText(dialog.selected)
def currentText(self):
return self.selection.text()
class MotionProfileWidget(QWidget):
parameters_changed = Signal(dict)
def __init__(self, parent):
super().__init__()
self.parent = parent
layout = QVBoxLayout(self)
self._parent_widget = parent
outer_layout = QVBoxLayout(self)
outer_layout.setContentsMargins(0, 0, 0, 0)
self.plot_pos = pg.PlotWidget()
self.plot_pos.setLabel("left", "Energy [eV]")
self.plot_pos.setLabel("bottom", "Time [s]")
self.plot_pos.setMouseEnabled(x=False, y=False)
self.plot_pos.setMenuEnabled(False)
self.plot_pos.hideButtons()
self.element = ElementSelector()
self.edge = ComboBox("Edge", "Edge", ["K", "L1", "L2", "L3"])
self.plot_vel = pg.PlotWidget()
self.plot_vel.setLabel("left", "Scan Velocity [eV/s]")
self.plot_vel.setLabel("bottom", "Energy [eV]")
self.plot_vel.setMouseEnabled(x=False, y=False)
self.plot_vel.setMenuEnabled(False)
self.plot_vel.hideButtons()
plot_container = QWidget()
plot_layout = QVBoxLayout(plot_container)
plot_layout.setContentsMargins(0, 0, 0, 0)
self.plot_pos = self._make_plot(left_label="Energy<br>[eV]", bottom_label="Time [s]")
self.plot_vel = self._make_plot(
left_label="Scan Velocity<br>[eV/s]", bottom_label="Energy [eV]"
)
self.edge_energy = 8000
plot_layout.addWidget(self.plot_pos)
plot_layout.addWidget(self.plot_vel)
plot_container.setMinimumHeight(
self.plot_pos.minimumHeight() + self.plot_vel.minimumHeight()
)
scroll = QScrollArea()
scroll.setWidgetResizable(True)
scroll.setFrameShape(QFrame.NoFrame)
scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
scroll.setWidget(plot_container)
self.edge_selection = QHBoxLayout()
self.edge_selection.addWidget(self.element)
self.edge_selection.addWidget(self.edge)
self.edge_selection.addStretch(1)
outer_layout.addLayout(self.edge_selection)
outer_layout.addWidget(scroll, stretch=1)
layout.addWidget(self.plot_pos)
layout.addWidget(self.plot_vel)
self.scan_name = None
self.scan_parameters = None
self.d_spacing = None
@staticmethod
def _make_plot(left_label: str, bottom_label: str) -> pg.PlotWidget:
plot = pg.PlotWidget()
plot.setLabel("left", left_label)
plot.setLabel("bottom", bottom_label)
plot.setMouseEnabled(x=False, y=False)
plot.setMenuEnabled(False)
plot.hideButtons()
plot.getAxis("left").enableAutoSIPrefix(False)
plot.getAxis("bottom").enableAutoSIPrefix(False)
plot.setMinimumHeight(220)
plot.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.MinimumExpanding)
return plot
def _add_edge_marker(self, plot, x_arr, y_arr, energy_arr):
if self.edge_energy is None:
return
idx = int(np.argmin(np.abs(np.asarray(energy_arr) - self.edge_energy)))
x_marker = x_arr[idx]
y_marker = y_arr[idx]
scatter = pg.ScatterPlotItem(
[x_marker], [y_marker], symbol="star", size=14, brush=pg.mkBrush("r"), pen=pg.mkPen("r")
)
plot.addItem(scatter, ignoreBounds=True)
label = pg.TextItem("edge", anchor=(0, 1), color="r")
label.setPos(x_marker, y_marker)
plot.addItem(label, ignoreBounds=True)
def update_plot(self, scan_name=None, scan_parameters=None, d_spacing=None):
if scan_name is not None:
self.scan_name = scan_name
@@ -146,17 +342,20 @@ class MotionProfileWidget(QWidget):
start_angle = self.energy_to_angle(self.scan_parameters["stop"])
stop_angle = self.energy_to_angle(self.scan_parameters["start"])
logger.info(f"start: {start_angle}, stop: {stop_angle}")
if "xas_simple_scan" in self.scan_name:
amp = (stop_angle - start_angle) / 2
two_pi_f = 2 * np.pi * 1 / (2 * self.scan_parameters["scan_time"])
pos = (start_angle + stop_angle) / 2 + amp * np.cos(x_time * two_pi_f)
logger.info(f"pos: {pos}")
vel = two_pi_f * amp * np.sin(x_time * two_pi_f)
energy = self.angle_to_energy(pos)
vel_e = self.velocity_to_energy_per_s(pos, vel)
self.plot_pos.plot(x_time, energy)
self._add_edge_marker(self.plot_pos, x_time, energy, energy)
self.plot_vel.plot(energy, -vel_e)
self._add_edge_marker(self.plot_vel, energy, -vel_e, energy)
if "xas_advanced_scan" in self.scan_name:
if (
self.scan_parameters["e_kink"] < self.scan_parameters["start"]
@@ -176,8 +375,43 @@ class MotionProfileWidget(QWidget):
x_time = np.cumsum(t) / 1000
energy = self.angle_to_energy(pos)
vel_e = self.velocity_to_energy_per_s(pos, vel)
self.plot_pos.plot(-np.flip(x_time) + x_time[-1], np.flip(energy))
time_flipped = -np.flip(x_time) + x_time[-1]
energy_flipped = np.flip(energy)
self.plot_pos.plot(time_flipped, energy_flipped)
self._add_edge_marker(self.plot_pos, time_flipped, energy_flipped, energy_flipped)
self.plot_vel.plot(energy, -vel_e)
self._add_edge_marker(self.plot_vel, energy, -vel_e, energy)
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.
"""
logger.info(f"theme update called with theme: {theme}")
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.plot_pos, self.plot_vel):
plot.setBackground(bg_color)
for axis in ("left", "bottom", "right", "top"):
ax = plot.getAxis(axis)
ax.setPen(pg.mkPen(color=fg_color))
ax.setTextPen(pg.mkPen(color=fg_color))
curves = plot.listDataItems()
colors = Colors.golden_angle_color(colormap="plasma", num=len(curves), format="HEX")
for curve, color in zip(curves, colors):
curve.setPen(pg.mkPen(color=color, width=2))
def energy_to_angle(self, energy: int | float):
if self.d_spacing is None:
@@ -212,6 +446,59 @@ class MotionProfileWidget(QWidget):
) * velocity
LABEL_WIDTH = 118
ROW_MARGINS = (4, 0, 4, 0)
ROW_SPACING = 6
class ComboBox(QWidget):
def __init__(self, identifier="", label="", enums=None):
super().__init__()
layout = QVBoxLayout(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)
self.label.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Preferred)
layout.addWidget(self.label)
self.value = QComboBox()
self.value.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
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)
if __name__ == "__main__":
app = QApplication(sys.argv)
apply_theme("light")