diff --git a/bec_widgets/cli.py b/bec_widgets/cli.py index 6db1c7fe..bcca01ef 100644 --- a/bec_widgets/cli.py +++ b/bec_widgets/cli.py @@ -6,6 +6,7 @@ from bec_lib.core import BECMessage, MessageEndpoints, RedisConnector from PyQt5 import uic from PyQt5.QtCore import pyqtSignal from PyQt5.QtWidgets import QApplication, QMainWindow +from scan2d_plot import BECScanPlot2D from .scan_plot import BECScanPlot @@ -42,6 +43,15 @@ class BEC_UI(QMainWindow): self.new_dap_data.connect(sp.redraw_dap) self.new_scan.connect(sp.clearData) + for sp in ui.findChildren(BECScanPlot2D): + for chan in (sp.x_channel, sp.y_channel, sp.z_channel): + self._scan_channels.add(chan) + + sp.initialize() + + self.new_scan_data.connect(sp.redraw_scan) + self.new_scan.connect(sp.clearData) + # Scan setup self._scan_id = None scan_lock = RLock() diff --git a/bec_widgets/scan2d_plot.py b/bec_widgets/scan2d_plot.py new file mode 100644 index 00000000..a9d20181 --- /dev/null +++ b/bec_widgets/scan2d_plot.py @@ -0,0 +1,124 @@ +import numpy as np +import pyqtgraph as pg +from bec_lib.core.logger import bec_logger +from PyQt5.QtCore import pyqtProperty, pyqtSlot + +logger = bec_logger.logger + + +pg.setConfigOptions(background="w", foreground="k", antialias=True) + + +class BECScanPlot2D(pg.GraphicsView): + def __init__(self, parent=None, background="default"): + super().__init__(parent, background) + + self._x_channel = "" + self._y_channel = "" + self._z_channel = "" + + self._xpos = [] + self._ypos = [] + + self._x_ind = None + self._y_ind = None + + self.plot_item = pg.PlotItem() + self.setCentralItem(self.plot_item) + self.plot_item.setAspectLocked(True) + + self.imageItem = pg.ImageItem() + self.plot_item.addItem(self.imageItem) + + def initialize(self): + self.plot_item.setLabel("bottom", self.x_channel) + self.plot_item.setLabel("left", self.y_channel) + + @pyqtSlot("PyQt_PyObject") + def clearData(self, msg): + # TODO: Do we reset in case of a scan type change? + self.imageItem.clear() + + # TODO: better to check the number of coordinates in metadata["positions"]? + if msg.metadata["scan_name"] != "grid_scan": + return + + positions = [sorted(set(pos)) for pos in zip(*msg.metadata["positions"])] + + motors = msg.metadata["scan_motors"] + if self.x_channel and self.y_channel: + self._x_ind = motors.index(self.x_channel) if self.x_channel in motors else None + self._y_ind = motors.index(self.y_channel) if self.y_channel in motors else None + elif not self.x_channel and not self.y_channel: + # Plot the first and second motors along x and y axes respectively + self._x_ind = 0 + self._y_ind = 1 + else: + logger.warning( + f"X and Y channels should be either both empty or both set in {self.objectName()}" + ) + + if self._x_ind is None or self._y_ind is None: + return + + xpos = positions[self._x_ind] + ypos = positions[self._y_ind] + + self._xpos = xpos + self._ypos = ypos + + self.imageItem.setImage(np.zeros(shape=(len(xpos), len(ypos)))) + + w = max(xpos) - min(xpos) + h = max(ypos) - min(ypos) + w_pix = w / (len(xpos) - 1) + h_pix = h / (len(ypos) - 1) + self.imageItem.setRect(min(xpos) - w_pix / 2, min(ypos) - h_pix / 2, w + w_pix, h + h_pix) + + self.plot_item.setLabel("bottom", motors[self._x_ind]) + self.plot_item.setLabel("left", motors[self._y_ind]) + + @pyqtSlot("PyQt_PyObject") + def redraw_scan(self, msg): + if not self.z_channel or msg.metadata["scan_name"] != "grid_scan": + return + + if self._x_ind is None or self._y_ind is None: + return + + point_id = msg.content["point_id"] + point_coord = msg.metadata["positions"][point_id] + + x_coord_ind = self._xpos.index(point_coord[self._x_ind]) + y_coord_ind = self._ypos.index(point_coord[self._y_ind]) + + data = msg.content["data"] + z_new = data[self.z_channel][self.z_channel]["value"] + + image = self.imageItem.image + image[x_coord_ind, y_coord_ind] = z_new + self.imageItem.setImage() + + @pyqtProperty(str) + def x_channel(self): + return self._x_channel + + @x_channel.setter + def x_channel(self, new_val): + self._x_channel = new_val + + @pyqtProperty(str) + def y_channel(self): + return self._y_channel + + @y_channel.setter + def y_channel(self, new_val): + self._y_channel = new_val + + @pyqtProperty(str) + def z_channel(self): + return self._z_channel + + @z_channel.setter + def z_channel(self, new_val): + self._z_channel = new_val diff --git a/bec_widgets/scan2d_plot_plugin.py b/bec_widgets/scan2d_plot_plugin.py new file mode 100644 index 00000000..4a6dcf19 --- /dev/null +++ b/bec_widgets/scan2d_plot_plugin.py @@ -0,0 +1,56 @@ +from PyQt5.QtDesigner import QPyDesignerCustomWidgetPlugin +from PyQt5.QtGui import QIcon + +from scan2d_plot import BECScanPlot2D + + +class BECScanPlot2DPlugin(QPyDesignerCustomWidgetPlugin): + def __init__(self, parent=None): + super().__init__(parent) + + self._initialized = False + + def initialize(self, formEditor): + if self._initialized: + return + + self._initialized = True + + def isInitialized(self): + return self._initialized + + def createWidget(self, parent): + return BECScanPlot2D(parent) + + def name(self): + return "BECScanPlot2D" + + def group(self): + return "BEC widgets" + + def icon(self): + return QIcon() + + def toolTip(self): + return "BEC plot for 2D scans" + + def whatsThis(self): + return "BEC plot for 2D scans" + + def isContainer(self): + return False + + def domXml(self): + return ( + '\n' + ' \n' + " BEC plot for 2D scans\n" + " \n" + ' \n' + " BEC plot for 2D scans in Python using PyQt.\n" + " \n" + "\n" + ) + + def includeFile(self): + return "scan2d_plot"