add energy and id to i05-shared

src/dodal/beamlines/i05_shared.py

@@ -1,10 +1,21 @@
 from dodal.device_manager import DeviceManager
-from dodal.devices.beamlines.i05.enums import Grating
+from dodal.devices.beamlines.i05_shared import (
+    APPLE_KNOT_EXCLUSION_ZONES,
+    Grating,
+    energy_to_gap_converter,
+    energy_to_phase_converter,
+)
 from dodal.devices.insertion_device import (
     Apple2,
     UndulatorGap,
     UndulatorLockedPhaseAxes,
 )
+from dodal.devices.insertion_device.apple_knot_controller import (
+    AppleKnotController,
+    AppleKnotPathFinder,
+)
+from dodal.devices.insertion_device.energy import BeamEnergy, InsertionDeviceEnergy
+from dodal.devices.insertion_device.polarisation import InsertionDevicePolarisation
 from dodal.devices.pgm import PlaneGratingMonochromator
 from dodal.devices.synchrotron import Synchrotron
 from dodal.utils import BeamlinePrefix, get_beamline_name
@@ -49,3 +60,36 @@ def id(
 ) -> Apple2[UndulatorLockedPhaseAxes]:
     """i05 insertion device."""
     return Apple2[UndulatorLockedPhaseAxes](id_gap=id_gap, id_phase=id_phase)
+
+
+@devices.factory()
+def id_controller(
+    id: Apple2[UndulatorLockedPhaseAxes],
+) -> AppleKnotController[UndulatorLockedPhaseAxes]:
+    return AppleKnotController[UndulatorLockedPhaseAxes](
+        apple=id,
+        gap_energy_motor_converter=energy_to_gap_converter,
+        phase_energy_motor_converter=energy_to_phase_converter,
+        path_finder=AppleKnotPathFinder(APPLE_KNOT_EXCLUSION_ZONES),
+    )
+
+
+@devices.factory()
+def id_energy(
+    id_controller: AppleKnotController[UndulatorLockedPhaseAxes],
+) -> InsertionDeviceEnergy:
+    return InsertionDeviceEnergy(id_controller=id_controller)
+
+
+@devices.factory()
+def id_polarisation(
+    id_controller: AppleKnotController[UndulatorLockedPhaseAxes],
+) -> InsertionDevicePolarisation:
+    return InsertionDevicePolarisation(id_controller=id_controller)
+
+
+@devices.factory()
+def energy(
+    id_energy: InsertionDeviceEnergy, pgm: PlaneGratingMonochromator
+) -> BeamEnergy:
+    return BeamEnergy(id_energy=id_energy, mono=pgm.energy)

src/dodal/devices/beamlines/i05/__init__.py

@@ -1,5 +1,3 @@
-from .compound_motors import PolynomCompoundMotors
-from .enums import Grating
 from .i05_motors import I05Goniometer
 
-__all__ = ["Grating", "PolynomCompoundMotors", "I05Goniometer"]
+__all__ = ["I05Goniometer"]

src/dodal/devices/beamlines/i05_shared/__init__.py

@@ -0,0 +1,15 @@
+from dodal.devices.beamlines.i05_shared.apple_knot_constants import (
+    APPLE_KNOT_EXCLUSION_ZONES,
+    energy_to_gap_converter,
+    energy_to_phase_converter,
+)
+from dodal.devices.beamlines.i05_shared.compound_motors import PolynomCompoundMotors
+from dodal.devices.beamlines.i05_shared.enums import Grating
+
+__all__ = [
+    "Grating",
+    "PolynomCompoundMotors",
+    "energy_to_gap_converter",
+    "energy_to_phase_converter",
+    "APPLE_KNOT_EXCLUSION_ZONES",
+]

src/dodal/devices/beamlines/i05_shared/apple_knot_constants.py

@@ -0,0 +1,92 @@
+from numpy.polynomial import Polynomial
+
+from dodal.common.maths import Rectangle2D
+from dodal.devices.insertion_device.enum import Pol
+
+# Define exclusion zones in phase-gap space
+APPLE_KNOT_EXCLUSION_ZONES = (
+    Rectangle2D(-65.5, 0.0, 65.5, 25.5),  # mechanical limit
+    Rectangle2D(-10.5, 0.0, 10.5, 37.5),  # power load limit
+)
+
+# Polynomials for energy to gap conversion
+LH_GAP_POLYNOMIAL = Polynomial(
+    [
+        12.464,
+        1.8417,
+        -0.030139,
+        0.00023511,
+        1.0158e-6,
+        -3.9229e-8,
+        3.6772e-10,
+        -1.7356e-12,
+        4.2103e-15,
+        -4.1724e-18,
+    ]
+)
+LV_GAP_POLYNOMIAL = Polynomial(
+    [
+        8.7456,
+        1.1344,
+        -0.024317,
+        0.00041143,
+        -5.0759e-6,
+        4.496e-8,
+        -2.7464e-10,
+        1.081e-12,
+        -2.4377e-15,
+        2.3749e-18,
+    ]
+)
+C_GAP_POLYNOMIAL = Polynomial(
+    [
+        9.1763,
+        1.4886,
+        -0.035968,
+        0.00064576,
+        -7.951e-6,
+        6.6281e-8,
+        -3.6547e-10,
+        1.2699e-12,
+        -2.5078e-15,
+        2.1392e-18,
+    ]
+)
+
+# Polynomial for energy to phase conversion
+C_PHASE_POLYNOMIAL = Polynomial(
+    [
+        34.431,
+        0.79535,
+        -0.022218,
+        0.00040781,
+        -4.921e-6,
+        3.9683e-8,
+        -2.1267e-10,
+        7.2752e-13,
+        -1.4341e-15,
+        1.2345e-18,
+    ]
+)
+
+
+def energy_to_gap_converter(energy: float, pol: Pol) -> float:
+    if pol == Pol.LH:
+        return float(LH_GAP_POLYNOMIAL(energy))
+    if pol == Pol.LV:
+        return float(LV_GAP_POLYNOMIAL(energy))
+    if pol == Pol.PC or pol == Pol.NC:
+        return float(C_GAP_POLYNOMIAL(energy))
+    return 0.0
+
+
+def energy_to_phase_converter(energy: float, pol: Pol) -> float:
+    if pol == Pol.LH:
+        return 0.0
+    if pol == Pol.LV:
+        return 70.0
+    if pol == Pol.PC:
+        return float(C_PHASE_POLYNOMIAL(energy))
+    if pol == Pol.NC:
+        return -float(C_PHASE_POLYNOMIAL(energy))
+    return 0.0

src/dodal/testing/fixtures/devices/apple_knot.py

@@ -0,0 +1,11 @@
+import pytest
+
+from dodal.devices.beamlines.i05_shared import (
+    APPLE_KNOT_EXCLUSION_ZONES,
+)
+from dodal.devices.insertion_device.apple_knot_controller import AppleKnotPathFinder
+
+
+@pytest.fixture
+def apple_knot_i05_path_finder() -> AppleKnotPathFinder:
+    return AppleKnotPathFinder(APPLE_KNOT_EXCLUSION_ZONES)

tests/devices/beamlines/i05/test_compound_motor.py → tests/devices/beamlines/i05_shared/test_compound_motor.py

@@ -6,7 +6,7 @@
 from ophyd_async.epics.motor import Motor
 from ophyd_async.testing import assert_configuration, assert_reading, partial_reading
 
-from dodal.devices.beamlines.i05 import PolynomCompoundMotors
+from dodal.devices.beamlines.i05_shared import PolynomCompoundMotors
 
 
 @pytest.fixture

tests/devices/insertion_device/test_apple_knot_path_finder.py

@@ -1,6 +1,5 @@
 import pytest
 
-from dodal.common.maths import Rectangle2D
 from dodal.devices.insertion_device import (
     AppleKnotPathFinder,
 )
@@ -10,13 +9,10 @@
 )
 
 # add mock_config_client, mock_id_gap, mock_phase and mock_jaw_phase_axes to pytest.
-pytest_plugins = ["dodal.testing.fixtures.devices.apple2"]
-
-# Define exclusion zones in phase-gap space
-TEST_APPLE_KNOT_EXCLUSION_ZONES = (
-    Rectangle2D(-65.5, 0.0, 65.5, 25.5),  # mechanical limit
-    Rectangle2D(-10.5, 0.0, 10.5, 37.5),  # power load limit
-)
+pytest_plugins = [
+    "dodal.testing.fixtures.devices.apple2",
+    "dodal.testing.fixtures.devices.apple_knot",
+]
 
 
 def get_pair_apple2_val(
@@ -42,11 +38,6 @@ def get_pair_apple2_val(
     return start_val, target_val
 
 
-@pytest.fixture
-def apple_knot_i05_path_finder() -> AppleKnotPathFinder:
-    return AppleKnotPathFinder(TEST_APPLE_KNOT_EXCLUSION_ZONES)
-
-
 @pytest.mark.parametrize(
     "start_phase, start_gap, target_phase, target_gap",
     [

tests/devices/insertion_device/test_apple_knot_undulator.py

@@ -1,10 +1,12 @@
-import numpy.polynomial as poly
 import pytest
 from ophyd_async.core import (
     set_mock_value,
 )
 
-from dodal.common.maths import Rectangle2D
+from dodal.devices.beamlines.i05_shared import (
+    energy_to_gap_converter,
+    energy_to_phase_converter,
+)
 from dodal.devices.insertion_device import (
     AppleKnotController,
     AppleKnotPathFinder,
@@ -16,47 +18,10 @@
 )
 
 # add mock_config_client, mock_id_gap, mock_phase and mock_jaw_phase_axes to pytest.
-pytest_plugins = ["dodal.testing.fixtures.devices.apple2"]
-
-# Define exclusion zones in phase-gap space
-TEST_APPLE_KNOT_EXCLUSION_ZONES = (
-    Rectangle2D(-65.5, 0.0, 65.5, 25.5),  # mechanical limit
-    Rectangle2D(-10.5, 0.0, 10.5, 37.5),  # power load limit
-)
-# Test polynomials for energy to gap/phase conversion
-TEST_LH_GAP_POLYNOMIAL = poly.Polynomial([12.46, 0.832, -0.002])
-TEST_LV_GAP_POLYNOMIAL = poly.Polynomial([8.7456, 0.412, -0.0024317])
-TEST_C_GAP_POLYNOMIAL = poly.Polynomial([9.1763, 0.312, -0.000305968])
-TEST_C_PHASE_POLYNOMIAL = poly.Polynomial(
-    [4.4, 0.79, -0.022, 0.00041, -4.921e-6, 3.9683e-8]
-)
-
-
-def energy_to_gap_converter(energy: float, pol: Pol) -> float:
-    if pol == Pol.LH:
-        return float(TEST_LH_GAP_POLYNOMIAL(energy))
-    if pol == Pol.LV:
-        return float(TEST_LV_GAP_POLYNOMIAL(energy))
-    if pol == Pol.PC or pol == Pol.NC:
-        return float(TEST_C_GAP_POLYNOMIAL(energy))
-    return 0.0
-
-
-def energy_to_phase_converter(energy: float, pol: Pol) -> float:
-    if pol == Pol.LH:
-        return 0.0
-    if pol == Pol.LV:
-        return 70.0
-    if pol == Pol.PC:
-        return float(TEST_C_PHASE_POLYNOMIAL(energy))
-    if pol == Pol.NC:
-        return -float(TEST_C_PHASE_POLYNOMIAL(energy))
-    return 0.0
-
-
-@pytest.fixture
-def apple_knot_i05_path_finder() -> AppleKnotPathFinder:
-    return AppleKnotPathFinder(TEST_APPLE_KNOT_EXCLUSION_ZONES)
+pytest_plugins = [
+    "dodal.testing.fixtures.devices.apple2",
+    "dodal.testing.fixtures.devices.apple_knot",
+]
 
 
 @pytest.fixture
@@ -139,6 +104,43 @@ async def test_id_set_pol(
     assert await mock_apple_knot_i05_controller.polarisation.get_value() == target_pol
 
 
+@pytest.mark.parametrize(
+    "initial_pol, initial_energy, target_pol",
+    [
+        (Pol.LV, 40.0, Pol.LA),
+        (Pol.LV, 40.0, Pol.LH3),
+        (Pol.LV, 80.0, Pol.LV3),
+        (Pol.LV, 80.0, Pol.NONE),
+    ],
+)
+async def test_id_set_pol_fails(
+    mock_apple_knot_i05_controller: AppleKnotController[UndulatorLockedPhaseAxes],
+    mock_locked_apple2: Apple2[UndulatorLockedPhaseAxes],
+    initial_pol: Pol,
+    initial_energy: float,
+    target_pol: Pol,
+):
+    mock_apple_knot_i05_controller._energy_set(initial_energy)
+    set_mock_value(
+        mock_locked_apple2.gap().user_readback,
+        energy_to_gap_converter(initial_energy, initial_pol),
+    )
+    set_mock_value(
+        mock_locked_apple2.phase().top_outer.user_readback,
+        energy_to_phase_converter(initial_energy, initial_pol),
+    )
+    set_mock_value(
+        mock_locked_apple2.phase().btm_inner.user_readback,
+        energy_to_phase_converter(initial_energy, initial_pol),
+    )
+    assert await mock_apple_knot_i05_controller.polarisation.get_value() == initial_pol
+    with pytest.raises(
+        RuntimeError,
+        match="No valid path found for move avoiding exclusion zones.",
+    ):
+        await mock_apple_knot_i05_controller.polarisation.set(target_pol)
+
+
 @pytest.mark.parametrize(
     "target_energy, initial_gap, initial_phase_top_outer",
     [