diff --git a/xtreme_bec/devices/pgm_otf.py b/xtreme_bec/devices/pgm_otf.py new file mode 100644 index 0000000..56d2419 --- /dev/null +++ b/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)