From 83e5b5f2b86903302de98e1029a3d986cc443f7c Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Fri, 7 Aug 2026 12:49:01 +0200 Subject: [PATCH] refactor(scatter_waveform): consume the data api for live and history --- .../scatter_waveform/scatter_waveform.py | 162 +++++++------- tests/unit_tests/test_scatter_waveform.py | 203 ++++++++++-------- 2 files changed, 183 insertions(+), 182 deletions(-) diff --git a/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py b/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py index 54640730..66350b21 100644 --- a/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py +++ b/bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py @@ -1,16 +1,15 @@ from __future__ import annotations -import numpy as np -import pyqtgraph as pg from bec_lib import bec_logger from bec_lib.endpoints import MessageEndpoints from pydantic import Field, ValidationError, field_validator -from qtpy.QtCore import QTimer, Signal +from qtpy.QtCore import Signal from qtpy.QtWidgets import QHBoxLayout, QMainWindow, QWidget from bec_widgets.utils.bec_connector import ConnectionConfig from bec_widgets.utils.colors import Colors from bec_widgets.utils.error_popups import SafeProperty, SafeSlot +from bec_widgets.utils.qt_data_subscription import QtDataSubscription from bec_widgets.utils.settings_dialog import SettingsDialog from bec_widgets.utils.toolbars.toolbar import MaterialIconAction from bec_widgets.widgets.plots.plot_base import PlotBase, UIMode @@ -65,7 +64,6 @@ class ScatterWaveform(PlotBase): "signal_z.setter", ] - sync_signal_update = Signal() new_scan = Signal() new_scan_id = Signal(str) scatter_waveform_property_changed = Signal() @@ -95,14 +93,11 @@ class ScatterWaveform(PlotBase): self.scan_id = None self.scan_item = None - # Scan status update loop + # Data delivery through the DataAPI; scan_status is only used for + # per-scan bookkeeping (reset, auto range, new_scan signals). + self._data_bridge: QtDataSubscription | None = None + self._data_bridge_scope: str | None = "live" self.bec_dispatcher.connect_slot(self.on_scan_status, MessageEndpoints.scan_status()) - self.bec_dispatcher.connect_slot(self.on_scan_progress, MessageEndpoints.scan_progress()) - - # Curve update loop - self.proxy_update_sync = pg.SignalProxy( - self.sync_signal_update, rateLimit=25, slot=self.update_sync_curves - ) self._init_scatter_curve_settings() @@ -291,7 +286,7 @@ class ScatterWaveform(PlotBase): self.update_labels() self.plot_item.addItem(self._main_curve) - self.sync_signal_update.emit() + self._setup_data_api_subscription() ################################################################################ # BEC Update Methods @@ -318,90 +313,76 @@ class ScatterWaveform(PlotBase): self.old_scan_id = self.scan_id self.scan_id = current_scan_id self.scan_item = self.queue.scan_storage.find_scan_by_ID(self.scan_id) + if self._data_bridge is None or self._data_bridge_scope != "live": + # The widget started idle and bound to the latest finished + # scan (or has no bridge yet); a scan is running now, so + # switch to a live-follow subscription. + self._setup_data_api_subscription() - # First trigger to update the scan curves - self.sync_signal_update.emit() + def _config_sources(self): + """Return the deduplicated (device, signal) sources of the curve config.""" + curve = self._main_curve + if curve is None or curve.config is None: + return None + try: + sources = [ + (curve.config.device_x.device, curve.config.device_x.signal), + (curve.config.device_y.device, curve.config.device_y.signal), + (curve.config.device_z.device, curve.config.device_z.signal), + ] + except AttributeError: + return None + if not all(dev for dev, _ in sources): + return None + return list(dict.fromkeys(sources)) - @SafeSlot(dict, dict) - def on_scan_progress(self, msg: dict, meta: dict): + def _setup_data_api_subscription(self, scan: str = "live"): + """(Re)create the DataAPI subscription for the configured sources.""" + self._cleanup_data_api_subscription() + self._data_bridge_scope = scan + sources = self._config_sources() + if sources is None: + return + try: + self._data_bridge = QtDataSubscription( + self.client, sources=sources, scan=scan, parent=self, min_emit_interval=0.1 + ) + self._data_bridge.updated.connect(self._on_data_update) + except Exception as exc: + logger.warning(f"Failed to configure scatter waveform data subscription: {exc}") + self._cleanup_data_api_subscription() + + def _cleanup_data_api_subscription(self): + if self._data_bridge is None: + return + try: + self._data_bridge.close() + finally: + self._data_bridge = None + + @SafeSlot(object) + def _on_data_update(self, update) -> None: """ - Slot for handling scan progress messages. Used for triggering the update of the sync curves. + Render one columnar DataAPI update (live or history). Args: - msg(dict): The message content. - meta(dict): The message metadata. + update (SubscriptionUpdate): Aligned full-state snapshot. """ - self.sync_signal_update.emit() - status = msg.get("done") - if status: - QTimer.singleShot(100, self.update_sync_curves) - QTimer.singleShot(300, self.update_sync_curves) - - @SafeSlot() - def update_sync_curves(self, _=None): - """ - Update the scan curves with the data from the scan segment. - """ - if self.scan_item is None: - logger.info("No scan executed so far; skipping device curves categorisation.") - return "none" - data, access_key = self._fetch_scan_data_and_access() - - if data == "none": - logger.info("No scan executed so far; skipping device curves categorisation.") - return "none" - + curve = self._main_curve + sources = self._config_sources() + if curve is None or sources is None: + return + columns = update.aligned() try: - device_x = self._main_curve.config.device_x.device - signal_x = self._main_curve.config.device_x.signal - device_y = self._main_curve.config.device_y.device - signal_y = self._main_curve.config.device_y.signal - device_z = self._main_curve.config.device_z.device - signal_z = self._main_curve.config.device_z.signal - except AttributeError: + config = curve.config + x_data = list(columns[(config.device_x.device, config.device_x.signal)]) + y_data = list(columns[(config.device_y.device, config.device_y.signal)]) + z_data = list(columns[(config.device_z.device, config.device_z.signal)]) + except (KeyError, AttributeError): return - - if access_key == "val": - x_data = data.get(device_x, {}).get(signal_x, {}).get(access_key, None) - y_data = data.get(device_y, {}).get(signal_y, {}).get(access_key, None) - z_data = data.get(device_z, {}).get(signal_z, {}).get(access_key, None) - else: - x_data = data.get(device_x, {}).get(signal_x, {}).read().get("value", None) - y_data = data.get(device_y, {}).get(signal_y, {}).read().get("value", None) - z_data = data.get(device_z, {}).get(signal_z, {}).read().get("value", None) - - if x_data is None or y_data is None or z_data is None: + if not x_data: return - - x_data, y_data, z_data = (np.atleast_1d(arr) for arr in (x_data, y_data, z_data)) - min_len = min(len(x_data), len(y_data), len(z_data)) - if min_len == 0: - return - self._main_curve.set_data(x=x_data[:min_len], y=y_data[:min_len], z=z_data[:min_len]) - - def _fetch_scan_data_and_access(self): - """ - Decide whether the widget is in live or historical mode - and return the appropriate data dict and access key. - - Returns: - data_dict (dict): The data structure for the current scan. - access_key (str): Either 'val' (live) or 'value' (history). - """ - if self.scan_item is None: - # Optionally fetch the latest from history if nothing is set - self.update_with_scan_history(-1) - if self.scan_item is None: - logger.info("No scan executed so far; skipping device curves categorisation.") - return "none", "none" - - if hasattr(self.scan_item, "live_data"): - # Live scan - return self.scan_item.live_data, "val" - else: - # Historical - scan_devices = self.scan_item.devices - return scan_devices, "value" + curve.set_data(x=x_data, y=y_data, z=z_data) ################################################################################ # Widget Specific Properties @@ -703,7 +684,7 @@ class ScatterWaveform(PlotBase): if scan_index is None: self.scan_id = scan_id self.scan_item = self.client.history.get_by_scan_id(scan_id) - self.sync_signal_update.emit() + self._setup_data_api_subscription(scan=scan_id) return if scan_index == -1: @@ -714,7 +695,7 @@ class ScatterWaveform(PlotBase): return self.scan_item = scan_item self.scan_id = scan_item.scan_id - self.sync_signal_update.emit() + self._setup_data_api_subscription(scan="live") return if len(self.client.history) == 0: @@ -725,7 +706,7 @@ class ScatterWaveform(PlotBase): metadata = self.scan_item.metadata self.scan_id = metadata["bec"]["scan_id"] - self.sync_signal_update.emit() + self._setup_data_api_subscription(scan=self.scan_id) ################################################################################ # Cleanup @@ -743,6 +724,7 @@ class ScatterWaveform(PlotBase): """ Cleanup the widget and disconnect all signals. """ + self._cleanup_data_api_subscription() if self.scatter_dialog is not None: self.scatter_dialog.close() self.scatter_dialog.deleteLater() diff --git a/tests/unit_tests/test_scatter_waveform.py b/tests/unit_tests/test_scatter_waveform.py index 04852368..c0d7689f 100644 --- a/tests/unit_tests/test_scatter_waveform.py +++ b/tests/unit_tests/test_scatter_waveform.py @@ -49,20 +49,73 @@ def test_scatter_waveform_color_map(qtbot, mocked_client): assert swf.color_map == "plasma" +def _fake_bridge_factory(monkeypatch): + created = [] + + class _FakeBridge: + def __init__(self, client, sources, scan="live", parent=None, min_emit_interval=0.1): + self.client = client + self.sources = list(sources) + self.scan = scan + self.healthy = True + self.closed = False + self.updated = MagicMock() + + def close(self): + self.closed = True + + def factory(client, sources, scan="live", parent=None, min_emit_interval=0.1): + bridge = _FakeBridge(client, sources, scan=scan) + created.append(bridge) + return bridge + + monkeypatch.setattr( + "bec_widgets.widgets.plots.scatter_waveform.scatter_waveform.QtDataSubscription", factory + ) + return created + + +def _make_update(scan_id="dummy"): + from bec_lib.data_api.models import SourceData, SubscriptionUpdate + + columns = {"samx": [10, 20, 30], "samy": [5, 10, 15], "bpm4i": [1, 2, 3]} + sources = {} + for dev, values in columns.items(): + sources[(dev, dev)] = SourceData( + device=dev, + entry=dev, + kind="monitored", + ordinals=(0, 1, 2), + values=tuple(values), + timestamps=(1.0, 2.0, 3.0), + complete=True, + ) + return SubscriptionUpdate( + scan_id=scan_id, + reason="live", + sources=sources, + aligned_ordinals=(0, 1, 2), + complete=True, + metadata={"group": "scan"}, + ) + + def test_scatter_waveform_update_with_scan_history(qtbot, mocked_client, monkeypatch): + bridges = _fake_bridge_factory(monkeypatch) swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) dummy_scan = create_dummy_scan_item() mocked_client.history = MagicMock() - # .get_by_scan_id() typically returns historical data, but we abuse it here - # to return mock live data mocked_client.history.get_by_scan_id.return_value = dummy_scan mocked_client.history.__getitem__.return_value = dummy_scan swf.plot("samx", "samy", "bpm4i", label="test_curve") swf.update_with_scan_history(scan_id="dummy") qtbot.waitUntil(lambda: swf.scan_item == dummy_scan, timeout=500) - qtbot.wait(200) + + # History flows through a scan-id-bound DataAPI subscription. + assert bridges[-1].scan == "dummy" + swf._on_data_update(_make_update("dummy")) x_data, y_data = swf.main_curve.getData() np.testing.assert_array_equal(x_data, [10, 20, 30]) @@ -70,50 +123,36 @@ def test_scatter_waveform_update_with_scan_history(qtbot, mocked_client, monkeyp def test_scatter_waveform_live_update(qtbot, mocked_client, monkeypatch): + bridges = _fake_bridge_factory(monkeypatch) swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) dummy_scan = create_dummy_scan_item() monkeypatch.setattr(swf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: dummy_scan) swf.plot("samx", "samy", "bpm4i", label="live_curve") + assert bridges[-1].scan == "live" + assert bridges[-1].sources == [("samx", "samx"), ("samy", "samy"), ("bpm4i", "bpm4i")] - # Simulate scan status indicating new scan start - msg = {"scan_id": "dummy"} - meta = {} - swf.on_scan_status(msg, meta) - + # Scan status only performs per-scan bookkeeping now. + swf.on_scan_status({"scan_id": "dummy"}, {}) assert swf.scan_id == "dummy" assert swf.scan_item == dummy_scan - qtbot.wait(500) - + swf._on_data_update(_make_update("dummy")) x_data, y_data = swf.main_curve.getData() np.testing.assert_array_equal(x_data, [10, 20, 30]) np.testing.assert_array_equal(y_data, [5, 10, 15]) -def test_scatter_waveform_scan_progress(qtbot, mocked_client, monkeypatch): +def test_scatter_waveform_plot_replaces_subscription(qtbot, mocked_client, monkeypatch): + bridges = _fake_bridge_factory(monkeypatch) swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) - dummy_scan = create_dummy_scan_item() - monkeypatch.setattr(swf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: dummy_scan) - swf.plot("samx", "samy", "bpm4i") - - # Simulate scan status indicating scan progress - swf.scan_id = "dummy" - swf.scan_item = dummy_scan - - msg = {"progress": 50} - meta = {} - swf.on_scan_progress(msg, meta) - qtbot.wait(500) - - # swf.update_sync_curves() - - x_data, y_data = swf.main_curve.getData() - np.testing.assert_array_equal(x_data, [10, 20, 30]) - np.testing.assert_array_equal(y_data, [5, 10, 15]) + first = bridges[-1] + swf.plot("samx", "samy", "samx") + assert first.closed is True + assert bridges[-1].sources == [("samx", "samx"), ("samy", "samy")] # deduplicated # def test_scatter_waveform_settings_popup(qtbot, mocked_client): @@ -574,69 +613,6 @@ def test_z_gradient_uses_shared_brush_pool(qtbot, mocked_client): assert brushes_first[-1].color().getRgb()[:3] == hi.getRgb()[:3] -def test_scatter_waveform_live_update_unequal_lengths(qtbot, mocked_client, monkeypatch): - """Device buffers fill independently during a live scan, so one of them can - be a point ahead when the update slot fires; the data must be trimmed to the - common length instead of handing pyqtgraph mismatched arrays.""" - swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) - - dummy_scan = create_dummy_scan_item() - # samx is one point ahead of samy and bpm4i - dummy_scan.live_data["samx"]["samx"].val = [10, 20, 30, 40] - monkeypatch.setattr(swf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: dummy_scan) - - swf.plot("samx", "samy", "bpm4i", label="live_curve") - swf.scan_id = "dummy" - swf.scan_item = dummy_scan - - swf.update_sync_curves() - - x_data, y_data = swf.main_curve.getData() - np.testing.assert_array_equal(x_data, [10, 20, 30]) - np.testing.assert_array_equal(y_data, [5, 10, 15]) - - -def test_scatter_waveform_live_update_incomplete_data(qtbot, mocked_client, monkeypatch): - """Missing or empty device buffers skip the update instead of erroring.""" - swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) - - dummy_scan = create_dummy_scan_item() - dummy_scan.live_data["samy"] = {} # y signal not delivered yet - monkeypatch.setattr(swf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: dummy_scan) - - swf.plot("samx", "samy", "bpm4i", label="live_curve") - swf.scan_id = "dummy" - swf.scan_item = dummy_scan - - swf.update_sync_curves() - assert swf.main_curve.getData() == (None, None) - - dummy_scan.live_data["samy"] = {"samy": DummyData(val=[], timestamps=[])} - swf.update_sync_curves() - assert swf.main_curve.getData() == (None, None) - - -def test_scatter_waveform_live_update_scalar_readback(qtbot, mocked_client, monkeypatch): - """The first live update can deliver bare scalar readbacks instead of buffers; - they must be plotted as a single point, not raise TypeError on len().""" - swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) - - dummy_scan = create_dummy_scan_item() - for device in ("samx", "samy", "bpm4i"): - dummy_scan.live_data[device][device] = DummyData(val=5.0, timestamps=100) - monkeypatch.setattr(swf.queue.scan_storage, "find_scan_by_ID", lambda scan_id: dummy_scan) - - swf.plot("samx", "samy", "bpm4i", label="live_curve") - swf.scan_id = "dummy" - swf.scan_item = dummy_scan - - swf.update_sync_curves() - - x_data, y_data = swf.main_curve.getData() - np.testing.assert_array_equal(x_data, [5.0]) - np.testing.assert_array_equal(y_data, [5.0]) - - def test_z_gradient_flat_and_empty_z(qtbot, mocked_client): """Degenerate z inputs keep returning None (no gradient).""" swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) @@ -649,3 +625,46 @@ def test_z_gradient_flat_and_empty_z(qtbot, mocked_client): assert curve._make_z_gradient(float("nan"), "plasma") is None mixed = curve._make_z_gradient([float("nan"), 1.0, 2.0], "plasma") assert mixed is not None and len(mixed) == 3 + + +def test_scatter_switches_history_bound_bridge_to_live_on_new_scan( + qtbot, mocked_client, monkeypatch +): + """Same live regression as the heatmap: a bridge bound to the latest + finished scan at idle startup must switch to live-follow when a scan + starts.""" + bridges = _fake_bridge_factory(monkeypatch) + swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) + swf.plot("samx", "samy", "bpm4i") + swf._setup_data_api_subscription(scan="old-finished-scan") + assert bridges[-1].scan == "old-finished-scan" + + swf.on_scan_status({"scan_id": "new-live-scan"}, {}) + assert bridges[-1].scan == "live" + + bridge_count = len(bridges) + swf.on_scan_status({"scan_id": "another-scan"}, {}) + assert len(bridges) == bridge_count + + +def test_scatter_waveform_update_skips_incomplete_columns(qtbot, mocked_client, monkeypatch): + """A source that has not delivered yet, or delivered nothing, must skip the + render instead of erroring. Unequal buffer lengths cannot reach the widget + any more: ``aligned()`` only returns ordinals present in every source.""" + from dataclasses import replace + + _fake_bridge_factory(monkeypatch) + swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) + swf.plot("samx", "samy", "bpm4i", label="live_curve") + + partial = _make_update("dummy") + partial.sources.pop(("samy", "samy")) + swf._on_data_update(partial) + assert swf.main_curve.getData() == (None, None) + + empty = _make_update("dummy") + for key, source in list(empty.sources.items()): + empty.sources[key] = replace(source, ordinals=(), values=(), timestamps=()) + object.__setattr__(empty, "aligned_ordinals", ()) + swf._on_data_update(empty) + assert swf.main_curve.getData() == (None, None)