wip

debye_bec/bec_widgets/widgets/scheduler/item_dialog.py

@@ -66,24 +66,15 @@ from qtpy.QtWidgets import (
     QLabel,
     QLineEdit,
     QMessageBox,
-    QPushButton,
     QScrollArea,
     QTabWidget,
     QVBoxLayout,
     QWidget,
 )
-
-from ....bec_ipython_client.plugins.digital_twin.digital_twin import DigitalTwinCore
-from ..digital_twin.calculations.calc_positions import calc_positions
-from ..digital_twin.digital_twin import DigitalTwin
-from ..digital_twin.offsets import Offsets
-from ..scan_control_xas.scan_control_xas import ScanControlXAS
 from .qt_widgets import MyButton
 
 logger = bec_logger.logger
 
-_DSPIN_RANGE = (-1e12, 1e12)
-
 # Tab indices, named instead of magic numbers now that there are four -
 # see _collect_result()/_apply_initial().
 _TAB_SCAN = 0
@@ -95,7 +86,7 @@ _TAB_OTHER = 3
 class ScheduleItemDialog(QDialog):
     """Add or edit one schedule item, via ScanControl, a move form, Digital Twin, or free text."""
 
-    def __init__(self, scans, dev, parent=None, initial: dict | None = None, client=None):
+    def __init__(self, scans, dev, parent=None, initial: dict | None = None, client=None, beamline:str|None=None):
         super().__init__(parent)
         self.setWindowTitle("Schedule item")
         self.setMinimumSize(520, 480)
@@ -104,13 +95,21 @@ class ScheduleItemDialog(QDialog):
         self._dev = dev
         self._client = client
 
+        self.beamline = beamline
+        if self.beamline in ['x01da', 'x10da']:
+            from ....bec_ipython_client.plugins.digital_twin.digital_twin import DigitalTwinCore
+            from ..digital_twin.digital_twin import DigitalTwin
+            from ..scan_control_xas.scan_control_xas import ScanControlXAS
+            from ..edge_selector import EdgeSelector
+
         layout = QVBoxLayout(self)
         self.tabs = QTabWidget()
         layout.addWidget(self.tabs)
 
         self._build_scan_tab()
         self._build_move_tab()
-        self._build_digital_twin_tab()
+        if self.beamline in ['x01da', 'x10da']:
+            self._build_digital_twin_tab()
         self._build_custom_tab()
 
         buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
@@ -138,7 +137,10 @@ class ScheduleItemDialog(QDialog):
         # client=None resolves to the same process-wide BEC client
         # (bec_dispatcher.client) our own widget uses - no second Redis
         # connection is opened.
-        self.scan_control = ScanControlXAS(parent=tab, client=self._client)
+        if self.beamline in ['x01da', 'x10da']:
+            self.scan_control = ScanControlXAS(parent=tab, client=self._client)
+        else:
+            self.scan_control = ScanControl(parent=tab, client=self._client)
         self.scan_control.button_run_scan.hide()
         layout.addWidget(self.scan_control)
 
@@ -301,17 +303,58 @@ class ScheduleItemDialog(QDialog):
         self.custom_edit.setPlaceholderText("scans.xas_simple_scan(12000, 14000, 2, 10)")
         layout.addWidget(self.custom_edit)
 
-        ic_form = self._create_ionization_chamber_form()
-        if ic_form is not None:
-            layout.addWidget(ic_form)
+        if self.beamline in ['x01da', 'x10da']:
+            abs_form = self._create_abs_form()
+            if abs_form is not None:
+                layout.addWidget(abs_form)
 
-        reffoil_form = self._create_reffoil_form()
-        if reffoil_form is not None:
-            layout.addWidget(reffoil_form)
+            ic_form = self._create_ionization_chamber_form()
+            if ic_form is not None:
+                layout.addWidget(ic_form)
+
+            reffoil_form = self._create_reffoil_form()
+            if reffoil_form is not None:
+                layout.addWidget(reffoil_form)
+
+            auto_gain_form = self._create_auto_gain_form()
+            if auto_gain_form is not None:
+                layout.addWidget(auto_gain_form)
 
         layout.addStretch(1)
         self.tabs.addTab(tab, "Other")
 
+    def _create_abs_form(self):
+        if "abs" in self._dev:
+            abs_group = QGroupBox("Frontend Absorber")
+            layout = QVBoxLayout(abs_group)
+            form = QFormLayout()
+            layout.addLayout(form)
+            form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint)
+            self.abs_selector = QComboBox()
+            self.abs_selector.addItems(["Open", "Force open", "Close"])
+
+            button_layout = QHBoxLayout()
+            generate_cmd = MyButton("Generate command", "default")
+            button_layout.addWidget(generate_cmd)
+            button_layout.addStretch(1)
+            layout.addLayout(button_layout)
+
+            generate_cmd.clicked.connect(self._generate_abs_command)
+
+            return abs_group
+        return None
+
+    def _generate_abs_command(self):
+        match self.abs_selector.currentText():
+            case 'Open':
+                suffix = 'open()'
+            case 'Force open':
+                suffix = 'open(force=True)'
+            case 'Close':
+                suffix = 'close()'
+        cmd = f"dev.abs.{suffix}"
+        self.custom_edit.setText(cmd)
+
     def _create_ionization_chamber_form(self):
         if all(key in self._dev for key in ("ic0", "ic1", "ic2")):
             ic_group = QGroupBox("Ionization chamber filling")
@@ -411,6 +454,73 @@ class ScheduleItemDialog(QDialog):
         cmd = f"dev.reffoilchanger.insert(ref='{self.reffoil_selector.currentText()}', wait=True)"
         self.custom_edit.setText(cmd)
 
+    def _create_auto_gain_form(self):
+        auto_gain_group = QGroupBox("Auto Gain")
+        layout = QVBoxLayout(auto_gain_group)
+
+        edge_selector_layout = QHBoxLayout()
+        edge_selector_label = QLabel("Absorption edge:")
+        self.edge_selector_button = MyButton("Choose", "default")
+        self.edge_label = QLabel("No edge selected")
+        edge_selector_layout.addWidget(edge_selector_label)
+        edge_selector_layout.addWidget(self.edge_selector_button)
+        edge_selector_layout.addWidget(self.edge_label)
+        self.edge_element = None
+        self.edge_edge = None
+
+        layout.addLayout(edge_selector_layout)
+
+        form = QFormLayout()
+        form.setFieldGrowthPolicy(QFormLayout.FieldGrowthPolicy.FieldsStayAtSizeHint)
+
+        self.ic0_check = QCheckBox("")
+        form.addRow("Auto gain on IC0", self.ic0_check)
+        self.ic1_check = QCheckBox("")
+        form.addRow("          on IC1", self.ic1_check)
+        self.ic2_check = QCheckBox("")
+        form.addRow("          on IC2", self.ic2_check)
+        self.pips_check = QCheckBox("")
+        form.addRow("          on PIPS", self.pips_check)
+
+        button_layout = QHBoxLayout()
+        generate_cmd = MyButton("Generate command", "default")
+        button_layout.addWidget(generate_cmd)
+        button_layout.addStretch(1)
+        layout.addLayout(button_layout)
+
+        self.edge_selector_button.clicked.connect(self._update_edge)
+        generate_cmd.clicked.connect(self._generate_auto_gain_command)
+
+        return auto_gain_group
+
+    def _update_edge(self, *_):
+        match self.beamline:
+            case "x01da":
+                dlg = EdgeSelector(self)
+            case "x10da":
+                dlg = EdgeSelector(self)
+            case _:
+                dlg = EdgeSelector(self)
+        if dlg.exec_():
+            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.edge_element = dlg.selected_element
+            self.edge_edge = dlg.selected_edge
+
+    def _generate_auto_gain_command(self):
+        if self.edge_edge is None or self.edge_element is None:
+            return
+        amplifiers = []
+        for amp, name in [(self.ic0_check, 'ic0'), (self.ic1_check, 'ic1'), (self.ic2_check, 'ic2'), (self.pips_check, 'pips')]:
+            if amp.isChecked():
+                amplifiers.append(name)
+        if amplifiers == []:
+            return
+        cmd = f"auto_gain.start(element={self.edge_element}, edge={self.edge_edge}, {amplifiers}, comp_ring_current=True)"
+        self.custom_edit.setText(cmd)
+
     def _collect_custom_result(self) -> dict:
         text = self.custom_edit.text().strip()
         if not text:

debye_bec/bec_widgets/widgets/scheduler/scheduler.py

@@ -45,7 +45,6 @@ from qtpy.QtWidgets import (
     QWidget,
 )
 
-from ....bec_ipython_client.plugins.digital_twin.digital_twin import DigitalTwinCore
 from .endpoints import schedule as schedule_endpoint
 from .enums import ScheduleItemStatus
 from .guard import SignalGuard
@@ -150,9 +149,20 @@ class Scheduler(BECWidget, QWidget):
         self.schedule_name = self.config.schedule_name
 
         self.get_bec_shortcuts()  # -> self.client, self.dev, self.scans, self.queue
+        self.beamline = self.get_beamline()
         self.connector = self.client.connector
 
-        self.digital_twin = DigitalTwinCore()
+        if self.beamline in ['x01da', 'x10da']:
+            logger.info(
+                'Scheduler running at X01DA or X10DA, import and load digital twin and auto-gain'
+            )
+            from ....bec_ipython_client.plugins.digital_twin.digital_twin import DigitalTwinCore
+            from ....bec_ipython_client.plugins.auto_gain import AutoGain
+            self.digital_twin = DigitalTwinCore()
+            self.auto_gain = AutoGain()
+        else:
+            self.digital_twin = None
+            self.auto_gain = None
 
         self._endpoint = schedule_endpoint(self.schedule_name)
         self._lock = threading.Lock()  # guards self.schedule (see module docstring)
@@ -190,6 +200,12 @@ class Scheduler(BECWidget, QWidget):
 
         self._closing = threading.Event()
 
+    def get_beamline(self):
+        hostname = self.client._hostname
+        start = hostname.find("x")
+        if start != -1:
+            return hostname[start : start + 5]
+
     # ------------------------------------------------------------------ #
     # UI
     # ------------------------------------------------------------------ #
@@ -576,12 +592,8 @@ class Scheduler(BECWidget, QWidget):
             data = self.schedule.model_dump(mode="json")
         default_name = "schedule.json"
 
-        hostname = self.client._hostname
-        start = hostname.find("x")
-        if start != -1:
-            beamline = hostname[start : start + 5]
-            active_account = self.client.active_account
-            default_name = f"/sls/{beamline}/data/{active_account}/raw/{default_name}"
+        active_account = self.client.active_account
+        default_name = f"/sls/{self.beamline}/data/{active_account}/raw/{default_name}"
 
         path, _ = QFileDialog.getSaveFileName(
             self, "Save schedule to file", default_name, "JSON files (*.json);;All files (*)"
@@ -608,12 +620,9 @@ class Scheduler(BECWidget, QWidget):
                 return
 
         start_folder = ""
-        hostname = self.client._hostname
-        start = hostname.find("x")
-        if start != -1:
-            beamline = hostname[start : start + 5]
-            active_account = self.client.active_account
-            start_folder = f"/sls/{beamline}/data/{active_account}/raw"
+
+        active_account = self.client.active_account
+        start_folder = f"/sls/{self.beamline}/data/{active_account}/raw"
 
         path, _ = QFileDialog.getOpenFileName(
             self, "Load schedule from file", start_folder, "JSON files (*.json);;All files (*)"
@@ -766,7 +775,7 @@ class Scheduler(BECWidget, QWidget):
     # ---- UI-triggered edit actions ---- #
     @SafeSlot()
     def _on_add_clicked(self):
-        dialog = ScheduleItemDialog(self.scans, self.dev, parent=self, client=self.client)
+        dialog = ScheduleItemDialog(self.scans, self.dev, parent=self, client=self.client, beamline=self.beamline)
         if dialog.exec_() != QDialog.Accepted:
             return
         result = dialog.result()
@@ -1147,6 +1156,7 @@ class Scheduler(BECWidget, QWidget):
             "scans": self.scans,
             "dev": self.dev,
             "digital_twin": self.digital_twin,
+            "auto_gain": self.auto_gain,
             "np": np,
         }
         while not self._abort_requested and not self._closing.is_set():
@@ -1289,7 +1299,7 @@ class Scheduler(BECWidget, QWidget):
 
 if __name__ == "__main__":
     app = QApplication(sys.argv)
-    dispatcher = BECDispatcher(gui_id="digital_twin")
+    dispatcher = BECDispatcher(gui_id="Scheduler")
     win = Scheduler()
     win.show()
     sys.exit(app.exec_())