feat: add BECScanPlot2D

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()

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

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 (
+            '<widget class="BECScanPlot2D" name="BECScanPlot2D">\n'
+            ' <property name="toolTip" >\n'
+            "  <string>BEC plot for 2D scans</string>\n"
+            " </property>\n"
+            ' <property name="whatsThis" >\n'
+            "  <string>BEC plot for 2D scans in Python using PyQt.</string>\n"
+            " </property>\n"
+            "</widget>\n"
+        )
+
+    def includeFile(self):
+        return "scan2d_plot"