mx-bluesky 1.2.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

mx_bluesky/common/parameters/components.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import os
+from abc import abstractmethod
+from collections.abc import Sequence
+from enum import StrEnum
+from pathlib import Path
+from typing import Literal, SupportsInt, cast
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import (
+    DetectorParams,
+    TriggerMode,
+)
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
+from pydantic_extra_types.semantic_version import SemanticVersion
+from scanspec.core import AxesPoints
+from semver import Version
+
+from mx_bluesky.common.parameters.constants import (
+    TEST_MODE,
+    DetectorParamConstants,
+    GridscanParamConstants,
+)
+
+PARAMETER_VERSION = Version.parse("5.2.0")
+
+
+class RotationAxis(StrEnum):
+    OMEGA = "omega"
+    PHI = "phi"
+    CHI = "chi"
+    KAPPA = "kappa"
+
+
+class XyzAxis(StrEnum):
+    X = "sam_x"
+    Y = "sam_y"
+    Z = "sam_z"
+
+
+class IspybExperimentType(StrEnum):
+    # Enum values from ispyb column data type
+    SAD = "SAD"  # at or slightly above the peak
+    SAD_INVERSE_BEAM = "SAD - Inverse Beam"
+    OSC = "OSC"  # "native" (in the absence of a heavy atom)
+    COLLECT_MULTIWEDGE = (
+        "Collect - Multiwedge"  # "poorly determined" ~ EDNA complex strategy???
+    )
+    MAD = "MAD"
+    HELICAL = "Helical"
+    MULTI_POSITIONAL = "Multi-positional"
+    MESH = "Mesh"
+    BURN = "Burn"
+    MAD_INVERSE_BEAM = "MAD - Inverse Beam"
+    CHARACTERIZATION = "Characterization"
+    DEHYDRATION = "Dehydration"
+    TOMO = "tomo"
+    EXPERIMENT = "experiment"
+    EM = "EM"
+    PDF = "PDF"
+    PDF_BRAGG = "PDF+Bragg"
+    BRAGG = "Bragg"
+    SINGLE_PARTICLE = "single particle"
+    SERIAL_FIXED = "Serial Fixed"
+    SERIAL_JET = "Serial Jet"
+    STANDARD = "Standard"  # Routine structure determination experiment
+    TIME_RESOLVED = "Time Resolved"  # Investigate the change of a system over time
+    DLS_ANVIL_HP = "Diamond Anvil High Pressure"  # HP sample environment pressure cell
+    CUSTOM = "Custom"  # Special or non-standard data collection
+    XRF_MAP = "XRF map"
+    ENERGY_SCAN = "Energy scan"
+    XRF_SPECTRUM = "XRF spectrum"
+    XRF_MAP_XAS = "XRF map xas"
+    MESH_3D = "Mesh3D"
+    SCREENING = "Screening"
+    STILL = "Still"
+    SSX_CHIP = "SSX-Chip"
+    SSX_JET = "SSX-Jet"
+
+    # Aliases for historic hyperion experiment type mapping
+    ROTATION = "SAD"
+    GRIDSCAN_2D = "mesh"
+    GRIDSCAN_3D = "Mesh3D"
+
+
+class MxBlueskyParameters(BaseModel):
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    def __hash__(self) -> int:
+        return self.model_dump_json().__hash__()
+
+    parameter_model_version: SemanticVersion
+
+    @field_validator("parameter_model_version")
+    @classmethod
+    def _validate_version(cls, version: Version):
+        assert (
+            version >= Version(major=PARAMETER_VERSION.major)
+        ), f"Parameter version too old! This version of hyperion uses {PARAMETER_VERSION}"
+        assert (
+            version <= Version(major=PARAMETER_VERSION.major + 1)
+        ), f"Parameter version too new! This version of hyperion uses {PARAMETER_VERSION}"
+        return version
+
+
+class WithSnapshot(BaseModel):
+    snapshot_directory: Path
+    snapshot_omegas_deg: list[float] | None = None
+
+    @property
+    def take_snapshots(self) -> bool:
+        return bool(self.snapshot_omegas_deg)
+
+
+class WithOptionalEnergyChange(BaseModel):
+    demand_energy_ev: float | None = Field(default=None, gt=0)
+
+
+class WithVisit(BaseModel):
+    beamline: str = Field(default="BL03I", pattern=r"BL\d{2}[BIJS]")
+    visit: str = Field(min_length=1)
+    det_dist_to_beam_converter_path: str = Field(
+        default=DetectorParamConstants.BEAM_XY_LUT_PATH
+    )
+    insertion_prefix: str = "SR03S" if TEST_MODE else "SR03I"
+    detector_distance_mm: float | None = Field(default=None, gt=0)
+
+
+class DiffractionExperiment(
+    MxBlueskyParameters, WithSnapshot, WithOptionalEnergyChange, WithVisit
+):
+    """For all experiments which use beam"""
+
+    file_name: str
+    exposure_time_s: float = Field(gt=0)
+    comment: str = Field(default="")
+    trigger_mode: TriggerMode = Field(default=TriggerMode.FREE_RUN)
+    run_number: int | None = Field(default=None, ge=0)
+    selected_aperture: ApertureValue | None = Field(default=None)
+    transmission_frac: float = Field(default=0.1)
+    ispyb_experiment_type: IspybExperimentType
+    storage_directory: str
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_directories(cls, values):
+        os.makedirs(values["storage_directory"], exist_ok=True)
+
+        values["snapshot_directory"] = values.get(
+            "snapshot_directory",
+            Path(values["storage_directory"], "snapshots").as_posix(),
+        )
+        return values
+
+    @property
+    def num_images(self) -> int:
+        return 0
+
+    @property
+    @abstractmethod
+    def detector_params(self) -> DetectorParams: ...
+
+
+class WithScan(BaseModel):
+    """For experiments where the scan is known"""
+
+    @property
+    @abstractmethod
+    def scan_points(self) -> AxesPoints: ...
+
+    @property
+    @abstractmethod
+    def num_images(self) -> int: ...
+
+
+class WithPandaGridScan(BaseModel):
+    """For experiments which use a PandA for constant-motion grid scans"""
+
+    panda_runup_distance_mm: float = Field(
+        default=GridscanParamConstants.PANDA_RUN_UP_DISTANCE_MM
+    )
+
+
+class SplitScan(BaseModel):
+    @property
+    @abstractmethod
+    def scan_indices(self) -> Sequence[SupportsInt]:
+        """Should return the first index of each scan (i.e. for each nexus file)"""
+        ...
+
+
+class WithSample(BaseModel):
+    sample_id: int
+    sample_puck: int | None = None
+    sample_pin: int | None = None
+
+
+class DiffractionExperimentWithSample(DiffractionExperiment, WithSample): ...
+
+
+class MultiXtalSelection(BaseModel):
+    name: str
+
+
+class TopNByMaxCountSelection(MultiXtalSelection):
+    name: Literal["TopNByMaxCount"] = "TopNByMaxCount"  #  pyright: ignore [reportIncompatibleVariableOverride]
+    n: int
+
+
+class WithCentreSelection(BaseModel):
+    select_centres: TopNByMaxCountSelection = Field(
+        discriminator="name", default=TopNByMaxCountSelection(n=1)
+    )
+
+    @property
+    def selection_params(self) -> MultiXtalSelection:
+        """A helper property because pydantic does not allow polymorphism with base classes
+        # only type unions"""
+        cast1 = cast(MultiXtalSelection, self.select_centres)
+        return cast1
+
+
+class OptionalXyzStarts(BaseModel):
+    x_start_um: float | None = None
+    y_start_um: float | None = None
+    z_start_um: float | None = None
+
+
+class XyzStarts(BaseModel):
+    x_start_um: float
+    y_start_um: float
+    z_start_um: float
+
+    def _start_for_axis(self, axis: XyzAxis) -> float:
+        match axis:
+            case XyzAxis.X:
+                return self.x_start_um
+            case XyzAxis.Y:
+                return self.y_start_um
+            case XyzAxis.Z:
+                return self.z_start_um
+
+
+class OptionalGonioAngleStarts(BaseModel):
+    omega_start_deg: float | None = None
+    phi_start_deg: float | None = None
+    chi_start_deg: float | None = None
+    kappa_start_deg: float | None = None

mx_bluesky/common/parameters/constants.py

@@ -0,0 +1,143 @@
+from enum import Enum
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import EIGER2_X_16M_SIZE
+from dodal.devices.zocalo.zocalo_constants import ZOCALO_ENV as ZOCALO_ENV_FROM_DODAL
+from dodal.utils import get_beamline_name
+from pydantic.dataclasses import dataclass
+
+BEAMLINE = get_beamline_name("test")
+TEST_MODE = BEAMLINE == "test"
+
+
+@dataclass(frozen=True)
+class DocDescriptorNames:
+    # Robot load event descriptor
+    ROBOT_LOAD = "robot_load"
+    # For callbacks to use
+    OAV_ROTATION_SNAPSHOT_TRIGGERED = "rotation_snapshot_triggered"
+    OAV_GRID_SNAPSHOT_TRIGGERED = "snapshot_to_ispyb"
+    HARDWARE_READ_PRE = "read_hardware_for_callbacks_pre_collection"
+    HARDWARE_READ_DURING = "read_hardware_for_callbacks_during_collection"
+    SAMPLE_HANDLING_EXCEPTION = "sample_handling_exception"
+    ZOCALO_HW_READ = "zocalo_read_hardware_plan"
+    FLYSCAN_RESULTS = "flyscan_results_obtained"
+
+
+@dataclass(frozen=True)
+class OavConstants:
+    OAV_CONFIG_JSON = (
+        "tests/test_data/test_OAVCentring.json"
+        if TEST_MODE
+        else (
+            f"/dls_sw/{BEAMLINE}/software/daq_configuration/json/OAVCentring_hyperion.json"
+        )
+    )
+
+
+@dataclass(frozen=True)
+class PlanNameConstants:
+    LOAD_CENTRE_COLLECT = "load_centre_collect"
+    # Robot load subplan
+    ROBOT_LOAD = "robot_load"
+    # Gridscan
+    GRID_DETECT_AND_DO_GRIDSCAN = "grid_detect_and_do_gridscan"
+    GRID_DETECT_INNER = "grid_detect"
+    GRIDSCAN_OUTER = "run_gridscan_move_and_tidy"
+    GRIDSCAN_AND_MOVE = "run_gridscan_and_move"
+    GRIDSCAN_MAIN = "run_gridscan"
+    DO_FGS = "do_fgs"
+    # IspyB callback activation
+    ISPYB_ACTIVATION = "ispyb_activation"
+    ROBOT_LOAD_AND_SNAPSHOTS = "robot_load_and_snapshots"
+    # Rotation scan
+    ROTATION_MULTI = "multi_rotation_wrapper"
+    ROTATION_OUTER = "rotation_scan_with_cleanup"
+    ROTATION_MAIN = "rotation_scan_main"
+    FLYSCAN_RESULTS = "xray_centre_results"
+
+
+@dataclass(frozen=True)
+class EnvironmentConstants:
+    ZOCALO_ENV = ZOCALO_ENV_FROM_DODAL
+
+
+@dataclass(frozen=True)
+class TriggerConstants:
+    ZOCALO = "trigger_zocalo_on"
+
+
+@dataclass(frozen=True)
+class HardwareConstants:
+    OAV_REFRESH_DELAY = 0.3
+    PANDA_FGS_RUN_UP_DEFAULT = 0.17
+    CRYOJET_MARGIN_MM = 0.2
+    THAWING_TIME = 20
+
+
+@dataclass(frozen=True)
+class GridscanParamConstants:
+    WIDTH_UM = 600.0
+    EXPOSURE_TIME_S = 0.004
+    USE_ROI = True
+    BOX_WIDTH_UM = 20.0
+    OMEGA_1 = 0.0
+    OMEGA_2 = 90.0
+    PANDA_RUN_UP_DISTANCE_MM = 0.2
+
+
+@dataclass(frozen=True)
+class RotationParamConstants:
+    DEFAULT_APERTURE_POSITION = ApertureValue.LARGE
+
+
+@dataclass(frozen=True)
+class DetectorParamConstants:
+    BEAM_XY_LUT_PATH = (
+        "tests/test_data/test_det_dist_converter.txt"
+        if TEST_MODE
+        else f"/dls_sw/{BEAMLINE}/software/daq_configuration/lookup/DetDistToBeamXYConverter.txt"
+    )
+    DETECTOR = EIGER2_X_16M_SIZE
+
+
+@dataclass(frozen=True)
+class ExperimentParamConstants:
+    DETECTOR = DetectorParamConstants()
+    GRIDSCAN = GridscanParamConstants()
+    ROTATION = RotationParamConstants()
+
+
+@dataclass(frozen=True)
+class PlanGroupCheckpointConstants:
+    # For places to synchronise / stop and wait in plans, use as bluesky group names
+    GRID_READY_FOR_DC = "grid_ready_for_data_collection"
+    ROTATION_READY_FOR_DC = "rotation_ready_for_data_collection"
+    MOVE_GONIO_TO_START = "move_gonio_to_start"
+    READY_FOR_OAV = "ready_for_oav"
+
+
+@dataclass(frozen=True)
+class SimConstants:
+    BEAMLINE = "BL03S"
+    INSERTION_PREFIX = "SR03S"
+    # this one is for unit tests
+    ISPYB_CONFIG = "tests/test_data/test_config.cfg"
+    # this one is for system tests
+    DEV_ISPYB_DATABASE_CFG = "/dls_sw/dasc/mariadb/credentials/ispyb-hyperion-dev.cfg"
+
+
+class Actions(Enum):
+    START = "start"
+    STOP = "stop"
+    SHUTDOWN = "shutdown"
+    STATUS = "status"
+
+
+class Status(Enum):
+    WARN = "Warn"
+    FAILED = "Failed"
+    SUCCESS = "Success"
+    BUSY = "Busy"
+    ABORTING = "Aborting"
+    IDLE = "Idle"

mx_bluesky/common/parameters/gridscan.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import os
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import (
+    DetectorParams,
+)
+from pydantic import Field
+
+from mx_bluesky.common.parameters.components import (
+    DiffractionExperimentWithSample,
+    IspybExperimentType,
+    OptionalGonioAngleStarts,
+    WithScan,
+    XyzStarts,
+)
+from mx_bluesky.common.parameters.constants import (
+    DetectorParamConstants,
+    GridscanParamConstants,
+    HardwareConstants,
+)
+from mx_bluesky.common.parameters.robot_load import RobotLoadAndEnergyChange
+
+
+class GridCommon(
+    DiffractionExperimentWithSample,
+    OptionalGonioAngleStarts,
+):
+    grid_width_um: float = Field(default=GridscanParamConstants.WIDTH_UM)
+    exposure_time_s: float = Field(default=GridscanParamConstants.EXPOSURE_TIME_S)
+    use_roi_mode: bool = Field(default=GridscanParamConstants.USE_ROI)
+
+    ispyb_experiment_type: IspybExperimentType = Field(
+        default=IspybExperimentType.GRIDSCAN_3D
+    )
+    selected_aperture: ApertureValue | None = Field(default=ApertureValue.SMALL)
+
+    @property
+    def detector_params(self):
+        self.det_dist_to_beam_converter_path = (
+            self.det_dist_to_beam_converter_path
+            or DetectorParamConstants.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
+        ), "Detector distance must be filled before generating DetectorParams"
+        os.makedirs(self.storage_directory, exist_ok=True)
+        return DetectorParams(
+            detector_size_constants=DetectorParamConstants.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=self.omega_start_deg or 0,
+            omega_increment=0,
+            num_images_per_trigger=1,
+            num_triggers=self.num_images,
+            use_roi_mode=self.use_roi_mode,
+            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
+            trigger_mode=self.trigger_mode,
+            **optional_args,
+        )
+
+
+class RobotLoadThenCentre(GridCommon):
+    thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)
+
+    def robot_load_params(self):
+        my_params = self.model_dump()
+        return RobotLoadAndEnergyChange(**my_params)
+
+    def pin_centre_then_xray_centre_params(self):
+        my_params = self.model_dump()
+        del my_params["thawing_time"]
+        return PinTipCentreThenXrayCentre(**my_params)
+
+
+class GridScanWithEdgeDetect(GridCommon):
+    box_size_um: float = Field(default=GridscanParamConstants.BOX_WIDTH_UM)
+
+
+class PinTipCentreThenXrayCentre(GridCommon):
+    tip_offset_um: float = 0
+
+
+class SpecifiedGrid(XyzStarts, WithScan):
+    """A specified grid is one which has defined values for the start position,
+    grid and box sizes, etc., as opposed to parameters for a plan which will create
+    those parameters at some point (e.g. through optical pin detection)."""

mx_bluesky/common/parameters/robot_load.py

@@ -0,0 +1,16 @@
+from pydantic import Field
+
+from mx_bluesky.common.parameters.components import (
+    MxBlueskyParameters,
+    WithOptionalEnergyChange,
+    WithSample,
+    WithSnapshot,
+    WithVisit,
+)
+from mx_bluesky.common.parameters.constants import HardwareConstants
+
+
+class RobotLoadAndEnergyChange(
+    MxBlueskyParameters, WithSample, WithSnapshot, WithOptionalEnergyChange, WithVisit
+):
+    thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)

mx_bluesky/common/plans/__init__.py

@@ -0,0 +1 @@
+"""Common MX plans which includes open and close data collection runs. These be used in isolation or as part of a larger plan"""

mx_bluesky/common/plans/do_fgs.py

@@ -0,0 +1,121 @@
+from collections.abc import Callable
+from time import time
+
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+from bluesky.utils import MsgGenerator
+from dodal.devices.eiger import EigerDetector
+from dodal.devices.fast_grid_scan import FastGridScanCommon
+from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.zocalo.zocalo_results import (
+    ZOCALO_STAGE_GROUP,
+)
+from dodal.log import LOGGER
+from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
+from scanspec.core import AxesPoints, Axis
+
+from mx_bluesky.common.device_setup_plans.read_hardware_for_setup import (
+    read_hardware_for_zocalo,
+)
+from mx_bluesky.common.parameters.constants import (
+    EnvironmentConstants,
+    PlanNameConstants,
+    TriggerConstants,
+)
+from mx_bluesky.common.utils.tracing import TRACER
+
+
+def _wait_for_zocalo_to_stage_then_do_fgs(
+    grid_scan_device: FastGridScanCommon,
+    detector: EigerDetector,
+    synchrotron: Synchrotron,
+    during_collection_plan: Callable[[], MsgGenerator] | None = None,
+):
+    expected_images = yield from bps.rd(grid_scan_device.expected_images)
+    exposure_sec_per_image = yield from bps.rd(detector.cam.acquire_time)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    LOGGER.info("waiting for topup if necessary...")
+    yield from check_topup_and_wait_if_necessary(
+        synchrotron,
+        expected_images * exposure_sec_per_image,
+        30.0,
+    )
+
+    # Make sure ZocaloResults queue is clear and ready to accept our new data. Zocalo MUST
+    # have been staged using ZOCALO_STAGE_GROUP prior to this
+    LOGGER.info("Waiting for Zocalo device queue to have been cleared...")
+    yield from bps.wait(ZOCALO_STAGE_GROUP)
+
+    # Triggers Zocalo if RE is subscribed to ZocaloCallback
+    yield from read_hardware_for_zocalo(detector)
+    LOGGER.info("Wait for all moves with no assigned group")
+    yield from bps.wait()
+
+    LOGGER.info("kicking off FGS")
+    yield from bps.kickoff(grid_scan_device, wait=True)
+    gridscan_start_time = time()
+    if during_collection_plan:
+        yield from during_collection_plan()
+    LOGGER.info("completing FGS")
+    yield from bps.complete(grid_scan_device, wait=True)
+    # Remove this logging statement once metrics have been added
+    LOGGER.info(
+        f"Grid scan motion program took {round(time()-gridscan_start_time,2)} to complete"
+    )
+
+
+def kickoff_and_complete_gridscan(
+    gridscan: FastGridScanCommon,
+    detector: EigerDetector,  # Once Eiger inherits from StandardDetector, use that type instead
+    synchrotron: Synchrotron,
+    scan_points: list[AxesPoints[Axis]],
+    scan_start_indices: list[int],
+    plan_during_collection: Callable[[], MsgGenerator] | None = None,
+    zocalo_environment: str = EnvironmentConstants.ZOCALO_ENV,
+):
+    """Triggers a grid scan motion program and waits for completion, accounting for synchrotron topup.
+    If the RunEngine is subscribed to ZocaloCallback, this plan will also trigger Zocalo.
+
+    Can be used for multiple successive grid scans, see Hyperion's usage
+
+    Args:
+        gridscan (FastGridScanCommon):          Device which can trigger a fast grid scan and wait for completion
+        detector (EigerDetector)                Detector device
+        synchrotron (Synchrotron):              Synchrotron device
+        scan_points (list[AxesPoints[Axis]]):   Each element in the list contains all the grid points for that grid scan.
+                                                Two elements in this list indicates that two grid scans will be done, eg for Hyperion's 3D grid scans.
+        scan_start_indices (list[int]):         Contains the first index of each grid scan
+        plan_during_collection (Optional, MsgGenerator): Generic plan called in between kickoff and completion,
+                                                eg waiting on zocalo.
+        zocalo_environment (Optional, str)      Used for zocalo connection
+    """
+
+    assert len(scan_points) == len(
+        scan_start_indices
+    ), "scan_points and scan_start_indices must be lists of the same length!"
+
+    plan_name = PlanNameConstants.DO_FGS
+
+    @TRACER.start_as_current_span(plan_name)
+    @bpp.set_run_key_decorator(plan_name)
+    @bpp.run_decorator(
+        md={
+            "subplan_name": plan_name,
+            TriggerConstants.ZOCALO: plan_name,
+            "scan_points": scan_points,
+            "scan_start_indices": scan_start_indices,
+            "zocalo_environment": zocalo_environment,
+        }
+    )
+    @bpp.contingency_decorator(
+        except_plan=lambda e: (yield from bps.stop(detector)),  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        else_plan=lambda: (yield from bps.unstage(detector)),  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+    )
+    def _decorated_do_fgs():
+        yield from _wait_for_zocalo_to_stage_then_do_fgs(
+            gridscan,
+            detector,
+            synchrotron,
+            during_collection_plan=plan_during_collection,
+        )
+
+    yield from _decorated_do_fgs()