feat: otf redesign

xtreme_bec/device_configs/optics.yaml

@@ -76,6 +76,32 @@ mono:
   deviceTags:
     - optics
 
+mirror_theta:
+  readoutPriority: baseline
+  description: "Mirror theta"
+  deviceClass: ophyd_devices.EpicsMotorVME
+  deviceConfig:
+    prefix: "X07MA:m1"
+  onFailure: retry
+  enabled: true
+  readOnly: false
+  softwareTrigger: false
+  deviceTags:
+    - optics
+
+grating_beta:
+  readoutPriority: baseline
+  description: "Grating beta"
+  deviceClass: ophyd_devices.EpicsMotorVME
+  deviceConfig:
+    prefix: "X07MA:m2"
+  onFailure: retry
+  enabled: true
+  readOnly: false
+  softwareTrigger: false
+  deviceTags:
+    - optics
+
 fmu_rox:
   readoutPriority: baseline
   description: "Focusing mirror rotation x"

xtreme_bec/devices/__init__.py

@@ -1 +1,2 @@
+from .pgm import PGMMonochromator
 from .x07ma_devices import *

xtreme_bec/devices/adc_signals.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import threading
+
+from ophyd import Component as Cpt
+from ophyd import Device, EpicsSignalRO, Kind, Signal
+from ophyd_devices import AsyncMultiSignal, PSIDeviceBase
+
+
+class NormSignal(Signal):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._metadata.update(write_access=False)
+
+    def wait_for_connection(self, timeout=0):
+        super().wait_for_connection(timeout)
+        self._metadata.update(connected=True)
+
+    def get(self, **kwargs):
+        val1 = self.parent.signal1.get()
+        val2 = self.parent.signal2.get()
+        return val1 / val2 if val2 != 0 else 0
+
+    def describe(self):
+        desc = {
+            "shape": [],
+            "dtype": "number",
+            "source": "PV: {} / {}".format(self.parent.signal1.pvname, self.parent.signal2.pvname),
+            "units": "",
+            "precision": self.parent.signal1.precision,
+        }
+        return desc
+
+
+class NormTEYSignals(Device):
+    signal1 = Cpt(EpicsSignalRO, name="signal1", read_pv="X07MA-ES1-AI:SIGNAL1", kind="omitted")
+    signal2 = Cpt(EpicsSignalRO, name="signal2", read_pv="X07MA-ES1-AI:SIGNAL2", kind="omitted")
+    norm = Cpt(NormSignal, name="norm", kind="hinted")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.norm.name = self.name
+
+
+class NormDIODESignals(Device):
+    signal1 = Cpt(EpicsSignalRO, name="signal1", read_pv="X07MA-ES1-AI:SIGNAL3", kind="omitted")
+    signal2 = Cpt(EpicsSignalRO, name="signal2", read_pv="X07MA-ES1-AI:SIGNAL2", kind="omitted")
+    norm = Cpt(NormSignal, name="norm", kind="hinted")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.norm.name = self.name
+
+
+class ADCSignals(PSIDeviceBase, Device):
+    """
+    ADC inputs
+    """
+
+    s1 = Cpt(EpicsSignalRO, "SIGNAL0", kind=Kind.hinted, auto_monitor=True)
+    s2 = Cpt(EpicsSignalRO, "SIGNAL1", kind=Kind.hinted, auto_monitor=True)
+    s3 = Cpt(EpicsSignalRO, "SIGNAL2", kind=Kind.hinted, auto_monitor=True)
+    s4 = Cpt(EpicsSignalRO, "SIGNAL3", kind=Kind.hinted, auto_monitor=True)
+    s5 = Cpt(EpicsSignalRO, "SIGNAL4", kind=Kind.hinted, auto_monitor=True)
+    s6 = Cpt(EpicsSignalRO, "SIGNAL5", kind=Kind.hinted, auto_monitor=True)
+    s7 = Cpt(EpicsSignalRO, "SIGNAL6", kind=Kind.hinted, auto_monitor=True)
+    norm_tey = Cpt(NormTEYSignals, name="norm_tey", kind=Kind.hinted)
+    norm_diode = Cpt(NormDIODESignals, name="norm_diode", kind=Kind.hinted)
+    energy_cerbk = Cpt(EpicsSignalRO, "X07MA-PGM:CERBK", kind=Kind.omitted, auto_monitor=True)
+    store_updates = Cpt(Signal, value=0, kind=Kind.config)
+    update_window = Cpt(Signal, value=0.1, kind=Kind.config)
+
+    data = Cpt(
+        AsyncMultiSignal,
+        name="data",
+        signals=[
+            "s1",
+            "s2",
+            "s3",
+            "s4",
+            "s5",
+            "s6",
+            "s7",
+            "norm_tey",
+            "norm_diode",
+            "energy_cerbk",
+        ],
+        ndim=1,
+        async_update={"type": "add", "max_shape": [None]},
+        max_size=1000,
+        docs="Aggregated data signals from the X07MA ADC inputs.",
+    )
+
+    def on_init(self) -> None:
+        """
+        Called when the device is initialized.
+        """
+        self._buffer = []
+        self._buffer_lock = threading.RLock()
+        self._flush_timer = None
+
+    def on_connected(self) -> None:
+        """
+        Called after the device is connected and its signals are connected.
+        Default values for signals should be set here.
+        """
+        self.energy_cerbk.subscribe(self._update_data, run=False)
+
+    def on_complete(self) -> None:
+        """
+        Called to inquire if a device has completed a scan.
+        """
+        if self._flush_timer is not None:
+            self._flush_timer.join()
+
+    def on_stop(self) -> None:
+        """
+        Called when the device is stopped.
+        """
+        self.store_updates.put(0)
+
+        if self._flush_timer is not None:
+            self._flush_timer.cancel()
+            self._flush_timer = None
+
+    def on_destroy(self) -> None:
+        """
+        Called when the device is destroyed.
+        """
+        with self._buffer_lock:
+            if self._flush_timer is not None:
+                self._flush_timer.cancel()
+                self._flush_timer = None
+
+    def _update_data(self, value, **kwargs):
+        """
+        Callback function to update the data signal with the latest readings from the ADC inputs.
+        This function is called whenever any of the monitored signals change.
+        Note that this function does not push updates to the data signal immediately. Instead, it
+        buffers the updates and pushes them in batches based on the update_window setting.
+
+        Args:
+            value: The new value of the signal that triggered this callback.
+            **kwargs: Additional keyword arguments passed by the signal subscription.
+        """
+        if not self.store_updates.get():
+            return
+
+        data: dict[str, float] = {
+            # type: ignore
+            "s1": self.s1.get(),
+            "s2": self.s2.get(),
+            "s3": self.s3.get(),
+            "s4": self.s4.get(),
+            "s5": self.s5.get(),
+            "s6": self.s6.get(),
+            "s7": self.s7.get(),
+            "energy_cerbk": self.energy_cerbk.get(),
+        }
+        data["norm_tey"] = data["s1"] / data["s2"] if data["s2"] != 0 else 0
+        data["norm_diode"] = data["s3"] / data["s2"] if data["s2"] != 0 else 0
+
+        with self._buffer_lock:
+            self._buffer.append(data)
+            if self._flush_timer is None:
+                self._flush_timer = threading.Timer(
+                    float(self.update_window.get()), self._flush_buffer
+                )
+                self._flush_timer.daemon = True
+                self._flush_timer.start()
+
+    def _flush_buffer(self) -> None:
+        """
+        Flush the buffered data and push it to the data signal.
+        This function is called by the timer set in the _update_data method. It collects all
+        buffered data, clears the buffer, and pushes the aggregated data to the data signal.
+        """
+        with self._buffer_lock:
+            buffered = self._buffer
+            self._buffer = []
+            self._flush_timer = None
+
+        if not buffered:
+            return
+
+        batched = {
+            key: [entry[key] for entry in buffered]
+            for key in (
+                "s1",
+                "s2",
+                "s3",
+                "s4",
+                "s5",
+                "s6",
+                "s7",
+                "norm_tey",
+                "norm_diode",
+                "energy_cerbk",
+            )
+        }
+        self.data.put(batched)

xtreme_bec/devices/pgm.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import time
+
+from bec_lib import bec_logger
+from ophyd import Component as Cpt
+from ophyd import EpicsSignal, EpicsSignalRO, Kind, PVPositioner, Signal
+from ophyd_devices import DeviceStatus, ProgressSignal, PSIDeviceBase, StatusBase
+
+logger = bec_logger.logger
+
+
+class PGMMonochromatorBase(PVPositioner):
+    """
+    PGM monochromator
+    """
+
+    setpoint = Cpt(EpicsSignal, "PHS-E:GO.A", auto_monitor=True)
+    readback = Cpt(EpicsSignalRO, "PGM:CERBK", kind=Kind.hinted, auto_monitor=True)
+    done = Cpt(EpicsSignalRO, "PHS:alldone", kind=Kind.omitted, auto_monitor=True)
+    mono_energy_readback = Cpt(EpicsSignalRO, "PGM:rbkenergy", kind=Kind.omitted, auto_monitor=True)
+
+    theta = Cpt(EpicsSignalRO, "PGM:theta", kind=Kind.normal, auto_monitor=True)
+    beta = Cpt(EpicsSignalRO, "PGM:beta", kind=Kind.normal, auto_monitor=True)
+
+    stop_signal = Cpt(EpicsSignal, "PGM:stop", kind=Kind.omitted)
+
+    cff = Cpt(EpicsSignal, "PGM:rbkcff", write_pv="PGM:cff.A", kind=Kind.config)
+    with_undulator = Cpt(EpicsSignal, "PHS-E:OPT", kind=Kind.config)
+    energy = Cpt(EpicsSignal, "PGM:energy", kind=Kind.omitted, auto_monitor=True)
+    undulator_energy = Cpt(EpicsSignal, "UIND:ENERGY-SP", kind=Kind.omitted, auto_monitor=True)
+    undulator_done = Cpt(EpicsSignalRO, "UIND:DONE", kind=Kind.omitted, auto_monitor=True)
+
+
+class PGMMonochromator(PSIDeviceBase, PGMMonochromatorBase):
+    """
+    PGM monochromator device that combines the functionality of PSIDeviceBase and PGMMonochromator.
+    """
+
+    USER_ACCESS = ["get_max_speeds", "set_progress_range"]
+
+    sync_with_undulator = Cpt(Signal, value=0, kind=Kind.config)
+    sync_update_period = Cpt(Signal, value=0.05, kind=Kind.config)
+    progress = Cpt(ProgressSignal)
+
+    def get_max_speeds(self, e1: float, e2: float, travel_time: float) -> tuple[float, float]:
+        """
+        Get the maximum speeds for the PGM monochromator axes (theta and beta)
+        based on the energy range and travel time.
+
+        Args:
+            e1 (float): Start energy in eV.
+            e2 (float): End energy in eV.
+            travel_time (float): Travel time in seconds.
+
+        Returns:
+            tuple: Maximum speeds for theta and beta axes in degrees per second.
+        """
+        initial_energy = self.energy.get()
+        self.energy.set(e1).wait()
+        theta1 = self.theta.get()
+        beta1 = self.beta.get()
+
+        self.energy.set(e2).wait()
+        theta2 = self.theta.get()
+        beta2 = self.beta.get()
+
+        self.energy.set(initial_energy).wait()
+
+        delta_theta = abs(theta2 - theta1)
+        delta_beta = abs(beta2 - beta1)
+        max_speed_theta = min(delta_theta / travel_time, self.max_speed)
+        max_speed_beta = min(delta_beta / travel_time, self.max_speed)
+        return max_speed_theta, max_speed_beta
+
+    def set_progress_range(self, e1: float, e2: float) -> None:
+        """
+        Set the energy range for progress calculation.
+
+        Args:
+            e1 (float): Start energy in eV.
+            e2 (float): End energy in eV.
+        """
+        self.energy_range = [e1, e2]
+
+    ########################################
+    #  Beamline Specific Implementations   #
+    ########################################
+
+    def on_init(self) -> None:
+        """
+        Called when the device is initialized.
+
+        No signals are connected at this point. If you like to
+        set default values on signals, please use on_connected instead.
+        """
+        self._sync_config_guard = False
+        self._last_sync_update_ts = 0.0
+        self._progress_done_emitted = False
+        self.readback.name = self.name
+        self.max_speed = 0.038  # maximum speed in degrees per second
+        self.energy_range = []
+
+    def on_connected(self) -> None:
+        """
+        Called after the device is connected and its signals are connected.
+        Default values for signals should be set here.
+        """
+        self.with_undulator.subscribe(self._on_with_undulator_changed, run=False)
+        self.sync_with_undulator.subscribe(self._on_sync_with_undulator_changed, run=False)
+        self.mono_energy_readback.subscribe(self._sync_undulator_to_mono, run=False)
+        self.readback.subscribe(self._update_progress_from_readback, run=False)
+
+    def on_stage(self) -> DeviceStatus | StatusBase | None:
+        """
+        Called while staging the device.
+        Information about the upcoming scan can be accessed from the scan_info (self.scan_info.msg) object.
+        """
+
+    def on_unstage(self) -> DeviceStatus | StatusBase | None:
+        """Called while unstaging the device."""
+
+    def on_pre_scan(self) -> DeviceStatus | StatusBase | None:
+        """Called right before the scan starts on all devices automatically."""
+
+    def on_trigger(self) -> DeviceStatus | StatusBase | None:
+        """Called when the device is triggered."""
+
+    def on_complete(self) -> DeviceStatus | StatusBase | None:
+        """Called to inquire if a device has completed a scans."""
+
+    def on_kickoff(self) -> DeviceStatus | StatusBase | None:
+        """Called to kickoff a device for a fly scan. Has to be called explicitly."""
+
+    def on_stop(self) -> None:
+        """Called when the device is stopped."""
+
+    def on_destroy(self) -> None:
+        """Called when the device is destroyed. Cleanup resources here."""
+
+    #######################################
+    #       Internal Helper Methods       #
+    #######################################
+
+    def _on_with_undulator_changed(self, value, **kwargs) -> None:
+        """
+        Called when the with_undulator signal changes.
+        If with_undulator is set to True, sync_with_undulator will be set to False to avoid conflicts.
+
+        Args:
+            value (bool): The new value of the with_undulator signal.
+        """
+        if self._sync_config_guard:
+            return
+
+        if bool(value) and bool(self.sync_with_undulator.get()):
+            self._sync_config_guard = True
+            try:
+                self.sync_with_undulator.put(0)
+            finally:
+                self._sync_config_guard = False
+
+    def _on_sync_with_undulator_changed(self, value, **kwargs) -> None:
+        """
+        Called when the sync_with_undulator signal changes.
+        If sync_with_undulator is set to True, with_undulator will be set to False to avoid conflicts.
+
+        Args:
+            value (bool): The new value of the sync_with_undulator signal.
+        """
+        if self._sync_config_guard:
+            return
+
+        if bool(value) and bool(self.with_undulator.get()):
+            self._sync_config_guard = True
+            try:
+                self.with_undulator.put(0)
+            finally:
+                self._sync_config_guard = False
+
+    def _sync_undulator_to_mono(self, value, old_value=None, **kwargs) -> None:
+        """
+        Syncs the undulator energy to the monochromator energy if the sync_with_undulator signal is enabled.
+        It updates the undulator energy only if the with_undulator signal is disabled, the undulator
+        is done, and the sync update period has elapsed since the last update.
+
+        Args:
+            value (float): The new value of the monochromator energy.
+            old_value (float, optional): The previous value of the monochromator energy.
+        """
+        if not bool(self.sync_with_undulator.get()):
+            return
+        if bool(self.with_undulator.get()):
+            return
+        if old_value == value:
+            return
+        if not bool(self.undulator_done.get()):
+            return
+        now = time.monotonic()
+        if now - self._last_sync_update_ts < float(self.sync_update_period.get()):
+            return
+
+        try:
+            self.undulator_energy.put(value)
+            self._last_sync_update_ts = now
+        except Exception:
+            logger.exception("Failed to keep undulator synced with mono for %s", self.name)
+
+    def _update_progress_from_readback(self, value, **kwargs) -> None:
+        """
+        Update the progress signal based on the monochromator readback value.
+        The progress is calculated based on the energy range (e2 - e1) and the current readback value.
+
+        Args:
+            value (float): The current readback value of the monochromator.
+        """
+        if len(self.energy_range) != 2:
+            return
+        if bool(self.done.get()):
+            if not self._progress_done_emitted:
+                self.progress.put(value=100, max_value=100, done=True)
+                self._progress_done_emitted = True
+            return
+
+        self._progress_done_emitted = False
+
+        e1, e2 = self.energy_range
+        if e1 == e2:
+            progress = 100.0
+        else:
+            progress = (float(value) - e1) / (e2 - e1) * 100.0
+            progress = max(0.0, min(100.0, progress))
+
+        self.progress.put(value=progress, max_value=100, done=False)

xtreme_bec/scans/otf_scan.py

@@ -70,7 +70,7 @@ class OTFScan(ScanBase):
         self.e2 = e2
         self.time = time
         self.mono = self.dev["mono"]
-        self.otf = self.dev["otf"]
+        self.max_speed = 0.038
 
         self.update_scan_info()
 
@@ -81,8 +81,26 @@ class OTFScan(ScanBase):
         before the scan is opened, such as preparing the positions (if not done already)
         or setting up the devices.
         """
+        theta_speed, beta_speed = self.dev.mono.get_max_speeds(
+            e1=self.e1, e2=self.e2, travel_time=self.time * 60
+        )
+        self.dev.mono.set_progress_range(e1=self.e1, e2=self.e2)
+        self.actions.set(
+            device=[
+                self.dev.mirror_theta.base_velocity,
+                self.dev.grating_beta.base_velocity,
+                self.dev.mirror_theta.velocity,
+                self.dev.grating_beta.velocity,
+            ],
+            value=[theta_speed / 10, beta_speed / 10, theta_speed, beta_speed],
+        )
+        # We set the readout priority of the adc signals to async so that the monitored readout
+        # does not trigger a readout of the adc signals. The adc signals are going through the
+        # async readout path.
+        self.actions.set_device_readout_priority(devices=[self.dev.signals], priority="async")
+
+        self.actions.add_scan_report_instruction_device_progress(device=self.mono)
 
-        self.actions.add_scan_report_instruction_device_progress(device=self.otf)
         self._baseline_readout_status = self.actions.read_baseline_devices(wait=False)
 
     @scan_hook
@@ -119,6 +137,7 @@ class OTFScan(ScanBase):
 
         # Move to the first energy
         self.mono.set(self.e1).wait()
+        self.mono.sync_with_undulator.put(1)
 
     @scan_hook
     def scan_core(self):
@@ -127,17 +146,8 @@ class OTFScan(ScanBase):
         This is where the main scan logic should be implemented.
         """
 
-        e1_status = self.otf.e1.set(self.e1)
-        e2_status = self.otf.e2.set(self.e2)
-        time_status = self.otf.time.set(self.time)
-
-        e1_status.wait()
-        e2_status.wait()
-        time_status.wait()
-
-        self.actions.kickoff(device=self.otf)
-
-        status = self.actions.complete(device=self.otf, wait=False)
+        self.dev.signals.store_updates.put(1)
+        status = self.mono.set(self.e2)
 
         while not status.done:
             self.at_each_point()
@@ -154,7 +164,9 @@ class OTFScan(ScanBase):
         """
         Post-scan steps to be executed after the main scan logic.
         """
+        self.dev.signals.store_updates.put(0)
         self.actions.complete_all_devices()
+        self.mono.sync_with_undulator.put(0)
 
     @scan_hook
     def unstage(self):
@@ -176,9 +188,9 @@ class OTFScan(ScanBase):
         This is a good place to implement any cleanup logic that needs to be executed in case of an exception,
         such as returning the devices to a safe state or moving the motors back to their starting position.
         """
+        self.dev.signals.store_updates.put(0)
+        self.mono.sync_with_undulator.put(0)
 
     #######################################################
     ######### Helper methods for the scan logic ###########
     #######################################################
-
-    # Implement scan-specific helper methods below.