feat: add pgmotf device

xtreme_bec/devices/pgm_otf.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from ophyd import Component as Cpt
+from ophyd import Device, EpicsSignal, EpicsSignalRO, Kind
+from ophyd_devices import AsyncSignal, DeviceStatus, PSIDeviceBase, StatusBase, TransitionStatus
+
+
+class PGMOtfControl(Device):
+    """
+    Ophyd device for controlling the PGM OTF (On-The-Fly) scan.
+    """
+
+    e1 = Cpt(EpicsSignal, "E1", kind=Kind.config)
+    e2 = Cpt(EpicsSignal, "E2", kind=Kind.config)
+    time = Cpt(EpicsSignal, "TIME", kind=Kind.config)
+    folder = Cpt(EpicsSignal, "FOLDER", kind=Kind.config)
+    file = Cpt(EpicsSignal, "FILE", kind=Kind.config)
+    acquire = Cpt(EpicsSignal, "START", auto_monitor=True)
+    raw_edata = Cpt(EpicsSignalRO, "EDATA", kind=Kind.hinted, auto_monitor=True)
+    raw_data = Cpt(EpicsSignalRO, "DATA", kind=Kind.hinted, auto_monitor=True)
+    raw_idata = Cpt(EpicsSignalRO, "IDATA", kind=Kind.hinted, auto_monitor=True)
+    raw_fdata = Cpt(EpicsSignalRO, "FDATA", kind=Kind.hinted, auto_monitor=True)
+    count = Cpt(EpicsSignalRO, "COUNT", kind=Kind.omitted, auto_monitor=True)
+
+
+class PGMOtf(PSIDeviceBase, PGMOtfControl):
+    """
+    Ophyd device for controlling the PGM OTF (On-The-Fly) scan.
+    Inherits from PSIDeviceBase and PGMOtfControl.
+    """
+
+    edata = Cpt(
+        AsyncSignal,
+        name="raw_edata",
+        ndim=1,
+        async_update={"type": "replace", "max_shape": [None]},
+        max_size=1000,
+        docs="Aggregated edata signal from the PGM OTF scan.",
+    )
+    data = Cpt(
+        AsyncSignal,
+        name="raw_data",
+        ndim=1,
+        async_update={"type": "replace", "max_shape": [None]},
+        max_size=1000,
+        docs="Aggregated data signal from the PGM OTF scan.",
+    )
+    idata = Cpt(
+        AsyncSignal,
+        name="raw_idata",
+        ndim=1,
+        async_update={"type": "replace", "max_shape": [None]},
+        max_size=1000,
+        docs="Aggregated idata signal from the PGM OTF scan.",
+    )
+    fdata = Cpt(
+        AsyncSignal,
+        name="raw_fdata",
+        ndim=1,
+        async_update={"type": "replace", "max_shape": [None]},
+        max_size=1000,
+        docs="Aggregated fdata signal from the PGM OTF scan.",
+    )
+
+    # FIXME: Once we have the proper energy readback, we should set up the progress update
+    # to use the energy progress between e1 and e2.
+    # energy = Cpt(EpicsSignalRO, "ENERGY_RBV", kind=Kind.omitted, auto_monitor=True)
+    # progress = Cpt(ProgressSignal)
+
+    ########################################
+    #  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.complete_status = 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.raw_edata.subscribe(self._update_data, run=False)
+        self.raw_data.subscribe(self._update_data, run=False)
+        self.raw_idata.subscribe(self._update_data, run=False)
+        self.raw_fdata.subscribe(self._update_data, 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.
+        """
+        if self.acquire.get() != 0:
+            raise RuntimeError("PGM OTF scan is already running. Please stop it before staging.")
+
+        self.complete_status = TransitionStatus(self.acquire, transitions=[0, 1, 0])
+        self.cancel_on_stop(self.complete_status)
+
+    def on_unstage(self) -> DeviceStatus | StatusBase | None:
+        """Called while unstaging the device."""
+        if self.complete_status is not None and not self.complete_status.done:
+            self.complete_status.set_finished()
+            self.acquire.put(0)
+
+        self.complete_status = None
+
+    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."""
+        return self.complete_status
+
+    def on_kickoff(self) -> DeviceStatus | StatusBase | None:
+        """Called to kickoff a device for a fly scan. Has to be called explicitly."""
+        self.acquire.put(1, use_complete=True)
+
+    def on_stop(self) -> None:
+        """Called when the device is stopped."""
+        self.acquire.put(0)
+
+    def on_destroy(self) -> None:
+        """Called when the device is destroyed. Cleanup resources here."""
+
+    def _update_data(self, value, obj, **kwargs):
+        """
+        Callback function to update the aggregated data signals when the raw signals change.
+        """
+        match obj.name:
+            case "raw_edata":
+                self.edata.put(value)
+            case "raw_data":
+                self.data.put(value)
+            case "raw_idata":
+                self.idata.put(value)
+            case "raw_fdata":
+                self.fdata.put(value)