i09 energy to gap function (#1660)

* add gap calcualation functions and tests

* remove scipy import - not found

* relocate to i09_1_shared

---------

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

Author: 3 people

src/dodal/devices/i09_1_shared/__init__.py

@@ -0,0 +1,3 @@
+from .hard_undulator_functions import calculate_gap_i09_hu, get_hu_lut_as_dict
+
+__all__ = ["calculate_gap_i09_hu", "get_hu_lut_as_dict"]

src/dodal/devices/i09_1_shared/hard_undulator_functions.py

@@ -0,0 +1,111 @@
+import numpy as np
+
+from dodal.devices.util.lookup_tables import energy_distance_table
+from dodal.log import LOGGER
+
+LUT_COMMENTS = ["#"]
+HU_SKIP_ROWS = 3
+
+# Physics constants
+ELECTRON_REST_ENERGY_MEV = 0.510999
+
+# Columns in the lookup table
+RING_ENERGY_COLUMN = 1
+MAGNET_FIELD_COLUMN = 2
+MIN_ENERGY_COLUMN = 3
+MAX_ENERGY_COLUMN = 4
+GAP_OFFSET_COLUMN = 7
+
+
+async def get_hu_lut_as_dict(lut_path: str) -> dict:
+    lut_dict: dict = {}
+    _lookup_table: np.ndarray = await energy_distance_table(
+        lut_path,
+        comments=LUT_COMMENTS,
+        skiprows=HU_SKIP_ROWS,
+    )
+    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}")
+    return lut_dict
+
+
+def calculate_gap_i09_hu(
+    photon_energy_kev: float,
+    look_up_table: dict[int, "np.ndarray"],
+    order: int = 1,
+    gap_offset: float = 0.0,
+    undulator_period_mm: int = 27,
+) -> float:
+    """
+    Calculate the undulator gap required to produce a given energy at a given harmonic order.
+    This algorithm was provided by the I09 beamline scientists, and is based on the physics of undulator radiation.
+    https://cxro.lbl.gov//PDF/X-Ray-Data-Booklet.pdf
+
+    Args:
+        photon_energy_kev (float): Requested photon energy in keV.
+        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 gap. 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 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
+
+    # 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...
+    # theta is the observation angle, assumed to be 0 here.
+    # Rearranging for K (the undulator parameter, related to magnetic field and gap)
+    # gives K^2 = 2*((2*n*gamma^2*lamda/lambda_u)-1)
+
+    undulator_parameter_sqr = (
+        4.959368e-6
+        * (order * gamma * gamma / (undulator_period_mm * photon_energy_kev))
+        - 2
+    )
+    if undulator_parameter_sqr < 0:
+        raise ValueError(
+            f"Diffraction parameter squared must be positive! Calculated value {undulator_parameter_sqr}."
+        )
+    undulator_parameter = np.sqrt(undulator_parameter_sqr)
+
+    # Undulator_parameter K is also defined as K = 0.934*B0[T]*lambda_u[cm],
+    # where B0[T] is a peak magnetic field that must depend on gap,
+    # but in our LUT it is does not depend on gap, so it's a factor,
+    # 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))
+    )
+
+    # Finnaly, rearranging the equation:
+    # undulator_parameter = undulator_parameter_max*exp(-pi*gap/lambda_u) for gap gives
+    gap = (
+        (undulator_period_mm / np.pi)
+        * np.log(undulator_parameter_max / undulator_parameter)
+        + look_up_table[order][GAP_OFFSET_COLUMN]
+        + gap_offset
+    )
+    LOGGER.debug(
+        f"Calculated gap is {gap}mm for energy {photon_energy_kev}keV at order {order}"
+    )
+
+    return gap

src/dodal/devices/util/lookup_tables.py

@@ -13,13 +13,19 @@
 from dodal.log import LOGGER
 
 
-async def energy_distance_table(lookup_table_path: str) -> np.ndarray:
+async def energy_distance_table(
+    lookup_table_path: str,
+    comments: Sequence[str] = ["#", "Units"],
+    skiprows: int = 0,
+) -> np.ndarray:
     """
     Returns a numpy formatted lookup table for required positions of an ID gap to
     provide emission at a given beam energy.
 
     Args:
         lookup_table_path: Path to lookup table
+        comments: Lines starting with any of these strings will be ignored
+        skiprows: Number of rows to skip at the start of the file
 
     Returns:
         ndarray: Lookup table
@@ -29,7 +35,7 @@ async def energy_distance_table(lookup_table_path: str) -> np.ndarray:
     # decodes the text
     async with aiofiles.open(lookup_table_path) as stream:
         raw_table = await stream.read()
-    return loadtxt(StringIO(raw_table), comments=["#", "Units"])
+    return loadtxt(StringIO(raw_table), comments=comments, skiprows=skiprows)
 
 
 def parse_lookup_table(filename: str) -> list[Sequence]:

tests/devices/i09_1/__init__.py

@@ -0,0 +1,9 @@
+from os.path import join
+from pathlib import Path
+
+TEST_DATA_PATH = Path(__file__).parent
+TEST_HARD_UNDULATOR_LUT = join(TEST_DATA_PATH, "test_lookuptable_i09_hu.txt")
+__all__ = [
+    "TEST_DATA_PATH",
+    "TEST_HARD_UNDULATOR_LUT",
+]

tests/devices/i09_1_shared/__init__.py

@@ -0,0 +1,26 @@
+#I09 Hard X-ray ID calibration parameters for I09 undulator
+ScannableNames	n	Ee	Br	Epmin	Epmax	Gmin	Gmax	Goffset
+ScannableUnits	ONE	GeV	T	KeV	Kev	mm	mm
+1	3.00089	0.98928	2.12	3.05	14.2650	23.7200	0.0
+2	3.04129	1.02504	2.50	2.80	5.05165	8.88007	0.0
+3	3.05798	1.03065	2.40	4.30	5.20000	8.99036	0.0
+4	3.03635	1.02332	3.20	5.70	5.26183	8.99640	0.0
+5	3.06334	1.03294	4.00	7.20	5.22735	9.02065	0.0
+6	3.04963	1.02913	4.70	8.60	5.13939	9.02527	0.0
+7	3.06515	1.03339	5.50	10.1	5.12684	9.02602	0.0
+8	3.05775	1.03223	6.30	11.5	5.16289	9.02873	0.0
+9	3.06829	1.03468	7.10	13.0	5.16357	9.03049	0.0
+10	3.06164	1.03328	7.90	14.4	5.17205	9.02845	0.0
+11	3.07056	1.03557	8.60	15.9	5.11350	9.04750	0.0
+12	3.06627	1.03482	9.40	17.3	5.12051	9.02826	0.0
+13	3.07176	1.03623	10.2	18.3	5.13027	8.84940	0.0
+14	3.06964	1.03587	11.0	18.3	5.13985	8.30146	0.0
+15	3.06515	1.03391	11.8	18.3	5.14643	7.82380	0.0
+16	3.08317	1.04191	12.6	18.3	5.15437	7.39118	0.0
+17	3.08294	1.04240	13.4	18.3	5.15906	7.00038	0.0
+18	3.09725	1.04819	14.2	18.3	5.16369	6.64312	0.0
+19	3.09824	1.04921	15.0	18.3	5.16749	6.31454	0.0
+20	3.11042	1.05404	15.7	18.3	5.13513	6.00821	0.0
+21	3.11112	1.05472	16.5	18.3	5.13986	5.72405	0.0
+22	3.11920	1.05799	17.3	18.3	5.14342	5.45745	0.0
+23	3.10218	1.05042	18.1	18.3	5.14564	5.20688	0.0

tests/devices/i09_1_shared/test_data/__init__.py

@@ -0,0 +1,53 @@
+import re
+
+import pytest
+
+from dodal.devices.i09_1_shared import calculate_gap_i09_hu, get_hu_lut_as_dict
+from tests.devices.i09_1_shared.test_data import TEST_HARD_UNDULATOR_LUT
+
+
+@pytest.fixture
+async def lut_dictionary() -> dict:
+    return await get_hu_lut_as_dict(TEST_HARD_UNDULATOR_LUT)
+
+
+@pytest.mark.parametrize(
+    "energy, order, expected_gap",
+    [
+        (2.13, 1, 12.81),
+        (2.78, 3, 6.05),
+        (6.24, 5, 7.95),
+    ],
+)
+async def test_calculate_gap_from_energy(
+    energy: float,
+    order: int,
+    expected_gap: float,
+    lut_dictionary: dict,
+):
+    assert calculate_gap_i09_hu(energy, lut_dictionary, order) == pytest.approx(
+        expected_gap, abs=0.01
+    )
+
+
+async def test_calculate_gap_from_energy_wrong_order(
+    lut_dictionary: dict,
+):
+    wrong_order = 100
+    with pytest.raises(
+        ValueError,
+        match=re.escape(f"Order parameter {wrong_order} not found in lookup table"),
+    ):
+        calculate_gap_i09_hu(30, lut_dictionary, wrong_order)
+
+
+async def test_calculate_gap_from_energy_wrong_k(
+    lut_dictionary: dict,
+):
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Diffraction parameter squared must be positive! Calculated value -1.78"
+        ),
+    ):
+        calculate_gap_i09_hu(30, lut_dictionary, 1)