mx-bluesky 1.4.9__py3-none-any.whl → 1.5.0__py3-none-any.whl

mx_bluesky/{hyperion/experiment_plans/grid_detect_then_xray_centre_plan.py → common/experiment_plans/common_grid_detect_then_xray_centre_plan.py}

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Protocol, TypeVar
 
-from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
 from bluesky import preprocessors as bpp
 from bluesky.preprocessors import subs_decorator
@@ -10,20 +10,27 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.backlight import BacklightPosition
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.oav.oav_parameters import OAVParameters
-from dodal.plans.preprocessors.verify_undulator_gap import (
-    verify_undulator_gap_before_run_decorator,
-)
 
 from mx_bluesky.common.device_setup_plans.manipulate_sample import (
     move_aperture_if_required,
 )
+from mx_bluesky.common.device_setup_plans.utils import (
+    start_preparing_data_collection_then_do_plan,
+)
 from mx_bluesky.common.experiment_plans.change_aperture_then_move_plan import (
     change_aperture_then_move_to_xtal,
 )
+from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
+    BeamlineSpecificFGSFeatures,
+    common_flyscan_xray_centre,
+)
 from mx_bluesky.common.experiment_plans.oav_grid_detection_plan import (
     OavGridDetectionComposite,
     grid_detection_plan,
 )
+from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
+    setup_beamline_for_OAV,
+)
 from mx_bluesky.common.external_interaction.callbacks.common.grid_detection_callback import (
     GridDetectionCallback,
     GridParamUpdate,
@@ -31,62 +38,97 @@ from mx_bluesky.common.external_interaction.callbacks.common.grid_detection_call
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     ispyb_activation_wrapper,
 )
-from mx_bluesky.common.parameters.constants import OavConstants
-from mx_bluesky.common.parameters.gridscan import GridCommon
-from mx_bluesky.common.plans.common_flyscan_xray_centre_plan import (
-    BeamlineSpecificFGSFeatures,
-    common_flyscan_xray_centre,
+from mx_bluesky.common.parameters.constants import (
+    OavConstants,
+    PlanGroupCheckpointConstants,
 )
-from mx_bluesky.common.preprocessors.preprocessors import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
+from mx_bluesky.common.parameters.device_composites import (
+    FlyScanEssentialDevices,
+    GridDetectThenXRayCentreComposite,
 )
-from mx_bluesky.common.utils.context import device_composite_from_context
+from mx_bluesky.common.parameters.gridscan import GridCommon, SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.xrc_result import XRayCentreEventHandler
-from mx_bluesky.hyperion.device_setup_plans.utils import (
-    start_preparing_data_collection_then_do_plan,
-)
-from mx_bluesky.hyperion.experiment_plans.hyperion_flyscan_xray_centre_plan import (
-    construct_hyperion_specific_features,
-)
-from mx_bluesky.hyperion.experiment_plans.oav_snapshot_plan import (
-    setup_beamline_for_OAV,
-)
-from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.device_composites import (
-    GridDetectThenXRayCentreComposite,
-    HyperionFlyScanXRayCentreComposite,
+
+TFlyScanEssentialDevices = TypeVar(
+    "TFlyScanEssentialDevices", bound=FlyScanEssentialDevices, contravariant=True
 )
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    HyperionSpecifiedThreeDGridScan,
+TSpecifiedThreeDGridScan = TypeVar(
+    "TSpecifiedThreeDGridScan", bound=SpecifiedThreeDGridScan, contravariant=True
 )
 
 
-def create_devices(context: BlueskyContext) -> GridDetectThenXRayCentreComposite:
-    return device_composite_from_context(context, GridDetectThenXRayCentreComposite)
+def grid_detect_then_xray_centre(
+    composite: GridDetectThenXRayCentreComposite,
+    parameters: GridCommon,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+    construct_beamline_specific: ConstructBeamlineSpecificFeatures,
+    oav_config: str = OavConstants.OAV_CONFIG_JSON,
+) -> MsgGenerator:
+    """
+    A plan which combines the collection of snapshots from the OAV and the determination
+    of the grid dimensions to use for the following grid scan.
+    """
 
+    eiger: EigerDetector = composite.eiger
 
-def create_parameters_for_flyscan_xray_centre(
-    grid_scan_with_edge_params: GridCommon,
-    grid_parameters: GridParamUpdate,
-) -> HyperionSpecifiedThreeDGridScan:
-    params_json = grid_scan_with_edge_params.model_dump()
-    params_json.update(grid_parameters)
-    flyscan_xray_centre_parameters = HyperionSpecifiedThreeDGridScan(**params_json)
-    LOGGER.info(f"Parameters for FGS: {flyscan_xray_centre_parameters}")
-    return flyscan_xray_centre_parameters
+    eiger.set_detector_parameters(parameters.detector_params)
+
+    oav_params = OAVParameters("xrayCentring", oav_config)
+
+    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,
+                xrc_params_type,
+                construct_beamline_specific,
+            ),
+            parameters,
+        )
+
+    yield from start_preparing_data_collection_then_do_plan(
+        composite.beamstop,
+        eiger,
+        composite.detector_motion,
+        parameters.detector_params.detector_distance,
+        plan_to_perform(),
+        group=PlanGroupCheckpointConstants.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,
+    )
 
 
 def detect_grid_and_do_gridscan(
     composite: GridDetectThenXRayCentreComposite,
     parameters: GridCommon,
     oav_params: OAVParameters,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+    construct_beamline_specific: ConstructBeamlineSpecificFeatures,
 ):
     snapshot_template = f"{parameters.detector_params.prefix}_{parameters.detector_params.run_number}_{{angle}}"
 
     grid_params_callback = GridDetectionCallback()
 
+    yield from setup_beamline_for_OAV(
+        composite.smargon,
+        composite.backlight,
+        composite.aperture_scatterguard,
+        wait=True,
+    )
+
     @bpp.subs_decorator([grid_params_callback])
     def run_grid_detection_plan(
         oav_params,
@@ -109,19 +151,12 @@ def detect_grid_and_do_gridscan(
             parameters.box_size_um,
         )
 
-    yield from setup_beamline_for_OAV(
-        composite.smargon,
-        composite.backlight,
-        composite.aperture_scatterguard,
-        wait=True,
-    )
-
     if parameters.selected_aperture:
         # Start moving the aperture/scatterguard into position without moving it in
         yield from bps.prepare(
             composite.aperture_scatterguard,
             parameters.selected_aperture,
-            group=CONST.WAIT.PREPARE_APERTURE,
+            group=PlanGroupCheckpointConstants.PREPARE_APERTURE,
         )
 
     yield from run_grid_detection_plan(
@@ -131,105 +166,42 @@ def detect_grid_and_do_gridscan(
     )
 
     yield from bps.abs_set(
-        composite.backlight, BacklightPosition.OUT, group=CONST.WAIT.GRID_READY_FOR_DC
+        composite.backlight,
+        BacklightPosition.OUT,
+        group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
     )
 
-    yield from bps.wait(CONST.WAIT.PREPARE_APERTURE)
+    yield from bps.wait(PlanGroupCheckpointConstants.PREPARE_APERTURE)
     yield from move_aperture_if_required(
         composite.aperture_scatterguard,
         parameters.selected_aperture,
-        group=CONST.WAIT.GRID_READY_FOR_DC,
+        group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,
     )
-
-    xrc_composite = HyperionFlyScanXRayCentreComposite(
-        aperture_scatterguard=composite.aperture_scatterguard,
-        attenuator=composite.attenuator,
-        backlight=composite.backlight,
-        eiger=composite.eiger,
-        panda_fast_grid_scan=composite.panda_fast_grid_scan,
-        flux=composite.flux,
-        s4_slit_gaps=composite.s4_slit_gaps,
-        smargon=composite.smargon,
-        undulator=composite.undulator,
-        synchrotron=composite.synchrotron,
-        xbpm_feedback=composite.xbpm_feedback,
-        zebra=composite.zebra,
-        zocalo=composite.zocalo,
-        panda=composite.panda,
-        zebra_fast_grid_scan=composite.zebra_fast_grid_scan,
-        dcm=composite.dcm,
-        robot=composite.robot,
-        sample_shutter=composite.sample_shutter,
+    xrc_params = create_parameters_for_flyscan_xray_centre(
+        parameters, grid_params_callback.get_grid_parameters(), xrc_params_type
     )
+    beamline_specific = construct_beamline_specific(composite, xrc_params)
 
-    params = create_parameters_for_flyscan_xray_centre(
-        parameters, grid_params_callback.get_grid_parameters()
-    )
+    yield from common_flyscan_xray_centre(composite, xrc_params, beamline_specific)
 
-    beamline_specific = construct_hyperion_specific_features(xrc_composite, params)
 
-    yield from _gridscan_with_undulator_checks(xrc_composite, params, beamline_specific)
-
-
-def _gridscan_with_undulator_checks(
-    composite: HyperionFlyScanXRayCentreComposite,
-    params: HyperionSpecifiedThreeDGridScan,
-    beamline_specific: BeamlineSpecificFGSFeatures,
+class ConstructBeamlineSpecificFeatures(
+    Protocol[TFlyScanEssentialDevices, TSpecifiedThreeDGridScan]
 ):
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite, params.transmission_frac
-    )
-    @verify_undulator_gap_before_run_decorator(composite)
-    def _inner():
-        yield from common_flyscan_xray_centre(composite, params, beamline_specific)
-
-    yield from _inner()
+    def __call__(
+        self,
+        xrc_composite: TFlyScanEssentialDevices,
+        xrc_parameters: TSpecifiedThreeDGridScan,
+    ) -> BeamlineSpecificFGSFeatures: ...
 
 
-def grid_detect_then_xray_centre(
-    composite: GridDetectThenXRayCentreComposite,
-    parameters: GridScanWithEdgeDetect,
-    oav_config: str = OavConstants.OAV_CONFIG_JSON,
-) -> MsgGenerator:
-    """
-    A plan which combines the collection of snapshots from the OAV and the determination
-    of the grid dimensions to use for the following grid scan.
-    """
-
-    eiger: EigerDetector = composite.eiger
-
-    eiger.set_detector_parameters(parameters.detector_params)
-
-    oav_params = OAVParameters("xrayCentring", oav_config)
-
-    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,
-        )
-
-    yield from start_preparing_data_collection_then_do_plan(
-        composite.beamstop,
-        eiger,
-        composite.detector_motion,
-        parameters.detector_params.detector_distance,
-        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,
-    )
+def create_parameters_for_flyscan_xray_centre(
+    parameters: GridCommon,
+    grid_parameters: GridParamUpdate,
+    xrc_params_type: type[SpecifiedThreeDGridScan],
+) -> SpecifiedThreeDGridScan:
+    params_json = parameters.model_dump()
+    params_json.update(grid_parameters)
+    flyscan_xray_centre_parameters = xrc_params_type(**params_json)
+    LOGGER.info(f"Parameters for FGS: {flyscan_xray_centre_parameters}")
+    return flyscan_xray_centre_parameters

mx_bluesky/common/{plans → experiment_plans}/inner_plans/do_fgs.py

@@ -14,10 +14,10 @@ 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.experiment_plans.read_hardware import read_hardware_for_zocalo
 from mx_bluesky.common.parameters.constants import (
     PlanNameConstants,
 )
-from mx_bluesky.common.plans.read_hardware import read_hardware_for_zocalo
 from mx_bluesky.common.utils.tracing import TRACER
 
 

mx_bluesky/common/experiment_plans/oav_grid_detection_plan.py

@@ -5,9 +5,7 @@ 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
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.oav.pin_image_recognition.utils import NONE_VALUE
@@ -17,7 +15,11 @@ from dodal.devices.smargon import Smargon
 from mx_bluesky.common.device_setup_plans.setup_oav import (
     pre_centring_setup_oav,
 )
-from mx_bluesky.common.parameters.constants import DocDescriptorNames, HardwareConstants
+from mx_bluesky.common.parameters.constants import (
+    DocDescriptorNames,
+    HardwareConstants,
+)
+from mx_bluesky.common.parameters.device_composites import OavGridDetectionComposite
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.exceptions import catch_exception_and_warn
 from mx_bluesky.common.utils.log import LOGGER
@@ -26,16 +28,6 @@ if TYPE_CHECKING:
     from dodal.devices.oav.oav_parameters import OAVParameters
 
 
-@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
-class OavGridDetectionComposite:
-    """All devices which are directly or indirectly required by this plan"""
-
-    backlight: Backlight
-    oav: OAV
-    smargon: Smargon
-    pin_tip_detection: PinTipDetection
-
-
 def create_devices(context: BlueskyContext) -> OavGridDetectionComposite:
     return device_composite_from_context(context, OavGridDetectionComposite)
 

mx_bluesky/{hyperion → common}/experiment_plans/oav_snapshot_plan.py

@@ -11,7 +11,10 @@ from dodal.devices.smargon import Smargon
 
 from mx_bluesky.common.device_setup_plans.setup_oav import setup_general_oav_params
 from mx_bluesky.common.parameters.components import WithSnapshot
-from mx_bluesky.hyperion.parameters.constants import CONST, DocDescriptorNames
+from mx_bluesky.common.parameters.constants import (
+    DocDescriptorNames,
+    PlanGroupCheckpointConstants,
+)
 
 OAV_SNAPSHOT_SETUP_SHOT = "oav_snapshot_setup_shot"
 OAV_SNAPSHOT_GROUP = "oav_snapshot_group"
@@ -28,7 +31,7 @@ def setup_beamline_for_OAV(
     smargon: Smargon,
     backlight: Backlight,
     aperture_scatterguard: ApertureScatterguard,
-    group=CONST.WAIT.READY_FOR_OAV,
+    group=PlanGroupCheckpointConstants.READY_FOR_OAV,
     wait=False,
 ):
     max_vel = yield from bps.rd(smargon.omega.max_velocity)

mx_bluesky/common/parameters/components.py

@@ -136,7 +136,7 @@ class WithSnapshot(BaseModel):
 
     @model_validator(mode="after")
     def _validate_omegas_with_grid_snapshots(self) -> Self:
-        assert not self.use_grid_snapshots or self.snapshot_omegas_deg is None, (
+        assert not self.use_grid_snapshots or not self.snapshot_omegas_deg, (
             "snapshot_omegas may not be specified with use_grid_snapshots"
         )
         return self

mx_bluesky/common/parameters/device_composites.py

@@ -0,0 +1,65 @@
+import pydantic
+from dodal.devices.aperturescatterguard import (
+    ApertureScatterguard,
+)
+from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
+from dodal.devices.backlight import Backlight
+from dodal.devices.common_dcm import BaseDCM
+from dodal.devices.detector.detector_motion import DetectorMotion
+from dodal.devices.eiger import EigerDetector
+from dodal.devices.fast_grid_scan import (
+    ZebraFastGridScan,
+)
+from dodal.devices.flux import Flux
+from dodal.devices.i03 import Beamstop
+from dodal.devices.oav.oav_detector import OAV
+from dodal.devices.oav.pin_image_recognition import PinTipDetection
+from dodal.devices.robot import BartRobot
+from dodal.devices.s4_slit_gaps import S4SlitGaps
+from dodal.devices.smargon import Smargon
+from dodal.devices.synchrotron import Synchrotron
+from dodal.devices.undulator import Undulator
+from dodal.devices.xbpm_feedback import XBPMFeedback
+from dodal.devices.zebra.zebra import Zebra
+from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
+from dodal.devices.zocalo import ZocaloResults
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class FlyScanEssentialDevices:
+    eiger: EigerDetector
+    synchrotron: Synchrotron
+    zocalo: ZocaloResults
+    smargon: Smargon
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class OavGridDetectionComposite:
+    """All devices which are directly or indirectly required by this plan"""
+
+    backlight: Backlight
+    oav: OAV
+    smargon: Smargon
+    pin_tip_detection: PinTipDetection
+
+
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class GridDetectThenXRayCentreComposite(FlyScanEssentialDevices):
+    """All devices which are directly or indirectly required by this plan"""
+
+    aperture_scatterguard: ApertureScatterguard
+    attenuator: BinaryFilterAttenuator
+    backlight: Backlight
+    beamstop: Beamstop
+    dcm: BaseDCM
+    detector_motion: DetectorMotion
+    zebra_fast_grid_scan: ZebraFastGridScan
+    flux: Flux
+    oav: OAV
+    pin_tip_detection: PinTipDetection
+    s4_slit_gaps: S4SlitGaps
+    undulator: Undulator
+    xbpm_feedback: XBPMFeedback
+    zebra: Zebra
+    robot: BartRobot
+    sample_shutter: ZebraShutter

mx_bluesky/hyperion/device_setup_plans/utils.py

@@ -1,20 +1,8 @@
-from collections.abc import Generator
-
 from bluesky import plan_stubs as bps
-from bluesky import preprocessors as bpp
-from bluesky.utils import Msg
 from dodal.devices.detector import (
     DetectorParams,
 )
-from dodal.devices.detector.detector_motion import DetectorMotion, ShutterState
-from dodal.devices.eiger import EigerDetector
 from dodal.devices.i03.dcm import DCM
-from dodal.devices.mx_phase1.beamstop import Beamstop, BeamstopPositions
-
-from mx_bluesky.common.device_setup_plans.position_detector import (
-    set_detector_z_position,
-    set_shutter,
-)
 
 
 def fill_in_energy_if_not_supplied(dcm: DCM, detector_params: DetectorParams):
@@ -22,39 +10,3 @@ def fill_in_energy_if_not_supplied(dcm: DCM, detector_params: DetectorParams):
         actual_energy_ev = 1000 * (yield from bps.rd(dcm.energy_in_kev))
         detector_params.expected_energy_ev = actual_energy_ev
     return detector_params
-
-
-def start_preparing_data_collection_then_do_plan(
-    beamstop: Beamstop,
-    eiger: EigerDetector,
-    detector_motion: DetectorMotion,
-    detector_distance_mm: float | None,
-    plan_to_run: Generator[Msg, None, None],
-    group="ready_for_data_collection",
-) -> Generator[Msg, None, None]:
-    """Starts preparing for the next data collection and then runs the
-    given plan.
-
-     Preparation consists of:
-     * Arming the Eiger
-     * Moving the detector to the specified position
-     * Opening the detect shutter
-     If the plan fails it will disarm the eiger.
-    """
-
-    def wrapped_plan():
-        yield from bps.abs_set(eiger.do_arm, 1, group=group)  # type: ignore # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
-        yield from bps.abs_set(
-            beamstop.selected_pos, BeamstopPositions.DATA_COLLECTION, group=group
-        )
-        if detector_distance_mm:
-            yield from set_detector_z_position(
-                detector_motion, detector_distance_mm, group
-            )
-        yield from set_shutter(detector_motion, ShutterState.OPEN, group)
-        yield from plan_to_run
-
-    yield from bpp.contingency_wrapper(
-        wrapped_plan(),
-        except_plan=lambda e: (yield from bps.stop(eiger)),  # type: ignore # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
-    )

mx_bluesky/hyperion/experiment_plans/__init__.py

@@ -3,8 +3,8 @@
 The __all__ list in here are the plans that are externally available from outside Hyperion.
 """
 
-from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
-    grid_detect_then_xray_centre,
+from mx_bluesky.hyperion.experiment_plans.hyperion_grid_detect_then_xray_centre_plan import (
+    hyperion_grid_detect_then_xray_centre,
 )
 from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
     load_centre_collect_full,
@@ -17,7 +17,7 @@ from mx_bluesky.hyperion.experiment_plans.rotation_scan_plan import (
 )
 
 __all__ = [
-    "grid_detect_then_xray_centre",
+    "hyperion_grid_detect_then_xray_centre",
     "pin_tip_centre_then_xray_centre",
     "rotation_scan",
     "load_centre_collect_full",

mx_bluesky/hyperion/experiment_plans/experiment_registry.py

@@ -5,7 +5,7 @@ from typing import TypedDict
 
 import mx_bluesky.hyperion.experiment_plans.rotation_scan_plan as rotation_scan_plan
 from mx_bluesky.hyperion.experiment_plans import (
-    grid_detect_then_xray_centre_plan,
+    hyperion_grid_detect_then_xray_centre_plan,
     load_centre_collect_full_plan,
     pin_centre_then_xray_centre_plan,
 )
@@ -38,8 +38,8 @@ class ExperimentRegistryEntry(TypedDict):
 
 
 PLAN_REGISTRY: dict[str, ExperimentRegistryEntry] = {
-    "grid_detect_then_xray_centre": {
-        "setup": grid_detect_then_xray_centre_plan.create_devices,
+    "hyperion_grid_detect_then_xray_centre": {
+        "setup": hyperion_grid_detect_then_xray_centre_plan.create_devices,
         "param_type": GridScanWithEdgeDetect,
     },
     "pin_tip_centre_then_xray_centre": {