Add AbstractDoseEstimation

mxcubecore/HardwareObjects/MAXIV/BioMAX/dose_estimator.py

@@ -0,0 +1,172 @@
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+from mxcubecore import HardwareRepository as HWR
+from mxcubecore.HardwareObjects.abstract.AbstractDoseEstimator import (
+    AbstractDoseEstimator,
+    DoseEstimateParameters,
+    DoseEstimation,
+    DoseEstimationError,
+    DoseEstimationOk,
+    ExperimentalGoal,
+)
+
+if TYPE_CHECKING:
+    from mxcubecore.HardwareObjects.MAXIV.BioMAX.BIOMAXFlux import BIOMAXFlux
+    from mxcubecore.HardwareObjects.MAXIV.Energy import Energy
+
+
+class _PolynomialSegment(BaseModel):
+    """
+    Represents the coefficients of piecewise cubic regression.
+
+    Flux value for given energy is being estimated using
+    parameters gathered during beamline operation at BioMAX.
+    Estimated energy is calculated differently for specific
+    thresholds.
+    """
+
+    max_energy_ev: int
+    coeffs: list[float]
+
+
+class DoseEstimator(AbstractDoseEstimator):
+    """Implementation of dose estimation for BioMAX.
+
+    This is an example YAML configuration
+    .. code-block:: yaml
+        class: MAXIV.BioMAX.dose_estimator.DoseEstimator
+        configuration:
+            flux_coefficients:
+                - max_energy_ev: 9060
+                    coeffs: [1.40919e14, -6.95909e10, 1.10638e7, -553.774]
+                - max_energy_ev: 12610
+                    coeffs: [1.40919e14, -6.95909e10, 1.10638e7, -553.774]
+                - max_energy_ev: 18750
+                    coeffs: [1.59363e14, -2.78613e10, 1.69424e6, -35.0571]
+                - max_energy_ev: 20400
+                    coeffs: [1.00206e12, 2.65477e9, -240059, 5.44578]
+                - max_energy_ev: 24000
+                    coeffs: [5.06378e14, -6.52503e10, 2.81565e6, -40.6289]
+            experimental_goals:
+                cryo_highres:
+                    label: Cryo High Resolution
+                    type: resolution_dependent
+                    mgy_per_angstrom: 10
+                room:
+                    label: Room Temperature
+                    type: static
+                    mgy: 0.2
+                cys_cys:
+                    label: Cys-Cys
+                    type: static
+                    mgy: 2
+                s_sad:
+                    label: S-SAD
+                    type: static
+                    mgy: 5
+                mad_sad:
+                    label: MAD/SAD
+                    type: static
+                    mgy: 6
+    """
+
+    class HOConfig(BaseModel):
+        flux_coefficients: list[_PolynomialSegment]
+        experimental_goals: dict[str, ExperimentalGoal]
+
+    def __init__(self, name: str) -> None:
+        self._energy_hwo: Energy | None = None
+        self._flux_hwo: BIOMAXFlux | None = None
+        super().__init__(name)
+
+    def init(self) -> None:
+        super().init()
+
+        self._energy_hwo = HWR.beamline.energy
+        self._flux_hwo = HWR.beamline.flux
+
+    def _get_flux_t_current(self, energy_ev: float) -> float:
+        buckets = self._config.flux_coefficients
+        for bucket in buckets:
+            # in here flux_coefficients is assumed to always have
+            # the max_energy_ev thresholds in ascending order.
+            if energy_ev <= bucket.max_energy_ev:
+                return sum(
+                    coeff * energy_ev**exponent
+                    for exponent, coeff in enumerate(bucket.coeffs)
+                )
+        max_energy_ev = max(bucket.max_energy_ev for bucket in buckets)
+        err_msg = f"Energy out of bounds. Max available value is {max_energy_ev} eV."
+        raise ValueError(err_msg)
+
+    def estimate_dose(  # noqa: PLR0911  many returns make it clearer :)
+        self,
+        params: DoseEstimateParameters,
+    ) -> DoseEstimation:
+        flux = self._flux_hwo.get_value()  # ph/s
+
+        if flux <= 0:
+            return DoseEstimationError(msg="No flux detected, can't estimate dose.")
+
+        beamline_flux_density = self._flux_hwo.flux_density
+
+        if beamline_flux_density == -1:
+            return DoseEstimationError(
+                msg="Flux has not been measured. Please measure it first"
+            )
+        if beamline_flux_density == 0:
+            return DoseEstimationError(
+                msg="No beam has been detected during flux measurements!"
+            )
+
+        flux_density_energy = self._flux_hwo.flux_density_energy
+
+        if flux_density_energy <= 0:
+            return DoseEstimationError(msg="Energy at the time of measurement unknown.")
+
+        beamline_energy = self._energy_hwo.get_value()  # keV
+
+        try:
+            flux_at_beamline_energy = self._get_flux_t_current(
+                energy_ev=beamline_energy * 1_000
+            )
+        except ValueError as ex:
+            return DoseEstimationError(msg=str(ex))
+
+        if flux_at_beamline_energy <= 0:
+            return DoseEstimationError(msg="Flux model predicts value <= 0")
+
+        # This is a relative error term for the flux measurements.
+        flux_scale = flux_at_beamline_energy / flux
+        try:
+            user_flux = (
+                self._get_flux_t_current(energy_ev=params.energy_kev * 1_000)
+                / flux_scale
+            )
+        except ValueError as ex:
+            return DoseEstimationError(msg=str(ex))
+
+        beam_size_x, beam_size_y = HWR.beamline.beam.get_beam_size()  # mm
+
+        user_flux_density = (user_flux * params.transmission_pct / 100) / (
+            beam_size_x * beam_size_y * 1e6
+        )
+        wavelength = self._energy_hwo.calculate_wavelength(energy=params.energy_kev)
+        # The 2000 constant comes from old code.
+        dose_rate = user_flux_density / (2000 / wavelength / wavelength)
+        estimated_dose_gy = dose_rate * params.exp_time_s * params.num_images
+        max_images = (
+            int((params.dose_limit_mgy * 1_000_000) / dose_rate / params.exp_time_s)
+            if params.dose_limit_mgy
+            else None
+        )
+        return DoseEstimationOk(
+            dose_mgy=estimated_dose_gy / 1_000_000,
+            max_images=max_images,
+        )
+
+    @property
+    def experimental_goals(self) -> dict[str, ExperimentalGoal]:
+        return self._config.experimental_goals

mxcubecore/HardwareObjects/abstract/AbstractDoseEstimator.py

@@ -0,0 +1,84 @@
+from abc import abstractmethod
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, Field
+
+from mxcubecore.BaseHardwareObjects import HardwareObject
+
+
+class DoseEstimationOk(BaseModel):
+    status: Literal["ok"] = "ok"
+    dose_mgy: float
+    max_images: int | None = None
+
+
+class DoseEstimationError(BaseModel):
+    status: Literal["error"] = "error"
+    msg: str
+
+
+DoseEstimation = Annotated[
+    Union[DoseEstimationOk, DoseEstimationError],
+    Field(discriminator="status"),
+]
+
+
+class StaticGoal(BaseModel):
+    type: Literal["static"]
+    label: str
+    mgy: float
+
+
+class ResolutionDependentGoal(BaseModel):
+    type: Literal["resolution_dependent"]
+    label: str
+    mgy_per_angstrom: float
+
+
+ExperimentalGoal = StaticGoal | ResolutionDependentGoal
+
+
+class DoseEstimateParameters(BaseModel):
+    num_images: int
+    exp_time_s: float
+    energy_kev: float
+    transmission_pct: float
+    resolution_a: float
+    dose_limit_mgy: float | None = None
+    experimental_goal: str | None = None
+
+
+class AbstractDoseEstimator(HardwareObject):
+    """Estimates absorbed radiation dose for given collection parameters."""
+
+    @abstractmethod
+    def estimate_dose(self, params: DoseEstimateParameters) -> DoseEstimation:
+        """
+        Performs the dose estimation.
+
+        Args:
+            num_images: Number of images taken during collection
+            exp_time_s: Total time the sample is exposed per detector image.
+            energy_kev: Beam energy in keV
+            transmission_pct: Transmission value as a percentage.
+            resolution_a: Resolution in Å.
+            dose_limit: Max acceptable total absorbed dose in MGy.
+                Can be used to e.g. derive maxImages.
+            experimental_goal: Describes experimental setting / goal.
+        """
+        ...
+
+    @property
+    def experimental_goals(self) -> dict[str, ExperimentalGoal]:
+        """Preset dose limits exposed to the user.
+
+        Default value of ~10MGy per Å comes from literature, e.g. here:
+        https://journals.iucr.org/d/issues/2010/04/00/ba5150/index.html#SEC5
+        """
+        return {
+            "cryo": ResolutionDependentGoal(
+                type="resolution_dependent",
+                label="Cryo High Resolution",
+                mgy_per_angstrom=10,
+            )
+        }

mxcubecore/HardwareObjects/mockup/dose_estimator.py

@@ -0,0 +1,92 @@
+"""Mock-up class to simulate dose estimation, used for testing."""
+
+import random
+
+from pydantic import BaseModel
+
+from mxcubecore.HardwareObjects.abstract.AbstractDoseEstimator import (
+    AbstractDoseEstimator,
+    DoseEstimateParameters,
+    DoseEstimation,
+    DoseEstimationError,
+    DoseEstimationOk,
+    ExperimentalGoal,
+    ResolutionDependentGoal,
+    StaticGoal,
+)
+
+_DEFAULT_EXPERIMENTAL_GOALS: dict[str, ExperimentalGoal] = {
+    "cryo_highres": ResolutionDependentGoal(
+        type="resolution_dependent",
+        label="Cryo High Resolution",
+        mgy_per_angstrom=10,
+    ),
+    "room": StaticGoal(type="static", label="Room Temperature", mgy=0.2),
+    "s_sad": StaticGoal(type="static", label="S-SAD", mgy=5),
+}
+
+
+class MockupDoseEstimator(AbstractDoseEstimator):
+    """Simulated dose estimation.
+
+    There is no physics here: the dose is simply proportional to the total
+    exposure and the transmission, with an arbitrary dose rate. That is enough
+    to see on the demo beamline that the estimation is wired up.
+
+    A fraction of the estimations fails at random, so that the error path shows
+    up in the demo on its own.
+
+    This is an example YAML configuration
+    .. code-block:: yaml
+        class: mxcubecore.HardwareObjects.mockup.dose_estimator.MockupDoseEstimator
+        configuration:
+            dose_rate_mgy_per_s: 1.0
+            error_probability: 0.1
+            experimental_goals:
+                cryo_highres:
+                    label: Cryo High Resolution
+                    type: resolution_dependent
+                    mgy_per_angstrom: 10
+                room:
+                    label: Room Temperature
+                    type: static
+                    mgy: 0.2
+    """
+
+    class HOConfig(BaseModel):
+        # Arbitrary dose rate at 100% transmission, in MGy/s.
+        dose_rate_mgy_per_s: float = 1.0
+        experimental_goals: dict[str, ExperimentalGoal] = _DEFAULT_EXPERIMENTAL_GOALS
+        # Fraction of the estimations that fail.
+        error_probability: float = 0.1
+
+    def estimate_dose(self, params: DoseEstimateParameters) -> DoseEstimation:
+        if random.random() < self._config.error_probability:  # noqa: S311
+            return DoseEstimationError(
+                msg="Randomly generated error, to simulate a failing estimation."
+            )
+
+        if params.num_images <= 0 or params.exp_time_s <= 0:
+            return DoseEstimationError(
+                msg="Number of images and exposure time must be positive."
+            )
+
+        dose_per_image_mgy = (
+            self._config.dose_rate_mgy_per_s
+            * params.exp_time_s
+            * params.transmission_pct
+            / 100
+        )
+        max_images = (
+            int(params.dose_limit_mgy / dose_per_image_mgy)
+            if params.dose_limit_mgy and dose_per_image_mgy > 0
+            else None
+        )
+        return DoseEstimationOk(
+            dose_mgy=dose_per_image_mgy * params.num_images,
+            max_images=max_images,
+        )
+
+    @property
+    def experimental_goals(self) -> dict[str, ExperimentalGoal]:
+        return self._config.experimental_goals