dls-dodal 1.67.0__py3-none-any.whl → 1.69.0__py3-none-any.whl

dodal/devices/insertion_device/energy.py

@@ -0,0 +1,88 @@
+import abc
+import asyncio
+
+from bluesky.protocols import Movable
+from ophyd_async.core import (
+    AsyncStatus,
+    Reference,
+    SignalRW,
+    StandardReadable,
+    StandardReadableFormat,
+    soft_signal_rw,
+)
+from ophyd_async.epics.motor import Motor
+
+from dodal.devices.insertion_device import MAXIMUM_MOVE_TIME, Apple2Controller
+from dodal.log import LOGGER
+
+
+class InsertionDeviceEnergyBase(abc.ABC, StandardReadable, Movable):
+    """Base class for ID energy movable device."""
+
+    def __init__(self, name: str = "") -> None:
+        self.energy: Reference[SignalRW[float]]
+        super().__init__(name=name)
+
+    @abc.abstractmethod
+    @AsyncStatus.wrap
+    async def set(self, energy: float) -> None: ...
+
+
+class BeamEnergy(StandardReadable, Movable[float]):
+    """
+    Compound device to set both ID and energy motor at the same time with an option to add an offset.
+    """
+
+    def __init__(
+        self, id_energy: InsertionDeviceEnergyBase, mono: Motor, name: str = ""
+    ) -> None:
+        """
+        Parameters
+        ----------
+
+        id_energy: InsertionDeviceEnergy
+            An InsertionDeviceEnergy device.
+        mono: Motor
+            A Motor(energy) device.
+        name:
+            New device name.
+        """
+        super().__init__(name=name)
+        self._id_energy = Reference(id_energy)
+        self._mono_energy = Reference(mono)
+
+        self.add_readables(
+            [
+                self._id_energy().energy(),
+                self._mono_energy().user_readback,
+            ],
+            StandardReadableFormat.HINTED_SIGNAL,
+        )
+
+        with self.add_children_as_readables(StandardReadableFormat.CONFIG_SIGNAL):
+            self.id_energy_offset = soft_signal_rw(float, initial_value=0)
+
+    @AsyncStatus.wrap
+    async def set(self, energy: float) -> None:
+        LOGGER.info(f"Moving f{self.name} energy to {energy}.")
+        await asyncio.gather(
+            self._id_energy().set(
+                energy=energy + await self.id_energy_offset.get_value()
+            ),
+            self._mono_energy().set(energy),
+        )
+
+
+class InsertionDeviceEnergy(InsertionDeviceEnergyBase):
+    """Apple2 ID energy movable device."""
+
+    def __init__(self, id_controller: Apple2Controller, name: str = "") -> None:
+        self.energy = Reference(id_controller.energy)
+        super().__init__(name=name)
+
+        self.add_readables([self.energy()], StandardReadableFormat.HINTED_SIGNAL)
+
+    @AsyncStatus.wrap
+    async def set(self, energy: float) -> None:
+        LOGGER.info(f"Setting insertion device energy to {energy}.")
+        await self.energy().set(energy, timeout=MAXIMUM_MOVE_TIME)

dodal/devices/insertion_device/energy_motor_lookup.py

@@ -2,7 +2,7 @@ from pathlib import Path
 
 from daq_config_server.client import ConfigServer
 
-from dodal.devices.insertion_device.apple2_undulator import Pol
+from dodal.devices.insertion_device.enum import Pol
 from dodal.devices.insertion_device.lookup_table_models import (
     LookupTable,
     LookupTableColumnConfig,

dodal/devices/insertion_device/enum.py

@@ -0,0 +1,17 @@
+from ophyd_async.core import StrictEnum
+
+
+class Pol(StrictEnum):
+    NONE = "None"
+    LH = "lh"
+    LV = "lv"
+    PC = "pc"
+    NC = "nc"
+    LA = "la"
+    LH3 = "lh3"
+    LV3 = "lv3"
+
+
+class UndulatorGateStatus(StrictEnum):
+    OPEN = "Open"
+    CLOSE = "Closed"

dodal/devices/insertion_device/lookup_table_models.py

@@ -39,7 +39,7 @@ from pydantic import (
     field_validator,
 )
 
-from dodal.devices.insertion_device.apple2_undulator import Pol
+from dodal.devices.insertion_device.enum import Pol
 
 DEFAULT_POLY_DEG = [
     "7th-order",
@@ -57,7 +57,7 @@ DEFAULT_GAP_FILE = "IDEnergy2GapCalibrations.csv"
 DEFAULT_PHASE_FILE = "IDEnergy2PhaseCalibrations.csv"
 
 ROW_PHASE_MOTOR_TOLERANCE = 0.004
-ROW_PHASE_CIRCULAR = 15
+ROW_PHASE_CIRCULAR = 15.0
 MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
 MAXIMUM_GAP_MOTOR_POSITION = 100
 
@@ -94,13 +94,19 @@ class LookupTableColumnConfig(BaseModel):
         description="When processing polarisation mode values, map their alias values to a real value.",
         default_factory=lambda: MODE_NAME_CONVERT,
     )
+    grating: A[
+        str | None, Field(description="Optional column name for entry grating.")
+    ] = None
 
 
 class EnergyCoverageEntry(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)  # So np.poly1d can be used.
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True, frozen=True
+    )  # arbitrary_types_allowed is True so np.poly1d can be used.
     min_energy: float
     max_energy: float
     poly: np.poly1d
+    grating: float | None = None
 
     @field_validator("poly", mode="before")
     @classmethod
@@ -119,7 +125,16 @@ class EnergyCoverageEntry(BaseModel):
 
 
 class EnergyCoverage(BaseModel):
-    energy_entries: list[EnergyCoverageEntry]
+    model_config = ConfigDict(frozen=True)
+    energy_entries: tuple[EnergyCoverageEntry, ...]
+
+    @field_validator("energy_entries", mode="after")
+    @classmethod
+    def _prepare_energy_entries(
+        cls, value: tuple[EnergyCoverageEntry, ...]
+    ) -> tuple[EnergyCoverageEntry, ...]:
+        """Convert incoming energy_entries to a sorted, immutable tuple."""
+        return tuple(sorted(value, key=lambda e: e.min_energy))
 
     @classmethod
     def generate(
@@ -128,7 +143,7 @@ class EnergyCoverage(BaseModel):
         max_energies: list[float],
         poly1d_params: list[list[float]],
     ) -> Self:
-        energy_entries = [
+        energy_entries = tuple(
             EnergyCoverageEntry(
                 min_energy=min_energy,
                 max_energy=max_energy,
@@ -137,16 +152,16 @@ class EnergyCoverage(BaseModel):
             for min_energy, max_energy, poly_params in zip(
                 min_energies, max_energies, poly1d_params, strict=True
             )
-        ]
+        )
         return cls(energy_entries=energy_entries)
 
     @property
     def min_energy(self) -> float:
-        return min(e.min_energy for e in self.energy_entries)
+        return self.energy_entries[0].min_energy
 
     @property
     def max_energy(self) -> float:
-        return max(e.max_energy for e in self.energy_entries)
+        return self.energy_entries[-1].max_energy
 
     def get_poly(self, energy: float) -> np.poly1d:
         """
@@ -157,25 +172,36 @@ class EnergyCoverage(BaseModel):
         energy:
             Energy value in the same units used to create the lookup table.
         """
-        # Cache initial values so don't do unnecessary work again
-        min_energy = self.min_energy
-        max_energy = self.max_energy
-        if energy < min_energy or energy > max_energy:
+
+        if not self.min_energy <= energy <= self.max_energy:
             raise ValueError(
-                f"Demanding energy must lie between {min_energy} and {max_energy}!"
+                f"Demanding energy must lie between {self.min_energy} and {self.max_energy}!"
             )
-        else:
-            for energy_coverage in self.energy_entries:
-                if (
-                    energy >= energy_coverage.min_energy
-                    and energy < energy_coverage.max_energy
-                ):
-                    return energy_coverage.poly
+
+        poly_index = self.get_energy_index(energy)
+        if poly_index is not None:
+            return self.energy_entries[poly_index].poly
         raise ValueError(
             "Cannot find polynomial coefficients for your requested energy."
             + " There might be gap in the calibration lookup table."
         )
 
+    def get_energy_index(self, energy: float) -> int | None:
+        """Binary search assumes self.energy_entries is sorted by min_energy.
+        Return index or None if not found."""
+        max_index = len(self.energy_entries) - 1
+        min_index = 0
+        while min_index <= max_index:
+            mid_index = (min_index + max_index) // 2
+            en_try = self.energy_entries[mid_index]
+            if en_try.min_energy <= energy <= en_try.max_energy:
+                return mid_index
+            elif energy < en_try.min_energy:
+                max_index = mid_index - 1
+            else:
+                min_index = mid_index + 1
+        return None
+
 
 class LookupTable(RootModel[dict[Pol, EnergyCoverage]]):
     """
@@ -183,6 +209,8 @@ class LookupTable(RootModel[dict[Pol, EnergyCoverage]]):
     values to Apple2 motor positions.
     """
 
+    model_config = ConfigDict(frozen=True)
+
     # Allow to auto specify a dict if one not provided
     def __init__(self, root: dict[Pol, EnergyCoverage] | None = None):
         super().__init__(root=root or {})
@@ -195,10 +223,8 @@ class LookupTable(RootModel[dict[Pol, EnergyCoverage]]):
     ) -> Self:
         """Generate a LookupTable containing multiple EnergyCoverage
         for provided polarisations."""
-        lut = cls()
-        for i in range(len(pols)):
-            lut.root[pols[i]] = energy_coverage[i]
-        return lut
+        root_data = dict(zip(pols, energy_coverage, strict=False))
+        return cls(root=root_data)
 
     def get_poly(
         self,
@@ -239,15 +265,15 @@ def convert_csv_to_lookup(
     -----------
     LookupTable
     """
+    temp_mode_entries: dict[Pol, list[EnergyCoverageEntry]] = {}
 
-    def process_row(row: dict[str, Any], lut: LookupTable) -> None:
-        """Process a single row from the CSV file and update the lookup table."""
+    def process_row(row: dict[str, Any]) -> None:
+        """Process a single row from the CSV file and update the temporary entry list."""
         raw_mode_value = str(row[lut_config.mode]).lower()
         mode_value = Pol(
             lut_config.mode_name_convert.get(raw_mode_value, raw_mode_value)
         )
 
-        # Create polynomial object for energy-to-gap/phase conversion
         coefficients = np.poly1d([float(row[coef]) for coef in lut_config.poly_deg])
 
         energy_entry = EnergyCoverageEntry(
@@ -255,27 +281,31 @@ def convert_csv_to_lookup(
             max_energy=float(row[lut_config.max_energy]),
             poly=coefficients,
         )
-        if mode_value not in lut.root:
-            lut.root[mode_value] = EnergyCoverage(energy_entries=[energy_entry])
-        else:
-            lut.root[mode_value].energy_entries.append(energy_entry)
+
+        if mode_value not in temp_mode_entries:
+            temp_mode_entries[mode_value] = []
+
+        temp_mode_entries[mode_value].append(energy_entry)
 
     reader = csv.DictReader(read_file_and_skip(file_contents, skip_line_start_with))
-    lut = LookupTable()
 
     for row in reader:
         source = lut_config.source
         # If there are multiple source only convert requested.
         if source is None or row[source.column] == source.value:
-            process_row(row=row, lut=lut)
-
+            process_row(row=row)
     # Check if our LookupTable is empty after processing, raise error if it is.
-    if not lut.root:
+    if not temp_mode_entries:
         raise RuntimeError(
             "LookupTable content is empty, failed to convert the file contents to "
             "a LookupTable!"
         )
-    return lut
+
+    final_lut_root: dict[Pol, EnergyCoverage] = {}
+    for pol, entries in temp_mode_entries.items():
+        final_lut_root[pol] = EnergyCoverage.model_validate({"energy_entries": entries})
+
+    return LookupTable(root=final_lut_root)
 
 
 def read_file_and_skip(file: str, skip_line_start_with: str = "#") -> Generator[str]:

dodal/devices/insertion_device/polarisation.py

@@ -0,0 +1,36 @@
+import asyncio
+
+from bluesky.protocols import Locatable, Location
+from ophyd_async.core import (
+    AsyncStatus,
+    Reference,
+    StandardReadable,
+    StandardReadableFormat,
+)
+
+from dodal.devices.insertion_device import MAXIMUM_MOVE_TIME, Apple2Controller
+from dodal.devices.insertion_device.enum import Pol
+from dodal.log import LOGGER
+
+
+class InsertionDevicePolarisation(StandardReadable, Locatable[Pol]):
+    """Apple2 ID polarisation movable device."""
+
+    def __init__(self, id_controller: Apple2Controller, name: str = "") -> None:
+        self.polarisation = Reference(id_controller.polarisation)
+        self.polarisation_setpoint = Reference(id_controller.polarisation_setpoint)
+        super().__init__(name=name)
+
+        self.add_readables([self.polarisation()], StandardReadableFormat.HINTED_SIGNAL)
+
+    @AsyncStatus.wrap
+    async def set(self, pol: Pol) -> None:
+        LOGGER.info(f"Setting insertion device polarisation to {pol.name}")
+        await self.polarisation().set(pol, timeout=MAXIMUM_MOVE_TIME)
+
+    async def locate(self) -> Location[Pol]:
+        """Return the current polarisation"""
+        setpoint, readback = await asyncio.gather(
+            self.polarisation_setpoint().get_value(), self.polarisation().get_value()
+        )
+        return Location(setpoint=setpoint, readback=readback)

dodal/devices/oav/oav_detector.py

@@ -1,13 +1,17 @@
+import asyncio
 from enum import IntEnum
 
 from bluesky.protocols import Movable
 from ophyd_async.core import (
     DEFAULT_TIMEOUT,
     AsyncStatus,
+    DeviceMock,
+    DeviceVector,
     LazyMock,
     SignalR,
     SignalRW,
     StandardReadable,
+    default_mock_class,
     derived_signal_r,
     soft_signal_rw,
 )
@@ -22,6 +26,7 @@ from dodal.devices.oav.oav_parameters import (
 )
 from dodal.devices.oav.snapshots.snapshot import Snapshot
 from dodal.devices.oav.snapshots.snapshot_with_grid import SnapshotWithGrid
+from dodal.log import LOGGER
 
 
 class Coords(IntEnum):
@@ -56,6 +61,35 @@ class NullZoomController(BaseZoomController):
             await self.level.set(value, wait=True)
 
 
+class BeamCentreForZoom(StandardReadable):
+    """These PVs hold the beam centre on the OAV at each zoom level.
+
+    When the zoom level is changed the IOC will update the OAV overlay PVs to be at these positions."""
+
+    def __init__(
+        self, prefix: str, level_name_pv_suffix: str, centre_value_pv_suffix: str
+    ) -> None:
+        self.level_name = epics_signal_r(
+            str, f"{prefix}MP:SELECT.{level_name_pv_suffix}"
+        )
+        self.x_centre = epics_signal_rw(
+            float, f"{prefix}PBCX:VAL{centre_value_pv_suffix}"
+        )
+        self.y_centre = epics_signal_rw(
+            float, f"{prefix}PBCY:VAL{centre_value_pv_suffix}"
+        )
+        super().__init__()
+
+
+class InstantMovingZoom(DeviceMock["ZoomController"]):
+    """Mock behaviour that instantly moves the zoom."""
+
+    async def connect(self, device: "ZoomController") -> None:
+        """Mock signals to do an instant move on setpoint write."""
+        device.DELAY_BETWEEN_MOTORS_AND_IMAGE_UPDATING_S = 0.001  # type:ignore
+
+
+@default_mock_class(InstantMovingZoom)
 class ZoomController(BaseZoomController):
     """
     Device to control the zoom level. This should be set like
@@ -63,19 +97,49 @@ class ZoomController(BaseZoomController):
         oav.zoom_controller.set("1.0x")
 
     Note that changing the zoom may change the AD wiring on the associated OAV, as such
-    you should wait on any zoom changs to finish before changing the OAV wiring.
+    you should wait on any zoom changes to finish before changing the OAV wiring.
     """
 
+    DELAY_BETWEEN_MOTORS_AND_IMAGE_UPDATING_S = 2
+
     def __init__(self, prefix: str, name: str = "") -> None:
         self.percentage = epics_signal_rw(float, f"{prefix}ZOOMPOSCMD")
 
         # Level is the string description of the zoom level e.g. "1.0x" or "1.0"
         self.level = epics_signal_rw(str, f"{prefix}MP:SELECT")
+
         super().__init__(name=name)
 
     @AsyncStatus.wrap
     async def set(self, value: str):
         await self.level.set(value, wait=True)
+        LOGGER.info(
+            "Waiting {self.DELAY_BETWEEN_MOTORS_AND_IMAGE_UPDATING_S} seconds for zoom to be noticeable"
+        )
+        await asyncio.sleep(self.DELAY_BETWEEN_MOTORS_AND_IMAGE_UPDATING_S)
+
+
+class ZoomControllerWithBeamCentres(ZoomController):
+    def __init__(self, prefix: str, name: str = "") -> None:
+        level_to_centre_mapping = [
+            ("ZRST", "A"),
+            ("ONST", "B"),
+            ("TWST", "C"),
+            ("THST", "D"),
+            ("FRST", "E"),
+            ("FVST", "F"),
+            ("SXST", "G"),
+            ("SVST", "H"),
+        ]
+
+        self.beam_centres = DeviceVector(
+            {
+                i: BeamCentreForZoom(prefix, *level_to_centre_mapping[i])
+                for i in range(len(level_to_centre_mapping))
+            }
+        )
+
+        super().__init__(prefix, name)
 
 
 class OAV(StandardReadable):
@@ -118,6 +182,7 @@ class OAV(StandardReadable):
             self.zoom_controller = zoom_controller
 
         self.cam = Cam(f"{prefix}CAM:", name=name)
+
         with self.add_children_as_readables():
             self.grid_snapshot = SnapshotWithGrid(
                 f"{prefix}{mjpeg_prefix}:", name, mjpg_x_size_pv, mjpg_y_size_pv

dodal/devices/oav/utils.py

@@ -2,6 +2,7 @@ from collections.abc import Generator
 from enum import IntEnum
 
 import bluesky.plan_stubs as bps
+import cv2
 import numpy as np
 from bluesky.utils import Msg
 
@@ -119,3 +120,19 @@ def wait_for_tip_to_be_found(
         raise PinNotFoundError(f"No pin found after {timeout} seconds")
 
     return Pixel((int(found_tip[0]), int(found_tip[1])))
+
+
+def convert_to_gray_and_blur(data: cv2.typing.MatLike) -> cv2.typing.MatLike:
+    """
+    Preprocess the image array data (convert to grayscale and apply a gaussian blur)
+    Image is converted to grayscale (using a weighted mean as green contributes more to brightness)
+    as we aren't interested in data relating to colour. A blur is then applied to mitigate
+    errors due to rogue hot pixels.
+    """
+
+    # kernel size describes how many of the neighbouring pixels are used for the blur,
+    # higher kernal size means more of a blur effect
+    kernel_size = (7, 7)
+
+    gray_arr = cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)
+    return cv2.GaussianBlur(gray_arr, kernel_size, 0)

dodal/devices/robot.py

@@ -1,6 +1,7 @@
 import asyncio
 from asyncio import FIRST_COMPLETED, CancelledError, Task, wait_for
 from dataclasses import dataclass
+from enum import IntEnum
 
 from bluesky.protocols import Movable
 from ophyd_async.core import (
@@ -21,8 +22,8 @@ from ophyd_async.epics.core import (
 
 from dodal.log import LOGGER
 
-WAIT_FOR_OLD_PIN_MSG = "Waiting on old pin unloaded"
-WAIT_FOR_NEW_PIN_MSG = "Waiting on new pin loaded"
+WAIT_FOR_BEAMLINE_DISABLE_MSG = "Waiting on beamline disable"
+WAIT_FOR_BEAMLINE_ENABLE_MSG = "Waiting on beamline enable"
 
 
 class RobotLoadError(Exception):
@@ -43,11 +44,19 @@ class SampleLocation:
     pin: int
 
 
+SAMPLE_LOCATION_EMPTY = SampleLocation(-1, -1)
+
+
 class PinMounted(StrictEnum):
     NO_PIN_MOUNTED = "No Pin Mounted"
     PIN_MOUNTED = "Pin Mounted"
 
 
+class BeamlineStatus(IntEnum):
+    ENABLED = 0
+    DISABLED = 1
+
+
 class ErrorStatus(Device):
     def __init__(self, prefix: str) -> None:
         self.str = epics_signal_r(str, prefix + "_ERR_MSG")
@@ -61,7 +70,7 @@ class ErrorStatus(Device):
             raise RobotLoadError(int(error_code), error_string) from raise_from
 
 
-class BartRobot(StandardReadable, Movable[SampleLocation | None]):
+class BartRobot(StandardReadable, Movable[SampleLocation]):
     """The sample changing robot."""
 
     # How long to wait for the robot if it is busy soaking/drying
@@ -86,6 +95,8 @@ class BartRobot(StandardReadable, Movable[SampleLocation | None]):
             self.current_puck = epics_signal_r(float, prefix + "CURRENT_PUCK_RBV")
             self.current_pin = epics_signal_r(float, prefix + "CURRENT_PIN_RBV")
 
+        self.beamline_disabled = epics_signal_r(int, prefix + "ROBOT_OP_16_BITS.B8")
+
         self.next_pin = epics_signal_rw_rbv(float, prefix + "NEXT_PIN")
         self.next_puck = epics_signal_rw_rbv(float, prefix + "NEXT_PUCK")
 
@@ -116,8 +127,8 @@ class BartRobot(StandardReadable, Movable[SampleLocation | None]):
         )
         super().__init__(name=name)
 
-    async def pin_state_or_error(self, expected_state=PinMounted.PIN_MOUNTED):
-        """This co-routine will finish when either the pin sensor reaches the specified
+    async def beamline_status_or_error(self, expected_state: BeamlineStatus):
+        """This co-routine will finish when either the beamline reaches the specified
         state or the robot gives an error (whichever happens first). In the case where
         there is an error a RobotLoadError error is raised.
         """
@@ -130,12 +141,12 @@ class BartRobot(StandardReadable, Movable[SampleLocation | None]):
             error_msg = await self.prog_error.str.get_value()
             raise RobotLoadError(error_code, error_msg)
 
-        async def wfv():
-            await wait_for_value(self.gonio_pin_sensor, expected_state, None)
+        async def wait_for_expected_state():
+            await wait_for_value(self.beamline_disabled, expected_state.value, None)
 
         tasks = [
             (Task(raise_if_error())),
-            (Task(wfv())),
+            (Task(wait_for_expected_state())),
         ]
         try:
             finished, unfinished = await asyncio.wait(
@@ -171,31 +182,37 @@ class BartRobot(StandardReadable, Movable[SampleLocation | None]):
             set_and_wait_for_value(self.next_pin, sample_location.pin),
         )
         await self.load.trigger()
-        if await self.gonio_pin_sensor.get_value() == PinMounted.PIN_MOUNTED:
-            LOGGER.info(WAIT_FOR_OLD_PIN_MSG)
-            await self.pin_state_or_error(PinMounted.NO_PIN_MOUNTED)
-        LOGGER.info(WAIT_FOR_NEW_PIN_MSG)
+        await self._wait_for_beamline_enabled_after_load_or_unload()
 
-        await self.pin_state_or_error()
+    async def _wait_for_beamline_enabled_after_load_or_unload(self):
+        if await self.beamline_disabled.get_value() == BeamlineStatus.ENABLED.value:
+            LOGGER.info(WAIT_FOR_BEAMLINE_DISABLE_MSG)
+            await self.beamline_status_or_error(BeamlineStatus.DISABLED)
+
+        LOGGER.info(WAIT_FOR_BEAMLINE_ENABLE_MSG)
+        await self.beamline_status_or_error(BeamlineStatus.ENABLED)
 
     @AsyncStatus.wrap
-    async def set(self, value: SampleLocation | None):
+    async def set(self, value: SampleLocation):
         """
         Perform a sample load from the specified sample location
         Args:
-            value: The pin and puck to load, or None to unload the sample.
+            value: The pin and puck to load, or SAMPLE_LOCATION_EMPTY to unload the sample.
         Raises:
-            RobotLoadError if a timeout occurs, or if an error occurs loading the smaple.
+            RobotLoadError if a timeout occurs, or if an error occurs loading the sample.
         """
         try:
-            if value is not None:
+            if value != SAMPLE_LOCATION_EMPTY:
                 await wait_for(
                     self._load_pin_and_puck(value),
                     timeout=self.LOAD_TIMEOUT + self.NOT_BUSY_TIMEOUT,
                 )
             else:
                 await self.unload.trigger(timeout=self.LOAD_TIMEOUT)
-                await wait_for_value(self.program_running, False, self.NOT_BUSY_TIMEOUT)
+                await wait_for(
+                    self._wait_for_beamline_enabled_after_load_or_unload(),
+                    timeout=self.LOAD_TIMEOUT + self.NOT_BUSY_TIMEOUT,
+                )
         except TimeoutError as e:
             await self.prog_error.raise_if_error(e)
             await self.controller_error.raise_if_error(e)

dodal/devices/selectable_source.py

@@ -0,0 +1,38 @@
+from typing import TypeVar
+
+from bluesky.protocols import Movable
+from ophyd_async.core import AsyncStatus, StandardReadable, StrictEnum, soft_signal_rw
+
+
+class SelectedSource(StrictEnum):
+    SOURCE1 = "source1"
+    SOURCE2 = "source2"
+
+
+T = TypeVar("T")
+
+
+def get_obj_from_selected_source(selected_source: SelectedSource, s1: T, s2: T) -> T:
+    """Util function that maps enum values for SelectedSource to two objects. It then
+    returns one of the objects that corrosponds to the selected_source value."""
+    match selected_source:
+        case SelectedSource.SOURCE1:
+            return s1
+        case SelectedSource.SOURCE2:
+            return s2
+
+
+class SourceSelector(StandardReadable, Movable[SelectedSource]):
+    """Device that holds a selected_source signal enum of SelectedSource. Useful for
+    beamlines with multiple sources to coordinate which energy source or shutter to use."""
+
+    def __init__(self, name: str = ""):
+        with self.add_children_as_readables():
+            self.selected_source = soft_signal_rw(
+                SelectedSource, SelectedSource.SOURCE1
+            )
+        super().__init__(name)
+
+    @AsyncStatus.wrap
+    async def set(self, value: SelectedSource):
+        await self.selected_source.set(value)