Fast energy scan plan

src/sm_bluesky/common/plan_stubs/__init__.py

@@ -1,10 +1,12 @@
-from .detectors import set_area_detector_acquire_time
+from .detection import fly_trigger_and_read, 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,7 @@
     "check_within_limit",
     "get_motor_positions",
     "get_velocity_and_step_size",
+    "cache_speed",
+    "restore_speed",
+    "fly_trigger_and_read",
 ]

src/sm_bluesky/common/plan_stubs/detectors.py → src/sm_bluesky/common/plan_stubs/detection.py

@@ -1,7 +1,14 @@
+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.protocols import Flyable
+from bluesky.utils import MsgGenerator, plan, short_uid
+from ophyd_async.core import FlyMotorInfo
 from ophyd_async.epics.adcore import AreaDetector, SingleTriggerDetector
 
+from sm_bluesky.log import LOGGER
+
 
 @plan
 def set_area_detector_acquire_time(
@@ -26,3 +33,19 @@ def set_area_detector_acquire_time(
     """
     drv = det.drv if isinstance(det, SingleTriggerDetector) else det.driver
     yield from abs_set(drv.acquire_time, acquire_time, wait=wait)
+
+
+@plan
+def fly_trigger_and_read(
+    motor: Flyable,
+    fly_info: FlyMotorInfo,
+    dets: list[Any],
+) -> MsgGenerator:
+    grp = short_uid("kickoff")
+    yield from bps.kickoff(motor, group=grp, wait=True)
+    LOGGER.info(f"flying motor =  {motor.name} at with info = {fly_info}")
+    done = yield from bps.complete(motor)
+    yield from bps.trigger_and_read(dets + [motor])
+    while not done.done:
+        yield from bps.trigger_and_read(dets + [motor])
+        yield from bps.checkpoint()

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
@@ -173,3 +174,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/__init__.py

@@ -12,7 +12,7 @@
     fast_scan_and_move_fit,
     step_scan_and_move_fit,
 )
-from .fast_scan import fast_scan_1d, fast_scan_grid
+from .fast_scan import fast_scan_1d, fast_scan_grid, soft_fly_energy_scan
 from .grid_scan import grid_fast_scan, grid_step_scan
 
 __all__ = [
@@ -25,4 +25,5 @@
     "grid_fast_scan",
     "grid_step_scan",
     "trigger_img",
+    "soft_fly_energy_scan",
 ]

src/sm_bluesky/common/plans/fast_scan.py

@@ -7,13 +7,19 @@
 )
 from bluesky.protocols import Readable
 from bluesky.utils import MsgGenerator, plan, short_uid
+from dodal.devices.insertion_device import BeamEnergy
 from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator
 from numpy import linspace
 from ophyd_async.core import FlyMotorInfo
 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,
+    fly_trigger_and_read,
+    restore_speed,
+)
 from sm_bluesky.log import LOGGER
 
 
@@ -181,13 +187,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.
@@ -234,7 +233,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],
@@ -244,7 +243,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}:"
@@ -258,16 +257,49 @@ def inner_fast_scan_1d(
             time_for_move=abs(start - end) / motor_speed,
         )
         yield from bps.prepare(motor, fly_info, group=grp, wait=True)
-        yield from bps.wait(group=grp)
-        yield from bps.kickoff(motor, group=grp, wait=True)
-        LOGGER.info(f"flying motor =  {motor.name} at speed = {motor_speed}")
-        done = yield from bps.complete(motor)
-        yield from bps.trigger_and_read(dets + [motor])
-        while not done.done:
-            yield from bps.trigger_and_read(dets + [motor])
-            yield from bps.checkpoint()
+        yield from fly_trigger_and_read(motor, fly_info, dets)
 
     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),
+    )
+
+
+@plan
+@attach_data_session_metadata_decorator()
+def soft_fly_energy_scan(
+    dets: list[Readable],
+    energy_device: BeamEnergy,
+    energy_start: float,
+    energy_end: float,
+    energy_step: float,
+    count_time: float,
+    md: dict[str, Any] | None = None,
+) -> MsgGenerator:
+    old_speeds = yield from cache_speed(
+        [
+            energy_device._mono_energy(),  # noqa: SLF001
+            energy_device._id_energy()._id_controller().apple2().gap(),  # noqa: SLF001
+        ]
+    )
+
+    fly_info = FlyMotorInfo(
+        start_position=energy_start,
+        end_position=energy_end,
+        time_for_move=abs(energy_end - energy_start) / energy_step * count_time,
+    )
+
+    @bpp.stage_decorator(dets)
+    @bpp.run_decorator(md=md)
+    def inn_fly_energy_scan(
+        energy_device: BeamEnergy,
+        fly_info: FlyMotorInfo,
+        dets: list[Readable],
+    ) -> MsgGenerator:
+        yield from bps.prepare(energy_device, fly_info, wait=True)
+        yield from fly_trigger_and_read(energy_device, fly_info, dets)
+
+    yield from finalize_wrapper(
+        plan=inn_fly_energy_scan(energy_device, fly_info, dets),
+        final_plan=restore_speed(old_speeds),
     )