diff --git a/src/pdftools/profiles/collection.py b/src/pdftools/profiles/collection.py index ad15c49..48e1ff3 100644 --- a/src/pdftools/profiles/collection.py +++ b/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.") diff --git a/src/pdftools/xrd.py b/src/pdftools/xrd.py index e5e5253..7ecbf71 100644 --- a/src/pdftools/xrd.py +++ b/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())) diff --git a/tests/test_xrd_jog.py b/tests/test_xrd_jog.py new file mode 100644 index 0000000..c12f7ff --- /dev/null +++ b/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