Merge branch 'main' into 788_pmac-use-kickoff

pyproject.toml

@@ -15,7 +15,7 @@ description = "Ophyd devices and other utils that could be used across DLS beaml
 dependencies = [
     "click",
     "ophyd",
-    "ophyd-async>=0.5.2",
+    "ophyd-async>=0.6,<0.7",
     "bluesky",
     "pyepics",
     "dataclasses-json",

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/i24/pmac.py

@@ -3,9 +3,9 @@
 
 from bluesky.protocols import Flyable, Triggerable
 from ophyd_async.core import (
+    CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
     AsyncStatus,
-    CalculateTimeout,
     SignalBackend,
     SignalR,
     SignalRW,
@@ -90,7 +90,7 @@ async def set(
         self,
         value: LaserSettings,
         wait=True,
-        timeout=CalculateTimeout,
+        timeout=CALCULATE_TIMEOUT,
     ):
         await self.signal.set(value.value, wait, timeout)
 
@@ -113,7 +113,7 @@ async def set(
         self,
         value: EncReset,
         wait=True,
-        timeout=CalculateTimeout,
+        timeout=CALCULATE_TIMEOUT,
     ):
         await self.signal.set(value.value, wait, timeout)
 

src/dodal/devices/tetramm.py

@@ -3,13 +3,13 @@
 
 from bluesky.protocols import Hints
 from ophyd_async.core import (
-    AsyncStatus,
     DatasetDescriber,
     DetectorControl,
     DetectorTrigger,
     Device,
     PathProvider,
     StandardDetector,
+    TriggerInfo,
     set_and_wait_for_value,
     soft_signal_r_and_setter,
 )
@@ -113,29 +113,24 @@ def get_deadtime(self, exposure: float | None) -> float:
         # 2 internal clock cycles. Best effort approximation
         return 2 / self.base_sample_rate
 
-    async def arm(
-        self,
-        num: int,
-        trigger: DetectorTrigger = DetectorTrigger.edge_trigger,
-        exposure: float | None = None,
-    ) -> AsyncStatus:
-        if exposure is None:
-            raise ValueError(
-                "Tetramm does not support arm without exposure time. "
-                "Is this a software scan? Tetramm only supports hardware scans."
-            )
-        self._validate_trigger(trigger)
+    async def prepare(self, trigger_info: TriggerInfo):
+        self._validate_trigger(trigger_info.trigger)
+        assert trigger_info.livetime is not None
 
         # trigger mode must be set first and on its own!
         await self._drv.trigger_mode.set(TetrammTrigger.ExtTrigger)
 
         await asyncio.gather(
-            self._drv.averaging_time.set(exposure), self.set_exposure(exposure)
+            self._drv.averaging_time.set(trigger_info.livetime),
+            self.set_exposure(trigger_info.livetime),
         )
 
-        status = await set_and_wait_for_value(self._drv.acquire, True)
+    async def arm(self):
+        self._arm_status = await set_and_wait_for_value(self._drv.acquire, True)
 
-        return status
+    async def wait_for_idle(self):
+        if self._arm_status:
+            await self._arm_status
 
     def _validate_trigger(self, trigger: DetectorTrigger) -> None:
         supported_trigger_types = {

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_gridscan.py

@@ -1,4 +1,4 @@
-from asyncio import wait_for
+from asyncio import TimeoutError, wait_for
 from dataclasses import dataclass
 
 import numpy as np