Merge pull request DiamondLightSource/hyperion#1237 from DiamondLight…

…Source/685_use_vmax_for_rotation

Use VMAX for resetting omega before and after rotation scans

setup.cfg

@@ -35,7 +35,7 @@ install_requires =
     xarray
     doct
     databroker
-    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@97e3cdc11b1b5092c7f12ab6bc5ea1d702401b68
+    dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git
     pydantic<2.0 # See https://github.com/DiamondLightSource/hyperion/issues/774
     scipy
     pyzmq<25 # See https://github.com/DiamondLightSource/hyperion/issues/1103

src/hyperion/experiment_plans/rotation_scan_plan.py

@@ -95,12 +95,14 @@ class RotationMotionProfile:
     shutter_opening_deg: float
     total_exposure_s: float
     distance_to_move_deg: float
+    max_velocity_deg_s: float
 
 
 def calculate_motion_profile(
     detector_params: DetectorParams,
     expt_params: RotationScanParams,
     motor_time_to_speed_s: float,
+    max_velocity_deg_s: float,
 ) -> RotationMotionProfile:
     """Calculates the various numbers needed for motions in the rotation scan.
     Rotates through "scan width" plus twice an "offset" to take into account
@@ -149,23 +151,19 @@ def calculate_motion_profile(
         shutter_opening_deg=shutter_opening_deg,
         total_exposure_s=total_exposure_s,
         distance_to_move_deg=distance_to_move_deg,
+        max_velocity_deg_s=max_velocity_deg_s,
     )
 
 
 def rotation_scan_plan(
-    composite: RotationScanComposite, params: RotationInternalParameters
+    composite: RotationScanComposite,
+    params: RotationInternalParameters,
+    motion_values: RotationMotionProfile,
 ):
     """A plan to collect diffraction images from a sample continuously rotating about
     a fixed axis - for now this axis is limited to omega. Only does the scan itself, no
     setup tasks."""
 
-    motor_time_to_speed = yield from bps.rd(composite.smargon.omega.acceleration)
-    motion_values = calculate_motion_profile(
-        params.hyperion_params.detector_params,
-        params.experiment_params,
-        motor_time_to_speed,
-    )
-
     @bpp.set_run_key_decorator(CONST.PLAN.ROTATION_MAIN)
     @bpp.run_decorator(
         md={
@@ -182,7 +180,9 @@ def _rotation_scan_plan(
         LOGGER.info(f"moving omega to beginning, {motion_values.start_scan_deg=}")
         # can move to start as fast as possible
         # TODO get VMAX, see https://github.com/bluesky/ophyd/issues/1122
-        yield from bps.abs_set(axis.velocity, DEFAULT_MAX_VELOCITY, wait=True)
+        yield from bps.abs_set(
+            axis.velocity, motion_values.max_velocity_deg_s, wait=True
+        )
         yield from bps.abs_set(
             axis,
             motion_values.start_motion_deg,
@@ -242,12 +242,10 @@ def _rotation_scan_plan(
     yield from _rotation_scan_plan(motion_values, composite)
 
 
-def cleanup_plan(composite: RotationScanComposite, **kwargs):
+def cleanup_plan(composite: RotationScanComposite, max_vel: float, **kwargs):
     LOGGER.info("Cleaning up after rotation scan")
     yield from cleanup_sample_environment(composite.detector_motion, group="cleanup")
-    yield from bps.abs_set(
-        composite.smargon.omega.velocity, DEFAULT_MAX_VELOCITY, group="cleanup"
-    )
+    yield from bps.abs_set(composite.smargon.omega.velocity, max_vel, group="cleanup")
     yield from make_trigger_safe(composite.zebra, group="cleanup")
     yield from bpp.finalize_wrapper(disarm_zebra(composite.zebra), bps.wait("cleanup"))
 
@@ -268,11 +266,23 @@ def rotation_scan(composite: RotationScanComposite, parameters: Any) -> MsgGener
     def rotation_scan_plan_with_stage_and_cleanup(
         params: RotationInternalParameters,
     ):
+        motor_time_to_speed = yield from bps.rd(composite.smargon.omega.acceleration)
+        max_vel = (
+            yield from bps.rd(composite.smargon.omega.max_velocity)
+            or DEFAULT_MAX_VELOCITY
+        )
+        motion_values = calculate_motion_profile(
+            params.hyperion_params.detector_params,
+            params.experiment_params,
+            motor_time_to_speed,
+            max_vel,
+        )
+
         eiger: EigerDetector = composite.eiger
         eiger.set_detector_parameters(params.hyperion_params.detector_params)
 
         @bpp.stage_decorator([eiger])
-        @bpp.finalize_decorator(lambda: cleanup_plan(composite=composite))
+        @bpp.finalize_decorator(lambda: cleanup_plan(composite, max_vel))
         def rotation_with_cleanup_and_stage(params: RotationInternalParameters):
             LOGGER.info("setting up sample environment...")
             yield from setup_sample_environment(
@@ -290,10 +300,10 @@ def rotation_with_cleanup_and_stage(params: RotationInternalParameters):
                 params.experiment_params.z,
                 group="move_x_y_z",
             )
-
             yield from rotation_scan_plan(
                 composite,
                 params,
+                motion_values,
             )
 
         LOGGER.info("setting up and staging eiger...")

tests/conftest.py

@@ -468,6 +468,8 @@ def fake_create_rotation_devices(
     smargon.omega.velocity.set = mock_omega_velocity_sets
     smargon.omega.set = mock_omega_sets
 
+    smargon.omega.max_velocity.sim_put(131)  # type: ignore
+
     return RotationScanComposite(
         attenuator=attenuator,
         backlight=backlight,

tests/unit_tests/experiment_plans/test_rotation_scan_plan.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 from unittest.mock import DEFAULT, MagicMock, call, patch
 
 import pytest
@@ -11,7 +11,7 @@
 from ophyd.status import Status
 
 from hyperion.experiment_plans.rotation_scan_plan import (
-    DEFAULT_MAX_VELOCITY,
+    RotationMotionProfile,
     RotationScanComposite,
     calculate_motion_profile,
     rotation_scan,
@@ -35,17 +35,15 @@ def do_rotation_main_plan_for_tests(
     run_eng: RunEngine,
     expt_params: RotationInternalParameters,
     devices: RotationScanComposite,
-    plan=rotation_scan_plan,
+    motion_values: RotationMotionProfile,
+    plan: Callable = rotation_scan_plan,
 ):
     with patch(
         "bluesky.preprocessors.__read_and_stash_a_motor",
         fake_read,
     ):
         run_eng(
-            plan(
-                devices,
-                expt_params,
-            ),
+            plan(devices, expt_params, motion_values),
         )
 
 
@@ -55,16 +53,31 @@ def run_full_rotation_plan(
     test_rotation_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
 ):
-    do_rotation_main_plan_for_tests(
-        RE, test_rotation_params, fake_create_rotation_devices, rotation_scan
+    with patch(
+        "bluesky.preprocessors.__read_and_stash_a_motor",
+        fake_read,
+    ):
+        RE(
+            rotation_scan(fake_create_rotation_devices, test_rotation_params),
+        )
+        return fake_create_rotation_devices
+
+
+@pytest.fixture
+def motion_values(test_rotation_params: RotationInternalParameters):
+    return calculate_motion_profile(
+        test_rotation_params.hyperion_params.detector_params,
+        test_rotation_params.experiment_params,
+        0.005,
+        222,
     )
-    return fake_create_rotation_devices
 
 
 def setup_and_run_rotation_plan_for_tests(
     RE: RunEngine,
     test_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
+    motion_values,
 ):
     smargon = fake_create_rotation_devices.smargon
 
@@ -79,9 +92,7 @@ def side_set_w_return(obj, *args):
 
     with patch("bluesky.plan_stubs.wait", autospec=True):
         do_rotation_main_plan_for_tests(
-            RE,
-            test_params,
-            fake_create_rotation_devices,
+            RE, test_params, fake_create_rotation_devices, motion_values
         )
 
     return {
@@ -97,11 +108,10 @@ def setup_and_run_rotation_plan_for_tests_standard(
     RE: RunEngine,
     test_rotation_params: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
+    motion_values,
 ):
     return setup_and_run_rotation_plan_for_tests(
-        RE,
-        test_rotation_params,
-        fake_create_rotation_devices,
+        RE, test_rotation_params, fake_create_rotation_devices, motion_values
     )
 
 
@@ -110,11 +120,10 @@ def setup_and_run_rotation_plan_for_tests_nomove(
     RE: RunEngine,
     test_rotation_params_nomove: RotationInternalParameters,
     fake_create_rotation_devices: RotationScanComposite,
+    motion_values,
 ):
     return setup_and_run_rotation_plan_for_tests(
-        RE,
-        test_rotation_params_nomove,
-        fake_create_rotation_devices,
+        RE, test_rotation_params_nomove, fake_create_rotation_devices, motion_values
     )
 
 
@@ -127,6 +136,7 @@ def test_rotation_scan_calculations(test_rotation_params: RotationInternalParame
         test_rotation_params.hyperion_params.detector_params,
         test_rotation_params.experiment_params,
         0.005,  # time for acceleration
+        224,
     )
 
     assert motion_values.direction == -1
@@ -202,6 +212,8 @@ def test_full_rotation_plan_smargon_settings(
     params: RotationInternalParameters = test_rotation_params
     expt_params = params.experiment_params
 
+    test_max_velocity = smargon.omega.max_velocity.get()
+
     omega_set: MagicMock = smargon.omega.set  # type: ignore
     omega_velocity_set: MagicMock = smargon.omega.velocity.set  # type: ignore
     rotation_speed = (
@@ -216,9 +228,9 @@ def test_full_rotation_plan_smargon_settings(
     assert omega_set.call_count == 2
     assert omega_velocity_set.call_count == 3
     assert omega_velocity_set.call_args_list == [
-        call(DEFAULT_MAX_VELOCITY),
+        call(test_max_velocity),
         call(rotation_speed),
-        call(DEFAULT_MAX_VELOCITY),
+        call(test_max_velocity),
     ]
 
 
@@ -249,6 +261,7 @@ def test_cleanup_happens(
     RE: RunEngine,
     test_rotation_params,
     fake_create_rotation_devices: RotationScanComposite,
+    motion_values: RotationMotionProfile,
 ):
 
     class MyTestException(Exception):
@@ -263,18 +276,12 @@ class MyTestException(Exception):
         with pytest.raises(MyTestException):
             RE(
                 rotation_scan_plan(
-                    fake_create_rotation_devices,
-                    test_rotation_params,
+                    fake_create_rotation_devices, test_rotation_params, motion_values
                 )
             )
             cleanup_plan.assert_not_called()
         # check that failure is handled in composite plan
         with pytest.raises(MyTestException) as exc:
-            RE(
-                rotation_scan(
-                    fake_create_rotation_devices,
-                    test_rotation_params,
-                )
-            )
+            RE(rotation_scan(fake_create_rotation_devices, test_rotation_params))
             assert "Experiment fails because this is a test" in exc.value.args[0]
             cleanup_plan.assert_called_once()