706 make undulator gap writeable i18 (#721)

* rebased

* fix tests

* fix the system test

* restore noqa

* fix the undulator test

* fix i03 undulator instantiation

* fix where the pv is called

* respond to feeedback

* fix undulator set method

* fix the test

* fix test import accesserror

* add i22 fix

* naming reframing

* add issue link

* fix imports from ophyd-async for undulator

* motor import fix

* remove asyncio gather

* remove asyncio gather

src/dodal/beamlines/i03.py

@@ -316,6 +316,7 @@ def undulator(
         wait_for_connection,
         fake_with_ophyd_sim,
         bl_prefix=False,
+        id_gap_lookup_table_path="/dls_sw/i03/software/daq_configuration/lookup/BeamLine_Undulator_toGap.txt",
     )
 
 
@@ -334,7 +335,6 @@ def undulator_dcm(
         undulator=undulator(wait_for_connection, fake_with_ophyd_sim),
         dcm=dcm(wait_for_connection, fake_with_ophyd_sim),
         daq_configuration_path=DAQ_CONFIGURATION_PATH,
-        id_gap_lookup_table_path="/dls_sw/i03/software/daq_configuration/lookup/BeamLine_Undulator_toGap.txt",
     )
 
 

src/dodal/beamlines/i04.py

@@ -308,6 +308,7 @@ def undulator(
         wait_for_connection,
         fake_with_ophyd_sim,
         bl_prefix=False,
+        id_gap_lookup_table_path="/dls_sw/i04/software/gda/config/lookupTables/BeamLine_Undulator_toGap.txt",
     )
 
 

src/dodal/beamlines/i22.py

@@ -199,6 +199,7 @@ def undulator(
         bl_prefix=False,
         poles=80,
         length=2.0,
+        id_gap_lookup_table_path="/dls_sw/i22/software/daq_configuration/lookup/BeamLine_Undulator_toGap.txt",
     )
 
 

src/dodal/devices/undulator.py

@@ -1,20 +1,51 @@
 from enum import Enum
 
-from ophyd_async.core import ConfigSignal, StandardReadable, soft_signal_r_and_setter
+import numpy as np
+from bluesky.protocols import Movable
+from numpy import argmin, ndarray
+from ophyd_async.core import (
+    AsyncStatus,
+    ConfigSignal,
+    StandardReadable,
+    soft_signal_r_and_setter,
+)
 from ophyd_async.epics.motor import Motor
 from ophyd_async.epics.signal import epics_signal_r
 
+from dodal.log import LOGGER
+
+from .util.lookup_tables import energy_distance_table
+
+
+class AccessError(Exception):
+    pass
+
+
+# Enable to allow testing when the beamline is down, do not change in production!
+TEST_MODE = False
+# will be made more generic in https://github.com/DiamondLightSource/dodal/issues/754
+
+
 # The acceptable difference, in mm, between the undulator gap and the DCM
 # energy, when the latter is converted to mm using lookup tables
 UNDULATOR_DISCREPANCY_THRESHOLD_MM = 2e-3
+STATUS_TIMEOUT_S: float = 10.0
 
 
 class UndulatorGapAccess(str, Enum):
     ENABLED = "ENABLED"
     DISABLED = "DISABLED"
 
 
-class Undulator(StandardReadable):
+def _get_closest_gap_for_energy(
+    dcm_energy_ev: float, energy_to_distance_table: ndarray
+) -> float:
+    table = energy_to_distance_table.transpose()
+    idx = argmin(np.abs(table[0] - dcm_energy_ev))
+    return table[1][idx]
+
+
+class Undulator(StandardReadable, Movable):
     """
     An Undulator-type insertion device, used to control photon emission at a given
     beam energy.
@@ -23,6 +54,7 @@ class Undulator(StandardReadable):
     def __init__(
         self,
         prefix: str,
+        id_gap_lookup_table_path: str,
         name: str = "",
         poles: int | None = None,
         length: float | None = None,
@@ -36,6 +68,7 @@ def __init__(
             name (str, optional): Name for device. Defaults to "".
         """
 
+        self.id_gap_lookup_table_path = id_gap_lookup_table_path
         with self.add_children_as_readables():
             self.gap_motor = Motor(prefix + "BLGAPMTR")
             self.current_gap = epics_signal_r(float, prefix + "CURRGAPD")
@@ -63,3 +96,59 @@ def __init__(
                 self.length = None
 
         super().__init__(name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: float):
+        """
+        Set the undulator gap to a given energy in keV
+
+        Args:
+            value: energy in keV
+        """
+        await self._set_undulator_gap(value)
+
+    async def _set_undulator_gap(self, energy_kev: float) -> None:
+        access_level = await self.gap_access.get_value()
+        if access_level is UndulatorGapAccess.DISABLED and not TEST_MODE:
+            raise AccessError("Undulator gap access is disabled. Contact Control Room")
+        LOGGER.info(f"Setting undulator gap to {energy_kev:.2f} kev")
+        target_gap = await self._get_gap_to_match_energy(energy_kev)
+
+        # Check if undulator gap is close enough to the value from the DCM
+        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"
+            )
+            if not TEST_MODE:
+                # Only move if the gap is sufficiently different to the value from the
+                # DCM lookup table AND we're not in TEST_MODE
+                await self.gap_motor.set(
+                    target_gap,
+                    timeout=STATUS_TIMEOUT_S,
+                )
+            else:
+                LOGGER.debug("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"
+            )
+
+    async def _get_gap_to_match_energy(self, energy_kev: float) -> float:
+        """
+        get a 2d np.array from lookup table that
+        converts energies to undulator gap distance
+        """
+        energy_to_distance_table: np.ndarray = await energy_distance_table(
+            self.id_gap_lookup_table_path
+        )
+
+        # Use the lookup table to get the undulator gap associated with this dcm energy
+        return _get_closest_gap_for_energy(
+            energy_kev * 1000,
+            energy_to_distance_table,
+        )

src/dodal/devices/undulator_dcm.py

@@ -1,19 +1,14 @@
 import asyncio
 
-import numpy as np
 from bluesky.protocols import Movable
-from numpy import argmin, ndarray
 from ophyd_async.core import AsyncStatus, StandardReadable
 
 from dodal.common.beamlines.beamline_parameters import get_beamline_parameters
-from dodal.log import LOGGER
 
 from .dcm import DCM
 from .undulator import Undulator, UndulatorGapAccess
-from .util.lookup_tables import energy_distance_table
 
 ENERGY_TIMEOUT_S: float = 30.0
-STATUS_TIMEOUT_S: float = 10.0
 
 # Enable to allow testing when the beamline is down, do not change in production!
 TEST_MODE = False
@@ -23,14 +18,6 @@ class AccessError(Exception):
     pass
 
 
-def _get_closest_gap_for_energy(
-    dcm_energy_ev: float, energy_to_distance_table: ndarray
-) -> float:
-    table = energy_to_distance_table.transpose()
-    idx = argmin(np.abs(table[0] - dcm_energy_ev))
-    return table[1][idx]
-
-
 class UndulatorDCM(StandardReadable, Movable):
     """
     Composite device to handle changing beamline energies, wraps the Undulator and the
@@ -48,7 +35,6 @@ def __init__(
         self,
         undulator: Undulator,
         dcm: DCM,
-        id_gap_lookup_table_path: str,
         daq_configuration_path: str,
         prefix: str = "",
         name: str = "",
@@ -61,11 +47,10 @@ def __init__(
         self.dcm = dcm
 
         # These attributes are just used by hyperion for lookup purposes
-        self.id_gap_lookup_table_path = id_gap_lookup_table_path
-        self.dcm_pitch_converter_lookup_table_path = (
+        self.pitch_energy_table_path = (
             daq_configuration_path + "/lookup/BeamLineEnergy_DCM_Pitch_converter.txt"
         )
-        self.dcm_roll_converter_lookup_table_path = (
+        self.roll_energy_table_path = (
             daq_configuration_path + "/lookup/BeamLineEnergy_DCM_Roll_converter.txt"
         )
         # I03 configures the DCM Perp as a side effect of applying this fixed value to the DCM Offset after an energy change
@@ -78,7 +63,7 @@ def __init__(
     async def set(self, value: float):
         await asyncio.gather(
             self._set_dcm_energy(value),
-            self._set_undulator_gap_if_required(value),
+            self.undulator.set(value),
         )
 
     async def _set_dcm_energy(self, energy_kev: float) -> None:
@@ -90,42 +75,3 @@ async def _set_dcm_energy(self, energy_kev: float) -> None:
             energy_kev,
             timeout=ENERGY_TIMEOUT_S,
         )
-
-    async def _set_undulator_gap_if_required(self, energy_kev: float) -> None:
-        LOGGER.info(f"Setting DCM energy to {energy_kev:.2f} kev")
-        gap_to_match_dcm_energy = await self._gap_to_match_dcm_energy(energy_kev)
-
-        # Check if undulator gap is close enough to the value from the DCM
-        current_gap = await self.undulator.current_gap.get_value()
-        tolerance = await self.undulator.gap_discrepancy_tolerance_mm.get_value()
-        if abs(gap_to_match_dcm_energy - current_gap) > tolerance:
-            LOGGER.info(
-                f"Undulator gap mismatch. {abs(gap_to_match_dcm_energy-current_gap):.3f}mm is outside tolerance.\
-                Moving gap to nominal value, {gap_to_match_dcm_energy:.3f}mm"
-            )
-            if not TEST_MODE:
-                # Only move if the gap is sufficiently different to the value from the
-                # DCM lookup table AND we're not in TEST_MODE
-                await self.undulator.gap_motor.set(
-                    gap_to_match_dcm_energy,
-                    timeout=STATUS_TIMEOUT_S,
-                )
-            else:
-                LOGGER.debug("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"
-            )
-
-    async def _gap_to_match_dcm_energy(self, energy_kev: float) -> float:
-        # Get 2d np.array converting energies to undulator gap distance, from lookup table
-        energy_to_distance_table = await energy_distance_table(
-            self.id_gap_lookup_table_path
-        )
-
-        # Use the lookup table to get the undulator gap associated with this dcm energy
-        return _get_closest_gap_for_energy(
-            energy_kev * 1000,
-            energy_to_distance_table,
-        )

tests/devices/system_tests/test_undulator_system.py

@@ -5,8 +5,15 @@
 
 SIM_INSERTION_PREFIX = "SR03S"
 
+ID_GAP_LOOKUP_TABLE_PATH: str = (
+    "./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt"
+)
+
 
 @pytest.mark.s03
 def test_undulator_connects():
     with DeviceCollector():
-        undulator = Undulator(f"{SIM_INSERTION_PREFIX}-MO-SERVC-01:")  # noqa: F841
+        undulator = Undulator(  # noqa: F841
+            f"{SIM_INSERTION_PREFIX}-MO-SERVC-01:",
+            id_gap_lookup_table_path=ID_GAP_LOOKUP_TABLE_PATH,
+        )

tests/devices/unit_tests/test_undulator.py

@@ -1,13 +1,24 @@
 from unittest.mock import ANY
 
+import numpy as np
 import pytest
 from ophyd_async.core import (
     DeviceCollector,
     assert_configuration,
     assert_reading,
+    set_mock_value,
 )
 
-from dodal.devices.undulator import Undulator, UndulatorGapAccess
+from dodal.devices.undulator import (
+    AccessError,
+    Undulator,
+    UndulatorGapAccess,
+    _get_closest_gap_for_energy,
+)
+
+ID_GAP_LOOKUP_TABLE_PATH: str = (
+    "./tests/devices/unit_tests/test_beamline_undulator_to_gap_lookup_table.txt"
+)
 
 
 @pytest.fixture
@@ -18,6 +29,7 @@ async def undulator() -> Undulator:
             name="undulator",
             poles=80,
             length=2.0,
+            id_gap_lookup_table_path=ID_GAP_LOOKUP_TABLE_PATH,
         )
     return undulator
 
@@ -84,6 +96,7 @@ async def test_poles_not_propagated_if_not_supplied():
             "UND-01",
             name="undulator",
             length=2.0,
+            id_gap_lookup_table_path=ID_GAP_LOOKUP_TABLE_PATH,
         )
     assert undulator.poles is None
     assert "undulator-poles" not in (await undulator.read_configuration())
@@ -95,6 +108,25 @@ async def test_length_not_propagated_if_not_supplied():
             "UND-01",
             name="undulator",
             poles=80,
+            id_gap_lookup_table_path=ID_GAP_LOOKUP_TABLE_PATH,
         )
     assert undulator.length is None
     assert "undulator-length" not in (await undulator.read_configuration())
+
+
+@pytest.mark.parametrize(
+    "energy, expected_output", [(5730, 5.4606), (7200, 6.045), (9000, 6.404)]
+)
+def test_correct_closest_distance_to_energy_from_table(energy, expected_output):
+    energy_to_distance_table = np.array([[5700, 5.4606], [7000, 6.045], [9700, 6.404]])
+    assert (
+        _get_closest_gap_for_energy(energy, energy_to_distance_table) == expected_output
+    )
+
+
+async def test_when_gap_access_is_disabled_set_energy_then_error_is_raised(
+    undulator,
+):
+    set_mock_value(undulator.gap_access, UndulatorGapAccess.DISABLED)
+    with pytest.raises(AccessError):
+        await undulator.set(5)