diff --git a/src/hyperion/device_setup_plans/manipulate_sample.py b/src/hyperion/device_setup_plans/manipulate_sample.py index d0149ac9d..ca3249026 100644 --- a/src/hyperion/device_setup_plans/manipulate_sample.py +++ b/src/hyperion/device_setup_plans/manipulate_sample.py @@ -16,15 +16,11 @@ def begin_sample_environment_setup( - detector_motion: DetectorMotion, attenuator: Attenuator, transmission_fraction: float, - detector_distance: float, group="setup_senv", ): """Start all sample environment changes that can be initiated before OAV snapshots are taken""" - yield from bps.abs_set(detector_motion.shutter, 1, group=group) - yield from bps.abs_set(detector_motion.z, detector_distance, group=group) yield from bps.abs_set(attenuator, transmission_fraction, group=group) diff --git a/src/hyperion/experiment_plans/rotation_scan_plan.py b/src/hyperion/experiment_plans/rotation_scan_plan.py index ea7cd504f..94338a048 100644 --- a/src/hyperion/experiment_plans/rotation_scan_plan.py +++ b/src/hyperion/experiment_plans/rotation_scan_plan.py @@ -40,6 +40,9 @@ make_trigger_safe, setup_zebra_for_rotation, ) +from hyperion.device_setup_plans.utils import ( + start_preparing_data_collection_then_do_plan, +) from hyperion.experiment_plans.oav_snapshot_plan import ( OavSnapshotComposite, oav_snapshot_plan, @@ -199,8 +202,7 @@ def _rotation_scan_plan( yield from bps.abs_set( axis, motion_values.start_motion_deg, - group="move_to_rotation_start", - wait=True, + group=CONST.WAIT.ROTATION_READY_FOR_DC, ) yield from setup_zebra_for_rotation( @@ -211,19 +213,18 @@ def _rotation_scan_plan( shutter_opening_deg=motion_values.shutter_opening_deg, shutter_opening_s=motion_values.shutter_time_s, group="setup_zebra", - wait=True, ) yield from setup_sample_environment( composite.aperture_scatterguard, params.selected_aperture, composite.backlight, + group=CONST.WAIT.ROTATION_READY_FOR_DC, ) LOGGER.info("Wait for any previous moves...") # wait for all the setup tasks at once - yield from bps.wait("setup_senv") - yield from bps.wait("move_to_rotation_start") + yield from bps.wait(CONST.WAIT.ROTATION_READY_FOR_DC) # get some information for the ispyb deposition and trigger the callback yield from read_hardware_for_zocalo(composite.eiger) @@ -312,16 +313,14 @@ def rotation_scan_plan_with_stage_and_cleanup( eiger: EigerDetector = composite.eiger eiger.set_detector_parameters(params.detector_params) - @bpp.stage_decorator([eiger]) @bpp.finalize_decorator(lambda: cleanup_plan(composite, max_vel)) def rotation_with_cleanup_and_stage(params: RotationScan): assert composite.aperture_scatterguard.aperture_positions is not None LOGGER.info("setting up sample environment...") yield from begin_sample_environment_setup( - composite.detector_motion, composite.attenuator, params.transmission_frac, - params.detector_params.detector_distance, + group=CONST.WAIT.ROTATION_READY_FOR_DC, ) LOGGER.info("moving to position (if specified)") yield from move_x_y_z( @@ -348,6 +347,13 @@ def rotation_with_cleanup_and_stage(params: RotationScan): yield from rotation_scan_plan(composite, params, motion_values) LOGGER.info("setting up and staging eiger...") - yield from rotation_with_cleanup_and_stage(params) + yield from start_preparing_data_collection_then_do_plan( + eiger, + composite.detector_motion, + params.detector_distance_mm, + rotation_with_cleanup_and_stage(params), + group=CONST.WAIT.ROTATION_READY_FOR_DC, + ) + yield from bps.unstage(eiger) yield from rotation_scan_plan_with_stage_and_cleanup(parameters) diff --git a/src/hyperion/parameters/constants.py b/src/hyperion/parameters/constants.py index f029aaea7..773de64a6 100644 --- a/src/hyperion/parameters/constants.py +++ b/src/hyperion/parameters/constants.py @@ -37,8 +37,8 @@ class PlanNameConstants: @dataclass(frozen=True) class PlanGroupCheckpointConstants: # For places to synchronise / stop and wait in plans, use as bluesky group names - # Gridscan - GRID_READY_FOR_DC = "ready_for_data_collection" + GRID_READY_FOR_DC = "grid_ready_for_data_collection" + ROTATION_READY_FOR_DC = "rotation_ready_for_data_collection" @dataclass(frozen=True) diff --git a/tests/unit_tests/experiment_plans/test_rotation_scan_plan.py b/tests/unit_tests/experiment_plans/test_rotation_scan_plan.py index 6068e2ac6..955b25f82 100644 --- a/tests/unit_tests/experiment_plans/test_rotation_scan_plan.py +++ b/tests/unit_tests/experiment_plans/test_rotation_scan_plan.py @@ -165,7 +165,7 @@ def test_rotation_scan( composite = fake_create_rotation_devices RE(rotation_scan(composite, test_rotation_params, oav_parameters_for_rotation)) - composite.eiger.stage.assert_called() # type: ignore + composite.eiger.do_arm.set.assert_called() # type: ignore composite.eiger.unstage.assert_called() # type: ignore @@ -208,7 +208,8 @@ async def test_full_rotation_plan_smargon_settings( assert await smargon.z.user_readback.get_value() == params.z_start_um assert ( # 4 * snapshots, restore omega, 1 * rotation sweep - omega_set.call_count == 4 + 1 + 1 + omega_set.call_count + == 4 + 1 + 1 ) # 1 to max vel in outer plan, 1 to max vel in setup_oav_snapshot_plan, 1 set before rotation, 1 restore in cleanup plan assert omega_velocity_set.call_count == 4 @@ -292,7 +293,6 @@ class MyTestException(Exception): def test_rotation_plan_reads_hardware( - RE: RunEngine, fake_create_rotation_devices: RotationScanComposite, test_rotation_params: RotationScan, motion_values: RotationMotionProfile, @@ -304,7 +304,6 @@ def test_rotation_plan_reads_hardware( fake_create_rotation_devices, test_rotation_params, motion_values ) ) - msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "create" @@ -322,70 +321,67 @@ def test_rotation_plan_reads_hardware( ) -def test_rotation_scan_initialises_detector_distance_shutter_and_tx_fraction( +@pytest.fixture +def rotation_scan_simulated_messages( sim_run_engine: RunEngineSimulator, fake_create_rotation_devices: RotationScanComposite, test_rotation_params: RotationScan, oav_parameters_for_rotation: OAVParameters, ): _add_sim_handlers_for_normal_operation(fake_create_rotation_devices, sim_run_engine) - - msgs = sim_run_engine.simulate_plan( + return sim_run_engine.simulate_plan( rotation_scan( fake_create_rotation_devices, test_rotation_params, oav_parameters_for_rotation, ) ) + + +def test_rotation_scan_initialises_detector_distance_shutter_and_tx_fraction( + rotation_scan_simulated_messages, + test_rotation_params: RotationScan, +): msgs = assert_message_and_return_remaining( - msgs, + rotation_scan_simulated_messages, lambda msg: msg.command == "set" - and msg.args[0] == 1 - and msg.obj.name == "detector_motion_shutter" - and msg.kwargs["group"] == "setup_senv", + and msg.args[0] == test_rotation_params.detector_distance_mm + and msg.obj.name == "detector_motion_z" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "set" - and msg.args[0] == test_rotation_params.detector_distance_mm - and msg.obj.name == "detector_motion_z" - and msg.kwargs["group"] == "setup_senv", + and msg.args[0] == 1 + and msg.obj.name == "detector_motion_shutter" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "set" and msg.obj.name == "attenuator" and msg.args[0] == test_rotation_params.transmission_frac - and msg.kwargs["group"] == "setup_senv", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "set" and msg.obj.name == "attenuator" and msg.args[0] == test_rotation_params.transmission_frac - and msg.kwargs["group"] == "setup_senv", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) assert_message_and_return_remaining( - msgs, lambda msg: msg.command == "wait" and msg.kwargs["group"] == "setup_senv" + msgs, + lambda msg: msg.command == "wait" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) def test_rotation_scan_moves_gonio_to_start_before_snapshots( - fake_create_rotation_devices: RotationScanComposite, - sim_run_engine: RunEngineSimulator, - test_rotation_params: RotationScan, - oav_parameters_for_rotation: OAVParameters, + rotation_scan_simulated_messages, ): - _add_sim_handlers_for_normal_operation(fake_create_rotation_devices, sim_run_engine) - msgs = sim_run_engine.simulate_plan( - rotation_scan( - fake_create_rotation_devices, - test_rotation_params, - oav_parameters_for_rotation, - ) - ) msgs = assert_message_and_return_remaining( - msgs, + rotation_scan_simulated_messages, lambda msg: msg.command == "wait" and msg.kwargs["group"] == "move_gonio_to_start", ) @@ -397,21 +393,10 @@ def test_rotation_scan_moves_gonio_to_start_before_snapshots( def test_rotation_scan_moves_aperture_in_backlight_out_after_snapshots_before_rotation( - fake_create_rotation_devices: RotationScanComposite, - sim_run_engine: RunEngineSimulator, - test_rotation_params: RotationScan, - oav_parameters_for_rotation: OAVParameters, + rotation_scan_simulated_messages, ): - _add_sim_handlers_for_normal_operation(fake_create_rotation_devices, sim_run_engine) - msgs = sim_run_engine.simulate_plan( - rotation_scan( - fake_create_rotation_devices, - test_rotation_params, - oav_parameters_for_rotation, - ) - ) msgs = assert_message_and_return_remaining( - msgs, + rotation_scan_simulated_messages, lambda msg: msg.command == "create" and msg.kwargs["name"] == DocDescriptorNames.OAV_ROTATION_SNAPSHOT_TRIGGERED, ) @@ -421,36 +406,27 @@ def test_rotation_scan_moves_aperture_in_backlight_out_after_snapshots_before_ro lambda msg: msg.command == "set" and msg.obj.name == "aperture_scatterguard" and msg.args[0] == AperturePositions.SMALL - and msg.kwargs["group"] == "setup_senv", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "set" and msg.obj.name == "backlight" and msg.args[0] == BacklightPosition.OUT - and msg.kwargs["group"] == "setup_senv", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) assert_message_and_return_remaining( - msgs, lambda msg: msg.command == "wait" and msg.kwargs["group"] == "setup_senv" + msgs, + lambda msg: msg.command == "wait" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) def test_rotation_scan_resets_omega_waits_for_sample_env_complete_after_snapshots_before_hw_read( - fake_create_rotation_devices: RotationScanComposite, - sim_run_engine: RunEngineSimulator, - test_rotation_params: RotationScan, - oav_parameters_for_rotation: OAVParameters, + test_rotation_params: RotationScan, rotation_scan_simulated_messages ): - _add_sim_handlers_for_normal_operation(fake_create_rotation_devices, sim_run_engine) - msgs = sim_run_engine.simulate_plan( - rotation_scan( - fake_create_rotation_devices, - test_rotation_params, - oav_parameters_for_rotation, - ) - ) msgs = assert_message_and_return_remaining( - msgs, + rotation_scan_simulated_messages, lambda msg: msg.command == "create" and msg.kwargs["name"] == DocDescriptorNames.OAV_ROTATION_SNAPSHOT_TRIGGERED, ) @@ -460,20 +436,12 @@ def test_rotation_scan_resets_omega_waits_for_sample_env_complete_after_snapshot lambda msg: msg.command == "set" and msg.obj.name == "smargon-omega" and msg.args[0] == test_rotation_params.omega_start_deg - and msg.kwargs["group"] == "move_to_rotation_start", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) msgs = assert_message_and_return_remaining( msgs, lambda msg: msg.command == "wait" - and msg.kwargs["group"] == "move_to_rotation_start", - ) - assert_message_and_return_remaining( - msgs, lambda msg: msg.command == "wait" and msg.kwargs["group"] == "setup_senv" - ) - assert_message_and_return_remaining( - msgs, - lambda msg: msg.command == "wait" - and msg.kwargs["group"] == "move_to_rotation_start", + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, ) assert_message_and_return_remaining( msgs, @@ -483,21 +451,10 @@ def test_rotation_scan_resets_omega_waits_for_sample_env_complete_after_snapshot def test_rotation_snapshot_setup_called_to_move_backlight_in_aperture_out_before_triggering( - fake_create_rotation_devices: RotationScanComposite, - sim_run_engine: RunEngineSimulator, - test_rotation_params: RotationScan, - oav_parameters_for_rotation: OAVParameters, + rotation_scan_simulated_messages, ): - _add_sim_handlers_for_normal_operation(fake_create_rotation_devices, sim_run_engine) - msgs = sim_run_engine.simulate_plan( - rotation_scan( - fake_create_rotation_devices, - test_rotation_params, - oav_parameters_for_rotation, - ) - ) msgs = assert_message_and_return_remaining( - msgs, + rotation_scan_simulated_messages, lambda msg: msg.command == "set" and msg.obj.name == "backlight" and msg.args[0] == BacklightPosition.IN @@ -549,7 +506,7 @@ def test_rotation_scan_skips_init_backlight_aperture_and_snapshots_if_snapshot_p for msg in msgs if msg.command == "set" and msg.obj.name == "smargon-omega" - and msg.kwargs["group"] == "move_to_rotation_start" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC ] ) == 1 @@ -572,3 +529,23 @@ def _add_sim_handlers_for_normal_operation( sim_run_engine.add_handler( "read", lambda msg: {"smargon-omega": {"value": -1}}, "smargon-omega" ) + + +def test_rotation_scan_arms_detector_and_takes_snapshots_whilst_arming( + rotation_scan_simulated_messages, +): + msgs = assert_message_and_return_remaining( + rotation_scan_simulated_messages, + lambda msg: msg.command == "set" + and msg.obj.name == "eiger_do_arm" + and msg.args[0] == 1 + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, + ) + msgs = assert_message_and_return_remaining( + msgs, lambda msg: msg.command == "trigger" and msg.obj.name == "oav_snapshot" + ) + msgs = assert_message_and_return_remaining( + msgs, + lambda msg: msg.command == "wait" + and msg.kwargs["group"] == CONST.WAIT.ROTATION_READY_FOR_DC, + )