make common cache speed plan

Author: Relm-Arrowny

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     # it will be auto-pinned to the latest release version by the pre-release workflow
     #
     "bluesky",
-    "dls-dodal>=1.56.0",
+    "dls-dodal@git+https://github.com/DiamondLightSource/dodal.git@1495-make-apple2-and-pgm-start-flying-in-sync",
     "ophyd-async[sim]",
     "scanspec",
 ]

src/sm_bluesky/common/plan_stubs/__init__.py

@@ -1,10 +1,12 @@
 from .detectors import set_area_detector_acquire_time
 from .motions import (
     MotorTable,
+    cache_speed,
     check_within_limit,
     get_motor_positions,
     get_velocity_and_step_size,
     move_motor_with_look_up,
+    restore_speed,
     set_slit_size,
 )
 
@@ -16,4 +18,6 @@
     "check_within_limit",
     "get_motor_positions",
     "get_velocity_and_step_size",
+    "cache_speed",
+    "restore_speed",
 ]

src/sm_bluesky/common/plan_stubs/motions.py

@@ -1,9 +1,10 @@
-from collections.abc import Hashable, Iterator
+import uuid
+from collections.abc import Generator, Hashable, Iterator
 from typing import Any
 
 import bluesky.plan_stubs as bps
 from bluesky.plan_stubs import abs_set
-from bluesky.utils import MsgGenerator, plan
+from bluesky.utils import Msg, MsgGenerator, plan
 from dodal.devices.slits import Slits
 from ophyd_async.epics.motor import Motor
 from pydantic import RootModel
@@ -172,3 +173,26 @@ def get_velocity_and_step_size(
         ideal_velocity = round(max_velocity, 3)
 
     return ideal_velocity, ideal_step_size
+
+
+def cache_speed(
+    devices_and_speeds: list[Motor],
+) -> Generator[Msg, Any, dict[Motor, float]]:
+    speeds = {}
+    for axis in devices_and_speeds:
+        speed = yield from bps.rd(axis.velocity)
+        speeds[axis] = speed
+    return speeds
+
+
+@plan
+def restore_speed(
+    devices_and_speeds: dict[Motor, float],
+    group: str | None = None,
+    wait_for_all: bool = True,
+) -> MsgGenerator:
+    reset_group = f"reset-{group if group else str(uuid.uuid4())[:6]}"
+    for device, speed in devices_and_speeds.items():
+        yield from bps.abs_set(device.velocity, speed, group=reset_group)
+    if wait_for_all:
+        yield from bps.wait(reset_group)

src/sm_bluesky/common/plans/fast_scan.py

@@ -14,7 +14,7 @@
 from ophyd_async.epics.motor import Motor
 
 from sm_bluesky.common.helper import add_extra_names_to_meta
-from sm_bluesky.common.plan_stubs import check_within_limit
+from sm_bluesky.common.plan_stubs import cache_speed, check_within_limit, restore_speed
 from sm_bluesky.log import LOGGER
 
 
@@ -182,13 +182,6 @@ def inner_fast_scan_grid(
     )
 
 
-@plan
-def reset_speed(old_speed, motor: Motor) -> MsgGenerator:
-    LOGGER.info(f"Clean up: setting motor speed to {old_speed}.")
-    if old_speed:
-        yield from bps.abs_set(motor.velocity, old_speed)
-
-
 def clean_up():
     LOGGER.info("Clean up")
     # possibly use to move back to starting position.
@@ -235,7 +228,7 @@ def _fast_scan_1d(
     """
 
     # read the current speed and store it
-    old_speed: float = yield from bps.rd(motor.velocity)
+    old_speed: dict[Motor, float] = yield from cache_speed([motor])
 
     def inner_fast_scan_1d(
         dets: list[Any],
@@ -245,7 +238,7 @@ def inner_fast_scan_1d(
         motor_speed: float | None = None,
     ):
         if not motor_speed:
-            motor_speed = old_speed
+            motor_speed = old_speed[motor]
 
         LOGGER.info(
             f"Starting 1d fly scan with {motor.name}:"
@@ -270,5 +263,5 @@ def inner_fast_scan_1d(
 
     yield from finalize_wrapper(
         plan=inner_fast_scan_1d(dets, motor, start, end, motor_speed),
-        final_plan=reset_speed(old_speed, motor),
+        final_plan=restore_speed(old_speed),
     )