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

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/common/xrc_result.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+from collections import defaultdict
 from collections.abc import Callable, Sequence
 from functools import partial
 
@@ -10,6 +11,7 @@ from event_model import RunStart
 
 from mx_bluesky.common.parameters.components import (
     MultiXtalSelection,
+    TopNByMaxCountForEachSampleSelection,
     TopNByMaxCountSelection,
 )
 
@@ -39,18 +41,21 @@ class XRayCentreResult:
             containing the crystal
         max_count: The maximum spot count encountered in any one grid box in the crystal
         total_count: The total count across all boxes in the crystal.
+        sample_id: The sample id associated with the centre.
     """
 
     centre_of_mass_mm: np.ndarray
     bounding_box_mm: tuple[np.ndarray, np.ndarray]
     max_count: int
     total_count: int
+    sample_id: int | None
 
     def __eq__(self, o):
         return (
             isinstance(o, XRayCentreResult)
             and o.max_count == self.max_count
             and o.total_count == self.total_count
+            and o.sample_id == self.sample_id
             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])
@@ -64,9 +69,27 @@ def top_n_by_max_count(
     return sorted_hits[:n]
 
 
+def top_n_by_max_count_for_each_sample(
+    unfiltered: Sequence[XRayCentreResult], n: int
+) -> Sequence[XRayCentreResult]:
+    xrc_results_by_sample_id: dict[int | None, list[XRayCentreResult]] = defaultdict(
+        list[XRayCentreResult]
+    )
+    for result in unfiltered:
+        xrc_results_by_sample_id[result.sample_id].append(result)
+    return [
+        result
+        for results in xrc_results_by_sample_id.values()
+        for result in sorted(results, key=lambda x: x.max_count, reverse=True)[: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)
+    match params:
+        case TopNByMaxCountSelection():
+            return partial(top_n_by_max_count, n=params.n)
+        case TopNByMaxCountForEachSampleSelection():
+            return partial(top_n_by_max_count_for_each_sample, n=params.n)
     raise ValueError(f"Invalid selection function {params.name}")

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": {

mx_bluesky/hyperion/experiment_plans/hyperion_flyscan_xray_centre_plan.py

@@ -9,7 +9,7 @@ from dodal.devices.fast_grid_scan import (
     set_fast_grid_scan_params,
 )
 
-from mx_bluesky.common.plans.common_flyscan_xray_centre_plan import (
+from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
     construct_beamline_specific_FGS_features,
 )
 from mx_bluesky.common.utils.log import LOGGER
@@ -34,51 +34,50 @@ class SmargonSpeedException(Exception):
 
 
 def construct_hyperion_specific_features(
-    fgs_composite: HyperionFlyScanXRayCentreComposite,
-    parameters: HyperionSpecifiedThreeDGridScan,
+    xrc_composite: HyperionFlyScanXRayCentreComposite,
+    xrc_parameters: HyperionSpecifiedThreeDGridScan,
 ):
     """
     Get all the information needed to do the Hyperion-specific parts of the XRC flyscan.
     """
-
     signals_to_read_pre_flyscan = [
-        fgs_composite.undulator.current_gap,
-        fgs_composite.synchrotron.synchrotron_mode,
-        fgs_composite.s4_slit_gaps.xgap,
-        fgs_composite.s4_slit_gaps.ygap,
-        fgs_composite.smargon.x,
-        fgs_composite.smargon.y,
-        fgs_composite.smargon.z,
-        fgs_composite.dcm.energy_in_kev,
+        xrc_composite.undulator.current_gap,
+        xrc_composite.synchrotron.synchrotron_mode,
+        xrc_composite.s4_slit_gaps.xgap,
+        xrc_composite.s4_slit_gaps.ygap,
+        xrc_composite.smargon.x,
+        xrc_composite.smargon.y,
+        xrc_composite.smargon.z,
+        xrc_composite.dcm.energy_in_kev,
     ]
 
     signals_to_read_during_collection = [
-        fgs_composite.aperture_scatterguard,
-        fgs_composite.attenuator.actual_transmission,
-        fgs_composite.flux.flux_reading,
-        fgs_composite.dcm.energy_in_kev,
-        fgs_composite.eiger.bit_depth,
+        xrc_composite.aperture_scatterguard,
+        xrc_composite.attenuator.actual_transmission,
+        xrc_composite.flux.flux_reading,
+        xrc_composite.dcm.energy_in_kev,
+        xrc_composite.eiger.bit_depth,
     ]
 
-    if parameters.features.use_panda_for_gridscan:
+    if xrc_parameters.features.use_panda_for_gridscan:
         setup_trigger_plan = _panda_triggering_setup
         tidy_plan = _panda_tidy
         set_flyscan_params_plan = partial(
             set_fast_grid_scan_params,
-            fgs_composite.panda_fast_grid_scan,
-            parameters.panda_FGS_params,
+            xrc_composite.panda_fast_grid_scan,
+            xrc_parameters.panda_FGS_params,
         )
-        fgs_motors = fgs_composite.panda_fast_grid_scan
+        fgs_motors = xrc_composite.panda_fast_grid_scan
 
     else:
         setup_trigger_plan = _zebra_triggering_setup
         tidy_plan = partial(_generic_tidy, group="flyscan_zebra_tidy", wait=True)
         set_flyscan_params_plan = partial(
             set_fast_grid_scan_params,
-            fgs_composite.zebra_fast_grid_scan,
-            parameters.FGS_params,
+            xrc_composite.zebra_fast_grid_scan,
+            xrc_parameters.FGS_params,
         )
-        fgs_motors = fgs_composite.zebra_fast_grid_scan
+        fgs_motors = xrc_composite.zebra_fast_grid_scan
     return construct_beamline_specific_FGS_features(
         setup_trigger_plan,
         tidy_plan,
@@ -91,47 +90,53 @@ def construct_hyperion_specific_features(
 
 
 def _generic_tidy(
-    fgs_composite: HyperionFlyScanXRayCentreComposite, group, wait=True
+    xrc_composite: HyperionFlyScanXRayCentreComposite, group, wait=True
 ) -> MsgGenerator:
     LOGGER.info("Tidying up Zebra")
     yield from tidy_up_zebra_after_gridscan(
-        fgs_composite.zebra, fgs_composite.sample_shutter, group=group, wait=wait
+        xrc_composite.zebra, xrc_composite.sample_shutter, group=group, wait=wait
     )
     LOGGER.info("Tidying up Zocalo")
     # make sure we don't consume any other results
-    yield from bps.unstage(fgs_composite.zocalo, group=group, wait=wait)
+    yield from bps.unstage(xrc_composite.zocalo, group=group, wait=wait)
 
     # Turn off dev/shm streaming to avoid filling disk, see https://github.com/DiamondLightSource/hyperion/issues/1395
     LOGGER.info("Turning off Eiger dev/shm streaming")
-    yield from bps.abs_set(fgs_composite.eiger.odin.fan.dev_shm_enable, 0)  # type: ignore # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
+    # Fix types in ophyd-async (https://github.com/DiamondLightSource/mx-bluesky/issues/855)
+    yield from bps.abs_set(
+        xrc_composite.eiger.odin.fan.dev_shm_enable,  # type: ignore
+        0,
+        group=group,
+        wait=wait,
+    )
 
 
-def _panda_tidy(fgs_composite: HyperionFlyScanXRayCentreComposite):
+def _panda_tidy(xrc_composite: HyperionFlyScanXRayCentreComposite):
     group = "panda_flyscan_tidy"
     LOGGER.info("Disabling panda blocks")
-    yield from disarm_panda_for_gridscan(fgs_composite.panda, group)
-    yield from _generic_tidy(fgs_composite, group, False)
+    yield from disarm_panda_for_gridscan(xrc_composite.panda, group)
+    yield from _generic_tidy(xrc_composite, group, False)
     yield from bps.wait(group, timeout=10)
-    yield from bps.unstage(fgs_composite.panda)
+    yield from bps.unstage(xrc_composite.panda)
 
 
 def _zebra_triggering_setup(
-    fgs_composite: HyperionFlyScanXRayCentreComposite,
+    xrc_composite: HyperionFlyScanXRayCentreComposite,
     parameters: HyperionSpecifiedThreeDGridScan,
 ) -> MsgGenerator:
     yield from setup_zebra_for_gridscan(
-        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+        xrc_composite.zebra, xrc_composite.sample_shutter, wait=True
     )
 
 
 def _panda_triggering_setup(
-    fgs_composite: HyperionFlyScanXRayCentreComposite,
+    xrc_composite: HyperionFlyScanXRayCentreComposite,
     parameters: HyperionSpecifiedThreeDGridScan,
 ) -> MsgGenerator:
     LOGGER.info("Setting up Panda for flyscan")
 
     run_up_distance_mm = yield from bps.rd(
-        fgs_composite.panda_fast_grid_scan.run_up_distance_mm
+        xrc_composite.panda_fast_grid_scan.run_up_distance_mm
     )
 
     # Set the time between x steps pv
@@ -140,7 +145,7 @@ def _panda_triggering_setup(
     time_between_x_steps_ms = (DEADTIME_S + parameters.exposure_time_s) * 1e3
 
     smargon_speed_limit_mm_per_s = yield from bps.rd(
-        fgs_composite.smargon.x.max_velocity
+        xrc_composite.smargon.x.max_velocity
     )
 
     sample_velocity_mm_per_s = (
@@ -161,7 +166,7 @@ def _panda_triggering_setup(
         )
 
     yield from bps.mv(
-        fgs_composite.panda_fast_grid_scan.time_between_x_steps_ms,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
+        xrc_composite.panda_fast_grid_scan.time_between_x_steps_ms,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
         time_between_x_steps_ms,  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
     )
 
@@ -169,9 +174,9 @@ def _panda_triggering_setup(
     yield from set_panda_directory(directory_provider_root)
 
     yield from setup_panda_for_flyscan(
-        fgs_composite.panda,
+        xrc_composite.panda,
         parameters.panda_FGS_params,
-        fgs_composite.smargon,
+        xrc_composite.smargon,
         parameters.exposure_time_s,
         time_between_x_steps_ms,
         sample_velocity_mm_per_s,
@@ -179,5 +184,5 @@ def _panda_triggering_setup(
 
     LOGGER.info("Setting up Zebra for panda flyscan")
     yield from setup_zebra_for_panda_flyscan(
-        fgs_composite.zebra, fgs_composite.sample_shutter, wait=True
+        xrc_composite.zebra, xrc_composite.sample_shutter, wait=True
     )

mx_bluesky/hyperion/experiment_plans/hyperion_grid_detect_then_xray_centre_plan.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator
+from dodal.plans.preprocessors.verify_undulator_gap import (
+    verify_undulator_gap_before_run_decorator,
+)
+
+from mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan import (
+    grid_detect_then_xray_centre,
+)
+from mx_bluesky.common.parameters.constants import OavConstants, PlanNameConstants
+from mx_bluesky.common.preprocessors.preprocessors import (
+    transmission_and_xbpm_feedback_for_collection_decorator,
+)
+from mx_bluesky.common.utils.context import device_composite_from_context
+from mx_bluesky.hyperion.experiment_plans.hyperion_flyscan_xray_centre_plan import (
+    construct_hyperion_specific_features,
+)
+from mx_bluesky.hyperion.parameters.device_composites import (
+    HyperionGridDetectThenXRayCentreComposite,
+)
+from mx_bluesky.hyperion.parameters.gridscan import (
+    GridScanWithEdgeDetect,
+    HyperionSpecifiedThreeDGridScan,
+)
+
+
+def create_devices(
+    context: BlueskyContext,
+) -> HyperionGridDetectThenXRayCentreComposite:
+    return device_composite_from_context(
+        context, HyperionGridDetectThenXRayCentreComposite
+    )
+
+
+def hyperion_grid_detect_then_xray_centre(
+    composite: HyperionGridDetectThenXRayCentreComposite,
+    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.
+    """
+
+    @verify_undulator_gap_before_run_decorator(composite)
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite, parameters.transmission_frac, PlanNameConstants.GRIDSCAN_OUTER
+    )
+    def plan_to_perform():
+        yield from grid_detect_then_xray_centre(
+            composite=composite,
+            parameters=parameters,
+            xrc_params_type=HyperionSpecifiedThreeDGridScan,
+            construct_beamline_specific=construct_hyperion_specific_features,
+            oav_config=oav_config,
+        )
+
+    yield from plan_to_perform()

mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py

@@ -76,17 +76,29 @@ def load_centre_collect_full(
             flyscan_event_handler,
         )
 
-        locations_to_collect_um: list[np.ndarray] = []
+        locations_to_collect_um: list[np.ndarray]
+        samples_to_collect: list[int]
 
         if flyscan_event_handler.xray_centre_results:
             selection_func = flyscan_result.resolve_selection_fn(
                 parameters.selection_params
             )
             hits = selection_func(flyscan_event_handler.xray_centre_results)
-            locations_to_collect_um = [hit.centre_of_mass_mm * 1000 for hit in hits]
+            hits_to_collect = []
+            for hit in hits:
+                if hit.sample_id is None:
+                    LOGGER.warning(
+                        f"Diffracting centre {hit} not collected because no sample id was assigned."
+                    )
+                else:
+                    hits_to_collect.append(hit)
 
+            locations_to_collect_um = [
+                hit.centre_of_mass_mm * 1000 for hit in hits_to_collect
+            ]
+            samples_to_collect = [hit.sample_id for hit in hits_to_collect]
             LOGGER.info(
-                f"Selected hits {hits} using {selection_func}, args={parameters.selection_params}"
+                f"Selected hits {hits_to_collect} using {selection_func}, args={parameters.selection_params}"
             )
         else:
             # If the xray centring hasn't found a result but has not thrown an error it
@@ -98,6 +110,7 @@ def load_centre_collect_full(
             locations_to_collect_um = [
                 np.array([initial_x_mm, initial_y_mm, initial_z_mm]) * 1000
             ]
+            samples_to_collect = [parameters.sample_id]
 
         multi_rotation = parameters.multi_rotation_scan
         rotation_template = multi_rotation.rotation_scans.copy()
@@ -108,9 +121,11 @@ def load_centre_collect_full(
 
         generator = rotation_scan_generator(is_alternating)
         next(generator)
-        for location in locations_to_collect_um:
+        for location, sample_id in zip(
+            locations_to_collect_um, samples_to_collect, strict=True
+        ):
             for rot in rotation_template:
-                combination = generator.send((rot, location))
+                combination = generator.send((rot, location, sample_id))
                 multi_rotation.rotation_scans.append(combination)
         multi_rotation = RotationScan.model_validate(multi_rotation)
 
@@ -125,8 +140,10 @@ def load_centre_collect_full(
 
 def rotation_scan_generator(
     is_alternating: bool,
-) -> Generator[RotationScanPerSweep, tuple[RotationScanPerSweep, np.ndarray], None]:
-    scan_template, location = yield  # type: ignore
+) -> Generator[
+    RotationScanPerSweep, tuple[RotationScanPerSweep, np.ndarray, int], None
+]:
+    scan_template, location, sample_id = yield  # type: ignore
     next_rotation_direction = scan_template.rotation_direction
     while True:
         scan = scan_template.model_copy()
@@ -135,6 +152,7 @@ def rotation_scan_generator(
             scan.y_start_um,
             scan.z_start_um,
         ) = location
+        scan.sample_id = sample_id
         if is_alternating:
             if next_rotation_direction != scan.rotation_direction:
                 # If originally specified direction of the current scan is different
@@ -146,4 +164,4 @@ def rotation_scan_generator(
                 scan.rotation_direction = next_rotation_direction
             next_rotation_direction = next_rotation_direction.opposite
 
-        scan_template, location = yield scan
+        scan_template, location, sample_id = yield scan

mx_bluesky/hyperion/experiment_plans/pin_centre_then_xray_centre_plan.py

@@ -9,9 +9,18 @@ from dodal.devices.eiger import EigerDetector
 from dodal.devices.oav.oav_parameters import OAVParameters
 
 from mx_bluesky.common.device_setup_plans.manipulate_sample import move_phi_chi_omega
+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_grid_detect_then_xray_centre_plan import (
+    detect_grid_and_do_gridscan,
+)
+from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
+    setup_beamline_for_OAV,
+)
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     ispyb_activation_wrapper,
 )
@@ -19,32 +28,33 @@ from mx_bluesky.common.parameters.constants import OavConstants
 from mx_bluesky.common.utils.context import device_composite_from_context
 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.grid_detect_then_xray_centre_plan import (
-    GridDetectThenXRayCentreComposite,
-    detect_grid_and_do_gridscan,
-)
-from mx_bluesky.hyperion.experiment_plans.oav_snapshot_plan import (
-    setup_beamline_for_OAV,
+from mx_bluesky.hyperion.experiment_plans.hyperion_flyscan_xray_centre_plan import (
+    construct_hyperion_specific_features,
 )
 from mx_bluesky.hyperion.experiment_plans.pin_tip_centring_plan import (
     PinTipCentringComposite,
     pin_tip_centre_plan,
 )
 from mx_bluesky.hyperion.parameters.constants import CONST
+from mx_bluesky.hyperion.parameters.device_composites import (
+    HyperionGridDetectThenXRayCentreComposite,
+)
 from mx_bluesky.hyperion.parameters.gridscan import (
     GridScanWithEdgeDetect,
+    HyperionSpecifiedThreeDGridScan,
     PinTipCentreThenXrayCentre,
 )
 
 
-def create_devices(context: BlueskyContext) -> GridDetectThenXRayCentreComposite:
+def create_devices(
+    context: BlueskyContext,
+) -> HyperionGridDetectThenXRayCentreComposite:
     """
-    GridDetectThenXRayCentreComposite contains all the devices we need, reuse that.
+    HyperionGridDetectThenXRayCentreComposite contains all the devices we need, reuse that.
     """
-    return device_composite_from_context(context, GridDetectThenXRayCentreComposite)
+    return device_composite_from_context(
+        context, HyperionGridDetectThenXRayCentreComposite
+    )
 
 
 def create_parameters_for_grid_detection(
@@ -60,7 +70,7 @@ def create_parameters_for_grid_detection(
 
 
 def pin_centre_then_flyscan_plan(
-    composite: GridDetectThenXRayCentreComposite,
+    composite: HyperionGridDetectThenXRayCentreComposite,
     parameters: PinTipCentreThenXrayCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ):
@@ -92,20 +102,21 @@ def pin_centre_then_flyscan_plan(
         )
 
         grid_detect_params = create_parameters_for_grid_detection(parameters)
-
         oav_params = OAVParameters("xrayCentring", oav_config_file)
 
         yield from detect_grid_and_do_gridscan(
             composite,
             grid_detect_params,
             oav_params,
+            HyperionSpecifiedThreeDGridScan,
+            construct_hyperion_specific_features,
         )
 
     yield from ispyb_activation_wrapper(_pin_centre_then_flyscan_plan(), parameters)
 
 
 def pin_tip_centre_then_xray_centre(
-    composite: GridDetectThenXRayCentreComposite,
+    composite: HyperionGridDetectThenXRayCentreComposite,
     parameters: PinTipCentreThenXrayCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ) -> MsgGenerator: