fix(auto_gain): testing and bugfixes auto_gain_core, auto_gain GUI, auto_gain implementation in scheduler

debye_bec/bec_ipython_client/plugins/auto_gain_core.py

@@ -70,7 +70,7 @@ class AutoGainCore:
     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
+            gains = AVAILABLE_GAINS.copy()
             for prohibited_gain in self.prohibited_gains[amplifier]:
                 gains.remove(prohibited_gain)
             return gains
@@ -224,6 +224,8 @@ class AutoGainCore:
                 ch["self.dev"].set_gain(lowest_gain)
                 ch["gain"] = lowest_gain
 
+            time.sleep(0.5)
+
             remeasure = True
             logger.info(f"Start measurement from {emin} eV to {emax} eV")
             while remeasure:
@@ -307,6 +309,7 @@ class AutoGainCore:
                             for gname in group_names:
                                 active_channels[gname]["self.dev"].set_gain(next_gain)
                                 active_channels[gname]["gain"] = next_gain
+                            time.sleep(0.5)
                             logger.info(f"Setting gain of {label} to {next_gain:.0e} and remeasure")
                             remeasure = True
                     else:

debye_bec/bec_widgets/widgets/auto_gain/auto_gain.py

@@ -26,28 +26,31 @@ class Worker(QObject):
     progress = Signal(str)
     completed = Signal(int)
 
-    @SafeSlot(int)
+    @SafeSlot(dict)
     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)
+        success = False
+        try:
+            core = AutoGainCore()
+            self.progress.emit(f"AutoGain started with config {config}")
+            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"],
+            )
+            lines = [
+                f"{amp} set to gain {res['gain']:.0e}, max signal {res['max_signal']:.3f} V"
+                for amp, res in results.items()
+            ]
+            self.progress.emit("\n".join(lines))
+            success = True
+        except Exception as exc:
+            logger.exception("Auto gain failed")
+            self.progress.emit(f"Error: {exc}")
+        finally:
+            self.completed.emit(success)
 
 
 class AutoGain(BECWidget, QWidget):
@@ -58,6 +61,8 @@ class AutoGain(BECWidget, QWidget):
     PLUGIN = True
     ICON_NAME = "tune"
 
+    start_worker = Signal(dict)
+
     def __init__(self, show_control: bool = True, *arg, parent=None, **kwargs):
         super().__init__(parent=parent, *arg, **kwargs)
         self.get_bec_shortcuts()
@@ -86,10 +91,8 @@ class AutoGain(BECWidget, QWidget):
         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")
 
@@ -118,10 +121,13 @@ class AutoGain(BECWidget, QWidget):
         self.start_button.clicked.connect(self._start)
 
         self.worker = Worker()
+        self.start_worker.connect(self.worker.start)
         self.worker_thread = QThread()
+        self.worker.moveToThread(self.worker_thread)
+
         self.worker.progress.connect(self._update_status)
         self.worker.completed.connect(self._complete)
-        self.worker.moveToThread(self.worker_thread)
+        self.worker.completed.connect(self.worker_thread.quit)
 
     def _update_edge(self, *_):
         dlg = EdgeSelector(self)
@@ -138,6 +144,7 @@ class AutoGain(BECWidget, QWidget):
             self.start_button.setEnabled(False)
             self.status.setColor(get_accent_colors().warning.name())
             self.worker_thread.start()
+            self.start_worker.emit(self.get_config())
 
     def _update_status(self, text):
         self.status.setText(text)
@@ -174,13 +181,13 @@ class AutoGain(BECWidget, QWidget):
             "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,
+            "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,
+            "link_ic0_ic1": self.link_ic0_ic1.isChecked(),
+            "compensate_ring_current": self.comp_ring_current.isChecked(),
             "prohibited_gains": self.prohibited_gains.selection(),
         }
 

debye_bec/bec_widgets/widgets/client.py

@@ -13,6 +13,7 @@ logger = bec_logger.logger
 
 
 _Widgets = {
+    "AutoGain": "AutoGain",
     "DataViewer": "DataViewer",
     "DigitalTwin": "DigitalTwin",
     "RestartServer": "RestartServer",
@@ -21,6 +22,30 @@ _Widgets = {
 }
 
 
+class AutoGain(RPCBase):
+    """Main widget of auto gain widget."""
+
+    _IMPORT_MODULE = "debye_bec.bec_widgets.widgets.auto_gain.auto_gain"
+
+    @rpc_call
+    def remove(self):
+        """
+        Cleanup the BECConnector
+        """
+
+    @rpc_call
+    def attach(self):
+        """
+        None
+        """
+
+    @rpc_call
+    def detach(self):
+        """
+        Detach the widget from its parent dock widget (if widget is in the dock), making it a floating widget.
+        """
+
+
 class DataViewer(RPCBase):
     """Main widget of Data Viewer"""
 
@@ -124,7 +149,7 @@ class Scheduler(RPCBase):
     _IMPORT_MODULE = "debye_bec.bec_widgets.widgets.scheduler.scheduler"
 
     @rpc_call
-    def run_schedule(self):
+    def run_schedule(self, first_item: "int" = 0):
         """
         Run the schedule, continuing from wherever it last left off.
         """
@@ -154,6 +179,7 @@ class Scheduler(RPCBase):
         index: "int | None" = None,
         kind: "str" = "custom",
         form_state: "dict | None" = None,
+        pause_on_failure: "bool" = False,
     ) -> "str":
         """
         RPC-exposed: insert a new, PENDING command into the schedule.
@@ -181,6 +207,7 @@ class Scheduler(RPCBase):
         command: "str",
         kind: "str" = "custom",
         form_state: "dict | None" = None,
+        pause_on_failure: "bool" = True,
     ):
         """
         RPC-exposed: change the command of an item that has not started

debye_bec/bec_widgets/widgets/designer_plugins.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 # pylint: skip-file
 
 designer_plugins = {
+    "AutoGain": ("debye_bec.bec_widgets.widgets.auto_gain.auto_gain", "AutoGain"),
     "DataViewer": ("debye_bec.bec_widgets.widgets.data_viewer.data_viewer", "DataViewer"),
     "DigitalTwin": ("debye_bec.bec_widgets.widgets.digital_twin.digital_twin", "DigitalTwin"),
     "RestartServer": (
@@ -19,6 +20,7 @@ designer_plugins = {
 }
 
 widget_icons = {
+    "AutoGain": "tune",
     "DataViewer": "find_in_page",
     "DigitalTwin": "lightbulb",
     "RestartServer": "restart_alt",

debye_bec/bec_widgets/widgets/digital_twin/beamline.py

@@ -1,51 +0,0 @@
-import socket
-
-from bec_lib import bec_logger
-
-from .types import BeamlineId
-
-logger = bec_logger.logger
-
-
-def get_beamline_id() -> BeamlineId:
-    """
-    Based on the bec servers hostname, tries to extract the beamline
-    identifier (e.g. x01da, x10da, etc).
-
-    Raises:
-        ValueError if beamline cannot be extracted from hostname or beamline not implemented.
-    """
-    bec_hostname = socket.gethostname()
-    start = bec_hostname.find("x")
-    if start != -1:
-        beamline = bec_hostname[start : start + 5]
-        match beamline:
-            case "x01da":
-                return BeamlineId.X01DA
-            case "x10da":
-                return BeamlineId.X10DA
-            case _:
-                raise ValueError(f"Not implemented beamline {beamline}")
-    else:
-        logger.warning(f"Failed to extract beamline from bec server hostname {bec_hostname}")
-        choice = input("Do you want to manually select a beamline? (yes/no): ").strip().lower()
-        if choice in ["yes", "y"]:
-            bl = input(f"Choose from: {[bl.value for bl in BeamlineId]}")
-            if bl in BeamlineId:
-                logger.info(f"Manually selected beamline {bl}")
-                return BeamlineId(bl)
-            else:
-                raise ValueError(f"Wrong selection {bl}")
-        else:
-            raise ValueError("Cannot open digital twin without a beamline")
-
-
-def get_parameters():
-    beamline = get_beamline_id()
-    if beamline == "x01da":
-        from . import x01da_parameters as parameters
-    elif beamline == "x10da":
-        from . import x10da_parameters as parameters
-    else:
-        raise ValueError(f"Unknown beamline: {beamline}")
-    return parameters

debye_bec/bec_widgets/widgets/scheduler/item_dialog.py

@@ -521,9 +521,10 @@ class ScheduleItemDialog(QDialog):
         config = self.auto_gain.get_config()
 
         cmd = (
-            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']})"
+            f"auto_gain.start(element='{config['element']}', edge='{config['edge']}', "
+            + f"amplifier={config['all_amplifiers']}, link_ic0_ic1={config['link_ic0_ic1']}, "
+            + f"comp_ring_current= {config['compensate_ring_current']}, "
+            + f"prohibited_gains={config['prohibited_gains']})"
         )
         return {
             "command": cmd,

debye_bec/bec_widgets/widgets/scheduler/scheduler.py

@@ -134,7 +134,15 @@ def build_item_display_text(item: ScheduleItem, beamline: str | None) -> str:
             if config["compensate_ring_current"]:
                 text += "\nCalculate for 400 mA ring current"
             if config["prohibited_gains"] != {}:
-                text += f"\nProhibited gains: {config['prohibited_gains']}"
+                prohibited_gains = (
+                    "{"
+                    + ", ".join(
+                        f"{amp}: [" + ", ".join(f"{g:.0e}" for g in gains) + "]"
+                        for amp, gains in config["prohibited_gains"].items()
+                    )
+                    + "}"
+                )
+                text += "\nProhibited gains: " + prohibited_gains
         elif item.command.startswith("dev.abs."):
             action = item.form_state["action"]
             text += f"{action} Absorber"
@@ -559,6 +567,7 @@ class Scheduler(BECWidget, QWidget):
     def _load_or_seed_schedule(self) -> Schedule:
         logger.info("load or seed schedule")
         msg: VariableMessage | None = self.connector.get(self._endpoint)
+        # msg = None  # Overwrite any existing schedule in redis
         if msg is not None:
             logger.info(f"Got msg from endpoint: {msg.value}")
             return Schedule.model_validate(msg.value)