feat(plot/waveform1d): BECWaveform1D can show z data of scatter coded to different detector like BECMonitor2DScatter; BECWaveform1D name changed to BECWaveform

bec_widgets/cli/client.py

@@ -142,9 +142,12 @@ class BECWaveform1D(RPCBase):
         self,
         x_name: "str",
         y_name: "str",
+        z_name: "Optional[str]" = None,
         x_entry: "Optional[str]" = None,
         y_entry: "Optional[str]" = None,
+        z_entry: "Optional[str]" = None,
         color: "Optional[str]" = None,
+        color_map_z: "Optional[str]" = "plasma",
         label: "Optional[str]" = None,
         validate_bec: "bool" = True,
         **kwargs
@@ -156,7 +159,10 @@ class BECWaveform1D(RPCBase):
             x_entry(str): Entry of the x signal.
             y_name(str): Name of the y signal.
             y_entry(str): Entry of the y signal.
+            z_name(str): Name of the z signal.
+            z_entry(str): Entry of the z signal.
             color(str, optional): Color of the curve. Defaults to None.
+            color_map_z(str): The color map to use for the z-axis.
             label(str, optional): Label of the curve. Defaults to None.
             **kwargs: Additional keyword arguments for the curve configuration.
 
@@ -404,11 +410,14 @@ class BECFigure(RPCBase, BECFigureClientMixin):
         self,
         x_name: "str" = None,
         y_name: "str" = None,
+        z_name: "str" = None,
         x_entry: "str" = None,
         y_entry: "str" = None,
+        z_entry: "str" = None,
         x: "list | np.ndarray" = None,
         y: "list | np.ndarray" = None,
         color: "Optional[str]" = None,
+        color_map_z: "Optional[str]" = "plasma",
         label: "Optional[str]" = None,
         validate: "bool" = True,
         row: "int" = None,
@@ -484,11 +493,14 @@ class BECFigure(RPCBase, BECFigureClientMixin):
         self,
         x_name: "str" = None,
         y_name: "str" = None,
+        z_name: "str" = None,
         x_entry: "str" = None,
         y_entry: "str" = None,
+        z_entry: "str" = None,
         x: "list | np.ndarray" = None,
         y: "list | np.ndarray" = None,
         color: "Optional[str]" = None,
+        color_map_z: "Optional[str]" = "plasma",
         label: "Optional[str]" = None,
         validate: "bool" = True,
         **axis_kwargs
@@ -498,11 +510,14 @@ class BECFigure(RPCBase, BECFigureClientMixin):
         Args:
             x_name(str): The name of the device for the x-axis.
             y_name(str): The name of the device for the y-axis.
+            z_name(str): The name of the device for the z-axis.
             x_entry(str): The name of the entry for the x-axis.
             y_entry(str): The name of the entry for the y-axis.
+            z_entry(str): The name of the entry for the z-axis.
             x(list | np.ndarray): Custom x data to plot.
             y(list | np.ndarray): Custom y data to plot.
             color(str): The color of the curve.
+            color_map_z(str): The color map to use for the z-axis.
             label(str): The label of the curve.
             validate(bool): If True, validate the device names and entries.
             **axis_kwargs: Additional axis properties to set on the widget after creation.
@@ -634,6 +649,14 @@ class BECCurve(RPCBase):
             symbol_color(str, optional): Color of the symbol. Defaults to None.
         """
 
+    @rpc_call
+    def set_colormap(self, colormap: "str"):
+        """
+        Set the colormap for the scatter plot z gradient.
+        Args:
+            colormap(str): Colormap for the scatter plot.
+        """
+
     @rpc_call
     def set_symbol(self, symbol: "str"):
         """

bec_widgets/cli/generate_cli.py

@@ -109,15 +109,15 @@ if __name__ == "__main__":  # pragma: no cover
 
     from bec_widgets.utils import BECConnector
     from bec_widgets.widgets.figure import BECFigure
-    from bec_widgets.widgets.plots import BECImageShow, BECMotorMap, BECPlotBase, BECWaveform1D
+    from bec_widgets.widgets.plots import BECImageShow, BECMotorMap, BECPlotBase, BECWaveform
     from bec_widgets.widgets.plots.image import BECImageItem
-    from bec_widgets.widgets.plots.waveform1d import BECCurve
+    from bec_widgets.widgets.plots.waveform import BECCurve
 
     current_path = os.path.dirname(__file__)
     client_path = os.path.join(current_path, "client.py")
     clss = [
         BECPlotBase,
-        BECWaveform1D,
+        BECWaveform,
         BECFigure,
         BECCurve,
         BECImageShow,

bec_widgets/cli/server.py

@@ -6,11 +6,11 @@ from qtpy.QtCore import QTimer
 from bec_widgets.utils import BECDispatcher
 from bec_widgets.utils.bec_connector import BECConnector
 from bec_widgets.widgets.figure import BECFigure
-from bec_widgets.widgets.plots import BECCurve, BECImageShow, BECWaveform1D
+from bec_widgets.widgets.plots import BECCurve, BECImageShow, BECWaveform
 
 
 class BECWidgetsCLIServer:
-    WIDGETS = [BECWaveform1D, BECFigure, BECCurve, BECImageShow]
+    WIDGETS = [BECWaveform, BECFigure, BECCurve, BECImageShow]
 
     def __init__(self, gui_id: str = None, dispatcher: BECDispatcher = None) -> None:
         self.dispatcher = BECDispatcher() if dispatcher is None else dispatcher

bec_widgets/widgets/__init__.py

@@ -10,5 +10,5 @@ from .motor_control import (
     MotorThread,
 )
 from .motor_map import MotorMap
-from .plots import BECCurve, BECMotorMap, BECWaveform1D
+from .plots import BECCurve, BECMotorMap, BECWaveform
 from .scan_control import ScanControl

bec_widgets/widgets/figure/figure.py

@@ -19,7 +19,7 @@ from bec_widgets.widgets.plots import (
     BECImageShow,
     BECMotorMap,
     BECPlotBase,
-    BECWaveform1D,
+    BECWaveform,
     Waveform1DConfig,
     WidgetConfig,
 )
@@ -44,7 +44,7 @@ class WidgetHandler:
     def __init__(self):
         self.widget_factory = {
             "PlotBase": (BECPlotBase, WidgetConfig),
-            "Waveform1D": (BECWaveform1D, Waveform1DConfig),
+            "Waveform1D": (BECWaveform, Waveform1DConfig),
             "ImShow": (BECImageShow, ImageConfig),
             "MotorMap": (BECMotorMap, MotorMapConfig),
         }
@@ -164,18 +164,21 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         self,
         x_name: str = None,
         y_name: str = None,
+        z_name: str = None,
         x_entry: str = None,
         y_entry: str = None,
+        z_entry: str = None,
         x: list | np.ndarray = None,
         y: list | np.ndarray = None,
         color: Optional[str] = None,
+        color_map_z: Optional[str] = "plasma",
         label: Optional[str] = None,
         validate: bool = True,
         row: int = None,
         col: int = None,
         config=None,
         **axis_kwargs,
-    ) -> BECWaveform1D:
+    ) -> BECWaveform:
         """
         Add a Waveform1D plot to the figure at the specified position.
         Args:
@@ -197,8 +200,8 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
 
         # TODO remove repetition from .plot method
 
-        # User wants to add scan curve
+        # User wants to add scan curve -> 1D Waveform
-        if x_name is not None and y_name is not None and x is None and y is None:
+        if x_name is not None and y_name is not None and z_name is None and x is None and y is None:
             waveform.add_curve_scan(
                 x_name=x_name,
                 y_name=y_name,
@@ -208,7 +211,27 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
                 color=color,
                 label=label,
             )
-            # User wants to add custom curve
+        # User wants to add scan curve -> 2D Waveform Scatter
+        if (
+            x_name is not None
+            and y_name is not None
+            and z_name is not None
+            and x is None
+            and y is None
+        ):
+            waveform.add_curve_scan(
+                x_name=x_name,
+                y_name=y_name,
+                z_name=z_name,
+                x_entry=x_entry,
+                y_entry=y_entry,
+                z_entry=z_entry,
+                color=color,
+                color_map=color_map_z,
+                label=label,
+                validate=validate,
+            )
+        # User wants to add custom curve
         elif x is not None and y is not None and x_name is None and y_name is None:
             waveform.add_curve_custom(
                 x=x,
@@ -223,52 +246,81 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget):
         self,
         x_name: str = None,
         y_name: str = None,
+        z_name: str = None,
         x_entry: str = None,
         y_entry: str = None,
+        z_entry: str = None,
         x: list | np.ndarray = None,
         y: list | np.ndarray = None,
         color: Optional[str] = None,
+        color_map_z: Optional[str] = "plasma",
         label: Optional[str] = None,
         validate: bool = True,
         **axis_kwargs,
-    ) -> BECWaveform1D:
+    ) -> BECWaveform:
         """
         Add a 1D waveform plot to the figure. Always access the first waveform widget in the figure.
         Args:
             x_name(str): The name of the device for the x-axis.
             y_name(str): The name of the device for the y-axis.
+            z_name(str): The name of the device for the z-axis.
             x_entry(str): The name of the entry for the x-axis.
             y_entry(str): The name of the entry for the y-axis.
+            z_entry(str): The name of the entry for the z-axis.
             x(list | np.ndarray): Custom x data to plot.
             y(list | np.ndarray): Custom y data to plot.
             color(str): The color of the curve.
+            color_map_z(str): The color map to use for the z-axis.
             label(str): The label of the curve.
             validate(bool): If True, validate the device names and entries.
             **axis_kwargs: Additional axis properties to set on the widget after creation.
 
         Returns:
-            BECWaveform1D: The waveform plot widget.
+            BECWaveform: The waveform plot widget.
         """
-        waveform = self._find_first_widget_by_class(BECWaveform1D, can_fail=True)
+        waveform = self._find_first_widget_by_class(BECWaveform, can_fail=True)
         if waveform is not None:
             if axis_kwargs:
                 waveform.set(**axis_kwargs)
         else:
             waveform = self.add_plot(**axis_kwargs)
 
-        # User wants to add scan curve
+        # User wants to add scan curve -> 1D Waveform
-        if x_name is not None and y_name is not None and x is None and y is None:
+        if x_name is not None and y_name is not None and z_name is None and x is None and y is None:
             waveform.add_curve_scan(
                 x_name=x_name,
                 y_name=y_name,
                 x_entry=x_entry,
                 y_entry=y_entry,
-                validate=validate,
                 color=color,
+                color_map_z="plasma",
                 label=label,
+                validate=validate,
+            )
+        # User wants to add scan curve -> 2D Waveform Scatter
+        elif (
+            x_name is not None
+            and y_name is not None
+            and z_name is not None
+            and x is None
+            and y is None
+        ):
+            waveform.add_curve_scan(
+                x_name=x_name,
+                y_name=y_name,
+                z_name=z_name,
+                x_entry=x_entry,
+                y_entry=y_entry,
+                z_entry=z_entry,
+                color=color,
+                color_map=color_map_z,
+                label=label,
+                validate=validate,
             )
         # User wants to add custom curve
-        elif x is not None and y is not None and x_name is None and y_name is None:
+        elif (
+            x is not None and y is not None and x_name is None and y_name is None and z_name is None
+        ):
             waveform.add_curve_custom(
                 x=x,
                 y=y,
@@ -817,7 +869,8 @@ class DebugWindow(QWidget):  # pragma: no cover:
         self.console.set_default_style("linux")
 
     def _init_figure(self):
-        self.figure.add_widget(widget_type="Waveform1D", row=0, col=0, title="Widget 1")
+        # self.figure.add_widget(widget_type="Waveform1D", row=0, col=0, title="Widget 1")
+        self.figure.plot("samx", "bpm4d")
         self.figure.add_widget(widget_type="Waveform1D", row=0, col=1, title="Widget 2")
         self.figure.add_image(
             title="Image", row=1, col=0, color_map="viridis", color_bar="simple", vrange=(0, 100)
@@ -830,14 +883,16 @@ class DebugWindow(QWidget):  # pragma: no cover:
         self.w4 = self.figure[1, 1]
 
         # curves for w1
-        self.w1.add_curve_scan("samx", "bpm4i", pen_style="dash")
+        self.w1.add_curve_scan("samx", "samy", "bpm4i", pen_style="dash")
-        self.w1.add_curve_custom(
+        self.w1.add_curve_scan("samx", "samy", "bpm3a", pen_style="dash")
-            x=[1, 2, 3, 4, 5],
+
-            y=[1, 2, 3, 4, 5],
+        # self.w1.add_curve_custom(
-            label="curve-custom",
+        #     x=[1, 2, 3, 4, 5],
-            color="blue",
+        #     y=[1, 2, 3, 4, 5],
-            pen_style="dashdot",
+        #     label="curve-custom",
-        )
+        #     color="blue",
+        #     pen_style="dashdot",
+        # )
         self.c1 = self.w1.get_config()
 
         # curves for w2

bec_widgets/widgets/monitor_scatter_2D/monitor_scatter_2D.py

@@ -22,7 +22,7 @@ CONFIG_DEFAULT = {
             "signals": {
                 "x": [{"name": "samx", "entry": "samx"}],
                 "y": [{"name": "samy", "entry": "samy"}],
-                "z": [{"name": "gauss_bpm", "entry": "gauss_bpm"}],
+                "z": [{"name": "bpm4i", "entry": "bpm4i"}],
             },
         },
         {

bec_widgets/widgets/plots/__init__.py

@@ -1,4 +1,4 @@
 from .image import BECImageItem, BECImageShow, ImageItemConfig
 from .motor_map import BECMotorMap, MotorMapConfig
 from .plot_base import AxisConfig, BECPlotBase, WidgetConfig
-from .waveform1d import BECCurve, BECWaveform1D, Waveform1DConfig
+from .waveform import BECCurve, BECWaveform, Waveform1DConfig

bec_widgets/widgets/plots/motor_map.py

@@ -14,7 +14,7 @@ from qtpy.QtWidgets import QWidget
 
 from bec_widgets.utils import EntryValidator
 from bec_widgets.widgets.plots.plot_base import BECPlotBase, WidgetConfig
-from bec_widgets.widgets.plots.waveform1d import Signal, SignalData
+from bec_widgets.widgets.plots.waveform import Signal, SignalData
 
 
 class MotorMapConfig(WidgetConfig):

bec_widgets/widgets/plots/waveform.py

@@ -34,6 +34,7 @@ class Signal(BaseModel):
     source: str
     x: SignalData  # TODO maybe add metadata for config gui later
     y: SignalData
+    z: Optional[SignalData] = None
 
 
 class CurveConfig(ConnectionConfig):
@@ -49,12 +50,13 @@ class CurveConfig(ConnectionConfig):
     )
     source: Optional[str] = Field(None, description="The source of the curve.")
     signals: Optional[Signal] = Field(None, description="The signal of the curve.")
+    colormap: Optional[str] = Field("plasma", description="The colormap of the curves z gradient.")
 
 
 class Waveform1DConfig(WidgetConfig):
     color_palette: Literal["plasma", "viridis", "inferno", "magma"] = Field(
         "plasma", description="The color palette of the figure widget."
-    )
+    )  # TODO can be extended to all colormaps from current pyqtgraph session
     curves: dict[str, CurveConfig] = Field(
         {}, description="The list of curves to be added to the 1D waveform widget."
     )
@@ -65,6 +67,7 @@ class BECCurve(BECConnector, pg.PlotDataItem):
         "set",
         "set_data",
         "set_color",
+        "set_colormap",
         "set_symbol",
         "set_symbol_color",
         "set_symbol_size",
@@ -135,6 +138,7 @@ class BECCurve(BECConnector, pg.PlotDataItem):
         # Mapping of keywords to setter methods
         method_map = {
             "color": self.set_color,
+            "colormap": self.set_colormap,
             "symbol": self.set_symbol,
             "symbol_color": self.set_symbol_color,
             "symbol_size": self.set_symbol_size,
@@ -203,6 +207,14 @@ class BECCurve(BECConnector, pg.PlotDataItem):
         self.config.pen_style = pen_style
         self.apply_config()
 
+    def set_colormap(self, colormap: str):
+        """
+        Set the colormap for the scatter plot z gradient.
+        Args:
+            colormap(str): Colormap for the scatter plot.
+        """
+        self.config.colormap = colormap
+
     def get_data(self) -> tuple[np.ndarray, np.ndarray]:
         """
         Get the data of the curve.
@@ -213,7 +225,7 @@ class BECCurve(BECConnector, pg.PlotDataItem):
         return x_data, y_data
 
 
-class BECWaveform1D(BECPlotBase):
+class BECWaveform(BECPlotBase):
     USER_ACCESS = [
         "add_curve_scan",
         "add_curve_custom",
@@ -467,9 +479,12 @@ class BECWaveform1D(BECPlotBase):
         self,
         x_name: str,
         y_name: str,
+        z_name: Optional[str] = None,
         x_entry: Optional[str] = None,
         y_entry: Optional[str] = None,
+        z_entry: Optional[str] = None,
         color: Optional[str] = None,
+        color_map_z: Optional[str] = "plasma",
         label: Optional[str] = None,
         validate_bec: bool = True,
         **kwargs,
@@ -481,7 +496,10 @@ class BECWaveform1D(BECPlotBase):
             x_entry(str): Entry of the x signal.
             y_name(str): Name of the y signal.
             y_entry(str): Entry of the y signal.
+            z_name(str): Name of the z signal.
+            z_entry(str): Entry of the z signal.
             color(str, optional): Color of the curve. Defaults to None.
+            color_map_z(str): The color map to use for the z-axis.
             label(str, optional): Label of the curve. Defaults to None.
             **kwargs: Additional keyword arguments for the curve configuration.
 
@@ -492,11 +510,14 @@ class BECWaveform1D(BECPlotBase):
         curve_source = "scan_segment"
 
         # Get entry if not provided and validate
-        x_entry, y_entry = self._validate_signal_entries(
+        x_entry, y_entry, z_entry = self._validate_signal_entries(
-            x_name, y_name, x_entry, y_entry, validate_bec
+            x_name, y_name, z_name, x_entry, y_entry, z_entry, validate_bec
         )
 
-        label = label or f"{y_name}-{y_entry}"
+        if z_name is not None and z_entry is not None:
+            label = label or f"{z_name}-{z_entry}"
+        else:
+            label = label or f"{y_name}-{y_entry}"
 
         curve_exits = self._check_curve_id(label, self._curves_data)
         if curve_exits:
@@ -515,11 +536,13 @@ class BECWaveform1D(BECPlotBase):
             parent_id=self.gui_id,
             label=label,
             color=color,
+            color_map=color_map_z,
             source=curve_source,
             signals=Signal(
                 source=curve_source,
                 x=SignalData(name=x_name, entry=x_entry),
                 y=SignalData(name=y_name, entry=y_entry),
+                z=SignalData(name=z_name, entry=z_entry) if z_name else None,
             ),
             **kwargs,
         )
@@ -530,28 +553,35 @@ class BECWaveform1D(BECPlotBase):
         self,
         x_name: str,
         y_name: str,
+        z_name: str | None,
         x_entry: str | None,
         y_entry: str | None,
+        z_entry: str | None,
         validate_bec: bool = True,
-    ) -> tuple[str, str]:
+    ) -> tuple[str, str, str | None]:
         """
         Validate the signal name and entry.
         Args:
             x_name(str): Name of the x signal.
             y_name(str): Name of the y signal.
+            z_name(str): Name of the z signal.
             x_entry(str|None): Entry of the x signal.
             y_entry(str|None): Entry of the y signal.
+            z_entry(str|None): Entry of the z signal.
             validate_bec(bool, optional): If True, validate the signal with BEC. Defaults to True.
         Returns:
-            tuple[str,str]: Validated x and y entries.
+            tuple[str,str,str|None]: Validated x, y, z entries.
         """
         if validate_bec:
             x_entry = self.entry_validator.validate_signal(x_name, x_entry)
             y_entry = self.entry_validator.validate_signal(y_name, y_entry)
+            if z_name:
+                z_entry = self.entry_validator.validate_signal(z_name, z_entry)
         else:
             x_entry = x_name if x_entry is None else x_entry
             y_entry = y_name if y_entry is None else y_entry
-        return x_entry, y_entry
+            z_entry = z_name if z_entry is None else z_entry
+        return x_entry, y_entry, z_entry
 
     def _check_curve_id(self, val: Any, dict_to_check: dict) -> bool:
         """
@@ -654,19 +684,54 @@ class BECWaveform1D(BECPlotBase):
         Args:
             data(ScanData): Data from the scan segment.
         """
+        data_x = None
+        data_y = None
+        data_z = None
         for curve_id, curve in self._curves_data["scan_segment"].items():
             x_name = curve.config.signals.x.name
             x_entry = curve.config.signals.x.entry
             y_name = curve.config.signals.y.name
             y_entry = curve.config.signals.y.entry
+            if curve.config.signals.z:
+                z_name = curve.config.signals.z.name
+                z_entry = curve.config.signals.z.entry
 
             try:
                 data_x = data[x_name][x_entry].val
                 data_y = data[y_name][y_entry].val
+                if curve.config.signals.z:
+                    data_z = data[z_name][z_entry].val
+                    color_z = self._make_z_gradient(
+                        data_z, curve.config.colormap
+                    )  # TODO decide how to implement custom gradient
             except TypeError:
                 continue
 
-            curve.setData(data_x, data_y)
+            if data_z is not None and color_z is not None:
+                curve.setData(x=data_x, y=data_y, symbolBrush=color_z)
+            else:
+                curve.setData(data_x, data_y)
+
+    def _make_z_gradient(self, data_z: list | np.ndarray, colormap: str) -> list | None:
+        """
+        Make a gradient color for the z values.
+        Args:
+            data_z(list|np.ndarray): Z values.
+            colormap(str): Colormap for the gradient color.
+
+        Returns:
+            list: List of colors 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:
+            return None
 
     def scan_history(self, scan_index: int = None, scan_id: str = None):
         """

tests/test_bec_figure.py

@@ -5,7 +5,7 @@ from unittest.mock import MagicMock
 import numpy as np
 import pytest
 
-from bec_widgets.widgets import BECFigure, BECMotorMap, BECWaveform1D
+from bec_widgets.widgets import BECFigure, BECMotorMap, BECWaveform
 from bec_widgets.widgets.plots import BECImageShow
 
 from .client_mocks import mocked_client
@@ -50,8 +50,8 @@ def test_bec_figure_add_remove_plot(bec_figure):
     assert "widget_1" in bec_figure._widgets
     assert "widget_2" in bec_figure._widgets
     assert "widget_3" in bec_figure._widgets
-    assert bec_figure._widgets["widget_1"].config.widget_class == "BECWaveform1D"
+    assert bec_figure._widgets["widget_1"].config.widget_class == "BECWaveform"
-    assert bec_figure._widgets["widget_2"].config.widget_class == "BECWaveform1D"
+    assert bec_figure._widgets["widget_2"].config.widget_class == "BECWaveform"
     assert bec_figure._widgets["widget_3"].config.widget_class == "BECPlotBase"
 
     # Check accessing positions by the grid in figure
@@ -64,7 +64,7 @@ def test_bec_figure_add_remove_plot(bec_figure):
     assert len(bec_figure._widgets) == initial_count + 2
     assert "widget_1" not in bec_figure._widgets
     assert "widget_3" in bec_figure._widgets
-    assert bec_figure._widgets["widget_2"].config.widget_class == "BECWaveform1D"
+    assert bec_figure._widgets["widget_2"].config.widget_class == "BECWaveform"
 
 
 def test_add_different_types_of_widgets(bec_figure):
@@ -72,7 +72,7 @@ def test_add_different_types_of_widgets(bec_figure):
     im = bec_figure.image("eiger")
     motor_map = bec_figure.motor_map("samx", "samy")
 
-    assert plt.__class__ == BECWaveform1D
+    assert plt.__class__ == BECWaveform
     assert im.__class__ == BECImageShow
     assert motor_map.__class__ == BECMotorMap
 

tests/test_bec_motor_map.py

@@ -2,7 +2,7 @@ import pytest
 
 from bec_widgets.widgets import BECMotorMap
 from bec_widgets.widgets.plots.motor_map import MotorMapConfig
-from bec_widgets.widgets.plots.waveform1d import Signal, SignalData
+from bec_widgets.widgets.plots.waveform import Signal, SignalData
 
 from .client_mocks import mocked_client
 

tests/test_waveform1d.py

@@ -4,7 +4,7 @@ from unittest.mock import MagicMock
 import numpy as np
 import pytest
 
-from bec_widgets.widgets.plots.waveform1d import CurveConfig, Signal, SignalData
+from bec_widgets.widgets.plots.waveform import CurveConfig, Signal, SignalData
 
 from .client_mocks import mocked_client
 from .test_bec_figure import bec_figure
@@ -49,7 +49,7 @@ def test_adding_curve_with_same_id(bec_figure):
 
 def test_create_waveform1D_by_config(bec_figure):
     w1_config_input = {
-        "widget_class": "BECWaveform1D",
+        "widget_class": "BECWaveform",
         "gui_id": "widget_1",
         "parent_id": "BECFigure_1708689320.788527",
         "row": 0,
@@ -73,6 +73,7 @@ def test_create_waveform1D_by_config(bec_figure):
                 "parent_id": "widget_1",
                 "label": "bpm4i-bpm4i",
                 "color": "#cc4778",
+                "colormap": "plasma",
                 "symbol": "o",
                 "symbol_color": None,
                 "symbol_size": 5,
@@ -95,6 +96,7 @@ def test_create_waveform1D_by_config(bec_figure):
                         "modifier": None,
                         "limits": None,
                     },
+                    "z": None,
                 },
             },
             "curve-custom": {
@@ -103,6 +105,7 @@ def test_create_waveform1D_by_config(bec_figure):
                 "parent_id": "widget_1",
                 "label": "curve-custom",
                 "color": "blue",
+                "colormap": "plasma",
                 "symbol": "o",
                 "symbol_color": None,
                 "symbol_size": 5,
@@ -232,6 +235,7 @@ def test_change_curve_appearance_methods(bec_figure, qtbot):
         "source": "scan_segment",
         "x": {"name": "samx", "entry": "samx", "unit": None, "modifier": None, "limits": None},
         "y": {"name": "bpm4i", "entry": "bpm4i", "unit": None, "modifier": None, "limits": None},
+        "z": None,
     }
 
 
@@ -260,6 +264,7 @@ def test_change_curve_appearance_args(bec_figure):
         "source": "scan_segment",
         "x": {"name": "samx", "entry": "samx", "unit": None, "modifier": None, "limits": None},
         "y": {"name": "bpm4i", "entry": "bpm4i", "unit": None, "modifier": None, "limits": None},
+        "z": None,
     }
 
 
@@ -343,6 +348,7 @@ def test_curve_add_by_config(bec_figure):
         "parent_id": "widget_1",
         "label": "bpm4i-bpm4i",
         "color": "#cc4778",
+        "colormap": "plasma",
         "symbol": "o",
         "symbol_color": None,
         "symbol_size": 5,
@@ -359,6 +365,7 @@ def test_curve_add_by_config(bec_figure):
                 "modifier": None,
                 "limits": None,
             },
+            "z": None,
         },
     }
 
@@ -428,3 +435,43 @@ def test_scan_history_with_val_access(bec_figure, qtbot):
 
     assert np.array_equal(x_data, [1, 2, 3])
     assert np.array_equal(y_data, [4, 5, 6])
+
+
+def test_scatter_2d_update(bec_figure, qtbot):
+    w1 = bec_figure.add_plot()
+
+    c1 = w1.add_curve_scan(x_name="samx", y_name="samx", z_name="bpm4i")
+
+    msg = {
+        "data": {
+            "samx": {"samx": {"value": [1, 2, 3]}},
+            "samy": {"samy": {"value": [4, 5, 6]}},
+            "bpm4i": {"bpm4i": {"value": [1, 3, 2]}},
+        },
+        "scan_id": 1,
+    }
+    msg_metadata = {"scan_name": "line_scan"}
+
+    mock_scan_data = MagicMock()
+    mock_scan_data.data = {
+        device_name: {
+            entry: MagicMock(val=msg["data"][device_name][entry]["value"])
+            for entry in msg["data"][device_name]
+        }
+        for device_name in msg["data"]
+    }
+
+    w1.queue.scan_storage.find_scan_by_ID.return_value = mock_scan_data
+
+    w1.on_scan_segment(msg, msg_metadata)
+    qtbot.wait(500)
+
+    data = c1.get_data()
+    expected_x_y_data = ([1, 2, 3], [1, 2, 3])
+    expected_z_colors = w1._make_z_gradient([1, 3, 2], "plasma")
+
+    scatter_points = c1.scatter.points()
+    colors = [point.brush().color() for point in scatter_points]
+
+    assert np.array_equal(data, expected_x_y_data)
+    assert colors == expected_z_colors