add gap to energy coverter to controller

src/dodal/beamlines/i06_shared.py

@@ -1,11 +1,14 @@
 from dodal.common.beamlines.beamline_utils import set_beamline as set_utils_beamline
 from dodal.device_manager import DeviceManager
-from dodal.devices.beamlines.i06_shared import I06Grating
+from dodal.devices.beamlines.i06_shared import I06EpicsPolynomialDevice, I06Grating
+from dodal.devices.beamlines.i06_shared.i06_apple2_controller import I06Apple2Controller
 from dodal.devices.insertion_device import (
     Apple2,
+    InsertionDevicePolarisation,
     UndulatorGap,
     UndulatorLockedPhaseAxes,
 )
+from dodal.devices.insertion_device.energy import InsertionDeviceEnergy
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.devices.synchrotron import Synchrotron
 from dodal.log import set_beamline as set_log_beamline
@@ -24,6 +27,22 @@ def synchrotron() -> Synchrotron:
     return Synchrotron()
 
 
+@devices.factory()
+def pgm() -> PlaneGratingMonochromator:
+    return PlaneGratingMonochromator(
+        prefix=f"{PREFIX.beamline_prefix}-OP-PGM-01:",
+        grating=I06Grating,
+        grating_pv="NLINES2",
+    )
+
+
+# Insertion Device
+# ------------- Downstream Insertion Device --------------------
+@devices.factory()
+def i06_idd_epics_polynomial_device() -> I06EpicsPolynomialDevice:
+    return I06EpicsPolynomialDevice(prefix=f"{PREFIX.beamline_prefix}-OP-IDD-01:")
+
+
 @devices.factory()
 def idd_gap() -> UndulatorGap:
     return UndulatorGap(prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-01:")
@@ -44,6 +63,25 @@ def idd(idd_gap: UndulatorGap, idd_phase: UndulatorLockedPhaseAxes) -> Apple2:
     return Apple2(id_gap=idd_gap, id_phase=idd_phase)
 
 
+@devices.factory()
+def idd_controller(
+    idd: Apple2, i06_idd_epics_polynomial_device: I06EpicsPolynomialDevice
+) -> I06Apple2Controller:
+    """I06 downstream insertion device controller."""
+    return I06Apple2Controller(
+        apple2=idd,
+        gap_energy_motor_lut=i06_idd_epics_polynomial_device.energy_gap_motor_lookup,
+        phase_energy_motor_lut=i06_idd_epics_polynomial_device.energy_phase_motor_lookup,  # need fix this too
+        gap_motor_energy_lut=i06_idd_epics_polynomial_device.gap_motor_energy_lookup,
+    )
+
+
+# -------------------- Upstream Insertion Device -------------------
+@devices.factory()
+def i06_idu_epics_polynomial_device() -> I06EpicsPolynomialDevice:
+    return I06EpicsPolynomialDevice(prefix=f"{PREFIX.beamline_prefix}-OP-IDU-01:")
+
+
 @devices.factory()
 def idu_gap() -> UndulatorGap:
     return UndulatorGap(prefix=f"{PREFIX.insertion_prefix}-MO-SERVC-21:")
@@ -65,9 +103,25 @@ def idu(idu_gap: UndulatorGap, idu_phase: UndulatorLockedPhaseAxes) -> Apple2:
 
 
 @devices.factory()
-def pgm() -> PlaneGratingMonochromator:
-    return PlaneGratingMonochromator(
-        prefix=f"{PREFIX.beamline_prefix}-OP-PGM-01:",
-        grating=I06Grating,
-        grating_pv="NLINES2",
+def idu_controller(
+    idu: Apple2, i06_idu_epics_polynomial_device: I06EpicsPolynomialDevice
+) -> I06Apple2Controller:
+    """I06 upstream insertion device controller."""
+    return I06Apple2Controller(
+        apple2=idu,
+        gap_energy_motor_lut=i06_idu_epics_polynomial_device.energy_gap_motor_lookup,
+        phase_energy_motor_lut=i06_idu_epics_polynomial_device.energy_phase_motor_lookup,
+        gap_motor_energy_lut=i06_idu_epics_polynomial_device.gap_motor_energy_lookup,
     )
+
+
+@devices.factory()
+def idu_energy(idu_controller: I06Apple2Controller) -> InsertionDeviceEnergy:
+    return InsertionDeviceEnergy(id_controller=idu_controller)
+
+
+@devices.factory()
+def idu_polarisation(
+    idu_controller: I06Apple2Controller,
+) -> InsertionDevicePolarisation:
+    return InsertionDevicePolarisation(id_controller=idu_controller)

src/dodal/devices/beamlines/i05_shared/apple_knot_constants.py

@@ -70,23 +70,23 @@
 )
 
 
-def energy_to_gap_converter(energy: float, pol: Pol) -> float:
+def energy_to_gap_converter(value: float, pol: Pol) -> float:
     if pol == Pol.LH:
-        return float(LH_GAP_POLYNOMIAL(energy))
+        return float(LH_GAP_POLYNOMIAL(value))
     if pol == Pol.LV:
-        return float(LV_GAP_POLYNOMIAL(energy))
+        return float(LV_GAP_POLYNOMIAL(value))
     if pol == Pol.PC or pol == Pol.NC:
-        return float(C_GAP_POLYNOMIAL(energy))
+        return float(C_GAP_POLYNOMIAL(value))
     return 0.0
 
 
-def energy_to_phase_converter(energy: float, pol: Pol) -> float:
+def energy_to_phase_converter(value: float, pol: Pol) -> float:
     if pol == Pol.LH:
         return 0.0
     if pol == Pol.LV:
         return 70.0
     if pol == Pol.PC:
-        return float(C_PHASE_POLYNOMIAL(energy))
+        return float(C_PHASE_POLYNOMIAL(value))
     if pol == Pol.NC:
-        return -float(C_PHASE_POLYNOMIAL(energy))
+        return -float(C_PHASE_POLYNOMIAL(value))
     return 0.0

src/dodal/devices/beamlines/i06_shared/__init__.py

@@ -1,3 +1,5 @@
+from .i06_apple2_controller import I06Apple2Controller
 from .i06_enum import I06Grating
+from .i06_epics_lookup import I06EpicsPolynomialDevice
 
-__all__ = ["I06Grating"]
+__all__ = ["I06Grating", "I06EpicsPolynomialDevice", "I06Apple2Controller"]

src/dodal/devices/beamlines/i06_shared/i06_apple2_controller.py

@@ -0,0 +1,61 @@
+from dodal.devices.insertion_device import (
+    Apple2,
+    Apple2Controller,
+    Apple2PhasesVal,
+    Apple2Val,
+    Pol,
+    UndulatorPhaseAxes,
+)
+from dodal.devices.insertion_device.energy_motor_lookup import EnergyMotorLookup
+
+ROW_PHASE_MOTOR_TOLERANCE = 0.004
+MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
+MAXIMUM_GAP_MOTOR_POSITION = 100
+DEFAULT_JAW_PHASE_POLY_PARAMS = [1.0 / 7.5, -120.0 / 7.5]
+
+
+class I06Apple2Controller(Apple2Controller[Apple2[UndulatorPhaseAxes]]):
+    """I17Apple2Controller is a extension of Apple2Controller which provide linear
+    arbitrary angle control.
+
+    Args:
+        apple2 (Apple2): An Apple2 device.
+        gap_energy_motor_lut (EnergyMotorLookup): The class that handles the gap
+            look up table logic for the insertion device.
+        phase_energy_motor_lut (EnergyMotorLookup): The class that handles the phase
+            look up table logic for the insertion device.
+        units (str, optional): The units of this device. Defaults to eV.
+        name (str, optional): New device name.
+    """
+
+    def __init__(
+        self,
+        apple2: Apple2[UndulatorPhaseAxes],
+        gap_energy_motor_lut: EnergyMotorLookup,
+        phase_energy_motor_lut: EnergyMotorLookup,
+        gap_motor_energy_lut: EnergyMotorLookup,
+        units: str = "eV",
+        name: str = "",
+    ) -> None:
+        self.gap_energy_motor_lut = gap_energy_motor_lut
+        self.phase_energy_motor_lut = phase_energy_motor_lut
+        self.gap_motor_energy_lut = gap_motor_energy_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,
+            gap_motor_energy_converter=gap_motor_energy_lut.find_value_in_lookup_table,
+            units=units,
+            name=name,
+        )
+
+    def _get_apple2_value(self, gap: float, phase: float, pol: Pol) -> Apple2Val:
+        return Apple2Val(
+            gap=gap,
+            phase=Apple2PhasesVal(
+                top_outer=phase,
+                top_inner=0.0,
+                btm_inner=phase,
+                btm_outer=0.0,
+            ),
+        )

src/dodal/devices/beamlines/i06_shared/i06_epics_lookup.py

@@ -0,0 +1,141 @@
+import asyncio
+
+from bluesky.protocols import Triggerable
+from ophyd_async.core import AsyncStatus, Device, DeviceMock, DeviceVector
+from ophyd_async.epics.core import epics_signal_rw
+
+from dodal.device_manager import DEFAULT_TIMEOUT
+from dodal.devices.insertion_device import (
+    EnergyCoverage,
+    EnergyMotorLookup,
+    LookupTable,
+    Pol,
+)
+from dodal.log import LOGGER
+
+ROW_PHASE_CIRCULAR = 15.0
+MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
+DEFAULT_POLY1D_PARAMETERS = {
+    Pol.LH: [0],
+    Pol.LV: [MAXIMUM_ROW_PHASE_MOTOR_POSITION],
+    Pol.PC: [ROW_PHASE_CIRCULAR],
+    Pol.PC3: [ROW_PHASE_CIRCULAR],
+    Pol.NC: [-ROW_PHASE_CIRCULAR],
+    Pol.NC3: [-ROW_PHASE_CIRCULAR],
+}
+
+
+class I06EpicsPolynomialDevice(Device, Triggerable):
+    def __init__(
+        self,
+        prefix: str,
+        max_energy: float = 2200,
+        min_energy: float = 70,
+        name: str = "",
+    ) -> None:
+        # Define mapping of polarization to PV suffix
+        self._pol_map = {
+            Pol.LH: "HZ",
+            Pol.LV: "VT",
+            Pol.PC: "PC",
+            Pol.NC: "NC",
+            Pol.PC3: "PC:HAR3",
+            Pol.NC3: "NC:HAR3",
+        }
+        self._inv_pol_map = {
+            Pol.LH: "BHZ",
+            Pol.LV: "BVT",
+            Pol.PC: "BPC",
+            Pol.NC: "BNC",
+            Pol.PC3: "BPC:HAR3",
+            Pol.NC3: "BNC:HAR3",
+        }
+        self.param_dict = {}
+        self.inv_param_dict = {}
+        self.min_energy = min_energy
+        self.max_energy = max_energy
+        # Initialize DeviceVectors
+        for pol, suffix in self._inv_pol_map.items():
+            attr_name = f"{pol.name.lower()}_inverse_params"
+            setattr(self, attr_name, self._make_params(f"{prefix}{suffix}"))
+            self.inv_param_dict[pol] = getattr(self, attr_name)
+
+        for pol, suffix in self._pol_map.items():
+            attr_name = f"{pol.name.lower()}_params"
+            setattr(self, attr_name, self._make_params(f"{prefix}{suffix}"))
+            self.param_dict[pol] = getattr(self, attr_name)
+        self.energy_gap_motor_lookup = EnergyMotorLookup()
+        self.energy_phase_motor_lookup = EnergyMotorLookup()
+        self.gap_motor_energy_lookup = EnergyMotorLookup()
+        super().__init__(name=name)
+
+    def _make_params(self, pv_prefix: str) -> DeviceVector:
+        return DeviceVector(
+            {
+                i: epics_signal_rw(float, read_pv=f"{pv_prefix}:C{i}")
+                for i in range(12, 0, -1)
+            }
+        )
+
+    @AsyncStatus.wrap
+    async def trigger(self) -> None:
+        """Triggering this device will update the lookup tables with the current PV values."""
+        await self.update_lookup()
+
+    async def _get_table_entries(
+        self,
+        param_dict: dict[Pol, DeviceVector | list[float]],
+        min_energy: float,
+        max_energy: float,
+    ) -> dict[Pol, EnergyCoverage]:
+        entries = {}
+        for pol, vector in param_dict.items():
+            if isinstance(vector, DeviceVector):
+                coeffs = await asyncio.gather(*(p.get_value() for p in vector.values()))
+            else:
+                coeffs = vector
+            entries[pol] = EnergyCoverage.generate(
+                min_energies=[min_energy],
+                max_energies=[max_energy],
+                poly1d_params=[coeffs],
+            )
+        return entries
+
+    async def update_lookup(self) -> None:
+        # Update gap lookup table
+        energy_entries = await self._get_table_entries(
+            self.param_dict, self.min_energy, self.max_energy
+        )
+        self.energy_gap_motor_lookup = EnergyMotorLookup(LookupTable(energy_entries))
+        # find gap range from energy range
+        min_gap = self.energy_gap_motor_lookup.find_value_in_lookup_table(
+            value=self.max_energy, pol=Pol.LH
+        )
+        max_gap = self.energy_gap_motor_lookup.find_value_in_lookup_table(
+            value=self.min_energy, pol=Pol.LH
+        )
+        # Update gap inverse lookup table
+        inv_energy_entries = await self._get_table_entries(
+            self.inv_param_dict, max_energy=max_gap, min_energy=min_gap
+        )
+        self.gap_motor_energy_lookup = EnergyMotorLookup(
+            LookupTable(inv_energy_entries)
+        )
+        # Update phase lookup table
+        energy_entries = await self._get_table_entries(
+            DEFAULT_POLY1D_PARAMETERS, max_energy=max_gap, min_energy=min_gap
+        )
+        self.energy_phase_motor_lookup = EnergyMotorLookup(LookupTable(energy_entries))
+        LOGGER.info("Updating lookup tables with new values from EPICS.")
+
+    async def connect(
+        self,
+        mock: bool | DeviceMock = False,
+        timeout: float = DEFAULT_TIMEOUT,
+        force_reconnect: bool = False,
+    ) -> None:
+        await super().connect(
+            mock=mock, timeout=timeout, force_reconnect=force_reconnect
+        )
+        # Highjack connect to update lookup tables on connection.
+        await self.update_lookup()

src/dodal/devices/insertion_device/__init__.py

@@ -30,6 +30,7 @@
 from .enum import Pol, UndulatorGateStatus
 from .lookup_table_models import (
     EnergyCoverage,
+    EnergyCoverageEntry,
     LookupTable,
     LookupTableColumnConfig,
     convert_csv_to_lookup,
@@ -67,4 +68,5 @@
     "EnergyMotorConvertor",
     "UnstoppableMotor",
     "InsertionDeviceEnergyBase",
+    "EnergyCoverageEntry",
 ]