From dc57a90357c709c4ff57d31350b7c06d0fb9f7f0 Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Wed, 12 Aug 2026 01:35:58 +0200 Subject: [PATCH] perf(scatter): pooled quantized z-brushes enable symbol-atlas caching --- .../plots/scatter_waveform/scatter_curve.py | 43 ++++++++++++++----- .../utility/widget_finder/widget_finder.py | 2 + tests/unit_tests/test_scatter_waveform.py | 42 ++++++++++++++++++ 3 files changed, 76 insertions(+), 11 deletions(-) diff --git a/bec_widgets/widgets/plots/scatter_waveform/scatter_curve.py b/bec_widgets/widgets/plots/scatter_waveform/scatter_curve.py index 5ca0e660..4303fa5a 100644 --- a/bec_widgets/widgets/plots/scatter_waveform/scatter_curve.py +++ b/bec_widgets/widgets/plots/scatter_waveform/scatter_curve.py @@ -7,6 +7,7 @@ import pyqtgraph as pg from bec_lib import bec_logger from pydantic import BaseModel, Field, ValidationError, field_validator from qtpy import QtCore +from qtpy.QtGui import QBrush, QColor from bec_widgets.utils.bec_connector import BECConnector, ConnectionConfig from bec_widgets.utils.colors import Colors @@ -16,6 +17,8 @@ if TYPE_CHECKING: # pragma: no cover logger = bec_logger.logger +Z_COLOR_LEVELS = 256 + # noinspection PyDataclass class ScatterDeviceSignal(BaseModel): @@ -60,6 +63,18 @@ class ScatterCurve(BECConnector, pg.PlotDataItem): USER_ACCESS = ["color_map"] + _brush_pools: dict[str, list[QBrush]] = {} + + @classmethod + def _brush_pool(cls, colormap: str) -> list[QBrush]: + """Return (creating once) the shared brush pool of one colormap.""" + pool = cls._brush_pools.get(colormap) + if pool is None: + lut = pg.colormap.get(colormap).getLookupTable(nPts=Z_COLOR_LEVELS, alpha=True) + pool = [QBrush(QColor(int(r), int(g), int(b), int(a))) for r, g, b, a in lut] + cls._brush_pools[colormap] = pool + return pool + def __init__( self, parent_item: ScatterWaveform, @@ -164,25 +179,31 @@ class ScatterCurve(BECConnector, pg.PlotDataItem): def _make_z_gradient(self, data_z: list | np.ndarray, colormap: str) -> list | None: """ - Make a gradient color for the z values. + Map the z values to brushes from the shared quantized colormap pool. Args: data_z(list|np.ndarray): Z values. colormap(str): Colormap for the gradient color. Returns: - list: List of colors for the z values. + list: List of shared QBrush objects for the z values. """ - # Normalize z_values for color mapping - z_min, z_max = np.min(data_z), np.max(data_z) - - if z_max != z_min: # Ensure that there is a range in the z values - z_values_norm = (data_z - z_min) / (z_max - z_min) - colormap = pg.colormap.get(colormap) # using colormap from global settings - colors = [colormap.map(z, mode="qcolor") for z in z_values_norm] - return colors - else: + data_z = np.atleast_1d(np.asarray(data_z, dtype=float)) + if data_z.size == 0: return None + finite = np.isfinite(data_z) + if not finite.any(): + return None + z_min, z_max = data_z[finite].min(), data_z[finite].max() + if z_max == z_min: + return None + z_values_norm = np.clip((data_z - z_min) / (z_max - z_min), 0.0, 1.0) + z_values_norm = np.nan_to_num(z_values_norm, nan=0.0) + indices = np.clip( + np.round(z_values_norm * (Z_COLOR_LEVELS - 1)).astype(int), 0, Z_COLOR_LEVELS - 1 + ) + pool = self._brush_pool(colormap) + return [pool[int(i)] for i in indices] def refresh_color_map(self, color_map: str): """ diff --git a/bec_widgets/widgets/utility/widget_finder/widget_finder.py b/bec_widgets/widgets/utility/widget_finder/widget_finder.py index bee2b1f0..3c6e81b7 100644 --- a/bec_widgets/widgets/utility/widget_finder/widget_finder.py +++ b/bec_widgets/widgets/utility/widget_finder/widget_finder.py @@ -15,6 +15,7 @@ from qtpy.QtWidgets import ( ) from bec_widgets import SafeProperty +from bec_widgets.utils.colors import apply_theme from bec_widgets.utils.widget_highlighter import WidgetHighlighter from bec_widgets.utils.widget_io import WidgetIO from bec_widgets.widgets.containers.main_window.main_window import BECMainWindowNoRPC @@ -202,6 +203,7 @@ if __name__ == "__main__": # pragma: no cover import sys app = QApplication(sys.argv) + apply_theme("dark") main_window = InspectorMainWindow() main_window.show() sys.exit(app.exec()) diff --git a/tests/unit_tests/test_scatter_waveform.py b/tests/unit_tests/test_scatter_waveform.py index a70f6f58..6c3dc87d 100644 --- a/tests/unit_tests/test_scatter_waveform.py +++ b/tests/unit_tests/test_scatter_waveform.py @@ -544,3 +544,45 @@ def test_scatter_curve_settings_fetch_all_properties(qtbot, mocked_client): assert settings.ui.device_x.currentText() == "samx" assert settings.ui.device_y.currentText() == "samy" assert settings.ui.device_z.currentText() == "bpm4i" + + +def test_z_gradient_uses_shared_brush_pool(qtbot, mocked_client): + """Brushes must be identity-stable pooled objects: pyqtgraph caches + rendered symbols by brush identity, and fresh per-point objects force a + full symbol-atlas rebuild on every update (the old dominant cost).""" + import pyqtgraph as pg + from qtpy.QtGui import QBrush + + from bec_widgets.widgets.plots.scatter_waveform.scatter_curve import Z_COLOR_LEVELS + + swf = create_widget(qtbot, ScatterWaveform, client=mocked_client) + swf.plot("samx", "samy", "bpm4i") + curve = swf.main_curve + + z = np.linspace(0, 1, 400) + brushes_first = curve._make_z_gradient(z, "plasma") + brushes_second = curve._make_z_gradient(z * 2 + 5, "plasma") # same normalized shape + + assert all(isinstance(b, QBrush) for b in brushes_first) + assert len({id(b) for b in brushes_first}) <= Z_COLOR_LEVELS + assert [id(b) for b in brushes_first] == [id(b) for b in brushes_second] + + cmap = pg.colormap.get("plasma") + lo = cmap.map(0.0, mode="qcolor") + hi = cmap.map(1.0, mode="qcolor") + assert brushes_first[0].color().getRgb()[:3] == lo.getRgb()[:3] + assert brushes_first[-1].color().getRgb()[:3] == hi.getRgb()[:3] + + +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) + swf.plot("samx", "samy", "bpm4i") + curve = swf.main_curve + assert curve._make_z_gradient([5.0, 5.0, 5.0], "plasma") is None + assert curve._make_z_gradient([], "plasma") is None + # live-path shapes that used to crash or warn: scalar, scalar NaN, mixed NaN + assert curve._make_z_gradient(3.7, "plasma") is None + 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