Profile conversion 5b: xrd_jog

src/pdftools/profiles/collection.py

@@ -35,7 +35,7 @@
     print_version_info,
     remove_md_keys,
 )
-from pdftools.xrd import xrd_count, xrd_dark, xrd_scan  # noqa: F401  for the session
+from pdftools.xrd import xrd_count, xrd_dark, xrd_jog, xrd_scan  # noqa: F401
 
 # Mechanism for passing arguments to profile.
 _parser = argparse.ArgumentParser("Collection profile for pdftools.")

src/pdftools/xrd.py

@@ -13,6 +13,7 @@
 from bluesky import plans as bp
 from bluesky import preprocessors as bpp
 from bluesky.protocols import Movable, Readable
+from bluesky.utils import short_uid
 from ophyd_async.core import DetectorTrigger, StandardDetector, TriggerInfo
 from ophyd_async.epics.adcore import ADImageMode, ContAcqDetector
 from ophyd_async.epics.motor import Motor
@@ -326,3 +327,96 @@ def _cleanup():
             delay=time_between_iterations,
         )
     )
+
+
+def xrd_jog(
+    detectors: Sequence[StandardDetector],
+    motor: Steppable,
+    start: float,
+    stop: float,
+    exposure_time: float,
+    images_to_average: int = 1,
+    num_images: int = 1,
+    sample_name: str | None = None,
+    description: str | None = None,
+    use_shutter: bool = True,
+    shutter: FastShutter | None = None,
+    md: dict | None = None,
+):
+    """Take XRD images while sweeping one motor across the sample.
+
+    Each image is taken while the motor travels from one end to the other, at
+    the speed that makes the sweep last the whole image (``exposure_time`` x
+    ``images_to_average``); the next image sweeps back. As pdf-profile-collection
+    ``jog``/``rocking_ct`` did; the motor's velocity is restored afterwards.
+
+    Parameters
+    ----------
+    detectors : list of detectors
+        the detectors to read, e.g. ``[pe1c]``
+    motor : Motor
+        the motor to sweep, e.g. a sample stage axis
+    start : float
+        one end of the sweep
+    stop : float
+        the other end; equal to ``start`` takes images without moving
+    exposure_time : float
+        exposure time per frame, in seconds
+    images_to_average : int
+        number of frames combined into each saved image
+    num_images : int
+        number of images, one sweep each
+    sample_name : str, optional
+        name of the sample
+    description : str, optional
+        description of the acquisition
+    use_shutter : bool
+        whether to open the fast shutter for the acquisition and close it after
+    shutter : FastShutter, optional
+        the fast shutter; the session's ``fs`` if not given
+    md : dict, optional
+        extra metadata for the start document
+    """
+    if use_shutter:
+        shutter = ensure_available(FastShutter, fs=shutter)
+    trigger_info = TriggerInfo(
+        trigger=DetectorTrigger.INTERNAL,
+        livetime=exposure_time,
+        exposures_per_collection=images_to_average,
+    )
+    sweep_time = exposure_time * images_to_average
+    ends = [start, stop]
+    to_prepare = list(detectors)
+    jog_md = _md("xrd_jog", sample_name, description, md)
+    jog_md["jog"] = {"motor": motor.name, "start": start, "stop": stop}
+
+    # bluesky's count passes the detectors as the first argument, by this name.
+    def _per_shot(detectors, take_reading=None):
+        # stage() at the start of the run clears what prepare() set
+        for det in to_prepare:
+            yield from bps.prepare(det, trigger_info, wait=True)
+        if start == stop:
+            return (yield from bps.trigger_and_read(list(detectors)))
+        here, there = ends
+        yield from bps.mv(motor, here)
+        yield from bps.mv(motor.velocity, abs(there - here) / sweep_time)
+        group = short_uid("jog")
+        yield from bps.abs_set(motor, there, group=group)
+        yield from bps.trigger_and_read(list(detectors))
+        yield from bps.wait(group=group)
+        ends.reverse()
+
+    def _body():
+        for det in detectors:
+            yield from ensure_cont_acq_exposure(det, exposure_time)
+        if use_shutter:
+            yield from bps.mv(shutter, "Open")
+        yield from bp.count(detectors, num=num_images, per_shot=_per_shot, md=jog_md)
+
+    def _cleanup():
+        if use_shutter:
+            yield from bps.mv(shutter, "Close")
+        yield from bps.mv(motor.velocity, velocity)
+
+    velocity = yield from bps.rd(motor.velocity)
+    return (yield from bpp.finalize_wrapper(_body(), _cleanup()))

tests/test_xrd_jog.py

@@ -0,0 +1,118 @@
+"""The jog plan against a mocked Perkin Elmer (conftest ``pe``) and fast shutter."""
+
+import asyncio
+from typing import Any
+
+import pytest
+from bluesky import plan_stubs as bps
+from bluesky.run_engine import RunEngine
+from ophyd_async.core import (
+    callback_on_mock_put,
+    get_mock_put,
+    init_devices,
+    set_mock_value,
+)
+from ophyd_async.epics.adcore import ContAcqDetector
+from ophyd_async.sim import SimMotor
+
+from pdftools.shutter import FastShutter, OnOff
+from pdftools.xrd import xrd_jog
+
+
+@pytest.fixture
+def fs(RE: RunEngine) -> FastShutter:
+    with init_devices(mock=True):
+        fs = FastShutter("FS:", name="fs", settle_time=0)
+    set_mock_value(fs.status, OnOff.ON)  # closed
+    callback_on_mock_put(fs.cmd, lambda value, **_: set_mock_value(fs.status, value))
+    return fs
+
+
+@pytest.fixture
+def grid_y(RE: RunEngine) -> SimMotor:
+    with init_devices():
+        grid_y = SimMotor(name="grid_y", instant=True)
+    RE(bps.mv(grid_y.velocity, 2.5))
+    return grid_y
+
+
+def _run(RE: RunEngine, plan) -> tuple[dict[str, list[Any]], list]:
+    docs: dict[str, list[Any]] = {}
+    msgs: list = []
+    RE.msg_hook = msgs.append  # pyright: ignore[reportAttributeAccessIssue]
+    try:
+        RE(plan, lambda name, doc: docs.setdefault(name, []).append(doc))
+    finally:
+        RE.msg_hook = None
+    return docs, msgs
+
+
+def _moves(msgs: list, obj) -> list[float]:
+    return [m.args[0] for m in msgs if m.command == "set" and m.obj is obj]
+
+
+async def test_xrd_jog_sweeps_back_and_forth_over_each_image(
+    RE: RunEngine, pe: ContAcqDetector, fs: FastShutter, grid_y: SimMotor
+):
+    docs, msgs = _run(
+        RE,
+        xrd_jog(
+            [pe],
+            grid_y,
+            0.0,
+            2.0,
+            0.1,
+            images_to_average=5,
+            num_images=3,
+            sample_name="Ni",
+            shutter=fs,
+        ),
+    )
+    (start,) = docs["start"]
+    assert start["plan_name"] == "xrd_jog"
+    assert start["sample_name"] == "Ni"
+    assert start["jog"] == {"motor": "grid_y", "start": 0.0, "stop": 2.0}
+    assert len(docs["event"]) == 3
+    # to start, then across; the next image sweeps back the other way
+    assert _moves(msgs, grid_y) == [0.0, 2.0, 2.0, 0.0, 0.0, 2.0]
+    # each sweep lasts one averaged image: 2.0 mm over 5 x 0.1 s
+    assert _moves(msgs, grid_y.velocity) == [pytest.approx(4.0)] * 3 + [2.5]
+    assert await grid_y.velocity.get_value() == 2.5
+    assert await pe.proc.num_filter.get_value() == 5  # pyright: ignore[reportAttributeAccessIssue]
+    assert [c.args[0] for c in get_mock_put(fs.cmd).call_args_list] == [
+        OnOff.OFF,
+        OnOff.ON,
+    ]
+
+
+async def test_xrd_jog_without_travel_does_not_move(
+    RE: RunEngine, pe: ContAcqDetector, fs: FastShutter, grid_y: SimMotor
+):
+    docs, msgs = _run(
+        RE, xrd_jog([pe], grid_y, 1.0, 1.0, 0.1, num_images=2, shutter=fs)
+    )
+    assert len(docs["event"]) == 2
+    assert _moves(msgs, grid_y) == []
+    assert await grid_y.velocity.get_value() == 2.5
+
+
+async def test_xrd_jog_restores_velocity_and_closes_the_shutter_on_failure(
+    RE: RunEngine, pe: ContAcqDetector, fs: FastShutter, grid_y: SimMotor
+):
+    tiff, cb = pe.tiff, pe.cb  # pyright: ignore[reportAttributeAccessIssue]
+    triggers: list[bool] = []
+
+    async def _fault_on_second_image(value):
+        if not value:
+            return
+        triggers.append(value)
+        if len(triggers) == 2:
+            raise RuntimeError("detector fault")
+        set_mock_value(tiff.num_captured, await tiff.num_captured.get_value() + 1)
+        asyncio.get_running_loop().call_later(0.01, set_mock_value, cb.trigger_, False)
+
+    callback_on_mock_put(cb.trigger_, _fault_on_second_image)
+    with pytest.raises(Exception, match="detector fault"):
+        RE(xrd_jog([pe], grid_y, 0.0, 2.0, 0.1, num_images=3, shutter=fs))
+    assert await grid_y.velocity.get_value() == 2.5
+    assert await fs.status.get_value() == OnOff.ON