diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py index 67b1ccf5..cba63b9d 100644 --- a/bec_widgets/cli/client.py +++ b/bec_widgets/cli/client.py @@ -31,6 +31,13 @@ class BECCurve(RPCBase): Remove the curve from the plot. """ + @property + @rpc_call + def dap_params(self): + """ + None + """ + @property @rpc_call def rpc_id(self) -> "str": @@ -143,6 +150,13 @@ class BECCurve(RPCBase): tuple[np.ndarray,np.ndarray]: X and Y data of the curve. """ + @property + @rpc_call + def dap_params(self): + """ + None + """ + class BECDock(RPCBase): @property @@ -457,6 +471,7 @@ class BECFigure(RPCBase): row: "int" = None, col: "int" = None, config=None, + dap: "str | None" = None, **axis_kwargs, ) -> "BECWaveform": """ @@ -550,6 +565,7 @@ class BECFigure(RPCBase): color_map_z: "str | None" = "plasma", label: "str | None" = None, validate: "bool" = True, + dap: "str | None" = None, **axis_kwargs, ) -> "BECWaveform": """ @@ -568,6 +584,7 @@ class BECFigure(RPCBase): 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. + dap(str): The DAP model to use for the curve. **axis_kwargs: Additional axis properties to set on the widget after creation. Returns: @@ -1467,6 +1484,7 @@ class BECWaveform(RPCBase): color_map_z: "str | None" = "plasma", label: "str | None" = None, validate: "bool" = True, + dap: "str | None" = None, ) -> "BECCurve": """ Plot a curve to the plot widget. @@ -1483,11 +1501,50 @@ class BECWaveform(RPCBase): 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. + dap(str): The dap model to use for the curve. If not specified, none will be added. Returns: BECCurve: The curve object. """ + @rpc_call + def add_dap( + self, + x_name: "str", + y_name: "str", + x_entry: "Optional[str]" = None, + y_entry: "Optional[str]" = None, + color: "Optional[str]" = None, + dap: "str" = "GaussianModel", + **kwargs, + ) -> "BECCurve": + """ + Add LMFIT dap model curve to the plot widget. + + Args: + x_name(str): Name of the x signal. + x_entry(str): Entry of the x signal. + y_name(str): Name of the y signal. + y_entry(str): Entry of the y 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. + dap(str): The dap model to use for the curve. + **kwargs: Additional keyword arguments for the curve configuration. + + Returns: + BECCurve: The curve object. + """ + + @rpc_call + def get_dap_params(self) -> "dict": + """ + Get the DAP parameters of all DAP curves. + + Returns: + dict: DAP parameters of all DAP curves. + """ + @rpc_call def remove_curve(self, *identifiers): """ diff --git a/bec_widgets/examples/jupyter_console/jupyter_console_window.py b/bec_widgets/examples/jupyter_console/jupyter_console_window.py index 174b19f3..4eb837af 100644 --- a/bec_widgets/examples/jupyter_console/jupyter_console_window.py +++ b/bec_widgets/examples/jupyter_console/jupyter_console_window.py @@ -14,21 +14,6 @@ from bec_widgets.widgets.dock.dock_area import BECDockArea from bec_widgets.widgets.figure import BECFigure from bec_widgets.widgets.jupyter_console.jupyter_console import BECJupyterConsole -# class JupyterConsoleWidget(RichJupyterWidget): # pragma: no cover: -# def __init__(self): -# super().__init__() -# -# self.kernel_manager = QtInProcessKernelManager() -# self.kernel_manager.start_kernel(show_banner=False) -# self.kernel_client = self.kernel_manager.client() -# self.kernel_client.start_channels() -# -# self.kernel_manager.kernel.shell.push({"np": np, "pg": pg}) -# -# def shutdown_kernel(self): -# self.kernel_client.stop_channels() -# self.kernel_manager.shutdown_kernel() - class JupyterConsoleWindow(QWidget): # pragma: no cover: """A widget that contains a Jupyter console linked to BEC Widgets with full API access (contains Qt and pyqtgraph API).""" @@ -61,6 +46,7 @@ class JupyterConsoleWindow(QWidget): # pragma: no cover: "fig0": self.fig0, "fig1": self.fig1, "fig2": self.fig2, + "plt": self.plt, "bar": self.bar, } ) @@ -115,7 +101,8 @@ 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.plot(x_name="samx", y_name="bpm4i") + self.plt = self.fig2.plot(x_name="samx", y_name="bpm3a") + self.plt.plot(x_name="samx", y_name="bpm4i", dap="GaussianModel") self.bar = self.d2.add_widget("RingProgressBar", row=0, col=1) self.bar.set_diameter(200) diff --git a/bec_widgets/utils/bec_connector.py b/bec_widgets/utils/bec_connector.py index 97aba79d..7362d077 100644 --- a/bec_widgets/utils/bec_connector.py +++ b/bec_widgets/utils/bec_connector.py @@ -228,6 +228,7 @@ class BECConnector: all_connections = self.rpc_register.list_all_connections() if len(all_connections) == 0: print("No more connections. Shutting down GUI BEC client.") + self.bec_dispatcher.disconnect_all() self.client.shutdown() # def closeEvent(self, event): diff --git a/bec_widgets/widgets/figure/figure.py b/bec_widgets/widgets/figure/figure.py index 8416f799..e25c8f00 100644 --- a/bec_widgets/widgets/figure/figure.py +++ b/bec_widgets/widgets/figure/figure.py @@ -195,10 +195,11 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): 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, + color: str | None = None, + color_map_z: str | None = "plasma", + label: str | None = None, validate: bool = True, + dap: str | None = None, ): """ Configure the waveform based on the provided parameters. @@ -217,6 +218,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): 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. + dap (str): The DAP model to use for the curve. """ if x is not None and y is None: if isinstance(x, np.ndarray): @@ -240,7 +242,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): 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( + waveform.plot( x_name=x_name, y_name=y_name, x_entry=x_entry, @@ -248,6 +250,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): validate=validate, color=color, label=label, + dap=dap, ) # User wants to add scan curve -> 2D Waveform Scatter if ( @@ -257,7 +260,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): and x is None and y is None ): - waveform.add_curve_scan( + waveform.plot( x_name=x_name, y_name=y_name, z_name=z_name, @@ -268,6 +271,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): color_map_z=color_map_z, label=label, validate=validate, + dap=dap, ) # 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: @@ -292,6 +296,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): row: int = None, col: int = None, config=None, + dap: str | None = None, **axis_kwargs, ) -> BECWaveform: """ @@ -339,6 +344,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): color_map_z=color_map_z, label=label, validate=validate, + dap=dap, ) return waveform @@ -357,6 +363,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): color_map_z: str | None = "plasma", label: str | None = None, validate: bool = True, + dap: str | None = None, **axis_kwargs, ) -> BECWaveform: """ @@ -375,6 +382,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): 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. + dap(str): The DAP model to use for the curve. **axis_kwargs: Additional axis properties to set on the widget after creation. Returns: @@ -403,6 +411,7 @@ class BECFigure(BECConnector, pg.GraphicsLayoutWidget): color_map_z=color_map_z, label=label, validate=validate, + dap=dap, ) # TODO remove repetition from .plot method return waveform diff --git a/bec_widgets/widgets/figure/plots/waveform/waveform.py b/bec_widgets/widgets/figure/plots/waveform/waveform.py index 1309dc44..46d925e9 100644 --- a/bec_widgets/widgets/figure/plots/waveform/waveform.py +++ b/bec_widgets/widgets/figure/plots/waveform/waveform.py @@ -1,10 +1,12 @@ from __future__ import annotations +import time from collections import defaultdict from typing import Any, Literal, Optional import numpy as np import pyqtgraph as pg +from bec_lib import messages from bec_lib.endpoints import MessageEndpoints from bec_lib.scan_data import ScanData from pydantic import Field, ValidationError @@ -36,6 +38,8 @@ class BECWaveform(BECPlotBase): "rpc_id", "config_dict", "plot", + "add_dap", + "get_dap_params", "remove_curve", "scan_history", "curves", @@ -57,6 +61,7 @@ class BECWaveform(BECPlotBase): "set_legend_label_size", ] scan_signal_update = pyqtSignal() + dap_params_update = pyqtSignal(dict) def __init__( self, @@ -73,6 +78,7 @@ class BECWaveform(BECPlotBase): ) self._curves_data = defaultdict(dict) + self.old_scan_id = None self.scan_id = None # Scan segment update proxy @@ -80,6 +86,9 @@ class BECWaveform(BECPlotBase): self.scan_signal_update, rateLimit=25, slot=self._update_scan_segment_plot ) + self.proxy_update_dap = pg.SignalProxy( + self.scan_signal_update, rateLimit=25, slot=self.refresh_dap + ) # Get bec shortcuts dev, scans, queue, scan_storage, dap self.get_bec_shortcuts() @@ -213,6 +222,7 @@ class BECWaveform(BECPlotBase): color_map_z: str | None = "plasma", label: str | None = None, validate: bool = True, + dap: str | None = None, # TODO add dap custom curve wrapper ) -> BECCurve: """ Plot a curve to the plot widget. @@ -229,6 +239,7 @@ class BECWaveform(BECPlotBase): 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. + dap(str): The dap model to use for the curve. If not specified, none will be added. Returns: BECCurve: The curve object. @@ -237,6 +248,8 @@ class BECWaveform(BECPlotBase): if x is not None and y is not None: return self.add_curve_custom(x=x, y=y, label=label, color=color) else: + if dap: + self.add_dap(x_name=x_name, y_name=y_name, dap=dap) return self.add_curve_scan( x_name=x_name, y_name=y_name, @@ -256,6 +269,7 @@ class BECWaveform(BECPlotBase): y: list | np.ndarray, label: str = None, color: str = None, + curve_source: str = "custom", **kwargs, ) -> BECCurve: """ @@ -266,12 +280,13 @@ class BECWaveform(BECPlotBase): y(list|np.ndarray): Y data of the curve. label(str, optional): Label of the curve. Defaults to None. color(str, optional): Color of the curve. Defaults to None. + curve_source(str, optional): Tag for source of the curve. Defaults to "custom". **kwargs: Additional keyword arguments for the curve configuration. Returns: BECCurve: The curve object. """ - curve_source = "custom" + curve_source = curve_source curve_id = label or f"Curve {len(self.plot_item.curves) + 1}" curve_exits = self._check_curve_id(curve_id, self._curves_data) @@ -314,10 +329,12 @@ class BECWaveform(BECPlotBase): color_map_z: Optional[str] = "plasma", label: Optional[str] = None, validate_bec: bool = True, + source: str = "scan_segment", + dap: Optional[str] = None, **kwargs, ) -> BECCurve: """ - Add a curve to the plot widget from the scan segment. + Add a curve to the plot widget from the scan segment. #TODO adapt docs to DAP Args: x_name(str): Name of the x signal. @@ -335,7 +352,7 @@ class BECWaveform(BECPlotBase): BECCurve: The curve object. """ # Check if curve already exists - curve_source = "scan_segment" + curve_source = source # Get entry if not provided and validate x_entry, y_entry, z_entry = self._validate_signal_entries( @@ -371,12 +388,74 @@ class BECWaveform(BECPlotBase): 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, + dap=dap, ), **kwargs, ) curve = self._add_curve_object(name=label, source=curve_source, config=curve_config) return curve + def add_dap( + self, + x_name: str, + y_name: str, + x_entry: Optional[str] = None, + y_entry: Optional[str] = None, + color: Optional[str] = None, + dap: str = "GaussianModel", + **kwargs, + ) -> BECCurve: + """ + Add LMFIT dap model curve to the plot widget. + + Args: + x_name(str): Name of the x signal. + x_entry(str): Entry of the x signal. + y_name(str): Name of the y signal. + y_entry(str): Entry of the y 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. + dap(str): The dap model to use for the curve. + **kwargs: Additional keyword arguments for the curve configuration. + + Returns: + BECCurve: The curve object. + """ + x_entry, y_entry, _ = self._validate_signal_entries( + x_name, y_name, None, x_entry, y_entry, None + ) + label = f"{y_name}-{y_entry}-{dap}" + curve = self.add_curve_scan( + x_name=x_name, + y_name=y_name, + x_entry=x_entry, + y_entry=y_entry, + color=color, + label=label, + source="DAP", + dap=dap, + pen_style="dash", + symbol="star", + **kwargs, + ) + + self.setup_dap(self.old_scan_id, self.scan_id) + self.refresh_dap() + return curve + + def get_dap_params(self) -> dict: + """ + Get the DAP parameters of all DAP curves. + + Returns: + dict: DAP parameters of all DAP curves. + """ + params = {} + for curve_id, curve in self._curves_data["DAP"].items(): + params[curve_id] = curve.dap_params + return params + def _add_curve_object( self, name: str, @@ -528,13 +607,75 @@ class BECWaveform(BECPlotBase): return if current_scan_id != self.scan_id: + self.old_scan_id = self.scan_id self.scan_id = current_scan_id self.scan_segment_data = self.queue.scan_storage.find_scan_by_ID( self.scan_id ) # TODO do scan access through BECFigure + self.setup_dap(self.old_scan_id, self.scan_id) self.scan_signal_update.emit() + def setup_dap(self, old_scan_id, new_scan_id): + """ + Setup DAP for the new scan. + + Args: + old_scan_id(str): old_scan_id, used to disconnect the previous dispatcher connection. + new_scan_id(str): new_scan_id, used to connect the new dispatcher connection. + + """ + self.bec_dispatcher.disconnect_slot( + self.update_dap, MessageEndpoints.dap_response(old_scan_id) + ) + if len(self._curves_data["DAP"]) > 0: + self.bec_dispatcher.connect_slot( + self.update_dap, MessageEndpoints.dap_response(new_scan_id) + ) + + def refresh_dap(self): + """ + Refresh the DAP curves with the latest data from the DAP model MessageEndpoints.dap_response(). + """ + for curve_id, curve in self._curves_data["DAP"].items(): + x_name = curve.config.signals.x.name + y_name = curve.config.signals.y.name + x_entry = curve.config.signals.x.entry + y_entry = curve.config.signals.y.entry + model_name = curve.config.signals.dap + model = getattr(self.dap, model_name) + + msg = messages.DAPRequestMessage( + dap_cls="LmfitService1D", + dap_type="on_demand", + config={ + "args": [self.scan_id, x_name, x_entry, y_name, y_entry], + "kwargs": {}, + "class_args": model._plugin_info["class_args"], + "class_kwargs": model._plugin_info["class_kwargs"], + }, + metadata={"RID": self.scan_id}, + ) + self.client.connector.set_and_publish(MessageEndpoints.dap_request(), msg) + + @pyqtSlot(dict, dict) + def update_dap(self, msg, metadata): + self.msg = msg + scan_id, x_name, x_entry, y_name, y_entry = msg["dap_request"].content["config"]["args"] + model = msg["dap_request"].content["config"]["class_kwargs"]["model"] + + curve_id_request = f"{y_name}-{y_entry}-{model}" + + for curve_id, curve in self._curves_data["DAP"].items(): + if curve_id == curve_id_request: + if msg["data"] is not None: + x = msg["data"][0]["x"] + y = msg["data"][0]["y"] + curve.setData(x, y) + curve.dap_params = msg["data"][1]["fit_parameters"] + self.dap_params_update.emit(curve.dap_params) + break + def _update_scan_segment_plot(self): """Update the plot with the data from the scan segment.""" data = self.scan_segment_data.data @@ -609,13 +750,17 @@ class BECWaveform(BECPlotBase): if scan_index is not None and scan_id is not None: raise ValueError("Only one of scan_id or scan_index can be provided.") + # Reset DAP connector + self.bec_dispatcher.disconnect_slot( + self.update_dap, MessageEndpoints.dap_response(self.scan_id) + ) if scan_index is not None: self.scan_id = self.queue.scan_storage.storage[scan_index].scan_id - data = self.queue.scan_storage.find_scan_by_ID(self.scan_id).data elif scan_id is not None: self.scan_id = scan_id - data = self.queue.scan_storage.find_scan_by_ID(self.scan_id).data + self.setup_dap(self.old_scan_id, self.scan_id) + data = self.queue.scan_storage.find_scan_by_ID(self.scan_id).data self._update_scan_curves(data) def get_all_data(self, output: Literal["dict", "pandas"] = "dict") -> dict | pd.DataFrame: @@ -661,6 +806,9 @@ class BECWaveform(BECPlotBase): def cleanup(self): """Cleanup the widget connection from BECDispatcher.""" self.bec_dispatcher.disconnect_slot(self.on_scan_segment, MessageEndpoints.scan_segment()) + self.bec_dispatcher.disconnect_slot( + self.update_dap, MessageEndpoints.dap_response(self.scan_id) + ) for curve in self.curves: curve.cleanup() super().cleanup() diff --git a/bec_widgets/widgets/figure/plots/waveform/waveform_curve.py b/bec_widgets/widgets/figure/plots/waveform/waveform_curve.py index 145bb18f..4bf3c1ad 100644 --- a/bec_widgets/widgets/figure/plots/waveform/waveform_curve.py +++ b/bec_widgets/widgets/figure/plots/waveform/waveform_curve.py @@ -31,6 +31,7 @@ class Signal(BaseModel): x: SignalData # TODO maybe add metadata for config gui later y: SignalData z: Optional[SignalData] = None + dap: Optional[str] = None model_config: dict = {"validate_assignment": True} @@ -63,6 +64,7 @@ class CurveConfig(ConnectionConfig): class BECCurve(BECConnector, pg.PlotDataItem): USER_ACCESS = [ "remove", + "dap_params", "rpc_id", "config_dict", "set", @@ -75,6 +77,7 @@ class BECCurve(BECConnector, pg.PlotDataItem): "set_pen_width", "set_pen_style", "get_data", + "dap_params", ] def __init__( @@ -96,6 +99,7 @@ class BECCurve(BECConnector, pg.PlotDataItem): self.parent_item = parent_item self.apply_config() + self.dap_params = None if kwargs: self.set(**kwargs) @@ -119,6 +123,14 @@ class BECCurve(BECConnector, pg.PlotDataItem): self.setSymbolSize(self.config.symbol_size) self.setSymbol(self.config.symbol) + @property + def dap_params(self): + return self._dap_params + + @dap_params.setter + def dap_params(self, value): + self._dap_params = value + def set_data(self, x, y): if self.config.source == "custom": self.setData(x, y) @@ -241,5 +253,6 @@ class BECCurve(BECConnector, pg.PlotDataItem): def remove(self): """Remove the curve from the plot.""" - self.parent_item.removeItem(self) + # self.parent_item.removeItem(self) + self.parent_item.remove_curve(self.name()) self.cleanup() diff --git a/tests/unit_tests/test_waveform1d.py b/tests/unit_tests/test_waveform1d.py index 6a461b16..b0b14580 100644 --- a/tests/unit_tests/test_waveform1d.py +++ b/tests/unit_tests/test_waveform1d.py @@ -85,6 +85,7 @@ def test_create_waveform1D_by_config(bec_figure): "pen_style": "dash", "source": "scan_segment", "signals": { + "dap": None, "source": "scan_segment", "x": { "name": "samx", @@ -248,6 +249,7 @@ def test_change_curve_appearance_methods(bec_figure, qtbot): assert c1.config.pen_style == "dashdot" assert c1.config.source == "scan_segment" assert c1.config.signals.model_dump() == { + "dap": None, "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}, @@ -277,6 +279,7 @@ def test_change_curve_appearance_args(bec_figure): assert c1.config.pen_style == "dashdot" assert c1.config.source == "scan_segment" assert c1.config.signals.model_dump() == { + "dap": None, "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}, @@ -384,6 +387,7 @@ def test_curve_add_by_config(bec_figure): "pen_style": "dash", "source": "scan_segment", "signals": { + "dap": None, "source": "scan_segment", "x": {"name": "samx", "entry": "samx", "unit": None, "modifier": None, "limits": None}, "y": {