diff --git a/bec_widgets/bec_dispatcher.py b/bec_widgets/bec_dispatcher.py new file mode 100644 index 00000000..d60b3c4b --- /dev/null +++ b/bec_widgets/bec_dispatcher.py @@ -0,0 +1,73 @@ +from collections import defaultdict +from threading import RLock + +from bec_lib.core import BECMessage, MessageEndpoints, RedisConnector +from PyQt5.QtCore import QObject, pyqtSignal + +bec_connector = RedisConnector("localhost:6379") + + +class _BECDispatcher(QObject): + scan_segment = pyqtSignal("PyQt_PyObject") + new_dap_data = pyqtSignal(dict) + new_scan = pyqtSignal("PyQt_PyObject") + + def __init__(self): + super().__init__() + # TODO: dap might not be a good fit to predefined slots, fix this inconsistency + self._slot_signal_map = { + "on_scan_segment": self.scan_segment, + "on_new_scan": self.new_scan, + } + self._daps = defaultdict(set) + + self._scan_id = None + scan_lock = RLock() + self._dap_threads = [] + + def _scan_cb(msg): + msg = BECMessage.ScanMessage.loads(msg.value)[0] + with scan_lock: + # TODO: use ScanStatusMessage instead? + scan_id = msg.content["scanID"] + if self._scan_id != scan_id: + self._scan_id = scan_id + self.new_scan.emit(msg) + self.scan_segment.emit(msg) + + scan_readback = MessageEndpoints.scan_segment() + self._scan_thread = bec_connector.consumer( + topics=scan_readback, + cb=_scan_cb, + ) + self._scan_thread.start() + + def connect(self, widget): + for slot_name, signal in self._slot_signal_map.items(): + slot = getattr(widget, slot_name, None) + if callable(slot): + signal.connect(slot) + + def connect_dap(self, slot, dap_name): + if dap_name not in self._daps: + + def _dap_cb(msg): + msg = BECMessage.ProcessedDataMessage.loads(msg.value) + self.new_dap_data.emit(msg.content["data"]) + + dap_ep = MessageEndpoints.processed_data(dap_name) + dap_thread = bec_connector.consumer(topics=dap_ep, cb=_dap_cb) + dap_thread.start() + self._dap_threads.append(dap_thread) + + self.new_dap_data.connect(slot) + self._daps[dap_name].add(slot) + + else: + # connect slot if it's not yet connected + if slot not in self._daps[dap_name]: + self._daps[dap_name].add(slot) + self.new_dap_data.connect(slot) + + +bec_dispatcher = _BECDispatcher() diff --git a/bec_widgets/cli.py b/bec_widgets/cli.py deleted file mode 100644 index 0ce775eb..00000000 --- a/bec_widgets/cli.py +++ /dev/null @@ -1,107 +0,0 @@ -import argparse -import os -from threading import RLock - -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 - - -class BEC_UI(QMainWindow): - new_scan_data = pyqtSignal("PyQt_PyObject") - new_dap_data = pyqtSignal(dict) # signal per proc instance? - new_scan = pyqtSignal("PyQt_PyObject") - - def __init__(self, uipath): - super().__init__() - self._scan_channels = set() - self._dap_channels = set() - - self._scan_thread = None - self._dap_threads = [] - - ui = uic.loadUi(uipath, self) - - _, fname = os.path.split(uipath) - self.setWindowTitle(fname) - - for sp in ui.findChildren(BECScanPlot): - for chan in (sp.x_channel, *sp.y_channel_list): - if chan.startswith("dap."): - chan = chan.partition("dap.")[-1] - self._dap_channels.add(chan) - else: - self._scan_channels.add(chan) - - sp.initialize() # TODO: move this elsewhere? - - self.new_scan_data.connect(sp.redraw_scan) # TODO: merge - 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() - - def _scan_cb(msg): - msg = BECMessage.ScanMessage.loads(msg.value)[0] - with scan_lock: - scan_id = msg.content["scanID"] - if self._scan_id != scan_id: - self._scan_id = scan_id - self.new_scan.emit(msg) - self.new_scan_data.emit(msg) - - bec_connector = RedisConnector("localhost:6379") - - if self._scan_channels: - scan_readback = MessageEndpoints.scan_segment() - self._scan_thread = bec_connector.consumer( - topics=scan_readback, - cb=_scan_cb, - ) - self._scan_thread.start() - - # DAP setup - def _proc_cb(msg): - msg = BECMessage.ProcessedDataMessage.loads(msg.value) - self.new_dap_data.emit(msg.content["data"]) - - if self._dap_channels: - for chan in self._dap_channels: - proc_ep = MessageEndpoints.processed_data(chan) - dap_thread = bec_connector.consumer(topics=proc_ep, cb=_proc_cb) - dap_thread.start() - self._dap_threads.append(dap_thread) - - self.show() - - -def main(): - parser = argparse.ArgumentParser( - prog="bec-pyqt", formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - - parser.add_argument("uipath", type=str, help="Path to a BEC ui file") - - args, rem = parser.parse_known_args() - - app = QApplication(rem) - BEC_UI(args.uipath) - app.exec_() - - -if __name__ == "__main__": - main() diff --git a/bec_widgets/scan2d_plot.py b/bec_widgets/scan2d_plot.py index a9d20181..2d84ed2f 100644 --- a/bec_widgets/scan2d_plot.py +++ b/bec_widgets/scan2d_plot.py @@ -3,6 +3,8 @@ import pyqtgraph as pg from bec_lib.core.logger import bec_logger from PyQt5.QtCore import pyqtProperty, pyqtSlot +from bec_widgets.bec_dispatcher import bec_dispatcher + logger = bec_logger.logger @@ -12,6 +14,7 @@ pg.setConfigOptions(background="w", foreground="k", antialias=True) class BECScanPlot2D(pg.GraphicsView): def __init__(self, parent=None, background="default"): super().__init__(parent, background) + bec_dispatcher.connect(self) self._x_channel = "" self._y_channel = "" @@ -30,12 +33,8 @@ class BECScanPlot2D(pg.GraphicsView): 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): + def on_new_scan(self, msg): # TODO: Do we reset in case of a scan type change? self.imageItem.clear() @@ -79,7 +78,7 @@ class BECScanPlot2D(pg.GraphicsView): self.plot_item.setLabel("left", motors[self._y_ind]) @pyqtSlot("PyQt_PyObject") - def redraw_scan(self, msg): + def on_scan_segment(self, msg): if not self.z_channel or msg.metadata["scan_name"] != "grid_scan": return @@ -106,6 +105,7 @@ class BECScanPlot2D(pg.GraphicsView): @x_channel.setter def x_channel(self, new_val): self._x_channel = new_val + self.plot_item.setLabel("bottom", new_val) @pyqtProperty(str) def y_channel(self): @@ -114,6 +114,7 @@ class BECScanPlot2D(pg.GraphicsView): @y_channel.setter def y_channel(self, new_val): self._y_channel = new_val + self.plot_item.setLabel("left", new_val) @pyqtProperty(str) def z_channel(self): diff --git a/bec_widgets/scan_plot.py b/bec_widgets/scan_plot.py index 3e68b200..4b9b410f 100644 --- a/bec_widgets/scan_plot.py +++ b/bec_widgets/scan_plot.py @@ -4,6 +4,8 @@ import pyqtgraph as pg from bec_lib.core.logger import bec_logger from PyQt5.QtCore import pyqtProperty, pyqtSlot +from bec_widgets.bec_dispatcher import bec_dispatcher + logger = bec_logger.logger @@ -14,6 +16,7 @@ COLORS = ["#fd7f6f", "#7eb0d5", "#b2e061", "#bd7ebe", "#ffb55a"] class BECScanPlot(pg.GraphicsView): def __init__(self, parent=None, background="default"): super().__init__(parent, background) + bec_dispatcher.connect(self) self.view = pg.PlotItem() self.setCentralItem(self.view) @@ -24,32 +27,13 @@ class BECScanPlot(pg.GraphicsView): self.scan_curves = {} self.dap_curves = {} - def initialize(self): - self.view.addLegend() - colors = itertools.cycle(COLORS) - - for y_chan in self.y_channel_list: - if y_chan.startswith("dap."): - y_chan = y_chan.partition("dap.")[-1] - curves = self.dap_curves - else: - curves = self.scan_curves - - curves[y_chan] = self.view.plot( - x=[], y=[], pen=pg.mkPen(color=next(colors), width=2), name=y_chan - ) - - self.view.setLabel("bottom", self._x_channel) - if len(self.scan_curves) == 1: - self.view.setLabel("left", next(iter(self.scan_curves))) - @pyqtSlot("PyQt_PyObject") - def clearData(self, _msg): + def on_new_scan(self, _msg): for plot_curve in {**self.scan_curves, **self.dap_curves}.values(): plot_curve.setData(x=[], y=[]) @pyqtSlot("PyQt_PyObject") - def redraw_scan(self, msg): + def on_scan_segment(self, msg): if not self.x_channel: return @@ -100,6 +84,28 @@ class BECScanPlot(pg.GraphicsView): def y_channel_list(self, new_list): self._y_channel_list = new_list + # Prepare plot for a potentially different list of y channels + self.view.clear() + + self.view.addLegend() + colors = itertools.cycle(COLORS) + + for y_chan in new_list: + # TODO: ideally, we dont want to care about dap/not dap here + if y_chan.startswith("dap."): + y_chan = y_chan.partition("dap.")[-1] + curves = self.dap_curves + bec_dispatcher.connect_dap(self.redraw_dap, y_chan) + else: + curves = self.scan_curves + + curves[y_chan] = self.view.plot( + x=[], y=[], pen=pg.mkPen(color=next(colors), width=2), name=y_chan + ) + + if len(new_list) == 1: + self.view.setLabel("left", new_list[0]) + @pyqtProperty(str) def x_channel(self): return self._x_channel @@ -107,6 +113,7 @@ class BECScanPlot(pg.GraphicsView): @x_channel.setter def x_channel(self, new_val): self._x_channel = new_val + self.view.setLabel("bottom", new_val) if __name__ == "__main__":