From 52bc322b2b8d3ef92ff3480e61bddaf32464f976 Mon Sep 17 00:00:00 2001 From: wyzula-jan Date: Thu, 6 Jun 2024 14:25:42 +0200 Subject: [PATCH] refactor(figure): logic for .add_image and .image consolidated; logic for .add_plot and .plot consolidated --- bec_widgets/cli/client.py | 17 +- .../jupyter_console/jupyter_console_window.py | 7 +- bec_widgets/widgets/figure/figure.py | 292 +++++++++++------- tests/end-2-end/test_bec_figure_rpc_e2e.py | 4 +- tests/end-2-end/test_rpc_register_e2e.py | 2 +- 5 files changed, 194 insertions(+), 128 deletions(-) diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index 78df1acf..608d5856 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -462,14 +462,14 @@ class BECFigure(RPCBase): @rpc_call def add_plot( self, + x: "list | np.ndarray" = None, + y: "list | np.ndarray" = None, 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, @@ -483,7 +483,18 @@ class BECFigure(RPCBase): Add a Waveform1D plot to the figure at the specified position. Args: - widget_id(str): The unique identifier of the widget. If not provided, a unique ID will be generated. + x(list | np.ndarray): Custom x data to plot. + y(list | np.ndarray): Custom y data to plot. + 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. + 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. row(int): The row coordinate of the widget in the figure. If not provided, the next empty row will be used. col(int): The column coordinate of the widget in the figure. If not provided, the next empty column will be used. config(dict): Additional configuration for the widget. diff --git a/bec_widgets/examples/jupyter_console/jupyter_console_window.py b/bec_widgets/examples/jupyter_console/jupyter_console_window.py index ab336dd0..87734d29 100644 --- a/bec_widgets/examples/jupyter_console/jupyter_console_window.py +++ b/bec_widgets/examples/jupyter_console/jupyter_console_window.py @@ -97,14 +97,16 @@ class JupyterConsoleWindow(QWidget): # pragma: no cover: self.w3 = self.figure[1, 0] # curves for w1 - self.w1.add_curve_scan("samx", "samy", "bpm4i", pen_style="dash") - self.w1.add_curve_scan("samx", "samy", "bpm3a", pen_style="dash") + self.w1.plot(x_name="samx", y_name="samy", z_name="bpm4i") + self.w1.plot(x_name="samx", y_name="samy", z_name="bpm3a") self.c1 = self.w1.get_config() def _init_dock(self): self.d0 = self.dock.add_dock(name="dock_0") self.fig0 = self.d0.add_widget("BECFigure") + data = np.random.rand(10, 2) + self.fig0.plot(data, label="2d Data") self.fig0.image("eiger", vrange=(0, 100)) self.d1 = self.dock.add_dock(name="dock_1", position="right") @@ -114,7 +116,6 @@ class JupyterConsoleWindow(QWidget): # pragma: no cover: self.d2 = self.dock.add_dock(name="dock_2", position="bottom") self.fig2 = self.d2.add_widget("BECFigure", row=0, col=0) - self.fig2.motor_map(x_name="samx", y_name="samy") self.fig2.plot(x_name="samx", y_name="bpm4i") self.bar = self.d2.add_widget("SpiralProgressBar", row=0, col=1) self.bar.set_diameter(200) diff --git a/bec_widgets/widgets/figure/figure.py b/bec_widgets/widgets/figure/figure.py index 7d08d4d3..75382c79 100644 --- a/bec_widgets/widgets/figure/figure.py +++ b/bec_widgets/widgets/figure/figure.py @@ -184,8 +184,9 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): """ self._widgets = value - def add_plot( + def _init_waveform( self, + waveform, x_name: str = None, y_name: str = None, z_name: str = None, @@ -198,33 +199,45 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): color_map_z: Optional[str] = "plasma", label: Optional[str] = None, validate: bool = True, - row: int = None, - col: int = None, - config=None, - **axis_kwargs, - ) -> BECWaveform: + ): """ - Add a Waveform1D plot to the figure at the specified position. + Configure the waveform based on the provided parameters. Args: - widget_id(str): The unique identifier of the widget. If not provided, a unique ID will be generated. - row(int): The row coordinate of the widget in the figure. If not provided, the next empty row will be used. - col(int): The column coordinate of the widget in the figure. If not provided, the next empty column will be used. - config(dict): Additional configuration for the widget. - **axis_kwargs(dict): Additional axis properties to set on the widget after creation. + waveform (BECWaveform): The waveform to configure. + x (list | np.ndarray): Custom x data to plot. + y (list | np.ndarray): Custom y data to plot. + 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. + 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. """ - widget_id = str(uuid.uuid4()) - waveform = self.add_widget( - widget_type="Waveform1D", - widget_id=widget_id, - row=row, - col=col, - config=config, - **axis_kwargs, - ) - - # TODO remove repetition from .plot method - + if x is not None and y is None: + if isinstance(x, np.ndarray): + if x.ndim == 1: + y = np.arange(x.size) + waveform.add_curve_custom(x=np.arange(x.size), y=x, color=color, label=label) + return waveform + if x.ndim == 2: + waveform.add_curve_custom(x=x[:, 0], y=x[:, 1], color=color, label=label) + return waveform + elif isinstance(x, list): + y = np.arange(len(x)) + waveform.add_curve_custom(x=np.arange(len(x)), y=x, color=color, label=label) + return waveform + else: + raise ValueError( + "Invalid input. Provide either device names (x_name, y_name) or custom data." + ) + if x is not None and y is not None: + waveform.add_curve_custom(x=x, y=y, color=color, label=label) + return waveform # 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( @@ -262,6 +275,73 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): return waveform + def add_plot( + self, + x: list | np.ndarray = None, + y: list | np.ndarray = None, + x_name: str = None, + y_name: str = None, + z_name: str = None, + x_entry: str = None, + y_entry: str = None, + z_entry: str = 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, + ) -> BECWaveform: + """ + Add a Waveform1D plot to the figure at the specified position. + + Args: + x(list | np.ndarray): Custom x data to plot. + y(list | np.ndarray): Custom y data to plot. + 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. + 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. + row(int): The row coordinate of the widget in the figure. If not provided, the next empty row will be used. + col(int): The column coordinate of the widget in the figure. If not provided, the next empty column will be used. + config(dict): Additional configuration for the widget. + **axis_kwargs(dict): Additional axis properties to set on the widget after creation. + """ + widget_id = str(uuid.uuid4()) + waveform = self.add_widget( + widget_type="Waveform1D", + widget_id=widget_id, + row=row, + col=col, + config=config, + **axis_kwargs, + ) + + waveform = self._init_waveform( + waveform=waveform, + x=x, + y=y, + 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_z=color_map_z, + label=label, + validate=validate, + ) + return waveform + @typechecked def plot( self, @@ -309,70 +389,61 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): else: waveform = self.add_plot(**axis_kwargs) - if x is not None and y is None: - if isinstance(x, np.ndarray): - if x.ndim == 1: - y = np.arange(x.size) - waveform.add_curve_custom(x=np.arange(x.size), y=x, color=color, label=label) - return waveform - if x.ndim == 2: - waveform.add_curve_custom(x=x[:, 0], y=x[:, 1], color=color, label=label) - return waveform - elif isinstance(x, list): - y = np.arange(len(x)) - waveform.add_curve_custom(x=np.arange(len(x)), y=x, color=color, label=label) - return waveform - else: - raise ValueError( - "Invalid input. Provide either device names (x_name, y_name) or custom data." - ) - if x is not None and y is not None: - waveform.add_curve_custom(x=x, y=y, color=color, label=label) - return waveform - - # 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, - 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_z=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 and z_name is None - ): - waveform.add_curve_custom(x=x, y=y, color=color, label=label) - else: - raise ValueError( - "Invalid input. Provide either device names (x_name, y_name) or custom data." - ) + waveform = self._init_waveform( + waveform=waveform, + x=x, + y=y, + 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_z=color_map_z, + label=label, + validate=validate, + ) + # TODO remove repetition from .plot method return waveform + def _init_image( + self, + image, + monitor: str = None, + color_bar: Literal["simple", "full"] = "full", + color_map: str = "magma", + data: np.ndarray = None, + vrange: tuple[float, float] = None, + ) -> BECImageShow: + """ + Configure the image based on the provided parameters. + + Args: + image (BECImageShow): The image to configure. + monitor (str): The name of the monitor to display. + color_bar (Literal["simple","full"]): The type of color bar to display. + color_map (str): The color map to use for the image. + data (np.ndarray): Custom data to display. + """ + if monitor is not None and data is None: + image.add_monitor_image( + monitor=monitor, color_map=color_map, vrange=vrange, color_bar=color_bar + ) + elif data is not None and monitor is None: + image.add_custom_image( + name="custom", data=data, color_map=color_map, vrange=vrange, color_bar=color_bar + ) + elif data is None and monitor is None: + # Setting appearance + if vrange is not None: + image.set_vrange(vmin=vrange[0], vmax=vrange[1]) + if color_map is not None: + image.set_color_map(color_map) + else: + raise ValueError("Invalid input. Provide either monitor name or custom data.") + return image + def image( self, monitor: str = None, @@ -405,23 +476,14 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): else: image = self.add_image(color_bar=color_bar, **axis_kwargs) - # Setting data #TODO check logic if monitor or data are already created - if monitor is not None and data is None: - image.add_monitor_image( - monitor=monitor, color_map=color_map, vrange=vrange, color_bar=color_bar - ) - elif data is not None and monitor is None: - image.add_custom_image( - name="custom", data=data, color_map=color_map, vrange=vrange, color_bar=color_bar - ) - elif data is None and monitor is None: - # Setting appearance - if vrange is not None: - image.set_vrange(vmin=vrange[0], vmax=vrange[1]) - if color_map is not None: - image.set_color_map(color_map) - else: - raise ValueError("Invalid input. Provide either monitor name or custom data.") + image = self._init_image( + image=image, + monitor=monitor, + color_bar=color_bar, + color_map=color_map, + data=data, + vrange=vrange, + ) return image def add_image( @@ -472,22 +534,14 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): config=config, **axis_kwargs, ) - # TODO remove repetition from .image method - if monitor is not None and data is None: - image.add_monitor_image( - monitor=monitor, color_map=color_map, vrange=vrange, color_bar=color_bar - ) - elif data is not None and monitor is None: - image.add_custom_image( - name="custom", data=data, color_map=color_map, vrange=vrange, color_bar=color_bar - ) - elif data is None and monitor is None: - # Setting appearance - if vrange is not None: - image.set_vrange(vmin=vrange[0], vmax=vrange[1]) - if color_map is not None: - image.set_color_map(color_map) - + image = self._init_image( + image=image, + monitor=monitor, + color_bar=color_bar, + color_map=color_map, + data=data, + vrange=vrange, + ) return image def motor_map(self, motor_x: str = None, motor_y: str = None, **axis_kwargs) -> BECMotorMap: diff --git a/tests/end-2-end/test_bec_figure_rpc_e2e.py b/tests/end-2-end/test_bec_figure_rpc_e2e.py index 88baac07..e02179da 100644 --- a/tests/end-2-end/test_bec_figure_rpc_e2e.py +++ b/tests/end-2-end/test_bec_figure_rpc_e2e.py @@ -9,7 +9,7 @@ def test_rpc_waveform1d_custom_curve(rpc_server_figure): fig = BECFigure(rpc_server_figure) ax = fig.add_plot() - curve = ax.add_curve_custom([1, 2, 3], [1, 2, 3]) + curve = ax.plot(x=[1, 2, 3], y=[1, 2, 3]) curve.set_color("red") curve = ax.curves[0] curve.set_color("blue") @@ -24,7 +24,7 @@ def test_rpc_plotting_shortcuts_init_configs(rpc_server_figure, qtbot): plt = fig.plot(x_name="samx", y_name="bpm4i") im = fig.image("eiger") motor_map = fig.motor_map("samx", "samy") - plt_z = fig.add_plot("samx", "samy", "bpm4i") + plt_z = fig.add_plot(x_name="samx", y_name="samy", z_name="bpm4i") # Checking if classes are correctly initialised assert len(fig.widgets) == 4 diff --git a/tests/end-2-end/test_rpc_register_e2e.py b/tests/end-2-end/test_rpc_register_e2e.py index 5580e452..2314d3cb 100644 --- a/tests/end-2-end/test_rpc_register_e2e.py +++ b/tests/end-2-end/test_rpc_register_e2e.py @@ -9,7 +9,7 @@ def test_rpc_register_list_connections(rpc_server_figure): plt = fig.plot(x_name="samx", y_name="bpm4i") im = fig.image("eiger") motor_map = fig.motor_map("samx", "samy") - plt_z = fig.add_plot("samx", "samy", "bpm4i") + plt_z = fig.add_plot(x_name="samx", y_name="samy", z_name="bpm4i") # keep only class names from objects, since objects on server and client are different # so the best we can do is to compare types (rpc register is unit-tested elsewhere)