feat(curve_dialog): add DAP functionality

bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.py

@@ -1,20 +1,20 @@
 from __future__ import annotations
 
 import os
+from typing import Literal
 
+from pydantic import BaseModel
 from PySide6.QtCore import QObject
 from PySide6.QtGui import QIcon
 from PySide6.QtWidgets import QComboBox, QLineEdit, QPushButton, QSpinBox, QTableWidget
-from pydantic import BaseModel
 from qtpy.QtCore import Slot
 from qtpy.QtWidgets import QVBoxLayout
 
 from bec_widgets.qt_utils.settings_dialog import SettingWidget
-from bec_widgets.utils import UILoader, Colors
+from bec_widgets.utils import BECConnector, Colors, UILoader
 from bec_widgets.widgets.color_button.color_button import ColorButton
 from bec_widgets.widgets.device_line_edit.device_line_edit import DeviceLineEdit
 from bec_widgets.widgets.figure.plots.plot_base import AxisConfig
-
 from bec_widgets.widgets.figure.plots.waveform.waveform_curve import CurveConfig
 
 
@@ -29,26 +29,46 @@ class CurveSettings(SettingWidget):
         self.layout.addWidget(self.ui)
 
         self.ui.add_curve.clicked.connect(self.add_curve)
+        self.ui.add_dap.clicked.connect(self.add_dap)
         self.ui.x_mode.currentIndexChanged.connect(self.set_x_mode)
-        self.ui.normalize_colors.clicked.connect(self.change_colormap)
+        self.ui.normalize_colors_scan.clicked.connect(lambda: self.change_colormap("scan"))
+        self.ui.normalize_colors_dap.clicked.connect(lambda: self.change_colormap("dap"))
 
     @Slot(dict)
     def display_current_settings(self, config: dict | BaseModel):
-        curves = config["scan_segment"]
 
-        # set mode of x axis box
+        # What elements should be enabled
         x_name = self.target_widget.waveform._x_axis_mode["name"]
         x_entry = self.target_widget.waveform._x_axis_mode["entry"]
-        self._setup_x_box(x_name, x_entry)
+        self._enable_ui_elements(x_name, x_entry)
         cm = self.target_widget.config.color_palette
-        self.ui.color_map_selector.combo.setCurrentText(cm)
+        self.ui.color_map_selector_scan.combo.setCurrentText(cm)
 
-        for label, curve in curves.items():
+        # Scan Curve Table
+        for label, curve in config["scan_segment"].items():
             row_count = self.ui.scan_table.rowCount()
             self.ui.scan_table.insertRow(row_count)
-            ScanRow(table_widget=self.ui.scan_table, row=row_count, config=curve.config)
+            DialogRow(
+                parent=self,
+                table_widget=self.ui.scan_table,
+                client=self.target_widget.client,
+                row=row_count,
+                config=curve.config,
+            ).add_scan_row()
 
-    def _setup_x_box(self, name, entry):
+        # Add DAP Curves
+        for label, curve in config["DAP"].items():
+            row_count = self.ui.dap_table.rowCount()
+            self.ui.dap_table.insertRow(row_count)
+            DialogRow(
+                parent=self,
+                table_widget=self.ui.dap_table,
+                client=self.target_widget.client,
+                row=row_count,
+                config=curve.config,
+            ).add_dap_row()
+
+    def _enable_ui_elements(self, name, entry):
         if name in ["index", "timestamp", "best_effort"]:
             self.ui.x_mode.setCurrentText(name)
             self.set_x_mode()
@@ -64,26 +84,38 @@ class CurveSettings(SettingWidget):
         if x_mode in ["index", "timestamp", "best_effort"]:
             self.ui.x_name.setEnabled(False)
             self.ui.x_entry.setEnabled(False)
+            self.ui.dap_table.setEnabled(False)
+            self.ui.add_dap.setEnabled(False)
         else:
             self.ui.x_name.setEnabled(True)
             self.ui.x_entry.setEnabled(True)
+            self.ui.dap_table.setEnabled(True)
+            self.ui.add_dap.setEnabled(True)
 
     @Slot()
-    def change_colormap(self):
+    def change_colormap(self, target: Literal["scan", "dap"]):
-        cm = self.ui.color_map_selector.combo.currentText()
+        if target == "scan":
-        rows = self.ui.scan_table.rowCount()
+            cm = self.ui.color_map_selector_scan.combo.currentText()
+            table = self.ui.scan_table
+        if target == "dap":
+            cm = self.ui.color_map_selector_dap.combo.currentText()
+            table = self.ui.dap_table
+        rows = table.rowCount()
         colors = Colors.golden_angle_color(colormap=cm, num=rows + 1, format="HEX")
+        color_button_col = 2 if target == "scan" else 3
         for row, color in zip(range(rows), colors):
-            self.ui.scan_table.cellWidget(row, 2).setColor(color)
+            table.cellWidget(row, color_button_col).setColor(color)
-        self.target_widget.set_colormap(cm)
 
     @Slot()
     def accept_changes(self):
-        self.accept_scan_curve_changes()
+        self.accept_curve_changes()
 
-    def accept_scan_curve_changes(self):
+    def accept_curve_changes(self):
-        old_curves = list(self.target_widget.waveform._curves_data["scan_segment"].values())
+        old_curves_scans = list(self.target_widget.waveform._curves_data["scan_segment"].values())
-        for curve in old_curves:
+        old_curves_dap = list(self.target_widget.waveform._curves_data["DAP"].values())
+        for curve in old_curves_scans:
+            curve.remove()
+        for curve in old_curves_dap:
             curve.remove()
         self.get_curve_params()
 
@@ -114,38 +146,85 @@ class CurveSettings(SettingWidget):
                 pen_width=width,
                 symbol_size=symbol_size,
             )
+
+        if x_mode not in ["index", "timestamp", "best_effort"]:
+            for row in range(self.ui.dap_table.rowCount()):
+                y_name = self.ui.dap_table.cellWidget(row, 0).text()
+                y_entry = self.ui.dap_table.cellWidget(row, 1).text()
+                dap = self.ui.dap_table.cellWidget(row, 2).currentText()
+                color = self.ui.dap_table.cellWidget(row, 3).get_color()
+                style = self.ui.dap_table.cellWidget(row, 4).currentText()
+                width = self.ui.dap_table.cellWidget(row, 5).value()
+                symbol_size = self.ui.dap_table.cellWidget(row, 6).value()
+
+                self.target_widget.add_dap(
+                    x_name=x_name,
+                    x_entry=x_entry,
+                    y_name=y_name,
+                    y_entry=y_entry,
+                    dap=dap,
+                    color=color,
+                    pen_style=style,
+                    pen_width=width,
+                    symbol_size=symbol_size,
+                )
         self.target_widget.scan_history(-1)
 
     def add_curve(self):
         row_count = self.ui.scan_table.rowCount()
         self.ui.scan_table.insertRow(row_count)
-        ScanRow(table_widget=self.ui.scan_table, row=row_count, config=None)
+        DialogRow(
+            parent=self,
+            table_widget=self.ui.scan_table,
+            client=self.target_widget.client,
+            row=row_count,
+            config=None,
+        ).add_scan_row()
+
+    def add_dap(self):
+        row_count = self.ui.dap_table.rowCount()
+        self.ui.dap_table.insertRow(row_count)
+        DialogRow(
+            parent=self,
+            table_widget=self.ui.dap_table,
+            client=self.target_widget.client,
+            row=row_count,
+            config=None,
+        ).add_dap_row()
 
 
-class ScanRow(QObject):
+class DialogRow(QObject):
     def __init__(
         self,
         parent=None,
         table_widget: QTableWidget = None,
-        row=None,
+        row: int = None,
-        config: dict | CurveConfig = None,
+        config: dict = None,
+        client=None,
     ):
-        super().__init__(parent=parent)
 
-        current_path = os.path.dirname(__file__)
+        super().__init__(parent=parent)
+        self.client = client
+
+        self.table_widget = table_widget
+        self.row = row
+        self.config = config
+        self.init_default_widgets()
+
+    def init_default_widgets(self):
+
         # Remove Button
-        icon_path = os.path.join(current_path, "remove.svg")
+        self.remove_button = RemoveButton()
-        self.remove_button = QPushButton()
-        self.remove_button.setIcon(QIcon(icon_path))
 
         # Name and Entry
         self.device_line_edit = DeviceLineEdit()
         self.entry_line_edit = QLineEdit()
 
+        self.dap_combo = QComboBox()
+        self.populate_dap_combobox()
+
         # Styling
-        default_color = Colors.golden_angle_color(colormap="magma", num=row + 1, format="HEX")[-1]
         self.color_button = ColorButton()
-        self.color_button.setColor(default_color)
         self.style_combo = StyleComboBox()
         self.width = QSpinBox()
         self.width.setMinimum(1)
@@ -157,27 +236,34 @@ class ScanRow(QObject):
         self.symbol_size.setMaximum(20)
         self.symbol_size.setValue(5)
 
-        self.table_widget = table_widget
-        self.row = row
-
         self.remove_button.clicked.connect(
             lambda: self.remove_row()
         )  # From some reason do not work without lambda
 
-        if config is not None:
+    def populate_dap_combobox(self):
-            self.fill_row_from_config(config)
+        available_models = [
+            attr
+            for attr in dir(self.client.dap)
+            if not attr.startswith("__")
+            and not callable(getattr(self.client.dap, attr))
+            and not attr.startswith("_")
+        ]
+        self.dap_combo.addItems(available_models)
 
-        self.add_row_to_table()
+    def add_scan_row(self):
+        if self.config is not None:
+            self.device_line_edit.setText(self.config.signals.y.name)
+            self.entry_line_edit.setText(self.config.signals.y.entry)
+            self.color_button.setColor(self.config.color)
+            self.style_combo.setCurrentText(self.config.pen_style)
+            self.width.setValue(self.config.pen_width)
+            self.symbol_size.setValue(self.config.symbol_size)
+        else:
+            default_color = Colors.golden_angle_color(
+                colormap="magma", num=self.row + 1, format="HEX"
+            )[-1]
+            self.color_button.setColor(default_color)
 
-    def fill_row_from_config(self, config):
-        self.device_line_edit.setText(config.signals.y.name)
-        self.entry_line_edit.setText(config.signals.y.entry)
-        self.color_button.setColor(config.color)
-        self.style_combo.setCurrentText(config.pen_style)
-        self.width.setValue(config.pen_width)
-        self.symbol_size.setValue(config.symbol_size)
-
-    def add_row_to_table(self):
         self.table_widget.setCellWidget(self.row, 0, self.device_line_edit)
         self.table_widget.setCellWidget(self.row, 1, self.entry_line_edit)
         self.table_widget.setCellWidget(self.row, 2, self.color_button)
@@ -186,6 +272,30 @@ class ScanRow(QObject):
         self.table_widget.setCellWidget(self.row, 5, self.symbol_size)
         self.table_widget.setCellWidget(self.row, 6, self.remove_button)
 
+    def add_dap_row(self):
+        if self.config is not None:
+            self.device_line_edit.setText(self.config.signals.y.name)
+            self.entry_line_edit.setText(self.config.signals.y.entry)
+            self.dap_combo.setCurrentText(self.config.signals.dap)
+            self.color_button.setColor(self.config.color)
+            self.style_combo.setCurrentText(self.config.pen_style)
+            self.width.setValue(self.config.pen_width)
+            self.symbol_size.setValue(self.config.symbol_size)
+        else:
+            default_color = Colors.golden_angle_color(
+                colormap="magma", num=self.row + 1, format="HEX"
+            )[-1]
+            self.color_button.setColor(default_color)
+
+        self.table_widget.setCellWidget(self.row, 0, self.device_line_edit)
+        self.table_widget.setCellWidget(self.row, 1, self.entry_line_edit)
+        self.table_widget.setCellWidget(self.row, 2, self.dap_combo)
+        self.table_widget.setCellWidget(self.row, 3, self.color_button)
+        self.table_widget.setCellWidget(self.row, 4, self.style_combo)
+        self.table_widget.setCellWidget(self.row, 5, self.width)
+        self.table_widget.setCellWidget(self.row, 6, self.symbol_size)
+        self.table_widget.setCellWidget(self.row, 7, self.remove_button)
+
     @Slot()
     def remove_row(self):
         row = self.table_widget.indexAt(self.remove_button.pos()).row()
@@ -200,3 +310,11 @@ class StyleComboBox(QComboBox):
     def __init__(self, parent=None):
         super().__init__(parent)
         self.addItems(["solid", "dash", "dot", "dashdot"])
+
+
+class RemoveButton(QPushButton):
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        current_path = os.path.dirname(__file__)
+        icon_path = os.path.join(current_path, "remove.svg")
+        self.setIcon(QIcon(icon_path))

bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.ui

@@ -31,7 +31,7 @@
      <property name="title">
       <string>X Axis</string>
      </property>
-     <layout class="QHBoxLayout" name="horizontalLayout" stretch="0,0,0,0,1,3,0,1,3">
+     <layout class="QHBoxLayout" name="horizontalLayout" stretch="0,0,0,0,0,1,3,0,1,3">
       <item>
        <widget class="QLabel" name="x_mode_label">
         <property name="text">
@@ -41,6 +41,12 @@
       </item>
       <item>
        <widget class="QComboBox" name="x_mode">
+        <property name="minimumSize">
+         <size>
+          <width>150</width>
+          <height>26</height>
+         </size>
+        </property>
         <item>
          <property name="text">
           <string>best_effort</string>
@@ -92,6 +98,19 @@
         </property>
        </widget>
       </item>
+      <item>
+       <spacer name="horizontalSpacer_4">
+        <property name="orientation">
+         <enum>Qt::Orientation::Horizontal</enum>
+        </property>
+        <property name="sizeHint" stdset="0">
+         <size>
+          <width>40</width>
+          <height>20</height>
+         </size>
+        </property>
+       </spacer>
+      </item>
       <item>
        <widget class="QLabel" name="x_name_label">
         <property name="text">
@@ -158,7 +177,7 @@
            </widget>
           </item>
           <item row="0" column="3">
-           <widget class="ColormapSelector" name="color_map_selector">
+           <widget class="ColormapSelector" name="color_map_selector_scan">
             <property name="minimumSize">
              <size>
               <width>150</width>
@@ -243,7 +262,7 @@
            </spacer>
           </item>
           <item row="0" column="2">
-           <widget class="QPushButton" name="normalize_colors">
+           <widget class="QPushButton" name="normalize_colors_scan">
             <property name="text">
              <string>Normalize Colors</string>
             </property>
@@ -268,27 +287,7 @@
           <property name="bottomMargin">
            <number>5</number>
           </property>
-          <item row="0" column="0">
+          <item row="1" column="0" colspan="4">
-           <widget class="QPushButton" name="add_dap">
-            <property name="text">
-             <string>Add DAP</string>
-            </property>
-           </widget>
-          </item>
-          <item row="0" column="1">
-           <spacer name="horizontalSpacer_2">
-            <property name="orientation">
-             <enum>Qt::Orientation::Horizontal</enum>
-            </property>
-            <property name="sizeHint" stdset="0">
-             <size>
-              <width>585</width>
-              <height>20</height>
-             </size>
-            </property>
-           </spacer>
-          </item>
-          <item row="1" column="0" colspan="2">
            <widget class="QTableWidget" name="dap_table">
             <attribute name="horizontalHeaderShowSortIndicator" stdset="0">
              <bool>true</bool>
@@ -338,6 +337,51 @@
             </column>
            </widget>
           </item>
+          <item row="0" column="2">
+           <widget class="QPushButton" name="normalize_colors_dap">
+            <property name="text">
+             <string>Normalize Colors</string>
+            </property>
+           </widget>
+          </item>
+          <item row="0" column="1">
+           <spacer name="horizontalSpacer_2">
+            <property name="orientation">
+             <enum>Qt::Orientation::Horizontal</enum>
+            </property>
+            <property name="sizeHint" stdset="0">
+             <size>
+              <width>585</width>
+              <height>20</height>
+             </size>
+            </property>
+           </spacer>
+          </item>
+          <item row="0" column="0">
+           <widget class="QPushButton" name="add_dap">
+            <property name="text">
+             <string>Add DAP</string>
+            </property>
+           </widget>
+          </item>
+          <item row="0" column="3">
+           <widget class="ColormapSelector" name="color_map_selector_dap">
+            <property name="minimumSize">
+             <size>
+              <width>150</width>
+              <height>0</height>
+             </size>
+            </property>
+            <property name="colormaps" stdset="0">
+             <stringlist>
+              <string>inferno</string>
+              <string>viridis</string>
+              <string>plasma</string>
+              <string>magma</string>
+             </stringlist>
+            </property>
+           </widget>
+          </item>
          </layout>
         </widget>
        </widget>

bec_widgets/widgets/waveform/waveform_widget.py

@@ -91,7 +91,7 @@ class BECWaveformWidget(BECConnector, QWidget):
         self._hook_actions()
 
         # TEst actions
-        self.plot(x_name="samx", y_name="bpm4i")
+        self.plot(x_name="samx", y_name="bpm4i", dap="GaussianModel")
         self.plot(x_name="samx", y_name="bpm3a")
         self.plot(x_name="samx", y_name="bpm6i")
 
@@ -261,6 +261,7 @@ class BECWaveformWidget(BECConnector, QWidget):
         y_entry: str | None = None,
         color: str | None = None,
         dap: str = "GaussianModel",
+        validate_bec: bool = True,
         **kwargs,
     ) -> BECCurve:
         """
@@ -272,9 +273,8 @@ class BECWaveformWidget(BECConnector, QWidget):
             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.
+            validate_bec(bool, optional): If True, validate the signal with BEC. Defaults to True.
             **kwargs: Additional keyword arguments for the curve configuration.
 
         Returns:
@@ -287,6 +287,7 @@ class BECWaveformWidget(BECConnector, QWidget):
             y_entry=y_entry,
             color=color,
             dap=dap,
+            validate_bec=validate_bec,
             **kwargs,
         )