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

mx_bluesky/hyperion/parameters/cli.py

@@ -0,0 +1,68 @@
+import argparse
+
+from pydantic.dataclasses import dataclass
+
+from mx_bluesky._version import version
+
+
+@dataclass
+class HyperionArgs:
+    dev_mode: bool = False
+    use_external_callbacks: bool = False
+    verbose_event_logging: bool = False
+    skip_startup_connection: bool = False
+
+
+def _add_callback_relevant_args(parser: argparse.ArgumentParser) -> None:
+    """adds arguments relevant to hyperion-callbacks."""
+    parser.add_argument(
+        "--dev",
+        action="store_true",
+        help="Use dev options, such as local graylog instances and S03",
+    )
+
+
+def parse_callback_dev_mode_arg() -> bool:
+    """Returns the bool representing the 'dev_mode' argument."""
+    parser = argparse.ArgumentParser()
+    _add_callback_relevant_args(parser)
+    args = parser.parse_args()
+    return args.dev
+
+
+def parse_cli_args() -> HyperionArgs:
+    """Parses all arguments relevant to hyperion. Returns an HyperionArgs dataclass with
+    the fields: (verbose_event_logging: bool,
+                 dev_mode: bool,
+                 skip_startup_connection: bool,
+                 external_callbacks: bool)"""
+    parser = argparse.ArgumentParser()
+    _add_callback_relevant_args(parser)
+    parser.add_argument(
+        "--verbose-event-logging",
+        action="store_true",
+        help="Log all bluesky event documents to graylog",
+    )
+    parser.add_argument(
+        "--skip-startup-connection",
+        action="store_true",
+        help="Skip connecting to EPICS PVs on startup",
+    )
+    parser.add_argument(
+        "--external-callbacks",
+        action="store_true",
+        help="Run the external hyperion-callbacks service and publish events over ZMQ",
+    )
+    parser.add_argument(
+        "--version",
+        help="Print hyperion version string",
+        action="version",
+        version=version,
+    )
+    args = parser.parse_args()
+    return HyperionArgs(
+        verbose_event_logging=args.verbose_event_logging or False,
+        dev_mode=args.dev or False,
+        skip_startup_connection=args.skip_startup_connection or False,
+        use_external_callbacks=args.external_callbacks or False,
+    )

mx_bluesky/{parameters → hyperion/parameters}/components.py

@@ -1,20 +1,32 @@
 from __future__ import annotations
 
+import datetime
+import json
 from abc import abstractmethod
 from collections.abc import Sequence
-from enum import Enum
+from enum import StrEnum
 from pathlib import Path
 from typing import SupportsInt, TypeVar
 
-from dodal.devices.aperturescatterguard import AperturePositionGDANames
+from dodal.devices.aperturescatterguard import ApertureValue
 from dodal.devices.detector import (
     DetectorParams,
     TriggerMode,
 )
-from pydantic import BaseModel, Field, root_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
 from scanspec.core import AxesPoints
 from semver import Version
 
+from mx_bluesky.hyperion.external_interaction.config_server import FeatureFlags
+from mx_bluesky.hyperion.parameters.constants import CONST
+
 T = TypeVar("T")
 
 
@@ -25,31 +37,24 @@ class ParameterVersion(Version):
             return version
         return cls.parse(version)
 
-    @classmethod
-    def __get_validators__(cls):
-        """Return a list of validator methods for pydantic models."""
-        yield cls._parse
 
-    @classmethod
-    def __modify_schema__(cls, field_schema):
-        """Inject/mutate the pydantic field schema in-place."""
-        field_schema.update(examples=["1.0.2", "2.15.3-alpha", "21.3.15-beta+12345"])
+PARAMETER_VERSION = ParameterVersion.parse("5.1.0")
 
 
-class RotationAxis(str, Enum):
+class RotationAxis(StrEnum):
     OMEGA = "omega"
     PHI = "phi"
     CHI = "chi"
     KAPPA = "kappa"
 
 
-class XyzAxis(str, Enum):
+class XyzAxis(StrEnum):
     X = "sam_x"
     Y = "sam_y"
     Z = "sam_z"
 
 
-class IspybExperimentType(str, Enum):
+class IspybExperimentType(StrEnum):
     # Enum values from ispyb column data type
     SAD = "SAD"  # at or slightly above the peak
     SAD_INVERSE_BEAM = "SAD - Inverse Beam"
@@ -94,36 +99,75 @@ class IspybExperimentType(str, Enum):
     GRIDSCAN_3D = "Mesh3D"
 
 
+class HyperionParameters(BaseModel):
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    def __hash__(self) -> int:
+        return self.json().__hash__()
+
+    features: FeatureFlags = Field(default=FeatureFlags())
+    parameter_model_version: ParameterVersion
+
+    @field_serializer("parameter_model_version")
+    def serialize_parameter_version(self, version: ParameterVersion):
+        return str(version)
+
+    @field_validator("parameter_model_version", mode="before")
+    @classmethod
+    def _validate_version(cls, version_str: str):
+        version = ParameterVersion.parse(version_str)
+        assert (
+            version >= ParameterVersion(major=PARAMETER_VERSION.major)
+        ), f"Parameter version too old! This version of hyperion uses {PARAMETER_VERSION}"
+        assert (
+            version <= ParameterVersion(major=PARAMETER_VERSION.major + 1)
+        ), f"Parameter version too new! This version of hyperion uses {PARAMETER_VERSION}"
+        return version
+
+    @classmethod
+    def from_json(cls, input: str | None):
+        assert input is not None
+        return cls(**json.loads(input))
+
+
 class WithSnapshot(BaseModel):
     snapshot_directory: Path
-    snapshot_omegas_deg: list[float] | None
+    snapshot_omegas_deg: list[float] | None = None
 
     @property
     def take_snapshots(self) -> bool:
         return bool(self.snapshot_omegas_deg)
 
 
-class DiffractionExperiment(WithSnapshot):
+class DiffractionExperiment(HyperionParameters, WithSnapshot):
     """For all experiments which use beam"""
 
-    visit: str = Field(min_length=5, regex=r"[\w]{2}[\d]+-[\d]+")
+    visit: str = Field(min_length=1)
     file_name: str
     exposure_time_s: float = Field(gt=0)
     comment: str = Field(default="")
-    beamline: str = Field(regex=r"BL\d{2}[BIJS]")
-    insertion_prefix: str = Field(regex=r"SR\d{2}[BIJS]")
-    det_dist_to_beam_converter_path: str
-    zocalo_environment: str
+    beamline: str = Field(default=CONST.I03.BEAMLINE, pattern=r"BL\d{2}[BIJS]")
+    insertion_prefix: str = Field(
+        default=CONST.I03.INSERTION_PREFIX, pattern=r"SR\d{2}[BIJS]"
+    )
+    det_dist_to_beam_converter_path: str = Field(
+        default=CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
+    )
+    zocalo_environment: str = Field(default=CONST.ZOCALO_ENV)
     trigger_mode: TriggerMode = Field(default=TriggerMode.FREE_RUN)
     detector_distance_mm: float | None = Field(default=None, gt=0)
     demand_energy_ev: float | None = Field(default=None, gt=0)
     run_number: int | None = Field(default=None, ge=0)
-    selected_aperture: AperturePositionGDANames | None = Field(default=None)
+    selected_aperture: ApertureValue | None = Field(default=None)
     transmission_frac: float = Field(default=0.1)
     ispyb_experiment_type: IspybExperimentType
     storage_directory: str
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def validate_snapshot_directory(cls, values):
         snapshot_dir = values.get(
             "snapshot_directory", Path(values["storage_directory"], "snapshots")
@@ -133,6 +177,12 @@ class DiffractionExperiment(WithSnapshot):
         )
         return values
 
+    @property
+    def visit_directory(self) -> Path:
+        return (
+            Path(CONST.I03.BASE_DATA_DIR) / str(datetime.date.today().year) / self.visit
+        )
+
     @property
     def num_images(self) -> int:
         return 0
@@ -168,8 +218,11 @@ class WithSample(BaseModel):
     sample_pin: int | None = None
 
 
+class DiffractionExperimentWithSample(DiffractionExperiment, WithSample): ...
+
+
 class WithOavCentring(BaseModel):
-    oav_centring_file: str
+    oav_centring_file: str = Field(default=CONST.I03.OAV_CENTRING_FILE)
 
 
 class OptionalXyzStarts(BaseModel):

mx_bluesky/hyperion/parameters/constants.py

@@ -0,0 +1,158 @@
+import os
+from enum import Enum
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import EIGER2_X_16M_SIZE
+from pydantic.dataclasses import dataclass
+
+TEST_MODE = os.environ.get("HYPERION_TEST_MODE")
+
+
+@dataclass(frozen=True)
+class SimConstants:
+    BEAMLINE = "BL03S"
+    INSERTION_PREFIX = "SR03S"
+    ZOCALO_ENV = "dev_artemis"
+    # 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"
+
+
+@dataclass(frozen=True)
+class PlanNameConstants:
+    # 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"
+    # Rotation scan
+    ROTATION_MULTI = "multi_rotation_wrapper"
+    ROTATION_OUTER = "rotation_scan_with_cleanup"
+    ROTATION_MAIN = "rotation_scan_main"
+
+
+@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 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"
+    ZOCALO_HW_READ = "zocalo_read_hardware_plan"
+
+
+@dataclass(frozen=True)
+class HardwareConstants:
+    OAV_REFRESH_DELAY = 0.3
+    PANDA_FGS_RUN_UP_DEFAULT = 0.17
+    CRYOJET_MARGIN_MM = 0.2
+
+
+@dataclass(frozen=True)
+class TriggerConstants:
+    ZOCALO = "trigger_zocalo_on"
+
+
+@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
+
+
+@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 "/dls_sw/i03/software/daq_configuration/lookup/DetDistToBeamXYConverter.txt"
+    )
+
+
+@dataclass(frozen=True)
+class ExperimentParamConstants:
+    DETECTOR = DetectorParamConstants()
+    GRIDSCAN = GridscanParamConstants()
+    ROTATION = RotationParamConstants()
+
+
+_test_oav_file = "tests/test_data/test_OAVCentring.json"
+_live_oav_file = "/dls_sw/i03/software/daq_configuration/json/OAVCentring_hyperion.json"
+
+
+@dataclass(frozen=True)
+class I03Constants:
+    BASE_DATA_DIR = "/tmp/dls/i03/data/" if TEST_MODE else "/dls/i03/data/"
+    BEAMLINE = "BL03S" if TEST_MODE else "BL03I"
+    DETECTOR = EIGER2_X_16M_SIZE
+    INSERTION_PREFIX = "SR03S" if TEST_MODE else "SR03I"
+    OAV_CENTRING_FILE = _test_oav_file if TEST_MODE else _live_oav_file
+    SHUTTER_TIME_S = 0.06
+    USE_PANDA_FOR_GRIDSCAN = False
+    USE_GPU_FOR_GRIDSCAN_ANALYSIS = False
+    THAWING_TIME = 20
+    USE_CPU_AND_GPU_ZOCALO = False
+
+
+@dataclass(frozen=True)
+class HyperionConstants:
+    HARDWARE = HardwareConstants()
+    I03 = I03Constants()
+    PARAM = ExperimentParamConstants()
+    PLAN = PlanNameConstants()
+    WAIT = PlanGroupCheckpointConstants()
+    SIM = SimConstants()
+    TRIGGER = TriggerConstants()
+    CALLBACK_0MQ_PROXY_PORTS = (5577, 5578)
+    DESCRIPTORS = DocDescriptorNames()
+    CONFIG_SERVER_URL = (
+        "http://fake-url-not-real"
+        if TEST_MODE
+        else "https://daq-config.diamond.ac.uk/api"
+    )
+    GRAYLOG_PORT = 12232
+    PARAMETER_SCHEMA_DIRECTORY = "src/hyperion/parameters/schemas/"
+    ZOCALO_ENV = "dev_artemis" if TEST_MODE else "artemis"
+
+
+CONST = HyperionConstants()
+
+
+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/hyperion/parameters/gridscan.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import os
+
+from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector import (
+    DetectorParams,
+)
+from dodal.devices.fast_grid_scan import (
+    PandAGridScanParams,
+    ZebraGridScanParams,
+)
+from pydantic import Field, PrivateAttr
+from scanspec.core import Path as ScanPath
+from scanspec.specs import Line, Static
+
+from mx_bluesky.hyperion.parameters.components import (
+    DiffractionExperimentWithSample,
+    IspybExperimentType,
+    OptionalGonioAngleStarts,
+    SplitScan,
+    WithOavCentring,
+    WithScan,
+    XyzStarts,
+)
+from mx_bluesky.hyperion.parameters.constants import CONST, I03Constants
+
+
+class GridCommon(
+    DiffractionExperimentWithSample, OptionalGonioAngleStarts, WithOavCentring
+):
+    grid_width_um: float = Field(default=CONST.PARAM.GRIDSCAN.WIDTH_UM)
+    exposure_time_s: float = Field(default=CONST.PARAM.GRIDSCAN.EXPOSURE_TIME_S)
+    use_roi_mode: bool = Field(default=CONST.PARAM.GRIDSCAN.USE_ROI)
+    panda_runup_distance_mm: float = Field(
+        default=CONST.HARDWARE.PANDA_FGS_RUN_UP_DEFAULT
+    )
+    use_panda: bool = Field(default=CONST.I03.USE_PANDA_FOR_GRIDSCAN)
+    use_gpu: bool = Field(default=CONST.I03.USE_GPU_FOR_GRIDSCAN_ANALYSIS)
+    use_cpu_and_gpu_zocalo: bool = Field(default=CONST.I03.USE_CPU_AND_GPU_ZOCALO)
+    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 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
+        ), "Detector distance must be filled before generating DetectorParams"
+        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=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,
+            enable_dev_shm=self.use_gpu,
+            **optional_args,
+        )
+
+
+class GridScanWithEdgeDetect(GridCommon): ...
+
+
+class PinTipCentreThenXrayCentre(GridCommon):
+    tip_offset_um: float = 0
+
+
+class RobotLoadThenCentre(GridCommon):
+    thawing_time: float = Field(default=CONST.I03.THAWING_TIME)
+
+    def pin_centre_then_xray_centre_params(self):
+        my_params = self.model_dump()
+        del my_params["thawing_time"]
+        return PinTipCentreThenXrayCentre(**my_params)
+
+
+class SpecifiedGridScan(GridCommon, XyzStarts, WithScan):
+    """A specified grid scan 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)."""
+
+    ...
+
+
+class ThreeDGridScan(SpecifiedGridScan, SplitScan):
+    """Parameters representing a so-called 3D grid scan, which consists of doing a
+    gridscan in X and Y, followed by one in X and Z."""
+
+    demand_energy_ev: float | None = Field(default=None)
+    grid1_omega_deg: float = Field(default=CONST.PARAM.GRIDSCAN.OMEGA_1)  # type: ignore
+    grid2_omega_deg: float = Field(default=CONST.PARAM.GRIDSCAN.OMEGA_2)
+    x_step_size_um: float = Field(default=CONST.PARAM.GRIDSCAN.BOX_WIDTH_UM)
+    y_step_size_um: float = Field(default=CONST.PARAM.GRIDSCAN.BOX_WIDTH_UM)
+    z_step_size_um: float = Field(default=CONST.PARAM.GRIDSCAN.BOX_WIDTH_UM)
+    y2_start_um: float
+    z2_start_um: float
+    x_steps: int = Field(gt=0)
+    y_steps: int = Field(gt=0)
+    z_steps: int = Field(gt=0)
+    _set_stub_offsets: bool = PrivateAttr(default_factory=lambda: False)
+
+    @property
+    def FGS_params(self) -> ZebraGridScanParams:
+        return ZebraGridScanParams(
+            x_steps=self.x_steps,
+            y_steps=self.y_steps,
+            z_steps=self.z_steps,
+            x_step_size=self.x_step_size_um,
+            y_step_size=self.y_step_size_um,
+            z_step_size=self.z_step_size_um,
+            x_start=self.x_start_um,
+            y1_start=self.y_start_um,
+            z1_start=self.z_start_um,
+            y2_start=self.y2_start_um,
+            z2_start=self.z2_start_um,
+            set_stub_offsets=self.features.set_stub_offsets,
+            dwell_time_ms=self.exposure_time_s * 1000,
+            transmission_fraction=self.transmission_frac,
+        )
+
+    @property
+    def panda_FGS_params(self) -> PandAGridScanParams:
+        if self.y_steps % 2 and self.z_steps > 0:
+            raise OddYStepsException(
+                "The number of Y steps must be even for a PandA gridscan"
+            )
+        return PandAGridScanParams(
+            x_steps=self.x_steps,
+            y_steps=self.y_steps,
+            z_steps=self.z_steps,
+            x_step_size=self.x_step_size_um,
+            y_step_size=self.y_step_size_um,
+            z_step_size=self.z_step_size_um,
+            x_start=self.x_start_um,
+            y1_start=self.y_start_um,
+            z1_start=self.z_start_um,
+            y2_start=self.y2_start_um,
+            z2_start=self.z2_start_um,
+            set_stub_offsets=self.features.set_stub_offsets,
+            run_up_distance_mm=self.panda_runup_distance_mm,
+            transmission_fraction=self.transmission_frac,
+        )
+
+    def do_set_stub_offsets(self, value: bool):
+        self._set_stub_offsets = value
+
+    @property
+    def grid_1_spec(self):
+        x_end = self.x_start_um + self.x_step_size_um * (self.x_steps - 1)
+        y1_end = self.y_start_um + self.y_step_size_um * (self.y_steps - 1)
+        grid_1_x = Line("sam_x", self.x_start_um, x_end, self.x_steps)
+        grid_1_y = Line("sam_y", self.y_start_um, y1_end, self.y_steps)
+        grid_1_z = Static("sam_z", self.z_start_um)
+        return grid_1_y.zip(grid_1_z) * ~grid_1_x
+
+    @property
+    def grid_2_spec(self):
+        x_end = self.x_start_um + self.x_step_size_um * (self.x_steps - 1)
+        z2_end = self.z2_start_um + self.z_step_size_um * (self.z_steps - 1)
+        grid_2_x = Line("sam_x", self.x_start_um, x_end, self.x_steps)
+        grid_2_z = Line("sam_z", self.z2_start_um, z2_end, self.z_steps)
+        grid_2_y = Static("sam_y", self.y2_start_um)
+        return grid_2_z.zip(grid_2_y) * ~grid_2_x
+
+    @property
+    def scan_indices(self):
+        """The first index of each gridscan, useful for writing nexus files/VDS"""
+        return [
+            0,
+            len(ScanPath(self.grid_1_spec.calculate()).consume().midpoints["sam_x"]),
+        ]
+
+    @property
+    def scan_spec(self):
+        """A fully specified ScanSpec object representing both grids, with x, y, z and
+        omega positions."""
+        return self.grid_1_spec.concat(self.grid_2_spec)
+
+    @property
+    def scan_points(self):
+        """A list of all the points in the scan_spec."""
+        return ScanPath(self.scan_spec.calculate()).consume().midpoints
+
+    @property
+    def scan_points_first_grid(self):
+        """A list of all the points in the first grid scan."""
+        return ScanPath(self.grid_1_spec.calculate()).consume().midpoints
+
+    @property
+    def scan_points_second_grid(self):
+        """A list of all the points in the second grid scan."""
+        return ScanPath(self.grid_2_spec.calculate()).consume().midpoints
+
+    @property
+    def num_images(self) -> int:
+        return len(self.scan_points["sam_x"])
+
+
+class OddYStepsException(Exception): ...