feat(auto_gain): Build own gui

debye_bec/bec_ipython_client/plugins/auto_gain_core.py

@@ -26,7 +26,6 @@ TIMEOUT_PUT_PV = 5  # Timeout to set a PV
 MONO_VELOCITY = 20  # Move velocity in deg/s
 TIMEOUT_MONO_MOVE = 30  # Timeout to finish a movement on the mono
 
-LINKED_CHANNELS = {"ic0", "ic1"}
 AVAILABLE_GAINS = [1e6, 1e7, 5e7, 1e8, 1e9]  # ascending order
 MIN_SIGNAL = 0.05  # Minimum signal to count as valid signal
 FULL_SCALE_V = 10.0  # NIDAQ AI full-scale range
@@ -53,7 +52,7 @@ class AutoGainError(Exception):
     """AutoGain specific error"""
 
 
-class AutoGain:
+class AutoGainCore:
     """Module to automatically set the gains for the selected amplifiers"""
 
     def __init__(self):
@@ -65,14 +64,27 @@ class AutoGain:
             raise AutoGainError("Did not get bec")
         self.dev = dev
         self.bec = bec
+        self.linked_channels = None
+        self.prohibited_gains = {}
+
+    def available_gains(self, amplifier):
+        """Returns a list of gains which are allowed for the amplifier"""
+        if amplifier in self.prohibited_gains:
+            gains = AVAILABLE_GAINS
+            for prohibited_gain in self.prohibited_gains[amplifier]:
+                gains.remove(prohibited_gain)
+            return gains
+        return AVAILABLE_GAINS
 
     def start(
         self,
         element: str,
         edge: str,
         amplifier: list[str] | None = None,
+        link_ic0_ic1: bool = True,
         comp_ring_current: bool = True,
-    ) -> None:
+        prohibited_gains: dict | None = None,
+    ) -> dict:
         """Start the auto-gain sequence. Measure the signals of the specified
         amplifiers and set the gains accordingly. Makes sure there is actually beam available.
 
@@ -85,8 +97,11 @@ class AutoGain:
             edge(str): Corresponding edge, e.g. 'L1'
             amplifier(list[str]): Amplifiers where auto-gain should be applied to
                                   Defaults to all amplifiers -> ['ic0', 'ic1', 'ic2', 'pips']
+            link_ic0_ic1(bool): Apply same gain to ic0 and ic1. Defaults to True
             comp_ring_current(bool): Respects the current ring current and calculates the gain(s)
                                      for a nominal ring current of 400 mA. Defaults to True
+            prohibited_gains(dict): Gains which are not allowed to be used.
+                                    E.g.: {'ic0': [1e6]}. Defaults to None
 
         Raises:
             If NIDAQ is not in measurement mode
@@ -101,6 +116,16 @@ class AutoGain:
             If the ring current drops to 0 mA during the measurement (beamdump)
         """
 
+        if link_ic0_ic1:
+            self.linked_channels = {"ic0", "ic1"}
+        else:
+            self.linked_channels = {}
+
+        if prohibited_gains is None:
+            self.prohibited_gains = {}
+        else:
+            self.prohibited_gains = prohibited_gains
+
         if amplifier is None:
             amplifier = ["ic0", "ic1", "ic2", "pips"]
 
@@ -180,6 +205,8 @@ class AutoGain:
 
         logger.info("All checks done, start preparing for measurement")
 
+        results = {}
+
         try:
             logger.info(f"Move mono to start of {emin} eV")
             self.dev.mo1_bragg.velocity.put(MONO_VELOCITY)
@@ -193,7 +220,7 @@ class AutoGain:
 
             # Set gains to lowest gain
             for name, ch in active_channels.items():
-                lowest_gain = AVAILABLE_GAINS[0]
+                lowest_gain = self.available_gains(name)[0]
                 ch["self.dev"].set_gain(lowest_gain)
                 ch["gain"] = lowest_gain
 
@@ -243,8 +270,8 @@ class AutoGain:
 
                     # ic0/ic1 must always end up with the same gain -> treat whichever of
                     # them are present (one or both) as a single group for this decision
-                    if name in LINKED_CHANNELS:
-                        group_names = [n for n in LINKED_CHANNELS if n in active_channels]
+                    if name in self.linked_channels:
+                        group_names = [n for n in self.linked_channels if n in active_channels]
                     else:
                         group_names = [name]
 
@@ -268,14 +295,14 @@ class AutoGain:
 
                     if raw_signal_min < MIN_SIGNAL:
                         logger.info(f"Raw signal for {label} is below {MIN_SIGNAL}")
-                        if gain_ref == AVAILABLE_GAINS[-1]:
+                        if gain_ref == self.available_gains(name)[-1]:
                             logger.warning(
                                 f"Amplifier of {label} at highest gain {gain_ref} and still not"
                                 + f" measured signal above {MIN_SIGNAL}"
                             )
                         else:
-                            next_gain = AVAILABLE_GAINS[
-                                bisect_right(AVAILABLE_GAINS, gain_ref * 100) - 1
+                            next_gain = self.available_gains(name)[
+                                bisect_right(self.available_gains(name), gain_ref * 100) - 1
                             ]
                             for gname in group_names:
                                 active_channels[gname]["self.dev"].set_gain(next_gain)
@@ -286,15 +313,19 @@ class AutoGain:
                         gain = max(
                             (
                                 g
-                                for g in AVAILABLE_GAINS
+                                for g in self.available_gains(name)
                                 if raw_signal_max / gain_ref * g <= FULL_SCALE_V * SAFETY_MARGIN
                             ),
-                            default=min(AVAILABLE_GAINS),
+                            default=min(self.available_gains(name)),
                         )
                         for gname in group_names:
                             active_channels[gname]["self.dev"].set_gain(gain)
                             active_channels[gname]["gain"] = gain
                         logger.info(f"Calculated final gain for {label} of {gain:.0e}")
+                        results[label] = {
+                            "gain": gain,
+                            "max_signal": raw_signal_max / gain_ref * gain,
+                        }
 
                     processed.update(group_names)
 
@@ -316,6 +347,8 @@ class AutoGain:
         self.dev.mo1_bragg.velocity.put(init_vel)
         wait_for_signal(self.dev.mo1_bragg.velocity, init_vel, timeout=TIMEOUT_PUT_PV)
 
+        return results
+
     def _get_ring_current(self) -> float:
         ring_current = 0
         retries = 0

debye_bec/bec_ipython_client/startup/post_startup.py

@@ -41,11 +41,11 @@ logger = bec_logger.logger
 
 logger.info("Using the Debye startup script.")
 
+from debye_bec.bec_ipython_client.plugins.auto_gain_core import AutoGainCore
 from debye_bec.bec_ipython_client.plugins.digital_twin_core.digital_twin_core import DigitalTwinCore
-from debye_bec.bec_ipython_client.plugins.auto_gain import AutoGain
 
 digital_twin = DigitalTwinCore()
 logger.success("Digital Twin Core loaded. Use 'digital_twin' to access it.")
 
-auto_gain = AutoGain()
+auto_gain = AutoGainCore()
 logger.success("Auto-Gain module loaded. Use 'auto_gain' to access it.")

debye_bec/bec_widgets/widgets/auto_gain/auto_gain.py

@@ -0,0 +1,194 @@
+"""
+Auto Gain: Custom BEC widget to automatically set gains.
+"""
+
+import sys
+
+from bec_lib import bec_logger
+from bec_lib.logger import bec_logger
+from bec_widgets.utils.bec_dispatcher import BECDispatcher
+from bec_widgets.utils.bec_widget import BECWidget
+from bec_widgets.utils.colors import apply_theme, get_accent_colors
+from bec_widgets.utils.error_popups import SafeSlot
+from qtpy.QtCore import QObject, QThread, Signal
+
+# pylint: disable=E0611
+from qtpy.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget
+
+from ....bec_ipython_client.plugins.auto_gain_core import AutoGainCore
+from ..edge_selector import EdgeSelector
+from .qt_widgets import CheckBox, Group, MatrixSelector, MyButton, TextIndicator
+
+logger = bec_logger.logger
+
+
+class Worker(QObject):
+    progress = Signal(str)
+    completed = Signal(int)
+
+    @SafeSlot(int)
+    def start(self, config: dict):
+        core = AutoGainCore()
+        self.progress.emit("AutoGain started")
+        results = core.start(
+            config["element"],
+            config["edge"],
+            amplifier=config["all_amplifiers"],
+            link_ic0_ic1=config["link_ic0_ic1"],
+            comp_ring_current=config["compensate_ring_current"],
+            prohibited_gains=config["prohibited_gains"],
+        )
+        results_string = ""
+        init = True
+        for amplifier, result in results.items():
+            if not init:
+                results_string += "\n"
+            else:
+                init = False
+            results_string += f"{amplifier} set to gain {result['gain']:.0e}, max signal {result['max_signal']:.3f} V"
+        self.progress.emit(results_string)
+        self.completed.emit(True)
+
+
+class AutoGain(BECWidget, QWidget):
+    """
+    Main widget of auto gain widget.
+    """
+
+    PLUGIN = True
+    ICON_NAME = "tune"
+
+    def __init__(self, show_control: bool = True, *arg, parent=None, **kwargs):
+        super().__init__(parent=parent, *arg, **kwargs)
+        self.get_bec_shortcuts()
+
+        central = QWidget()
+        self.root_layout = QVBoxLayout(central)
+
+        self.edge_selector_button = MyButton("Choose", "default")
+        self.edge_label = QLabel("No edge selected")
+        self.edge_group = Group(
+            "Edge selection", [self.edge_selector_button, self.edge_label], orientation="horizontal"
+        )
+        self.edge_element = None
+        self.edge_edge = None
+        self.edge_energy = None
+
+        self.root_layout.addWidget(self.edge_group)
+
+        self.ic0_check = CheckBox("ic0", "IC0")
+        self.ic1_check = CheckBox("ic1", "IC1")
+        self.ic2_check = CheckBox("ic2", "IC2")
+        self.pips_check = CheckBox("pips", "PIPS")
+        self.apply_group = Group(
+            "Apply auto gain on", [self.ic0_check, self.ic1_check, self.ic2_check, self.pips_check]
+        )
+        self.root_layout.addWidget(self.apply_group)
+
+        self.link_ic0_ic1 = CheckBox("link_ic0_ic1", "Link IC0/IC1 gains")
+        self.link_ic0_ic1.setChecked(True)
+
+        self.comp_ring_current = CheckBox("comp_ring_current", "Calculate for 400 mA ring current")
+        self.comp_ring_current.setChecked(True)
+
+        self.prohibited_gains = MatrixSelector("Prohibited gains")
+
+        self.expert_group = Group(
+            "Expert settings", [self.link_ic0_ic1, self.comp_ring_current, self.prohibited_gains]
+        )
+
+        self.root_layout.addWidget(self.expert_group)
+
+        self.start_button = MyButton("Start", "success")
+        self.status = TextIndicator(label="Status")
+        self.status.setText("Standby")
+        self.status.setColor(get_accent_colors().default.name())
+
+        self.control_group = Group("Control", [self.start_button, self.status])
+
+        if show_control:
+            self.root_layout.addWidget(self.control_group)
+
+        self.root_layout.addStretch()
+
+        self.setLayout(self.root_layout)
+        self.setWindowTitle("Auto Gain")
+
+        self.edge_selector_button.clicked.connect(self._update_edge)
+        self.start_button.clicked.connect(self._start)
+
+        self.worker = Worker()
+        self.worker_thread = QThread()
+        self.worker.progress.connect(self._update_status)
+        self.worker.completed.connect(self._complete)
+        self.worker.moveToThread(self.worker_thread)
+
+    def _update_edge(self, *_):
+        dlg = EdgeSelector(self)
+        if dlg.exec_():
+            self.edge_element = dlg.selected_element
+            self.edge_edge = dlg.selected_edge
+            self.edge_energy = dlg.selected_energy
+            self.edge_label.setText(
+                f"{dlg.selected_element}, {dlg.selected_edge}-edge, {dlg.selected_energy:0.1f} eV"
+            )
+
+    def _start(self):
+        if not self.worker_thread.isRunning():
+            self.start_button.setEnabled(False)
+            self.status.setColor(get_accent_colors().warning.name())
+            self.worker_thread.start()
+
+    def _update_status(self, text):
+        self.status.setText(text)
+
+    def _complete(self, *_):
+        self.start_button.setEnabled(True)
+        self.status.setColor(get_accent_colors().success.name())
+
+    def load_config(self, config: dict):
+        self.edge_element = config["element"]
+        self.edge_edge = config["edge"]
+        self.edge_energy = config["energy"]
+        self.ic0_check.setChecked(config["ic0"])
+        self.ic1_check.setChecked(config["ic1"])
+        self.ic2_check.setChecked(config["ic2"])
+        self.pips_check.setChecked(config["pips"])
+        self.link_ic0_ic1.setChecked(config["link_ic0_ic1"])
+        self.comp_ring_current.setChecked(config["compensate_ring_current"])
+        self.prohibited_gains.set_selection(config["prohibited_gains"])
+
+    def get_config(self) -> dict:
+        if self.edge_element is None or self.edge_edge is None:
+            raise ValueError("Choose an absorption edge first.")
+        amplifier_checks = [
+            (self.ic0_check, "ic0"),
+            (self.ic1_check, "ic1"),
+            (self.ic2_check, "ic2"),
+            (self.pips_check, "pips"),
+        ]
+        amplifiers = [name for check, name in amplifier_checks if check.isChecked()]
+        if not amplifiers:
+            raise ValueError("Select at least one amplifier for auto gain.")
+        return {
+            "element": self.edge_element,
+            "edge": self.edge_edge,
+            "energy": self.edge_energy,
+            "ic0": self.ic0_check.isChecked,
+            "ic1": self.ic1_check.isChecked,
+            "ic2": self.ic2_check.isChecked,
+            "pips": self.pips_check.isChecked,
+            "all_amplifiers": amplifiers,
+            "link_ic0_ic1": self.link_ic0_ic1.isChecked,
+            "compensate_ring_current": self.comp_ring_current.isChecked,
+            "prohibited_gains": self.prohibited_gains.selection(),
+        }
+
+
+if __name__ == "__main__":
+    app = QApplication(sys.argv)
+    apply_theme("light")
+    dispatcher = BECDispatcher(gui_id="auto_gain")
+    win = AutoGain()
+    win.show()
+    sys.exit(app.exec_())

debye_bec/bec_widgets/widgets/auto_gain/auto_gain.pyproject

@@ -0,0 +1 @@
+{'files': ['auto_gain.py']}

debye_bec/bec_widgets/widgets/auto_gain/auto_gain_plugin.py

@@ -0,0 +1,57 @@
+# Copyright (C) 2022 The Qt Company Ltd.
+# SPDX-License-Identifier: LicenseRef-Qt-Commercial OR BSD-3-Clause
+
+from bec_widgets.utils.bec_designer import designer_material_icon
+from qtpy.QtDesigner import QDesignerCustomWidgetInterface
+from qtpy.QtWidgets import QWidget
+
+from .auto_gain import AutoGain
+
+DOM_XML = """
+<ui language='c++'>
+    <widget class='AutoGain' name='auto_gain'>
+    </widget>
+</ui>
+"""
+
+
+class AutoGainPlugin(QDesignerCustomWidgetInterface):  # pragma: no cover
+    def __init__(self):
+        super().__init__()
+        self._form_editor = None
+
+    def createWidget(self, parent):
+        if parent is None:
+            return QWidget()
+        t = AutoGain(parent)
+        return t
+
+    def domXml(self):
+        return DOM_XML
+
+    def group(self):
+        return ""
+
+    def icon(self):
+        return designer_material_icon(AutoGain.ICON_NAME)
+
+    def includeFile(self):
+        return "auto_gain"
+
+    def initialize(self, form_editor):
+        self._form_editor = form_editor
+
+    def isContainer(self):
+        return False
+
+    def isInitialized(self):
+        return self._form_editor is not None
+
+    def name(self):
+        return "AutoGain"
+
+    def toolTip(self):
+        return "AutoGain"
+
+    def whatsThis(self):
+        return self.toolTip()

debye_bec/bec_widgets/widgets/auto_gain/qt_widgets.py

@@ -0,0 +1,183 @@
+from bec_widgets.utils.colors import get_accent_colors
+from qtpy.QtCore import Qt
+
+# pylint: disable=E0611
+from qtpy.QtWidgets import (
+    QCheckBox,
+    QGridLayout,
+    QGroupBox,
+    QHBoxLayout,
+    QLabel,
+    QLayout,
+    QPushButton,
+    QVBoxLayout,
+    QWidget,
+)
+
+LABEL_WIDTH = 118
+ROW_MARGINS = (4, 0, 4, 0)
+ROW_SPACING = 6
+
+
+class Group(QGroupBox):
+    def __init__(self, label, objs, orientation="vertical"):
+        super().__init__(label)
+        if orientation == "vertical":
+            self._layout = QVBoxLayout(self)
+        elif orientation == "horizontal":
+            self._layout = QHBoxLayout(self)
+        else:
+            raise ValueError(f"Orientation {orientation} is not supported!")
+        for obj in objs:
+            if isinstance(obj, QWidget):
+                self._layout.addWidget(obj)
+            elif isinstance(obj, QLayout):
+                self._layout.addLayout(obj)
+
+
+class CheckBox(QWidget):
+    def __init__(self, identifier="", label=""):
+        super().__init__()
+        layout = QHBoxLayout(self)
+        layout.setContentsMargins(*ROW_MARGINS)
+        layout.setSpacing(ROW_SPACING)
+        self.identifier = identifier
+        self.label = QLabel(label)
+        self.label.setFixedWidth(LABEL_WIDTH)
+        self.label.setWordWrap(True)
+        layout.addWidget(self.label)
+        self.checkbox = QCheckBox()
+        layout.addWidget(self.checkbox)
+
+    def setChecked(self, set_checked: bool = True):
+        self.checkbox.setChecked(set_checked)
+
+    def isChecked(self):
+        return self.checkbox.isChecked()
+
+
+class MatrixSelector(QWidget):
+    """Checkbox matrix for selecting amplifier/value combinations.
+
+    Displays a grid with one row per amplifier (IC0, IC1, IC2, PIPS) and one
+    column per value (1e6, 1e7, 5e7, 1e8, 1e9). Each cell is a checkbox, so any
+    combination of rows and columns can be selected independently.
+
+    Use `selection()` to read the current state as a dict mapping the
+    lowercase row name to a list of the selected column values (as floats),
+    e.g. ``{"ic0": [1e6], "ic1": [1e6, 5e7]}``. Rows with no checked boxes are
+    omitted.
+
+    Attributes:
+        ROWS (list[str]): Row titles, in display order.
+        COLS (list[str]): Column titles, in display order. These are also
+            parsed with ``float()`` to produce the values returned by
+            `selection()`.
+        checks (dict[tuple[str, str], QCheckBox]): Checkboxes keyed by
+            ``(row, col)`` using the original title strings.
+
+    Example:
+        >>> selector = MatrixSelector()
+        >>> my_layout.addWidget(selector)
+        >>> selector.selection()
+        {'ic0': [1000000.0], 'ic1': [1000000.0, 50000000.0]}
+    """
+
+    ROWS = ["IC0", "IC1", "IC2", "PIPS"]
+    COLS = ["1e6", "1e7", "5e7", "1e8", "1e9"]
+
+    def __init__(self, label: str = "", parent=None):
+        super().__init__(parent)
+        layout = QHBoxLayout(self)
+        layout.setContentsMargins(*ROW_MARGINS)
+        layout.setSpacing(ROW_SPACING)
+        self.label = QLabel(label)
+        self.label.setFixedWidth(LABEL_WIDTH)
+        self.label.setWordWrap(True)
+        layout.addWidget(self.label)
+
+        grid_layout = QGridLayout()
+
+        # Column headers
+        for c, col in enumerate(self.COLS, start=1):
+            grid_layout.addWidget(QLabel(col), 0, c, alignment=Qt.AlignmentFlag.AlignCenter)
+
+        # Row headers + checkboxes
+        self.checks = {}
+        for r, row in enumerate(self.ROWS, start=1):
+            grid_layout.addWidget(QLabel(row), r, 0)
+            for c, col in enumerate(self.COLS, start=1):
+                cb = QCheckBox()
+                grid_layout.addWidget(cb, r, c, alignment=Qt.AlignmentFlag.AlignCenter)
+                self.checks[(row, col)] = cb
+
+        layout.addLayout(grid_layout)
+        layout.addStretch()
+
+    def selection(self):
+        """Return e.g. {"ic0": [1e6], "ic1": [1e6, 5e7]}; rows with no selection are omitted."""
+        result = {}
+        for row in self.ROWS:
+            values = [float(col) for col in self.COLS if self.checks[(row, col)].isChecked()]
+            if values:
+                result[row.lower()] = values
+        return result
+
+    def set_selection(self, selection):
+        """Set the checkboxes from a dict as returned by `selection()`.
+
+        Rows not present in `selection` are cleared. Values that don't match
+        any column are ignored.
+        """
+        # Normalize keys to lowercase and values to floats for comparison
+        wanted = {key.lower(): {float(v) for v in values} for key, values in selection.items()}
+        for row in self.ROWS:
+            row_values = wanted.get(row.lower(), set())
+            for col in self.COLS:
+                cb = self.checks[(row, col)]
+                cb.setChecked(float(col) in row_values)
+
+
+class MyButton(QPushButton):
+    def __init__(self, text="", color="default", parent=None):
+        self.color = color
+        super().__init__(text, parent)
+        self.apply_theme()
+
+    def apply_theme(self):
+        if self.isEnabled():
+            colors = get_accent_colors()
+            color = getattr(colors, self.color).name()
+            self.setStyleSheet(f"QPushButton {{ background-color: {color}; color: white; }}")
+        else:
+            self.setStyleSheet(
+                "QPushButton {{background-color: rgb(120, 120, 120); color: white;}}"
+            )
+
+    def setEnabled(self, enable: bool = True):
+        super().setEnabled(enable)
+        self.apply_theme()
+
+
+class TextIndicator(QWidget):
+    def __init__(self, label):
+        super().__init__()
+        layout = QHBoxLayout(self)
+        layout.setContentsMargins(*ROW_MARGINS)
+        layout.setSpacing(ROW_SPACING)
+        self.label = QLabel(label)
+        self.label.setFixedWidth(LABEL_WIDTH)
+        self.label.setWordWrap(True)
+        layout.addWidget(self.label)
+        self.text = QLabel("-")
+        self.text.setAlignment(Qt.AlignTop)  # type: ignore
+        layout.addWidget(self.text)
+
+    def setLabel(self, label) -> None:
+        self.label.setText(label)
+
+    def setText(self, text):
+        self.text.setText(text)
+
+    def setColor(self, color: str):
+        self.text.setStyleSheet(f"QLabel {{color:{color}}}")

debye_bec/bec_widgets/widgets/auto_gain/register_auto_gain.py

@@ -0,0 +1,15 @@
+def main():  # pragma: no cover
+    from qtpy import PYSIDE6
+
+    if not PYSIDE6:
+        print("PYSIDE6 is not available in the environment. Cannot patch designer.")
+        return
+    from PySide6.QtDesigner import QPyDesignerCustomWidgetCollection
+
+    from .auto_gain_plugin import AutoGainPlugin
+
+    QPyDesignerCustomWidgetCollection.addCustomWidget(AutoGainPlugin())
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()

debye_bec/bec_widgets/widgets/restart_server/restart_server.py

@@ -21,10 +21,6 @@ logger = bec_logger.logger
 class RestartServer(BECWidget, QWidget):
     """
     Main widget of server restart widget.
-
-    This widget has no visible UI of its own - it only drives the confirmation/
-    progress dialog. The surrounding dock panel is hidden as soon as we can find
-    it, so the user only ever sees the dialog, never an empty "Restart Server" tab.
     """
 
     PLUGIN = True

debye_bec/bec_widgets/widgets/scheduler/item_dialog.py

@@ -69,12 +69,14 @@ class ScheduleItemDialog(QDialog):
             from ....bec_ipython_client.plugins.digital_twin_core.digital_twin_core import (
                 DigitalTwinCore,
             )
+            from ..auto_gain.auto_gain import AutoGain
             from ..digital_twin.digital_twin import DigitalTwin
             from ..edge_selector import EdgeSelector
             from ..scan_control_xas.scan_control_xas import ScanControlXAS
 
             self.DigitalTwinCore = DigitalTwinCore
             self.DigitalTwin = DigitalTwin
+            self.AutoGain = AutoGain
             self.EdgeSelector = EdgeSelector
             self.ScanControlXAS = ScanControlXAS
 
@@ -496,102 +498,41 @@ class ScheduleItemDialog(QDialog):
     # ------------------------------------------------------------------ #
     def _build_auto_gain_page(self):
         page = QWidget()
-        layout = QVBoxLayout(page)
+        outer = QVBoxLayout(page)
 
-        edge_selector_layout = QHBoxLayout()
-        edge_selector_layout.addWidget(QLabel("Absorption edge:"))
-        self.edge_selector_button = MyButton("Choose", "default")
-        self.edge_label = QLabel("No edge selected")
-        edge_selector_layout.addWidget(self.edge_selector_button)
-        edge_selector_layout.addWidget(self.edge_label)
-        edge_selector_layout.addStretch()
-        self.edge_element = None
-        self.edge_edge = None
-        self.edge_energy = None
+        scroll = QScrollArea()
+        scroll.setWidgetResizable(True)
+        self.auto_gain = self.AutoGain(show_control=False, parent=scroll, client=self._client)
 
-        layout.addLayout(edge_selector_layout)
+        scroll.setWidget(self.auto_gain)
+        outer.addWidget(scroll)
 
-        layout.addWidget(QLabel("Auto gain on"))
+        hint = QLabel(
+            "Configure the auto-gain functionality, then confirm with Add below - the "
+            "auto-gain will only be executed when this schedule item runs, not immediately."
+        )
+        hint.setWordWrap(True)
+        hint.setStyleSheet("color: gray;")
+        outer.addWidget(hint)
 
-        form = QFormLayout()
-        layout.addLayout(form)
-        form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint)
-
-        self.ic0_check = QCheckBox("")
-        form.addRow("IC0", self.ic0_check)
-        self.ic1_check = QCheckBox("")
-        form.addRow("IC1", self.ic1_check)
-        self.ic2_check = QCheckBox("")
-        form.addRow("IC2", self.ic2_check)
-        self.pips_check = QCheckBox("")
-        form.addRow("PIPS", self.pips_check)
-
-        layout.addStretch()
-
-        self.edge_selector_button.clicked.connect(self._update_edge)
-
-        self._add_page("auto_gain", "Auto-Gain", page)
-
-    def _update_edge(self, *_):
-        # NOTE: EdgeSelector is currently the same dialog for every
-        # beamline; this hook is here so a beamline-specific selector can
-        # be swapped in later without touching callers.
-        dlg = self.EdgeSelector(self)
-        if dlg.exec_():
-            self.edge_element = dlg.selected_element
-            self.edge_edge = dlg.selected_edge
-            self.edge_energy = dlg.selected_energy
-            self.edge_label.setText(
-                f"{dlg.selected_element}, {dlg.selected_edge}-edge, {dlg.selected_energy:0.1f} eV"
-            )
+        self._add_page("auto_gain", "Auto Gain", page)
 
     def _collect_auto_gain_result(self) -> dict:
-        if self.edge_element is None or self.edge_edge is None:
-            raise ValueError("Choose an absorption edge first.")
-        amplifier_checks = [
-            (self.ic0_check, "ic0"),
-            (self.ic1_check, "ic1"),
-            (self.ic2_check, "ic2"),
-            (self.pips_check, "pips"),
-        ]
-        amplifiers = [name for check, name in amplifier_checks if check.isChecked()]
-        if not amplifiers:
-            raise ValueError("Select at least one amplifier for auto gain.")
+        config = self.auto_gain.get_config()
+
         cmd = (
-            f"auto_gain.start(element={self.edge_element}, edge={self.edge_edge}, "
-            f"amplifier={amplifiers}, comp_ring_current=True)"
+            f"auto_gain.start(element={config['element']}, edge={config['edge']}, "
+            f"amplifier={config['all_amplifiers']}, comp_ring_current="
+            f"{config['compensate_ring_current']}, prohibited_gains={config['prohibited_gains']})"
         )
         return {
             "command": cmd,
             "kind": "custom",
-            "form_state": {
-                "text": cmd,
-                "source": "auto_gain",
-                "element": self.edge_element,
-                "edge": self.edge_edge,
-                "energy": self.edge_energy,
-                "ic0": self.ic0_check.isChecked(),
-                "ic1": self.ic1_check.isChecked(),
-                "ic2": self.ic2_check.isChecked(),
-                "pips": self.pips_check.isChecked(),
-            },
+            "form_state": {"text": cmd, "source": "auto_gain", "config": config},
         }
 
     def _prefill_auto_gain_form(self, state: dict):
-        element = state.get("element")
-        edge = state.get("edge")
-        if element is not None and edge is not None:
-            self.edge_element = element
-            self.edge_edge = edge
-            self.edge_energy = state.get("energy")
-            if self.edge_energy is not None:
-                self.edge_label.setText(f"{element}, {edge}-edge, {self.edge_energy:0.1f} eV")
-            else:
-                self.edge_label.setText(f"{element}, {edge}-edge")
-        self.ic0_check.setChecked(bool(state.get("ic0", False)))
-        self.ic1_check.setChecked(bool(state.get("ic1", False)))
-        self.ic2_check.setChecked(bool(state.get("ic2", False)))
-        self.pips_check.setChecked(bool(state.get("pips", False)))
+        self.auto_gain.load_config(state["config"])
 
     # ------------------------------------------------------------------ #
     # Other

debye_bec/bec_widgets/widgets/scheduler/scheduler.py

@@ -126,16 +126,15 @@ def build_item_display_text(item: ScheduleItem, beamline: str | None) -> str:
             if beamline == "x01da":
                 text += f"\n     SMPL-XRD: {config['smpl_to_xrd']:.3f} mm"
         elif item.command.startswith("auto_gain.start("):
-            config = item.form_state
-            text += f"Auto-configure gains for {item.form_state['element']}, {item.form_state['edge']}-edge"
-            devices = ""
-            for device in ["ic0", "ic1", "ic2", "pips"]:
-                if config[device] is True:
-                    if len(devices) == 0:
-                        devices += f"\n      for devices: {device}"
-                    else:
-                        devices += f", {device}"
-            text += devices
+            config = item.form_state["config"]
+            text += f"Auto-configure gains for {config['element']}, {config['edge']}-edge"
+            text += f"\nfor devices {config['all_amplifiers']}"
+            if config["link_ic0_ic1"]:
+                text += "\nLink IC0 and IC1 gains"
+            if config["compensate_ring_current"]:
+                text += "\nCalculate for 400 mA ring current"
+            if config["prohibited_gains"] != {}:
+                text += f"\nProhibited gains: {config['prohibited_gains']}"
         elif item.command.startswith("dev.abs."):
             action = item.form_state["action"]
             text += f"{action} Absorber"
@@ -331,13 +330,13 @@ class Scheduler(BECWidget, QWidget):
             logger.info(
                 "Scheduler running at X01DA or X10DA, import and load digital twin and auto-gain"
             )
-            from ....bec_ipython_client.plugins.auto_gain import AutoGain
+            from ....bec_ipython_client.plugins.auto_gain_core import AutoGainCore
             from ....bec_ipython_client.plugins.digital_twin_core.digital_twin_core import (
                 DigitalTwinCore,
             )
 
             self.digital_twin = DigitalTwinCore()
-            self.auto_gain = AutoGain()
+            self.auto_gain = AutoGainCore()
         else:
             self.digital_twin = None
             self.auto_gain = None