I09 Add reverse function to calculate energy from gap, harmonics and lookuptable.

src/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"]

src/dodal/devices/i09_1_shared/hard_undulator_functions.py

@@ -14,11 +14,16 @@
 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 = {}
+    lut_dict: dict = {int: "np.ndarray"}
     _lookup_table: np.ndarray = await energy_distance_table(
         lut_path,
         comments=LUT_COMMENTS,
@@ -30,6 +35,37 @@ async def get_hu_lut_as_dict(lut_path: str) -> 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"],
@@ -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

tests/devices/i09_1_shared/test_undulator_functions.py

@@ -3,6 +3,7 @@
 import pytest
 
 from dodal.devices.i09_1_shared import calculate_gap_i09_hu, get_hu_lut_as_dict
+from dodal.devices.i09_1_shared.hard_undulator_functions import calculate_energy_i09_hu
 from tests.devices.i09_1_shared.test_data import TEST_HARD_UNDULATOR_LUT
 
 
@@ -14,9 +15,9 @@ async def lut_dictionary() -> dict:
 @pytest.mark.parametrize(
     "energy, order, expected_gap",
     [
-        (2.13, 1, 12.81),
-        (2.78, 3, 6.05),
-        (6.24, 5, 7.95),
+        (2.13, 1, 12.8146),
+        (2.78, 3, 6.0537),
+        (6.24, 5, 7.9561),
     ],
 )
 async def test_calculate_gap_from_energy(
@@ -26,7 +27,26 @@ async def test_calculate_gap_from_energy(
     lut_dictionary: dict,
 ):
     assert calculate_gap_i09_hu(energy, lut_dictionary, order) == pytest.approx(
-        expected_gap, abs=0.01
+        expected_gap, abs=0.0001
+    )
+
+
+@pytest.mark.parametrize(
+    "energy, order, gap",
+    [
+        (2.1454, 1, 12.91),
+        (2.6410, 3, 5.75),
+        (6.1355, 5, 7.84),
+    ],
+)
+async def test_calculate_energy_from_gap(
+    energy: float,
+    order: int,
+    gap: float,
+    lut_dictionary: dict,
+):
+    assert calculate_energy_i09_hu(gap, lut_dictionary, order) == pytest.approx(
+        energy, abs=0.0001
     )