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

mx_bluesky/hyperion/experiment_plans/grid_detect_then_xray_centre_plan.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
-import dataclasses
 from pathlib import Path
 
-from blueapi.core import BlueskyContext, MsgGenerator
+import pydantic
+from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
 from bluesky import preprocessors as bpp
+from bluesky.preprocessors import subs_decorator
+from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.backlight import Backlight, BacklightPosition
@@ -29,17 +31,22 @@ from dodal.devices.zocalo import ZocaloResults
 from ophyd_async.fastcs.panda import HDFPanda
 
 from mx_bluesky.common.parameters.constants import OavConstants
+from mx_bluesky.common.parameters.gridscan import GridScanWithEdgeDetect
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import (
     move_aperture_if_required,
 )
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
 )
+from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
+    change_aperture_then_move_to_xtal,
+)
 from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
     FlyScanXRayCentreComposite as FlyScanXRayCentreComposite,
 )
 from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
-    flyscan_xray_centre,
+    XRayCentreEventHandler,
+    flyscan_xray_centre_no_move,
 )
 from mx_bluesky.hyperion.experiment_plans.oav_grid_detection_plan import (
     OavGridDetectionComposite,
@@ -55,13 +62,12 @@ from mx_bluesky.hyperion.external_interaction.callbacks.xray_centre.ispyb_callba
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
 from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    ThreeDGridScan,
+    HyperionThreeDGridScan,
 )
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class GridDetectThenXRayCentreComposite:
     """All devices which are directly or indirectly required by this plan"""
 
@@ -87,6 +93,10 @@ class GridDetectThenXRayCentreComposite:
     robot: BartRobot
     sample_shutter: ZebraShutter
 
+    @property
+    def sample_motors(self):
+        return self.smargon
+
 
 def create_devices(context: BlueskyContext) -> GridDetectThenXRayCentreComposite:
     return device_composite_from_context(context, GridDetectThenXRayCentreComposite)
@@ -95,10 +105,10 @@ def create_devices(context: BlueskyContext) -> GridDetectThenXRayCentreComposite
 def create_parameters_for_flyscan_xray_centre(
     grid_scan_with_edge_params: GridScanWithEdgeDetect,
     grid_parameters: GridParamUpdate,
-) -> ThreeDGridScan:
+) -> HyperionThreeDGridScan:
     params_json = grid_scan_with_edge_params.model_dump()
     params_json.update(grid_parameters)
-    flyscan_xray_centre_parameters = ThreeDGridScan(**params_json)
+    flyscan_xray_centre_parameters = HyperionThreeDGridScan(**params_json)
     LOGGER.info(f"Parameters for FGS: {flyscan_xray_centre_parameters}")
     return flyscan_xray_centre_parameters
 
@@ -150,7 +160,7 @@ def detect_grid_and_do_gridscan(
         group=CONST.WAIT.GRID_READY_FOR_DC,
     )
 
-    yield from flyscan_xray_centre(
+    yield from flyscan_xray_centre_no_move(
         FlyScanXRayCentreComposite(
             aperture_scatterguard=composite.aperture_scatterguard,
             attenuator=composite.attenuator,
@@ -193,19 +203,33 @@ def grid_detect_then_xray_centre(
 
     oav_params = OAVParameters("xrayCentring", oav_config)
 
-    plan_to_perform = ispyb_activation_wrapper(
-        detect_grid_and_do_gridscan(
-            composite,
+    flyscan_event_handler = XRayCentreEventHandler()
+
+    @subs_decorator(flyscan_event_handler)
+    def plan_to_perform():
+        yield from ispyb_activation_wrapper(
+            detect_grid_and_do_gridscan(
+                composite,
+                parameters,
+                oav_params,
+            ),
             parameters,
-            oav_params,
-        ),
-        parameters,
-    )
+        )
 
-    return start_preparing_data_collection_then_do_plan(
+    yield from start_preparing_data_collection_then_do_plan(
         eiger,
         composite.detector_motion,
         parameters.detector_params.detector_distance,
-        plan_to_perform,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809 and MsgGenerator should allow for return values
+        plan_to_perform(),
         group=CONST.WAIT.GRID_READY_FOR_DC,
     )
+
+    assert (
+        flyscan_event_handler.xray_centre_results
+    ), "Flyscan result event not received or no crystal found and exception not raised"
+
+    yield from change_aperture_then_move_to_xtal(
+        flyscan_event_handler.xray_centre_results[0],
+        composite.smargon,
+        composite.aperture_scatterguard,
+    )

mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py

@@ -1,25 +1,43 @@
-import dataclasses
+from __future__ import annotations
 
+from collections.abc import Sequence
+
+import pydantic
 from blueapi.core import BlueskyContext
+from bluesky.preprocessors import run_decorator, set_run_key_decorator, subs_wrapper
+from bluesky.utils import MsgGenerator
 from dodal.devices.oav.oav_parameters import OAVParameters
+from dodal.devices.smargon import Smargon
 
+import mx_bluesky.hyperion.experiment_plans.common.xrc_result as flyscan_result
+from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
+    XRayCentreEventHandler,
+)
 from mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan import (
     RobotLoadThenCentreComposite,
-    robot_load_then_centre,
+    robot_load_then_xray_centre,
 )
 from mx_bluesky.hyperion.experiment_plans.rotation_scan_plan import (
+    MultiRotationScan,
     RotationScanComposite,
     multi_rotation_scan,
 )
+from mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback import (
+    sample_handling_callback_decorator,
+)
+from mx_bluesky.hyperion.log import LOGGER
+from mx_bluesky.hyperion.parameters.constants import CONST
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class LoadCentreCollectComposite(RobotLoadThenCentreComposite, RotationScanComposite):
     """Composite that provides access to the required devices."""
 
-    pass
+    @property
+    def sample_motors(self) -> Smargon:
+        return self.smargon
 
 
 def create_devices(context: BlueskyContext) -> LoadCentreCollectComposite:
@@ -27,11 +45,11 @@ def create_devices(context: BlueskyContext) -> LoadCentreCollectComposite:
     return device_composite_from_context(context, LoadCentreCollectComposite)
 
 
-def load_centre_collect_full_plan(
+def load_centre_collect_full(
     composite: LoadCentreCollectComposite,
-    params: LoadCentreCollect,
+    parameters: LoadCentreCollect,
     oav_params: OAVParameters | None = None,
-):
+) -> MsgGenerator:
     """Attempt a complete data collection experiment, consisting of the following:
     * Load the sample if necessary
     * Move to the specified goniometer start angles
@@ -41,6 +59,54 @@ def load_centre_collect_full_plan(
     """
     if not oav_params:
         oav_params = OAVParameters(context="xrayCentring")
-    yield from robot_load_then_centre(composite, params.robot_load_then_centre)
 
-    yield from multi_rotation_scan(composite, params.multi_rotation_scan, oav_params)
+    @set_run_key_decorator(CONST.PLAN.LOAD_CENTRE_COLLECT)
+    @run_decorator(
+        md={
+            "metadata": {"sample_id": parameters.sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    @sample_handling_callback_decorator()
+    def plan_with_callback_subs():
+        flyscan_event_handler = XRayCentreEventHandler()
+        yield from subs_wrapper(
+            robot_load_then_xray_centre(composite, parameters.robot_load_then_centre),
+            flyscan_event_handler,
+        )
+
+        assert flyscan_event_handler.xray_centre_results, "Flyscan result event not received or no crystal found and exception not raised"
+
+        selection_func = flyscan_result.resolve_selection_fn(
+            parameters.selection_params
+        )
+        hits: Sequence[flyscan_result.XRayCentreResult] = selection_func(
+            flyscan_event_handler.xray_centre_results
+        )
+        LOGGER.info(
+            f"Selected hits {hits} using {selection_func}, args={parameters.selection_params}"
+        )
+
+        multi_rotation = parameters.multi_rotation_scan
+        rotation_template = multi_rotation.rotation_scans.copy()
+
+        multi_rotation.rotation_scans.clear()
+
+        for hit in hits:
+            for rot in rotation_template:
+                combination = rot.model_copy()
+                (
+                    combination.x_start_um,
+                    combination.y_start_um,
+                    combination.z_start_um,
+                ) = (axis * 1000 for axis in hit.centre_of_mass_mm)
+                multi_rotation.rotation_scans.append(combination)
+        multi_rotation = MultiRotationScan.model_validate(multi_rotation)
+
+        assert (
+            multi_rotation.demand_energy_ev
+            == parameters.robot_load_then_centre.demand_energy_ev
+        ), "Setting a different energy for gridscan and rotation is not supported"
+        yield from multi_rotation_scan(composite, multi_rotation, oav_params)
+
+    yield from plan_with_callback_subs()

mx_bluesky/hyperion/experiment_plans/oav_grid_detection_plan.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
-import dataclasses
 import math
 from typing import TYPE_CHECKING
 
 import bluesky.plan_stubs as bps
 import numpy as np
+import pydantic
 from blueapi.core import BlueskyContext
 from dodal.devices.backlight import Backlight
 from dodal.devices.oav.oav_detector import OAV
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
     from dodal.devices.oav.oav_parameters import OAVParameters
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class OavGridDetectionComposite:
     """All devices which are directly or indirectly required by this plan"""
 

mx_bluesky/hyperion/experiment_plans/oav_snapshot_plan.py

@@ -1,8 +1,8 @@
 from datetime import datetime
 from typing import Protocol
 
-from blueapi.core import MsgGenerator
 from bluesky import plan_stubs as bps
+from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
 from dodal.devices.backlight import Backlight, BacklightPosition
 from dodal.devices.oav.oav_detector import OAV

mx_bluesky/hyperion/experiment_plans/optimise_attenuation_plan.py

@@ -1,9 +1,9 @@
-import dataclasses
 from enum import Enum
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
 import numpy as np
+import pydantic
 from blueapi.core import BlueskyContext
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.xspress3.xspress3 import Xspress3
@@ -22,7 +22,7 @@ class Direction(Enum):
     NEGATIVE = "negative"
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class OptimizeAttenuationComposite:
     """All devices which are directly or indirectly required by this plan"""
 

mx_bluesky/hyperion/experiment_plans/pin_centre_then_xray_centre_plan.py

@@ -2,15 +2,27 @@ from __future__ import annotations
 
 import json
 
-from blueapi.core import BlueskyContext, MsgGenerator
+import bluesky.preprocessors as bpp
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.oav.oav_parameters import OAVParameters
 
 from mx_bluesky.common.parameters.constants import OavConstants
+from mx_bluesky.common.parameters.gridscan import (
+    GridScanWithEdgeDetect,
+    PinTipCentreThenXrayCentre,
+)
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_phi_chi_omega
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
 )
+from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
+    change_aperture_then_move_to_xtal,
+)
+from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
+    XRayCentreEventHandler,
+)
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     GridDetectThenXRayCentreComposite,
     detect_grid_and_do_gridscan,
@@ -27,10 +39,6 @@ from mx_bluesky.hyperion.external_interaction.callbacks.xray_centre.ispyb_callba
 )
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    PinTipCentreThenXrayCentre,
-)
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
@@ -53,13 +61,12 @@ def create_parameters_for_grid_detection(
     return grid_detect_and_xray_centre
 
 
-def pin_centre_then_xray_centre_plan(
+def pin_centre_then_flyscan_plan(
     composite: GridDetectThenXRayCentreComposite,
     parameters: PinTipCentreThenXrayCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ):
-    """Plan that perfoms a pin tip centre followed by an xray centre to completely
-    centre the sample"""
+    """Plan that performs a pin tip centre followed by a flyscan to determine the centres of interest"""
 
     pin_tip_centring_composite = PinTipCentringComposite(
         oav=composite.oav,
@@ -68,7 +75,7 @@ def pin_centre_then_xray_centre_plan(
         pin_tip_detection=composite.pin_tip_detection,
     )
 
-    def _pin_centre_then_xray_centre_plan():
+    def _pin_centre_then_flyscan_plan():
         yield from setup_beamline_for_OAV(
             composite.smargon, composite.backlight, composite.aperture_scatterguard
         )
@@ -96,7 +103,7 @@ def pin_centre_then_xray_centre_plan(
             oav_params,
         )
 
-    yield from ispyb_activation_wrapper(_pin_centre_then_xray_centre_plan(), parameters)
+    yield from ispyb_activation_wrapper(_pin_centre_then_flyscan_plan(), parameters)
 
 
 def pin_tip_centre_then_xray_centre(
@@ -105,15 +112,27 @@ def pin_tip_centre_then_xray_centre(
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ) -> MsgGenerator:
     """Starts preparing for collection then performs the pin tip centre and xray centre"""
-
     eiger: EigerDetector = composite.eiger
 
     eiger.set_detector_parameters(parameters.detector_params)
 
-    return start_preparing_data_collection_then_do_plan(
-        eiger,
-        composite.detector_motion,
-        parameters.detector_params.detector_distance,
-        pin_centre_then_xray_centre_plan(composite, parameters, oav_config_file),
-        group=CONST.WAIT.GRID_READY_FOR_DC,
+    flyscan_event_handler = XRayCentreEventHandler()
+
+    @bpp.subs_decorator(flyscan_event_handler)
+    def pin_centre_flyscan_then_fetch_results() -> MsgGenerator:
+        yield from start_preparing_data_collection_then_do_plan(
+            eiger,
+            composite.detector_motion,
+            parameters.detector_params.detector_distance,
+            pin_centre_then_flyscan_plan(composite, parameters, oav_config_file),
+            group=CONST.WAIT.GRID_READY_FOR_DC,
+        )
+
+    yield from pin_centre_flyscan_then_fetch_results()
+    flyscan_results = flyscan_event_handler.xray_centre_results
+    assert (
+        flyscan_results
+    ), "Flyscan result event not received or no crystal found and exception not raised"
+    yield from change_aperture_then_move_to_xtal(
+        flyscan_results[0], composite.smargon, composite.aperture_scatterguard
     )

mx_bluesky/hyperion/experiment_plans/pin_tip_centring_plan.py

@@ -1,7 +1,7 @@
-import dataclasses
 from collections.abc import Generator
 
 import bluesky.plan_stubs as bps
+import pydantic
 from blueapi.core import BlueskyContext
 from bluesky.utils import Msg
 from dodal.devices.backlight import Backlight
@@ -19,7 +19,7 @@ from mx_bluesky.hyperion.device_setup_plans.setup_oav import pre_centring_setup_
 from mx_bluesky.hyperion.device_setup_plans.smargon import (
     move_smargon_warn_on_out_of_range,
 )
-from mx_bluesky.hyperion.exceptions import WarningException
+from mx_bluesky.hyperion.exceptions import SampleException
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
@@ -27,7 +27,7 @@ from mx_bluesky.hyperion.utils.context import device_composite_from_context
 DEFAULT_STEP_SIZE = 0.5
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class PinTipCentringComposite:
     """All devices which are directly or indirectly required by this plan"""
 
@@ -68,7 +68,7 @@ def move_pin_into_view(
         max_steps (int, optional): The number of steps to search with. Defaults to 2.
 
     Raises:
-        WarningException: Error if the pin tip is never found
+        SampleException: Error if the pin tip is never found
 
     Returns:
         Tuple[int, int]: The location of the pin tip in pixels
@@ -105,7 +105,7 @@ def move_pin_into_view(
     tip_xy_px = yield from trigger_and_return_pin_tip(pin_tip_device)
 
     if not pin_tip_valid(tip_xy_px):
-        raise WarningException(
+        raise SampleException(
             "Pin tip centring failed - pin too long/short/bent and out of range"
         )
     else:

mx_bluesky/hyperion/experiment_plans/robot_load_and_change_energy.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import dataclasses
 from collections.abc import Generator
 from datetime import datetime
 from pathlib import Path
@@ -8,6 +7,7 @@ from typing import cast
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
+import pydantic
 from blueapi.core import BlueskyContext
 from bluesky.utils import Msg
 from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
@@ -22,18 +22,18 @@ from dodal.devices.thawer import Thawer
 from dodal.devices.undulator_dcm import UndulatorDCM
 from dodal.devices.webcam import Webcam
 from dodal.devices.xbpm_feedback import XBPMFeedback
-from dodal.plans.motor_util_plans import MoveTooLarge, home_and_reset_wrapper
+from dodal.plan_stubs.motor_utils import MoveTooLarge, home_and_reset_wrapper
 
+from mx_bluesky.common.parameters.robot_load import RobotLoadAndEnergyChange
 from mx_bluesky.hyperion.experiment_plans.set_energy_plan import (
     SetEnergyComposite,
     set_energy_plan,
 )
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.robot_load import RobotLoadAndEnergyChange
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class RobotLoadAndEnergyChangeComposite:
     # SetEnergyComposite fields
     vfm: FocusingMirrorWithStripes
@@ -128,11 +128,7 @@ def do_robot_load(
         group="robot_load",
     )
 
-    if demand_energy_ev:
-        yield from set_energy_plan(
-            demand_energy_ev / 1000,
-            cast(SetEnergyComposite, composite),
-        )
+    yield from set_energy_plan(demand_energy_ev, cast(SetEnergyComposite, composite))
 
     yield from bps.wait("robot_load")
 
@@ -218,24 +214,28 @@ def robot_load_and_change_energy_plan(
     yield from prepare_for_robot_load(
         composite.aperture_scatterguard, composite.smargon
     )
-    yield from bpp.run_wrapper(
-        robot_load_and_snapshots(
-            composite,
-            sample_location,
-            params.snapshot_directory,
-            params.thawing_time,
-            params.demand_energy_ev,
-        ),
-        md={
-            "subplan_name": CONST.PLAN.ROBOT_LOAD,
-            "metadata": {
-                "visit": params.visit,
-                "sample_id": params.sample_id,
-                "sample_puck": sample_location.puck,
-                "sample_pin": sample_location.pin,
+
+    yield from bpp.set_run_key_wrapper(
+        bpp.run_wrapper(
+            robot_load_and_snapshots(
+                composite,
+                sample_location,
+                params.snapshot_directory,
+                params.thawing_time,
+                params.demand_energy_ev,
+            ),
+            md={
+                "subplan_name": CONST.PLAN.ROBOT_LOAD,
+                "metadata": {
+                    "visit": params.visit,
+                    "sample_id": params.sample_id,
+                    "sample_puck": sample_location.puck,
+                    "sample_pin": sample_location.pin,
+                },
+                "activate_callbacks": [
+                    "RobotLoadISPyBCallback",
+                ],
             },
-            "activate_callbacks": [
-                "RobotLoadISPyBCallback",
-            ],
-        },
+        ),
+        CONST.PLAN.ROBOT_LOAD_AND_SNAPSHOTS,
     )