diff --git a/mxcubecore/HardwareObjects/MAXIV/BioMAX/dose_estimator.py b/mxcubecore/HardwareObjects/MAXIV/BioMAX/dose_estimator.py new file mode 100644 index 0000000000..84d5822793 --- /dev/null +++ b/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 diff --git a/mxcubecore/HardwareObjects/abstract/AbstractDoseEstimator.py b/mxcubecore/HardwareObjects/abstract/AbstractDoseEstimator.py new file mode 100644 index 0000000000..d8f20d0625 --- /dev/null +++ b/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, + ) + } diff --git a/mxcubecore/HardwareObjects/mockup/dose_estimator.py b/mxcubecore/HardwareObjects/mockup/dose_estimator.py new file mode 100644 index 0000000000..19dae5f51d --- /dev/null +++ b/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