perf(scatter): pooled quantized z-brushes enable symbol-atlas caching

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):
         """

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

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