Add hard undulator energy and hard energy classes (#1713)

* add classes

* add beamline objects

* remove lut object parameter for now

* add energy_readback and tests

* fix readables in hardEnergy class

* add hard energy tests

* remove lut from i09_1

* rename module

* fix typo

* and reverse function for energy from gap calculation and update tests

* fix tests

* add test

* test

* test2

* test_win_laptop_commit

* implement review  suggestions

* mainly remove names parameter from devices, change undulator PV

* added issue for LUT

---------

Co-authored-by: eir17846 <[email protected]>

Author: Villtord

src/dodal/beamlines/i09_1.py

@@ -10,6 +10,11 @@
 from dodal.devices.electron_analyser import EnergySource
 from dodal.devices.electron_analyser.specs import SpecsDetector
 from dodal.devices.i09_1 import LensMode, PsuMode
+from dodal.devices.i09_1_shared.hard_energy import HardEnergy, HardInsertionDeviceEnergy
+from dodal.devices.i09_1_shared.hard_undulator_functions import (
+    calculate_energy_i09_hu,
+    calculate_gap_i09_hu,
+)
 from dodal.devices.synchrotron import Synchrotron
 from dodal.devices.undulator import UndulatorInMm, UndulatorOrder
 from dodal.log import set_beamline as set_log_beamline
@@ -52,9 +57,28 @@ def analyser() -> SpecsDetector[LensMode, PsuMode]:
 
 @device_factory()
 def undulator() -> UndulatorInMm:
-    return UndulatorInMm(prefix=f"{PREFIX.beamline_prefix}-MO-UND-01:")
+    return UndulatorInMm(prefix=f"{BeamlinePrefix(BL).insertion_prefix}-MO-SERVC-01:")
 
 
 @device_factory()
 def harmonics() -> UndulatorOrder:
-    return UndulatorOrder(name="harmonics")
+    return UndulatorOrder()
+
+
+@device_factory()
+def hu_id_energy() -> HardInsertionDeviceEnergy:
+    return HardInsertionDeviceEnergy(
+        undulator_order=harmonics(),
+        undulator=undulator(),
+        lut={},  # ToDo https://github.com/DiamondLightSource/sm-bluesky/issues/239
+        gap_to_energy_func=calculate_energy_i09_hu,
+        energy_to_gap_func=calculate_gap_i09_hu,
+    )
+
+
+@device_factory()
+def hu_energy() -> HardEnergy:
+    return HardEnergy(
+        dcm=dcm(),
+        undulator_energy=hu_id_energy(),
+    )

src/dodal/devices/i09_1_shared/__init__.py

@@ -1,7 +1,14 @@
+from .hard_energy import HardEnergy, HardInsertionDeviceEnergy
 from .hard_undulator_functions import (
     calculate_energy_i09_hu,
     calculate_gap_i09_hu,
     get_hu_lut_as_dict,
 )
 
-__all__ = ["calculate_gap_i09_hu", "get_hu_lut_as_dict", "calculate_energy_i09_hu"]
+__all__ = [
+    "calculate_gap_i09_hu",
+    "get_hu_lut_as_dict",
+    "calculate_energy_i09_hu",
+    "HardInsertionDeviceEnergy",
+    "HardEnergy",
+]

src/dodal/devices/i09_1_shared/hard_energy.py

@@ -0,0 +1,112 @@
+from asyncio import gather
+from collections.abc import Callable
+
+from bluesky.protocols import Locatable, Location, Movable
+from numpy import ndarray
+from ophyd_async.core import (
+    AsyncStatus,
+    Reference,
+    StandardReadable,
+    StandardReadableFormat,
+    derived_signal_rw,
+    soft_signal_rw,
+)
+
+from dodal.devices.common_dcm import DoubleCrystalMonochromatorBase
+from dodal.devices.i09_1_shared.hard_undulator_functions import (
+    MAX_ENERGY_COLUMN,
+    MIN_ENERGY_COLUMN,
+)
+from dodal.devices.undulator import UndulatorInMm, UndulatorOrder
+
+
+class HardInsertionDeviceEnergy(StandardReadable, Movable[float]):
+    """
+    Compound device to link hard x-ray undulator gap and order to photon energy.
+    Setting the energy adjusts the undulator gap accordingly.
+    """
+
+    def __init__(
+        self,
+        undulator_order: UndulatorOrder,
+        undulator: UndulatorInMm,
+        lut: dict[int, ndarray],
+        gap_to_energy_func: Callable[..., float],
+        energy_to_gap_func: Callable[..., float],
+        name: str = "",
+    ) -> None:
+        self._lut = lut
+        self.gap_to_energy_func = gap_to_energy_func
+        self.energy_to_gap_func = energy_to_gap_func
+        self._undulator_order_ref = Reference(undulator_order)
+        self._undulator_ref = Reference(undulator)
+
+        self.add_readables([undulator_order, undulator.current_gap])
+        with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
+            self.energy_demand = soft_signal_rw(float)
+            self.energy = derived_signal_rw(
+                raw_to_derived=self._read_energy,
+                set_derived=self._set_energy,
+                current_gap=self._undulator_ref().gap_motor.user_readback,
+                current_order=self._undulator_order_ref().value,
+                derived_units="keV",
+            )
+        super().__init__(name=name)
+
+    def _read_energy(self, current_gap: float, current_order: int) -> float:
+        return self.gap_to_energy_func(
+            gap=current_gap,
+            look_up_table=self._lut,
+            order=current_order,
+        )
+
+    async def _set_energy(self, energy: float) -> None:
+        current_order = await self._undulator_order_ref().value.get_value()
+        min_energy, max_energy = self._lut[current_order][
+            MIN_ENERGY_COLUMN : MAX_ENERGY_COLUMN + 1
+        ]
+        if not (min_energy <= energy <= max_energy):
+            raise ValueError(
+                f"Requested energy {energy} keV is out of range for harmonic {current_order}: "
+                f"[{min_energy}, {max_energy}] keV"
+            )
+
+        target_gap = self.energy_to_gap_func(
+            photon_energy_kev=energy, look_up_table=self._lut, order=current_order
+        )
+        await self._undulator_ref().set(target_gap)
+
+    @AsyncStatus.wrap
+    async def set(self, value: float) -> None:
+        self.energy_demand.set(value)
+        await self.energy.set(value)
+
+
+class HardEnergy(StandardReadable, Locatable[float]):
+    """
+    Energy compound device that provides combined change of both DCM energy and undulator gap accordingly.
+    """
+
+    def __init__(
+        self,
+        dcm: DoubleCrystalMonochromatorBase,
+        undulator_energy: HardInsertionDeviceEnergy,
+        name: str = "",
+    ) -> None:
+        self._dcm_ref = Reference(dcm)
+        self._undulator_energy_ref = Reference(undulator_energy)
+        self.add_readables([undulator_energy, dcm.energy_in_keV])
+        super().__init__(name=name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: float) -> None:
+        await gather(
+            self._dcm_ref().energy_in_keV.set(value),
+            self._undulator_energy_ref().set(value),
+        )
+
+    async def locate(self) -> Location[float]:
+        return Location(
+            setpoint=await self._dcm_ref().energy_in_keV.user_setpoint.get_value(),
+            readback=await self._dcm_ref().energy_in_keV.user_readback.get_value(),
+        )