Start arming asynchronously for rotation and then take snapshots. Als…

…o simplifies the waiting logic in the rotation plan

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)
 
 

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)

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)

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,
+    )