diff --git a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py b/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py index a6ef4400f..4ed8f6382 100755 --- a/src/mx_bluesky/hyperion/experiment_plans/flyscan_xray_centre_plan.py +++ b/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]}, ) diff --git a/src/mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py b/src/mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py index 8154bdafc..fa63c86d9 100644 --- a/src/mx_bluesky/hyperion/experiment_plans/load_centre_collect_full_plan.py +++ b/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) diff --git a/tests/unit_tests/hyperion/experiment_plans/test_load_centre_collect_full_plan.py b/tests/unit_tests/hyperion/experiment_plans/test_load_centre_collect_full_plan.py index 65e4e59ae..4d32d1abe 100644 --- a/tests/unit_tests/hyperion/experiment_plans/test_load_centre_collect_full_plan.py +++ b/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