mx-bluesky 0.0.2__py3-none-any.whl → 1.1.0__py3-none-any.whl

mx_bluesky/hyperion/parameters/rotation.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator
+from itertools import accumulate
+from typing import Annotated, Any
+
+from annotated_types import Len
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import DetectorParams
+from dodal.devices.zebra import (
+    RotationDirection,
+)
+from dodal.log import LOGGER
+from pydantic import Field, field_validator, model_validator
+from scanspec.core import AxesPoints
+from scanspec.core import Path as ScanPath
+from scanspec.specs import Line
+
+from mx_bluesky.hyperion.parameters.components import (
+    DiffractionExperimentWithSample,
+    IspybExperimentType,
+    OptionalGonioAngleStarts,
+    OptionalXyzStarts,
+    RotationAxis,
+    SplitScan,
+    WithScan,
+)
+from mx_bluesky.hyperion.parameters.constants import (
+    CONST,
+    I03Constants,
+)
+
+
+class RotationScanPerSweep(OptionalGonioAngleStarts, OptionalXyzStarts):
+    omega_start_deg: float = Field(default=0)  # type: ignore
+    rotation_axis: RotationAxis = Field(default=RotationAxis.OMEGA)
+    scan_width_deg: float = Field(default=360, gt=0)
+    rotation_direction: RotationDirection = Field(default=RotationDirection.NEGATIVE)
+    nexus_vds_start_img: int = Field(default=0, ge=0)
+
+
+class RotationExperiment(DiffractionExperimentWithSample):
+    shutter_opening_time_s: float = Field(default=CONST.I03.SHUTTER_TIME_S)
+    rotation_increment_deg: float = Field(default=0.1, gt=0)
+    ispyb_experiment_type: IspybExperimentType = Field(
+        default=IspybExperimentType.ROTATION
+    )
+
+    def _detector_params(self, omega_start_deg: float):
+        self.det_dist_to_beam_converter_path = (
+            self.det_dist_to_beam_converter_path
+            or CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
+        )
+        optional_args = {}
+        if self.run_number:
+            optional_args["run_number"] = self.run_number
+        assert self.detector_distance_mm is not None
+        os.makedirs(self.storage_directory, exist_ok=True)
+        return DetectorParams(
+            detector_size_constants=I03Constants.DETECTOR,
+            expected_energy_ev=self.demand_energy_ev,
+            exposure_time=self.exposure_time_s,
+            directory=self.storage_directory,
+            prefix=self.file_name,
+            detector_distance=self.detector_distance_mm,
+            omega_start=omega_start_deg,
+            omega_increment=self.rotation_increment_deg,
+            num_images_per_trigger=self.num_images,
+            num_triggers=1,
+            use_roi_mode=False,
+            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
+            **optional_args,
+        )
+
+    @field_validator("selected_aperture")
+    @classmethod
+    def _set_default_aperture_position(cls, aperture_position: ApertureValue | None):
+        if not aperture_position:
+            default_aperture = CONST.PARAM.ROTATION.DEFAULT_APERTURE_POSITION
+            LOGGER.warning(
+                f"No aperture position selected. Defaulting to {default_aperture}"
+            )
+            return default_aperture
+        else:
+            return aperture_position
+
+
+class RotationScan(WithScan, RotationScanPerSweep, RotationExperiment):
+    @property
+    def detector_params(self):
+        return self._detector_params(self.omega_start_deg)
+
+    @property
+    def scan_points(self) -> AxesPoints:
+        scan_spec = Line(
+            axis="omega",
+            start=self.omega_start_deg,
+            stop=(
+                self.omega_start_deg
+                + (self.scan_width_deg - self.rotation_increment_deg)
+            ),
+            num=self.num_images,
+        )
+        scan_path = ScanPath(scan_spec.calculate())
+        return scan_path.consume().midpoints
+
+    @property
+    def num_images(self) -> int:
+        return int(self.scan_width_deg / self.rotation_increment_deg)
+
+
+class MultiRotationScan(RotationExperiment, SplitScan):
+    rotation_scans: Annotated[list[RotationScanPerSweep], Len(min_length=1)]
+
+    def _single_rotation_scan(self, scan: RotationScanPerSweep) -> RotationScan:
+        # self has everything from RotationExperiment
+        params = self.model_dump()
+        del params["rotation_scans"]
+        # provided `scan` has everything from RotationScanPerSweep
+        params.update(scan.model_dump())
+        # together they have everything for RotationScan
+        return RotationScan(**params)
+
+    @model_validator(mode="after")
+    @classmethod
+    def correct_start_vds(cls, values: Any) -> Any:
+        assert isinstance(values, MultiRotationScan)
+        start_img = 0.0
+        for scan in values.rotation_scans:
+            scan.nexus_vds_start_img = int(start_img)
+            start_img += scan.scan_width_deg / values.rotation_increment_deg
+        return values
+
+    @property
+    def single_rotation_scans(self) -> Iterator[RotationScan]:
+        for scan in self.rotation_scans:
+            yield self._single_rotation_scan(scan)
+
+    def _num_images_per_scan(self):
+        return [
+            int(scan.scan_width_deg / self.rotation_increment_deg)
+            for scan in self.rotation_scans
+        ]
+
+    @property
+    def num_images(self):
+        return sum(self._num_images_per_scan())
+
+    @property
+    def scan_indices(self):
+        return list(accumulate([0, *self._num_images_per_scan()]))
+
+    @property
+    def detector_params(self):
+        return self._detector_params(self.rotation_scans[0].omega_start_deg)
+
+    @property
+    def ispyb_params(self):  # pyright: ignore
+        raise ValueError("Please get ispyb params from one of the individual scans")