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

mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -1,19 +1,20 @@
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from functools import partial
 from pathlib import Path
-from time import time
 from typing import Protocol
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
 import numpy as np
-from blueapi.core import BlueskyContext, MsgGenerator
+import pydantic
+from blueapi.core import BlueskyContext
+from bluesky.callbacks import CallbackBase
+from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import (
     ApertureScatterguard,
-    ApertureValue,
 )
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.backlight import Backlight
@@ -30,7 +31,7 @@ from dodal.devices.fast_grid_scan import (
 from dodal.devices.flux import Flux
 from dodal.devices.robot import BartRobot
 from dodal.devices.s4_slit_gaps import S4SlitGaps
-from dodal.devices.smargon import Smargon, StubPosition
+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
@@ -39,17 +40,17 @@ from dodal.devices.zebra_controlled_shutter import ZebraShutter
 from dodal.devices.zocalo.zocalo_results import (
     ZOCALO_READING_PLAN_NAME,
     ZOCALO_STAGE_GROUP,
+    XrcResult,
     ZocaloResults,
-    get_processing_result,
+    get_full_processing_results,
 )
-from dodal.plans.check_topup import check_topup_and_wait_if_necessary
+from event_model import RunStart
 from ophyd_async.fastcs.panda import HDFPanda
-from scanspec.core import AxesPoints, Axis
 
-from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.common.plans.do_fgs import kickoff_and_complete_gridscan
+from mx_bluesky.common.utils.tracing import TRACER
 from mx_bluesky.hyperion.device_setup_plans.read_hardware_for_setup import (
     read_hardware_during_collection,
-    read_hardware_for_zocalo,
     read_hardware_pre_collection,
 )
 from mx_bluesky.hyperion.device_setup_plans.setup_panda import (
@@ -65,11 +66,17 @@ from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
 from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
     transmission_and_xbpm_feedback_for_collection_decorator,
 )
-from mx_bluesky.hyperion.exceptions import WarningException
+from mx_bluesky.hyperion.exceptions import CrystalNotFoundException, SampleException
+from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
+    change_aperture_then_move_to_xtal,
+)
+from mx_bluesky.hyperion.experiment_plans.common.xrc_result import XRayCentreResult
+from mx_bluesky.hyperion.external_interaction.callbacks.xray_centre.ispyb_callback import (
+    ispyb_activation_wrapper,
+)
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import ThreeDGridScan
-from mx_bluesky.hyperion.tracing import TRACER
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
@@ -77,13 +84,7 @@ class SmargonSpeedException(Exception):
     pass
 
 
-class CrystalNotFoundException(WarningException):
-    """Raised if grid detection completed normally but no crystal was found."""
-
-    pass
-
-
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class FlyScanXRayCentreComposite:
     """All devices which are directly or indirectly required by this plan"""
 
@@ -112,30 +113,33 @@ class FlyScanXRayCentreComposite:
         return self.smargon
 
 
+class XRayCentreEventHandler(CallbackBase):
+    def __init__(self):
+        super().__init__()
+        self.xray_centre_results: Sequence[XRayCentreResult] | None = None
+
+    def start(self, doc: RunStart) -> RunStart | None:
+        if "xray_centre_results" in doc:
+            self.xray_centre_results = [
+                XRayCentreResult(**result_dict)
+                for result_dict in doc["xray_centre_results"]  # type: ignore
+            ]
+        return doc
+
+
 def create_devices(context: BlueskyContext) -> FlyScanXRayCentreComposite:
     """Creates the devices required for the plan and connect to them"""
     return device_composite_from_context(context, FlyScanXRayCentreComposite)
 
 
-def flyscan_xray_centre(
-    composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+def flyscan_xray_centre_no_move(
+    composite: FlyScanXRayCentreComposite, parameters: HyperionThreeDGridScan
 ) -> MsgGenerator:
-    """Create the plan to run the grid scan based on provided parameters.
-
-    The ispyb handler should be added to the whole gridscan as we want to capture errors
-    at any point in it.
-
-    Args:
-        parameters (ThreeDGridScan): The parameters to run the scan.
-
-    Returns:
-        Generator: The plan for the gridscan
-    """
+    """Perform a flyscan and determine the centres of interest"""
     parameters.features.update_self_from_server()
     composite.eiger.set_detector_parameters(parameters.detector_params)
-    composite.zocalo.zocalo_environment = parameters.zocalo_environment
-    composite.zocalo.use_cpu_and_gpu = parameters.use_cpu_and_gpu_zocalo
+    composite.zocalo.zocalo_environment = CONST.ZOCALO_ENV
+    composite.zocalo.use_cpu_and_gpu = parameters.features.compare_cpu_and_gpu_zocalo
 
     feature_controlled = _get_feature_controlled(composite, parameters)
 
@@ -143,8 +147,6 @@ def flyscan_xray_centre(
     @bpp.run_decorator(  # attach experiment metadata to the start document
         md={
             "subplan_name": CONST.PLAN.GRIDSCAN_OUTER,
-            CONST.TRIGGER.ZOCALO: CONST.PLAN.DO_FGS,
-            "zocalo_environment": parameters.zocalo_environment,
             "hyperion_parameters": parameters.model_dump_json(),
             "activate_callbacks": [
                 "GridscanNexusFileCallback",
@@ -157,25 +159,66 @@ def flyscan_xray_centre(
         composite.attenuator,
         parameters.transmission_frac,
     )
-    def run_gridscan_and_move_and_tidy(
+    def run_gridscan_and_fetch_and_tidy(
         fgs_composite: FlyScanXRayCentreComposite,
-        params: ThreeDGridScan,
+        params: HyperionThreeDGridScan,
         feature_controlled: _FeatureControlled,
-    ):
-        yield from run_gridscan_and_move(fgs_composite, params, feature_controlled)
+    ) -> MsgGenerator:
+        yield from run_gridscan_and_fetch_results(
+            fgs_composite, params, feature_controlled
+        )
+
+    yield from run_gridscan_and_fetch_and_tidy(
+        composite, parameters, feature_controlled
+    )
+
+
+def flyscan_xray_centre(
+    composite: FlyScanXRayCentreComposite,
+    parameters: HyperionThreeDGridScan,
+) -> MsgGenerator:
+    """Create the plan to run the grid scan based on provided parameters.
+
+    The ispyb handler should be added to the whole gridscan as we want to capture errors
+    at any point in it.
+
+    Args:
+        parameters (ThreeDGridScan): The parameters to run the scan.
+
+    Returns:
+        Generator: The plan for the gridscan
+    """
+    xrc_event_handler = XRayCentreEventHandler()
+
+    @bpp.subs_decorator(xrc_event_handler)
+    def flyscan_and_fetch_results() -> MsgGenerator:
+        yield from ispyb_activation_wrapper(
+            flyscan_xray_centre_no_move(composite, parameters), parameters
+        )
 
-    return run_gridscan_and_move_and_tidy(composite, parameters, feature_controlled)
+    yield from flyscan_and_fetch_results()
+
+    xray_centre_results = xrc_event_handler.xray_centre_results
+    assert (
+        xray_centre_results
+    ), "Flyscan result event not received or no crystal found and exception not raised"
+    yield from change_aperture_then_move_to_xtal(
+        xray_centre_results[0],
+        composite.smargon,
+        composite.aperture_scatterguard,
+        parameters,
+    )
 
 
 @bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_AND_MOVE)
 @bpp.run_decorator(md={"subplan_name": CONST.PLAN.GRIDSCAN_AND_MOVE})
-def run_gridscan_and_move(
+def run_gridscan_and_fetch_results(
     fgs_composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+    parameters: HyperionThreeDGridScan,
     feature_controlled: _FeatureControlled,
 ) -> MsgGenerator:
     """A multi-run plan which runs a gridscan, gets the results from zocalo
-    and moves to the centre of mass determined by zocalo"""
+    and fires an event with the centres of mass determined by zocalo"""
 
     # We get the initial motor positions so we can return to them on zocalo failure
     initial_xyz = np.array(
@@ -202,51 +245,66 @@ def run_gridscan_and_move(
                 [fgs_composite.zocalo], name=ZOCALO_READING_PLAN_NAME
             )
             LOGGER.info("Zocalo triggered and read, interpreting results.")
-            xray_centre, bbox_size = yield from get_processing_result(
-                fgs_composite.zocalo
-            )
-            LOGGER.info(f"Got xray centre: {xray_centre}, bbox size: {bbox_size}")
-            if xray_centre is not None:
-                xray_centre = parameters.FGS_params.grid_position_to_motor_position(
-                    xray_centre
-                )
+            xrc_results = yield from get_full_processing_results(fgs_composite.zocalo)
+            LOGGER.info(f"Got xray centres, top 5: {xrc_results[:5]}")
+            if xrc_results:
+                flyscan_results = [
+                    _xrc_result_in_boxes_to_result_in_mm(xr, parameters)
+                    for xr in xrc_results
+                ]
             else:
                 LOGGER.warning("No X-ray centre received")
                 raise CrystalNotFoundException()
-            if bbox_size is not None:
-                with TRACER.start_span("change_aperture"):
-                    yield from set_aperture_for_bbox_size(
-                        fgs_composite.aperture_scatterguard, bbox_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 = xray_centre
-            yield from move_x_y_z(fgs_composite.sample_motors, x, y, z, wait=True)
-
-        if parameters.FGS_params.set_stub_offsets:
-            LOGGER.info("Recentring smargon co-ordinate system to this point.")
-            yield from bps.mv(
-                fgs_composite.sample_motors.stub_offsets, StubPosition.CURRENT_AS_CENTER
-            )
+            yield from _fire_xray_centre_result_event(flyscan_results)
+
     finally:
         # 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)
+        yield from bps.abs_set(fgs_composite.eiger.odin.fan.dev_shm_enable, 0)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 
         # Wait on everything before returning to GDA (particularly apertures), can be removed
         # when we do not return to GDA here
         yield from bps.wait()
 
 
+def _xrc_result_in_boxes_to_result_in_mm(
+    xrc_result: XrcResult, parameters: HyperionThreeDGridScan
+) -> XRayCentreResult:
+    fgs_params = parameters.FGS_params
+    xray_centre = fgs_params.grid_position_to_motor_position(
+        np.array(xrc_result["centre_of_mass"])
+    )
+    return XRayCentreResult(
+        centre_of_mass_mm=xray_centre,
+        bounding_box_mm=(
+            fgs_params.grid_position_to_motor_position(
+                np.array(xrc_result["bounding_box"][0])
+            ),
+            fgs_params.grid_position_to_motor_position(
+                np.array(xrc_result["bounding_box"][1])
+            ),
+        ),
+        max_count=xrc_result["max_count"],
+        total_count=xrc_result["total_count"],
+    )
+
+
+@bpp.set_run_key_decorator(CONST.PLAN.FLYSCAN_RESULTS)
+def _fire_xray_centre_result_event(results: Sequence[XRayCentreResult]):
+    def empty_plan():
+        return iter([])
+
+    yield from bpp.run_wrapper(
+        empty_plan(),
+        md={"xray_centre_results": [dataclasses.asdict(r) for r in results]},
+    )
+
+
 @bpp.set_run_key_decorator(CONST.PLAN.GRIDSCAN_MAIN)
 @bpp.run_decorator(md={"subplan_name": CONST.PLAN.GRIDSCAN_MAIN})
 def run_gridscan(
     fgs_composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+    parameters: HyperionThreeDGridScan,
     feature_controlled: _FeatureControlled,
     md={  # noqa
         "plan_name": CONST.PLAN.GRIDSCAN_MAIN,
@@ -266,7 +324,7 @@ def run_gridscan(
             fgs_composite.undulator,
             fgs_composite.synchrotron,
             fgs_composite.s4_slit_gaps,
-            fgs_composite.robot,
+            fgs_composite.dcm,
             fgs_composite.smargon,
         )
 
@@ -287,7 +345,7 @@ def run_gridscan(
 
     LOGGER.info("Waiting for arming to finish")
     yield from bps.wait(CONST.WAIT.GRID_READY_FOR_DC)
-    yield from bps.stage(fgs_composite.eiger)
+    yield from bps.stage(fgs_composite.eiger)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 
     yield from kickoff_and_complete_gridscan(
         feature_controlled.fgs_motors,
@@ -295,66 +353,11 @@ def run_gridscan(
         fgs_composite.synchrotron,
         [parameters.scan_points_first_grid, parameters.scan_points_second_grid],
         parameters.scan_indices,
-        do_during_run=read_during_collection,
+        plan_during_collection=read_during_collection,
     )
     yield from bps.abs_set(feature_controlled.fgs_motors.z_steps, 0, wait=False)
 
 
-def kickoff_and_complete_gridscan(
-    gridscan: FastGridScanCommon,
-    eiger: EigerDetector,
-    synchrotron: Synchrotron,
-    scan_points: list[AxesPoints[Axis]],
-    scan_start_indices: list[int],
-    do_during_run: Callable[[], MsgGenerator] | None = None,
-):
-    @TRACER.start_as_current_span(CONST.PLAN.DO_FGS)
-    @bpp.set_run_key_decorator(CONST.PLAN.DO_FGS)
-    @bpp.run_decorator(
-        md={
-            "subplan_name": CONST.PLAN.DO_FGS,
-            "scan_points": scan_points,
-            "scan_start_indices": scan_start_indices,
-        }
-    )
-    @bpp.contingency_decorator(
-        except_plan=lambda e: (yield from bps.stop(eiger)),
-        else_plan=lambda: (yield from bps.unstage(eiger)),
-    )
-    def do_fgs():
-        # Check topup gate
-        expected_images = yield from bps.rd(gridscan.expected_images)
-        exposure_sec_per_image = yield from bps.rd(eiger.cam.acquire_time)
-        LOGGER.info("waiting for topup if necessary...")
-        yield from check_topup_and_wait_if_necessary(
-            synchrotron,
-            expected_images * exposure_sec_per_image,
-            30.0,
-        )
-        yield from read_hardware_for_zocalo(eiger)
-        LOGGER.info("Wait for all moves with no assigned group")
-        yield from bps.wait()
-        LOGGER.info("kicking off FGS")
-        yield from bps.kickoff(gridscan, wait=True)
-        gridscan_start_time = time()
-        LOGGER.info("Waiting for Zocalo device queue to have been cleared...")
-        yield from bps.wait(
-            ZOCALO_STAGE_GROUP
-        )  # Make sure ZocaloResults queue is clear and ready to accept our new data
-        if do_during_run:
-            LOGGER.info(f"Running {do_during_run} during FGS")
-            yield from do_during_run()
-        LOGGER.info("completing FGS")
-        yield from bps.complete(gridscan, wait=True)
-
-        # Remove this logging statement once metrics have been added
-        LOGGER.info(
-            f"Gridscan motion program took {round(time()-gridscan_start_time,2)} to complete"
-        )
-
-    yield from do_fgs()
-
-
 def wait_for_gridscan_valid(fgs_motors: FastGridScanCommon, timeout=0.5):
     LOGGER.info("Waiting for valid fgs_params")
     SLEEP_PER_CHECK = 0.1
@@ -369,32 +372,7 @@ def wait_for_gridscan_valid(fgs_motors: FastGridScanCommon, timeout=0.5):
             LOGGER.info("Gridscan scan valid and position counter reset")
             return
         yield from bps.sleep(SLEEP_PER_CHECK)
-    raise WarningException("Scan invalid - pin too long/short/bent and out of range")
-
-
-def set_aperture_for_bbox_size(
-    aperture_device: ApertureScatterguard,
-    bbox_size: list[int] | np.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()
+    raise SampleException("Scan invalid - pin too long/short/bent and out of range")
 
 
 @dataclasses.dataclass
@@ -408,7 +386,7 @@ class _FeatureControlled:
         def __call__(
             self,
             fgs_composite: FlyScanXRayCentreComposite,
-            parameters: ThreeDGridScan,
+            parameters: HyperionThreeDGridScan,
             initial_xyz: np.ndarray,
         ) -> MsgGenerator: ...
 
@@ -420,7 +398,7 @@ class _FeatureControlled:
 
 def _get_feature_controlled(
     fgs_composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+    parameters: HyperionThreeDGridScan,
 ):
     if parameters.features.use_panda_for_gridscan:
         return _FeatureControlled(
@@ -469,7 +447,7 @@ def _panda_tidy(fgs_composite: FlyScanXRayCentreComposite):
 
 def _zebra_triggering_setup(
     fgs_composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+    parameters: HyperionThreeDGridScan,
     initial_xyz: np.ndarray,
 ):
     yield from setup_zebra_for_gridscan(
@@ -479,7 +457,7 @@ def _zebra_triggering_setup(
 
 def _panda_triggering_setup(
     fgs_composite: FlyScanXRayCentreComposite,
-    parameters: ThreeDGridScan,
+    parameters: HyperionThreeDGridScan,
     initial_xyz: np.ndarray,
 ):
     LOGGER.info("Setting up Panda for flyscan")
@@ -498,15 +476,15 @@ def _panda_triggering_setup(
     )
 
     sample_velocity_mm_per_s = (
-        parameters.panda_FGS_params.x_step_size * 1e3 / time_between_x_steps_ms
+        parameters.panda_FGS_params.x_step_size_mm * 1e3 / time_between_x_steps_ms
     )
     if sample_velocity_mm_per_s > smargon_speed_limit_mm_per_s:
         raise SmargonSpeedException(
             f"Smargon speed was calculated from x step size\
-            {parameters.panda_FGS_params.x_step_size} and\
+            {parameters.panda_FGS_params.x_step_size_mm}mm and\
             time_between_x_steps_ms {time_between_x_steps_ms} as\
-            {sample_velocity_mm_per_s}. The smargon's speed limit is\
-            {smargon_speed_limit_mm_per_s} mm/s."
+            {sample_velocity_mm_per_s}mm/s. The smargon's speed limit is\
+            {smargon_speed_limit_mm_per_s}mm/s."
         )
     else:
         LOGGER.info(
@@ -515,8 +493,8 @@ def _panda_triggering_setup(
         )
 
     yield from bps.mv(
-        fgs_composite.panda_fast_grid_scan.time_between_x_steps_ms,
-        time_between_x_steps_ms,
+        fgs_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
     )
 
     directory_provider_root = Path(parameters.storage_directory)

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
@@ -15,7 +17,7 @@ from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import PandAFastGridScan, ZebraFastGridScan
 from dodal.devices.flux import Flux
 from dodal.devices.oav.oav_detector import OAV
-from dodal.devices.oav.oav_parameters import OAV_CONFIG_JSON, OAVParameters
+from dodal.devices.oav.oav_parameters import OAVParameters
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
 from dodal.devices.robot import BartRobot
 from dodal.devices.s4_slit_gaps import S4SlitGaps
@@ -28,17 +30,23 @@ from dodal.devices.zebra_controlled_shutter import ZebraShutter
 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,
@@ -54,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"""
 
@@ -86,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)
@@ -94,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
 
@@ -109,7 +120,7 @@ def detect_grid_and_do_gridscan(
 ):
     snapshot_template = f"{parameters.detector_params.prefix}_{parameters.detector_params.run_number}_{{angle}}"
 
-    grid_params_callback = GridDetectionCallback(composite.oav.parameters)
+    grid_params_callback = GridDetectionCallback()
 
     @bpp.subs_decorator([grid_params_callback])
     def run_grid_detection_plan(
@@ -129,7 +140,8 @@ def detect_grid_and_do_gridscan(
             oav_params,
             snapshot_template,
             str(snapshot_dir),
-            grid_width_microns=parameters.grid_width_um,
+            parameters.grid_width_um,
+            parameters.box_size_um,
         )
 
     yield from run_grid_detection_plan(
@@ -148,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,
@@ -178,7 +190,7 @@ def detect_grid_and_do_gridscan(
 def grid_detect_then_xray_centre(
     composite: GridDetectThenXRayCentreComposite,
     parameters: GridScanWithEdgeDetect,
-    oav_config: str = OAV_CONFIG_JSON,
+    oav_config: str = OavConstants.OAV_CONFIG_JSON,
 ) -> MsgGenerator:
     """
     A plan which combines the collection of snapshots from the OAV and the determination
@@ -191,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,
+        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,
+    )