feat(waveform1d): dap LMFit model can be added to plot

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

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)
 

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

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: str | None = None,
-        color_map_z: Optional[str] = "plasma",
+        color_map_z: str | None = "plasma",
-        label: Optional[str] = None,
+        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

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

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

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