From 82bf740580e6febd6aa1ca2eecb6926623de8a5f Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Tue, 11 Nov 2025 14:08:49 +0000 Subject: [PATCH 01/16] Wrap plans and plan stubs --- src/dodal/plan_stubs/__init__.py | 22 +- src/dodal/plan_stubs/wrapped.py | 26 +- src/dodal/plans/__init__.py | 4 +- src/dodal/plans/wrapped.py | 108 ++++++++ tests/plan_stubs/test_wrapped_stubs.py | 10 + tests/plans/test_wrapped.py | 363 ++++++++++++++++++++++++- 6 files changed, 526 insertions(+), 7 deletions(-) diff --git a/src/dodal/plan_stubs/__init__.py b/src/dodal/plan_stubs/__init__.py index 2c84ded8654..25f7a5e97c4 100644 --- a/src/dodal/plan_stubs/__init__.py +++ b/src/dodal/plan_stubs/__init__.py @@ -1,3 +1,21 @@ -from .wrapped import move, move_relative, set_absolute, set_relative, sleep, wait +from .wrapped import ( + move, + move_relative, + rd, + set_absolute, + set_relative, + sleep, + stop, + wait, +) -__all__ = ["move", "move_relative", "set_absolute", "set_relative", "sleep", "wait"] +__all__ = [ + "move", + "move_relative", + "rd", + "set_absolute", + "set_relative", + "sleep", + "stop", + "wait", +] diff --git a/src/dodal/plan_stubs/wrapped.py b/src/dodal/plan_stubs/wrapped.py index 43555dbbc16..1d134c2fb7c 100644 --- a/src/dodal/plan_stubs/wrapped.py +++ b/src/dodal/plan_stubs/wrapped.py @@ -3,7 +3,7 @@ from typing import Annotated, TypeVar import bluesky.plan_stubs as bps -from bluesky.protocols import Movable +from bluesky.protocols import Movable, Readable, Stoppable from bluesky.utils import MsgGenerator """ @@ -146,3 +146,27 @@ def wait( """ return (yield from bps.wait(group, timeout=timeout)) + + +def rd(readable: Readable) -> MsgGenerator: + """Reads a single-value non-triggered object, wrapper for `bp.rd`. + + Args: + readable (Readable): The device to be read + + Returns: + Iterator[MsgGenerator]: Bluesky messages + """ + return (yield from bps.rd(readable)) + + +def stop(stoppable: Stoppable) -> MsgGenerator: + """Stop a device, wrapper for `bp.stop`. + + Args: + stoppable (Stoppable): Device to be stopped + + Returns: + Iterator[MsgGenerator]: Bluesky messages + """ + return (yield from bps.stop(stoppable)) diff --git a/src/dodal/plans/__init__.py b/src/dodal/plans/__init__.py index fb402459694..46ac9e7f4d6 100644 --- a/src/dodal/plans/__init__.py +++ b/src/dodal/plans/__init__.py @@ -1,4 +1,4 @@ from .scanspec import spec_scan -from .wrapped import count +from .wrapped import count, grid_scan, rel_grid_scan, rel_scan, scan -__all__ = ["count", "spec_scan"] +__all__ = ["count", "grid_scan", "rel_grid_scan", "rel_scan", "scan", "spec_scan"] diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 48875c52353..66069922046 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -55,3 +55,111 @@ def count( metadata = metadata or {} metadata["shape"] = (num,) yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) + + +def scan( + detectors: Annotated[ + set[Readable] | list[Readable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + args: Annotated[ + tuple, + Field( + description="For one or more dimensions, 'motor1, start1, stop1, ..., " + "motorN, startN, stopN'. Motors can be any 'settable' object (motor, " + "temp controller, etc.)" + ), + ], + num: Annotated[int, Field(description="Number of points")], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one multi-motor trajectory. + Wraps bluesky.plans.scan(det, *args, num, md=metadata)""" + metadata = metadata or {} + metadata["shape"] = (num,) + yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) + + +def rel_scan( + detectors: Annotated[ + set[Readable] | list[Readable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + args: Annotated[ + tuple, + Field( + description="For one or more dimensions, 'motor1, start1, stop1, ..., " + "motorN, startN, stopN'. Motors can be any 'settable' object (motor, " + "temp controller, etc.)" + ), + ], + num: Annotated[int, Field(description="Number of points")], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one multi-motor trajectory, relative to current position. + Wraps bluesky.plans.rel_scan(det, *args, num, md=metadata)""" + metadata = metadata or {} + metadata["shape"] = (num,) + yield from bp.rel_scan(tuple(detectors), *args, num=num, md=metadata) + + +def grid_scan( + detectors: Annotated[ + set[Readable] | list[Readable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + args: Annotated[ + tuple, + Field( + description="Patterend like (motor1, start1, stop1, num1, ..., motorN, " + "startN, stopN, numN). The first motor is the 'slowest', the outer loop. " + "For all motors except the first motor, there is a 'snake' argument: a" + "boolean indicating whether to follow snake-like, winding trajectory or a" + "simple left to right." + ), + ], + snake_axes: list | bool | None = None, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over a mesh; each motor is on an independent trajectory. + Wraps bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata)""" + metadata = metadata or {} + yield from bp.grid_scan(tuple(detectors), *args, snake_axes=snake_axes, md=metadata) + + +def rel_grid_scan( + detectors: Annotated[ + set[Readable] | list[Readable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + args: Annotated[ + tuple, + Field( + description="Patterend like (motor1, start1, stop1, num1, ..., motorN, " + "startN, stopN, numN). The first motor is the 'slowest', the outer loop. " + "For all motors except the first motor, there is a 'snake' argument: a" + "boolean indicating whether to follow snake-like, winding trajectory or a" + "simple left to right." + ), + ], + snake_axes: list | bool | None = None, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over a mesh relative to current position. + Wraps bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata)""" + metadata = metadata or {} + yield from bp.rel_grid_scan( + tuple(detectors), *args, snake_axes=snake_axes, md=metadata + ) diff --git a/tests/plan_stubs/test_wrapped_stubs.py b/tests/plan_stubs/test_wrapped_stubs.py index 9af2e5d5dd7..cbc77ac9bd2 100644 --- a/tests/plan_stubs/test_wrapped_stubs.py +++ b/tests/plan_stubs/test_wrapped_stubs.py @@ -10,9 +10,11 @@ from dodal.plan_stubs.wrapped import ( move, move_relative, + rd, set_absolute, set_relative, sleep, + stop, wait, ) @@ -147,3 +149,11 @@ def test_wait_group_and_timeout(): assert list(wait("foo", 5.0)) == [ Msg("wait", group="foo", timeout=5.0, error_on_timeout=True, watch=_EMPTY) ] + + +def test_rd(x_axis: SimMotor): + assert list(rd(x_axis)) == [Msg("locate", obj=x_axis)] + + +def test_stop(x_axis: SimMotor): + assert list(stop(x_axis)) == [Msg("stop", obj=x_axis)] diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index b6b76dbc88c..bca9902dc53 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from typing import cast +from typing import Any, cast import pytest from bluesky.protocols import Readable @@ -17,7 +17,8 @@ ) from pydantic import ValidationError -from dodal.plans.wrapped import count +from dodal.devices.motors import Motor +from dodal.plans.wrapped import count, grid_scan, rel_grid_scan, rel_scan, scan @pytest.fixture @@ -157,3 +158,361 @@ def test_plan_produces_expected_datums( docs = documents_from_num.get("stream_datum") data_keys = [det.name, f"{det.name}-sum"] assert docs and len(docs) == len(data_keys) * length + + +@pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) +def test_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + num: int, +): + run_engine(scan(detectors={det}, args=(x_axis, x_start, x_stop), num=num)) + + +@pytest.mark.parametrize( + "x_start, x_stop, y_start, y_stop, num", ([0, 2, 2, 0, 5], [-1, 1, -1, 1, 3]) +) +def test_scan_with_two_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + y_axis: Motor, + y_start: Any, + y_stop: Any, + num: int, +): + run_engine( + scan( + detectors={det}, + args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + num=num, + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, y_start, y_stop, num", ([-1, 1, 2, 0, 0], [-1, 1, -1, 1, 3.5]) +) +def test_scan_fails_when_given_bad_info( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + y_axis: Motor, + y_start: Any, + y_stop: Any, + num: int, +): + with pytest.raises(ValueError): + run_engine( + scan( + detectors={det}, + args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + num=num, + ) + ) + + +@pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) +def test_rel_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + num: int, +): + run_engine(rel_scan(detectors={det}, args=(x_axis, x_start, x_stop), num=num)) + + +@pytest.mark.parametrize( + "x_start, x_stop, y_start, y_stop, num", ([0, 2, 2, 0, 5], [-1, 1, -1, 1, 3]) +) +def test_rel_scan_with_two_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + y_axis: Motor, + y_start: Any, + y_stop: Any, + num: int, +): + run_engine( + rel_scan( + detectors={det}, + args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + num=num, + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, y_start, y_stop, num", ([-1, 1, 2, 0, 0], [-1, 1, -1, 1, 3.5]) +) +def test_rel_scan_fails_when_given_bad_info( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + y_axis: Motor, + y_start: Any, + y_stop: Any, + num: int, +): + with pytest.raises(ValueError): + run_engine( + rel_scan( + detectors={det}, + args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + num=num, + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_grid_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_grid_scan_when_snaking( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + snake_axes=True, + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_grid_scan_when_snaking_subset_of_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + snake_axes=[x_axis], + ) + ) + + +def test_grid_scan_fails_when_snaking_slow_axis( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 3, x_axis, 0, 2, 3), + snake_axes=[y_axis], + ) + ) + + +def test_grid_scan_fails_when_given_length_of_zero( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(RuntimeError): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 0, x_axis, 0, 2, 3), + ) + ) + + +def test_grid_scan_fails_when_given_non_integer_length( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(TypeError): + run_engine( + grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 3.5, x_axis, 0, 2, 3), + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_rel_grid_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_rel_grid_scan_when_snaking( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + snake_axes=True, + ) + ) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), +) +def test_rel_grid_scan_when_snaking_subset_of_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_start: Any, + x_stop: Any, + x_num: int, + y_axis: Motor, + y_start: Any, + y_stop: Any, + y_num: int, +): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + snake_axes=[x_axis], + ) + ) + + +def test_rel_grid_scan_fails_when_snaking_slow_axis( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 3, x_axis, 0, 2, 3), + snake_axes=[y_axis], + ) + ) + + +def test_rel_grid_scan_fails_when_given_length_of_zero( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(RuntimeError): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 0, x_axis, 0, 2, 3), + ) + ) + + +def test_rel_grid_scan_fails_when_given_non_integer_length( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(TypeError): + run_engine( + rel_grid_scan( + detectors={det}, + args=(y_axis, 0, 2, 3.5, x_axis, 0, 2, 3), + ) + ) From 1b1bd158d6524cc34979c9b62d315b1275b82f44 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Wed, 12 Nov 2025 13:10:37 +0000 Subject: [PATCH 02/16] Make movable devices injectable --- src/dodal/plans/wrapped.py | 93 +++++++++++++++++++++++++------------ tests/plans/test_wrapped.py | 88 +++++++++++++++++++++++++++-------- 2 files changed, 131 insertions(+), 50 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 66069922046..0bf20c320a7 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -1,8 +1,10 @@ +import itertools from collections.abc import Sequence from typing import Annotated, Any import bluesky.plans as bp -from bluesky.protocols import Readable +from bluesky.protocols import Movable, Readable +from ophyd_async.core import AsyncReadable from pydantic import Field, NonNegativeFloat, validate_call from dodal.common import MsgGenerator @@ -27,7 +29,7 @@ @validate_call(config={"arbitrary_types_allowed": True}) def count( detectors: Annotated[ - set[Readable], + set[Readable] | set[AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -57,20 +59,39 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) +def _make_args(movables, params, num_params): + movables_len = len(movables) + params_len = len(params) + if params_len % movables_len != 0 or params_len % num_params != 0: + raise ValueError(f"params must contain {num_params} values for each movable") + + args = [] + it = iter(params) + param_chunks = iter(lambda: tuple(itertools.islice(it, num_params)), ()) + for movable, param_chunk in zip(movables, param_chunks, strict=False): + args.append(movable) + args.extend(param_chunk) + return args + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) def scan( detectors: Annotated[ - set[Readable] | list[Readable], + set[Readable] | set[AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - args: Annotated[ - tuple, + movables: Annotated[ + list[Movable], Field(description="One or more movable to move during the scan.") + ], + params: Annotated[ + list[float], Field( - description="For one or more dimensions, 'motor1, start1, stop1, ..., " - "motorN, startN, stopN'. Motors can be any 'settable' object (motor, " - "temp controller, etc.)" + description="Start and stop points for each movable, 'start1, stop1, ...," + "startN, stopN' for every movable in `movables`." ), ], num: Annotated[int, Field(description="Number of points")], @@ -80,23 +101,28 @@ def scan( Wraps bluesky.plans.scan(det, *args, num, md=metadata)""" metadata = metadata or {} metadata["shape"] = (num,) + args = _make_args(movables=movables, params=params, num_params=2) yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) def rel_scan( detectors: Annotated[ - set[Readable] | list[Readable], + set[Readable] | set[AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - args: Annotated[ - tuple, + movables: Annotated[ + list[Movable], Field(description="One or more movable to move during the scan.") + ], + params: Annotated[ + list[float], Field( - description="For one or more dimensions, 'motor1, start1, stop1, ..., " - "motorN, startN, stopN'. Motors can be any 'settable' object (motor, " - "temp controller, etc.)" + description="Start and stop points for each movable, 'start1, stop1, ...," + "startN, stopN' for every movable in `movables`." ), ], num: Annotated[int, Field(description="Number of points")], @@ -106,25 +132,28 @@ def rel_scan( Wraps bluesky.plans.rel_scan(det, *args, num, md=metadata)""" metadata = metadata or {} metadata["shape"] = (num,) + args = _make_args(movables=movables, params=params, num_params=2) yield from bp.rel_scan(tuple(detectors), *args, num=num, md=metadata) +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) def grid_scan( detectors: Annotated[ - set[Readable] | list[Readable], + set[Readable] | set[AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - args: Annotated[ - tuple, + movables: Annotated[ + list[Movable], Field(description="One or more movable to move during the scan.") + ], + params: Annotated[ + list[float | int], Field( - description="Patterend like (motor1, start1, stop1, num1, ..., motorN, " - "startN, stopN, numN). The first motor is the 'slowest', the outer loop. " - "For all motors except the first motor, there is a 'snake' argument: a" - "boolean indicating whether to follow snake-like, winding trajectory or a" - "simple left to right." + description="Start and stop points for each movable, 'start1, stop1, ...," + "startN, stopN' for every movable in `movables`." ), ], snake_axes: list | bool | None = None, @@ -133,25 +162,28 @@ def grid_scan( """Scan over a mesh; each motor is on an independent trajectory. Wraps bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata)""" metadata = metadata or {} + args = _make_args(movables=movables, params=params, num_params=3) yield from bp.grid_scan(tuple(detectors), *args, snake_axes=snake_axes, md=metadata) +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) def rel_grid_scan( detectors: Annotated[ - set[Readable] | list[Readable], + set[Readable] | set[AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - args: Annotated[ - tuple, + movables: Annotated[ + list[Movable], Field(description="One or more movable to move during the scan.") + ], + params: Annotated[ + list[float | int], Field( - description="Patterend like (motor1, start1, stop1, num1, ..., motorN, " - "startN, stopN, numN). The first motor is the 'slowest', the outer loop. " - "For all motors except the first motor, there is a 'snake' argument: a" - "boolean indicating whether to follow snake-like, winding trajectory or a" - "simple left to right." + description="Start and stop points for each movable, 'start1, stop1, ...," + "startN, stopN' for every movable in `movables`." ), ], snake_axes: list | bool | None = None, @@ -160,6 +192,7 @@ def rel_grid_scan( """Scan over a mesh relative to current position. Wraps bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata)""" metadata = metadata or {} + args = _make_args(movables=movables, params=params, num_params=3) yield from bp.rel_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata ) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index bca9902dc53..eccbd84c229 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -18,7 +18,14 @@ from pydantic import ValidationError from dodal.devices.motors import Motor -from dodal.plans.wrapped import count, grid_scan, rel_grid_scan, rel_scan, scan +from dodal.plans.wrapped import ( + _make_args, + count, + grid_scan, + rel_grid_scan, + rel_scan, + scan, +) @pytest.fixture @@ -160,6 +167,20 @@ def test_plan_produces_expected_datums( assert docs and len(docs) == len(data_keys) * length +@pytest.mark.parametrize( + "num_params, params", ([2, [1, 2, 3, 4]], [3, [1, 2, 3, 3, 4, 3]]) +) +def test_make_args(x_axis: Motor, y_axis: Motor, num_params: int, params: list[float]): + movables = [x_axis, y_axis] + args = _make_args(movables=movables, params=params, num_params=num_params) + print(args) + assert len(args) == len(movables) + len(params) + assert args[0] == x_axis + assert args[(num_params + 1)] == y_axis + assert args[1] == 1 + assert args[(num_params + 2)] == 3 + + @pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) def test_scan( run_engine: RunEngine, @@ -169,7 +190,9 @@ def test_scan( x_stop: Any, num: int, ): - run_engine(scan(detectors={det}, args=(x_axis, x_start, x_stop), num=num)) + run_engine( + scan(detectors={det}, movables=[x_axis], params=[x_start, x_stop], num=num) + ) @pytest.mark.parametrize( @@ -189,12 +212,22 @@ def test_scan_with_two_axes( run_engine( scan( detectors={det}, - args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + movables=[x_axis, y_axis], + params=[x_start, x_stop, y_start, y_stop], num=num, ) ) +def test_scan_fails_when_given_wrong_number_of_params( + run_engine: RunEngine, det: StandardDetector, x_axis: Motor, y_axis: Motor +): + with pytest.raises(ValueError): + run_engine( + scan(detectors={det}, movables=[x_axis, y_axis], params=[0, 1, 2], num=3) + ) + + @pytest.mark.parametrize( "x_start, x_stop, y_start, y_stop, num", ([-1, 1, 2, 0, 0], [-1, 1, -1, 1, 3.5]) ) @@ -213,7 +246,8 @@ def test_scan_fails_when_given_bad_info( run_engine( scan( detectors={det}, - args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + movables=[x_axis, y_axis], + params=[x_start, x_stop, y_start, y_stop], num=num, ) ) @@ -228,7 +262,9 @@ def test_rel_scan( x_stop: Any, num: int, ): - run_engine(rel_scan(detectors={det}, args=(x_axis, x_start, x_stop), num=num)) + run_engine( + rel_scan(detectors={det}, movables=[x_axis], params=[x_start, x_stop], num=num) + ) @pytest.mark.parametrize( @@ -248,7 +284,8 @@ def test_rel_scan_with_two_axes( run_engine( rel_scan( detectors={det}, - args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + movables=[x_axis, y_axis], + params=[x_start, x_stop, y_start, y_stop], num=num, ) ) @@ -272,7 +309,8 @@ def test_rel_scan_fails_when_given_bad_info( run_engine( rel_scan( detectors={det}, - args=(x_axis, x_start, x_stop, y_axis, y_start, y_stop), + movables=[x_axis, y_axis], + params=[x_start, x_stop, y_start, y_stop], num=num, ) ) @@ -297,7 +335,8 @@ def test_grid_scan( run_engine( grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], ) ) @@ -321,7 +360,8 @@ def test_grid_scan_when_snaking( run_engine( grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=True, ) ) @@ -346,7 +386,8 @@ def test_grid_scan_when_snaking_subset_of_axes( run_engine( grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=[x_axis], ) ) @@ -362,7 +403,8 @@ def test_grid_scan_fails_when_snaking_slow_axis( run_engine( grid_scan( detectors={det}, - args=(y_axis, 0, 2, 3, x_axis, 0, 2, 3), + movables=[y_axis, x_axis], + params=[0, 2, 3, 0, 2, 3], snake_axes=[y_axis], ) ) @@ -378,7 +420,8 @@ def test_grid_scan_fails_when_given_length_of_zero( run_engine( grid_scan( detectors={det}, - args=(y_axis, 0, 2, 0, x_axis, 0, 2, 3), + movables=[y_axis, x_axis], + params=[0, 2, 0, 0, 2, 3], ) ) @@ -393,7 +436,8 @@ def test_grid_scan_fails_when_given_non_integer_length( run_engine( grid_scan( detectors={det}, - args=(y_axis, 0, 2, 3.5, x_axis, 0, 2, 3), + movables=[y_axis, x_axis], + params=[0, 2, 3.5, 0, 2, 3], ) ) @@ -417,7 +461,8 @@ def test_rel_grid_scan( run_engine( rel_grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], ) ) @@ -441,7 +486,8 @@ def test_rel_grid_scan_when_snaking( run_engine( rel_grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=True, ) ) @@ -466,7 +512,8 @@ def test_rel_grid_scan_when_snaking_subset_of_axes( run_engine( rel_grid_scan( detectors={det}, - args=(y_axis, y_start, y_stop, y_num, x_axis, x_start, x_stop, x_num), + movables=[y_axis, x_axis], + params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=[x_axis], ) ) @@ -482,7 +529,8 @@ def test_rel_grid_scan_fails_when_snaking_slow_axis( run_engine( rel_grid_scan( detectors={det}, - args=(y_axis, 0, 2, 3, x_axis, 0, 2, 3), + movables=[y_axis, x_axis], + params=[0, 2, 3, 0, 2, 3], snake_axes=[y_axis], ) ) @@ -498,7 +546,8 @@ def test_rel_grid_scan_fails_when_given_length_of_zero( run_engine( rel_grid_scan( detectors={det}, - args=(y_axis, 0, 2, 0, x_axis, 0, 2, 3), + movables=[y_axis, x_axis], + params=[0, 2, 0, 0, 2, 3], ) ) @@ -512,7 +561,6 @@ def test_rel_grid_scan_fails_when_given_non_integer_length( with pytest.raises(TypeError): run_engine( rel_grid_scan( - detectors={det}, - args=(y_axis, 0, 2, 3.5, x_axis, 0, 2, 3), + detectors={det}, movables=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3] ) ) From e91bd12fddc734aa5406574972e0bedcc47aa4f0 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Wed, 12 Nov 2025 13:25:32 +0000 Subject: [PATCH 03/16] Fix detector types --- src/dodal/plans/wrapped.py | 10 +++++----- tests/plans/test_wrapped.py | 3 ++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 0bf20c320a7..d80f6d00bfd 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -29,7 +29,7 @@ @validate_call(config={"arbitrary_types_allowed": True}) def count( detectors: Annotated[ - set[Readable] | set[AsyncReadable], + set[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -78,7 +78,7 @@ def _make_args(movables, params, num_params): @validate_call(config={"arbitrary_types_allowed": True}) def scan( detectors: Annotated[ - set[Readable] | set[AsyncReadable], + set[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -109,7 +109,7 @@ def scan( @validate_call(config={"arbitrary_types_allowed": True}) def rel_scan( detectors: Annotated[ - set[Readable] | set[AsyncReadable], + set[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -140,7 +140,7 @@ def rel_scan( @validate_call(config={"arbitrary_types_allowed": True}) def grid_scan( detectors: Annotated[ - set[Readable] | set[AsyncReadable], + set[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -170,7 +170,7 @@ def grid_scan( @validate_call(config={"arbitrary_types_allowed": True}) def rel_grid_scan( detectors: Annotated[ - set[Readable] | set[AsyncReadable], + set[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index eccbd84c229..21017498dc4 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -13,6 +13,7 @@ StreamResource, ) from ophyd_async.core import ( + AsyncReadable, StandardDetector, ) from pydantic import ValidationError @@ -63,7 +64,7 @@ def test_count_delay_validation(det: StandardDetector, run_engine: RunEngine): def test_count_detectors_validation(run_engine: RunEngine): - args: dict[str, set[Readable]] = { + args: dict[str, set[Readable | AsyncReadable]] = { # No device to read "Set should have at least 1 item after validation, not 0": set(), # Not Readable From b167b669919cf8e0ea482094f8e07a4fd9d0567c Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Wed, 12 Nov 2025 16:14:53 +0000 Subject: [PATCH 04/16] add list_scan, rel_list_scan, list_grid_scan, rel_list_grid_scan --- src/dodal/plans/wrapped.py | 189 ++++++++++++++++++---- tests/plans/test_wrapped.py | 310 +++++++++++++++++++++++++++++++----- 2 files changed, 429 insertions(+), 70 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index d80f6d00bfd..638dbbc984d 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -8,6 +8,7 @@ from pydantic import Field, NonNegativeFloat, validate_call from dodal.common import MsgGenerator +from dodal.devices.motors import Motor from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator """This module wraps plan(s) from bluesky.plans until required handling for them is @@ -29,7 +30,7 @@ @validate_call(config={"arbitrary_types_allowed": True}) def count( detectors: Annotated[ - set[Readable | AsyncReadable], + Sequence[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, @@ -59,16 +60,20 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) -def _make_args(movables, params, num_params): - movables_len = len(movables) +def _make_args( + movers: Sequence[Movable | Motor], + params: list[Any] | Sequence[Any], + num_params: int, +): + movers_len = len(movers) params_len = len(params) - if params_len % movables_len != 0 or params_len % num_params != 0: + if params_len % movers_len != 0 or params_len % num_params != 0: raise ValueError(f"params must contain {num_params} values for each movable") args = [] it = iter(params) param_chunks = iter(lambda: tuple(itertools.islice(it, num_params)), ()) - for movable, param_chunk in zip(movables, param_chunks, strict=False): + for movable, param_chunk in zip(movers, param_chunks, strict=False): args.append(movable) args.extend(param_chunk) return args @@ -78,20 +83,21 @@ def _make_args(movables, params, num_params): @validate_call(config={"arbitrary_types_allowed": True}) def scan( detectors: Annotated[ - set[Readable | AsyncReadable], + Sequence[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - movables: Annotated[ - list[Movable], Field(description="One or more movable to move during the scan.") + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), ], params: Annotated[ list[float], Field( description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movables`." + "startN, stopN' for every movable in `movers`." ), ], num: Annotated[int, Field(description="Number of points")], @@ -101,7 +107,7 @@ def scan( Wraps bluesky.plans.scan(det, *args, num, md=metadata)""" metadata = metadata or {} metadata["shape"] = (num,) - args = _make_args(movables=movables, params=params, num_params=2) + args = _make_args(movers=movers, params=params, num_params=2) yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) @@ -109,20 +115,21 @@ def scan( @validate_call(config={"arbitrary_types_allowed": True}) def rel_scan( detectors: Annotated[ - set[Readable | AsyncReadable], + Sequence[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - movables: Annotated[ - list[Movable], Field(description="One or more movable to move during the scan.") + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), ], params: Annotated[ list[float], Field( description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movables`." + "startN, stopN' for every movable in `movers`." ), ], num: Annotated[int, Field(description="Number of points")], @@ -132,7 +139,7 @@ def rel_scan( Wraps bluesky.plans.rel_scan(det, *args, num, md=metadata)""" metadata = metadata or {} metadata["shape"] = (num,) - args = _make_args(movables=movables, params=params, num_params=2) + args = _make_args(movers=movers, params=params, num_params=2) yield from bp.rel_scan(tuple(detectors), *args, num=num, md=metadata) @@ -140,20 +147,21 @@ def rel_scan( @validate_call(config={"arbitrary_types_allowed": True}) def grid_scan( detectors: Annotated[ - set[Readable | AsyncReadable], + Sequence[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - movables: Annotated[ - list[Movable], Field(description="One or more movable to move during the scan.") + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), ], params: Annotated[ - list[float | int], + Sequence[float | int], Field( description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movables`." + "startN, stopN' for every movable in `movers`." ), ], snake_axes: list | bool | None = None, @@ -162,7 +170,7 @@ def grid_scan( """Scan over a mesh; each motor is on an independent trajectory. Wraps bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata)""" metadata = metadata or {} - args = _make_args(movables=movables, params=params, num_params=3) + args = _make_args(movers=movers, params=params, num_params=3) yield from bp.grid_scan(tuple(detectors), *args, snake_axes=snake_axes, md=metadata) @@ -170,20 +178,21 @@ def grid_scan( @validate_call(config={"arbitrary_types_allowed": True}) def rel_grid_scan( detectors: Annotated[ - set[Readable | AsyncReadable], + Sequence[Readable | AsyncReadable], Field( description="Set of readable devices, will take a reading at each point", min_length=1, ), ], - movables: Annotated[ - list[Movable], Field(description="One or more movable to move during the scan.") + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), ], params: Annotated[ - list[float | int], + Sequence[float | int], Field( description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movables`." + "startN, stopN' for every movable in `movers`." ), ], snake_axes: list | bool | None = None, @@ -192,7 +201,133 @@ def rel_grid_scan( """Scan over a mesh relative to current position. Wraps bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata)""" metadata = metadata or {} - args = _make_args(movables=movables, params=params, num_params=3) + args = _make_args(movers=movers, params=params, num_params=3) yield from bp.rel_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata ) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def list_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[list[Any]], + Field( + description="List of points for each movable, '[point1, point2, ..., ], " + "[point1, point2, ...], ...' for every movable in `movers`." + ), + ], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one or more variables in steps simultaneously. + Wraps bluesky.plans.list_scan(det, *args, md=metadata).""" + metadata = metadata or {} + args = _make_args(movers=movers, params=params, num_params=1) + yield from bp.list_scan(tuple(detectors), *args, md=metadata) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def rel_list_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[list[Any]], + Field( + description="List of points for each movable, '[point1, point2, ..., ], " + "[point1, point2, ...], ...' for every movable in `movers`." + ), + ], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one or more variables simultaneously relative to current position. + Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata).""" + metadata = metadata or {} + args = _make_args(movers=movers, params=params, num_params=1) + yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def list_grid_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[list[Any]], + Field( + description="List of points for each movable, '[point1, point2, ..., ], " + "[point1, point2, ...], ...' for every movable in `movers`." + ), + ], + snake_axes: bool = False, # Currently specifying axes to snake is not supported + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one or more variables for each given point on independent trajectories. + Wraps bluesky.plans.list_grid_scan(det, *args, md=metadata).""" + metadata = metadata or {} + args = _make_args(movers=movers, params=params, num_params=1) + yield from bp.list_grid_scan( + tuple(detectors), *args, snake_axes=snake_axes, md=metadata + ) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def rel_list_grid_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[list[Any]], + Field( + description="List of points for each movable, '[point1, point2, ..., ], " + "[point1, point2, ...], ...' for every movable in `movers`." + ), + ], + snake_axes: bool = False, # Currently specifying axes to snake is not supported + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over some variables for each given point relative to current position. + Wraps bluesky.plans.rel_list_grid_scan(det, *args, md=metadata).""" + metadata = metadata or {} + args = _make_args(movers=movers, params=params, num_params=1) + yield from bp.rel_list_grid_scan( + tuple(detectors), *args, snake_axes=snake_axes, md=metadata + ) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 21017498dc4..f643eef98e4 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -23,7 +23,11 @@ _make_args, count, grid_scan, + list_grid_scan, + list_scan, rel_grid_scan, + rel_list_grid_scan, + rel_list_scan, rel_scan, scan, ) @@ -35,7 +39,7 @@ def documents_from_num( ) -> dict[str, list[Document]]: docs: dict[str, list[Document]] = {} run_engine( - count({det}, num=request.param), + count([det], num=request.param), lambda name, doc: docs.setdefault(name, []).append(doc), ) return docs @@ -59,16 +63,16 @@ def test_count_delay_validation(det: StandardDetector, run_engine: RunEngine): } for delay, reason in args.items(): with pytest.raises((ValidationError, AssertionError), match=reason): - run_engine(count({det}, num=3, delay=delay)) + run_engine(count([det], num=3, delay=delay)) print(delay) def test_count_detectors_validation(run_engine: RunEngine): - args: dict[str, set[Readable | AsyncReadable]] = { + args: dict[str, Sequence[Readable | AsyncReadable]] = { # No device to read - "Set should have at least 1 item after validation, not 0": set(), + "1 validation error for count": set(), # Not Readable - "Input should be an instance of Readable": set("foo"), # type: ignore + "Input should be an instance of Sequence": set("foo"), # type: ignore } for reason, dets in args.items(): with pytest.raises(ValidationError, match=reason): @@ -83,7 +87,7 @@ def test_count_num_validation(det: StandardDetector, run_engine: RunEngine): } for num, reason in args.items(): with pytest.raises(ValidationError, match=reason): - run_engine(count({det}, num=num)) + run_engine(count([det], num=num)) @pytest.mark.parametrize( @@ -172,16 +176,28 @@ def test_plan_produces_expected_datums( "num_params, params", ([2, [1, 2, 3, 4]], [3, [1, 2, 3, 3, 4, 3]]) ) def test_make_args(x_axis: Motor, y_axis: Motor, num_params: int, params: list[float]): - movables = [x_axis, y_axis] - args = _make_args(movables=movables, params=params, num_params=num_params) + movers = [x_axis, y_axis] + args = _make_args(movers=movers, params=params, num_params=num_params) print(args) - assert len(args) == len(movables) + len(params) + assert len(args) == len(movers) + len(params) assert args[0] == x_axis assert args[(num_params + 1)] == y_axis assert args[1] == 1 assert args[(num_params + 2)] == 3 +def test_make_args_when_given_lists(x_axis: Motor, y_axis: Motor): + args = _make_args( + movers=[x_axis, y_axis], params=[[1, 2, 3, 4], [3, 4, 5, 6]], num_params=1 + ) + print(args) + assert len(args) == 4 + assert args[0] == x_axis + assert args[2] == y_axis + assert args[1][0] == 1 + assert args[3][0] == 3 + + @pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) def test_scan( run_engine: RunEngine, @@ -192,7 +208,7 @@ def test_scan( num: int, ): run_engine( - scan(detectors={det}, movables=[x_axis], params=[x_start, x_stop], num=num) + scan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) ) @@ -212,8 +228,8 @@ def test_scan_with_two_axes( ): run_engine( scan( - detectors={det}, - movables=[x_axis, y_axis], + detectors=[det], + movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], num=num, ) @@ -225,7 +241,7 @@ def test_scan_fails_when_given_wrong_number_of_params( ): with pytest.raises(ValueError): run_engine( - scan(detectors={det}, movables=[x_axis, y_axis], params=[0, 1, 2], num=3) + scan(detectors=[det], movers=[x_axis, y_axis], params=[0, 1, 2], num=3) ) @@ -246,8 +262,8 @@ def test_scan_fails_when_given_bad_info( with pytest.raises(ValueError): run_engine( scan( - detectors={det}, - movables=[x_axis, y_axis], + detectors=[det], + movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], num=num, ) @@ -264,7 +280,7 @@ def test_rel_scan( num: int, ): run_engine( - rel_scan(detectors={det}, movables=[x_axis], params=[x_start, x_stop], num=num) + rel_scan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) ) @@ -284,8 +300,8 @@ def test_rel_scan_with_two_axes( ): run_engine( rel_scan( - detectors={det}, - movables=[x_axis, y_axis], + detectors=[det], + movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], num=num, ) @@ -309,8 +325,8 @@ def test_rel_scan_fails_when_given_bad_info( with pytest.raises(ValueError): run_engine( rel_scan( - detectors={det}, - movables=[x_axis, y_axis], + detectors=[det], + movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], num=num, ) @@ -335,8 +351,8 @@ def test_grid_scan( ): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], ) ) @@ -360,8 +376,8 @@ def test_grid_scan_when_snaking( ): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=True, ) @@ -386,8 +402,8 @@ def test_grid_scan_when_snaking_subset_of_axes( ): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=[x_axis], ) @@ -403,8 +419,8 @@ def test_grid_scan_fails_when_snaking_slow_axis( with pytest.raises(ValueError): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[0, 2, 3, 0, 2, 3], snake_axes=[y_axis], ) @@ -420,8 +436,8 @@ def test_grid_scan_fails_when_given_length_of_zero( with pytest.raises(RuntimeError): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[0, 2, 0, 0, 2, 3], ) ) @@ -436,8 +452,8 @@ def test_grid_scan_fails_when_given_non_integer_length( with pytest.raises(TypeError): run_engine( grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3], ) ) @@ -461,8 +477,8 @@ def test_rel_grid_scan( ): run_engine( rel_grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], ) ) @@ -486,8 +502,8 @@ def test_rel_grid_scan_when_snaking( ): run_engine( rel_grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=True, ) @@ -512,8 +528,8 @@ def test_rel_grid_scan_when_snaking_subset_of_axes( ): run_engine( rel_grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], snake_axes=[x_axis], ) @@ -529,8 +545,8 @@ def test_rel_grid_scan_fails_when_snaking_slow_axis( with pytest.raises(ValueError): run_engine( rel_grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[0, 2, 3, 0, 2, 3], snake_axes=[y_axis], ) @@ -546,8 +562,8 @@ def test_rel_grid_scan_fails_when_given_length_of_zero( with pytest.raises(RuntimeError): run_engine( rel_grid_scan( - detectors={det}, - movables=[y_axis, x_axis], + detectors=[det], + movers=[y_axis, x_axis], params=[0, 2, 0, 0, 2, 3], ) ) @@ -562,6 +578,214 @@ def test_rel_grid_scan_fails_when_given_non_integer_length( with pytest.raises(TypeError): run_engine( rel_grid_scan( - detectors={det}, movables=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3] + detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3] + ) + ) + + +@pytest.mark.parametrize("x_list", ([[0, 1, 2, 3]], [[1.1, 2.2, 3.3]])) +def test_list_scan( + run_engine: RunEngine, det: StandardDetector, x_axis: Motor, x_list: Any +): + run_engine(list_scan(detectors=[det], movers=[x_axis], params=x_list)) + + +@pytest.mark.parametrize( + "x_list, y_list", + ( + [[3, 2, 1], [1, 2, 3]], + [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], + ), +) +def test_list_scan_with_two_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_list: list, + y_axis: Motor, + y_list: list, +): + run_engine( + list_scan(detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list]) + ) + + +def test_list_scan_with_two_axes_fails_when_given_differnt_list_lengths( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine( + list_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], + ) + ) + + +@pytest.mark.parametrize("x_list", ([[0, 1, 2, 3]], [[1.1, 2.2, 3.3]])) +def test_rel_list_scan( + run_engine: RunEngine, det: StandardDetector, x_axis: Motor, x_list: Any +): + run_engine(rel_list_scan(detectors=[det], movers=[x_axis], params=x_list)) + + +@pytest.mark.parametrize( + "x_list, y_list", + ( + [[3, 2, 1], [1, 2, 3]], + [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], + ), +) +def test_rel_list_scan_with_two_axes( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_list: list, + y_axis: Motor, + y_list: list, +): + run_engine( + rel_list_scan(detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list]) + ) + + +def test_rel_list_scan_with_two_axes_fails_when_given_differnt_list_lengths( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine( + rel_list_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], + ) + ) + + +@pytest.mark.parametrize( + "x_list, y_list", + ( + [[3, 2, 1], [1, 2, 3]], + [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], + ), +) +def test_list_grid_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_list: list, + y_axis: Motor, + y_list: list, +): + run_engine( + list_grid_scan( + detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list] + ) + ) + + +def test_list_grid_scan_when_given_differnt_list_lengths( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + run_engine( + list_grid_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], + ) + ) + + +def test_list_grid_scan_when_given_bad_info( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(TypeError): + run_engine( + list_grid_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], ["one", 2, 3, 4, 5]], + ) + ) + + +@pytest.mark.parametrize( + "x_list, y_list", + ( + [[3, 2, 1], [1, 2, 3]], + [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], + ), +) +def test_rel_list_grid_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + x_list: list, + y_axis: Motor, + y_list: list, +): + run_engine( + rel_list_grid_scan( + detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list] + ) + ) + + +def test_rel_list_grid_scan_with_two_axes_when_snaking( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + run_engine( + rel_list_grid_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]], + snake_axes=True, + ) + ) + + +def test_rel_list_grid_scan_when_given_differnt_list_lengths( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + run_engine( + rel_list_grid_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], + ) + ) + + +def test_rel_list_grid_scan_when_given_bad_info( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(TypeError): + run_engine( + list_grid_scan( + detectors=[det], + movers=[x_axis, y_axis], + params=[[1, 2, 3, 4, 5], ["one", 2, 3, 4, 5]], ) ) From 395c7e2cf90748fa8ac91952562925a467623db0 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Thu, 13 Nov 2025 18:00:27 +0000 Subject: [PATCH 05/16] rename scans and fix adsim test --- src/dodal/plans/__init__.py | 25 +++++++- src/dodal/plans/wrapped.py | 12 ++-- system_tests/test_adsim.py | 2 +- tests/plans/test_wrapped.py | 116 ++++++++++++++++++------------------ 4 files changed, 88 insertions(+), 67 deletions(-) diff --git a/src/dodal/plans/__init__.py b/src/dodal/plans/__init__.py index 46ac9e7f4d6..089e5f83779 100644 --- a/src/dodal/plans/__init__.py +++ b/src/dodal/plans/__init__.py @@ -1,4 +1,25 @@ from .scanspec import spec_scan -from .wrapped import count, grid_scan, rel_grid_scan, rel_scan, scan +from .wrapped import ( + count, + grid_num_rscan, + grid_num_scan, + list_grid_rscan, + list_grid_scan, + list_rscan, + list_scan, + num_rscan, + num_scan, +) -__all__ = ["count", "grid_scan", "rel_grid_scan", "rel_scan", "scan", "spec_scan"] +__all__ = [ + "count", + "grid_num_rscan", + "grid_num_scan", + "list_grid_rscan", + "list_grid_scan", + "list_rscan", + "list_scan", + "num_rscan", + "num_scan", + "spec_scan", +] diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 638dbbc984d..3e95562cbdb 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -81,7 +81,7 @@ def _make_args( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def scan( +def num_scan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( @@ -113,7 +113,7 @@ def scan( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def rel_scan( +def num_rscan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( @@ -145,7 +145,7 @@ def rel_scan( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def grid_scan( +def grid_num_scan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( @@ -176,7 +176,7 @@ def grid_scan( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def rel_grid_scan( +def grid_num_rscan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( @@ -239,7 +239,7 @@ def list_scan( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def rel_list_scan( +def list_rscan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( @@ -302,7 +302,7 @@ def list_grid_scan( @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) -def rel_list_grid_scan( +def list_grid_rscan( detectors: Annotated[ Sequence[Readable | AsyncReadable], Field( diff --git a/system_tests/test_adsim.py b/system_tests/test_adsim.py index 665a4175c96..9c5bf8af73f 100644 --- a/system_tests/test_adsim.py +++ b/system_tests/test_adsim.py @@ -87,7 +87,7 @@ def documents_from_num( ) -> dict[str, list[DocumentType]]: docs: dict[str, list[DocumentType]] = {} run_engine( - count({det}, num=request.param), + count([det], num=request.param), lambda name, doc: docs.setdefault(name, []).append(doc), ) return docs diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index f643eef98e4..bda4b926539 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -22,14 +22,14 @@ from dodal.plans.wrapped import ( _make_args, count, - grid_scan, + grid_num_rscan, + grid_num_scan, + list_grid_rscan, list_grid_scan, + list_rscan, list_scan, - rel_grid_scan, - rel_list_grid_scan, - rel_list_scan, - rel_scan, - scan, + num_rscan, + num_scan, ) @@ -199,7 +199,7 @@ def test_make_args_when_given_lists(x_axis: Motor, y_axis: Motor): @pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) -def test_scan( +def test_num_scan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -208,14 +208,14 @@ def test_scan( num: int, ): run_engine( - scan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) + num_scan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) ) @pytest.mark.parametrize( "x_start, x_stop, y_start, y_stop, num", ([0, 2, 2, 0, 5], [-1, 1, -1, 1, 3]) ) -def test_scan_with_two_axes( +def test_num_scan_with_two_axes( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -227,7 +227,7 @@ def test_scan_with_two_axes( num: int, ): run_engine( - scan( + num_scan( detectors=[det], movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], @@ -236,19 +236,19 @@ def test_scan_with_two_axes( ) -def test_scan_fails_when_given_wrong_number_of_params( +def test_num_scan_fails_when_given_wrong_number_of_params( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, y_axis: Motor ): with pytest.raises(ValueError): run_engine( - scan(detectors=[det], movers=[x_axis, y_axis], params=[0, 1, 2], num=3) + num_scan(detectors=[det], movers=[x_axis, y_axis], params=[0, 1, 2], num=3) ) @pytest.mark.parametrize( "x_start, x_stop, y_start, y_stop, num", ([-1, 1, 2, 0, 0], [-1, 1, -1, 1, 3.5]) ) -def test_scan_fails_when_given_bad_info( +def test_num_scan_fails_when_given_bad_info( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -261,7 +261,7 @@ def test_scan_fails_when_given_bad_info( ): with pytest.raises(ValueError): run_engine( - scan( + num_scan( detectors=[det], movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], @@ -271,7 +271,7 @@ def test_scan_fails_when_given_bad_info( @pytest.mark.parametrize("x_start, x_stop, num", ([0, 2, 5], [1, -1, 3])) -def test_rel_scan( +def test_num_rscan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -280,14 +280,14 @@ def test_rel_scan( num: int, ): run_engine( - rel_scan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) + num_rscan(detectors=[det], movers=[x_axis], params=[x_start, x_stop], num=num) ) @pytest.mark.parametrize( "x_start, x_stop, y_start, y_stop, num", ([0, 2, 2, 0, 5], [-1, 1, -1, 1, 3]) ) -def test_rel_scan_with_two_axes( +def test_num_rscan_with_two_axes( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -299,7 +299,7 @@ def test_rel_scan_with_two_axes( num: int, ): run_engine( - rel_scan( + num_rscan( detectors=[det], movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], @@ -311,7 +311,7 @@ def test_rel_scan_with_two_axes( @pytest.mark.parametrize( "x_start, x_stop, y_start, y_stop, num", ([-1, 1, 2, 0, 0], [-1, 1, -1, 1, 3.5]) ) -def test_rel_scan_fails_when_given_bad_info( +def test_num_rscan_fails_when_given_bad_info( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -324,7 +324,7 @@ def test_rel_scan_fails_when_given_bad_info( ): with pytest.raises(ValueError): run_engine( - rel_scan( + num_rscan( detectors=[det], movers=[x_axis, y_axis], params=[x_start, x_stop, y_start, y_stop], @@ -337,7 +337,7 @@ def test_rel_scan_fails_when_given_bad_info( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_grid_scan( +def test_grid_num_scan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -350,7 +350,7 @@ def test_grid_scan( y_num: int, ): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -362,7 +362,7 @@ def test_grid_scan( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_grid_scan_when_snaking( +def test_grid_num_scan_when_snaking( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -375,7 +375,7 @@ def test_grid_scan_when_snaking( y_num: int, ): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -388,7 +388,7 @@ def test_grid_scan_when_snaking( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_grid_scan_when_snaking_subset_of_axes( +def test_grid_num_scan_when_snaking_subset_of_axes( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -401,7 +401,7 @@ def test_grid_scan_when_snaking_subset_of_axes( y_num: int, ): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -410,7 +410,7 @@ def test_grid_scan_when_snaking_subset_of_axes( ) -def test_grid_scan_fails_when_snaking_slow_axis( +def test_grid_num_scan_fails_when_snaking_slow_axis( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -418,7 +418,7 @@ def test_grid_scan_fails_when_snaking_slow_axis( ): with pytest.raises(ValueError): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 3, 0, 2, 3], @@ -427,7 +427,7 @@ def test_grid_scan_fails_when_snaking_slow_axis( ) -def test_grid_scan_fails_when_given_length_of_zero( +def test_grid_num_scan_fails_when_given_length_of_zero( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -435,7 +435,7 @@ def test_grid_scan_fails_when_given_length_of_zero( ): with pytest.raises(RuntimeError): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 0, 0, 2, 3], @@ -443,7 +443,7 @@ def test_grid_scan_fails_when_given_length_of_zero( ) -def test_grid_scan_fails_when_given_non_integer_length( +def test_grid_num_scan_fails_when_given_non_integer_length( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -451,7 +451,7 @@ def test_grid_scan_fails_when_given_non_integer_length( ): with pytest.raises(TypeError): run_engine( - grid_scan( + grid_num_scan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3], @@ -463,7 +463,7 @@ def test_grid_scan_fails_when_given_non_integer_length( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_rel_grid_scan( +def test_grid_num_rscan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -476,7 +476,7 @@ def test_rel_grid_scan( y_num: int, ): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -488,7 +488,7 @@ def test_rel_grid_scan( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_rel_grid_scan_when_snaking( +def test_grid_num_rscan_when_snaking( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -501,7 +501,7 @@ def test_rel_grid_scan_when_snaking( y_num: int, ): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -514,7 +514,7 @@ def test_rel_grid_scan_when_snaking( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([0, 2, 3, 0, 2, 3], [-1, 1, 5, 1, -1, 5]), ) -def test_rel_grid_scan_when_snaking_subset_of_axes( +def test_grid_num_rscan_when_snaking_subset_of_axes( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -527,7 +527,7 @@ def test_rel_grid_scan_when_snaking_subset_of_axes( y_num: int, ): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[y_start, y_stop, y_num, x_start, x_stop, x_num], @@ -536,7 +536,7 @@ def test_rel_grid_scan_when_snaking_subset_of_axes( ) -def test_rel_grid_scan_fails_when_snaking_slow_axis( +def test_grid_num_rscan_fails_when_snaking_slow_axis( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -544,7 +544,7 @@ def test_rel_grid_scan_fails_when_snaking_slow_axis( ): with pytest.raises(ValueError): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 3, 0, 2, 3], @@ -553,7 +553,7 @@ def test_rel_grid_scan_fails_when_snaking_slow_axis( ) -def test_rel_grid_scan_fails_when_given_length_of_zero( +def test_grid_num_rscan_fails_when_given_length_of_zero( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -561,7 +561,7 @@ def test_rel_grid_scan_fails_when_given_length_of_zero( ): with pytest.raises(RuntimeError): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 0, 0, 2, 3], @@ -569,7 +569,7 @@ def test_rel_grid_scan_fails_when_given_length_of_zero( ) -def test_rel_grid_scan_fails_when_given_non_integer_length( +def test_grid_num_rscan_fails_when_given_non_integer_length( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -577,7 +577,7 @@ def test_rel_grid_scan_fails_when_given_non_integer_length( ): with pytest.raises(TypeError): run_engine( - rel_grid_scan( + grid_num_rscan( detectors=[det], movers=[y_axis, x_axis], params=[0, 2, 3.5, 0, 2, 3] ) ) @@ -627,10 +627,10 @@ def test_list_scan_with_two_axes_fails_when_given_differnt_list_lengths( @pytest.mark.parametrize("x_list", ([[0, 1, 2, 3]], [[1.1, 2.2, 3.3]])) -def test_rel_list_scan( +def test_list_rscan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, x_list: Any ): - run_engine(rel_list_scan(detectors=[det], movers=[x_axis], params=x_list)) + run_engine(list_rscan(detectors=[det], movers=[x_axis], params=x_list)) @pytest.mark.parametrize( @@ -640,7 +640,7 @@ def test_rel_list_scan( [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], ), ) -def test_rel_list_scan_with_two_axes( +def test_list_rscan_with_two_axes( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -649,11 +649,11 @@ def test_rel_list_scan_with_two_axes( y_list: list, ): run_engine( - rel_list_scan(detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list]) + list_rscan(detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list]) ) -def test_rel_list_scan_with_two_axes_fails_when_given_differnt_list_lengths( +def test_list_rscan_with_two_axes_fails_when_given_differnt_list_lengths( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -661,7 +661,7 @@ def test_rel_list_scan_with_two_axes_fails_when_given_differnt_list_lengths( ): with pytest.raises(ValueError): run_engine( - rel_list_scan( + list_rscan( detectors=[det], movers=[x_axis, y_axis], params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], @@ -729,7 +729,7 @@ def test_list_grid_scan_when_given_bad_info( [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], ), ) -def test_rel_list_grid_scan( +def test_list_grid_rscan( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -738,20 +738,20 @@ def test_rel_list_grid_scan( y_list: list, ): run_engine( - rel_list_grid_scan( + list_grid_rscan( detectors=[det], movers=[x_axis, y_axis], params=[x_list, y_list] ) ) -def test_rel_list_grid_scan_with_two_axes_when_snaking( +def test_list_grid_rscan_with_two_axes_when_snaking( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, y_axis: Motor, ): run_engine( - rel_list_grid_scan( + list_grid_rscan( detectors=[det], movers=[x_axis, y_axis], params=[[1, 2, 3, 4, 5], [1, 2, 3, 4, 5]], @@ -760,14 +760,14 @@ def test_rel_list_grid_scan_with_two_axes_when_snaking( ) -def test_rel_list_grid_scan_when_given_differnt_list_lengths( +def test_list_grid_rscan_when_given_differnt_list_lengths( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, y_axis: Motor, ): run_engine( - rel_list_grid_scan( + list_grid_rscan( detectors=[det], movers=[x_axis, y_axis], params=[[1, 2, 3, 4, 5], [1, 2, 3, 4]], @@ -775,7 +775,7 @@ def test_rel_list_grid_scan_when_given_differnt_list_lengths( ) -def test_rel_list_grid_scan_when_given_bad_info( +def test_list_grid_rscan_when_given_bad_info( run_engine: RunEngine, det: StandardDetector, x_axis: Motor, @@ -783,7 +783,7 @@ def test_rel_list_grid_scan_when_given_bad_info( ): with pytest.raises(TypeError): run_engine( - list_grid_scan( + list_grid_rscan( detectors=[det], movers=[x_axis, y_axis], params=[[1, 2, 3, 4, 5], ["one", 2, 3, 4, 5]], From 1a4e80238bb331c7cf73af202b11c6a827512b92 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Mon, 17 Nov 2025 15:46:43 +0000 Subject: [PATCH 06/16] add shape to list_scan variations --- src/dodal/plans/wrapped.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 3e95562cbdb..3ecd898c25e 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -233,6 +233,8 @@ def list_scan( """Scan over one or more variables in steps simultaneously. Wraps bluesky.plans.list_scan(det, *args, md=metadata).""" metadata = metadata or {} + shape = [len(positions) for positions in params] + metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.list_scan(tuple(detectors), *args, md=metadata) @@ -263,6 +265,8 @@ def list_rscan( """Scan over one or more variables simultaneously relative to current position. Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata).""" metadata = metadata or {} + shape = [len(positions) for positions in params] + metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) @@ -294,6 +298,8 @@ def list_grid_scan( """Scan over one or more variables for each given point on independent trajectories. Wraps bluesky.plans.list_grid_scan(det, *args, md=metadata).""" metadata = metadata or {} + shape = [len(positions) for positions in params] + metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.list_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata @@ -327,6 +333,8 @@ def list_grid_rscan( """Scan over some variables for each given point relative to current position. Wraps bluesky.plans.rel_list_grid_scan(det, *args, md=metadata).""" metadata = metadata or {} + shape = [len(positions) for positions in params] + metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.rel_list_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata From 0c4eccc3981db387e63795647714d44a4b69eda8 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Tue, 18 Nov 2025 15:01:55 +0000 Subject: [PATCH 07/16] add step_scan variations and tests --- src/dodal/plans/wrapped.py | 244 +++++++++++++++++++++++++++++++++--- tests/plans/test_wrapped.py | 195 ++++++++++++++++++++++++++++ 2 files changed, 419 insertions(+), 20 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 3ecd898c25e..a338ef95501 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -3,6 +3,7 @@ from typing import Annotated, Any import bluesky.plans as bp +import numpy as np from bluesky.protocols import Movable, Readable from ophyd_async.core import AsyncReadable from pydantic import Field, NonNegativeFloat, validate_call @@ -67,15 +68,15 @@ def _make_args( ): movers_len = len(movers) params_len = len(params) - if params_len % movers_len != 0 or params_len % num_params != 0: - raise ValueError(f"params must contain {num_params} values for each movable") - args = [] - it = iter(params) - param_chunks = iter(lambda: tuple(itertools.islice(it, num_params)), ()) - for movable, param_chunk in zip(movers, param_chunks, strict=False): - args.append(movable) - args.extend(param_chunk) + if params_len % movers_len == 0 and params_len % num_params == 0: + it = iter(params) + param_chunks = iter(lambda: tuple(itertools.islice(it, num_params)), ()) + for movable, param_chunk in zip(movers, param_chunks, strict=False): + args.append(movable) + args.extend(param_chunk) + else: + raise ValueError(f"params must contain {num_params} values for each movable") return args @@ -160,8 +161,8 @@ def grid_num_scan( params: Annotated[ Sequence[float | int], Field( - description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movers`." + description="Start and stop and number of points for each movable, 'start1," + "stop1, num1, ..., startN, stopN, numN' for every movable in `movers`." ), ], snake_axes: list | bool | None = None, @@ -191,8 +192,8 @@ def grid_num_rscan( params: Annotated[ Sequence[float | int], Field( - description="Start and stop points for each movable, 'start1, stop1, ...," - "startN, stopN' for every movable in `movers`." + description="Start and stop and number of points for each movable, 'start1," + "stop1, num1, ..., startN, stopN, numN' for every movable in `movers`." ), ], snake_axes: list | bool | None = None, @@ -233,8 +234,6 @@ def list_scan( """Scan over one or more variables in steps simultaneously. Wraps bluesky.plans.list_scan(det, *args, md=metadata).""" metadata = metadata or {} - shape = [len(positions) for positions in params] - metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.list_scan(tuple(detectors), *args, md=metadata) @@ -265,8 +264,6 @@ def list_rscan( """Scan over one or more variables simultaneously relative to current position. Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata).""" metadata = metadata or {} - shape = [len(positions) for positions in params] - metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) @@ -298,8 +295,6 @@ def list_grid_scan( """Scan over one or more variables for each given point on independent trajectories. Wraps bluesky.plans.list_grid_scan(det, *args, md=metadata).""" metadata = metadata or {} - shape = [len(positions) for positions in params] - metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.list_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata @@ -333,9 +328,218 @@ def list_grid_rscan( """Scan over some variables for each given point relative to current position. Wraps bluesky.plans.rel_list_grid_scan(det, *args, md=metadata).""" metadata = metadata or {} - shape = [len(positions) for positions in params] - metadata["shape"] = (shape,) args = _make_args(movers=movers, params=params, num_params=1) yield from bp.rel_list_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata ) + + +def _make_stepped_list( + params: list[Any] | Sequence[Any], + num: int | None = None, +): + start = params[0] + if len(params) == 3: + stop = params[1] + step = params[2] + stepped_list = np.arange(start=start, stop=stop, step=step).tolist() + if abs((stepped_list[-1] + step) - stop) <= (step * 0.01): + stepped_list.append(stepped_list[-1] + step) + elif len(params) == 2 and num: + step = params[1] + stepped_list = [start + (n * step) for n in range(num)] + else: + stepped_list = [] + return stepped_list + + +def _make_concurrently_stepped_lists( + movers_len: int, + params: list[Any] | Sequence[Any], +): + # separate out the first three elements + # separate the rest of the elements into chunks of two + stepped_lists = [] + params_len = len(params) + if movers_len == 1 and params_len == 3: + stepped_lists.append(_make_stepped_list(params)) + elif movers_len > 1 and (params_len - 3) % 2 == 0: + leader_params = params[:3] + follower_params = params[3:] + stepped_lists.append(_make_stepped_list(leader_params)) + num = len(stepped_lists[0]) + it = iter(follower_params) + param_chunks = iter(lambda: tuple(itertools.islice(it, 2)), ()) + for param_chunk in param_chunks: + stepped_lists.append(_make_stepped_list(param_chunk, num=num)) + else: + raise ValueError( + "params must contain 3 values for the first movable and 2 values for each " + "successive movable." + ) + return stepped_lists + + +def _make_independently_stepped_lists( + movers_len: int, + params: list[Any] | Sequence[Any], +): + num_params = 3 # It will always be start stop step for each axis + stepped_lists = [] + params_len = len(params) + if params_len % movers_len == 0 and params_len % num_params == 0: + it = iter(params) + param_chunks = iter(lambda: tuple(itertools.islice(it, num_params)), ()) + for param_chunk in param_chunks: + stepped_lists.append(_make_stepped_list(param_chunk)) + else: + raise ValueError(f"params must contain {num_params} values for each movable") + return stepped_lists + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def step_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[Any], + Field( + description="Start, stop, and step values for the first movable, and start " + "and step values for each successive movable, 'start1, stop1, step1, " + "start2, step 2, ..., startN, stepN' for every movable in `movers`." + ), + ], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over one multi-motor trajectory with specified step size. + Generates lists of points for each trajectory for bp.list_scan. + """ + movers_len = len(movers) + stepped_lists = _make_concurrently_stepped_lists( + movers_len=movers_len, params=params + ) + metadata = metadata or {} + args = _make_args(movers=movers, params=stepped_lists, num_params=1) + yield from bp.list_scan(tuple(detectors), *args, md=metadata) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def step_rscan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[Any], + Field( + description="Start, stop, and step values for the first movable, and start " + "and step values for each successive movable, 'start1, stop1, step1, " + "start2, step 2, ..., startN, stepN' for every movable in `movers`." + ), + ], + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over multi-motor trajectory relative to position with specified step size. + Generates lists of points for each trajectory for bp.rel_list_scan. + """ + movers_len = len(movers) + stepped_lists = _make_concurrently_stepped_lists( + movers_len=movers_len, params=params + ) + metadata = metadata or {} + args = _make_args(movers=movers, params=stepped_lists, num_params=1) + yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def step_grid_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[Any], + Field( + description="Start, stop, and step values 'start1, stop1, step1, ..., " + "startN, stepN' for every movable in `movers`." + ), + ], + snake_axes: bool = False, # Currently specifying axes to snake is not supported + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over independent multi-motor trajectories with specified step size. + Generates lists of points for each trajectory for bp.list_grid_scan. + """ + movers_len = len(movers) + stepped_lists = _make_independently_stepped_lists( + movers_len=movers_len, params=params + ) + metadata = metadata or {} + args = _make_args(movers=movers, params=stepped_lists, num_params=1) + yield from bp.list_grid_scan( + tuple(detectors), *args, snake_axes=snake_axes, md=metadata + ) + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def step_grid_rscan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + movers: Annotated[ + Sequence[Movable | Motor], + Field(description="One or more movable to move during the scan."), + ], + params: Annotated[ + list[Any], + Field( + description="Start, stop, and step values 'start1, stop1, step1, ..., " + "startN, stepN' for every movable in `movers`." + ), + ], + snake_axes: bool = False, # Currently specifying axes to snake is not supported + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan over independent multi-motor relative trajectories with specified step size. + Generates lists of points for each trajectory for bp.rel_list_grid_scan. + """ + movers_len = len(movers) + stepped_lists = _make_independently_stepped_lists( + movers_len=movers_len, params=params + ) + metadata = metadata or {} + args = _make_args(movers=movers, params=stepped_lists, num_params=1) + yield from bp.rel_list_grid_scan( + tuple(detectors), *args, snake_axes=snake_axes, md=metadata + ) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index bda4b926539..c5c5e99e171 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -21,6 +21,9 @@ from dodal.devices.motors import Motor from dodal.plans.wrapped import ( _make_args, + _make_concurrently_stepped_lists, + _make_independently_stepped_lists, + _make_stepped_list, count, grid_num_rscan, grid_num_scan, @@ -30,6 +33,10 @@ list_scan, num_rscan, num_scan, + step_grid_rscan, + step_grid_scan, + step_rscan, + step_scan, ) @@ -789,3 +796,191 @@ def test_list_grid_rscan_when_given_bad_info( params=[[1, 2, 3, 4, 5], ["one", 2, 3, 4, 5]], ) ) + + +def test_make_stepped_list_when_given_three_params(): + stepped_list = _make_stepped_list(params=[0, 1, 0.1]) + assert len(stepped_list) == 11 + assert stepped_list[0] == 0 + assert stepped_list[-1] == 1 + + +def test_make_stepped_list_when_given_two_params(): + stepped_list = _make_stepped_list(params=[0, 0.1], num=11) + assert len(stepped_list) == 11 + assert stepped_list[0] == 0 + assert stepped_list[-1] == 1 + + +def test_make_concurrently_stepped_lists(): + stepped_lists = _make_concurrently_stepped_lists( + movers_len=2, params=[0, 1, 0.1, 0, 1] + ) + assert len(stepped_lists) == 2 + assert len(stepped_lists[0]) == 11 and len(stepped_lists[1]) == 11 + assert stepped_lists[0][-1] == 1 + assert stepped_lists[1][-1] == 10 + + +def test_make_independently_stepped_lists(): + stepped_lists = _make_independently_stepped_lists( + movers_len=2, params=[0, 1, 0.1, -10, 10, 1] + ) + assert len(stepped_lists) == 2 + assert len(stepped_lists[0]) == 11 and len(stepped_lists[1]) == 21 + assert stepped_lists[0][-1] == 1 + assert stepped_lists[1][-1] == 10 + + +@pytest.mark.parametrize("params", ([0, 1, 0.1], [-1, 1, 0.1], [0, 10, 1])) +def test_step_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + params: list[Any], +): + run_engine(step_scan(detectors=[det], movers=[x_axis], params=params)) + + +@pytest.mark.parametrize( + "params", ([0, 1, 0.1, 0, 0.1], [-1, 1, 0.1, -1, 0.1], [0, 10, 1, 0, 1]) +) +def test_step_scan_with_multiple_movers( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine(step_scan(detectors=[det], movers=[x_axis, y_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 1, 0.1, 0])) +def test_step_scan_fails_when_given_incorrect_number_of_params( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + with pytest.raises(ValueError): + run_engine(step_scan(detectors=[det], movers=[x_axis, y_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1], [-1, 1, 0.1], [0, 10, 1])) +def test_step_rscan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + params: list[Any], +): + run_engine(step_rscan(detectors=[det], movers=[x_axis], params=params)) + + +@pytest.mark.parametrize( + "params", ([0, 1, 0.1, 0, 0.1], [-1, 1, 0.1, -1, 0.1], [0, 10, 1, 0, 1]) +) +def test_step_rscan_with_multiple_movers( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine(step_rscan(detectors=[det], movers=[x_axis, y_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 1, 0.1, 0])) +def test_step_rscan_fails_when_given_incorrect_number_of_params( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + with pytest.raises(ValueError): + run_engine(step_rscan(detectors=[det], movers=[x_axis, y_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +def test_step_grid_scan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine(step_grid_scan(detectors=[det], movers=[y_axis, x_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +def test_step_grid_scan_when_snaking( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine( + step_grid_scan( + detectors=[det], movers=[y_axis, x_axis], params=params, snake_axes=True + ) + ) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1], [0, 10, 1, 0])) +def test_step_grid_scan_fails_when_given_incorrect_number_of_params( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + with pytest.raises(ValueError): + run_engine( + step_grid_scan( + detectors=[det], movers=[y_axis, x_axis], params=params, snake_axes=True + ) + ) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +def test_step_grid_rscan( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine(step_grid_rscan(detectors=[det], movers=[y_axis, x_axis], params=params)) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +def test_step_grid_rscan_when_snaking( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + run_engine( + step_grid_rscan( + detectors=[det], movers=[y_axis, x_axis], params=params, snake_axes=True + ) + ) + + +@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1], [0, 10, 1, 0])) +def test_step_grid_rscan_fails_when_given_incorrect_number_of_params( + run_engine: RunEngine, + det: StandardDetector, + x_axis: Motor, + y_axis: Motor, + params: list[Any], +): + with pytest.raises(ValueError): + run_engine( + step_grid_rscan( + detectors=[det], movers=[y_axis, x_axis], params=params, snake_axes=True + ) + ) From f6be2eb419d7550c04a6fffe0ddb3f028053f7f2 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Wed, 19 Nov 2025 08:15:39 +0000 Subject: [PATCH 08/16] Add new plans to init --- src/dodal/plans/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/dodal/plans/__init__.py b/src/dodal/plans/__init__.py index 089e5f83779..b94c77f4163 100644 --- a/src/dodal/plans/__init__.py +++ b/src/dodal/plans/__init__.py @@ -9,6 +9,10 @@ list_scan, num_rscan, num_scan, + step_grid_rscan, + step_grid_scan, + step_rscan, + step_scan, ) __all__ = [ @@ -22,4 +26,8 @@ "num_rscan", "num_scan", "spec_scan", + "step_grid_rscan", + "step_grid_scan", + "step_rscan", + "step_scan", ] From 30e1608e9aa6b74962cfd4518f74d59933eca251 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Thu, 20 Nov 2025 15:44:15 +0000 Subject: [PATCH 09/16] add rounding to avoid floating point errors --- src/dodal/plans/wrapped.py | 15 ++++++++++++--- tests/plans/test_wrapped.py | 28 ++++++++++++++++------------ 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index a338ef95501..fb2af9d5ad2 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -1,5 +1,6 @@ import itertools from collections.abc import Sequence +from decimal import Decimal from typing import Annotated, Any import bluesky.plans as bp @@ -338,19 +339,27 @@ def _make_stepped_list( params: list[Any] | Sequence[Any], num: int | None = None, ): + def round_list_elements(stepped_list, step): + d = Decimal(str(step)) + exponent = d.as_tuple().exponent + decimal_places = -exponent # type: ignore + return np.round(stepped_list, decimals=decimal_places).tolist() + start = params[0] if len(params) == 3: stop = params[1] step = params[2] stepped_list = np.arange(start=start, stop=stop, step=step).tolist() - if abs((stepped_list[-1] + step) - stop) <= (step * 0.01): + if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.01): stepped_list.append(stepped_list[-1] + step) + rounded_stepped_list = round_list_elements(stepped_list=stepped_list, step=step) elif len(params) == 2 and num: step = params[1] stepped_list = [start + (n * step) for n in range(num)] + rounded_stepped_list = round_list_elements(stepped_list=stepped_list, step=step) else: - stepped_list = [] - return stepped_list + rounded_stepped_list = [] + return rounded_stepped_list def _make_concurrently_stepped_lists( diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index c5c5e99e171..9fd5615d768 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -798,18 +798,22 @@ def test_list_grid_rscan_when_given_bad_info( ) -def test_make_stepped_list_when_given_three_params(): - stepped_list = _make_stepped_list(params=[0, 1, 0.1]) - assert len(stepped_list) == 11 - assert stepped_list[0] == 0 - assert stepped_list[-1] == 1 - - -def test_make_stepped_list_when_given_two_params(): - stepped_list = _make_stepped_list(params=[0, 0.1], num=11) - assert len(stepped_list) == 11 - assert stepped_list[0] == 0 - assert stepped_list[-1] == 1 +@pytest.mark.parametrize( + "params", ([-1, 1, 0.1], [-2, 2, 0.2], [1, -1, -0.1], [2, -2, -0.2]) +) +def test_make_stepped_list_when_given_three_params(params: list[Any]): + stepped_list = _make_stepped_list(params=params) + assert len(stepped_list) == 21 + assert stepped_list[0] / stepped_list[-1] == -1 + assert stepped_list[10] == 0 + + +@pytest.mark.parametrize("params", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) +def test_make_stepped_list_when_given_two_params(params: list[Any]): + stepped_list = _make_stepped_list(params=params, num=21) + assert len(stepped_list) == 21 + assert stepped_list[0] / stepped_list[-1] == -1 + assert stepped_list[10] == 0 def test_make_concurrently_stepped_lists(): From 81e437c8b4f404c6ce780e91c21e9e05b29e9ea5 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Thu, 20 Nov 2025 15:57:41 +0000 Subject: [PATCH 10/16] Increase step size in tests --- tests/plans/test_wrapped.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 9fd5615d768..88cabc8606e 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -871,7 +871,7 @@ def test_step_scan_fails_when_given_incorrect_number_of_params( run_engine(step_scan(detectors=[det], movers=[x_axis, y_axis], params=params)) -@pytest.mark.parametrize("params", ([0, 1, 0.1], [-1, 1, 0.1], [0, 10, 1])) +@pytest.mark.parametrize("params", ([0, 1, 0.25], [-1, 1, 0.5], [0, 10, 2.5])) def test_step_rscan( run_engine: RunEngine, det: StandardDetector, @@ -882,7 +882,7 @@ def test_step_rscan( @pytest.mark.parametrize( - "params", ([0, 1, 0.1, 0, 0.1], [-1, 1, 0.1, -1, 0.1], [0, 10, 1, 0, 1]) + "params", ([0, 1, 0.25, 0, 0.25], [-1, 1, 0.5, -1, 0.5], [0, 10, 2.5, 0, 2.5]) ) def test_step_rscan_with_multiple_movers( run_engine: RunEngine, @@ -894,7 +894,7 @@ def test_step_rscan_with_multiple_movers( run_engine(step_rscan(detectors=[det], movers=[x_axis, y_axis], params=params)) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 1, 0.1, 0])) +@pytest.mark.parametrize("params", ([0, 1, 0.5, 0, 1, 0.5], [0, 1, 0.5, 0])) def test_step_rscan_fails_when_given_incorrect_number_of_params( run_engine: RunEngine, det: StandardDetector, @@ -906,7 +906,7 @@ def test_step_rscan_fails_when_given_incorrect_number_of_params( run_engine(step_rscan(detectors=[det], movers=[x_axis, y_axis], params=params)) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +@pytest.mark.parametrize("params", ([0, 1, 0.5, 0, 1, 0.5], [0, 10, 5, 0, 10, 5])) def test_step_grid_scan( run_engine: RunEngine, det: StandardDetector, @@ -917,7 +917,7 @@ def test_step_grid_scan( run_engine(step_grid_scan(detectors=[det], movers=[y_axis, x_axis], params=params)) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +@pytest.mark.parametrize("params", ([0, 1, 0.5, 0, 1, 0.5], [0, 10, 5, 0, 10, 5])) def test_step_grid_scan_when_snaking( run_engine: RunEngine, det: StandardDetector, @@ -932,7 +932,7 @@ def test_step_grid_scan_when_snaking( ) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1], [0, 10, 1, 0])) +@pytest.mark.parametrize("params", ([0, 1, 0.25, 0, 1], [0, 10, 2.5, 0])) def test_step_grid_scan_fails_when_given_incorrect_number_of_params( run_engine: RunEngine, det: StandardDetector, @@ -948,7 +948,7 @@ def test_step_grid_scan_fails_when_given_incorrect_number_of_params( ) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +@pytest.mark.parametrize("params", ([0, 1, 0.5, 0, 1, 0.5], [0, 10, 5, 0, 10, 5])) def test_step_grid_rscan( run_engine: RunEngine, det: StandardDetector, @@ -959,7 +959,7 @@ def test_step_grid_rscan( run_engine(step_grid_rscan(detectors=[det], movers=[y_axis, x_axis], params=params)) -@pytest.mark.parametrize("params", ([0, 1, 0.1, 0, 1, 0.1], [0, 10, 1, 0, 10, 1])) +@pytest.mark.parametrize("params", ([0, 1, 0.5, 0, 1, 0.5], [0, 10, 5, 0, 10, 5])) def test_step_grid_rscan_when_snaking( run_engine: RunEngine, det: StandardDetector, From 1f4718a45b3b75da08a69f3a3321dd03c34e6bbe Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Thu, 20 Nov 2025 16:23:44 +0000 Subject: [PATCH 11/16] add test for empty stepped_list --- tests/plans/test_wrapped.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 88cabc8606e..b4cd3485059 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -816,6 +816,11 @@ def test_make_stepped_list_when_given_two_params(params: list[Any]): assert stepped_list[10] == 0 +def test_make_stepped_list_when_given_wrong_number_of_params(): + stepped_list = _make_stepped_list(params=[1]) + assert stepped_list == [] + + def test_make_concurrently_stepped_lists(): stepped_lists = _make_concurrently_stepped_lists( movers_len=2, params=[0, 1, 0.1, 0, 1] From 88bbce0470981ff7258a20b93325c14fc23b4d05 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Thu, 20 Nov 2025 18:52:18 +0000 Subject: [PATCH 12/16] Try dictionary for args, rather than separate movables and parameters --- src/dodal/plans/wrapped.py | 31 +++++++++++++++++++++++++++++++ tests/plans/test_wrapped.py | 19 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index fb2af9d5ad2..ef6b7f1ebc1 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -62,6 +62,37 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def mapping_num_scan( + detectors: Annotated[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + min_length=1, + ), + ], + params: dict[Movable | Motor, list[float | int]], + num: Annotated[int, Field(description="Number of points")], + metadata: dict[str, Any] | None = None, +): + """Scan over one multi-motor trajectory. + Wraps bluesky.plans.scan(det, *args, num, md=metadata)""" + metadata = metadata or {} + metadata["shape"] = (num,) + + args = _make_new_args(params) + yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) + + +def _make_new_args(params: dict[Movable | Motor, list[float | int]]): + args = [] + for param in params: + args.append(param) + args.extend(params[param]) + return args + + def _make_args( movers: Sequence[Movable | Motor], params: list[Any] | Sequence[Any], diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index b4cd3485059..84a0d5a568f 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -16,6 +16,7 @@ AsyncReadable, StandardDetector, ) +from ophyd_async.sim import SimMotor from pydantic import ValidationError from dodal.devices.motors import Motor @@ -23,6 +24,7 @@ _make_args, _make_concurrently_stepped_lists, _make_independently_stepped_lists, + _make_new_args, _make_stepped_list, count, grid_num_rscan, @@ -31,6 +33,7 @@ list_grid_scan, list_rscan, list_scan, + mapping_num_scan, num_rscan, num_scan, step_grid_rscan, @@ -179,6 +182,22 @@ def test_plan_produces_expected_datums( assert docs and len(docs) == len(data_keys) * length +def test_make_new_args(x_axis: Motor, y_axis: Motor): + args = _make_new_args({x_axis: [0, 1], y_axis: [2, 3]}) + assert args[0] == x_axis + assert args[1] == 0 + assert args[2] == 1 + assert args[3] == y_axis + assert args[4] == 2 + assert args[5] == 3 + + +def test_mapping_num_scan( + run_engine: RunEngine, det: StandardDetector, x_axis: SimMotor +): + run_engine(mapping_num_scan(detectors=[det], params={x_axis: [1, 2]}, num=3)) + + @pytest.mark.parametrize( "num_params, params", ([2, [1, 2, 3, 4]], [3, [1, 2, 3, 3, 4, 3]]) ) From 26996852345fac7006bd69fdded2691e1c19df47 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Mon, 24 Nov 2025 15:31:29 +0000 Subject: [PATCH 13/16] Remove test scan --- src/dodal/plans/wrapped.py | 31 ------------------------------- tests/plans/test_wrapped.py | 8 -------- 2 files changed, 39 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index ef6b7f1ebc1..fb2af9d5ad2 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -62,37 +62,6 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) -@attach_data_session_metadata_decorator() -@validate_call(config={"arbitrary_types_allowed": True}) -def mapping_num_scan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - min_length=1, - ), - ], - params: dict[Movable | Motor, list[float | int]], - num: Annotated[int, Field(description="Number of points")], - metadata: dict[str, Any] | None = None, -): - """Scan over one multi-motor trajectory. - Wraps bluesky.plans.scan(det, *args, num, md=metadata)""" - metadata = metadata or {} - metadata["shape"] = (num,) - - args = _make_new_args(params) - yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) - - -def _make_new_args(params: dict[Movable | Motor, list[float | int]]): - args = [] - for param in params: - args.append(param) - args.extend(params[param]) - return args - - def _make_args( movers: Sequence[Movable | Motor], params: list[Any] | Sequence[Any], diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 84a0d5a568f..f5e2a3f734a 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -16,7 +16,6 @@ AsyncReadable, StandardDetector, ) -from ophyd_async.sim import SimMotor from pydantic import ValidationError from dodal.devices.motors import Motor @@ -33,7 +32,6 @@ list_grid_scan, list_rscan, list_scan, - mapping_num_scan, num_rscan, num_scan, step_grid_rscan, @@ -192,12 +190,6 @@ def test_make_new_args(x_axis: Motor, y_axis: Motor): assert args[5] == 3 -def test_mapping_num_scan( - run_engine: RunEngine, det: StandardDetector, x_axis: SimMotor -): - run_engine(mapping_num_scan(detectors=[det], params={x_axis: [1, 2]}, num=3)) - - @pytest.mark.parametrize( "num_params, params", ([2, [1, 2, 3, 4]], [3, [1, 2, 3, 3, 4, 3]]) ) From 7afd8a0f0cd7290a2704f345215e37d4c3c7ba33 Mon Sep 17 00:00:00 2001 From: Emily Arnold <222046505+EmsArnold@users.noreply.github.com> Date: Mon, 24 Nov 2025 15:46:03 +0000 Subject: [PATCH 14/16] Remove unused import from test_wrapped.py Removed unused import of _make_new_args. --- tests/plans/test_wrapped.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index f5e2a3f734a..3f1782a420c 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -23,7 +23,6 @@ _make_args, _make_concurrently_stepped_lists, _make_independently_stepped_lists, - _make_new_args, _make_stepped_list, count, grid_num_rscan, From 54c58463562e947402d1fb25b7bbdf967e731168 Mon Sep 17 00:00:00 2001 From: Emily Arnold <222046505+EmsArnold@users.noreply.github.com> Date: Mon, 24 Nov 2025 15:51:30 +0000 Subject: [PATCH 15/16] Remove unused test from test_wrapped.py --- tests/plans/test_wrapped.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 3f1782a420c..b4cd3485059 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -179,16 +179,6 @@ def test_plan_produces_expected_datums( assert docs and len(docs) == len(data_keys) * length -def test_make_new_args(x_axis: Motor, y_axis: Motor): - args = _make_new_args({x_axis: [0, 1], y_axis: [2, 3]}) - assert args[0] == x_axis - assert args[1] == 0 - assert args[2] == 1 - assert args[3] == y_axis - assert args[4] == 2 - assert args[5] == 3 - - @pytest.mark.parametrize( "num_params, params", ([2, [1, 2, 3, 4]], [3, [1, 2, 3, 3, 4, 3]]) ) From 17052134b1e6f8e4f4bd375788637a4c5340ae42 Mon Sep 17 00:00:00 2001 From: Emily Arnold Date: Tue, 25 Nov 2025 12:56:01 +0000 Subject: [PATCH 16/16] reduce reliance on sign of step --- src/dodal/plans/wrapped.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index fb2af9d5ad2..46cdbfcfbdf 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -349,6 +349,9 @@ def round_list_elements(stepped_list, step): if len(params) == 3: stop = params[1] step = params[2] + if abs(step) > abs(stop - start): + step = abs(stop - start) + step = abs(step) * np.sign(stop - start) stepped_list = np.arange(start=start, stop=stop, step=step).tolist() if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.01): stepped_list.append(stepped_list[-1] + step)