Fast energy scan plan

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 .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,8 +1,15 @@
+from typing import Any
+
+import bluesky.plan_stubs as bps
 from blueapi.core import MsgGenerator
 from bluesky.plan_stubs import abs_set
-from bluesky.utils import plan
+from bluesky.protocols import Flyable
+from bluesky.utils import 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(
@@ -27,3 +34,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
@@ -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/__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

@@ -8,13 +8,19 @@
 )
 from bluesky.protocols import Readable
 from bluesky.utils import plan, short_uid
+from dodal.devices.apple2_undulator import EnergySetter
 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
 
 
@@ -182,13 +188,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 +234,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 +244,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}:"
@@ -259,16 +258,46 @@ 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: EnergySetter,
+    energy_start: float,
+    energy_end: float,
+    energy_step: float,
+    count_time: float,
+    md: dict[str, Any] | None = None,
+):
+    old_speeds = yield from cache_speed(
+        [energy_device.pgm_ref().energy, energy_device.id.gap]
+    )
+
+    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: EnergySetter,
+        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),
     )

tests/common/plans/test_fast_energy_scan.py

@@ -0,0 +1,112 @@
+from collections import defaultdict
+from unittest.mock import ANY
+
+import pytest
+from bluesky.run_engine import RunEngine
+from dodal.devices.apple2_undulator import (
+    EnergySetter,
+    UndulatorGateStatus,
+)
+from dodal.devices.i10.i10_apple2 import I10Apple2
+from dodal.devices.pgm import PGM
+from ophyd_async.core import (
+    Device,
+    StrictEnum,
+    init_devices,
+)
+from ophyd_async.testing import assert_emitted, set_mock_value
+
+from sm_bluesky.common.plans import soft_fly_energy_scan
+
+from ...sim_devices.sim_stage import SimMotorExtra
+from ...test_data.common import (
+    LOOKUP_TABLE_PATH,
+)
+
+
+class Grating(StrictEnum):
+    TESTING = "Grating"
+
+
+class FakePGM(Device):
+    def __init__(self, name=""):
+        self.energy = SimMotorExtra(instant=False)
+        super().__init__(name=name)
+
+
+@pytest.fixture
+async def mock_pgm(prefix: str = "BLXX-EA-DET-007:") -> FakePGM:
+    async with init_devices(mock=True):
+        mock_pgm = FakePGM()
+    return mock_pgm
+
+
+@pytest.fixture
+async def mock_energy(mock_pgm: PGM) -> EnergySetter:
+    async with init_devices(mock=True):
+        mock_energy = EnergySetter(
+            id=I10Apple2(
+                look_up_table_dir=LOOKUP_TABLE_PATH,
+                source=("Source", "idu"),
+                prefix="BLWOW-MO-SERVC-01:",
+            ),
+            pgm=mock_pgm,
+        )
+
+    set_mock_value(mock_energy.id.gap.gate, UndulatorGateStatus.CLOSE)
+    set_mock_value(mock_energy.id.gap.high_limit_travel, 200)
+    set_mock_value(mock_energy.id.gap.low_limit_travel, 10)
+    set_mock_value(mock_energy.id.gap.acceleration_time, 0.2)
+    set_mock_value(mock_energy.id.phase.gate, UndulatorGateStatus.CLOSE)
+    set_mock_value(mock_energy.id.id_jaw_phase.gate, UndulatorGateStatus.CLOSE)
+    set_mock_value(mock_energy.id.id_jaw_phase.jaw_phase.velocity, 1)
+    set_mock_value(mock_energy.id.gap.velocity, 2)
+    set_mock_value(mock_energy.id.gap.max_velocity, 200)
+    set_mock_value(mock_energy.id.gap.min_velocity, 0.0)
+
+    set_mock_value(mock_energy.id.phase.btm_inner.velocity, 1)
+    set_mock_value(mock_energy.id.phase.top_inner.velocity, 1)
+    set_mock_value(mock_energy.id.phase.btm_outer.velocity, 1)
+    set_mock_value(mock_energy.id.phase.top_outer.velocity, 1)
+    set_mock_value(mock_energy.pgm_ref().energy.acceleration_time, 0.1)
+    set_mock_value(mock_energy.pgm_ref().energy.user_readback, 500)
+    set_mock_value(mock_energy.pgm_ref().energy.user_setpoint, 500)
+    set_mock_value(mock_energy.pgm_ref().energy.max_velocity, 50)
+    set_mock_value(mock_energy.pgm_ref().energy.high_limit_travel, 1700)
+    set_mock_value(mock_energy.pgm_ref().energy.low_limit_travel, 400)
+    return mock_energy
+
+
+async def test_soft_fly_energy_scan_success(
+    mock_energy: EnergySetter, RE: RunEngine, det
+):
+    docs = defaultdict(list)
+    det.start_simulation()
+
+    def capture_emitted(name, doc):
+        docs[name].append(doc)
+
+    RE(
+        soft_fly_energy_scan([det], mock_energy, 700, 800, 0.2, 1e-3),
+        capture_emitted,
+        wait=True,
+    )
+
+    assert_emitted(docs, start=1, descriptor=1, event=ANY, stop=1)
+    # Number of event depend how fast motor is moving, it has to be more than 1
+    assert len(docs["event"]) > 1
+    # check the starting point
+    assert docs["event"][0]["data"] == {
+        "rand": ANY,
+        "mock_energy-id-energy": 750,
+        "mock_pgm-energy": 700.0,
+    }
+    # check end point
+    assert docs["event"][-1]["data"] == {
+        "rand": ANY,
+        "mock_energy-id-energy": 750,
+        "mock_pgm-energy": 810.0,
+    }
+    # speed reset
+    assert await mock_energy.pgm_ref().energy.velocity.get_value() == 1.0
+    assert await mock_energy.id.gap.velocity.get_value() == 2.0

tests/common/plans/test_fast_scan.py

@@ -5,7 +5,6 @@
 from bluesky.run_engine import RunEngine
 from dodal.devices.motors import XYZStage
 from numpy import linspace
-from ophyd.sim import SynPeriodicSignal
 from ophyd_async.testing import assert_emitted, get_mock_put
 
 from sm_bluesky.common.plans.fast_scan import fast_scan_1d, fast_scan_grid
@@ -15,12 +14,6 @@
 A_BIT = 0.001
 
 
-@pytest.fixture
-def det():
-    det = SynPeriodicSignal(name="rand", labels={"detectors"})
-    return det
-
-
 async def test_fast_scan_1d_fail_limit_check(sim_motor: XYZStage, RE: RunEngine, det):
     """Testing both high and low limits making sure nothing get run if it is exceeded"""
     docs = defaultdict(list)

tests/conftest.py

@@ -13,6 +13,7 @@
 )
 from dodal.devices.motors import XYZStage
 from dodal.utils import make_all_devices
+from ophyd.sim import SynPeriodicSignal
 from ophyd_async.core import (
     FilenameProvider,
     StaticFilenameProvider,
@@ -197,3 +198,9 @@ async def andor2_point() -> SingleTriggerDetector:
         andor2_point = SingleTriggerDetector(drv=ADBaseIO("p99"))
 
     return andor2_point
+
+
+@pytest.fixture
+def det():
+    det = SynPeriodicSignal(name="rand", labels={"detectors"})
+    return det