dls-dodal 1.65.0__py3-none-any.whl → 1.66.0__py3-none-any.whl

dodal/devices/i04/beamsize.py

@@ -0,0 +1,45 @@
+from ophyd_async.core import Reference, derived_signal_r
+
+from dodal.devices.aperturescatterguard import ApertureScatterguard
+from dodal.devices.beamsize.beamsize import BeamsizeBase
+from dodal.devices.i04.transfocator import Transfocator
+
+
+class Beamsize(BeamsizeBase):
+    def __init__(
+        self,
+        transfocator: Transfocator,
+        aperture_scatterguard: ApertureScatterguard,
+        name="",
+    ):
+        super().__init__(name=name)
+        self._transfocator_ref = Reference(transfocator)
+        self._aperture_scatterguard_ref = Reference(aperture_scatterguard)
+
+        with self.add_children_as_readables():
+            self.x_um = derived_signal_r(
+                self._get_beamsize_x,
+                transfocator_size_x=self._transfocator_ref().current_horizontal_size_rbv,
+                aperture_radius=self._aperture_scatterguard_ref().radius,
+                derived_units="µm",
+            )
+            self.y_um = derived_signal_r(
+                self._get_beamsize_y,
+                transfocator_size_y=self._transfocator_ref().current_vertical_size_rbv,
+                aperture_radius=self._aperture_scatterguard_ref().radius,
+                derived_units="µm",
+            )
+
+    def _get_beamsize_x(
+        self,
+        transfocator_size_x: float,
+        aperture_radius: float,
+    ) -> float:
+        return min(transfocator_size_x, aperture_radius)
+
+    def _get_beamsize_y(
+        self,
+        transfocator_size_y: float,
+        aperture_radius: float,
+    ) -> float:
+        return min(transfocator_size_y, aperture_radius)

dodal/devices/i04/murko_results.py

@@ -1,5 +1,6 @@
 import json
 import pickle
+import time
 from dataclasses import dataclass
 from enum import Enum
 from typing import TypedDict
@@ -21,6 +22,7 @@ from dodal.devices.oav.oav_calculations import (
 from dodal.log import LOGGER
 
 NO_MURKO_RESULT = (-1, -1)
+RESULTS_COMPLETE_MESSAGE = "murko_results_complete"
 
 
 class MurkoMetadata(TypedDict):
@@ -71,7 +73,8 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
     solutions for y and z can be calculated using numpy's linear algebra library.
     """
 
-    TIMEOUT_S = 2
+    GET_MESSAGE_TIMEOUT_S = 2
+    RESULTS_COMPLETE_TIMEOUT_S = 5
     PERCENTAGE_TO_USE = 25
     LEFTMOST_PIXEL_TO_USE = 10
     NUMBER_OF_WRONG_RESULTS_TO_LOG = 5
@@ -82,7 +85,6 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
         redis_password=RedisConstants.REDIS_PASSWORD,
         redis_db=RedisConstants.MURKO_REDIS_DB,
         name="",
-        stop_angle=350,
     ):
         self.redis_client = StrictRedis(
             host=redis_host,
@@ -91,7 +93,6 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
         )
         self.pubsub = self.redis_client.pubsub()
         self.sample_id = soft_signal_rw(str)  # Should get from redis
-        self.stop_angle = stop_angle
 
         self._reset()
 
@@ -103,7 +104,7 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
         super().__init__(name=name)
 
     def _reset(self):
-        self._last_omega = 0
+        self._last_omega = None
         self._results: list[MurkoResult] = []
 
     @AsyncStatus.wrap
@@ -120,14 +121,27 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
 
     @AsyncStatus.wrap
     async def trigger(self):
-        # Wait for results
         sample_id = await self.sample_id.get_value()
-        while self._last_omega < self.stop_angle:
+        t_last_result = time.time()
+        while True:
+            if time.time() - t_last_result > self.RESULTS_COMPLETE_TIMEOUT_S:
+                LOGGER.warning(
+                    f"Time since last result > {self.RESULTS_COMPLETE_TIMEOUT_S}, expected to receive {RESULTS_COMPLETE_MESSAGE}"
+                )
+                break
             # waits here for next batch to be received
-            message = await self.pubsub.get_message(timeout=self.TIMEOUT_S)
-            if message is None:
-                continue
-            await self.process_batch(message, sample_id)
+            message = await self.pubsub.get_message(timeout=self.GET_MESSAGE_TIMEOUT_S)
+            if message and message["type"] == "message":
+                t_last_result = time.time()
+                data = pickle.loads(message["data"])
+
+                if data == RESULTS_COMPLETE_MESSAGE:
+                    LOGGER.info(
+                        f"Received results complete message: {RESULTS_COMPLETE_MESSAGE}"
+                    )
+                    break
+
+                await self.process_batch(data, sample_id)
 
         if not self._results:
             raise NoResultsFoundError("No results retrieved from Murko")
@@ -155,22 +169,18 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
                 f"murko:{sample_id}:metadata", result.uuid, json.dumps(result.metadata)
             )
 
-    async def process_batch(self, message: dict | None, sample_id: str):
-        if message and message["type"] == "message":
-            batch_results: list[dict] = pickle.loads(message["data"])
-            for results in batch_results:
-                for uuid, result in results.items():
-                    if metadata_str := await self.redis_client.hget(  # type: ignore
-                        f"murko:{sample_id}:metadata", uuid
-                    ):
-                        LOGGER.info(
-                            f"Found metadata for uuid {uuid}, processing result"
-                        )
-                        self.process_result(
-                            result, MurkoMetadata(json.loads(metadata_str))
-                        )
-                    else:
-                        LOGGER.info(f"Found no metadata for uuid {uuid}")
+    async def process_batch(
+        self, batch_results: list[tuple[str, dict]], sample_id: str
+    ):
+        for result_with_uuid in batch_results:
+            uuid, result = result_with_uuid
+            if metadata_str := await self.redis_client.hget(  # type: ignore
+                f"murko:{sample_id}:metadata", uuid
+            ):
+                LOGGER.info(f"Found metadata for uuid {uuid}, processing result")
+                self.process_result(result, MurkoMetadata(json.loads(metadata_str)))
+            else:
+                LOGGER.info(f"Found no metadata for uuid {uuid}")
 
     def process_result(self, result: dict, metadata: MurkoMetadata):
         """Uses the 'most_likely_click' coordinates from Murko to calculate the

dodal/devices/i04/transfocator.py

@@ -1,5 +1,4 @@
 import asyncio
-import math
 
 from ophyd_async.core import (
     AsyncStatus,
@@ -9,6 +8,7 @@ from ophyd_async.core import (
 )
 from ophyd_async.epics.core import epics_signal_r, epics_signal_rw
 
+from dodal.common.device_utils import periodic_reminder
 from dodal.log import LOGGER
 
 
@@ -23,15 +23,15 @@ class Transfocator(StandardReadable):
     """
 
     def __init__(self, prefix: str, name: str = ""):
+        self._vert_size_calc_sp = epics_signal_rw(float, prefix + "VERT_REQ")
+        self._num_lenses_calc_rbv = epics_signal_r(float, prefix + "LENS_PRED")
+        self.start = epics_signal_rw(int, prefix + "START.PROC")
+        self.start_rbv = epics_signal_r(int, prefix + "START_RBV")
+
         with self.add_children_as_readables():
-            self.beamsize_set_microns = epics_signal_rw(float, prefix + "VERT_REQ")
-            self.predicted_vertical_num_lenses = epics_signal_rw(
-                float, prefix + "LENS_PRED"
-            )
             self.number_filters_sp = epics_signal_rw(int, prefix + "NUM_FILTERS")
-            self.start = epics_signal_rw(int, prefix + "START.PROC")
-            self.start_rbv = epics_signal_r(int, prefix + "START_RBV")
-            self.vertical_lens_rbv = epics_signal_r(float, prefix + "VER")
+            self.current_horizontal_size_rbv = epics_signal_r(float, prefix + "HOR")
+            self.current_vertical_size_rbv = epics_signal_r(float, prefix + "VER")
 
         self.TIMEOUT = 120
 
@@ -41,14 +41,10 @@ class Transfocator(StandardReadable):
         # We can only put an integer number of lenses in the beam but the
         # calculation in the IOC returns the theoretical float number of lenses
         value = round(value)
-        LOGGER.info(f"Transfocator setting {value} filters")
         await self.number_filters_sp.set(value)
         await self.start.set(1)
-        LOGGER.info("Waiting for start_rbv to change to 1")
         await wait_for_value(self.start_rbv, 1, self.TIMEOUT)
-        LOGGER.info("Waiting for start_rbv to change to 0")
         await wait_for_value(self.start_rbv, 0, self.TIMEOUT)
-        self.latest_pred_vertical_num_lenses = value
 
     @AsyncStatus.wrap
     async def set(self, value: float):
@@ -59,29 +55,27 @@ class Transfocator(StandardReadable):
         4. Start the device moving
         5. Wait for the start_rbv goes high and low again
         """
-        self.latest_pred_vertical_num_lenses = (
-            await self.predicted_vertical_num_lenses.get_value()
-        )
-
         LOGGER.info(f"Transfocator setting {value} beamsize")
 
-        if await self.beamsize_set_microns.get_value() != value:
-            # Logic in the IOC calculates predicted_vertical_num_lenses when beam_set_microns changes
+        # Logic in the IOC calculates _num_lenses_calc_rbv when _vert_size_calc_sp changes
+
+        # Register an observer before setting _vert_size_calc_sp to ensure we don't miss changes
+        num_lenses_calc_iterator = observe_value(
+            self._num_lenses_calc_rbv, timeout=self.TIMEOUT
+        )
+
+        await anext(num_lenses_calc_iterator)
+        await self._vert_size_calc_sp.set(value)
+        calc_lenses = await anext(num_lenses_calc_iterator)
 
-            # Register an observer before setting beamsize_set_microns to ensure we don't miss changes
-            predicted_vertical_num_lenses_iterator = observe_value(
-                self.predicted_vertical_num_lenses, timeout=self.TIMEOUT
-            )
-            # Keep initial prediction before setting to later compare with change after setting
-            current_prediction = await anext(predicted_vertical_num_lenses_iterator)
-            await self.beamsize_set_microns.set(value)
-            accepted_prediction = await anext(predicted_vertical_num_lenses_iterator)
-            if not math.isclose(current_prediction, accepted_prediction, abs_tol=1e-8):
-                await self.set_based_on_prediction(accepted_prediction)
+        async with periodic_reminder(
+            f"Waiting for transfocator to insert {calc_lenses} into beam"
+        ):
+            await self.set_based_on_prediction(calc_lenses)
 
         number_filters_rbv, vertical_lens_size_rbv = await asyncio.gather(
             self.number_filters_sp.get_value(),
-            self.vertical_lens_rbv.get_value(),
+            self.current_vertical_size_rbv.get_value(),
         )
 
         LOGGER.info(

dodal/devices/i07/id.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+from dodal.devices.undulator import UndulatorInKeV, UndulatorOrder
+from dodal.devices.util.lookup_tables import energy_distance_table
+
+
+class InsertionDevice(UndulatorInKeV):
+    """
+    Insertion device for i07 including beamline-specific energy-gap lookup behaviour
+    """
+
+    def __init__(
+        self,
+        name: str,
+        prefix: str,
+        harmonic: UndulatorOrder,
+        id_gap_lookup_table_path: str = "/dls_sw/i07/software/gda/config/lookupTables/"
+        + "IIDCalibrationTable.txt",
+    ) -> None:
+        super().__init__(prefix, id_gap_lookup_table_path, name=name)
+        self.harmonic = harmonic
+
+    async def _get_gap_to_match_energy(self, energy_kev: float) -> float:
+        """
+        i07's energy scans remain on a particular harmonic while changing energy.  The
+        calibration table has one row for each harmonic, row contains max and min
+        energies and their corresponding ID gaps.  The requested energy is used to
+        interpolate between these values, assuming a linear relationship on the relevant
+        scale.
+        """
+        energy_to_distance_table: np.ndarray = await energy_distance_table(
+            self.id_gap_lookup_table_path, comments="#", skiprows=2
+        )
+        harmonic_value: int = await self.harmonic.value.get_value()
+
+        row: np.ndarray = energy_to_distance_table[harmonic_value - 1, :]
+        gap = np.interp(energy_kev, [row[1], row[2]], [row[3], row[4]])
+        return gap

dodal/devices/i09_1_shared/__init__.py

@@ -1,3 +1,7 @@
-from .hard_undulator_functions import calculate_gap_i09_hu, get_hu_lut_as_dict
+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"]
+__all__ = ["calculate_gap_i09_hu", "get_hu_lut_as_dict", "calculate_energy_i09_hu"]

dodal/devices/i09_1_shared/hard_undulator_functions.py

@@ -14,11 +14,16 @@ RING_ENERGY_COLUMN = 1
 MAGNET_FIELD_COLUMN = 2
 MIN_ENERGY_COLUMN = 3
 MAX_ENERGY_COLUMN = 4
+MIN_GAP_COLUMN = 5
+MAX_GAP_COLUMN = 6
 GAP_OFFSET_COLUMN = 7
 
+MAGNET_BLOCKS_PER_PERIOD = 4
+MAGNTE_BLOCK_HEIGHT_MM = 16
 
-async def get_hu_lut_as_dict(lut_path: str) -> dict:
-    lut_dict: dict = {}
+
+async def get_hu_lut_as_dict(lut_path: str) -> dict[int, np.ndarray]:
+    lut_dict: dict[int, np.ndarray] = {}
     _lookup_table: np.ndarray = await energy_distance_table(
         lut_path,
         comments=LUT_COMMENTS,
@@ -26,13 +31,44 @@ async def get_hu_lut_as_dict(lut_path: str) -> dict:
     )
     for i in range(_lookup_table.shape[0]):
         lut_dict[_lookup_table[i][0]] = _lookup_table[i]
-        LOGGER.debug(f"Loaded lookup table:\n {lut_dict}")
+    LOGGER.debug(f"Loaded lookup table: {lut_dict}")
     return lut_dict
 
 
+def _validate_order(order: int, look_up_table: dict[int, "np.ndarray"]) -> None:
+    """Validate that the harmonic order exists in the lookup table."""
+    if order not in look_up_table.keys():
+        raise ValueError(f"Order parameter {order} not found in lookup table")
+
+
+def _calculate_gamma(look_up_table: dict[int, "np.ndarray"], order: int) -> float:
+    """Calculate the Lorentz factor gamma from the lookup table."""
+    return 1000 * look_up_table[order][RING_ENERGY_COLUMN] / ELECTRON_REST_ENERGY_MEV
+
+
+def _calculate_undulator_parameter_max(
+    magnet_field: float, undulator_period_mm: int
+) -> float:
+    """
+    Calculate the maximum undulator parameter.
+    """
+    return (
+        (
+            2
+            * 0.0934
+            * undulator_period_mm
+            * magnet_field
+            * MAGNET_BLOCKS_PER_PERIOD
+            / np.pi
+        )
+        * np.sin(np.pi / MAGNET_BLOCKS_PER_PERIOD)
+        * (1 - np.exp(-2 * np.pi * MAGNTE_BLOCK_HEIGHT_MM / undulator_period_mm))
+    )
+
+
 def calculate_gap_i09_hu(
     photon_energy_kev: float,
-    look_up_table: dict[int, "np.ndarray"],
+    look_up_table: dict[int, np.ndarray],
     order: int = 1,
     gap_offset: float = 0.0,
     undulator_period_mm: int = 27,
@@ -52,13 +88,9 @@ def calculate_gap_i09_hu(
     Returns:
         float: Calculated undulator gap in millimeters.
     """
-    magnet_blocks_per_period = 4
-    magnet_block_height_mm = 16
 
-    if order not in look_up_table.keys():
-        raise ValueError(f"Order parameter {order} not found in lookup table")
-
-    gamma = 1000 * look_up_table[order][RING_ENERGY_COLUMN] / ELECTRON_REST_ENERGY_MEV
+    _validate_order(order, look_up_table)
+    gamma = _calculate_gamma(look_up_table, order)
 
     # Constructive interference of radiation emitted at different poles
     # lamda = (lambda_u/2*gamma^2)*(1+K^2/2 + gamma^2*theta^2)/n for n=1,2,3...
@@ -83,17 +115,8 @@ def calculate_gap_i09_hu(
     # leading to K = 0.934*B0[T]*lambda_u[cm]*exp(-pi*gap/lambda_u) or
     # K = undulator_parameter_max*exp(-pi*gap/lambda_u)
     # Calculating undulator_parameter_max gives:
-    undulator_parameter_max = (
-        (
-            2
-            * 0.0934
-            * undulator_period_mm
-            * look_up_table[order][MAGNET_FIELD_COLUMN]
-            * magnet_blocks_per_period
-            / np.pi
-        )
-        * np.sin(np.pi / magnet_blocks_per_period)
-        * (1 - np.exp(-2 * np.pi * magnet_block_height_mm / undulator_period_mm))
+    undulator_parameter_max = _calculate_undulator_parameter_max(
+        look_up_table[order][MAGNET_FIELD_COLUMN], undulator_period_mm
     )
 
     # Finnaly, rearranging the equation:
@@ -109,3 +132,44 @@ def calculate_gap_i09_hu(
     )
 
     return gap
+
+
+def calculate_energy_i09_hu(
+    gap: float,
+    look_up_table: dict[int, "np.ndarray"],
+    order: int = 1,
+    gap_offset: float = 0.0,
+    undulator_period_mm: int = 27,
+) -> float:
+    """
+    Calculate the photon energy produced by the undulator at a given gap and harmonic order.
+    Reverse of the calculate_gap_i09_hu function.
+
+    Args:
+        gap (float): Undulator gap in millimeters.
+        look_up_table (dict[int, np.ndarray]): Lookup table containing undulator and beamline parameters for each harmonic order.
+        order (int, optional): Harmonic order for which to calculate the energy. Defaults to 1.
+        gap_offset (float, optional): Additional gap offset to apply (in mm). Defaults to 0.0.
+        undulator_period_mm (int, optional): Undulator period in mm. Defaults to 27.
+
+    Returns:
+        float: Calculated photon energy in keV.
+    """
+    _validate_order(order, look_up_table)
+
+    gamma = _calculate_gamma(look_up_table, order)
+    undulator_parameter_max = _calculate_undulator_parameter_max(
+        look_up_table[order][MAGNET_FIELD_COLUMN], undulator_period_mm
+    )
+
+    undulator_parameter = undulator_parameter_max / np.exp(
+        (gap - look_up_table[order][GAP_OFFSET_COLUMN] - gap_offset)
+        / (undulator_period_mm / np.pi)
+    )
+    energy_kev = (
+        4.959368e-6
+        * order
+        * np.square(gamma)
+        / (undulator_period_mm * (np.square(undulator_parameter) + 2))
+    )
+    return energy_kev