diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index d8dbdc1c..f489f954 100644 --- a/bec_widgets/cli/client.py +++ b/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"): """ diff --git a/bec_widgets/cli/generate_cli.py b/bec_widgets/cli/generate_cli.py index 46353569..9e3cc876 100644 --- a/bec_widgets/cli/generate_cli.py +++ b/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, diff --git a/bec_widgets/cli/server.py b/bec_widgets/cli/server.py index c7142a83..46f7abf9 100644 --- a/bec_widgets/cli/server.py +++ b/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 diff --git a/bec_widgets/widgets/__init__.py b/bec_widgets/widgets/__init__.py index bd648a71..be077df9 100644 --- a/bec_widgets/widgets/__init__.py +++ b/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 diff --git a/bec_widgets/widgets/figure/figure.py b/bec_widgets/widgets/figure/figure.py index b28d03e8..c3beda6f 100644 --- a/bec_widgets/widgets/figure/figure.py +++ b/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 - if x_name is not None and y_name is not None and x is None and y is None: + # User wants to add scan curve -> 1D Waveform + 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 - if x_name is not None and y_name is not None and x is None and y is None: + # User wants to add scan curve -> 1D Waveform + 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_custom( - x=[1, 2, 3, 4, 5], - y=[1, 2, 3, 4, 5], - label="curve-custom", - color="blue", - pen_style="dashdot", - ) + self.w1.add_curve_scan("samx", "samy", "bpm4i", pen_style="dash") + self.w1.add_curve_scan("samx", "samy", "bpm3a", pen_style="dash") + + # self.w1.add_curve_custom( + # x=[1, 2, 3, 4, 5], + # y=[1, 2, 3, 4, 5], + # label="curve-custom", + # color="blue", + # pen_style="dashdot", + # ) self.c1 = self.w1.get_config() # curves for w2 diff --git a/bec_widgets/widgets/monitor_scatter_2D/monitor_scatter_2D.py b/bec_widgets/widgets/monitor_scatter_2D/monitor_scatter_2D.py index 4a7a0636..e1f02922 100644 --- a/bec_widgets/widgets/monitor_scatter_2D/monitor_scatter_2D.py +++ b/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"}], }, }, { diff --git a/bec_widgets/widgets/plots/__init__.py b/bec_widgets/widgets/plots/__init__.py index 7b0cbbd0..a2f182c7 100644 --- a/bec_widgets/widgets/plots/__init__.py +++ b/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 diff --git a/bec_widgets/widgets/plots/motor_map.py b/bec_widgets/widgets/plots/motor_map.py index ee669c04..ea90c665 100644 --- a/bec_widgets/widgets/plots/motor_map.py +++ b/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): diff --git a/bec_widgets/widgets/plots/waveform1d.py b/bec_widgets/widgets/plots/waveform.py similarity index 88% rename from bec_widgets/widgets/plots/waveform1d.py rename to bec_widgets/widgets/plots/waveform.py index 2508bfc0..c0cea0ab 100644 --- a/bec_widgets/widgets/plots/waveform1d.py +++ b/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_name, y_name, x_entry, y_entry, validate_bec + x_entry, y_entry, z_entry = self._validate_signal_entries( + 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): """ diff --git a/tests/test_bec_figure.py b/tests/test_bec_figure.py index fb56db22..bda8355f 100644 --- a/tests/test_bec_figure.py +++ b/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_2"].config.widget_class == "BECWaveform1D" + assert bec_figure._widgets["widget_1"].config.widget_class == "BECWaveform" + 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 diff --git a/tests/test_bec_motor_map.py b/tests/test_bec_motor_map.py index 1545375a..41cd12af 100644 --- a/tests/test_bec_motor_map.py +++ b/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 diff --git a/tests/test_waveform1d.py b/tests/test_waveform1d.py index f9e20be3..9bf42f5a 100644 --- a/tests/test_waveform1d.py +++ b/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