Collect where we are if no centering needed

src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py

@@ -119,10 +119,10 @@ def __init__(self):
         self.xray_centre_results: Sequence[XRayCentreResult] | None = None
 
     def start(self, doc: RunStart) -> RunStart | None:
-        if "xray_centre_results" in doc:
+        if CONST.PLAN.FLYSCAN_RESULTS in doc:
             self.xray_centre_results = [
                 XRayCentreResult(**result_dict)
-                for result_dict in doc["xray_centre_results"]  # type: ignore
+                for result_dict in doc[CONST.PLAN.FLYSCAN_RESULTS]  # type: ignore
             ]
         return doc
 
@@ -304,7 +304,7 @@ def empty_plan():
 
     yield from bpp.run_wrapper(
         empty_plan(),
-        md={"xray_centre_results": [dataclasses.asdict(r) for r in results]},
+        md={CONST.PLAN.FLYSCAN_RESULTS: [dataclasses.asdict(r) for r in results]},
     )
 
 

src/mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
-
+import bluesky.plan_stubs as bps
+import numpy as np
 import pydantic
 from blueapi.core import BlueskyContext
 from bluesky.preprocessors import run_decorator, set_run_key_decorator, subs_wrapper
@@ -66,33 +66,40 @@ def plan_with_callback_subs():
             flyscan_event_handler,
         )
 
-        assert flyscan_event_handler.xray_centre_results, (
-            "Flyscan result event not received or no crystal found and exception not raised"
-        )
+        locations_to_collect_um: list[np.ndarray] = []
 
-        selection_func = flyscan_result.resolve_selection_fn(
-            parameters.selection_params
-        )
-        hits: Sequence[flyscan_result.XRayCentreResult] = selection_func(
-            flyscan_event_handler.xray_centre_results
-        )
-        LOGGER.info(
-            f"Selected hits {hits} using {selection_func}, args={parameters.selection_params}"
-        )
+        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]
+
+            LOGGER.info(
+                f"Selected hits {hits} using {selection_func}, args={parameters.selection_params}"
+            )
+        else:
+            # If the xray centring hasn't found a result but has not thrown an error it
+            # means that we do not need to recentre and can collect where we are
+            initial_x = yield from bps.rd(composite.smargon.x.user_readback)
+            initial_y = yield from bps.rd(composite.smargon.y.user_readback)
+            initial_z = yield from bps.rd(composite.smargon.z.user_readback)
+
+            locations_to_collect_um = [np.array([initial_x, initial_y, initial_z])]
 
         multi_rotation = parameters.multi_rotation_scan
         rotation_template = multi_rotation.rotation_scans.copy()
 
         multi_rotation.rotation_scans.clear()
 
-        for hit in hits:
+        for location in locations_to_collect_um:
             for rot in rotation_template:
                 combination = rot.model_copy()
                 (
                     combination.x_start_um,
                     combination.y_start_um,
                     combination.z_start_um,
-                ) = (axis * 1000 for axis in hit.centre_of_mass_mm)
+                ) = location
                 multi_rotation.rotation_scans.append(combination)
         multi_rotation = MultiRotationScan.model_validate(multi_rotation)
 

tests/unit_tests/hyperion/experiment_plans/test_load_centre_collect_full_plan.py

@@ -54,7 +54,7 @@ def composite(
     robot_load_composite,
     fake_create_rotation_devices,
     pin_tip_detection_with_found_pin,
-    sim_run_engine,
+    sim_run_engine: RunEngineSimulator,
 ) -> LoadCentreCollectComposite:
     rlaec_args = {
         field.name: getattr(robot_load_composite, field.name)
@@ -285,7 +285,7 @@ def test_collect_full_plan_happy_path_invokes_all_steps_and_centres_on_best_flys
     composite: LoadCentreCollectComposite,
     load_centre_collect_params: LoadCentreCollect,
     oav_parameters_for_rotation: OAVParameters,
-    sim_run_engine,
+    sim_run_engine: RunEngineSimulator,
 ):
     sim_run_engine.add_handler_for_callback_subscribes()
     sim_fire_event_on_open_run(sim_run_engine, CONST.PLAN.FLYSCAN_RESULTS)
@@ -360,12 +360,12 @@ def test_collect_full_plan_happy_path_invokes_all_steps_and_centres_on_best_flys
     "mx_bluesky.hyperion.experiment_plans.robot_load_and_change_energy.set_energy_plan",
     new=MagicMock(),
 )
-def test_load_centre_collect_ful_skips_collect_if_pin_tip_not_found(
+def test_load_centre_collect_full_skips_collect_if_pin_tip_not_found(
     mock_rotation_scan: MagicMock,
     composite: LoadCentreCollectComposite,
     load_centre_collect_params: LoadCentreCollect,
     oav_parameters_for_rotation: OAVParameters,
-    sim_run_engine,
+    sim_run_engine: RunEngineSimulator,
 ):
     sim_run_engine.add_read_handler_for(
         composite.pin_tip_detection.triggered_tip, PinTipDetection.INVALID_POSITION
@@ -396,8 +396,7 @@ def test_load_centre_collect_full_plan_skips_collect_if_no_diffraction(
     composite: LoadCentreCollectComposite,
     load_centre_collect_params: LoadCentreCollect,
     oav_parameters_for_rotation: OAVParameters,
-    sim_run_engine,
-    grid_detection_callback_with_detected_grid,
+    sim_run_engine: RunEngineSimulator,
 ):
     sim_run_engine.add_read_handler_for(composite.oav.microns_per_pixel_x, 1.58)
     sim_run_engine.add_read_handler_for(composite.oav.microns_per_pixel_y, 1.58)
@@ -612,11 +611,44 @@ def test_box_size_passed_through_to_gridscan(
     RE: RunEngine,
 ):
     load_centre_collect_params.robot_load_then_centre.box_size_um = 25
-    with pytest.raises(AssertionError, match="Flyscan result event not received.*"):
-        RE(
-            load_centre_collect_full(
-                composite, load_centre_collect_params, oav_parameters_for_rotation
-            )
+
+    RE(
+        load_centre_collect_full(
+            composite, load_centre_collect_params, oav_parameters_for_rotation
         )
+    )
     detect_grid_call = mock_detect_grid.mock_calls[0]
     assert detect_grid_call.args[1].box_size_um == 25
+
+
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan.multi_rotation_scan",
+    return_value=iter([]),
+)
+@patch(
+    "mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan.robot_load_then_xray_centre",
+    return_value=iter([]),
+)
+def test_load_centre_collect_full_collects_at_current_location_if_no_xray_centring_required(
+    _: MagicMock,
+    mock_rotation_scan: MagicMock,
+    composite: LoadCentreCollectComposite,
+    load_centre_collect_params: LoadCentreCollect,
+    oav_parameters_for_rotation: OAVParameters,
+    sim_run_engine: RunEngineSimulator,
+):
+    sim_run_engine.add_read_handler_for(composite.smargon.x, 1.1)
+    sim_run_engine.add_read_handler_for(composite.smargon.y, 2.2)
+    sim_run_engine.add_read_handler_for(composite.smargon.z, 3.3)
+
+    sim_run_engine.simulate_plan(
+        load_centre_collect_full(
+            composite, load_centre_collect_params, oav_parameters_for_rotation
+        )
+    )
+
+    rotation_scans = mock_rotation_scan.call_args.args[1].rotation_scans
+    assert len(rotation_scans) == 1
+    assert rotation_scans[0].x_start_um == 1.1
+    assert rotation_scans[0].y_start_um == 2.2
+    assert rotation_scans[0].z_start_um == 3.3