dls-dodal 1.64.0__py3-none-any.whl → 1.66.0__py3-none-any.whl

dodal/devices/thawer.py

@@ -1,59 +1,21 @@
-from asyncio import CancelledError, Task, create_task, sleep
-
 from bluesky.protocols import Movable, Stoppable
 from ophyd_async.core import (
     AsyncStatus,
-    Device,
     OnOff,
-    Reference,
-    SignalRW,
     StandardReadable,
 )
 from ophyd_async.epics.core import epics_signal_rw
 
-from dodal.log import LOGGER
-
-
-class ThawingTimer(Device, Stoppable, Movable[float]):
-    def __init__(self, control_signal: SignalRW[OnOff]) -> None:
-        self._control_signal_ref = Reference(control_signal)
-        self._thawing_task: Task | None = None
-        super().__init__("thaw_for_time_s")
 
-    @AsyncStatus.wrap
-    async def set(self, value: float):
-        if self._thawing_task:
-            LOGGER.info("Thawing task already in progress, resetting timer")
-            self._thawing_task.cancel()
-        else:
-            LOGGER.info("Thawing started")
-            await self._control_signal_ref().set(OnOff.ON)
-        self._thawing_task = create_task(sleep(value))
-        try:
-            await self._thawing_task
-        except CancelledError:
-            LOGGER.info("Timer task cancelled.")
-            raise
-        else:
-            LOGGER.info("Thawing completed")
-            await self._control_signal_ref().set(OnOff.OFF)
-
-    @AsyncStatus.wrap
-    async def stop(self, *args, **kwargs):
-        if self._thawing_task:
-            self._thawing_task.cancel()
-            self._thawing_task = None
-        LOGGER.info("Thawer stopped.")
-        await self._control_signal_ref().set(OnOff.OFF)
-
-
-class Thawer(StandardReadable, Stoppable):
+class Thawer(StandardReadable, Stoppable, Movable[OnOff]):
     def __init__(self, prefix: str, name: str = "") -> None:
-        self.control = epics_signal_rw(OnOff, prefix + ":CTRL")
-        self.thaw_for_time_s = ThawingTimer(self.control)
+        self._control = epics_signal_rw(OnOff, prefix + ":CTRL")
         super().__init__(name)
 
+    @AsyncStatus.wrap
+    async def set(self, value: OnOff):
+        await self._control.set(value)
+
     @AsyncStatus.wrap
     async def stop(self, *args, **kwargs):
-        await self.thaw_for_time_s.stop()
-        await self.control.set(OnOff.OFF)
+        await self._control.set(OnOff.OFF)

dodal/devices/undulator.py

@@ -1,7 +1,8 @@
 import os
+from abc import ABC, abstractmethod
 
 import numpy as np
-from bluesky.protocols import Movable
+from bluesky.protocols import Locatable, Location, Movable
 from numpy import ndarray
 from ophyd_async.core import (
     AsyncStatus,
@@ -9,6 +10,7 @@ from ophyd_async.core import (
     StandardReadable,
     StandardReadableFormat,
     soft_signal_r_and_setter,
+    soft_signal_rw,
 )
 from ophyd_async.epics.core import epics_signal_r
 from ophyd_async.epics.motor import Motor
@@ -38,111 +40,177 @@ def _get_gap_for_energy(
     )
 
 
-class Undulator(StandardReadable, Movable[float]):
+class BaseUndulator(StandardReadable, Movable[float], ABC):
     """
-    An Undulator-type insertion device, used to control photon emission at a given
-    beam energy.
+    Base class for undulator devices providing gap control and access management.
+    This class expects target gap value [mm] passed in set method.
     """
 
     def __init__(
         self,
         prefix: str,
-        id_gap_lookup_table_path: str = os.devnull,
-        name: str = "",
         poles: int | None = None,
         length: float | None = None,
+        undulator_period: int | None = None,
         baton: Baton | None = None,
+        name: str = "",
     ) -> None:
-        """Constructor
-
+        """
         Args:
             prefix: PV prefix
-            poles (int): Number of magnetic poles built into the undulator
-            length (float): Length of the undulator in meters
+            poles (int, optional): Number of magnetic poles built into the undulator
+            length (float, optional): Length of the undulator in meters
+            undulator_period(int, optional): Undulator period
+            baton (optional): Baton object if provided.
             name (str, optional): Name for device. Defaults to "".
         """
-
         self.baton_ref = Reference(baton) if baton else None
-        self.id_gap_lookup_table_path = id_gap_lookup_table_path
+
         with self.add_children_as_readables():
+            self.gap_access = epics_signal_r(EnabledDisabledUpper, prefix + "IDBLENA")
             self.gap_motor = Motor(prefix + "BLGAPMTR")
             self.current_gap = epics_signal_r(float, prefix + "CURRGAPD")
-            self.gap_access = epics_signal_r(EnabledDisabledUpper, prefix + "IDBLENA")
 
         with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL):
             self.gap_discrepancy_tolerance_mm, _ = soft_signal_r_and_setter(
                 float,
                 initial_value=UNDULATOR_DISCREPANCY_THRESHOLD_MM,
             )
-            if poles is not None:
-                self.poles, _ = soft_signal_r_and_setter(
-                    int,
-                    initial_value=poles,
-                )
-            else:
-                self.poles = None
-
-            if length is not None:
-                self.length, _ = soft_signal_r_and_setter(
-                    float,
-                    initial_value=length,
-                )
-            else:
-                self.length = None
-
-        super().__init__(name)
+            self.poles = self._make_signal_if_not_none(poles, int)
+            self.length = self._make_signal_if_not_none(length, float)
+            self.undulator_period = self._make_signal_if_not_none(undulator_period, int)
+
+        super().__init__(name=name)
+
+    def _make_signal_if_not_none(self, initial_value, type):
+        if initial_value is None:
+            return None
+        signal, _ = soft_signal_r_and_setter(type, initial_value=initial_value)
+        return signal
 
+    @abstractmethod
     @AsyncStatus.wrap
-    async def set(self, value: float):
+    async def set(self, value: float) -> None:
         """
-        Set the undulator gap to a given energy in keV
+        Move undulator to a given position.
+        Abstract method - must be implemented by subclasses.
 
         Args:
-            value: energy in keV
+            value: target position - units depend on implementation
         """
-        await self._set_undulator_gap(value)
+        ...
 
-    async def raise_if_not_enabled(self):
-        access_level = await self.gap_access.get_value()
-        commissioning_mode = await self._is_commissioning_mode_enabled()
-        if access_level is EnabledDisabledUpper.DISABLED and not commissioning_mode:
-            raise AccessError("Undulator gap access is disabled. Contact Control Room")
+    async def _set_gap(self, value: float) -> None:
+        """
+        Set the undulator gap to a given value in mm.
+
+        Args:
+            value: gap in mm
+        """
+        await self.raise_if_not_enabled()  # Check access
+        if await self._check_gap_within_threshold(value):
+            LOGGER.debug(
+                "Gap is already in the correct place, no need to ask it to move"
+            )
+            return
 
-    async def _set_undulator_gap(self, energy_kev: float) -> None:
-        await self.raise_if_not_enabled()
-        target_gap = await self._get_gap_to_match_energy(energy_kev)
         LOGGER.info(
-            f"Setting undulator gap to {target_gap:.3f}mm based on {energy_kev:.2f}kev"
+            f"Undulator gap mismatch. Moving gap to nominal value, {value:.3f}mm"
         )
+        commissioning_mode = await self._is_commissioning_mode_enabled()
+        if not commissioning_mode:
+            # Only move if the gap is sufficiently different to the value from the
+            # DCM lookup table AND we're not in commissioning mode
+            await self.gap_motor.set(
+                value,
+                timeout=STATUS_TIMEOUT_S,
+            )
+        else:
+            LOGGER.warning("In test mode, not moving ID gap")
+
+    async def _check_gap_within_threshold(self, target_gap: float) -> bool:
+        """
+        Check if the undulator gap is within the acceptable threshold of the target gap.
 
-        # Check if undulator gap is close enough to the value from the DCM
+        Args:
+            target_gap: target gap in mm
+        Returns:
+            True if the gap is within the threshold, False otherwise
+        """
         current_gap = await self.current_gap.get_value()
         tolerance = await self.gap_discrepancy_tolerance_mm.get_value()
-        difference = abs(target_gap - current_gap)
-        if difference > tolerance:
-            LOGGER.info(
-                f"Undulator gap mismatch. {difference:.3f}mm is outside tolerance.\
-                Moving gap to nominal value, {target_gap:.3f}mm"
-            )
-            commissioning_mode = await self._is_commissioning_mode_enabled()
-            if not commissioning_mode:
-                # Only move if the gap is sufficiently different to the value from the
-                # DCM lookup table AND we're not in commissioning mode
-                await self.gap_motor.set(
-                    target_gap,
-                    timeout=STATUS_TIMEOUT_S,
-                )
-            else:
-                LOGGER.warning("In test mode, not moving ID gap")
-        else:
-            LOGGER.debug(
-                "Gap is already in the correct place for the new energy value "
-                f"{energy_kev}, no need to ask it to move"
-            )
+        return abs(target_gap - current_gap) <= tolerance
 
-    async def _is_commissioning_mode_enabled(self):
+    async def _is_commissioning_mode_enabled(self) -> bool | None:
+        """
+        Asynchronously checks if commissioning mode is enabled via the baton reference.
+        """
         return self.baton_ref and await self.baton_ref().commissioning.get_value()
 
+    async def raise_if_not_enabled(self) -> AccessError | None:
+        """
+        Asynchronously raises AccessError if gap access is disabled and not in commissioning mode.
+        """
+        access_level = await self.gap_access.get_value()
+        commissioning_mode = await self._is_commissioning_mode_enabled()
+        if access_level is EnabledDisabledUpper.DISABLED and not commissioning_mode:
+            raise AccessError("Undulator gap access is disabled. Contact Control Room")
+
+
+class UndulatorInKeV(BaseUndulator):
+    """
+    An Undulator-type insertion device, used to control photon emission at a given beam energy.
+    This class expects energy [keV] passed in set method and does conversion to gap
+    internally, for which it requires path to lookup table file in constructor.
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        id_gap_lookup_table_path: str = os.devnull,
+        poles: int | None = None,
+        length: float | None = None,
+        undulator_period: int | None = None,
+        baton: Baton | None = None,
+        name: str = "",
+    ) -> None:
+        """Constructor
+
+        Args:
+            prefix: PV prefix
+            id_gap_lookup_table_path (str): Path to a lookup table file
+            poles (int, optional): Number of magnetic poles built into the undulator
+            length (float, optional): Length of the undulator in meters
+            undulator_period(int, optional): Undulator period
+            baton (optional): Baton object if provided.
+            name (str, optional): Name for device. Defaults to "".
+        """
+
+        self.id_gap_lookup_table_path = id_gap_lookup_table_path
+        super().__init__(
+            prefix=prefix,
+            poles=poles,
+            length=length,
+            undulator_period=undulator_period,
+            baton=baton,
+            name=name,
+        )
+
+    @AsyncStatus.wrap
+    async def set(self, value: float):
+        """
+        Check conditions and Set undulator gap to a given energy in keV
+
+        Args:
+            value: energy in keV
+        """
+        # Convert energy in keV to gap in mm first
+        target_gap = await self._get_gap_to_match_energy(value)
+        LOGGER.info(
+            f"Setting undulator gap to {target_gap:.3f}mm based on {value:.2f}kev"
+        )
+        await self._set_gap(target_gap)
+
     async def _get_gap_to_match_energy(self, energy_kev: float) -> float:
         """
         get a 2d np.array from lookup table that
@@ -157,3 +225,47 @@ class Undulator(StandardReadable, Movable[float]):
             energy_kev * 1000,
             energy_to_distance_table,
         )
+
+
+class UndulatorInMm(BaseUndulator):
+    """
+    An Undulator-type insertion device, used to control photon emission.
+    This class expects gap [mm] passed in set method.
+    """
+
+    @AsyncStatus.wrap
+    async def set(self, value: float):
+        """
+        Check conditions and Set undulator gap to a given value in mm
+
+        Args:
+            value: value in mm
+        """
+        await self._set_gap(value)
+
+
+class UndulatorOrder(StandardReadable, Locatable[int]):
+    """
+    Represents the order of an undulator device. Allows setting and locating the order.
+    """
+
+    def __init__(self, name: str = "") -> None:
+        """
+        Args:
+            name: Name for device. Defaults to ""
+        """
+        with self.add_children_as_readables():
+            self.value = soft_signal_rw(int, initial_value=3)
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: int) -> None:
+        if (value >= 0) and isinstance(value, int):
+            await self.value.set(value)
+        else:
+            raise ValueError(
+                f"Undulator order must be a positive integer. Requested value: {value}"
+            )
+
+    async def locate(self) -> Location[int]:
+        return await self.value.locate()