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

dodal/devices/insertion_device/__init__.py

@@ -0,0 +1,58 @@
+from .apple2_undulator import (
+    DEFAULT_MOTOR_MIN_TIMEOUT,
+    MAXIMUM_MOVE_TIME,
+    Apple2,
+    Apple2Controller,
+    Apple2EnforceLHMoveController,
+    Apple2LockedPhasesVal,
+    Apple2PhasesVal,
+    Apple2Val,
+    BeamEnergy,
+    EnabledDisabledUpper,
+    EnergyMotorConvertor,
+    InsertionDeviceEnergy,
+    InsertionDevicePolarisation,
+    UndulatorGap,
+    UndulatorJawPhase,
+    UndulatorLockedPhaseAxes,
+    UndulatorPhaseAxes,
+)
+from .energy_motor_lookup import (
+    ConfigServerEnergyMotorLookup,
+    EnergyMotorLookup,
+)
+from .id_enum import Pol, UndulatorGateStatus
+from .lookup_table_models import (
+    EnergyCoverage,
+    LookupTable,
+    LookupTableColumnConfig,
+    convert_csv_to_lookup,
+)
+
+__all__ = [
+    "Apple2",
+    "Apple2Controller",
+    "Apple2EnforceLHMoveController",
+    "UndulatorGap",
+    "UndulatorPhaseAxes",
+    "UndulatorJawPhase",
+    "Apple2Val",
+    "Apple2PhasesVal",
+    "MAXIMUM_MOVE_TIME",
+    "LookupTable",
+    "LookupTableColumnConfig",
+    "convert_csv_to_lookup",
+    "InsertionDeviceEnergy",
+    "InsertionDevicePolarisation",
+    "BeamEnergy",
+    "UndulatorLockedPhaseAxes",
+    "EnergyCoverage",
+    "Pol",
+    "DEFAULT_MOTOR_MIN_TIMEOUT",
+    "EnabledDisabledUpper",
+    "UndulatorGateStatus",
+    "Apple2LockedPhasesVal",
+    "EnergyMotorLookup",
+    "ConfigServerEnergyMotorLookup",
+    "EnergyMotorConvertor",
+]

dodal/devices/{apple2_undulator.py → insertion_device/apple2_undulator.py}

@@ -14,7 +14,6 @@ from ophyd_async.core import (
     SignalW,
     StandardReadable,
     StandardReadableFormat,
-    StrictEnum,
     derived_signal_rw,
     soft_signal_r_and_setter,
     soft_signal_rw,
@@ -24,6 +23,8 @@ from ophyd_async.epics.core import epics_signal_r, epics_signal_rw, epics_signal
 from ophyd_async.epics.motor import Motor
 
 from dodal.common.enums import EnabledDisabledUpper
+from dodal.devices.insertion_device.energy_motor_lookup import EnergyMotorLookup
+from dodal.devices.insertion_device.id_enum import Pol, UndulatorGateStatus
 from dodal.log import LOGGER
 
 T = TypeVar("T")
@@ -32,11 +33,6 @@ DEFAULT_MOTOR_MIN_TIMEOUT = 10
 MAXIMUM_MOVE_TIME = 550  # There is no useful movements take longer than this.
 
 
-class UndulatorGateStatus(StrictEnum):
-    OPEN = "Open"
-    CLOSE = "Closed"
-
-
 @dataclass
 class Apple2LockedPhasesVal:
     top_outer: str
@@ -58,16 +54,6 @@ class Apple2Val:
         return self.phase
 
 
-class Pol(StrictEnum):
-    NONE = "None"
-    LH = "lh"
-    LV = "lv"
-    PC = "pc"
-    NC = "nc"
-    LA = "la"
-    LH3 = "lh3"
-
-
 ROW_PHASE_MOTOR_TOLERANCE = 0.004
 MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
 MAXIMUM_GAP_MOTOR_POSITION = 100
@@ -398,7 +384,7 @@ class Apple2(StandardReadable, Movable[Apple2Val], Generic[PhaseAxesType]):
 
 
 class EnergyMotorConvertor(Protocol):
-    def __call__(self, energy: float, pol: Pol) -> tuple[float, float]:
+    def __call__(self, energy: float, pol: Pol) -> float:
         """Protocol to provide energy to motor position conversion"""
         ...
 
@@ -426,19 +412,18 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         Soft signal for the polarisation setpoint.
     polarisation : derived_signal_rw
         Hardware-backed signal for polarisation readback and control.
-    energy_to_motor : EnergyMotorConvertor
-        Callable that converts energy and polarisation to motor positions.
+    gap_energy_to_motor_converter : EnergyMotorConvertor
+        Callable that converts energy and polarisation to gap motor positions.
+    phase_energy_to_motor_converter : EnergyMotorConvertor
+        Callable that converts energy and polarisation to phase motor positions.
 
     Abstract Methods
     ----------------
-    _set_motors_from_energy(value: float) -> None
-        Abstract method to set motor positions for a given energy and polarisation.
-    energy_to_motor : EnergyMotorConvertor
-        A callable that converts energy and polarisation to motor positions.
-
+    _get_apple2_value(gap: float, phase: float) -> Apple2Val
+        Abstract method to return the Apple2Val used to set the apple2 with.
     Notes
     -----
-    - Subclasses must implement `_set_motors_from_energy` for beamline-specific logic.
+    - Subclasses must implement `_get_apple2_value` for beamline-specific logic.
     - LH3 polarisation is indistinguishable from LH in hardware; special handling is provided.
     - Supports multiple polarisation modes, including linear horizontal (LH), linear vertical (LV),
       positive circular (PC), negative circular (NC), and linear arbitrary (LA).
@@ -448,7 +433,9 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
     def __init__(
         self,
         apple2: Apple2Type,
-        energy_to_motor_converter: EnergyMotorConvertor,
+        gap_energy_motor_converter: EnergyMotorConvertor,
+        phase_energy_motor_converter: EnergyMotorConvertor,
+        units: str = "eV",
         name: str = "",
     ) -> None:
         """
@@ -460,19 +447,20 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         name: str
             Name of the device.
         """
-        self.energy_to_motor = energy_to_motor_converter
         self.apple2 = Reference(apple2)
+        self.gap_energy_motor_converter = gap_energy_motor_converter
+        self.phase_energy_motor_converter = phase_energy_motor_converter
 
         # Store the set energy for readback.
         self._energy, self._energy_set = soft_signal_r_and_setter(
-            float, initial_value=None, units="eV"
+            float, initial_value=None, units=units
         )
         with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
             self.energy = derived_signal_rw(
                 raw_to_derived=self._read_energy,
                 set_derived=self._set_energy,
                 energy=self._energy,
-                derived_units="eV",
+                derived_units=units,
             )
 
         # Store the polarisation for setpoint. And provide readback for LH3.
@@ -480,10 +468,11 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         self.polarisation_setpoint, self._polarisation_setpoint_set = (
             soft_signal_r_and_setter(Pol)
         )
+        phase = self.apple2().phase()
         # check if undulator phase is unlocked.
-        if isinstance(self.apple2().phase(), UndulatorPhaseAxes):
-            top_inner = self.apple2().phase().top_inner.user_readback
-            btm_outer = self.apple2().phase().btm_outer.user_readback
+        if isinstance(phase, UndulatorPhaseAxes):
+            top_inner = phase.top_inner.user_readback
+            btm_outer = phase.btm_outer.user_readback
         else:
             # If locked phase axes make the locked phase 0.
             top_inner = btm_outer = soft_signal_rw(float, initial_value=0.0)
@@ -495,24 +484,35 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
                 raw_to_derived=self._read_pol,
                 set_derived=self._set_pol,
                 pol=self.polarisation_setpoint,
-                top_outer=self.apple2().phase().top_outer.user_readback,
+                top_outer=phase.top_outer.user_readback,
                 top_inner=top_inner,
-                btm_inner=self.apple2().phase().btm_inner.user_readback,
+                btm_inner=phase.btm_inner.user_readback,
                 btm_outer=btm_outer,
                 gap=self.apple2().gap().user_readback,
             )
         super().__init__(name)
 
     @abc.abstractmethod
-    async def _set_motors_from_energy(self, value: float) -> None:
+    def _get_apple2_value(self, gap: float, phase: float, pol: Pol) -> Apple2Val:
         """
         This method should be implemented by the beamline specific ID class as the
         motor positions will be different for each beamline depending on the
-        undulator design and the lookup table used.
+        undulator design.
         """
 
+    async def _set_motors_from_energy_and_polarisation(
+        self, energy: float, pol: Pol
+    ) -> None:
+        """Set the undulator motors for a given energy and polarisation."""
+        gap = self.gap_energy_motor_converter(energy=energy, pol=pol)
+        phase = self.phase_energy_motor_converter(energy=energy, pol=pol)
+        apple2_val = self._get_apple2_value(gap, phase, pol)
+        LOGGER.info(f"Setting polarisation to {pol}, with values: {apple2_val}")
+        await self.apple2().set(id_motor_values=apple2_val)
+
     async def _set_energy(self, energy: float) -> None:
-        await self._set_motors_from_energy(energy)
+        pol = await self._check_and_get_pol_setpoint()
+        await self._set_motors_from_energy_and_polarisation(energy, pol)
         self._energy_set(energy)
 
     def _read_energy(self, energy: float) -> float:
@@ -524,7 +524,6 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         Check the polarisation setpoint and if it is NONE try to read it from
         hardware.
         """
-
         pol = await self.polarisation_setpoint.get_value()
 
         if pol == Pol.NONE:
@@ -653,6 +652,70 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         return Pol.NONE, 0.0
 
 
+class Apple2EnforceLHMoveController(Apple2Controller[Apple2]):
+    """The latest Apple2 version allows unrestricted motor movement.
+    However, because of the high forces involved in polarization changes,
+    all movements must be performed using the Linear Horizontal (LH) mode.
+    A look-up table must also be used to determine the highest energy that can
+    be reached in LH mode."""
+
+    def __init__(
+        self,
+        apple2: Apple2,
+        gap_energy_motor_lut: EnergyMotorLookup,
+        phase_energy_motor_lut: EnergyMotorLookup,
+        units: str = "eV",
+        name: str = "",
+    ) -> None:
+        self.gap_energy_motor_lu = gap_energy_motor_lut
+        self.phase_energy_motor_lu = phase_energy_motor_lut
+        super().__init__(
+            apple2=apple2,
+            gap_energy_motor_converter=gap_energy_motor_lut.find_value_in_lookup_table,
+            phase_energy_motor_converter=phase_energy_motor_lut.find_value_in_lookup_table,
+            units=units,
+            name=name,
+        )
+
+    def _get_apple2_value(self, gap: float, phase: float, pol: Pol) -> Apple2Val:
+        apple2_val = Apple2Val(
+            gap=f"{gap:.6f}",
+            phase=Apple2PhasesVal(
+                top_outer=f"{phase:.6f}",
+                top_inner=f"{0.0:.6f}",
+                btm_inner=f"{phase:.6f}",
+                btm_outer=f"{0.0:.6f}",
+            ),
+        )
+        LOGGER.info(f"Getting apple2 value for pol={pol}, gap={gap}, phase={phase}.")
+        LOGGER.info(f"Apple2 motor values: {apple2_val}.")
+
+        return apple2_val
+
+    async def _set_pol(
+        self,
+        value: Pol,
+    ) -> None:
+        # I09/I21 require all polarisation change to go via LH.
+        current_pol = await self.polarisation.get_value()
+        if current_pol == value:
+            LOGGER.info(f"Polarisation already at {value}")
+        else:
+            target_energy = await self.energy.get_value()
+            if (value is not Pol.LH) and (current_pol is not Pol.LH):
+                self._polarisation_setpoint_set(Pol.LH)
+                max_lh_energy = float(
+                    self.gap_energy_motor_lu.lut.root[Pol("lh")].max_energy
+                )
+                lh_setpoint = (
+                    max_lh_energy if target_energy > max_lh_energy else target_energy
+                )
+                LOGGER.info(f"Changing polarisation to {value} via {Pol.LH}")
+                await self.energy.set(lh_setpoint, timeout=MAXIMUM_MOVE_TIME)
+            self._polarisation_setpoint_set(value)
+            await self.energy.set(target_energy, timeout=MAXIMUM_MOVE_TIME)
+
+
 class InsertionDeviceEnergyBase(abc.ABC, StandardReadable, Movable):
     """Base class for ID energy movable device."""
 
@@ -717,12 +780,7 @@ class InsertionDeviceEnergy(InsertionDeviceEnergyBase):
         self.energy = Reference(id_controller.energy)
         super().__init__(name=name)
 
-        self.add_readables(
-            [
-                self.energy(),
-            ],
-            StandardReadableFormat.HINTED_SIGNAL,
-        )
+        self.add_readables([self.energy()], StandardReadableFormat.HINTED_SIGNAL)
 
     @AsyncStatus.wrap
     async def set(self, energy: float) -> None:

dodal/devices/insertion_device/energy_motor_lookup.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+
+from daq_config_server.client import ConfigServer
+
+from dodal.devices.insertion_device.id_enum import Pol
+from dodal.devices.insertion_device.lookup_table_models import (
+    LookupTable,
+    LookupTableColumnConfig,
+    convert_csv_to_lookup,
+)
+
+
+class EnergyMotorLookup:
+    """
+    Handles a lookup table for Apple2 ID, converting energy/polarisation to a motor
+    position.
+
+    After update_lookup_table() has populated the lookup table, `find_value_in_lookup_table()`
+    can be used to compute gap / phase for a requested energy / polarisation pair.
+    """
+
+    def __init__(self, lut: LookupTable | None = None):
+        if lut is None:
+            lut = LookupTable()
+        self.lut = lut
+
+    def update_lookup_table(self) -> None:
+        """Do nothing by default. Sub classes may override this method to provide logic
+        on what updating lookup table does."""
+        pass
+
+    def find_value_in_lookup_table(self, energy: float, pol: Pol) -> float:
+        """
+        Convert energy and polarisation to a value from the lookup table.
+
+        Parameters:
+        -----------
+        energy : float
+            Desired energy.
+        pol : Pol
+            Polarisation mode.
+
+        Returns:
+        ----------
+        float
+            gap / phase motor position from the lookup table.
+        """
+        # if lut is empty, force an update to pull updated lut incase subclasses have
+        # implemented it.
+        if not self.lut.root:
+            self.update_lookup_table()
+        poly = self.lut.get_poly(energy=energy, pol=pol)
+        return poly(energy)
+
+
+class ConfigServerEnergyMotorLookup(EnergyMotorLookup):
+    """Fetches and parses lookup table (csv) from a config server, supports dynamic
+    updates, and validates input."""
+
+    def __init__(
+        self,
+        config_client: ConfigServer,
+        lut_config: LookupTableColumnConfig,
+        path: Path,
+    ):
+        """
+        Parameters:
+        -----------
+        config_client:
+            The config server client to fetch the look up table data.
+        lut_config:
+            Configuration that defines how to process file contents into a LookupTable
+        path:
+            File path to the lookup table.
+        """
+        self.path = path
+        self.config_client = config_client
+        self.lut_config = lut_config
+        super().__init__()
+
+    def read_lut(self) -> LookupTable:
+        file_contents = self.config_client.get_file_contents(
+            self.path, reset_cached_result=True
+        )
+        return convert_csv_to_lookup(file_contents, lut_config=self.lut_config)
+
+    def update_lookup_table(self) -> None:
+        self.lut = self.read_lut()

dodal/devices/insertion_device/id_enum.py

@@ -0,0 +1,17 @@
+from ophyd_async.core import StrictEnum
+
+
+class Pol(StrictEnum):
+    NONE = "None"
+    LH = "lh"
+    LV = "lv"
+    PC = "pc"
+    NC = "nc"
+    LA = "la"
+    LH3 = "lh3"
+    LV3 = "lv3"
+
+
+class UndulatorGateStatus(StrictEnum):
+    OPEN = "Open"
+    CLOSE = "Closed"