dls-dodal 1.57.0__py3-none-any.whl → 1.59.1__py3-none-any.whl

dodal/devices/i04/murko_results.py

@@ -1,5 +1,6 @@
 import json
 import pickle
+from dataclasses import dataclass
 from enum import Enum
 from typing import TypedDict
 
@@ -21,9 +22,6 @@ from dodal.log import LOGGER
 
 NO_MURKO_RESULT = (-1, -1)
 
-MurkoResult = dict
-FullMurkoResults = dict[str, list[MurkoResult]]
-
 
 class MurkoMetadata(TypedDict):
     zoom_percentage: float
@@ -42,6 +40,19 @@ class Coord(Enum):
     z = 2
 
 
+@dataclass
+class MurkoResult:
+    centre_px: tuple
+    x_dist_mm: float
+    y_dist_mm: float
+    omega: float
+    uuid: str
+
+
+class NoResultsFound(ValueError):
+    pass
+
+
 class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
     """Device that takes crystal centre values from Murko and uses them to set the
     x, y, z coordinate of the sample to be in line with the beam centre.
@@ -59,6 +70,8 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
     """
 
     TIMEOUT_S = 2
+    PERCENTAGE_TO_USE = 25
+    NUMBER_OF_WRONG_RESULTS_TO_LOG = 5
 
     def __init__(
         self,
@@ -74,12 +87,10 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
             db=redis_db,
         )
         self.pubsub = self.redis_client.pubsub()
-        self._last_omega = 0
         self.sample_id = soft_signal_rw(str)  # Should get from redis
         self.stop_angle = stop_angle
-        self.x_dists_mm = []
-        self.y_dists_mm = []
-        self.omegas = []
+
+        self._reset()
 
         with self.add_children_as_readables():
             # Diffs from current x/y/z
@@ -88,6 +99,10 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
             self.z_mm, self._z_mm_setter = soft_signal_r_and_setter(float)
         super().__init__(name=name)
 
+    def _reset(self):
+        self._last_omega = 0
+        self.results: list[MurkoResult] = []
+
     @AsyncStatus.wrap
     async def stage(self):
         await self.pubsub.subscribe("murko-results")
@@ -97,6 +112,7 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
 
     @AsyncStatus.wrap
     async def unstage(self):
+        self._reset()
         await self.pubsub.unsubscribe()
 
     @AsyncStatus.wrap
@@ -106,19 +122,26 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
         while self._last_omega < self.stop_angle:
             # waits here for next batch to be received
             message = await self.pubsub.get_message(timeout=self.TIMEOUT_S)
-            if message is None:  # No more messages to process
-                break
+            if message is None:
+                continue
             await self.process_batch(message, sample_id)
 
-        for i in range(len(self.omegas)):
-            LOGGER.debug(
-                f"omega: {round(self.omegas[i], 2)}, x: {round(self.x_dists_mm[i], 2)}, y: {round(self.y_dists_mm[i], 2)}"
-            )
+        if not self.results:
+            raise NoResultsFound("No results retrieved from Murko")
+
+        for result in self.results:
+            LOGGER.debug(result)
+
+        self.filter_outliers()
 
-        LOGGER.info(f"Using average of x beam distances: {self.x_dists_mm}")
-        avg_x = float(np.mean(self.x_dists_mm))
-        LOGGER.info(f"Finding least square y and z from y distances: {self.y_dists_mm}")
-        best_y, best_z = get_yz_least_squares(self.y_dists_mm, self.omegas)
+        x_dists_mm = [result.x_dist_mm for result in self.results]
+        y_dists_mm = [result.y_dist_mm for result in self.results]
+        omegas = [result.omega for result in self.results]
+
+        LOGGER.info(f"Using average of x beam distances: {x_dists_mm}")
+        avg_x = float(np.mean(x_dists_mm))
+        LOGGER.info(f"Finding least square y and z from y distances: {y_dists_mm}")
+        best_y, best_z = get_yz_least_squares(y_dists_mm, omegas)
         # x, y, z are relative to beam centre. Need to move negative these values to get centred.
         self._x_mm_setter(-avg_x)
         self._y_mm_setter(-best_y)
@@ -163,15 +186,38 @@ class MurkoResultsDevice(StandardReadable, Triggerable, Stageable):
                 centre_px[0],
                 centre_px[1],
             )
-            self.x_dists_mm.append(
-                beam_dist_px[0] * metadata["microns_per_x_pixel"] / 1000
-            )
-            self.y_dists_mm.append(
-                beam_dist_px[1] * metadata["microns_per_y_pixel"] / 1000
+            self.results.append(
+                MurkoResult(
+                    centre_px=centre_px,
+                    x_dist_mm=beam_dist_px[0] * metadata["microns_per_x_pixel"] / 1000,
+                    y_dist_mm=beam_dist_px[1] * metadata["microns_per_y_pixel"] / 1000,
+                    omega=omega,
+                    uuid=metadata["uuid"],
+                )
             )
-            self.omegas.append(omega)
             self._last_omega = omega
 
+    def filter_outliers(self):
+        """Whilst murko is not fully trained it often gives us poor results.
+        When it is wrong it usually picks up the base of the pin, rather than the tip,
+        meaning that by keeping only a percentage of the results with the smallest X we
+        remove many of the outliers.
+        """
+        LOGGER.info(f"Number of results before filtering: {len(self.results)}")
+        sorted_results = sorted(self.results, key=lambda item: item.centre_px[0])
+
+        worst_results = [
+            r.uuid for r in sorted_results[-self.NUMBER_OF_WRONG_RESULTS_TO_LOG :]
+        ]
+
+        LOGGER.info(
+            f"Worst {self.NUMBER_OF_WRONG_RESULTS_TO_LOG} murko results were {worst_results}"
+        )
+        cutoff = max(1, int(len(sorted_results) * self.PERCENTAGE_TO_USE / 100))
+        smallest_x = sorted_results[:cutoff]
+        self.results = smallest_x
+        LOGGER.info(f"Number of results after filtering: {len(self.results)}")
+
 
 def get_yz_least_squares(vertical_dists: list, omegas: list) -> tuple[float, float]:
     """Get the least squares solution for y and z from the vertical distances and omega angles.

dodal/devices/i10/i10_apple2.py

@@ -1,11 +1,13 @@
 import asyncio
 import csv
+import io
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, SupportsFloat
 
 import numpy as np
 from bluesky.protocols import Movable
+from daq_config_server.client import ConfigServer
 from ophyd_async.core import (
     AsyncStatus,
     Device,
@@ -15,13 +17,14 @@ from ophyd_async.core import (
     soft_signal_r_and_setter,
     soft_signal_rw,
 )
+from pydantic import BaseModel, ConfigDict, RootModel
 
 from dodal.log import LOGGER
 
 from ..apple2_undulator import (
     Apple2,
     Apple2Val,
-    Lookuptable,
+    EnergyMotorConvertor,
     Pol,
     UndulatorGap,
     UndulatorJawPhase,
@@ -34,6 +37,7 @@ MAXIMUM_ROW_PHASE_MOTOR_POSITION = 24.0
 MAXIMUM_GAP_MOTOR_POSITION = 100
 DEFAULT_JAW_PHASE_POLY_PARAMS = [1.0 / 7.5, -120.0 / 7.5]
 ALPHA_OFFSET = 180
+MAXIMUM_MOVE_TIME = 550  # There is no useful movements take longer than this.
 
 
 # data class to store the lookup table configuration that is use in convert_csv_to_lookup
@@ -53,36 +57,75 @@ class LookupTableConfig:
     poly_deg: list | None
 
 
-class I10Apple2(Apple2):
-    """I10Apple2 is the i10 version of Apple2 ID, set and update_lookuptable function
-    should be the only part that is I10 specific.
+class EnergyMinMax(BaseModel):
+    Minimum: float
+    Maximum: float
 
-    A pair of look up tables are needed to provide the conversion betwApple 2 ID/undulator has 4 extra degrees of freedom compare to the standard Undulator,
-    each bank of magnet can move independently to each other,
-    which allow the production of different x-ray polarisation as well as energy.
-    This type of ID is use on I10, I21, I09, I17 and I06 for soft x-ray.een motor position and energy.
 
-    Set is in energy(eV).
+class EnergyCoverageEntry(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    Low: float
+    High: float
+    Poly: np.poly1d
+
+
+class EnergyCoverage(RootModel):
+    root: dict[str, EnergyCoverageEntry]
+
+
+class LookupTableEntries(BaseModel):
+    Energies: EnergyCoverage
+    Limit: EnergyMinMax
+
+
+class Lookuptable(RootModel):
+    """BaseModel class for the lookup table.
+    Apple2 lookup table should be in this format.
+
+    {mode: {'Energies': {Any: {'Low': float,
+                            'High': float,
+                            'Poly':np.poly1d
+                            }
+                        }
+            'Limit': {'Minimum': float,
+                    'Maximum': float
+                    }
+        }
+    }
+    """
+
+    root: dict[str, LookupTableEntries]
+
+
+class I10EnergyMotorLookup:
+    """
+    Handles lookup tables for I10 Apple2 ID, converting energy and polarisation to gap
+     and phase. Fetches and parses lookup tables from a config server, supports dynamic
+     updates, and validates input.
     """
 
     def __init__(
         self,
         look_up_table_dir: str,
         source: tuple[str, str],
-        prefix: str,
+        config_client: ConfigServer,
         mode: str = "Mode",
         min_energy: str = "MinEnergy",
         max_energy: str = "MaxEnergy",
+        gap_file_name: str = "IDEnergy2GapCalibrations.csv",
+        phase_file_name: str = "IDEnergy2PhaseCalibrations.csv",
         poly_deg: list | None = None,
-        name: str = "",
-    ) -> None:
-        """
+    ):
+        """Initialise the I10EnergyMotorLookup class with lookup table headers provided.
+
         Parameters
         ----------
         look_up_table_dir:
             The path to look up table.
         source:
-            The column name and the name of the source in look up table. e.g. ("source", "idu")
+            The column name and the name of the source in look up table. e.g. ( "source", "idu")
+        config_client:
+            The config server client to fetch the look up table.
         mode:
             The column name of the mode in look up table.
         min_energy:
@@ -91,17 +134,14 @@ class I10Apple2(Apple2):
             The column name that contain the maximum energy in look up table.
         poly_deg:
             The column names for the parameters for the energy conversion polynomial, starting with the least significant.
-        prefix:
-            epic pv for id
-        Name:
-            Name of the device
-        """
 
-        energy_gap_table_path = Path(look_up_table_dir, "IDEnergy2GapCalibrations.csv")
-        energy_phase_table_path = Path(
-            look_up_table_dir, "IDEnergy2PhaseCalibrations.csv"
-        )
-        # A dataclass contains the path to the look up table and the expected column names.
+        """
+        self.lookup_tables: dict[str, dict[str | None, dict[str, dict[str, Any]]]] = {
+            "Gap": {},
+            "Phase": {},
+        }
+        energy_gap_table_path = Path(look_up_table_dir, gap_file_name)
+        energy_phase_table_path = Path(look_up_table_dir, phase_file_name)
         self.lookup_table_config = LookupTableConfig(
             path=LookupPath(Gap=energy_gap_table_path, Phase=energy_phase_table_path),
             source=source,
@@ -110,6 +150,212 @@ class I10Apple2(Apple2):
             max_energy=max_energy,
             poly_deg=poly_deg,
         )
+        self.config_client = config_client
+        self._available_pol = []
+
+    @property
+    def available_pol(self) -> list[str | None]:
+        return self._available_pol
+
+    @available_pol.setter
+    def available_pol(self, value: list[str | None]) -> None:
+        self._available_pol = value
+
+    def update_lookuptable(self):
+        """
+        Update lookup tables from files and validate their format.
+        """
+        LOGGER.info("Updating lookup dictionary from file.")
+        for key, path in self.lookup_table_config.path.__dict__.items():
+            self.lookup_tables[key] = self.convert_csv_to_lookup(
+                file=path,
+                source=self.lookup_table_config.source,
+                mode=self.lookup_table_config.mode,
+                min_energy=self.lookup_table_config.min_energy,
+                max_energy=self.lookup_table_config.max_energy,
+                poly_deg=self.lookup_table_config.poly_deg,
+            )
+            Lookuptable.model_validate(self.lookup_tables[key])
+
+        self.available_pol = list(self.lookup_tables["Gap"].keys())
+
+    def get_motor_from_energy(self, energy: float, pol: Pol) -> tuple[float, float]:
+        """
+        Convert energy and polarisation to gap and phase motor positions.
+
+        Parameters
+        ----------
+        energy : float
+            Desired energy in eV.
+        pol : Pol
+            Polarisation mode.
+
+        Returns
+        -------
+        tuple[float, float]
+            (gap, phase) motor positions.
+
+        """
+        if self.available_pol == []:
+            self.update_lookuptable()
+
+        gap_poly = self._get_poly(
+            lookup_table=self.lookup_tables["Gap"], energy=energy, pol=pol
+        )
+        phase_poly = self._get_poly(
+            lookup_table=self.lookup_tables["Phase"], energy=energy, pol=pol
+        )
+        return gap_poly(energy), phase_poly(energy)
+
+    def _get_poly(
+        self,
+        energy: float,
+        pol: Pol,
+        lookup_table: dict[str | None, dict[str, dict[str, Any]]],
+    ) -> np.poly1d:
+        """
+        Get polynomial for a given energy and polarisation.
+
+        Raises
+        ------
+        ValueError
+            If energy is out of bounds or coefficients are missing.
+        """
+        if (
+            energy < lookup_table[pol]["Limit"]["Minimum"]
+            or energy > lookup_table[pol]["Limit"]["Maximum"]
+        ):
+            raise ValueError(
+                "Demanding energy must lie between {} and {} eV!".format(
+                    lookup_table[pol]["Limit"]["Minimum"],
+                    lookup_table[pol]["Limit"]["Maximum"],
+                )
+            )
+        else:
+            for energy_range in lookup_table[pol]["Energies"].values():
+                if energy >= energy_range["Low"] and energy < energy_range["High"]:
+                    return energy_range["Poly"]
+
+        raise ValueError(
+            """Cannot find polynomial coefficients for your requested energy.
+        There might be gap in the calibration lookup table."""
+        )
+
+    def convert_csv_to_lookup(
+        self,
+        file: str,
+        source: tuple[str, str],
+        mode: str | None = "Mode",
+        min_energy: str | None = "MinEnergy",
+        max_energy: str | None = "MaxEnergy",
+        poly_deg: list | None = None,
+    ) -> dict[str | None, dict[str, dict[str, dict[str, Any]]]]:
+        """
+        Convert a CSV file to a lookup table dictionary.
+
+        Returns
+        -------
+        dict
+            Dictionary in Apple2 lookup table format.
+
+        Raises
+        ------
+        RuntimeError
+            If the CSV cannot be converted.
+
+        """
+        if poly_deg is None:
+            poly_deg = [
+                "7th-order",
+                "6th-order",
+                "5th-order",
+                "4th-order",
+                "3rd-order",
+                "2nd-order",
+                "1st-order",
+                "b",
+            ]
+        lookup_table = {}
+        polarisations = set()
+
+        def process_row(row: dict) -> None:
+            """Process a single row from the CSV file and update the lookup table."""
+            mode_value = row[mode]
+            if mode_value not in polarisations:
+                polarisations.add(mode_value)
+                lookup_table[mode_value] = {
+                    "Energies": {},
+                    "Limit": {
+                        "Minimum": float(row[min_energy]),
+                        "Maximum": float(row[max_energy]),
+                    },
+                }
+
+            # Create polynomial object for energy-to-gap/phase conversion
+            coefficients = [float(row[coef]) for coef in poly_deg]
+            polynomial = np.poly1d(coefficients)
+
+            lookup_table[mode_value]["Energies"][row[min_energy]] = {
+                "Low": float(row[min_energy]),
+                "High": float(row[max_energy]),
+                "Poly": polynomial,
+            }
+
+            # Update energy limits
+            lookup_table[mode_value]["Limit"]["Minimum"] = min(
+                lookup_table[mode_value]["Limit"]["Minimum"], float(row[min_energy])
+            )
+            lookup_table[mode_value]["Limit"]["Maximum"] = max(
+                lookup_table[mode_value]["Limit"]["Maximum"], float(row[max_energy])
+            )
+
+        csv_file = self.config_client.get_file_contents(file, reset_cached_result=True)
+        reader = csv.DictReader(io.StringIO(csv_file))
+        for row in reader:
+            # If there are multiple source only convert requested.
+            if row[source[0]] == source[1]:
+                process_row(row=row)
+        if not lookup_table:
+            raise RuntimeError(f"Unable to convert lookup table:\t{file}")
+        return lookup_table
+
+
+class I10Apple2(Apple2):
+    """I10Apple2 is the i10 version of Apple2 ID, set and energy_motor_convertor
+     should be the only part that is I10 specific.
+
+    A EnergyMotorConvertor function is needed to provide the conversion between
+     x-ray motor position and energy.
+
+    Set is in energy(eV).
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        energy_motor_convertor: EnergyMotorConvertor,
+        name: str = "",
+    ) -> None:
+        """
+        Parameters
+        ----------
+        look_up_table_dir:
+            The path to look up table.
+        source:
+            The column name and the name of the source in look up table. e.g. ("source", "idu")
+        mode:
+            The column name of the mode in look up table.
+        min_energy:
+            The column name that contain the maximum energy in look up table.
+        max_energy:
+            The column name that contain the maximum energy in look up table.
+        poly_deg:
+            The column names for the parameters for the energy conversion polynomial, starting with the least significant.
+        prefix:
+            epic pv for id
+        Name:
+            Name of the device
+        """
 
         with self.add_children_as_readables():
             super().__init__(
@@ -122,6 +368,7 @@ class I10Apple2(Apple2):
                     btm_inner="RPQ3",
                     btm_outer="RPQ4",
                 ),
+                energy_motor_convertor=energy_motor_convertor,
                 name=name,
             )
             self.id_jaw_phase = UndulatorJawPhase(
@@ -150,7 +397,7 @@ class I10Apple2(Apple2):
                 )
 
             self._set_pol_setpoint(pol)
-        gap, phase = await self._get_id_gap_phase(value)
+        gap, phase = self.energy_to_motor(energy=value, pol=pol)
         phase3 = phase * (-1 if pol == Pol.LA else 1)
         id_set_val = Apple2Val(
             top_outer=f"{phase:.6f}",
@@ -167,29 +414,6 @@ class I10Apple2(Apple2):
             await self.id_jaw_phase.set_move.set(1)
         LOGGER.info(f"Energy set to {value} eV successfully.")
 
-    def update_lookuptable(self):
-        """
-        Update the stored lookup tabled from file.
-
-        """
-        LOGGER.info("Updating lookup dictionary from file.")
-        for key, path in self.lookup_table_config.path.__dict__.items():
-            if path.exists():
-                self.lookup_tables[key] = convert_csv_to_lookup(
-                    file=path,
-                    source=self.lookup_table_config.source,
-                    mode=self.lookup_table_config.mode,
-                    min_energy=self.lookup_table_config.min_energy,
-                    max_energy=self.lookup_table_config.max_energy,
-                    poly_deg=self.lookup_table_config.poly_deg,
-                )
-                # ensure the importing lookup table is the correct format
-                Lookuptable.model_validate(self.lookup_tables[key])
-            else:
-                raise FileNotFoundError(f"{key} look up table is not in path: {path}")
-
-        self._available_pol = list(self.lookup_tables["Gap"].keys())
-
 
 class EnergySetter(StandardReadable, Movable[float]):
     """
@@ -250,7 +474,8 @@ class I10Apple2Pol(StandardReadable, Movable[Pol]):
     @AsyncStatus.wrap
     async def set(self, value: Pol) -> None:
         LOGGER.info(f"Changing f{self.name} polarisation to {value}.")
-        await self.id_ref().polarisation.set(value)
+        # Timeout is determined internally by the set method later, so we set it to max here.
+        await self.id_ref().polarisation.set(value, timeout=MAXIMUM_MOVE_TIME)
 
 
 class LinearArbitraryAngle(StandardReadable, Movable[SupportsFloat]):
@@ -320,6 +545,7 @@ class I10Id(Device):
         prefix: str,
         look_up_table_dir: str,
         source: tuple[str, str],
+        config_client: ConfigServer,
         jaw_phase_limit=12.0,
         jaw_phase_poly_param=DEFAULT_JAW_PHASE_POLY_PARAMS,
         angle_threshold_deg=30.0,
@@ -341,12 +567,16 @@ class I10Id(Device):
         linear_arbitrary_angle : LinearArbitraryAngle
             A device for controlling the linear arbitrary polarization angle.
         """
+        self.lookup_table_client = I10EnergyMotorLookup(
+            look_up_table_dir=look_up_table_dir,
+            source=source,
+            config_client=config_client,
+        )
         self.energy = EnergySetter(
             id=I10Apple2(
-                look_up_table_dir=look_up_table_dir,
-                name="id_energy",
-                source=source,
                 prefix=prefix,
+                energy_motor_convertor=self.lookup_table_client.get_motor_from_energy,
+                name="id_energy",
             ),
             pgm=pgm,
             name="energy",
@@ -361,91 +591,3 @@ class I10Id(Device):
         )
 
         super().__init__(name=name)
-
-
-def convert_csv_to_lookup(
-    file: str,
-    source: tuple[str, str],
-    mode: str | None = "Mode",
-    min_energy: str | None = "MinEnergy",
-    max_energy: str | None = "MaxEnergy",
-    poly_deg: list | None = None,
-) -> dict[str | None, dict[str, dict[str, dict[str, Any]]]]:
-    """
-    Convert a CSV file to a dictionary compatible with the Apple2 lookup table format.
-
-    Parameters
-    ----------
-    file : str
-        Path to the CSV file.
-    source : tuple[str, str]
-        Tuple specifying the column name and source name (e.g., ("Source", "idu")).
-    mode : str, optional
-        Column name for the available modes (e.g., "lv", "lh", "pc", "nc"), by default "Mode".
-    min_energy : str, optional
-        Column name for the minimum energy, by default "MinEnergy".
-    max_energy : str, optional
-        Column name for the maximum energy, by default "MaxEnergy".
-    poly_deg : list, optional
-        Column names for polynomial coefficients, starting with the least significant term.
-
-    Returns
-    -------
-    dict
-        A dictionary conforming to the Apple2 lookup table format.
-
-    """
-    if poly_deg is None:
-        poly_deg = [
-            "7th-order",
-            "6th-order",
-            "5th-order",
-            "4th-order",
-            "3rd-order",
-            "2nd-order",
-            "1st-order",
-            "b",
-        ]
-    lookup_table = {}
-    polarisations = set()
-
-    def process_row(row: dict) -> None:
-        """Process a single row from the CSV file and update the lookup table."""
-        mode_value = row[mode]
-        if mode_value not in polarisations:
-            polarisations.add(mode_value)
-            lookup_table[mode_value] = {
-                "Energies": {},
-                "Limit": {
-                    "Minimum": float(row[min_energy]),
-                    "Maximum": float(row[max_energy]),
-                },
-            }
-
-        # Create polynomial object for energy-to-gap/phase conversion
-        coefficients = [float(row[coef]) for coef in poly_deg]
-        polynomial = np.poly1d(coefficients)
-
-        lookup_table[mode_value]["Energies"][row[min_energy]] = {
-            "Low": float(row[min_energy]),
-            "High": float(row[max_energy]),
-            "Poly": polynomial,
-        }
-
-        # Update energy limits
-        lookup_table[mode_value]["Limit"]["Minimum"] = min(
-            lookup_table[mode_value]["Limit"]["Minimum"], float(row[min_energy])
-        )
-        lookup_table[mode_value]["Limit"]["Maximum"] = max(
-            lookup_table[mode_value]["Limit"]["Maximum"], float(row[max_energy])
-        )
-
-    with open(file, newline="") as csvfile:
-        reader = csv.DictReader(csvfile)
-        for row in reader:
-            # If there are multiple source only convert requested.
-            if row[source[0]] == source[1]:
-                process_row(row=row)
-    if not lookup_table:
-        raise RuntimeError(f"Unable to convert lookup table:/n/t{file}")
-    return lookup_table