bec/bec_widgets
1
0
Fork 0
mirror of https://github.com/bec-project/bec_widgets.git synced 2026-09-06 16:40:56 +02:00
Code Activity

refactor(scatter_waveform): consume the data api for live and history

Browse Source
This commit is contained in:
wyzula_j
2026-08-14 12:00:27 +02:00
parent 6b1b508fd6
commit 83e5b5f2b8
2 changed files with 183 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bec_widgets/widgets/plots/scatter_waveform/scatter_waveform.py
+72 -90
View File
@@ -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()
tests/unit_tests/test_scatter_waveform.py
+111 -92
View File
@@ -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)