From b14f2c0fe3dec387b9d4b6918f350f697d0aa5c2 Mon Sep 17 00:00:00 2001 From: x01da Date: Mon, 4 May 2026 06:52:48 +0200 Subject: [PATCH] wip: digital twin --- .../widgets/digital_twin/calc_surfaces.py | 119 +++++ .../widgets/digital_twin/digital_twin.py | 451 ++++++++++++++++-- debye_bec/bec_widgets/widgets/qt_widgets.py | 3 + 3 files changed, 535 insertions(+), 38 deletions(-) create mode 100644 debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py diff --git a/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py b/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py new file mode 100644 index 0000000..25a129b --- /dev/null +++ b/debye_bec/bec_widgets/widgets/digital_twin/calc_surfaces.py @@ -0,0 +1,119 @@ +import os +import re +import numpy as np +from bec_lib import bec_logger + +os.environ["USE_XRT"] = "False" +import debye_bec.bec_widgets.widgets.digital_twin.x01da_parameters as bl + +def calc_surfaces(cfg): + + out = { + '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 = (bl.cm.limOptX[0][index] + bl.cm.limOptX[1][index]) / 2 + + 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) + + xtalPos = bl.mo1.xtalOffsetX[index] + xtalLength1 = bl.mo1.xtalLength1[index] + xtalLength2 = bl.mo1.xtalLength2[index] + + widthBeam = 2 * bl.mo1.center[1] * np.tan(cfg['h_acc']) + + heightBeam = 2 * bl.cm.center[1] * np.tan(cfg['v_acc']) + w = heightBeam / np.sin(cfg['mo1_bragg']) + + if cfg['mo1_mode'] in 'Monochromatic': + out['mo1_1']['x'] = [xtalPos-widthBeam/2, xtalPos+widthBeam/2, xtalPos+widthBeam/2, xtalPos-widthBeam/2] + out['mo1_1']['y'] = [xtalLength1/2-c-w/2, xtalLength1/2-c-w/2, xtalLength1/2-c+w/2, xtalLength1/2-c+w/2] + out['mo1_2']['x'] = [xtalPos-widthBeam/2, xtalPos+widthBeam/2, xtalPos+widthBeam/2, xtalPos-widthBeam/2] + out['mo1_2']['y'] = [-xtalLength2/2+e-w/2, -xtalLength2/2+e-w/2, -xtalLength2/2+e+w/2, -xtalLength2/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) + off = (bl.fm.limOptXToroid[0][index] + bl.fm.limOptXToroid[1][index]) / 2 + r = bl.fm.r[index] + else: + surface = bl.fm.surfaceFlat + stripe = re.sub(r'\s*\(.*?\)', '', cfg['fm_stripe']).strip() + index = surface.index(stripe) + off = (bl.fm.limOptXFlat[0][index] + bl.fm.limOptXFlat[1][index]) / 2 + r = bl.fm.r[index] + + widthBeam = 2 * bl.fm.center[1] * np.tan(cfg['h_acc']) + + if cfg['fm_stripe'] in ('Rh (toroid)', 'Pt (toroid)'): + + l = heightBeam/np.sin(cfg['fm_pitch']) + alpha = np.arccos(1-widthBeam**2/(2*r**2)) + h = r-(r*np.cos(alpha/2)) + z = h/np.tan(cfg['fm_pitch']) + + x = [off-widthBeam/2, off-widthBeam/2] + y = [l/2-z/2, -l/2-z/2] + + res = 20 + xElipse = np.linspace(0, np.pi, res) + yElipse = np.linspace(0, np.pi, res) + xElipse = [-widthBeam/2*np.cos(i)+off for i in xElipse] + yElipse = [widthBeam*np.sin(i)*z/widthBeam-l/2-z/2 for i in yElipse] + + x.extend(xElipse) + y.extend(yElipse) + + x.extend([off+widthBeam/2, off+widthBeam/2]) + y.extend([-l/2-z/2, l/2-z/2]) + + res = 50 + xElipse = np.linspace(np.pi, 0, res) + yElipse = np.linspace(np.pi, 0, res) + xElipse = [-widthBeam/2*np.cos(i)+off for i in xElipse] + yElipse = [widthBeam*np.sin(i)*z/widthBeam+l/2-z/2 for i in yElipse] + + x.extend(xElipse) + y.extend(yElipse) + + out['fm']['x'] = x + out['fm']['y'] = y + + else: # flat surface, no toroid + l = heightBeam/np.sin(cfg['fm_pitch']) + + w1 = 2 * (bl.fm.center[1]-l/2) * np.tan(cfg['h_acc']) + w2 = 2 * (bl.fm.center[1]+l/2) * np.tan(cfg['h_acc']) + + out['fm']['x'] = [off-w1/2, off+w1/2, off+w2/2, off-w2/2] + out['fm']['y'] = [-l/2, -l/2, l/2, l/2] + + return out diff --git a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py index 97cf118..24258c6 100644 --- a/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py +++ b/debye_bec/bec_widgets/widgets/digital_twin/digital_twin.py @@ -1,11 +1,13 @@ import sys +import re import datetime import numpy as np +from scipy.interpolate import UnivariateSpline from bec_lib import bec_logger # pylint: disable=E0611 from qtpy.QtWidgets import ( QWidget, QVBoxLayout, QHBoxLayout, QLabel, - QApplication, QLayout + QApplication, QLayout, QGroupBox ) # pylint: disable=E0611 from qtpy.QtCore import Qt, QTimer @@ -21,6 +23,7 @@ from debye_bec.bec_widgets.widgets.qt_widgets import InputNumberField, ComboBox, from debye_bec.bec_widgets.widgets.digital_twin.calculate_positions import calc_positions from debye_bec.bec_widgets.widgets.digital_twin.calculate_sideview import calc_sideview +from debye_bec.bec_widgets.widgets.digital_twin.calc_surfaces import calc_surfaces import debye_bec.bec_widgets.widgets.digital_twin.x01da_parameters as bl @@ -43,18 +46,23 @@ class DigitalTwin(BECWidget, QWidget): central = QWidget() self.root_layout = QHBoxLayout(central) + self.plot_widget = QWidget() + self.plot_layout = QVBoxLayout(self.plot_widget) self.input = InputPanel() - self.plot_widget = PlotWidget() + self.sideview_plot = SideviewPlot() + self.surface_plots = SurfacePlots() self.positions = PositionsPanel() self.root_layout.addWidget(self.input, stretch=1, alignment=Qt.AlignTop) # type: ignore + self.plot_layout.addWidget(self.sideview_plot) # type: ignore + self.plot_layout.addWidget(self.surface_plots) # type: ignore self.root_layout.addWidget(self.plot_widget, stretch=1, alignment=Qt.AlignTop) # type: ignore self.root_layout.addWidget(self.positions, stretch=1, alignment=Qt.AlignTop) # type: ignore self.setLayout(self.root_layout) self.setWindowTitle("Digital Twin") - # self.resize(1500, 800) + self.resize(1800, 800) self.input.energy.value_changed_connect(self.calc_bragg_angle) self.input.sldi_hacc.value_changed_connect(self.calc_positions) @@ -76,12 +84,25 @@ class DigitalTwin(BECWidget, QWidget): self.input.fm_stripe.activated_connect(self.calc_ideal_fm_pitch) self.input.smpl.value_changed_connect(self.calc_ideal_fm_pitch) + self.input.energy.value_changed_connect(self.calc_cm_reflectivity) + self.input.cm_pitch.value_changed_connect(self.calc_cm_reflectivity) + self.input.cm_stripe.activated_connect(self.calc_cm_reflectivity) + self.input.fm_pitch.value_changed_connect(self.calc_fm_reflectivity) + self.input.fm_stripe.activated_connect(self.calc_fm_reflectivity) + self.input.energy.value_changed_connect(self.calc_cm_reflectivity) + + self.input.energy.value_changed_connect(self.calc_energy_resolution) + self.input.mo1_xtal.activated_connect(self.calc_energy_resolution) + self.bragg_angle = 0 self.calc_bragg_angle() self.calc_ideal_fm_pitch() self.calc_crit_angle() self.calc_assistant_sideview() self.calc_reality_sideview() + self.calc_cm_reflectivity() + self.calc_fm_reflectivity() + self.calc_energy_resolution() # Timer: update plot every 1 second self._timer = QTimer(self) @@ -90,6 +111,67 @@ class DigitalTwin(BECWidget, QWidget): # TODO: Check if I need to stop the timer if the widget is closed? self._timer.start() + def calc_energy_resolution(self, *args): + xtal = self.input.mo1_xtal.currentText().translate(str.maketrans('', '', '()')) # Remove brackets from xtal name to conform with parameters + index = bl.mo1.xtal.index(xtal) + crystal = bl.mo1.material1[index] + E = self.input.energy.value() + + dtheta = np.linspace(-30, 90, 601) + theta = crystal.get_Bragg_angle(E) + dtheta * 1e-6 + refl = np.abs(crystal.get_amplitude(E, 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) + r1, r2 = spline.roots() + fwhm_rad = (r2 - r1) * 1e-6 # µrad → rad + + # Energy resolution + theta_B = crystal.get_Bragg_angle(E) + dE_over_E = fwhm_rad / np.tan(theta_B) + dE = dE_over_E * E + + 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") + + self.input.mo1_eres.setValue(dE) + + def calc_cm_reflectivity(self, *args): + index = bl.cm.surface.index(self.input.cm_stripe.currentText()) + rs, rp = bl.cm.material[index].get_amplitude( + self.input.energy.value(), + np.sin(-self.input.cm_pitch.value() * 1e-3) + )[0:2] + self.input.cm_refl.setValue(100 * abs(rs)**2) + self.input.cm_refl.setLabel(f"Refl. at {self.input.energy.value():.0f} eV") + rs, rp = bl.cm.material[index].get_amplitude( + 2 * self.input.energy.value(), + np.sin(-self.input.cm_pitch.value() * 1e-3) + )[0:2] + self.input.cm_refl_harm.setValue(100 * abs(rs)**2) + self.input.cm_refl_harm.setLabel(f"Refl. at {2 * self.input.energy.value():.0f} eV") + + def calc_fm_reflectivity(self, *args): + if self.input.fm_stripe.currentText() in ('Rh (toroid)', 'Pt (toroid)'): + surface = bl.fm.surfaceToroid + material = bl.fm.materialToroid + stripe = re.sub(r'\s*\(.*?\)', '', self.input.fm_stripe.currentText()).strip() + index = surface.index(stripe) + else: + surface = bl.fm.surfaceFlat + material = bl.fm.materialFlat + stripe = re.sub(r'\s*\(.*?\)', '', self.input.fm_stripe.currentText()).strip() + index = surface.index(stripe) + rs, rp = material[index].get_amplitude( + self.input.energy.value(), + np.sin(-self.input.fm_pitch.value() * 1e-3) + )[0:2] + self.input.fm_refl.setValue(100 * abs(rs)**2) + self.input.fm_refl.setLabel(f"Refl. at {self.input.energy.value():.0f} eV") + def get_assistant_config(self): config = { # Config in SI units! 'energy' : self.input.energy.value(), @@ -110,24 +192,42 @@ class DigitalTwin(BECWidget, QWidget): # TODO Needs to run in a loop in a separate thread due to the long time it takes to get the values from self.dev... def get_reality_config(self): - if abs(self.dev.mo1_trx.position) > 5: + if abs(self.dev.mo1_trx.read()['mo1_trx']['value']) > 5: mo1_mode = 'Monochromatic' else: mo1_mode = 'Pinkbeam' - # TODO: stripe detection, mo1_bragg and acceptance + mo1_bragg = self.dev.mo1_bragg.read() + sldi_gapx = self.dev.sldi_gapx.read()['sldi_gapx']['value'] + sldi_gapy = self.dev.sldi_gapy.read()['sldi_gapy']['value'] + d1 = bl.feSlits.center1[1] + h_acc = np.tan(sldi_gapx / (2 * d1)) + v_acc = np.tan(sldi_gapy / (2 * d1)) + cm_trx = -self.dev.cm_trx.read()['cm_trx']['value'] + 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 + fm_trx = -self.dev.fm_trx.read()['fm_trx']['value'] + 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)' config = { # Config in SI units! - 'energy' : self.dev.mo1_bragg.position, - 'h_acc' : self.input.sldi_hacc.value() * 1e-3, - 'v_acc' : self.input.sldi_vacc.value() * 1e-3, - 'cm_pitch' : -self.dev.cm_rotx.position * 1e-3, - 'cm_stripe' : self.input.cm_stripe.currentText(), + 'energy' : mo1_bragg['mo1_bragg']['value'], + 'h_acc' : h_acc, + 'v_acc' : v_acc, + 'cm_pitch' : -self.dev.cm_rotx.read()['cm_rotx']['value'] * 1e-3, + 'cm_stripe' : cm_stripe, 'mo1_mode' : mo1_mode, - 'mo1_xtal' : self.dev.mo1_bragg.crystal.current_xtal_string.get(), - 'mo1_bragg' : self.dev.mo1_bragg.angle.get(), - 'fm_pitch' : -self.dev.fm_rotx.position * 1e-3, - 'fm_stripe' : self.input.fm_stripe.currentText(), + 'mo1_xtal' : mo1_bragg['mo1_bragg_crystal_current_xtal_string']['value'], + 'mo1_bragg' : mo1_bragg['mo1_bragg_angle']['value']/180*np.pi, + 'fm_pitch' : -self.dev.fm_rotx.read()['fm_rotx']['value'] * 1e-3, + 'fm_stripe' : fm_stripe, 'fm_gain_height' : 1, - 'smpl' : self.dev.ot_es1_trz.position, + 'smpl' : self.dev.ot_es1_trz.read()['ot_es1_trz']['value'], } # logger.info(f'Config created: {config}') return config @@ -135,17 +235,30 @@ class DigitalTwin(BECWidget, QWidget): @SafeSlot() def calc_assistant_sideview(self, *args): beam = calc_sideview(self.get_assistant_config()) - self.plot_widget.data['assistant'][0] = beam['Z'] - self.plot_widget.data['assistant'][1] = beam['Y'] - self.plot_widget.update_curves() + self.sideview_plot.data['assistant']['x'] = beam['Z'] + self.sideview_plot.data['assistant']['y'] = beam['Y'] + self.sideview_plot.update_curves() @SafeSlot() def calc_reality_sideview(self): - logger.info('Update reality plot') + # logger.info('Update reality plot') beam = calc_sideview(self.get_reality_config()) - self.plot_widget.data['reality'][0] = beam['Z'] - self.plot_widget.data['reality'][1] = beam['Y'] - self.plot_widget.update_curves() + self.sideview_plot.data['reality']['x'] = beam['Z'] + self.sideview_plot.data['reality']['y'] = beam['Y'] + self.sideview_plot.update_curves() + + # TODO Move to different place + self.calc_reality_surfaces() + + @SafeSlot() + def calc_assistant_surfaces(self, *args): + surfaces = calc_surfaces(self.get_assistant_config()) + self.surface_plots.update_surfaces(scene='assistant', data=surfaces) + + @SafeSlot() + def calc_reality_surfaces(self, *args): + surfaces = calc_surfaces(self.get_reality_config()) + self.surface_plots.update_surfaces(scene='reality', data=surfaces) @SafeSlot() def calc_positions(self, *args): @@ -174,6 +287,7 @@ class DigitalTwin(BECWidget, QWidget): # TODO move to somewhere else! self.calc_assistant_sideview() + self.calc_assistant_surfaces() @SafeSlot() def calc_bragg_angle(self, *args): @@ -197,11 +311,20 @@ class DigitalTwin(BECWidget, QWidget): self.bragg_angle = 0 if val > -1 and val < 1: self.bragg_angle = np.asin(val) + + cm_pitch = -self.dev.cm_rotx.read()['cm_rotx']['value'] * 1e-3 + if self.input.mo1_mode.currentText() in 'Monochromatic': + # Add 2x CM pitch to the bragg angle + bragg_angle_cor = ((2 * cm_pitch) + self.bragg_angle) / np.pi * 180 + elif self.input.mo1_mode.currentText() in 'Pinkbeam': + # Align xtal surfaces parallel to beam + bragg_angle_cor = (2 * cm_pitch) / np.pi * 180 + + self.input.mo1_bragg_angle.setValue(bragg_angle_cor) self.calc_positions() @SafeSlot() def update_mono_mode(self, *args): - logger.info(f'Got args {args}') if self.input.mo1_mode.currentText() in 'Monochromatic': self.input.mo1_xtal.setDisabled(False) else: @@ -247,7 +370,7 @@ class InputPanel(QWidget): self._layout.setSizeConstraint(QLayout.SetFixedSize) # type: ignore # Energy - self.energy = InputNumberField('Energy [keV]', init=8979, decimals=0, single_step=100, ll=4000, hl=65000) + self.energy = InputNumberField('Energy [eV]', init=8979, decimals=0, single_step=100, ll=4000, hl=65000) # FE Slits Acceptance self.sldi_hacc = InputNumberField('Horizontal [± mrad]', init=0.25, decimals=2, single_step=0.01, ll=-0.1, hl=0.9) @@ -264,23 +387,31 @@ class InputPanel(QWidget): self.cm_stripe = ComboBox('Stripe', ['Si', 'Rh', 'Pt']) self.cm_pitch_critical = NumberIndicator('Critical Pitch', 'mrad', decimals=3) self.cm_pitch = InputNumberField('Pitch [mrad]', init=-2.391, decimals=3, single_step=0.01, ll=-4.6, hl=-1.2) + self.cm_refl = NumberIndicator('Refl. at x eV', '%', decimals=0) + self.cm_refl_harm = NumberIndicator('Refl. at x eV', '%', decimals=0) self.cm_ass_group = Group( 'Collimating Mirror', [ self.cm_stripe, self.cm_pitch_critical, self.cm_pitch, + self.cm_refl, + self.cm_refl_harm, ] ) # Monochromator self.mo1_mode = ComboBox('Mode', ['Monochromatic', 'Pinkbeam']) self.mo1_xtal = ComboBox('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, ] ) @@ -288,12 +419,14 @@ class InputPanel(QWidget): self.fm_stripe = ComboBox('Stripe', ['Rh (toroid)', 'Rh (flat)', 'Pt (toroid)', 'Pt (flat)']) self.fm_pitch_ideal = NumberIndicator('Ideal Pitch', 'mrad', decimals=3) self.fm_pitch = InputNumberField('Pitch [mrad]', init=-2.391, decimals=3, single_step=0.01, ll=-10, hl=2) + self.fm_refl = NumberIndicator('Refl. at x eV', '%', decimals=0) self.fm_ass_group = Group( 'Focusing Mirror', [ self.fm_stripe, self.fm_pitch_ideal, self.fm_pitch, + self.fm_refl, ] ) @@ -446,7 +579,242 @@ class PositionsPanel(QWidget): self._layout .addWidget(self.position_group) self._layout .addStretch() -class PlotWidget(QWidget): +class SurfacePlots(QWidget): + """Plot widget with two curves and legend.""" + + def __init__(self, parent=None): + super().__init__(parent) + self._layout = QHBoxLayout(self) + # self._layout.setSizeConstraint(QLayout.SetFixedSize) # type: ignore + + self.surfaces = { + '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 = { + 'cm': {}, + 'mo1_1': {}, + 'mo1_2': {}, + 'fm': {}, + } + + app = QApplication.instance() + theme = app.theme.theme # type: ignore + if theme == "light": + self.color_impenetrable = (30, 30, 30) + self.colors = [(79, 163, 224), (240, 128, 60)] + else: # dark theme + self.color_impenetrable = (220, 220, 220) + self.colors = [(26, 111, 173), (212, 83, 10)] + + + # 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.setXRange(0, 25000, padding=0.1) + # plot_widget.setYRange(-20, 120, padding=0.1) + 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, device in self.surfaces[scene].items(): + brush = pg.mkBrush(color=self.colors[idx] + (150,)) + widget = self.plots[name] + self.plots[name][scene] = widget['widget'].plot( + [], + [], + pen=None, + name=scene, + brush=brush, + fillLevel=0, + ) + self.plots[name][scene].setZValue(1) + + # self._layout.addStretch() + + logger.info(f'Created surfaces: {self.surfaces}') + logger.info(f'Created plots: {self.plots}') + + self.plot_walls() + # self.update_curves() + + def plot_walls(self): + + def plot_mirror_stripe(widget, surface, limOptX, limOptY): + for sf, lx, hx, ly, hy in zip(surface, limOptX[0], limOptX[1], limOptY[0], limOptY[1]): + rect = pg.QtWidgets.QGraphicsRectItem( # pylint: disable=E1101 + lx, + ly, + hx - lx, + hy - ly, + ) + rect.setBrush(pg.QtGui.QBrush(pg.QtGui.QColor(*self.color_impenetrable))) # pylint: disable=E1101 + rect.setPen(pg.mkPen(color=self.color_impenetrable, width=2)) + widget.addItem(rect) + text = pg.TextItem(sf, color='w', anchor=(0.5, 0.5)) # TODO: CHange color according to theme + widget.addItem(text) + text.setPos((hx+lx)/2, (hy+ly)/2) + text.setZValue(2) + + def plot_mono_surface(widget, xtal, xtalWidth, xtalOffsetX, xtalLength): + for sf, w, offx, len in zip(xtal, xtalWidth, xtalOffsetX, xtalLength): + rect = pg.QtWidgets.QGraphicsRectItem( # pylint: disable=E1101 + offx - w/2, + -len/2, + w, + len, + ) + rect.setBrush(pg.QtGui.QBrush(pg.QtGui.QColor(*self.color_impenetrable))) # pylint: disable=E1101 + rect.setPen(pg.mkPen(color=self.color_impenetrable, width=2)) + widget.addItem(rect) + text = pg.TextItem(sf, color='w', anchor=(0.5, 0.5)) # TODO: CHange color according to theme + widget.addItem(text) + text.setPos(offx, 0) + text.setZValue(2) + + for name, plot in self.plots.items(): + if name in 'cm': + plot_mirror_stripe(plot['widget'], bl.cm.surface, bl.cm.limOptX, bl.cm.limOptY) + elif name in 'mo1_1': + plot_mono_surface(plot['widget'], bl.mo1.xtal, bl.mo1.xtalWidth, bl.mo1.xtalOffsetX, bl.mo1.xtalLength1) + elif name in 'mo1_2': + plot_mono_surface(plot['widget'], bl.mo1.xtal, bl.mo1.xtalWidth, bl.mo1.xtalOffsetX, bl.mo1.xtalLength2) + elif name in 'fm': + plot_mirror_stripe(plot['widget'], bl.fm.surfaceFlat, bl.fm.limOptXFlat, bl.fm.limOptYFlat) + plot_mirror_stripe(plot['widget'], bl.fm.surfaceToroid, bl.fm.limOptXToroid, bl.fm.limOptYToroid) + else: + raise Exception(f'Plot {name} not found!') + for name, plot in self.plots.items(): + plot['widget'].disableAutoRange() + + def impenetrable_color(self): + app = QApplication.instance() + theme = app.theme.theme # type: ignore + if theme == "light": + return (30, 30, 30) + else: + return (220, 220, 220) + + def golden_angle_color( + self, + colormap: str, + num: int, + format="QColor", + theme_offset=0.2, + theme=None, + ) -> list: + """ + Extract num colors from the specified colormap following golden angle distribution and return them in the specified format. + + Args: + colormap (str): Name of the colormap. + num (int): Number of requested colors. + format (Literal["QColor","HEX","RGB"]): The format of the returned colors ('RGB', 'HEX', 'QColor'). + theme_offset (float): Has to be between 0-1. Offset to avoid colors too close to white or black with light or dark theme respectively for pyqtgraph plot background. + + Returns: + list: List of colors in the specified format. + + Raises: + ValueError: If theme_offset is not between 0 and 1. + """ + + cmap = pg.colormap.get(colormap) + phi = (1 + np.sqrt(5)) / 2 # Golden ratio + golden_angle_conjugate = 1 - (1 / phi) # Approximately 0.38196601125 + + min_pos, max_pos = self.set_theme_offset(theme, theme_offset) + + # Generate positions within the acceptable range + positions = np.mod(np.arange(num) * golden_angle_conjugate, 1) + positions = min_pos + positions * (max_pos - min_pos) + + # Sample colors from the colormap at the calculated positions + colors = cmap.map(positions, mode="float") # type: ignore + color_list = [] + + for color in colors: # type: ignore + if format.upper() == "HEX": + color_list.append(QColor.fromRgbF(*color).name()) + elif format.upper() == "RGB": + color_list.append(tuple((np.array(color) * 255).astype(int))) + elif format.upper() == "QCOLOR": + color_list.append(QColor.fromRgbF(*color)) + else: + raise ValueError("Unsupported format. Please choose 'RGB', 'HEX', or 'QColor'.") + return color_list + + def set_theme_offset(self, theme = None, offset=0.2) -> tuple: + """ + Set the theme offset to avoid colors too close to white or black with light or dark theme respectively for pyqtgraph plot background. + + Args: + theme(str): The theme to be applied. + offset(float): Offset to avoid colors too close to white or black with light or dark theme respectively for pyqtgraph plot background. + + Returns: + tuple: Tuple of min_pos and max_pos. + + Raises: + ValueError: If theme_offset is not between 0 and 1. + """ + + if offset < 0 or offset > 1: + raise ValueError("theme_offset must be between 0 and 1") + + if theme is None: + app = QApplication.instance() + if hasattr(app, "theme"): + theme = app.theme.theme # type: ignore + + if theme == "light": + min_pos = 0.0 + max_pos = 1 - offset + else: + min_pos = 0.0 + offset + max_pos = 1.0 + + return min_pos, max_pos + + def update_surfaces(self, scene, data): + self.surfaces[scene] = data + for name, device in self.surfaces[scene].items(): + 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) + # fill = pg.FillBetweenItem(curve, widget.plot(device['x'], np.zeros(len(device['x'])), pen=None), brush=pg.mkBrush('b')) + # widget.addItem(fill) + logger.info(self.surfaces) + +class SideviewPlot(QWidget): """Plot widget with two curves and legend.""" def __init__(self, parent=None): @@ -458,28 +826,35 @@ class PlotWidget(QWidget): self.plot_widget.getAxis('bottom').enableAutoSIPrefix(False) self.plot_widget.addLegend() + app = QApplication.instance() + theme = app.theme.theme # type: ignore + if theme == "light": + self.color_impenetrable = (30, 30, 30) + self.colors = [(26, 111, 173), (212, 83, 10)] + else: # dark theme + self.color_impenetrable = (220, 220, 220) + self.colors = [(79, 163, 224), (240, 128, 60)] + self.curves = [] - colors = self.golden_angle_color( - colormap='plasma', - num=2, - format="HEX", - ) - self.color_impenetrable = self.impenetrable_color() self.data = { - 'assistant': [[0, 1000, 2000], [0, 20, 30]], - 'reality': [[0, 1000, 2000], [0, 18, 36]], + 'assistant': {'x': [0, 1000, 2000], 'y': [0, 20, 30]}, + 'reality': {'x': [0, 1000, 2000], 'y': [0, 15, 50]}, } self.pipes = [] self.walls = [] - for idx, element in enumerate(self.data): + for idx, name in enumerate(self.data.keys()): + if name in "assistant": + pen = pg.mkPen(color=self.colors[idx], width=2, style=Qt.DashLine) + else: + pen = pg.mkPen(color=self.colors[idx], width=2) self.curves.append( self.plot_widget.plot( [], [], - pen=pg.mkPen(color=colors[idx], width=2), - name=element, + pen=pen, + name=name, ) ) @@ -624,8 +999,8 @@ class PlotWidget(QWidget): def update_curves(self): for idx, element in enumerate(self.data): self.curves[idx].setData( - x=np.array(self.data[element][0]), - y=np.array(self.data[element][1]), + x=np.array(self.data[element]['x']), + y=np.array(self.data[element]['y']), ) diff --git a/debye_bec/bec_widgets/widgets/qt_widgets.py b/debye_bec/bec_widgets/widgets/qt_widgets.py index 2ccb69b..8a9d1df 100644 --- a/debye_bec/bec_widgets/widgets/qt_widgets.py +++ b/debye_bec/bec_widgets/widgets/qt_widgets.py @@ -66,6 +66,9 @@ class NumberIndicator(QWidget): 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}'