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

dodal/devices/i19/access_controlled/attenuator_motor_squad.py

@@ -0,0 +1,61 @@
+from typing import Annotated, Any, Self
+
+from ophyd_async.core import AsyncStatus
+from pydantic import BaseModel, model_validator
+from pydantic.types import PositiveInt, StringConstraints
+
+from dodal.devices.i19.access_controlled.blueapi_device import (
+    OpticsBlueAPIDevice,
+)
+from dodal.devices.i19.access_controlled.hutch_access import ACCESS_DEVICE_NAME
+
+PermittedKeyStr = Annotated[str, StringConstraints(pattern="^[A-Za-z0-9-_]*$")]
+
+
+class AttenuatorMotorPositionDemands(BaseModel):
+    continuous_demands: dict[PermittedKeyStr, float] = {}
+    indexed_demands: dict[PermittedKeyStr, PositiveInt] = {}
+
+    @model_validator(mode="after")
+    def no_keys_clash(self) -> Self:
+        common_keys = set(self.continuous_demands).intersection(self.indexed_demands)
+        common_key_count = sum(1 for _ in common_keys)
+        if common_key_count < 1:
+            return self
+        else:
+            ks: str = "key" if common_key_count == 1 else "keys"
+            error_msg = f"Common {ks} found in distinct motor demands: {common_keys}"
+            raise ValueError(error_msg)
+
+    def validated_complete_demand(self) -> dict[PermittedKeyStr, Any]:
+        return self.continuous_demands | self.indexed_demands
+
+
+class AttenuatorMotorSquad(OpticsBlueAPIDevice):
+    """ I19-specific proxy device which requests absorber position changes in the x-ray attenuator.
+
+    Sends REST call to blueapi controlling optics on the I19 cluster.
+     The hutch in use is compared against the hutch which sent the REST call.
+    Only the hutch in use will be permitted to execute a plan (requesting motor moves).
+    As the two hutches are located in series, checking the hutch in use is necessary to \
+    avoid accidentally operating optics devices from one hutch while the other has beam time.
+
+    The name of the hutch that wants to operate the optics device is passed to the \
+    access controlled device upon instantiation of the latter.
+
+    For details see the architecture described in \
+    https://github.com/DiamondLightSource/i19-bluesky/issues/30.
+    """
+
+    @AsyncStatus.wrap
+    async def set(self, value: AttenuatorMotorPositionDemands):
+        request_params = {
+            "name": "operate_motor_squad_plan",
+            "params": {
+                "experiment_hutch": self._invoking_hutch,
+                "access_device": ACCESS_DEVICE_NAME,
+                "attenuator_demands": value.validated_complete_demand(),
+            },
+            "instrument_session": self.instrument_session,
+        }
+        await super().set(request_params)

dodal/devices/i19/access_controlled/blueapi_device.py

@@ -29,11 +29,19 @@ class OpticsBlueAPIDevice(StandardReadable, Movable[D]):
     https://github.com/DiamondLightSource/i19-bluesky/issues/30.
     """
 
-    def __init__(self, name: str = "") -> None:
+    def __init__(
+        self, hutch: HutchState, instrument_session: str = "", name: str = ""
+    ) -> None:
+        self.hutch_request = hutch
+        self.instrument_session = instrument_session
         self.url = OPTICS_BLUEAPI_URL
         self.headers = HEADERS
         super().__init__(name)
 
+    @property
+    def _invoking_hutch(self) -> str:
+        return self.hutch_request.value
+
     @AsyncStatus.wrap
     async def set(self, value: D):
         """ On set send a POST request to the optics blueapi with the name and \

dodal/devices/i19/access_controlled/shutter.py

@@ -36,16 +36,14 @@ class AccessControlledShutter(OpticsBlueAPIDevice):
         # see https://github.com/DiamondLightSource/blueapi/issues/1187
         with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
             self.shutter_status = epics_signal_r(ShutterState, f"{prefix}STA")
-        self.hutch_request = hutch
-        self.instrument_session = instrument_session
-        super().__init__(name)
+        super().__init__(hutch=hutch, instrument_session=instrument_session, name=name)
 
     @AsyncStatus.wrap
     async def set(self, value: ShutterDemand):
         request_params = {
             "name": "operate_shutter_plan",
             "params": {
-                "experiment_hutch": self.hutch_request.value,
+                "experiment_hutch": self._invoking_hutch,
                 "access_device": ACCESS_DEVICE_NAME,
                 "shutter_demand": value.value,
             },

dodal/devices/{apple2_undulator.py → insertion_device/apple2_undulator.py}

@@ -398,7 +398,7 @@ class Apple2(StandardReadable, Movable[Apple2Val], Generic[PhaseAxesType]):
 
 
 class EnergyMotorConvertor(Protocol):
-    def __call__(self, energy: float, pol: Pol) -> tuple[float, float]:
+    def __call__(self, energy: float, pol: Pol) -> float:
         """Protocol to provide energy to motor position conversion"""
         ...
 
@@ -426,19 +426,18 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         Soft signal for the polarisation setpoint.
     polarisation : derived_signal_rw
         Hardware-backed signal for polarisation readback and control.
-    energy_to_motor : EnergyMotorConvertor
-        Callable that converts energy and polarisation to motor positions.
+    gap_energy_to_motor_converter : EnergyMotorConvertor
+        Callable that converts energy and polarisation to gap motor positions.
+    phase_energy_to_motor_converter : EnergyMotorConvertor
+        Callable that converts energy and polarisation to phase motor positions.
 
     Abstract Methods
     ----------------
-    _set_motors_from_energy(value: float) -> None
-        Abstract method to set motor positions for a given energy and polarisation.
-    energy_to_motor : EnergyMotorConvertor
-        A callable that converts energy and polarisation to motor positions.
-
+    _get_apple2_value(gap: float, phase: float) -> Apple2Val
+        Abstract method to return the Apple2Val used to set the apple2 with.
     Notes
     -----
-    - Subclasses must implement `_set_motors_from_energy` for beamline-specific logic.
+    - Subclasses must implement `_get_apple2_value` for beamline-specific logic.
     - LH3 polarisation is indistinguishable from LH in hardware; special handling is provided.
     - Supports multiple polarisation modes, including linear horizontal (LH), linear vertical (LV),
       positive circular (PC), negative circular (NC), and linear arbitrary (LA).
@@ -448,7 +447,9 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
     def __init__(
         self,
         apple2: Apple2Type,
-        energy_to_motor_converter: EnergyMotorConvertor,
+        gap_energy_motor_converter: EnergyMotorConvertor,
+        phase_energy_motor_converter: EnergyMotorConvertor,
+        units: str = "eV",
         name: str = "",
     ) -> None:
         """
@@ -460,19 +461,20 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         name: str
             Name of the device.
         """
-        self.energy_to_motor = energy_to_motor_converter
         self.apple2 = Reference(apple2)
+        self.gap_energy_motor_converter = gap_energy_motor_converter
+        self.phase_energy_motor_converter = phase_energy_motor_converter
 
         # Store the set energy for readback.
         self._energy, self._energy_set = soft_signal_r_and_setter(
-            float, initial_value=None, units="eV"
+            float, initial_value=None, units=units
         )
         with self.add_children_as_readables(StandardReadableFormat.HINTED_SIGNAL):
             self.energy = derived_signal_rw(
                 raw_to_derived=self._read_energy,
                 set_derived=self._set_energy,
                 energy=self._energy,
-                derived_units="eV",
+                derived_units=units,
             )
 
         # Store the polarisation for setpoint. And provide readback for LH3.
@@ -480,10 +482,11 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         self.polarisation_setpoint, self._polarisation_setpoint_set = (
             soft_signal_r_and_setter(Pol)
         )
+        phase = self.apple2().phase()
         # check if undulator phase is unlocked.
-        if isinstance(self.apple2().phase(), UndulatorPhaseAxes):
-            top_inner = self.apple2().phase().top_inner.user_readback
-            btm_outer = self.apple2().phase().btm_outer.user_readback
+        if isinstance(phase, UndulatorPhaseAxes):
+            top_inner = phase.top_inner.user_readback
+            btm_outer = phase.btm_outer.user_readback
         else:
             # If locked phase axes make the locked phase 0.
             top_inner = btm_outer = soft_signal_rw(float, initial_value=0.0)
@@ -495,24 +498,35 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
                 raw_to_derived=self._read_pol,
                 set_derived=self._set_pol,
                 pol=self.polarisation_setpoint,
-                top_outer=self.apple2().phase().top_outer.user_readback,
+                top_outer=phase.top_outer.user_readback,
                 top_inner=top_inner,
-                btm_inner=self.apple2().phase().btm_inner.user_readback,
+                btm_inner=phase.btm_inner.user_readback,
                 btm_outer=btm_outer,
                 gap=self.apple2().gap().user_readback,
             )
         super().__init__(name)
 
     @abc.abstractmethod
-    async def _set_motors_from_energy(self, value: float) -> None:
+    def _get_apple2_value(self, gap: float, phase: float, pol: Pol) -> Apple2Val:
         """
         This method should be implemented by the beamline specific ID class as the
         motor positions will be different for each beamline depending on the
-        undulator design and the lookup table used.
+        undulator design.
         """
 
+    async def _set_motors_from_energy_and_polarisation(
+        self, energy: float, pol: Pol
+    ) -> None:
+        """Set the undulator motors for a given energy and polarisation."""
+        gap = self.gap_energy_motor_converter(energy=energy, pol=pol)
+        phase = self.phase_energy_motor_converter(energy=energy, pol=pol)
+        apple2_val = self._get_apple2_value(gap, phase, pol)
+        LOGGER.info(f"Setting polarisation to {pol}, with values: {apple2_val}")
+        await self.apple2().set(id_motor_values=apple2_val)
+
     async def _set_energy(self, energy: float) -> None:
-        await self._set_motors_from_energy(energy)
+        pol = await self._check_and_get_pol_setpoint()
+        await self._set_motors_from_energy_and_polarisation(energy, pol)
         self._energy_set(energy)
 
     def _read_energy(self, energy: float) -> float:
@@ -524,7 +538,6 @@ class Apple2Controller(abc.ABC, StandardReadable, Generic[Apple2Type]):
         Check the polarisation setpoint and if it is NONE try to read it from
         hardware.
         """
-
         pol = await self.polarisation_setpoint.get_value()
 
         if pol == Pol.NONE:
@@ -717,12 +730,7 @@ class InsertionDeviceEnergy(InsertionDeviceEnergyBase):
         self.energy = Reference(id_controller.energy)
         super().__init__(name=name)
 
-        self.add_readables(
-            [
-                self.energy(),
-            ],
-            StandardReadableFormat.HINTED_SIGNAL,
-        )
+        self.add_readables([self.energy()], StandardReadableFormat.HINTED_SIGNAL)
 
     @AsyncStatus.wrap
     async def set(self, energy: float) -> None:

dodal/devices/insertion_device/energy_motor_lookup.py

@@ -0,0 +1,88 @@
+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.lookup_table_models import (
+    LookupTable,
+    LookupTableColumnConfig,
+    convert_csv_to_lookup,
+)
+
+
+class EnergyMotorLookup:
+    """
+    Handles a lookup table for Apple2 ID, converting energy/polarisation to a motor
+    position.
+
+    After update_lookup_table() has populated the lookup table, `find_value_in_lookup_table()`
+    can be used to compute gap / phase for a requested energy / polarisation pair.
+    """
+
+    def __init__(self, lut: LookupTable | None = None):
+        if lut is None:
+            lut = LookupTable()
+        self.lut = lut
+
+    def update_lookup_table(self) -> None:
+        """Do nothing by default. Sub classes may override this method to provide logic
+        on what updating lookup table does."""
+        pass
+
+    def find_value_in_lookup_table(self, energy: float, pol: Pol) -> float:
+        """
+        Convert energy and polarisation to a value from the lookup table.
+
+        Parameters:
+        -----------
+        energy : float
+            Desired energy.
+        pol : Pol
+            Polarisation mode.
+
+        Returns:
+        ----------
+        float
+            gap / phase motor position from the lookup table.
+        """
+        # if lut is empty, force an update to pull updated lut incase subclasses have
+        # implemented it.
+        if not self.lut.root:
+            self.update_lookup_table()
+        poly = self.lut.get_poly(energy=energy, pol=pol)
+        return poly(energy)
+
+
+class ConfigServerEnergyMotorLookup(EnergyMotorLookup):
+    """Fetches and parses lookup table (csv) from a config server, supports dynamic
+    updates, and validates input."""
+
+    def __init__(
+        self,
+        config_client: ConfigServer,
+        lut_config: LookupTableColumnConfig,
+        path: Path,
+    ):
+        """
+        Parameters:
+        -----------
+        config_client:
+            The config server client to fetch the look up table data.
+        lut_config:
+            Configuration that defines how to process file contents into a LookupTable
+        path:
+            File path to the lookup table.
+        """
+        self.path = path
+        self.config_client = config_client
+        self.lut_config = lut_config
+        super().__init__()
+
+    def read_lut(self) -> LookupTable:
+        file_contents = self.config_client.get_file_contents(
+            self.path, reset_cached_result=True
+        )
+        return convert_csv_to_lookup(file_contents, lut_config=self.lut_config)
+
+    def update_lookup_table(self) -> None:
+        self.lut = self.read_lut()

dodal/devices/insertion_device/lookup_table_models.py

@@ -0,0 +1,287 @@
+"""Apple2 lookup table utilities and CSV converter.
+
+This module provides helpers to read, validate and convert Apple2 insertion-device
+lookup tables (energy -> gap/phase polynomials) from CSV sources into an
+in-memory dictionary format used by the Apple2 controllers.
+
+Data format produced
+The lookup-table dictionary created by convert_csv_to_lookup() follows this
+structure:
+
+{
+  "POL_MODE": {
+    "energy_entries": [
+        {
+            "low": <float>,
+            "high": <float>,
+            "poly": <numpy.poly1d>
+        },
+      ...
+    ]
+  },
+  ...
+}
+"""
+
+import csv
+import io
+from collections.abc import Generator
+from typing import Annotated as A
+from typing import Any, NamedTuple, Self
+
+import numpy as np
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    RootModel,
+    field_serializer,
+    field_validator,
+)
+
+from dodal.devices.insertion_device.apple2_undulator import Pol
+
+DEFAULT_POLY_DEG = [
+    "7th-order",
+    "6th-order",
+    "5th-order",
+    "4th-order",
+    "3rd-order",
+    "2nd-order",
+    "1st-order",
+    "b",
+]
+
+MODE_NAME_CONVERT = {"cr": "pc", "cl": "nc"}
+DEFAULT_GAP_FILE = "IDEnergy2GapCalibrations.csv"
+DEFAULT_PHASE_FILE = "IDEnergy2PhaseCalibrations.csv"
+
+ROW_PHASE_MOTOR_TOLERANCE = 0.004
+ROW_PHASE_CIRCULAR = 15
+MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
+MAXIMUM_GAP_MOTOR_POSITION = 100
+
+DEFAULT_POLY1D_PARAMETERS = {
+    Pol.LH: [0],
+    Pol.LV: [MAXIMUM_ROW_PHASE_MOTOR_POSITION],
+    Pol.PC: [ROW_PHASE_CIRCULAR],
+    Pol.NC: [-ROW_PHASE_CIRCULAR],
+    Pol.LH3: [0],
+}
+
+
+class Source(NamedTuple):
+    column: str
+    value: str
+
+
+class LookupTableColumnConfig(BaseModel):
+    """Configuration on how to process a csv file columns into a LookupTable data model."""
+
+    source: A[
+        Source | None,
+        Field(
+            description="If not None, only process the row if the source column name match the value."
+        ),
+    ] = None
+    mode: A[str, Field(description="Polarisation mode column name.")] = "Mode"
+    min_energy: A[str, Field(description="Minimum energy column name.")] = "MinEnergy"
+    max_energy: A[str, Field(description="Maximum energy column name.")] = "MaxEnergy"
+    poly_deg: list[str] = Field(
+        description="Polynomial column names.", default_factory=lambda: DEFAULT_POLY_DEG
+    )
+    mode_name_convert: dict[str, str] = Field(
+        description="When processing polarisation mode values, map their alias values to a real value.",
+        default_factory=lambda: MODE_NAME_CONVERT,
+    )
+
+
+class EnergyCoverageEntry(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)  # So np.poly1d can be used.
+    min_energy: float
+    max_energy: float
+    poly: np.poly1d
+
+    @field_validator("poly", mode="before")
+    @classmethod
+    def validate_and_convert_poly(
+        cls: type[Self], value: np.poly1d | list
+    ) -> np.poly1d:
+        """If reading from serialized data, it will be using a list. Convert to np.poly1d"""
+        if isinstance(value, list):
+            return np.poly1d(value)
+        return value
+
+    @field_serializer("poly", mode="plain")
+    def serialize_poly(self, value: np.poly1d) -> list:
+        """Allow np.poly1d to work when serializing."""
+        return value.coefficients.tolist()
+
+
+class EnergyCoverage(BaseModel):
+    energy_entries: list[EnergyCoverageEntry]
+
+    @classmethod
+    def generate(
+        cls: type[Self],
+        min_energies: list[float],
+        max_energies: list[float],
+        poly1d_params: list[list[float]],
+    ) -> Self:
+        energy_entries = [
+            EnergyCoverageEntry(
+                min_energy=min_energy,
+                max_energy=max_energy,
+                poly=np.poly1d(poly_params),
+            )
+            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)
+
+    @property
+    def max_energy(self) -> float:
+        return max(e.max_energy for e in self.energy_entries)
+
+    def get_poly(self, energy: float) -> np.poly1d:
+        """
+        Return the numpy.poly1d polynomial applicable for the given energy.
+
+        Parameters:
+        -----------
+        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:
+            raise ValueError(
+                f"Demanding energy must lie between {min_energy} and {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
+        raise ValueError(
+            "Cannot find polynomial coefficients for your requested energy."
+            + " There might be gap in the calibration lookup table."
+        )
+
+
+class LookupTable(RootModel[dict[Pol, EnergyCoverage]]):
+    """
+    Specialised lookup table for insertion devices to relate the energy and polarisation
+    values to Apple2 motor positions.
+    """
+
+    # Allow to auto specify a dict if one not provided
+    def __init__(self, root: dict[Pol, EnergyCoverage] | None = None):
+        super().__init__(root=root or {})
+
+    @classmethod
+    def generate(
+        cls: type[Self],
+        pols: list[Pol],
+        energy_coverage: list[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
+
+    def get_poly(
+        self,
+        energy: float,
+        pol: Pol,
+    ) -> np.poly1d:
+        """
+        Return the numpy.poly1d polynomial applicable for the given energy and polarisation.
+
+        Parameters:
+        -----------
+        energy:
+            Energy value in the same units used to create the lookup table.
+        pol:
+            Polarisation mode (Pol enum).
+        """
+        return self.root[pol].get_poly(energy)
+
+
+def convert_csv_to_lookup(
+    file_contents: str,
+    lut_config: LookupTableColumnConfig,
+    skip_line_start_with: str = "#",
+) -> LookupTable:
+    """
+    Convert CSV content into the Apple2 lookup-table dictionary.
+
+    Parameters:
+    -----------
+    file_contents:
+        The CSV file contents as string.
+    lut_config:
+        The configuration that how to process the file_contents into a LookupTable.
+    skip_line_start_with
+        Lines beginning with this prefix are skipped (default "#").
+
+    Returns:
+    -----------
+    LookupTable
+    """
+
+    def process_row(row: dict[str, Any], lut: LookupTable) -> None:
+        """Process a single row from the CSV file and update the lookup table."""
+        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(
+            min_energy=float(row[lut_config.min_energy]),
+            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)
+
+    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)
+
+    # Check if our LookupTable is empty after processing, raise error if it is.
+    if not lut.root:
+        raise RuntimeError(
+            "LookupTable content is empty, failed to convert the file contents to "
+            "a LookupTable!"
+        )
+    return lut
+
+
+def read_file_and_skip(file: str, skip_line_start_with: str = "#") -> Generator[str]:
+    """Yield non-comment lines from the CSV content string."""
+    for line in io.StringIO(file):
+        if line.startswith(skip_line_start_with):
+            continue
+        else:
+            yield line

dodal/devices/motors.py

@@ -122,6 +122,20 @@ class XYPitchStage(XYStage):
         super().__init__(prefix, name, x_infix, y_infix)
 
 
+class XYRollStage(XYStage):
+    def __init__(
+        self,
+        prefix: str,
+        x_infix: str = _X,
+        y_infix: str = _Y,
+        roll_infix: str = "ROLL",
+        name: str = "",
+    ) -> None:
+        with self.add_children_as_readables():
+            self.roll = Motor(prefix + roll_infix)
+        super().__init__(prefix, name, x_infix, y_infix)
+
+
 class XYZPitchYawStage(XYZStage):
     def __init__(
         self,