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

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.id_enum import Pol
 
 DEFAULT_POLY_DEG = [
     "7th-order",
@@ -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/testing/electron_analyser/device_factory.py

@@ -1,16 +1,16 @@
 from typing import Any, get_args, get_origin
 
-from dodal.devices.electron_analyser.abstract import (
+from dodal.devices.electron_analyser.base.base_detector import TElectronAnalyserDetector
+from dodal.devices.electron_analyser.base.base_driver_io import (
     TAbstractAnalyserDriverIO,
 )
-from dodal.devices.electron_analyser.detector import TElectronAnalyserDetector
 from dodal.devices.electron_analyser.vgscienta import (
     VGScientaAnalyserDriverIO,
     VGScientaDetector,
 )
 
 
-async def create_driver(
+def create_driver(
     driver_class: type[TAbstractAnalyserDriverIO],
     **kwargs: Any,
 ) -> TAbstractAnalyserDriverIO:
@@ -34,7 +34,7 @@ async def create_driver(
     return driver_class(**(parameters | kwargs))
 
 
-async def create_detector(
+def create_detector(
     detector_class: type[TElectronAnalyserDetector],
     **kwargs: Any,
 ) -> TElectronAnalyserDetector:

dodal/testing/fixtures/devices/apple2.py

@@ -7,7 +7,7 @@ from ophyd_async.core import (
     set_mock_value,
 )
 
-from dodal.devices.insertion_device.apple2_undulator import (
+from dodal.devices.insertion_device import (
     EnabledDisabledUpper,
     UndulatorGap,
     UndulatorGateStatus,

dodal/testing/fixtures/run_engine.py

@@ -3,6 +3,7 @@ Allow external repos to reuse these fixtures so defined in single place.
 """
 
 import asyncio
+import copy
 import os
 import threading
 import time
@@ -34,10 +35,13 @@ async def _ensure_running_bluesky_event_loop(_global_run_engine):
 
 @pytest.fixture()
 async def run_engine(_global_run_engine: RunEngine) -> AsyncGenerator[RunEngine, None]:
+    initial_md = copy.deepcopy(_global_run_engine.md)
     try:
         yield _global_run_engine
     finally:
+        # Clear subscriptions, cache, and reset metadata
         _global_run_engine.reset()
+        _global_run_engine.md = initial_md
 
 
 @pytest_asyncio.fixture(scope="session", loop_scope="session")

dodal/devices/electron_analyser/abstract/__init__.py

@@ -1,25 +0,0 @@
-from .base_detector import (
-    BaseElectronAnalyserDetector,
-)
-from .base_driver_io import AbstractAnalyserDriverIO, TAbstractAnalyserDriverIO
-from .base_region import (
-    AbstractBaseRegion,
-    AbstractBaseSequence,
-    TAbstractBaseRegion,
-    TAbstractBaseSequence,
-    TAcquisitionMode,
-    TLensMode,
-)
-
-__all__ = [
-    "AbstractBaseRegion",
-    "AbstractBaseSequence",
-    "TAbstractBaseRegion",
-    "TAbstractBaseSequence",
-    "TAcquisitionMode",
-    "TLensMode",
-    "AbstractAnalyserDriverIO",
-    "BaseElectronAnalyserDetector",
-    "AbstractAnalyserDriverIO",
-    "TAbstractAnalyserDriverIO",
-]

dodal/devices/electron_analyser/abstract/base_detector.py

@@ -1,63 +0,0 @@
-from typing import Generic
-
-from bluesky.protocols import Reading, Triggerable
-from event_model import DataKey
-from ophyd_async.core import (
-    AsyncConfigurable,
-    AsyncReadable,
-    AsyncStatus,
-    Device,
-)
-from ophyd_async.epics.adcore import ADBaseController
-
-from dodal.devices.electron_analyser.abstract.base_driver_io import (
-    TAbstractAnalyserDriverIO,
-)
-
-
-class BaseElectronAnalyserDetector(
-    Device,
-    Triggerable,
-    AsyncReadable,
-    AsyncConfigurable,
-    Generic[TAbstractAnalyserDriverIO],
-):
-    """
-    Detector for data acquisition of electron analyser. Can only acquire using settings
-    already configured for the device.
-
-    If possible, this should be changed to inherit from a StandardDetector. Currently,
-    StandardDetector forces you to use a file writer which doesn't apply here.
-    See issue https://github.com/bluesky/ophyd-async/issues/888
-    """
-
-    def __init__(
-        self,
-        controller: ADBaseController[TAbstractAnalyserDriverIO],
-        name: str = "",
-    ):
-        self._controller = controller
-        super().__init__(name)
-
-    @AsyncStatus.wrap
-    async def trigger(self) -> None:
-        await self._controller.arm()
-        await self._controller.wait_for_idle()
-
-    async def read(self) -> dict[str, Reading]:
-        return await self._controller.driver.read()
-
-    async def describe(self) -> dict[str, DataKey]:
-        data = await self._controller.driver.describe()
-        # Correct the shape for image
-        prefix = self._controller.driver.name + "-"
-        energy_size = len(await self._controller.driver.energy_axis.get_value())
-        angle_size = len(await self._controller.driver.angle_axis.get_value())
-        data[prefix + "image"]["shape"] = [angle_size, energy_size]
-        return data
-
-    async def read_configuration(self) -> dict[str, Reading]:
-        return await self._controller.driver.read_configuration()
-
-    async def describe_configuration(self) -> dict[str, DataKey]:
-        return await self._controller.driver.describe_configuration()

dodal/devices/electron_analyser/abstract/types.py

@@ -1,12 +0,0 @@
-from typing import TypeVar
-
-from ophyd_async.core import StrictEnum, SupersetEnum
-
-TAcquisitionMode = TypeVar("TAcquisitionMode", bound=StrictEnum)
-# Allow SupersetEnum. Specs analysers can connect to Lens and Psu mode separately to the
-# analyser which leaves the enum to either be "Not connected" OR the available enums
-# when connected.
-TLensMode = TypeVar("TLensMode", bound=SupersetEnum | StrictEnum)
-TPsuMode = TypeVar("TPsuMode", bound=SupersetEnum | StrictEnum)
-TPassEnergy = TypeVar("TPassEnergy", bound=StrictEnum | float)
-TPassEnergyEnum = TypeVar("TPassEnergyEnum", bound=StrictEnum)

dodal/devices/electron_analyser/detector.py

@@ -1,143 +0,0 @@
-from typing import Generic, TypeVar
-
-from bluesky.protocols import Stageable
-from ophyd_async.core import AsyncStatus
-from ophyd_async.epics.adcore import ADBaseController
-
-from dodal.common.data_util import load_json_file_to_class
-from dodal.devices.controllers import ConstantDeadTimeController
-from dodal.devices.electron_analyser.abstract.base_detector import (
-    BaseElectronAnalyserDetector,
-)
-from dodal.devices.electron_analyser.abstract.base_driver_io import (
-    TAbstractAnalyserDriverIO,
-)
-from dodal.devices.electron_analyser.abstract.base_region import (
-    TAbstractBaseRegion,
-    TAbstractBaseSequence,
-)
-
-
-class ElectronAnalyserRegionDetector(
-    BaseElectronAnalyserDetector[TAbstractAnalyserDriverIO],
-    Generic[TAbstractAnalyserDriverIO, TAbstractBaseRegion],
-):
-    """
-    Extends electron analyser detector to configure specific region settings before data
-    acquisition. It is designed to only exist inside a plan.
-    """
-
-    def __init__(
-        self,
-        controller: ADBaseController[TAbstractAnalyserDriverIO],
-        region: TAbstractBaseRegion,
-        name: str = "",
-    ):
-        self.region = region
-        super().__init__(controller, name)
-
-    @AsyncStatus.wrap
-    async def trigger(self) -> None:
-        # Configure region parameters on the driver first before data collection.
-        await self._controller.driver.set(self.region)
-        await super().trigger()
-
-
-TElectronAnalyserRegionDetector = TypeVar(
-    "TElectronAnalyserRegionDetector",
-    bound=ElectronAnalyserRegionDetector,
-)
-
-
-class ElectronAnalyserDetector(
-    BaseElectronAnalyserDetector[TAbstractAnalyserDriverIO],
-    Stageable,
-    Generic[
-        TAbstractAnalyserDriverIO,
-        TAbstractBaseSequence,
-        TAbstractBaseRegion,
-    ],
-):
-    """
-    Electron analyser detector with the additional functionality to load a sequence file
-    and create a list of temporary ElectronAnalyserRegionDetector objects. These will
-    setup configured region settings before data acquisition.
-    """
-
-    def __init__(
-        self,
-        sequence_class: type[TAbstractBaseSequence],
-        driver: TAbstractAnalyserDriverIO,
-        name: str = "",
-    ):
-        # Save driver as direct child so participates with connect()
-        self.driver = driver
-        self._sequence_class = sequence_class
-        controller = ConstantDeadTimeController[TAbstractAnalyserDriverIO](driver, 0)
-        super().__init__(controller, name)
-
-    @AsyncStatus.wrap
-    async def stage(self) -> None:
-        """
-        Prepare the detector for use by ensuring it is idle and ready.
-
-        This method asynchronously stages the detector by first disarming the controller
-        to ensure the detector is not actively acquiring data, then invokes the driver's
-        stage procedure. This ensures the detector is in a known, ready state
-        before use.
-
-        Raises:
-            Any exceptions raised by the driver's stage or controller's disarm methods.
-        """
-        await self._controller.disarm()
-        await self.driver.stage()
-
-    @AsyncStatus.wrap
-    async def unstage(self) -> None:
-        """Disarm the detector."""
-        await self._controller.disarm()
-        await self.driver.unstage()
-
-    def load_sequence(self, filename: str) -> TAbstractBaseSequence:
-        """
-        Load the sequence data from a provided json file into a sequence class.
-
-        Args:
-            filename: Path to the sequence file containing the region data.
-
-        Returns:
-            Pydantic model representing the sequence file.
-        """
-        return load_json_file_to_class(self._sequence_class, filename)
-
-    def create_region_detector_list(
-        self, filename: str, enabled_only=True
-    ) -> list[
-        ElectronAnalyserRegionDetector[TAbstractAnalyserDriverIO, TAbstractBaseRegion]
-    ]:
-        """
-        Create a list of detectors equal to the number of regions in a sequence file.
-        Each detector is responsible for setting up a specific region.
-
-        Args:
-            filename:     Path to the sequence file containing the region data.
-            enabled_only: If true, only include the region if enabled is True.
-
-        Returns:
-            List of ElectronAnalyserRegionDetector, equal to the number of regions in
-            the sequence file.
-        """
-        seq = self.load_sequence(filename)
-        regions = seq.get_enabled_regions() if enabled_only else seq.regions
-        return [
-            ElectronAnalyserRegionDetector(
-                self._controller, r, self.name + "_" + r.name
-            )
-            for r in regions
-        ]
-
-
-TElectronAnalyserDetector = TypeVar(
-    "TElectronAnalyserDetector",
-    bound=ElectronAnalyserDetector,
-)

dodal/devices/electron_analyser/specs/detector.py

@@ -1,34 +0,0 @@
-from typing import Generic
-
-from dodal.devices.electron_analyser.abstract.types import TLensMode, TPsuMode
-from dodal.devices.electron_analyser.detector import (
-    ElectronAnalyserDetector,
-)
-from dodal.devices.electron_analyser.energy_sources import (
-    DualEnergySource,
-    EnergySource,
-)
-from dodal.devices.electron_analyser.specs.driver_io import SpecsAnalyserDriverIO
-from dodal.devices.electron_analyser.specs.region import SpecsRegion, SpecsSequence
-
-
-class SpecsDetector(
-    ElectronAnalyserDetector[
-        SpecsAnalyserDriverIO[TLensMode, TPsuMode],
-        SpecsSequence[TLensMode, TPsuMode],
-        SpecsRegion[TLensMode, TPsuMode],
-    ],
-    Generic[TLensMode, TPsuMode],
-):
-    def __init__(
-        self,
-        prefix: str,
-        lens_mode_type: type[TLensMode],
-        psu_mode_type: type[TPsuMode],
-        energy_source: DualEnergySource | EnergySource,
-        name: str = "",
-    ):
-        driver = SpecsAnalyserDriverIO[TLensMode, TPsuMode](
-            prefix, lens_mode_type, psu_mode_type, energy_source
-        )
-        super().__init__(SpecsSequence[lens_mode_type, psu_mode_type], driver, name)

dodal/devices/electron_analyser/types.py

@@ -1,57 +0,0 @@
-from ophyd_async.core import StrictEnum, SupersetEnum
-
-from dodal.devices.electron_analyser.abstract.base_driver_io import (
-    AbstractAnalyserDriverIO,
-)
-from dodal.devices.electron_analyser.abstract.base_region import (
-    AbstractBaseRegion,
-    AbstractBaseSequence,
-)
-from dodal.devices.electron_analyser.detector import (
-    ElectronAnalyserDetector,
-    ElectronAnalyserRegionDetector,
-)
-from dodal.devices.electron_analyser.specs.detector import (
-    SpecsAnalyserDriverIO,
-    SpecsDetector,
-)
-from dodal.devices.electron_analyser.vgscienta.detector import (
-    VGScientaAnalyserDriverIO,
-    VGScientaDetector,
-)
-
-AnyAcqMode = StrictEnum
-AnyLensMode = SupersetEnum | StrictEnum
-AnyPsuMode = SupersetEnum | StrictEnum
-AnyPassEnergy = StrictEnum | float
-AnyPassEnergyEnum = StrictEnum
-
-# Electron analyser types that encompasses all implementations, useful for tests and
-# plans
-ElectronAnalyserDetectorImpl = (
-    VGScientaDetector[AnyLensMode, AnyPsuMode, AnyPassEnergyEnum]
-    | SpecsDetector[AnyLensMode, AnyPsuMode]
-)
-ElectronAnalyserDriverImpl = (
-    VGScientaAnalyserDriverIO[AnyLensMode, AnyPsuMode, AnyPassEnergyEnum]
-    | SpecsAnalyserDriverIO[AnyLensMode, AnyPsuMode]
-)
-
-# Short hand the type so less verbose
-AbstractBaseRegion = AbstractBaseRegion[AnyAcqMode, AnyLensMode, AnyPassEnergy]
-
-# Generic electron analyser types that supports full typing with the abstract classes.
-GenericElectronAnalyserDetector = ElectronAnalyserDetector[
-    AbstractAnalyserDriverIO[
-        AbstractBaseRegion, AnyAcqMode, AnyLensMode, AnyPsuMode, AnyPassEnergy
-    ],
-    AbstractBaseSequence[AbstractBaseRegion],
-    AbstractBaseRegion,
-]
-
-GenericElectronAnalyserRegionDetector = ElectronAnalyserRegionDetector[
-    AbstractAnalyserDriverIO[
-        AbstractBaseRegion, AnyAcqMode, AnyLensMode, AnyPsuMode, AnyPassEnergy
-    ],
-    AbstractBaseRegion,
-]

dodal/devices/electron_analyser/vgscienta/detector.py

@@ -1,48 +0,0 @@
-from typing import Generic
-
-from dodal.devices.electron_analyser.abstract.types import (
-    TLensMode,
-    TPassEnergyEnum,
-    TPsuMode,
-)
-from dodal.devices.electron_analyser.detector import (
-    ElectronAnalyserDetector,
-)
-from dodal.devices.electron_analyser.energy_sources import (
-    DualEnergySource,
-    EnergySource,
-)
-from dodal.devices.electron_analyser.vgscienta.driver_io import (
-    VGScientaAnalyserDriverIO,
-)
-from dodal.devices.electron_analyser.vgscienta.region import (
-    VGScientaRegion,
-    VGScientaSequence,
-)
-
-
-class VGScientaDetector(
-    ElectronAnalyserDetector[
-        VGScientaAnalyserDriverIO[TLensMode, TPsuMode, TPassEnergyEnum],
-        VGScientaSequence[TLensMode, TPsuMode, TPassEnergyEnum],
-        VGScientaRegion[TLensMode, TPassEnergyEnum],
-    ],
-    Generic[TLensMode, TPsuMode, TPassEnergyEnum],
-):
-    def __init__(
-        self,
-        prefix: str,
-        lens_mode_type: type[TLensMode],
-        psu_mode_type: type[TPsuMode],
-        pass_energy_type: type[TPassEnergyEnum],
-        energy_source: DualEnergySource | EnergySource,
-        name: str = "",
-    ):
-        driver = VGScientaAnalyserDriverIO[TLensMode, TPsuMode, TPassEnergyEnum](
-            prefix, lens_mode_type, psu_mode_type, pass_energy_type, energy_source
-        )
-        super().__init__(
-            VGScientaSequence[lens_mode_type, psu_mode_type, pass_energy_type],
-            driver,
-            name,
-        )