Stage zocalo and centre outside of common gridscan (#1557)

* Stage and unstage zocalo results device outside of common gridscan

* Separate XRC from grid detect

* Fix tests and add new tests

Author: olliesilvester

docs/user/hyperion/advanced/concurrent_operations.puml

@@ -63,7 +63,7 @@ activate RE #ffeeee
     RE -> RE++ : start_preparing_data_collection_then_do_plan
         RE -> eiger ++ : arm the eiger
         RE -> GridscanISPyB ++ : send parameters to grdscan callback
-            RE -> RE ++ #ffdddd : pin_centre_then_flyscan_plan
+            RE -> RE ++ #ffdddd : pin_centre_then_xray_centre_plan
                 RE -> RE: setup_beamline_for_oav
                 RE -> RE: pin_tip_centre_plan
                 ...do pin tip centring...
@@ -121,7 +121,7 @@ snapshots later
                         RE -> zebra : tidy zebra
                     deactivate RE /'common_flyscan_xray_centre'/
                 deactivate RE /' detect_grid_and_do_gridscan'/
-            return /' pin_centre_then_flyscan_plan'/
+            return /' pin_centre_then_xray_centre_plan'/
         deactivate GridscanISPyB
     deactivate RE /' start_preparing_data_collection_then_do_plan '/
 deactivate RE /'robot_load_then_xray_centre'/

src/mx_bluesky/beamlines/i04/experiment_plans/i04_grid_detect_then_xray_centre_plan.py

@@ -43,13 +43,19 @@
     setup_zebra_for_gridscan,
     tidy_up_zebra_after_gridscan,
 )
+from mx_bluesky.common.experiment_plans.change_aperture_then_move_plan import (
+    get_results_then_change_aperture_and_move_to_xtal,
+)
 from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
     BeamlineSpecificFGSFeatures,
     construct_beamline_specific_fast_gridscan_features,
 )
 from mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan import (
     grid_detect_then_xray_centre,
 )
+from mx_bluesky.common.experiment_plans.inner_plans.xrc_results_utils import (
+    zocalo_stage_decorator,
+)
 from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
     setup_beamline_for_oav,
 )
@@ -85,6 +91,7 @@
 from mx_bluesky.common.utils.utils import (
     fix_transmission_and_exposure_time_for_current_wavelength,
 )
+from mx_bluesky.common.utils.xrc_result import XRayCentreEventHandler
 
 DEFAULT_XRC_BEAMSIZE_MICRONS = 20
 
@@ -154,10 +161,10 @@ def i04_default_grid_detect_and_xray_centre(
     composite = GridDetectThenXRayCentreComposite(
         eiger,
         synchrotron,
-        zocalo,
         smargon,
         aperture_scatterguard,
         attenuator,
+        zocalo,
         backlight,
         beamstop,
         beamsize,
@@ -194,12 +201,14 @@ def tidy_beamline():
                 composite.detector_motion,
             )
 
+    @zocalo_stage_decorator(composite.zocalo)
     @bpp.finalize_decorator(tidy_beamline)
     def _inner_grid_detect_then_xrc():
         # These callbacks let us talk to ISPyB and Nexgen. They aren't included in the common plan because
         # Hyperion handles its callbacks differently to BlueAPI-managed plans, see
         # https://github.com/DiamondLightSource/mx-bluesky/issues/1117
-        callbacks = create_gridscan_callbacks()
+        flyscan_event_handler = XRayCentreEventHandler()
+        callbacks = *create_gridscan_callbacks(), flyscan_event_handler
 
         @bpp.subs_decorator(callbacks)
         @verify_undulator_gap_before_run_decorator(composite)
@@ -219,6 +228,14 @@ def grid_detect_then_xray_centre_with_callbacks():
 
         try:
             yield from grid_detect_then_xray_centre_with_callbacks()
+            assert isinstance(
+                grid_common_params.specified_grid_params, SpecifiedThreeDGridScan
+            ), "Specified grid params couldn't be found after grid detection"
+            yield from get_results_then_change_aperture_and_move_to_xtal(
+                composite,
+                grid_common_params.specified_grid_params,
+                flyscan_event_handler,
+            )
         except CrystalNotFoundError:
             yield from bps.mv(
                 smargon.x, initial_x, smargon.y, initial_y, smargon.z, initial_z
@@ -316,7 +333,6 @@ def construct_i04_specific_features(
         fgs_motors,
         signals_to_read_pre_flyscan,
         signals_to_read_during_collection,  # type: ignore # until https://github.com/DiamondLightSource/mx-bluesky/issues/1076
-        get_xrc_results_from_zocalo=True,
     )
 
 

src/mx_bluesky/common/experiment_plans/change_aperture_then_move_plan.py

@@ -4,10 +4,32 @@
 from dodal.devices.smargon import Smargon, StubPosition
 
 from mx_bluesky.common.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.common.experiment_plans.inner_plans.xrc_results_utils import (
+    fetch_xrc_results_from_zocalo,
+)
 from mx_bluesky.common.parameters.constants import PlanGroupCheckpointConstants
+from mx_bluesky.common.parameters.device_composites import (
+    GridDetectThenXRayCentreComposite,
+)
+from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.utils.tracing import TRACER
-from mx_bluesky.common.xrc_result import XRayCentreResult
+from mx_bluesky.common.utils.xrc_result import XRayCentreEventHandler, XRayCentreResult
+
+
+def get_results_then_change_aperture_and_move_to_xtal(
+    composite: GridDetectThenXRayCentreComposite,
+    parameters: SpecifiedThreeDGridScan,
+    flyscan_event_handler: XRayCentreEventHandler,
+):
+    yield from fetch_xrc_results_from_zocalo(composite.zocalo, parameters)
+    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
+    )
 
 
 def change_aperture_then_move_to_xtal(
@@ -23,7 +45,7 @@ def change_aperture_then_move_to_xtal(
     bounding_box_size = numpy.abs(
         best_hit.bounding_box_mm[1] - best_hit.bounding_box_mm[0]
     )
-    yield from set_aperture_for_bbox_mm(
+    yield from _set_aperture_for_bbox_mm(
         aperture_scatterguard,
         bounding_box_size,
     )
@@ -41,7 +63,7 @@ def change_aperture_then_move_to_xtal(
         yield from bps.mv(smargon.stub_offsets, StubPosition.CURRENT_AS_CENTER)
 
 
-def set_aperture_for_bbox_mm(
+def _set_aperture_for_bbox_mm(
     aperture_device: ApertureScatterguard,
     bbox_size_mm: list[float] | numpy.ndarray,
     group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC,

src/mx_bluesky/common/experiment_plans/common_flyscan_xray_centre_plan.py

@@ -6,43 +6,33 @@
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
-import numpy as np
 from bluesky.protocols import Readable
 from bluesky.utils import FailedStatus, MsgGenerator
-from dodal.common.beamlines.commissioning_mode import read_commissioning_mode
+from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import (
     FastGridScanCommon,
     FastGridScanThreeD,
     GridScanInvalidError,
 )
-from dodal.devices.zocalo import ZocaloResults
-from dodal.devices.zocalo.zocalo_results import (
-    XrcResult,
-    get_full_processing_results,
-)
 
 from mx_bluesky.common.experiment_plans.inner_plans.do_fgs import (
-    ZOCALO_STAGE_GROUP,
     kickoff_and_complete_gridscan,
 )
 from mx_bluesky.common.experiment_plans.inner_plans.read_hardware import (
     read_hardware_plan,
 )
 from mx_bluesky.common.parameters.constants import (
     DocDescriptorNames,
-    GridscanParamConstants,
     PlanGroupCheckpointConstants,
     PlanNameConstants,
 )
 from mx_bluesky.common.parameters.device_composites import FlyScanEssentialDevices
 from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
 from mx_bluesky.common.utils.exceptions import (
-    CrystalNotFoundError,
     SampleError,
 )
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.utils.tracing import TRACER
-from mx_bluesky.common.xrc_result import XRayCentreResult
 
 
 @dataclasses.dataclass
@@ -55,26 +45,19 @@ class BeamlineSpecificFGSFeatures:
         ..., MsgGenerator
     ]  # Eventually replace with https://github.com/DiamondLightSource/mx-bluesky/issues/819
     read_during_collection_plan: Callable[..., MsgGenerator]
-    get_xrc_results_from_zocalo: bool
-
 
-def generic_tidy(xrc_composite: FlyScanEssentialDevices, wait=True) -> MsgGenerator:
-    """Tidy Zocalo and turn off Eiger dev/shm. Ran after the beamline-specific tidy plan"""
 
-    LOGGER.info("Tidying up Zocalo")
-    group = "generic_tidy"
-    # make sure we don't consume any other results
-    yield from bps.unstage(xrc_composite.zocalo, group=group)
+def tidy_eiger(eiger: EigerDetector) -> MsgGenerator:
+    """Turn off Eiger dev/shm. Ran after the beamline-specific tidy plan"""
 
     # 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")
     # 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 # until https://github.com/DiamondLightSource/mx-bluesky/issues/1076
+        eiger.odin.fan.dev_shm_enable,  # type: ignore # until https://github.com/DiamondLightSource/mx-bluesky/issues/1076
         0,
-        group=group,
+        wait=True,
     )
-    yield from bps.wait(group)
 
 
 def construct_beamline_specific_fast_gridscan_features(
@@ -84,7 +67,6 @@ def construct_beamline_specific_fast_gridscan_features(
     fgs_motors: FastGridScanCommon,
     signals_to_read_pre_flyscan: Sequence[Readable],
     signals_to_read_during_collection: Sequence[Readable],
-    get_xrc_results_from_zocalo: bool = False,
 ) -> BeamlineSpecificFGSFeatures:
     """Construct the class needed to do beamline-specific parts of the XRC FGS
 
@@ -104,9 +86,6 @@ def construct_beamline_specific_fast_gridscan_features(
 
         signals_to_read_during_collection (Callable): Signals which will be read and saved as a bluesky event
         document whilst the gridscan motion is in progress
-
-        get_xrc_results_from_zocalo (bool): If true, fetch grid scan results from zocalo after completion, as well as
-        update the ispyb comment field with information about the results. See _fetch_xrc_results_from_zocalo
     """
     read_pre_flyscan_plan = partial(
         read_hardware_plan,
@@ -127,7 +106,6 @@ def construct_beamline_specific_fast_gridscan_features(
         fgs_motors,
         read_pre_flyscan_plan,
         read_during_collection_plan,
-        get_xrc_results_from_zocalo,
     )
 
 
@@ -159,7 +137,7 @@ def common_flyscan_xray_centre(
 
     def _overall_tidy():
         yield from beamline_specific.tidy_plan()
-        yield from generic_tidy(composite)
+        yield from tidy_eiger(composite.eiger)
 
     def _decorated_flyscan():
         @bpp.set_run_key_decorator(PlanNameConstants.GRIDSCAN_OUTER)
@@ -181,81 +159,16 @@ def run_gridscan_and_tidy(
             yield from beamline_specific.setup_trigger_plan(fgs_composite, parameters)
 
             LOGGER.info("Starting grid scan")
-            yield from bps.stage(
-                fgs_composite.zocalo, group=ZOCALO_STAGE_GROUP
-            )  # connect to zocalo and make sure the queue is clear
             yield from run_gridscan(fgs_composite, params, beamline_specific)
 
             LOGGER.info("Grid scan finished")
 
-            if beamline_specific.get_xrc_results_from_zocalo:
-                yield from _fetch_xrc_results_from_zocalo(composite.zocalo, parameters)
-
         yield from run_gridscan_and_tidy(composite, parameters, beamline_specific)
 
     composite.eiger.set_detector_parameters(parameters.detector_params)
     yield from _decorated_flyscan()
 
 
-def _fetch_xrc_results_from_zocalo(
-    zocalo_results: ZocaloResults,
-    parameters: SpecifiedThreeDGridScan,
-) -> MsgGenerator:
-    """
-    Get XRC results from the ZocaloResults device which was staged during a grid scan,
-    and store them in XRayCentreEventHandler.xray_centre_results by firing an event.
-
-    The RunEngine must be subscribed to XRayCentreEventHandler for this plan to work.
-    """
-
-    LOGGER.info("Getting X-ray center Zocalo results...")
-
-    yield from bps.trigger(zocalo_results, wait=True)
-    LOGGER.info("Zocalo triggered and read, interpreting results.")
-    xrc_results = yield from get_full_processing_results(zocalo_results)
-    LOGGER.info(f"Got xray centres, top 5: {xrc_results[:5]}")
-    filtered_results = [
-        result
-        for result in xrc_results
-        if result["total_count"]
-        >= GridscanParamConstants.ZOCALO_MIN_TOTAL_COUNT_THRESHOLD
-    ]
-    discarded_count = len(xrc_results) - len(filtered_results)
-    if discarded_count > 0:
-        LOGGER.info(f"Removed {discarded_count} results because below threshold")
-    if filtered_results:
-        flyscan_results = [
-            _xrc_result_in_boxes_to_result_in_mm(xr, parameters)
-            for xr in filtered_results
-        ]
-    else:
-        commissioning_mode = yield from read_commissioning_mode()
-        if commissioning_mode:
-            LOGGER.info("Commissioning mode enabled, returning dummy result")
-            flyscan_results = [_generate_dummy_xrc_result(parameters)]
-        else:
-            LOGGER.warning("No X-ray centre received")
-            raise CrystalNotFoundError()
-    yield from _fire_xray_centre_result_event(flyscan_results)
-
-
-def _generate_dummy_xrc_result(params: SpecifiedThreeDGridScan) -> XRayCentreResult:
-    com = [params.x_steps / 2, params.y_steps / 2, params.z_steps / 2]
-    max_voxel = [round(p) for p in com]
-    return _xrc_result_in_boxes_to_result_in_mm(
-        XrcResult(
-            centre_of_mass=com,
-            max_voxel=max_voxel,
-            bounding_box=[max_voxel, [p + 1 for p in max_voxel]],
-            n_voxels=1,
-            max_count=10000,
-            total_count=100000,
-            sample_id=params.sample_id,
-        ),
-        params,
-    )
-
-
 def run_gridscan(
     fgs_composite: FlyScanEssentialDevices,
     parameters: SpecifiedThreeDGridScan,
@@ -296,49 +209,3 @@ def run_gridscan(
     # in a GDA-happy state.
     if isinstance(beamline_specific.fgs_motors, FastGridScanThreeD):
         yield from bps.abs_set(beamline_specific.fgs_motors.z_steps, 0, wait=False)
-
-
-def _xrc_result_in_boxes_to_result_in_mm(
-    xrc_result: XrcResult, parameters: SpecifiedThreeDGridScan
-) -> XRayCentreResult:
-    fgs_params = parameters.fast_gridscan_params
-    xray_centre = fgs_params.grid_position_to_motor_position(
-        np.array(xrc_result["centre_of_mass"])
-    )
-    # A correction is applied to the bounding box to map discrete grid coordinates to
-    # the corners of the box in motor-space; we do not apply this correction
-    # to the xray-centre as it is already in continuous space and the conversion has
-    # been performed already
-    # In other words, xrc_result["bounding_box"] contains the position of the box centre,
-    # so we subtract half a box to get the corner of the box
-    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]) - 0.5
-            ),
-            fgs_params.grid_position_to_motor_position(
-                np.array(xrc_result["bounding_box"][1]) - 0.5
-            ),
-        ),
-        max_count=xrc_result["max_count"],
-        total_count=xrc_result["total_count"],
-        sample_id=xrc_result["sample_id"],
-    )
-
-
-def _fire_xray_centre_result_event(results: Sequence[XRayCentreResult]):
-    def empty_plan():
-        return iter([])
-
-    yield from bpp.set_run_key_wrapper(
-        bpp.run_wrapper(
-            empty_plan(),
-            md={
-                PlanNameConstants.FLYSCAN_RESULTS: [
-                    dataclasses.asdict(r) for r in results
-                ]
-            },
-        ),
-        PlanNameConstants.FLYSCAN_RESULTS,
-    )