Create ConfigServerEnergyMotorLookup and decouple gap and phase (#1733)

* add i09 look up table praiser

* move get_poly to lookuptable

* extracted EnergyMotorLookup base class from i10EnergyMotorLookup

* added new make_phase_tables function

* add test for correct table output

* remove_i09

* add docstring

* add test for helper functions

* add test for skipping

* group and reuse fixture, add J09 controller and tests

* there is no reason to set mock value to zero

* spacing fix

* remove print

* add gap tests

* add controller, energy and polarisation to i09_2 configuration

* move path to top

* add model_validation

* remove debug

* Update docstring for energy_jid function

* Update nc parameter and phase calculation logic

Change 'nc' parameter to negative ROW_PHASE_CIRCULAR and adjust phase calculation based on pol.

* Refactor Lookuptable class documentation

Removed outdated docstring from Lookuptable class and updated initialization docstring for I10EnergyMotorLookup class.

* update expected lookup table after sign change

* change lookup table schema to snake case

* replace dictionary with basemodel

* add tying

* add gap and phase

* fat finger correction

* Update ID lookup logic to use type checking

* Fix some tests

* Improve models to not use shared defaults

* undo syntax error

* Fixed test to check for success on loading i10 lut

* Simplified lut logic

* Added back generate_lookup_table function

* Fixed all tests but polarisation

* Updated doc strings

* Moved generic test from i10 to test_lookup_table_apple2

* Renamed lut_column_config to lut_config

* Updated more doc strings

* Fix default phase_file name and thus tests

* Use Pol in LookupTable rather than string

* Updated test_convert_csv_to_lookup_overwrite_name_convert_default to use Pol

* Update tests to use Pol rather str value

* fix tests

* fix i09 controller and test

* reuse fixture

* Improve EnergyCoverageEntry to have a poly serializer and EnergyCoverage to use float as key

* Add type checking to i10Apple2 phase

* Removed commente out code

* Update i10 id tests to use json files rather than pickle files so they are human readable

* Remove comments

* Fixed poly test

* Fixed test_make_phase_tables_multiple_entries

* Added test_lookup_table_is_serialisable

* Update src/dodal/devices/util/lookup_tables_apple2.py

Co-authored-by: Raymond Fan <[email protected]>

* Fixed formatting

* Improved code coverage

* Simplified ID lookup table logic

* Updated doc strings, tidy tests and variable names

* Made default files a constant

* added gap and phase lookup

* Update logging messages

* Removed duplicate test logic

* Decoupled path from LookupTableConfig, updated convert_csv_to_lookup to use file_contents again

* Updated doc strings

* Clean up imports

* fixed test

* update and reuse fixture

* change I09 to use json data

* fix controller

* add test for Path not given

* correct test

* typing

* modify _setpol to go to lh first for i09 id

* Rename class and update comments for clarity

* Remove POLY_DEG list from test_i09_apple2.py

Removed the POLY_DEG list from the test file.

* fix lint

* remove I09EnergyMotorLookup and move phase generation into EnergyMotorlookup

* move lookup table path into lut_config

* revert lookuptable with path

* introduce spacing to data

* Update src/dodal/devices/util/lookup_tables_apple2.py

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

* Update src/dodal/devices/util/lookup_tables_apple2.py

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

* remove J09defaultlookuptable

* remove default in test

* lint

* update ophyd

* Apply suggestions from code review

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

* fix typo

* fix file name and test

* Create AbstractEnergyMotorLookup

* Moved apple2 fixture to dodal.testing.fixtures.devices.apple2

* Fix test

* Remove init for DummyEnergyMotorLookup

* Removed unneeded pytest plugins

* Adjust i17 implementation

* Fix tests

* Add PhaseAxes type to Apple2 for i17

* Update lookup_tables_apple2 doc strings

* Renamed update_lut to update_lookup_table

* Further improve type checking for i09 and i17 apple2

* Reworked GeneratePoly1DFromFileEnergyMotorLookup to ConfiguredEnergyMotorLookup. Update i09 to use phase lookup table

* Updated variable names for generate_lookup_table to be lut

* Restructured controller logic to reuse more components

* Updated doc strings and function names

* Add missing types to EnergyMotorLookup

* Correct return type for process_row

* Remove ConfiguredEnergyMotorLookup and changed this to EnergyMotorLookup, removed abstract

* Rename _id_set_value to _get_apple2_value

* Renamed get_motor_from_energy to find_value_in_lookup_table

* Remove available_pol from EnergyMotorLookup

* Fix tests

* Renamed parameters from feedback

* Renamed configured_energy_motor_lookup to energy_motor_lookup

* Removed pol arg from _set_apple2, i10 now extends _set_motors_from_energy_and_polarisation instead

* Add csv to ConfigServerEnergyMotorLookup doc stirng

* Update ConfigServerEnergyMotorLookup __init__ doc string to remove reference to gap

* Removed unneeded test setup

* Fixed test

* Updated docs for Apple2Controller

* Removed _set_apple2 as no longer needed

---------

Co-authored-by: Relm-Arrowny <[email protected]>

Author: oliwenmandiamond

src/dodal/beamlines/i09_2.py

@@ -16,19 +16,23 @@
 )
 from dodal.devices.i09.enums import Grating
 from dodal.devices.i09_2_shared.i09_apple2 import (
-    J09_POLY_DEG,
-    EnergyMotorLookup,
+    J09_GAP_POLY_DEG_COLUMNS,
+    J09_PHASE_POLY_DEG_COLUMNS,
     J09Apple2Controller,
 )
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.devices.synchrotron import Synchrotron
-from dodal.devices.util.lookup_tables_apple2 import LookupTableConfig
+from dodal.devices.util.lookup_tables_apple2 import (
+    ConfigServerEnergyMotorLookup,
+    LookupTableConfig,
+)
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
 J09_CONF_CLIENT = ConfigServer(url="https://daq-config.diamond.ac.uk")
 LOOK_UPTABLE_DIR = "/dls_sw/i09-2/software/gda/workspace_git/gda-diamond.git/configurations/i09-2-shared/lookupTables/"
 GAP_LOOKUP_FILE_NAME = "JIDEnergy2GapCalibrations.csv"
+PHASE_LOOKUP_FILE_NAME = "JIDEnergy2PhaseCalibrations.csv"
 
 BL = get_beamline_name("i09-2")
 PREFIX = BeamlinePrefix(BL, suffix="J")
@@ -79,10 +83,15 @@ def jid_controller() -> J09Apple2Controller:
     """J09 insertion device controller."""
     return J09Apple2Controller(
         apple2=jid(),
-        energy_motor_lut=EnergyMotorLookup(
-            lut_config=LookupTableConfig(poly_deg=J09_POLY_DEG),
+        gap_energy_motor_lut=ConfigServerEnergyMotorLookup(
+            lut_config=LookupTableConfig(poly_deg=J09_GAP_POLY_DEG_COLUMNS),
+            config_client=J09_CONF_CLIENT,
+            path=Path(LOOK_UPTABLE_DIR, GAP_LOOKUP_FILE_NAME),
+        ),
+        phase_energy_motor_lut=ConfigServerEnergyMotorLookup(
+            lut_config=LookupTableConfig(poly_deg=J09_PHASE_POLY_DEG_COLUMNS),
             config_client=J09_CONF_CLIENT,
-            gap_path=Path(LOOK_UPTABLE_DIR, GAP_LOOKUP_FILE_NAME),
+            path=Path(LOOK_UPTABLE_DIR, PHASE_LOOKUP_FILE_NAME),
         ),
     )
 

src/dodal/beamlines/i10_optics.py

@@ -39,7 +39,7 @@
 from dodal.devices.util.lookup_tables_apple2 import (
     DEFAULT_GAP_FILE,
     DEFAULT_PHASE_FILE,
-    EnergyMotorLookup,
+    ConfigServerEnergyMotorLookup,
     LookupTableConfig,
 )
 from dodal.log import set_beamline as set_log_beamline
@@ -50,8 +50,6 @@
 set_utils_beamline(BL)
 PREFIX = BeamlinePrefix(BL)
 
-LOOK_UPTABLE_DIR = "/dls_sw/i10/software/blueapi/scratch/i10-config/lookupTables/"
-
 
 @device_factory()
 def synchrotron() -> Synchrotron:
@@ -123,13 +121,22 @@ def idd() -> I10Apple2:
 @device_factory()
 def idd_controller() -> I10Apple2Controller:
     """I10 downstream insertion device controller."""
-    idd_energy_motor_lut = EnergyMotorLookup(
+    source = ("Source", "idd")
+    idd_gap_energy_motor_lut = ConfigServerEnergyMotorLookup(
+        config_client=I10_CONF_CLIENT,
+        lut_config=LookupTableConfig(source=source),
+        path=Path(LOOK_UPTABLE_DIR, DEFAULT_GAP_FILE),
+    )
+    idd_phase_energy_motor_lut = ConfigServerEnergyMotorLookup(
         config_client=I10_CONF_CLIENT,
-        lut_config=LookupTableConfig(source=("Source", "idd")),
-        gap_path=Path(LOOK_UPTABLE_DIR, DEFAULT_GAP_FILE),
-        phase_path=Path(LOOK_UPTABLE_DIR, DEFAULT_PHASE_FILE),
+        lut_config=LookupTableConfig(source=source),
+        path=Path(LOOK_UPTABLE_DIR, DEFAULT_PHASE_FILE),
+    )
+    return I10Apple2Controller(
+        apple2=idd(),
+        gap_energy_motor_lut=idd_gap_energy_motor_lut,
+        phase_energy_motor_lut=idd_phase_energy_motor_lut,
     )
-    return I10Apple2Controller(apple2=idd(), energy_motor_lut=idd_energy_motor_lut)
 
 
 @device_factory()
@@ -188,13 +195,22 @@ def idu() -> I10Apple2:
 @device_factory()
 def idu_controller() -> I10Apple2Controller:
     """I10 upstream insertion device controller."""
-    idu_energy_motor_lut = EnergyMotorLookup(
+    source = ("Source", "idu")
+    idu_gap_energy_motor_lut = ConfigServerEnergyMotorLookup(
         config_client=I10_CONF_CLIENT,
-        lut_config=LookupTableConfig(source=("Source", "idu")),
-        gap_path=Path(LOOK_UPTABLE_DIR, DEFAULT_GAP_FILE),
-        phase_path=Path(LOOK_UPTABLE_DIR, DEFAULT_PHASE_FILE),
+        lut_config=LookupTableConfig(source=source),
+        path=Path(LOOK_UPTABLE_DIR, DEFAULT_GAP_FILE),
+    )
+    idu_phase_energy_motor_lut = ConfigServerEnergyMotorLookup(
+        config_client=I10_CONF_CLIENT,
+        lut_config=LookupTableConfig(source=source),
+        path=Path(LOOK_UPTABLE_DIR, DEFAULT_PHASE_FILE),
+    )
+    return I10Apple2Controller(
+        apple2=idd(),
+        gap_energy_motor_lut=idu_gap_energy_motor_lut,
+        phase_energy_motor_lut=idu_phase_energy_motor_lut,
     )
-    return I10Apple2Controller(apple2=idd(), energy_motor_lut=idu_energy_motor_lut)
 
 
 @device_factory()

src/dodal/beamlines/i17.py

@@ -19,6 +19,10 @@
 from dodal.devices.i17.i17_apple2 import I17Apple2Controller
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.util.lookup_tables_apple2 import (
+    EnergyMotorLookup,
+    LookupTable,
+)
 from dodal.log import set_beamline as set_log_beamline
 from dodal.utils import BeamlinePrefix, get_beamline_name
 
@@ -76,7 +80,9 @@ def id() -> Apple2:
 def id_controller() -> Apple2Controller:
     """I17 insertion device controller with dummy energy to motor_converter."""
     return I17Apple2Controller(
-        apple2=id(), energy_to_motor_converter=lambda energy, pol: (0.0, 0.0)
+        apple2=id(),
+        gap_energy_motor_lut=EnergyMotorLookup(lut=LookupTable()),
+        phase_energy_motor_lut=EnergyMotorLookup(lut=LookupTable()),
     )
 
 

src/dodal/devices/apple2_undulator.py

@@ -398,7 +398,7 @@ async def set(self, id_motor_values: Apple2Val) -> None:
 
 
 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 +426,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 +447,8 @@ 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:
@@ -461,8 +461,9 @@ def __init__(
         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(
@@ -481,10 +482,11 @@ def __init__(
         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)
@@ -496,24 +498,35 @@ def __init__(
                 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:
@@ -525,7 +538,6 @@ async def _check_and_get_pol_setpoint(self) -> Pol:
         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:

src/dodal/devices/i09_2_shared/i09_apple2.py

@@ -5,13 +5,11 @@
     Apple2PhasesVal,
     Apple2Val,
     Pol,
+    UndulatorPhaseAxes,
 )
-from dodal.devices.util.lookup_tables_apple2 import (
-    EnergyMotorLookup,
-)
-from dodal.log import LOGGER
+from dodal.devices.util.lookup_tables_apple2 import EnergyMotorLookup
 
-J09_POLY_DEG = [
+J09_GAP_POLY_DEG_COLUMNS = [
     "9th-order",
     "8th-order",
     "7th-order",
@@ -24,31 +22,44 @@
     "0th-order",
 ]
 
+J09_PHASE_POLY_DEG_COLUMNS = ["0th-order"]
+
 
-class J09Apple2Controller(Apple2Controller[Apple2]):
+class J09Apple2Controller(Apple2Controller[Apple2[UndulatorPhaseAxes]]):
     def __init__(
         self,
-        apple2: Apple2,
-        energy_motor_lut: EnergyMotorLookup,
+        apple2: Apple2[UndulatorPhaseAxes],
+        gap_energy_motor_lut: EnergyMotorLookup,
+        phase_energy_motor_lut: EnergyMotorLookup,
         units: str = "keV",
         name: str = "",
     ) -> None:
-        self.lookup_table_client = energy_motor_lut
+        """
+        Parameters:
+        -----------
+        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:
+            the units of this device. Defaults to eV.
+        name : str, optional
+            New device name.
+        """
+        self.gap_energy_motor_lut = gap_energy_motor_lut
+        self.phase_energy_motor_lut = phase_energy_motor_lut
         super().__init__(
             apple2=apple2,
-            energy_to_motor_converter=self.lookup_table_client.get_motor_from_energy,
+            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,
         )
 
-    async def _set_motors_from_energy(self, value: float) -> None:
-        """
-        Set the undulator motors for a given energy and polarisation.
-        """
-
-        pol = await self._check_and_get_pol_setpoint()
-        gap, phase = self.energy_to_motor(energy=value, pol=pol)
-        id_set_val = Apple2Val(
+    def _get_apple2_value(self, gap: float, phase: float, pol: Pol) -> Apple2Val:
+        return Apple2Val(
             gap=f"{gap:.6f}",
             phase=Apple2PhasesVal(
                 top_outer=f"{phase:.6f}",
@@ -58,9 +69,6 @@ async def _set_motors_from_energy(self, value: float) -> None:
             ),
         )
 
-        LOGGER.info(f"Setting polarisation to {pol}, with values: {id_set_val}")
-        await self.apple2().set(id_motor_values=id_set_val)
-
     async def _set_pol(
         self,
         value: Pol,
@@ -70,7 +78,7 @@ async def _set_pol(
         if value is not Pol.LH:
             self._polarisation_setpoint_set(Pol.LH)
             max_lh_energy = float(
-                self.lookup_table_client.lookup_tables.gap.root[Pol.LH].limit.maximum
+                self.gap_energy_motor_lut.lut.root[Pol.LH].limit.maximum
             )
             lh_setpoint = (
                 max_lh_energy if target_energy > max_lh_energy else target_energy