mx-bluesky 1.4.0__py3-none-any.whl → 1.4.1a0__py3-none-any.whl

mx_bluesky/common/external_interaction/config_server.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+from functools import cache
+
+from daq_config_server.client import ConfigServer
+from pydantic import BaseModel, Field, model_validator
+
+
+class FeatureFlags(BaseModel, ABC):
+    """Abstract class to use ConfigServer to toggle features for an experiment
+
+    A module wanting to use FeatureFlags should inherit this class, add boolean features
+    as attributes, and implement a get_config_server method, which returns a cached creation of
+    ConfigServer. See HyperionFeatureFlags for an example
+
+    Values supplied upon class instantiation will always take priority over the config server. If connection to the server cannot
+    be made AND values were not supplied, attributes will use their default values
+    """
+
+    # Feature values supplied at construction will override values from the config server
+    overriden_features: dict = Field(default_factory=dict, exclude=True)
+
+    @staticmethod
+    @cache
+    @abstractmethod
+    def get_config_server() -> ConfigServer: ...
+
+    @model_validator(mode="before")
+    @classmethod
+    def mark_overridden_features(cls, values):
+        assert isinstance(values, dict)
+        values["overriden_features"] = values.copy()
+        return values
+
+    def _get_flags(self):
+        flags = type(self).get_config_server().best_effort_get_all_feature_flags()
+        return {f: flags[f] for f in flags if f in self.model_fields.keys()}
+
+    def update_self_from_server(self):
+        """Used to update the feature flags from the server during a plan. Where there are flags which were explicitly set from externally supplied parameters, these values will be used instead."""
+        for flag, value in self._get_flags().items():
+            updated_value = (
+                value
+                if flag not in self.overriden_features.keys()
+                else self.overriden_features[flag]
+            )
+            setattr(self, flag, updated_value)

mx_bluesky/common/parameters/components.py

@@ -1,21 +1,33 @@
-import json
+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 SupportsInt
+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 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
+from mx_bluesky.common.parameters.constants import (
+    TEST_MODE,
+    DetectorParamConstants,
+    GridscanParamConstants,
+)
 
 PARAMETER_VERSION = Version.parse("5.2.0")
 
@@ -100,11 +112,6 @@ class MxBlueskyParameters(BaseModel):
         ), 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
@@ -146,12 +153,12 @@ class DiffractionExperiment(
 
     @model_validator(mode="before")
     @classmethod
-    def validate_snapshot_directory(cls, values):
-        snapshot_dir = values.get(
-            "snapshot_directory", Path(values["storage_directory"], "snapshots")
-        )
-        values["snapshot_directory"] = (
-            snapshot_dir if isinstance(snapshot_dir, Path) else Path(snapshot_dir)
+    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
 
@@ -176,6 +183,14 @@ class WithScan(BaseModel):
     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
@@ -193,6 +208,28 @@ class WithSample(BaseModel):
 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

mx_bluesky/common/parameters/constants.py

@@ -1,6 +1,7 @@
 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
@@ -19,6 +20,7 @@ class DocDescriptorNames:
     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"
+    FLYSCAN_RESULTS = "flyscan_results_obtained"
 
 
 @dataclass(frozen=True)
@@ -47,6 +49,7 @@ class PlanNameConstants:
     ROTATION_MULTI = "multi_rotation_wrapper"
     ROTATION_OUTER = "rotation_scan_with_cleanup"
     ROTATION_MAIN = "rotation_scan_main"
+    FLYSCAN_RESULTS = "xray_centre_results"
 
 
 @dataclass(frozen=True)
@@ -64,6 +67,7 @@ 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)
@@ -87,8 +91,9 @@ class DetectorParamConstants:
     BEAM_XY_LUT_PATH = (
         "tests/test_data/test_det_dist_converter.txt"
         if TEST_MODE
-        else "/dls_sw/{BEAMLINE}/software/daq_configuration/lookup/DetDistToBeamXYConverter.txt"
+        else f"/dls_sw/{BEAMLINE}/software/daq_configuration/lookup/DetDistToBeamXYConverter.txt"
     )
+    DETECTOR = EIGER2_X_16M_SIZE
 
 
 @dataclass(frozen=True)

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/{hyperion → common}/parameters/robot_load.py

@@ -7,10 +7,10 @@ from mx_bluesky.common.parameters.components import (
     WithSnapshot,
     WithVisit,
 )
-from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.common.parameters.constants import HardwareConstants
 
 
 class RobotLoadAndEnergyChange(
     MxBlueskyParameters, WithSample, WithSnapshot, WithOptionalEnergyChange, WithVisit
 ):
-    thawing_time: float = Field(default=CONST.I03.THAWING_TIME)
+    thawing_time: float = Field(default=HardwareConstants.THAWING_TIME)

mx_bluesky/common/plans/do_fgs.py

@@ -3,7 +3,7 @@ from time import time
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
-from blueapi.core import MsgGenerator
+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
@@ -11,7 +11,7 @@ from dodal.devices.zocalo.zocalo_results import (
     ZOCALO_STAGE_GROUP,
 )
 from dodal.log import LOGGER
-from dodal.plans.check_topup import check_topup_and_wait_if_necessary
+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 (

mx_bluesky/common/utils/log.py

@@ -12,6 +12,8 @@ from dodal.log import (
 )
 from dodal.log import LOGGER as dodal_logger
 
+LOGGER = logging.getLogger("mx-bluesky")
+
 __logger_handlers: DodalLogHandlers | None = None
 
 

mx_bluesky/hyperion/__main__.py

@@ -7,9 +7,10 @@ from queue import Queue
 from traceback import format_exception
 from typing import Any
 
-from blueapi.core import BlueskyContext, MsgGenerator
+from blueapi.core import BlueskyContext
 from bluesky.callbacks.zmq import Publisher
 from bluesky.run_engine import RunEngine
+from bluesky.utils import MsgGenerator
 from flask import Flask, request
 from flask_restful import Api, Resource
 from pydantic.dataclasses import dataclass

mx_bluesky/hyperion/device_setup_plans/setup_panda.py

@@ -4,7 +4,7 @@ from importlib import resources
 from pathlib import Path
 
 import bluesky.plan_stubs as bps
-from blueapi.core import MsgGenerator
+from bluesky.utils import MsgGenerator
 from dodal.common.beamlines.beamline_utils import get_path_provider
 from dodal.devices.fast_grid_scan import PandAGridScanParams
 from ophyd_async.core import load_device
@@ -63,12 +63,12 @@ def _get_seq_table(
         An instance of SeqTable describing the panda sequencer table
     """
 
-    start_of_grid_x_counts = int(parameters.x_start * MM_TO_ENCODER_COUNTS)
+    start_of_grid_x_counts = int(parameters.x_start_mm * MM_TO_ENCODER_COUNTS)
 
     # x_start is the first trigger point, so we need to travel to x_steps-1 for the final trigger point
     end_of_grid_x_counts = int(
         start_of_grid_x_counts
-        + (parameters.x_step_size * (parameters.x_steps - 1) * MM_TO_ENCODER_COUNTS)
+        + (parameters.x_step_size_mm * (parameters.x_steps - 1) * MM_TO_ENCODER_COUNTS)
     )
 
     exposure_distance_x_counts = int(exposure_distance_mm * MM_TO_ENCODER_COUNTS)
@@ -140,7 +140,7 @@ def setup_panda_for_flyscan(
     """
     assert parameters.x_steps > 0
     assert time_between_x_steps_ms * 1000 >= exposure_time_s
-    assert sample_velocity_mm_per_s * exposure_time_s < parameters.x_step_size
+    assert sample_velocity_mm_per_s * exposure_time_s < parameters.x_step_size_mm
 
     yield from bps.stage(panda, group="panda-config")
 

mx_bluesky/hyperion/device_setup_plans/setup_zebra.py

@@ -3,7 +3,7 @@ from functools import wraps
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
-from blueapi.core import MsgGenerator
+from bluesky.utils import MsgGenerator
 from dodal.devices.zebra import (
     AUTO_SHUTTER_GATE,
     AUTO_SHUTTER_INPUT_1,

mx_bluesky/hyperion/experiment_plans/__init__.py

@@ -9,6 +9,9 @@ from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     grid_detect_then_xray_centre,
 )
+from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
+    load_centre_collect_full,
+)
 from mx_bluesky.hyperion.experiment_plans.pin_centre_then_xray_centre_plan import (
     pin_tip_centre_then_xray_centre,
 )
@@ -27,4 +30,5 @@ __all__ = [
     "pin_tip_centre_then_xray_centre",
     "multi_rotation_scan",
     "robot_load_then_centre",
+    "load_centre_collect_full",
 ]

mx_bluesky/hyperion/experiment_plans/change_aperture_then_move_plan.py

@@ -0,0 +1,83 @@
+import bluesky.plan_stubs as bps
+import bluesky.preprocessors as bpp
+import numpy
+from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
+from dodal.devices.smargon import Smargon, StubPosition
+
+from mx_bluesky.common.utils.tracing import TRACER
+from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.hyperion.experiment_plans.common.xrc_result import XRayCentreResult
+from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
+
+
+def change_aperture_then_move_to_xtal(
+    best_hit: XRayCentreResult,
+    smargon: Smargon,
+    aperture_scatterguard: ApertureScatterguard,
+    parameters: HyperionThreeDGridScan | None = None,
+):
+    """For the given x-ray centring result,
+    * Change the aperture so that the beam size is comparable to the crystal size
+    * Centre on the centre-of-mass
+    * Reset the stub offsets if specified by params"""
+    if best_hit.bounding_box_mm is not None:
+        bounding_box_size = numpy.abs(
+            best_hit.bounding_box_mm[1] - best_hit.bounding_box_mm[0]
+        )
+        with TRACER.start_span("change_aperture"):
+            yield from _set_aperture_for_bbox_mm(
+                aperture_scatterguard, bounding_box_size
+            )
+    else:
+        LOGGER.warning("No bounding box size received")
+
+    # once we have the results, go to the appropriate position
+    LOGGER.info("Moving to centre of mass.")
+    with TRACER.start_span("move_to_result"):
+        x, y, z = best_hit.centre_of_mass_mm
+        yield from move_x_y_z(smargon, x, y, z, wait=True)
+
+    # TODO support for setting stub offsets in multipin
+    # https://github.com/DiamondLightSource/mx-bluesky/issues/552
+    if parameters and parameters.FGS_params.set_stub_offsets:
+        LOGGER.info("Recentring smargon co-ordinate system to this point.")
+        yield from bps.mv(
+            # See: https://github.com/bluesky/bluesky/issues/1809
+            smargon.stub_offsets,  # type: ignore
+            StubPosition.CURRENT_AS_CENTER,  # type: ignore
+        )
+
+
+def _set_aperture_for_bbox_mm(
+    aperture_device: ApertureScatterguard, bbox_size_mm: list[float] | numpy.ndarray
+):
+    # TODO confirm correction factor see https://github.com/DiamondLightSource/mx-bluesky/issues/618
+    ASSUMED_BOX_SIZE_MM = 0.020
+    bbox_size_boxes = [round(mm / ASSUMED_BOX_SIZE_MM) for mm in bbox_size_mm]
+    yield from set_aperture_for_bbox_size(aperture_device, bbox_size_boxes)
+
+
+def set_aperture_for_bbox_size(
+    aperture_device: ApertureScatterguard,
+    bbox_size: list[int] | numpy.ndarray,
+):
+    # bbox_size is [x,y,z], for i03 we only care about x
+    new_selected_aperture = (
+        ApertureValue.MEDIUM if bbox_size[0] < 2 else ApertureValue.LARGE
+    )
+    LOGGER.info(
+        f"Setting aperture to {new_selected_aperture} based on bounding box size {bbox_size}."
+    )
+
+    @bpp.set_run_key_decorator("change_aperture")
+    @bpp.run_decorator(
+        md={
+            "subplan_name": "change_aperture",
+            "aperture_size": new_selected_aperture.value,
+        }
+    )
+    def set_aperture():
+        yield from bps.abs_set(aperture_device, new_selected_aperture)
+
+    yield from set_aperture()

mx_bluesky/hyperion/experiment_plans/common/xrc_result.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable, Sequence
+from functools import partial
+
+import numpy as np
+
+from mx_bluesky.common.parameters.components import (
+    MultiXtalSelection,
+    TopNByMaxCountSelection,
+)
+
+
+@dataclasses.dataclass
+class XRayCentreResult:
+    """Represents information about a hit from an X-ray centring."""
+
+    centre_of_mass_mm: np.ndarray
+    bounding_box_mm: tuple[np.ndarray, np.ndarray]
+    max_count: int
+    total_count: int
+
+    def __eq__(self, o):
+        return (
+            isinstance(o, XRayCentreResult)
+            and o.max_count == self.max_count
+            and o.total_count == self.total_count
+            and all(o.centre_of_mass_mm == self.centre_of_mass_mm)
+            and all(o.bounding_box_mm[0] == self.bounding_box_mm[0])
+            and all(o.bounding_box_mm[1] == self.bounding_box_mm[1])
+        )
+
+
+def top_n_by_max_count(
+    unfiltered: Sequence[XRayCentreResult], n: int
+) -> Sequence[XRayCentreResult]:
+    sorted_hits = sorted(unfiltered, key=lambda result: result.max_count, reverse=True)
+    return sorted_hits[:n]
+
+
+def resolve_selection_fn(
+    params: MultiXtalSelection,
+) -> Callable[[Sequence[XRayCentreResult]], Sequence[XRayCentreResult]]:
+    if isinstance(params, TopNByMaxCountSelection):
+        return partial(top_n_by_max_count, n=params.n)
+    raise ValueError(f"Invalid selection function {params.name}")

mx_bluesky/hyperion/experiment_plans/experiment_registry.py

@@ -5,6 +5,11 @@ from typing import TypedDict
 
 import mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan as flyscan_xray_centre_plan
 import mx_bluesky.hyperion.experiment_plans.rotation_scan_plan as rotation_scan_plan
+from mx_bluesky.common.parameters.gridscan import (
+    GridScanWithEdgeDetect,
+    PinTipCentreThenXrayCentre,
+    RobotLoadThenCentre,
+)
 from mx_bluesky.hyperion.experiment_plans import (
     grid_detect_then_xray_centre_plan,
     load_centre_collect_full_plan,
@@ -18,12 +23,7 @@ from mx_bluesky.hyperion.external_interaction.callbacks.common.callback_util imp
     create_robot_load_and_centre_callbacks,
     create_rotation_callbacks,
 )
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    PinTipCentreThenXrayCentre,
-    RobotLoadThenCentre,
-    ThreeDGridScan,
-)
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
 from mx_bluesky.hyperion.parameters.rotation import MultiRotationScan, RotationScan
 
@@ -39,7 +39,7 @@ def do_nothing():
 class ExperimentRegistryEntry(TypedDict):
     setup: Callable
     param_type: type[
-        ThreeDGridScan
+        HyperionThreeDGridScan
         | GridScanWithEdgeDetect
         | RotationScan
         | MultiRotationScan
@@ -53,7 +53,7 @@ class ExperimentRegistryEntry(TypedDict):
 PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
     "flyscan_xray_centre": {
         "setup": flyscan_xray_centre_plan.create_devices,
-        "param_type": ThreeDGridScan,
+        "param_type": HyperionThreeDGridScan,
         "callbacks_factory": create_gridscan_callbacks,
     },
     "grid_detect_then_xray_centre": {
@@ -81,7 +81,7 @@ PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
         "param_type": MultiRotationScan,
         "callbacks_factory": create_rotation_callbacks,
     },
-    "load_centre_collect_full_plan": {
+    "load_centre_collect_full": {
         "setup": load_centre_collect_full_plan.create_devices,
         "param_type": LoadCentreCollect,
         "callbacks_factory": create_load_centre_collect_callbacks,