From 0adc946fe5dbc49d9b76b2c6df869f46f9a6cd7f Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 1 May 2026 13:19:18 +0000 Subject: [PATCH 01/26] Refactor scans so they are flatter --- src/dodal/plans/wrapped.py | 327 +++++++++++++++++++++--------------- tests/plans/test_wrapped.py | 215 ++++++++++++------------ 2 files changed, 299 insertions(+), 243 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 1ea01906db8..043f7ad922f 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -62,31 +62,34 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) -def _make_num_scan_args( - params: list[tuple[Movable, list[float | int]]], num: int | None = None -): - shape = [] - if num: - shape = [num] - for param in params: - if len(param[1]) == 2: - pass - else: - raise ValueError("You must provide 'start stop' for each motor.") - else: - for param in params: - if len(param[1]) == 3: - shape.append(param[1][-1]) - else: - raise ValueError( - "You must provide 'start stop num' for each motor in a grid scan." - ) - - args = [] - for param in params: - args.append(param[0]) - args.extend(param[1]) - return args, shape +# def _make_num_scan_args( +# params: Sequence[Movable | float | int], num: int | None = None +# ) -> list[float]: +# # shape = [] +# # if num: +# # shape = [num] +# # for param in params: +# # if len(param[1]) == 2: +# # pass +# # else: +# # raise ValueError("You must provide 'start stop' for each motor.") +# # else: +# # for param in params: +# # if len(param[1]) == 3: +# # shape.append(param[1][-1]) +# # else: +# # raise ValueError( +# # "You must provide 'start stop num' for each motor in a grid scan." +# # ) + +# # args = [] +# # for param in params: +# # args.append(param[0]) +# # args.extend(param[1]) +# for param in params: +# if isinstance(param, Movable): + +# return args, shape @validate_call(config={"arbitrary_types_allowed": True}) @@ -98,11 +101,11 @@ def num_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For concurrent " - "trajectories, provide '[(movable1, [start1, stop1]), (movable2, [start2, " - "stop2]), ... , (movableN, [startN, stopN])]'." + "trajectories, provide '[movable1, start1, stop1, movable2, start2, stop2, " + "... , movableN, startN, stopN]'." ), ], num: int, @@ -113,12 +116,10 @@ def num_scan( The scan is defined by number of points along scan trajector(y/ies). Wraps bluesky.plans.scan(det, *args, num, md=metadata). """ - # TODO: move to using Range spec and spec_scan when stable and tested at v1.0 - args, shape = _make_num_scan_args(params, num) metadata = metadata or {} - metadata["shape"] = shape + metadata["shape"] = (num,) - yield from bp.scan(tuple(detectors), *args, num=num, md=metadata) + yield from bp.scan(tuple(detectors), *params, num=num, md=metadata) @validate_call(config={"arbitrary_types_allowed": True}) @@ -130,7 +131,7 @@ def num_grid_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ trajectories, provide '[(movable1, [start1, stop1, num1]), (movable2, \ @@ -146,12 +147,9 @@ def num_grid_scan( axes by default (all axes but the first axis provided). Wraps bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata). """ - # TODO: move to using Range spec and spec_scan when stable and tested at v1.0 - args, shape = _make_num_scan_args(params) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.grid_scan(tuple(detectors), *args, snake_axes=snake_axes, md=metadata) + yield from bp.grid_scan( + tuple(detectors), *params, snake_axes=snake_axes, md=metadata + ) @validate_call(config={"arbitrary_types_allowed": True}) @@ -163,11 +161,11 @@ def num_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For concurrent \ - trajectories, provide '[(movable1, [start1, stop1]), (movable2, [start2, \ - stop2]), ... , (movableN, [startN, stopN])]'." + trajectories, provide '[movable1, start1, stop1, movable2, start2, stop2, \ + ... , movableN, startN, stopN]'." ), ], num: int | None = None, @@ -178,12 +176,10 @@ def num_rscan( The scan is defined by number of points along scan trajector(y/ies). Wraps bluesky.plans.rel_scan(det, *args, num, md=metadata). """ - # TODO: move to using Range spec and spec_scan when stable and tested at v1.0 - args, shape = _make_num_scan_args(params, num) metadata = metadata or {} - metadata["shape"] = shape + metadata["shape"] = (num,) - yield from bp.rel_scan(tuple(detectors), *args, num=num, md=metadata) + yield from bp.rel_scan(tuple(detectors), *params, num=num, md=metadata) @validate_call(config={"arbitrary_types_allowed": True}) @@ -195,7 +191,7 @@ def num_grid_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ trajectories, provide '[(movable1, [start1, stop1, num1]), (movable2, \ @@ -211,30 +207,20 @@ def num_grid_rscan( axes by default (all axes but the first axis provided). Wraps bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata). """ - # TODO: move to using Range spec and spec_scan when stable and tested at v1.0 - args, shape = _make_num_scan_args(params) - metadata = metadata or {} - metadata["shape"] = shape - yield from bp.rel_grid_scan( - tuple(detectors), *args, snake_axes=snake_axes, md=metadata + tuple(detectors), *params, snake_axes=snake_axes, md=metadata ) -def _make_list_scan_args(params: list[tuple[Movable, list[float | int]]], grid: bool): - shape = [] - args = [] +def _make_list_scan_shape( + params: Sequence[Movable | list[float | int]], +) -> tuple[int, ...]: for param in params: - shape.append(len(param[1])) - args.append(param[0]) - args.append(param[1]) - - if not grid: - shape = list(set(shape)) - if len(shape) > 1: - raise ValueError("Lists of motor positions are not equal in length.") - - return args, shape + # List arg must all be same size. If list missing or not same size, this will + # be validated by bp.list_scan. + if isinstance(param, list): + return (len(param),) + return () @validate_call(config={"arbitrary_types_allowed": True}) @@ -246,7 +232,7 @@ def list_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + list[Movable | list[float | int]], Field( description="List of tuples (device, positions). For concurrent \ trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ @@ -261,11 +247,11 @@ def list_scan( The scan is defined by providing a list of points for each scan trajectory. Wraps bluesky.plans.list_scan(det, *args, md=metadata). """ - args, shape = _make_list_scan_args(params=params, grid=False) metadata = metadata or {} - metadata["shape"] = shape + metadata["shape"] = _make_list_scan_shape(params) - yield from bp.list_scan(tuple(detectors), *args, md=metadata) + # Not sure about this one + yield from bp.list_scan(tuple(detectors), *tuple(params), md=metadata) # type: ignore @validate_call(config={"arbitrary_types_allowed": True}) @@ -277,7 +263,7 @@ def list_grid_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | list[float | int]], Field( description="List of tuples (device, positions). For independent \ trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ @@ -293,12 +279,15 @@ def list_grid_scan( all fast axes by default (all axes but the first axis provided). Wraps bluesky.plans.list_grid_scan(det, *args, md=metadata). """ - args, shape = _make_list_scan_args(params=params, grid=True) metadata = metadata or {} - metadata["shape"] = shape + shape = [] + for param in params: + if isinstance(param, list): + shape.append(len(param)) + metadata["shape"] = tuple(shape) yield from bp.list_grid_scan( - tuple(detectors), *args, snake_axes=snake_axes, md=metadata + tuple(detectors), *params, snake_axes=snake_axes, md=metadata ) @@ -311,7 +300,7 @@ def list_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | list[float | int]], Field( description="List of tuples (device, positions). For concurrent \ trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ @@ -326,11 +315,9 @@ def list_rscan( The scan is defined by providing a list of points for each scan trajectory. Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). """ - args, shape = _make_list_scan_args(params=params, grid=False) metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) + metadata["shape"] = _make_list_scan_shape(params) + yield from bp.rel_list_scan(tuple(detectors), *params, md=metadata) @validate_call(config={"arbitrary_types_allowed": True}) @@ -342,7 +329,7 @@ def list_grid_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | list[float | int]], Field( description="List of tuples (device, positions). For independent \ trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ @@ -358,16 +345,16 @@ def list_grid_rscan( all fast axes by default (all axes but the first axis provided). Wraps bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). """ - args, shape = _make_list_scan_args(params=params, grid=True) metadata = metadata or {} - metadata["shape"] = shape - + metadata["shape"] = _make_list_scan_shape(params) yield from bp.rel_list_grid_scan( - tuple(detectors), *args, snake_axes=snake_axes, md=metadata + tuple(detectors), *params, snake_axes=snake_axes, md=metadata ) -def _round_list_elements(stepped_list, params) -> list[float]: +def _round_list_elements( + stepped_list: list[float | int], params: list[float | int] +) -> list[float | int]: decimals = [Decimal(str(param)) for param in params] exponents = [d.as_tuple().exponent for d in decimals] decimal_places = [-exponent for exponent in exponents] # type: ignore @@ -375,7 +362,9 @@ def _round_list_elements(stepped_list, params) -> list[float]: return np.round(stepped_list, decimals=max_decimal_places).tolist() -def _make_stepped_list_step(start: float, stop: float, step: float) -> list: +def _make_stepped_list_step( + start: float, stop: float, step: float +) -> list[float | int]: if start == stop: raise ValueError( f"Start ({start}) and stop ({stop}) values cannot be the same." @@ -392,7 +381,9 @@ def _make_stepped_list_step(start: float, stop: float, step: float) -> list: return rounded_stepped_list -def _make_stepped_list_num(start, step, num) -> list: +def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | int]: + if num <= 0: + raise ValueError() stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements( stepped_list=stepped_list, params=[start, step] @@ -400,49 +391,108 @@ def _make_stepped_list_num(start, step, num) -> list: return rounded_stepped_list -def _make_step_scan_args( - params: list[tuple[Movable, list[float | int]]], grid: bool -) -> tuple[list[Any], list[float]]: - args = [] +def _make_step_scan_args_and_shape( + params: Sequence[Movable | float | int], grid: bool +) -> tuple[Sequence[Movable | list[float | int]], tuple[int, ...]]: + + args: list[Movable | list[float | int]] = [] shape = [] - stepped_list_length = None - - first_movable_param, *additional_movable_params = params - if len(first_movable_param[1]) == 3: - start, stop, step = first_movable_param[1] - stepped_list = _make_stepped_list_step(start, stop, step) - stepped_list_length = len(stepped_list) - args.append(first_movable_param[0]) - args.append(stepped_list) - shape.append(stepped_list_length) - else: - raise ValueError( - f"You provided {len(first_movable_param[1])} parameters for {first_movable_param[0]}, rather than 3." - ) - for param in additional_movable_params: - if grid: - if len(param[1]) == 3: - start, stop, step = param[1] - stepped_list = _make_stepped_list_step(start, stop, step) - args.append(param[0]) - args.append(stepped_list) - shape.append(len(stepped_list)) - else: - raise ValueError( - f"You provided {len(param[1])} parameters for {param[0]}, rather than 3." - ) - else: - if len(param[1]) == 2: - start, step = param[1] - stepped_list = _make_stepped_list_num(start, step, stepped_list_length) - args.append(param[0]) - args.append(stepped_list) - else: - raise ValueError( - f"You provided {len(param[1])} parameters {param[0]}, rather than 2." - ) - - return args, shape + stepped_list_length = 0 + + try: + for i, param in enumerate(params): + if isinstance(param, Movable): + movable = param + start = params[i + 1] + if not isinstance(start, (float, int)): + raise ValueError( + f"You provided movable {movable} with no start, stop, step." + ) + + stop = params[i + 2] + if not isinstance(stop, (float, int)): + raise ValueError( + f"You provided movable {movable} with start value {start} but no stop and step." + ) + + step = params[i + 3] + if not isinstance(step, (float, int)): + raise ValueError( + f"You provided movable {movable} with start value {start}, stop value {stop} but no step value." + ) + + movable_values = _make_stepped_list_step(start, stop, step) + stepped_list_length = len(movable_values) + args.append(movable) + args.append(movable_values) + + if not grid: + break + + if not grid: + # Skip first 4 values as already done them. + for i, param in enumerate(params[4:]): + if isinstance(param, Movable): + movable = param + start = params[i + 1] + if not isinstance(start, (float, int)): + raise ValueError( + f"You provided movable {movable} with no start, stop, step." + ) + + step = params[i + 2] + if not isinstance(step, (float, int)): + raise ValueError( + f"You provided movable {movable} with start value {start}, stop value {step} but no step value." + ) + + movable_values = _make_stepped_list_num( + start, step, stepped_list_length + ) + args.append(movable) + args.append(movable_values) + # shape.append(len(values)) + except IndexError as e: + raise ValueError("Incorrect parameters provided.") from e + + return args, tuple(shape) + + # first_movable_param, *additional_movable_params = params + # if len(first_movable_param[1]) == 3: + # start, stop, step = first_movable_param[1] + # stepped_list = _make_stepped_list_step(start, stop, step) + # stepped_list_length = len(stepped_list) + # args.append(first_movable_param[0]) + # args.append(stepped_list) + # shape.append(stepped_list_length) + # else: + # raise ValueError( + # f"You provided {len(first_movable_param[1])} parameters for {first_movable_param[0]}, rather than 3." + # ) + # for param in additional_movable_params: + # if grid: + # if len(param[1]) == 3: + # start, stop, step = param[1] + # stepped_list = _make_stepped_list_step(start, stop, step) + # args.append(param[0]) + # args.append(stepped_list) + # shape.append(len(stepped_list)) + # else: + # raise ValueError( + # f"You provided {len(param[1])} parameters for {param[0]}, rather than 3." + # ) + # else: + # if len(param[1]) == 2: + # start, step = param[1] + # stepped_list = _make_stepped_list_num(start, step, stepped_list_length) + # args.append(param[0]) + # args.append(stepped_list) + # else: + # raise ValueError( + # f"You provided {len(param[1])} parameters {param[0]}, rather than 2." + # ) + + # return args, shape @validate_call(config={"arbitrary_types_allowed": True}) @@ -454,7 +504,7 @@ def step_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For concurrent \ trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ @@ -469,11 +519,12 @@ def step_scan( bluesky.plans.list_scan(det, *args, md=metadata). """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args(params, grid=False) + args, shape = _make_step_scan_args_and_shape(params, grid=False) + print(args) metadata = metadata or {} metadata["shape"] = shape - yield from bp.list_scan(tuple(detectors), *args, md=metadata) + yield from bp.list_scan(tuple(detectors), *tuple(args), md=metadata) @validate_call(config={"arbitrary_types_allowed": True}) @@ -485,7 +536,7 @@ def step_grid_scan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ @@ -502,7 +553,7 @@ def step_grid_scan( default (all axes but the first axis provided). """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args(params, grid=True) + args, shape = _make_step_scan_args_and_shape(params, grid=True) metadata = metadata or {} metadata["shape"] = shape @@ -520,7 +571,7 @@ def step_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For concurrent \ trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ @@ -535,7 +586,7 @@ def step_rscan( bluesky.plans.rel_list_scan(det, *args, md=metadata). """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args(params, grid=False) + args, shape = _make_step_scan_args_and_shape(params, grid=False) metadata = metadata or {} metadata["shape"] = shape @@ -551,7 +602,7 @@ def step_grid_rscan( ), ], params: Annotated[ - list[tuple[Movable, list[float | int]]], + Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ @@ -568,7 +619,7 @@ def step_grid_rscan( default (all axes but the first axis provided). """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args(params, grid=True) + args, shape = _make_step_scan_args_and_shape(params, grid=True) metadata = metadata or {} metadata["shape"] = shape diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 43bf8b21ec3..4335e3799bf 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -20,9 +20,7 @@ from dodal.devices.motors import Motor from dodal.plans.wrapped import ( - _make_list_scan_args, - _make_num_scan_args, - _make_step_scan_args, + _make_step_scan_args_and_shape, _make_stepped_list_num, _make_stepped_list_step, _round_list_elements, @@ -187,26 +185,26 @@ def test_count_with_no_detector_raise_error(run_engine: RunEngine): run_engine(count([])) -@pytest.mark.parametrize( - "x_list, y_list, num, final_shape, final_length", - ( - [[0.0, 1.1], [2.2, 3.3], 3, [3], 6], - [[0.0, 1.1, 2], [2.2, 3.3, 3], None, [2, 3], 8], - ), -) -def test_make_num_scan_args( - x_axis: Motor, - y_axis: Motor, - x_list: list[float | int], - y_list: list[float | int], - num: int | None, - final_shape: list[int], - final_length: int, -): - args, shape = _make_num_scan_args([(x_axis, x_list), (y_axis, y_list)], num=num) - assert shape == final_shape - assert len(args) == final_length - assert args[0] == x_axis +# @pytest.mark.parametrize( +# "x_list, y_list, num, final_shape, final_length", +# ( +# [[0.0, 1.1], [2.2, 3.3], 3, [3], 6], +# [[0.0, 1.1, 2], [2.2, 3.3, 3], None, [2, 3], 8], +# ), +# ) +# def test_make_num_scan_args( +# x_axis: Motor, +# y_axis: Motor, +# x_list: list[float | int], +# y_list: list[float | int], +# num: int | None, +# final_shape: list[int], +# final_length: int, +# ): +# # args, shape = _make_num_scan_args([(x_axis, x_list), (y_axis, y_list)], num=num) +# assert shape == final_shape +# assert len(args) == final_length +# assert args[0] == x_axis def _assert_emitted( @@ -254,8 +252,9 @@ def test_num_scan_with_one_axis( x_list: list[float | int], num: int, ): - run_engine(num_scan(detectors=detectors, params=[(x_axis, x_list)], num=num)) + run_engine(num_scan(detectors=detectors, params=[x_axis, *x_list], num=num)) _assert_emitted(run_engine_documents, detectors, num) + print(run_engine_documents["start"][0]["shape"]) @pytest.mark.parametrize( @@ -274,7 +273,7 @@ def test_num_scan_with_two_axes( run_engine( num_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], num=num, ) ) @@ -285,7 +284,7 @@ def test_num_scan_fails_when_given_wrong_number_of_params( run_engine: RunEngine, detectors: Sequence[StandardDetector], x_axis: Motor ): with pytest.raises(ValueError): - run_engine(num_scan(detectors=detectors, params=[(x_axis, [-1, 1, 5])], num=5)) + run_engine(num_scan(detectors=detectors, params=[x_axis, -1, 1, 5], num=5)) @pytest.mark.parametrize( @@ -305,7 +304,7 @@ def test_num_scan_fails_when_given_bad_info( run_engine( num_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], num=num, ) ) @@ -327,7 +326,7 @@ def test_num_grid_scan( run_engine( num_grid_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], ) ) _assert_emitted(run_engine_documents, detectors, num) @@ -349,7 +348,7 @@ def test_num_grid_scan_when_not_snaking( run_engine( num_grid_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=False, ) ) @@ -364,9 +363,7 @@ def test_num_grid_scan_fails_when_given_wrong_number_of_params( ): with pytest.raises(ValueError): run_engine( - num_grid_scan( - detectors=detectors, params=[(x_axis, [0, 1.1, 2]), (y_axis, [1.1])] - ) + num_grid_scan(detectors=detectors, params=[x_axis, 0, 1.1, 2, y_axis, 1.1]) ) @@ -385,7 +382,7 @@ def test_num_scan_fails_when_asked_to_snake_slow_axis( run_engine( num_grid_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=[x_axis], ) ) @@ -400,7 +397,7 @@ def test_num_rscan( x_list: list[float | int], num: int, ): - run_engine(num_rscan(detectors=detectors, params=[(x_axis, x_list)], num=num)) + run_engine(num_rscan(detectors=detectors, params=[x_axis, *x_list], num=num)) _assert_emitted(run_engine_documents, detectors, num) @@ -419,7 +416,7 @@ def test_num_rscan_with_two_axes( ): run_engine( num_rscan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)], num=num + detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list], num=num ) ) _assert_emitted(run_engine_documents, detectors, num) @@ -441,7 +438,7 @@ def test_num_rscan_fails_when_given_bad_info( run_engine( num_rscan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], num=num, ) ) @@ -463,7 +460,7 @@ def test_num_grid_rscan( run_engine( num_grid_rscan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], ) ) _assert_emitted(run_engine_documents, detectors, num) @@ -485,7 +482,7 @@ def test_num_grid_rscan_when_not_snaking( run_engine( num_grid_rscan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=False, ) ) @@ -507,7 +504,7 @@ def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( run_engine( num_grid_rscan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=[x_axis], ) ) @@ -526,8 +523,8 @@ def test_make_list_scan_args( final_shape: list, final_length: int, ): - args, shape = _make_list_scan_args( - params=[(x_axis, x_list), (y_axis, y_list)], grid=grid + args, shape = _make_step_scan_args_and_shape( + params=[x_axis, *x_list, y_axis, *y_list], grid=grid ) assert len(args) == final_length assert shape == final_shape @@ -538,8 +535,8 @@ def test_make_list_scan_args_fails_when_lists_are_different_lengths( y_axis: Motor, ): with pytest.raises(ValueError): - _make_list_scan_args( - params=[(x_axis, [0, 1, 2]), (y_axis, [0, 1, 2, 3])], grid=False + _make_step_scan_args_and_shape( + params=[x_axis, 0, 1, 2, y_axis, 0, 1, 2, 3], grid=False ) @@ -553,7 +550,7 @@ def test_list_scan( ): num = int(len(x_list)) - run_engine(list_scan(detectors=detectors, params=[(x_axis, x_list)])) + run_engine(list_scan(detectors=detectors, params=[x_axis, x_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -574,9 +571,7 @@ def test_list_scan_with_two_axes( y_list: list, ): num = int(len(x_list)) - run_engine( - list_scan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) - ) + run_engine(list_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -590,7 +585,7 @@ def test_list_scan_fails_with_differnt_list_lengths( run_engine( list_scan( detectors=detectors, - params=[(x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4])], + params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], ) ) @@ -613,7 +608,7 @@ def test_list_grid_scan( ): num = int(len(x_list) * len(y_list)) run_engine( - list_grid_scan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) + list_grid_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -627,7 +622,7 @@ def test_list_rscan( x_list: list, ): num = int(len(x_list)) - run_engine(list_rscan(detectors=detectors, params=[(x_axis, x_list)])) + run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -649,9 +644,7 @@ def test_list_rscan_with_two_axes( ): num = int(len(x_list)) - run_engine( - list_rscan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) - ) + run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -665,7 +658,7 @@ def test_list_rscan_fails_with_differnt_list_lengths( run_engine( list_rscan( detectors=detectors, - params=[(x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4])], + params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], ) ) @@ -689,9 +682,7 @@ def test_list_grid_rscan( num = int(len(x_list) * len(y_list)) run_engine( - list_grid_rscan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)] - ) + list_grid_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -756,10 +747,25 @@ def test_make_stepped_list_fails_when_given_equal_start_and_stop_values(): @pytest.mark.parametrize( "x_list, y_list, grid, final_shape, final_length", ( - [[0, 1, 0.25], [0, 0.1], False, [5], 4], - [[0, 1, 0.25], [0, 1, 0.2], True, [5, 6], 4], - [[0, -1, -0.25], [0, -0.1], False, [5], 4], - [[0, -1, -0.25], [0, -1, -0.2], True, [5, 6], 4], + [[0, 1, 0.25], [0, 0.1], False, (5,), 4], + [ + [0, 1, 0.25], + [0, 1, 0.2], + True, + (5, 6), + 4, + ], + [[0, -1, -0.25], [0, -0.1], False, (5), 4], + [ + [0, -1, -0.25], + [0, -1, -0.2], + True, + ( + 5, + 6, + ), + 4, + ], ), ) def test_make_step_scan_args( @@ -771,13 +777,14 @@ def test_make_step_scan_args( final_shape: list, final_length: int, ): - args, shape = _make_step_scan_args( - params=[(x_axis, x_list), (y_axis, y_list)], grid=grid + args, shape = _make_step_scan_args_and_shape( + params=[x_axis, *x_list, y_axis, *y_list], grid=grid ) assert shape == final_shape - assert len(args) == final_length - assert args[0] == x_axis - assert args[2] == y_axis + assert args == [x_axis] + # assert len(args) == final_length + # assert args[0] == x_axis + # assert args[2] == y_axis @pytest.mark.parametrize( @@ -791,16 +798,16 @@ def test_make_step_scan_args( ) def test_make_step_scan_args_fails_when_given_incorrect_number_of_parameters( x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], z_axis: Motor, - z_list: list, + z_list: list[float | int], grid: bool, ): with pytest.raises(ValueError): - _make_step_scan_args( - params=[(x_axis, x_list), (y_axis, y_list), (z_axis, z_list)], grid=grid + _make_step_scan_args_and_shape( + params=[x_axis, *x_list, y_axis, *y_list, z_axis, *z_list], grid=grid ) @@ -812,10 +819,10 @@ def test_step_scan( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], num, ): - run_engine(step_scan(detectors=detectors, params=[(x_axis, x_list)])) + run_engine(step_scan(detectors=detectors, params=[x_axis, *x_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -832,13 +839,13 @@ def test_step_scan_with_multiple_axes( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], num, ): run_engine( - step_scan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) + step_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -856,13 +863,13 @@ def test_step_grid_scan( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], num, ): run_engine( - step_grid_scan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) + step_grid_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -879,15 +886,15 @@ def test_step_grid_scan_when_not_snaking( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], num, ): run_engine( step_grid_scan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=False, ) ) @@ -901,14 +908,14 @@ def test_step_grid_scan_fails_when_given_incorrect_number_of_params( run_engine: RunEngine, detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], ): with pytest.raises(ValueError): run_engine( step_grid_scan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)] + detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] ) ) @@ -921,10 +928,10 @@ def test_step_rscan( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, - num, + x_list: list[float | int], + num: int, ): - run_engine(step_rscan(detectors=detectors, params=[(x_axis, x_list)])) + run_engine(step_rscan(detectors=detectors, params=[x_axis, *x_list])) _assert_emitted(run_engine_documents, detectors, num) @@ -941,13 +948,13 @@ def test_step_rscan_with_multiple_axes( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, - num, + y_list: list[float | int], + num: int, ): run_engine( - step_rscan(detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)]) + step_rscan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -965,15 +972,13 @@ def test_step_grid_rscan( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, - num, + y_list: list[float | int], + num: int, ): run_engine( - step_grid_rscan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)] - ) + step_grid_rscan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) @@ -990,15 +995,15 @@ def test_step_grid_rscan_when_not_snaking( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, - num, + y_list: list[float | int], + num: int, ): run_engine( step_grid_rscan( detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], + params=[x_axis, *x_list, y_axis, *y_list], snake_axes=False, ) ) @@ -1012,13 +1017,13 @@ def test_step_grid_rscan_fails_when_given_incorrect_number_of_params( run_engine: RunEngine, detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], ): with pytest.raises(ValueError): run_engine( step_grid_rscan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)] + detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] ) ) From 9ad324aa4fde7a53ab7f4344fa808443ba7fc88b Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Thu, 7 May 2026 10:07:38 +0000 Subject: [PATCH 02/26] Update step scan logic, tests pass --- src/dodal/plans/wrapped.py | 214 ++++++++++++++---------------------- tests/plans/test_wrapped.py | 51 +-------- 2 files changed, 88 insertions(+), 177 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 043f7ad922f..714506a2682 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -1,6 +1,6 @@ from collections.abc import Sequence from decimal import Decimal -from typing import Annotated, Any +from typing import Annotated, Any, TypeVar import bluesky.plans as bp import numpy as np @@ -25,6 +25,8 @@ - Limits and metadata (e.g. units). """ +T = TypeVar("T") + @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) @@ -62,36 +64,6 @@ def count( yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) -# def _make_num_scan_args( -# params: Sequence[Movable | float | int], num: int | None = None -# ) -> list[float]: -# # shape = [] -# # if num: -# # shape = [num] -# # for param in params: -# # if len(param[1]) == 2: -# # pass -# # else: -# # raise ValueError("You must provide 'start stop' for each motor.") -# # else: -# # for param in params: -# # if len(param[1]) == 3: -# # shape.append(param[1][-1]) -# # else: -# # raise ValueError( -# # "You must provide 'start stop num' for each motor in a grid scan." -# # ) - -# # args = [] -# # for param in params: -# # args.append(param[0]) -# # args.extend(param[1]) -# for param in params: -# if isinstance(param, Movable): - -# return args, shape - - @validate_call(config={"arbitrary_types_allowed": True}) def num_scan( detectors: Annotated[ @@ -392,108 +364,85 @@ def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | def _make_step_scan_args_and_shape( - params: Sequence[Movable | float | int], grid: bool -) -> tuple[Sequence[Movable | list[float | int]], tuple[int, ...]]: - - args: list[Movable | list[float | int]] = [] - shape = [] - stepped_list_length = 0 - - try: - for i, param in enumerate(params): - if isinstance(param, Movable): - movable = param - start = params[i + 1] - if not isinstance(start, (float, int)): - raise ValueError( - f"You provided movable {movable} with no start, stop, step." - ) - - stop = params[i + 2] - if not isinstance(stop, (float, int)): - raise ValueError( - f"You provided movable {movable} with start value {start} but no stop and step." - ) - - step = params[i + 3] - if not isinstance(step, (float, int)): - raise ValueError( - f"You provided movable {movable} with start value {start}, stop value {stop} but no step value." - ) - - movable_values = _make_stepped_list_step(start, stop, step) - stepped_list_length = len(movable_values) - args.append(movable) - args.append(movable_values) - - if not grid: - break - - if not grid: - # Skip first 4 values as already done them. - for i, param in enumerate(params[4:]): - if isinstance(param, Movable): - movable = param - start = params[i + 1] - if not isinstance(start, (float, int)): - raise ValueError( - f"You provided movable {movable} with no start, stop, step." - ) - - step = params[i + 2] - if not isinstance(step, (float, int)): - raise ValueError( - f"You provided movable {movable} with start value {start}, stop value {step} but no step value." - ) - - movable_values = _make_stepped_list_num( - start, step, stepped_list_length - ) - args.append(movable) - args.append(movable_values) - # shape.append(len(values)) - except IndexError as e: - raise ValueError("Incorrect parameters provided.") from e + params: Sequence[Movable | float | int], + grid: bool, +) -> tuple[list[Movable | list[float]], tuple[int, ...]]: + + def require( + value: object, + expected: type[T] | tuple[type, ...], + name: str, + ) -> T: + expected_tuple = expected if isinstance(expected, tuple) else (expected,) + if not isinstance(value, expected_tuple): + allowed = ", ".join(t.__name__ for t in expected_tuple) + raise ValueError( + f"Parameter {name} must be one of type ({allowed}), got {type(value).__name__}" + ) + return value # type: ignore[return-value] + + def parse_full_axis( + values: Sequence[Movable | float | int], + ) -> tuple[Movable, float, float, float]: + if len(values) != 4: + raise ValueError( + f"Full axis must be movable, start, stop, step. You provided {values}" + ) + movable = require(values[0], Movable, "movable") + start = require(values[1], (int, float), "start") + stop = require(values[2], (int, float), "stop") + step = require(values[3], (int, float), "step") + + return movable, start, stop, step + + def parse_relative_axis( + values: Sequence[Movable | float | int], + ) -> tuple[Movable, float, float]: + if len(values) != 3: + raise ValueError( + f"Relative axis must be movable, start, step. You provided {values}" + ) + movable = require(values[0], Movable, "movable") + start = require(values[1], (int, float), "start") + step = require(values[2], (int, float), "step") + + return movable, start, step + + if len(params) < 4: + raise ValueError("At least one axis must provide (movable, start, stop, step)") + + args: list[Movable | list[float]] = [] + shape: list[int] = [] + + # First axis defines scan length + movable, start, stop, step = parse_full_axis(params[:4]) + + values = _make_stepped_list_step(start, stop, step) + stepped_list_length = len(values) + + args.extend([movable, values]) + shape.append(stepped_list_length) + + remaining = params[4:] + + chunk_size = 4 if grid else 3 + + if len(remaining) % chunk_size != 0: + raise ValueError("Incorrect number of parameters for additional axes") + + for i in range(0, len(remaining), chunk_size): + chunk = remaining[i : i + chunk_size] + if grid: + movable, start, stop, step = parse_full_axis(chunk) + values = _make_stepped_list_step(start, stop, step) + shape.append(len(values)) + else: + movable, start, step = parse_relative_axis(chunk) + values = _make_stepped_list_num(start, step, stepped_list_length) + args.extend([movable, values]) return args, tuple(shape) - # first_movable_param, *additional_movable_params = params - # if len(first_movable_param[1]) == 3: - # start, stop, step = first_movable_param[1] - # stepped_list = _make_stepped_list_step(start, stop, step) - # stepped_list_length = len(stepped_list) - # args.append(first_movable_param[0]) - # args.append(stepped_list) - # shape.append(stepped_list_length) - # else: - # raise ValueError( - # f"You provided {len(first_movable_param[1])} parameters for {first_movable_param[0]}, rather than 3." - # ) - # for param in additional_movable_params: - # if grid: - # if len(param[1]) == 3: - # start, stop, step = param[1] - # stepped_list = _make_stepped_list_step(start, stop, step) - # args.append(param[0]) - # args.append(stepped_list) - # shape.append(len(stepped_list)) - # else: - # raise ValueError( - # f"You provided {len(param[1])} parameters for {param[0]}, rather than 3." - # ) - # else: - # if len(param[1]) == 2: - # start, step = param[1] - # stepped_list = _make_stepped_list_num(start, step, stepped_list_length) - # args.append(param[0]) - # args.append(stepped_list) - # else: - # raise ValueError( - # f"You provided {len(param[1])} parameters {param[0]}, rather than 2." - # ) - - # return args, shape - @validate_call(config={"arbitrary_types_allowed": True}) def step_scan( @@ -520,11 +469,10 @@ def step_scan( """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 args, shape = _make_step_scan_args_and_shape(params, grid=False) - print(args) metadata = metadata or {} metadata["shape"] = shape - yield from bp.list_scan(tuple(detectors), *tuple(args), md=metadata) + yield from bp.list_scan(tuple(detectors), *tuple(args), md=metadata) # type: ignore @validate_call(config={"arbitrary_types_allowed": True}) @@ -623,6 +571,8 @@ def step_grid_rscan( metadata = metadata or {} metadata["shape"] = shape + print(args) + 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 4335e3799bf..3664a47181e 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -512,9 +512,12 @@ def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( @pytest.mark.parametrize( "x_list, y_list, grid, final_shape, final_length", - ([[0, 1, 2], [3, 4, 5], False, [3], 4], [[0, 1, 2], [3, 4, 5, 6], True, [3, 4], 4]), + ( + [[0, 10, 1], [0, 5], False, (11,), 4], + [[0, 10, 1], [0, 5, 1], True, (11, 6), 4], + ), ) -def test_make_list_scan_args( +def test_make_step_scan_args_and_shape( x_axis: Motor, x_list: list, y_axis: Motor, @@ -526,6 +529,7 @@ def test_make_list_scan_args( args, shape = _make_step_scan_args_and_shape( params=[x_axis, *x_list, y_axis, *y_list], grid=grid ) + print(args) assert len(args) == final_length assert shape == final_shape @@ -744,49 +748,6 @@ def test_make_stepped_list_fails_when_given_equal_start_and_stop_values(): _make_stepped_list_step(start=1.1, stop=1.1, step=0.25) -@pytest.mark.parametrize( - "x_list, y_list, grid, final_shape, final_length", - ( - [[0, 1, 0.25], [0, 0.1], False, (5,), 4], - [ - [0, 1, 0.25], - [0, 1, 0.2], - True, - (5, 6), - 4, - ], - [[0, -1, -0.25], [0, -0.1], False, (5), 4], - [ - [0, -1, -0.25], - [0, -1, -0.2], - True, - ( - 5, - 6, - ), - 4, - ], - ), -) -def test_make_step_scan_args( - x_axis: Motor, - x_list: list[float], - y_axis: Motor, - y_list: list[float], - grid: bool, - final_shape: list, - final_length: int, -): - args, shape = _make_step_scan_args_and_shape( - params=[x_axis, *x_list, y_axis, *y_list], grid=grid - ) - assert shape == final_shape - assert args == [x_axis] - # assert len(args) == final_length - # assert args[0] == x_axis - # assert args[2] == y_axis - - @pytest.mark.parametrize( "x_list, y_list, z_list, grid", ( From 8a476aa37e7b3e921ca0844a09f6f869926a8cf5 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Thu, 7 May 2026 10:22:27 +0000 Subject: [PATCH 03/26] Add error message to _make_stepped_list_num --- src/dodal/plans/wrapped.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 714506a2682..332b228099b 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -355,7 +355,7 @@ def _make_stepped_list_step( def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | int]: if num <= 0: - raise ValueError() + raise ValueError("Number of steps must be greater than zero.") stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements( stepped_list=stepped_list, params=[start, step] @@ -428,7 +428,7 @@ def parse_relative_axis( chunk_size = 4 if grid else 3 if len(remaining) % chunk_size != 0: - raise ValueError("Incorrect number of parameters for additional axes") + raise ValueError("Incorrect number of parameters for additional axes.") for i in range(0, len(remaining), chunk_size): chunk = remaining[i : i + chunk_size] From 7a948f81478e38158d0566452a8b133a57f4ff89 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 12:54:12 +0000 Subject: [PATCH 04/26] Update doc strings and _make_step_scan_args_and_shape logic --- src/dodal/plans/wrapped.py | 176 ++++++++++++++++++------------------- 1 file changed, 85 insertions(+), 91 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 332b228099b..5ff8cf5b218 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -237,9 +237,9 @@ def list_grid_scan( params: Annotated[ Sequence[Movable | list[float | int]], Field( - description="List of tuples (device, positions). For independent \ - trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ - [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'." + description="For independent trajectories, provide" + "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ..., " + "movableN, [point1, point2, ...]]'." ), ], snake_axes: bool = True, # Currently specifying axes to snake is not supported @@ -274,9 +274,9 @@ def list_rscan( params: Annotated[ Sequence[Movable | list[float | int]], Field( - description="List of tuples (device, positions). For concurrent \ - trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ - [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'. Number \ + description="For concurrent trajectories, provide " + "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ..., " + "movableN, [point1, point2, ...]]'. Number \ of points for each movable must be equal." ), ], @@ -303,9 +303,9 @@ def list_grid_rscan( params: Annotated[ Sequence[Movable | list[float | int]], Field( - description="List of tuples (device, positions). For independent \ - trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ - [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'." + description="For independent trajectories, provide " + "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ... , " + "movableN, [point1, point2, ...]]'." ), ], snake_axes: bool = True, # Currently specifying axes to snake is not supported @@ -363,85 +363,81 @@ def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | return rounded_stepped_list -def _make_step_scan_args_and_shape( - params: Sequence[Movable | float | int], - grid: bool, -) -> tuple[list[Movable | list[float]], tuple[int, ...]]: - - def require( - value: object, - expected: type[T] | tuple[type, ...], - name: str, - ) -> T: - expected_tuple = expected if isinstance(expected, tuple) else (expected,) - if not isinstance(value, expected_tuple): - allowed = ", ".join(t.__name__ for t in expected_tuple) - raise ValueError( - f"Parameter {name} must be one of type ({allowed}), got {type(value).__name__}" - ) - return value # type: ignore[return-value] - - def parse_full_axis( - values: Sequence[Movable | float | int], - ) -> tuple[Movable, float, float, float]: - if len(values) != 4: - raise ValueError( - f"Full axis must be movable, start, stop, step. You provided {values}" - ) - movable = require(values[0], Movable, "movable") - start = require(values[1], (int, float), "start") - stop = require(values[2], (int, float), "stop") - step = require(values[3], (int, float), "step") - - return movable, start, stop, step - - def parse_relative_axis( - values: Sequence[Movable | float | int], - ) -> tuple[Movable, float, float]: - if len(values) != 3: - raise ValueError( - f"Relative axis must be movable, start, step. You provided {values}" - ) - movable = require(values[0], Movable, "movable") - start = require(values[1], (int, float), "start") - step = require(values[2], (int, float), "step") - - return movable, start, step - - if len(params) < 4: - raise ValueError("At least one axis must provide (movable, start, stop, step)") - - args: list[Movable | list[float]] = [] - shape: list[int] = [] +def require( + value: object, + expected: type[T] | tuple[type[T], ...], + name: str, +) -> T: + expected_tuple = expected if isinstance(expected, tuple) else (expected,) + if not isinstance(value, expected_tuple): + allowed = ", ".join(t.__name__ for t in expected_tuple) + raise ValueError( + f"Parameter {name} must be one of type ({allowed}), got {type(value).__name__}" + ) + return value # type: ignore[return-value] - # First axis defines scan length - movable, start, stop, step = parse_full_axis(params[:4]) - values = _make_stepped_list_step(start, stop, step) - stepped_list_length = len(values) +def parse_full_axis( + values: Sequence[Movable | float | int], +) -> tuple[Movable, float, float, float]: + if len(values) != 4: + raise ValueError( + f"The axis must be movable, start, stop, step. You provided {values}" + ) + movable = require(values[0], Movable, "movable") + start = require(values[1], (int, float), "start") + stop = require(values[2], (int, float), "stop") + step = require(values[3], (int, float), "step") + return movable, start, stop, step - args.extend([movable, values]) - shape.append(stepped_list_length) - remaining = params[4:] +def parse_relative_axis( + values: Sequence[Movable | float | int], +) -> tuple[Movable, float, float]: + if len(values) != 3: + raise ValueError( + f"The axis must be movable, start, step. You provided {', '.join(map(str, values))}" + ) + movable = require(values[0], Movable, "movable") + start = require(values[1], (int, float), "start") + step = require(values[2], (int, float), "step") + return movable, start, step - chunk_size = 4 if grid else 3 - if len(remaining) % chunk_size != 0: - raise ValueError("Incorrect number of parameters for additional axes.") +def _make_step_scan_args_and_shape( + params: Sequence[Movable | float | int], grid: bool +) -> tuple[list[Movable | list[float]], tuple[int, ...]]: + """Convert [x, 1, 4, 1, ...] to [x, [1, 2, 3, 4], ...].""" + list_of_movable_with_values: list[list[Movable | float | int]] = [] + current_list: list[Movable | float | int] = [] + for param in params: + if isinstance(param, Movable): + current_list = [param] + list_of_movable_with_values.append(current_list) + elif isinstance(param, (int, float)): + current_list.append(param) + else: + raise ValueError( + f'Scan syntax only takes movables or numbers for params. You provided "{param}".' + ) - for i in range(0, len(remaining), chunk_size): - chunk = remaining[i : i + chunk_size] - if grid: - movable, start, stop, step = parse_full_axis(chunk) - values = _make_stepped_list_step(start, stop, step) - shape.append(len(values)) + step_scan_args: list[Movable | list[float]] = [] + shape = [] + first_axis = True + for movable_with_values in list_of_movable_with_values: + if first_axis or grid: + movable, start, stop, step = parse_full_axis(movable_with_values) + movable_values = _make_stepped_list_step(start, stop, step) + shape.append(len(movable_values)) + first_axis = False else: - movable, start, step = parse_relative_axis(chunk) - values = _make_stepped_list_num(start, step, stepped_list_length) - args.extend([movable, values]) + # If not a grid scan, expects start, stop for all other axes and use the + # first axis shape for the number of steps. + movable, start, step = parse_relative_axis(movable_with_values) + movable_values = _make_stepped_list_num(start, step, shape[0]) + step_scan_args.extend([movable, movable_values]) - return args, tuple(shape) + return step_scan_args, tuple(shape) @validate_call(config={"arbitrary_types_allowed": True}) @@ -455,9 +451,9 @@ def step_scan( params: Annotated[ Sequence[Movable | float | int], Field( - description="List of tuples (device, parameter). For concurrent \ - trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ - [start2, step2]), ... , (movableN, [startN, stepN])]'." + description="For concurrent trajectories, provide " + "'[movable1, start1, stop1, step1, movable2, start2, step2, ... , " + "movableN, startN, stepN]'." ), ], metadata: dict[str, Any] | None = None, @@ -487,8 +483,8 @@ def step_grid_scan( Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ - trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ - [start2, stop2, step2]), ... , (movableN, [startN, stopN, stepN])]'." + trajectories, provide '[movable1, start1, stop1, step1, movable2, start2, " + "stop2, step2, ... , movableN, startN, stopN, stepN]'." ), ], snake_axes: bool = True, # Currently specifying axes to snake is not supported @@ -521,9 +517,9 @@ def step_rscan( params: Annotated[ Sequence[Movable | float | int], Field( - description="List of tuples (device, parameter). For concurrent \ - trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ - [start2, step2]), ... , (movableN, [startN, stepN])]'." + description="For concurrent trajectories, provide " + "'[movable1, start1, stop1, step1, movable2, start2, step2, ... , " + "movableN, startN, stepN]'." ), ], metadata: dict[str, Any] | None = None, @@ -553,8 +549,8 @@ def step_grid_rscan( Sequence[Movable | float | int], Field( description="List of tuples (device, parameter). For independent \ - trajectories, provide '[(movable1, [start1, stop1, step1]), (movable2, \ - [start2, stop2, step2]), ... , (movableN, [startN, stopN, stepN])]'." + trajectories, provide '[movable1, start1, stop1, step1, movable2, \ + start2, stop2, step2, ... , movableN, startN, stopN, stepN]'." ), ], snake_axes: bool = True, # Currently specifying axes to snake is not supported @@ -571,8 +567,6 @@ def step_grid_rscan( metadata = metadata or {} metadata["shape"] = shape - print(args) - yield from bp.rel_list_grid_scan( tuple(detectors), *args, snake_axes=snake_axes, md=metadata ) From 4473ec5e10a4bf878b327818b5b9c97dcb4da2dc Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 12:54:24 +0000 Subject: [PATCH 05/26] Add stronger tests for invalid arguments --- tests/plans/test_wrapped.py | 89 ++++++++++++++++++++++++++----------- 1 file changed, 64 insertions(+), 25 deletions(-) diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 3664a47181e..140295c2587 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -185,28 +185,6 @@ def test_count_with_no_detector_raise_error(run_engine: RunEngine): run_engine(count([])) -# @pytest.mark.parametrize( -# "x_list, y_list, num, final_shape, final_length", -# ( -# [[0.0, 1.1], [2.2, 3.3], 3, [3], 6], -# [[0.0, 1.1, 2], [2.2, 3.3, 3], None, [2, 3], 8], -# ), -# ) -# def test_make_num_scan_args( -# x_axis: Motor, -# y_axis: Motor, -# x_list: list[float | int], -# y_list: list[float | int], -# num: int | None, -# final_shape: list[int], -# final_length: int, -# ): -# # args, shape = _make_num_scan_args([(x_axis, x_list), (y_axis, y_list)], num=num) -# assert shape == final_shape -# assert len(args) == final_length -# assert args[0] == x_axis - - def _assert_emitted( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], @@ -971,10 +949,30 @@ def test_step_grid_rscan_when_not_snaking( _assert_emitted(run_engine_documents, detectors, num) +@pytest.mark.parametrize("x_list, y_list", ([[0, 1], [0, 1, 0.1]], [[0], [0, 1, 0.1]])) +def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axes( + run_engine: RunEngine, + detectors: Sequence[StandardDetector], + x_axis: Motor, + x_list: list[float | int], + y_axis: Motor, + y_list: list[float | int], +): + with pytest.raises( + ValueError, + match="The axis must be movable, start, stop, step.", + ): + run_engine( + step_grid_scan( + detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] + ) + ) + + @pytest.mark.parametrize( "x_list, y_list", ([[0, 1, 0.1], [0, 1, 0.1, 1]], [[0, 1, 0.1], [0]]) ) -def test_step_grid_rscan_fails_when_given_incorrect_number_of_params( +def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_second_axes( run_engine: RunEngine, detectors: Sequence[StandardDetector], x_axis: Motor, @@ -982,9 +980,50 @@ def test_step_grid_rscan_fails_when_given_incorrect_number_of_params( y_axis: Motor, y_list: list[float | int], ): - with pytest.raises(ValueError): + with pytest.raises( + ValueError, + match="The axis must be movable, start, stop, step.", + ): run_engine( - step_grid_rscan( + step_grid_scan( detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] ) ) + + +@pytest.mark.parametrize( + "x_list, y_list", ([[0, 1, 0.1], [0, 1, 0.1]], [[0, 1, 0.1], [0]]) +) +def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( + run_engine: RunEngine, + detectors: Sequence[StandardDetector], + x_axis: Motor, + x_list: list[float | int], + y_axis: Motor, + y_list: list[float | int], +): + with pytest.raises( + ValueError, + match="The axis must be movable, start, stop.", + ): + run_engine( + step_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) + ) + + +def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises( + ValueError, + match="Scan syntax only takes movables or numbers for params.", + ): + _make_step_scan_args_and_shape( + [x_axis, 1, "3", 1, y_axis, 1, "4", 1], # type: ignore + grid=True, + ) + _make_step_scan_args_and_shape( + [x_axis, 1, "3", 1, y_axis, 1, "4"], # type: ignore + grid=False, + ) From d93ba1a75126da32807bd0e35b1ce80ae7e2c707 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 13:44:22 +0000 Subject: [PATCH 06/26] Added additional test to check for scan shape --- src/dodal/plans/wrapped.py | 25 +++--- tests/plans/test_wrapped.py | 150 +++++++++++++++++------------------- 2 files changed, 82 insertions(+), 93 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 5ff8cf5b218..9e789638ba2 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -185,14 +185,19 @@ def num_grid_rscan( def _make_list_scan_shape( - params: Sequence[Movable | list[float | int]], + params: Sequence[Movable | list[float | int]], grid: bool ) -> tuple[int, ...]: + shape = [] for param in params: # List arg must all be same size. If list missing or not same size, this will # be validated by bp.list_scan. if isinstance(param, list): - return (len(param),) - return () + dim = len(param) + shape.append(dim) + if not grid: + break + + return tuple(shape) @validate_call(config={"arbitrary_types_allowed": True}) @@ -220,7 +225,7 @@ def list_scan( Wraps bluesky.plans.list_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params) + metadata["shape"] = _make_list_scan_shape(params, grid=False) # Not sure about this one yield from bp.list_scan(tuple(detectors), *tuple(params), md=metadata) # type: ignore @@ -252,11 +257,7 @@ def list_grid_scan( bluesky.plans.list_grid_scan(det, *args, md=metadata). """ metadata = metadata or {} - shape = [] - for param in params: - if isinstance(param, list): - shape.append(len(param)) - metadata["shape"] = tuple(shape) + metadata["shape"] = _make_list_scan_shape(params, grid=False) yield from bp.list_grid_scan( tuple(detectors), *params, snake_axes=snake_axes, md=metadata @@ -288,7 +289,7 @@ def list_rscan( Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params) + metadata["shape"] = _make_list_scan_shape(params, grid=False) yield from bp.rel_list_scan(tuple(detectors), *params, md=metadata) @@ -318,7 +319,7 @@ def list_grid_rscan( bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params) + metadata["shape"] = _make_list_scan_shape(params, grid=True) yield from bp.rel_list_grid_scan( tuple(detectors), *params, snake_axes=snake_axes, md=metadata ) @@ -396,7 +397,7 @@ def parse_relative_axis( ) -> tuple[Movable, float, float]: if len(values) != 3: raise ValueError( - f"The axis must be movable, start, step. You provided {', '.join(map(str, values))}" + f"The axis must be movable, start, stop. You provided {', '.join(map(str, values))}" ) movable = require(values[0], Movable, "movable") start = require(values[1], (int, float), "start") diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 140295c2587..836c6f314da 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -40,6 +40,13 @@ ) +def assert_expected_shape( + run_engine_documents: Mapping[str, list[dict]], expected_shape: tuple[int, ...] +) -> None: + start = run_engine_documents["start"][0] + assert start["shape"] == expected_shape + + def test_count_delay_validation(det: StandardDetector, run_engine: RunEngine): args: dict[float | Sequence[float], str] = { # type: ignore # List wrong length @@ -59,7 +66,6 @@ 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)) - print(delay) def test_count_detectors_validation(run_engine: RunEngine): @@ -97,10 +103,10 @@ def test_count_plan_produces_expected_start_document( start = run_engine_documents.get("start") assert start and len(start) == 1 run_start = cast(RunStart, start[0]) - assert run_start.get("shape") == shape assert (hints := run_start.get("hints")) and ( hints.get("dimensions") == [(("time",), "primary")] ) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) @@ -232,7 +238,7 @@ def test_num_scan_with_one_axis( ): run_engine(num_scan(detectors=detectors, params=[x_axis, *x_list], num=num)) _assert_emitted(run_engine_documents, detectors, num) - print(run_engine_documents["start"][0]["shape"]) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -256,6 +262,7 @@ def test_num_scan_with_two_axes( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) def test_num_scan_fails_when_given_wrong_number_of_params( @@ -289,16 +296,16 @@ def test_num_scan_fails_when_given_bad_info( @pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) + "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) ) def test_num_grid_scan( run_engine: RunEngine, run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list[float | int], + x_list: tuple[float, float, int], y_axis: Motor, - y_list: list[float | int], + y_list: tuple[float, float, int], ): num = int(x_list[-1] * y_list[-1]) run_engine( @@ -308,19 +315,20 @@ def test_num_grid_scan( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) @pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) + "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) ) def test_num_grid_scan_when_not_snaking( run_engine: RunEngine, run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list[float | int], + x_list: tuple[float, float, int], y_axis: Motor, - y_list: list[float | int], + y_list: tuple[float, float, int], ): num = int(x_list[-1] * y_list[-1]) run_engine( @@ -331,6 +339,7 @@ def test_num_grid_scan_when_not_snaking( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) def test_num_grid_scan_fails_when_given_wrong_number_of_params( @@ -377,6 +386,7 @@ def test_num_rscan( ): run_engine(num_rscan(detectors=detectors, params=[x_axis, *x_list], num=num)) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -398,6 +408,7 @@ def test_num_rscan_with_two_axes( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -423,16 +434,16 @@ def test_num_rscan_fails_when_given_bad_info( @pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) + "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) ) def test_num_grid_rscan( run_engine: RunEngine, run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list[float | int], + x_list: tuple[float, float, int], y_axis: Motor, - y_list: list[float | int], + y_list: tuple[float, float, int], ): num = int(x_list[-1] * y_list[-1]) run_engine( @@ -442,19 +453,20 @@ def test_num_grid_rscan( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) @pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) + "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) ) def test_num_grid_rscan_when_not_snaking( run_engine: RunEngine, run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list[float | int], + x_list: tuple[float, float, int], y_axis: Motor, - y_list: list[float | int], + y_list: tuple[float, float, int], ): num = int(x_list[-1] * y_list[-1]) run_engine( @@ -465,6 +477,7 @@ def test_num_grid_rscan_when_not_snaking( ) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) @pytest.mark.parametrize( @@ -507,7 +520,6 @@ def test_make_step_scan_args_and_shape( args, shape = _make_step_scan_args_and_shape( params=[x_axis, *x_list, y_axis, *y_list], grid=grid ) - print(args) assert len(args) == final_length assert shape == final_shape @@ -530,10 +542,10 @@ def test_list_scan( x_axis: Motor, x_list: list, ): - num = int(len(x_list)) - + num = len(x_list) run_engine(list_scan(detectors=detectors, params=[x_axis, x_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -555,6 +567,7 @@ def test_list_scan_with_two_axes( num = int(len(x_list)) run_engine(list_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) def test_list_scan_fails_with_differnt_list_lengths( @@ -584,15 +597,16 @@ def test_list_grid_scan( run_engine_documents: Mapping[str, list[dict]], detectors: Sequence[StandardDetector], x_axis: Motor, - x_list: list, + x_list: list[float | int], y_axis: Motor, - y_list: list, + y_list: list[float | int], ): num = int(len(x_list) * len(y_list)) run_engine( list_grid_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (len(x_list), len(y_list))) @pytest.mark.parametrize("x_list", ([0, 1, 2, 3], [1.1, 2.2, 3.3])) @@ -606,6 +620,7 @@ def test_list_rscan( num = int(len(x_list)) run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (len(x_list),)) @pytest.mark.parametrize( @@ -628,6 +643,7 @@ def test_list_rscan_with_two_axes( run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) def test_list_rscan_fails_with_differnt_list_lengths( @@ -667,6 +683,7 @@ def test_list_grid_rscan( list_grid_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (len(x_list), len(y_list))) @pytest.mark.parametrize( @@ -763,6 +780,7 @@ def test_step_scan( ): run_engine(step_scan(detectors=detectors, params=[x_axis, *x_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -787,14 +805,18 @@ def test_step_scan_with_multiple_axes( step_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( - "x_list, y_list, num", + "x_list, expected_num_x, y_list, expected_num_y, snake", ( - [[0, 1, 0.25], [0, 2, 0.5], 25], - [[-1, 1, 0.25], [1, -1, -0.5], 45], - [[0, 10, 2.5], [0, -10, -2.5], 25], + [[0, 1, 0.25], 5, [0, 2, 0.5], 5, True], + [[0, 1, 0.25], 5, [0, 2, 0.5], 5, False], + [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, True], + [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, False], + [[0, 10, 2.5], 5, [0, -10, -2.5], 5, True], + [[0, 10, 2.5], 5, [0, -10, -2.5], 5, False], ), ) def test_step_grid_scan( @@ -803,41 +825,21 @@ def test_step_grid_scan( detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], + expected_num_x: int, y_axis: Motor, y_list: list[float | int], - num, -): - run_engine( - step_grid_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - - -@pytest.mark.parametrize( - "x_list, y_list, num", - ( - [[0, 1, 0.25], [0, 2, 0.5], 25], - [[-1, 1, 0.25], [1, -1, -0.5], 45], - ), -) -def test_step_grid_scan_when_not_snaking( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num, + expected_num_y: int, + snake: bool, ): run_engine( step_grid_scan( detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=False, + snake_axes=snake, ) ) - _assert_emitted(run_engine_documents, detectors, num) + _assert_emitted(run_engine_documents, detectors, expected_num_x * expected_num_y) + assert_expected_shape(run_engine_documents, (expected_num_x, expected_num_y)) @pytest.mark.parametrize( @@ -860,7 +862,8 @@ def test_step_grid_scan_fails_when_given_incorrect_number_of_params( @pytest.mark.parametrize( - "x_list, num", ([[0, 1, 0.1], 11], [[-1, 1, 0.1], 21], [[0, 10, 1], 11]) + "x_list, num", + ([[0, 1, 0.1], 11], [[-1, 1, 0.1], 21], [[0, 10, 1], 11]), ) def test_step_rscan( run_engine: RunEngine, @@ -872,6 +875,7 @@ def test_step_rscan( ): run_engine(step_rscan(detectors=detectors, params=[x_axis, *x_list])) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( @@ -896,14 +900,18 @@ def test_step_rscan_with_multiple_axes( step_rscan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) ) _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) @pytest.mark.parametrize( - "x_list, y_list, num", + "x_list, expected_num_x, y_list, expected_num_y, snake", ( - [[0, 1, 0.25], [0, 2, 0.5], 25], - [[-1, 1, 0.25], [1, -1, -0.5], 45], - [[0, 10, 2.5], [0, -10, -2.5], 25], + [[0, 1, 0.25], 5, [0, 2, 0.5], 5, True], + [[0, 1, 0.25], 5, [0, 2, 0.5], 5, False], + [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, True], + [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, False], + [[0, 10, 2.5], 5, [0, -10, -2.5], 5, True], + [[0, 10, 2.5], 5, [0, -10, -2.5], 5, False], ), ) def test_step_grid_rscan( @@ -912,41 +920,21 @@ def test_step_grid_rscan( detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], + expected_num_x: int, y_axis: Motor, y_list: list[float | int], - num: int, -): - run_engine( - step_grid_rscan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - - -@pytest.mark.parametrize( - "x_list, y_list, num", - ( - [[0, 1, 0.25], [0, 2, 0.5], 25], - [[-1, 1, 0.25], [1, -1, -0.5], 45], - ), -) -def test_step_grid_rscan_when_not_snaking( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, + expected_num_y: int, + snake: bool, ): run_engine( step_grid_rscan( detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=False, + snake_axes=snake, ) ) - _assert_emitted(run_engine_documents, detectors, num) + _assert_emitted(run_engine_documents, detectors, expected_num_x * expected_num_y) + assert_expected_shape(run_engine_documents, (expected_num_x, expected_num_y)) @pytest.mark.parametrize("x_list, y_list", ([[0, 1], [0, 1, 0.1]], [[0], [0, 1, 0.1]])) From 650398268f0a644fe6ffd453430665d5d236b523 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 13:52:21 +0000 Subject: [PATCH 07/26] Fix shape test --- src/dodal/plans/wrapped.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 9e789638ba2..ceabf0ae66d 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -11,15 +11,16 @@ from dodal.common import MsgGenerator 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 -moved into bluesky or better handled in downstream services. - -Required decorators are installed on plan import +"""This module wraps plan(s) from bluesky.plans so they are comptaible with blueapi. +Required decorators are installed on plan import. https://github.com/DiamondLightSource/blueapi/issues/474 -Non-serialisable fields are ignored when they are optional +Non-serialisable fields are ignored when they are optional. https://github.com/DiamondLightSource/blueapi/issues/711 +Using *args in plans is currently not supported. +https://github.com/DiamondLightSource/blueapi/issues/1450 + We may also need other adjustments for UI purposes, e.g. - Forcing uniqueness or orderedness of Readables. - Limits and metadata (e.g. units). @@ -257,7 +258,7 @@ def list_grid_scan( bluesky.plans.list_grid_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params, grid=False) + metadata["shape"] = _make_list_scan_shape(params, grid=True) yield from bp.list_grid_scan( tuple(detectors), *params, snake_axes=snake_axes, md=metadata From ed2c7e2f24cbbdd49468d38c182fd8589713765c Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 14:34:59 +0000 Subject: [PATCH 08/26] Add test for require --- src/dodal/plans/wrapped.py | 6 ++++-- tests/plans/test_wrapped.py | 38 ++++++++++++++++++++++++++++++++++--- 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index ceabf0ae66d..da93e0983a8 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -343,6 +343,8 @@ def _make_stepped_list_step( raise ValueError( f"Start ({start}) and stop ({stop}) values cannot be the same." ) + if step <= 0: + raise ValueError("Step size must be greater than zero.") if abs(step) > abs(stop - start): step = stop - start step = abs(step) * np.sign(stop - start) @@ -357,7 +359,7 @@ def _make_stepped_list_step( def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | int]: if num <= 0: - raise ValueError("Number of steps must be greater than zero.") + raise ValueError("Number of points must be greater than zero.") stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements( stepped_list=stepped_list, params=[start, step] @@ -374,7 +376,7 @@ def require( if not isinstance(value, expected_tuple): allowed = ", ".join(t.__name__ for t in expected_tuple) raise ValueError( - f"Parameter {name} must be one of type ({allowed}), got {type(value).__name__}" + f"Parameter {name} must be one of type {allowed}, got {type(value).__name__}." ) return value # type: ignore[return-value] diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 836c6f314da..901bda4519b 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -1,3 +1,4 @@ +import re from collections.abc import Mapping, Sequence from typing import cast @@ -33,6 +34,7 @@ num_grid_scan, num_rscan, num_scan, + require, step_grid_rscan, step_grid_scan, step_rscan, @@ -738,9 +740,27 @@ def test_make_stepped_list_num(start: float, step: float): assert stepped_list[10] == 0 -def test_make_stepped_list_fails_when_given_equal_start_and_stop_values(): - with pytest.raises(ValueError): - _make_stepped_list_step(start=1.1, stop=1.1, step=0.25) +def test_make_stepped_list_num_fails_when_num_is_zero(): + start = stop = 1.1 + with pytest.raises( + ValueError, + match=re.escape( + f"Start ({start}) and stop ({stop}) values cannot be the same." + ), + ): + _make_stepped_list_step(start=start, stop=stop, step=0.25) + + +def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): + with pytest.raises(ValueError, match="Number of points must be greater than zero."): + _make_stepped_list_num(start=1, step=0.1, num=0) + + +def test_require_raises_error_if_not_correct_type(): + with pytest.raises( + ValueError, match="Parameter test must be one of type str, got int." + ): + require(value=5, expected=str, name="test") @pytest.mark.parametrize( @@ -1015,3 +1035,15 @@ def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( [x_axis, 1, "3", 1, y_axis, 1, "4"], # type: ignore grid=False, ) + + +def test_step_scan_fails_with_step_size_zero( + run_engine: RunEngine, + detectors: Sequence[StandardDetector], + x_axis: Motor, +): + with pytest.raises( + ValueError, + match="Step size must be greater than zero.", + ): + run_engine(step_scan(detectors=detectors, params=[x_axis, 1, 5, 0])) From 6ee9bfd07ea3fb3018d3e43ea754f98e6706726d Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 15:09:53 +0000 Subject: [PATCH 09/26] Fix and optimise tests --- src/dodal/plans/wrapped.py | 10 ++++--- tests/plans/test_wrapped.py | 54 +++++++++++++------------------------ 2 files changed, 25 insertions(+), 39 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index da93e0983a8..8d85520d8ef 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -343,8 +343,8 @@ def _make_stepped_list_step( raise ValueError( f"Start ({start}) and stop ({stop}) values cannot be the same." ) - if step <= 0: - raise ValueError("Step size must be greater than zero.") + if step == 0: + raise ValueError(f"Step size {step} cannot be zero.") if abs(step) > abs(stop - start): step = stop - start step = abs(step) * np.sign(stop - start) @@ -358,8 +358,10 @@ def _make_stepped_list_step( def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | int]: - if num <= 0: - raise ValueError("Number of points must be greater than zero.") + if num == 0 or step == 0: + raise ValueError( + f"Number of points ({num}) and number of steps ({step}) cannot be zero." + ) stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements( stepped_list=stepped_list, params=[start, step] diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 901bda4519b..3d169c40fe5 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -268,10 +268,10 @@ def test_num_scan_with_two_axes( def test_num_scan_fails_when_given_wrong_number_of_params( - run_engine: RunEngine, detectors: Sequence[StandardDetector], x_axis: Motor + run_engine: RunEngine, x_axis: Motor ): with pytest.raises(ValueError): - run_engine(num_scan(detectors=detectors, params=[x_axis, -1, 1, 5], num=5)) + run_engine(num_scan(detectors=[], params=[x_axis, -1, 1, 5], num=5)) @pytest.mark.parametrize( @@ -280,7 +280,6 @@ def test_num_scan_fails_when_given_wrong_number_of_params( ) def test_num_scan_fails_when_given_bad_info( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -290,7 +289,7 @@ def test_num_scan_fails_when_given_bad_info( with pytest.raises(ValueError): run_engine( num_scan( - detectors=detectors, + detectors=[], params=[x_axis, *x_list, y_axis, *y_list], num=num, ) @@ -346,14 +345,11 @@ def test_num_grid_scan_when_not_snaking( def test_num_grid_scan_fails_when_given_wrong_number_of_params( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, y_axis: Motor, ): with pytest.raises(ValueError): - run_engine( - num_grid_scan(detectors=detectors, params=[x_axis, 0, 1.1, 2, y_axis, 1.1]) - ) + run_engine(num_grid_scan(detectors=[], params=[x_axis, 0, 1.1, 2, y_axis, 1.1])) @pytest.mark.parametrize( @@ -361,7 +357,6 @@ def test_num_grid_scan_fails_when_given_wrong_number_of_params( ) def test_num_scan_fails_when_asked_to_snake_slow_axis( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -370,7 +365,7 @@ def test_num_scan_fails_when_asked_to_snake_slow_axis( with pytest.raises(ValueError): run_engine( num_grid_scan( - detectors=detectors, + detectors=[], params=[x_axis, *x_list, y_axis, *y_list], snake_axes=[x_axis], ) @@ -418,7 +413,6 @@ def test_num_rscan_with_two_axes( ) def test_num_rscan_fails_when_given_bad_info( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -428,7 +422,7 @@ def test_num_rscan_fails_when_given_bad_info( with pytest.raises(ValueError): run_engine( num_rscan( - detectors=detectors, + detectors=[], params=[x_axis, *x_list, y_axis, *y_list], num=num, ) @@ -487,7 +481,6 @@ def test_num_grid_rscan_when_not_snaking( ) def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -496,7 +489,7 @@ def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( with pytest.raises(ValueError): run_engine( num_grid_rscan( - detectors=detectors, + detectors=[], params=[x_axis, *x_list, y_axis, *y_list], snake_axes=[x_axis], ) @@ -574,14 +567,13 @@ def test_list_scan_with_two_axes( def test_list_scan_fails_with_differnt_list_lengths( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, y_axis: Motor, ): with pytest.raises(ValueError): run_engine( list_scan( - detectors=detectors, + detectors=[], params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], ) ) @@ -650,14 +642,13 @@ def test_list_rscan_with_two_axes( def test_list_rscan_fails_with_differnt_list_lengths( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, y_axis: Motor, ): with pytest.raises(ValueError): run_engine( list_rscan( - detectors=detectors, + detectors=[], params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], ) ) @@ -752,8 +743,11 @@ def test_make_stepped_list_num_fails_when_num_is_zero(): def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): - with pytest.raises(ValueError, match="Number of points must be greater than zero."): - _make_stepped_list_num(start=1, step=0.1, num=0) + with pytest.raises( + ValueError, + match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), + ): + _make_stepped_list_num(start=1, step=0, num=0) def test_require_raises_error_if_not_correct_type(): @@ -960,7 +954,6 @@ def test_step_grid_rscan( @pytest.mark.parametrize("x_list, y_list", ([[0, 1], [0, 1, 0.1]], [[0], [0, 1, 0.1]])) def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axes( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -971,9 +964,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axes( match="The axis must be movable, start, stop, step.", ): run_engine( - step_grid_scan( - detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] - ) + step_grid_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list]) ) @@ -982,7 +973,6 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axes( ) def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_second_axes( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -993,9 +983,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_second_axes( match="The axis must be movable, start, stop, step.", ): run_engine( - step_grid_scan( - detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] - ) + step_grid_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list]) ) @@ -1004,7 +992,6 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_second_axes( ) def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, x_list: list[float | int], y_axis: Motor, @@ -1014,9 +1001,7 @@ def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( ValueError, match="The axis must be movable, start, stop.", ): - run_engine( - step_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) - ) + run_engine(step_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list])) def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( @@ -1039,11 +1024,10 @@ def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( def test_step_scan_fails_with_step_size_zero( run_engine: RunEngine, - detectors: Sequence[StandardDetector], x_axis: Motor, ): with pytest.raises( ValueError, - match="Step size must be greater than zero.", + match="Step size 0 cannot be zero.", ): - run_engine(step_scan(detectors=detectors, params=[x_axis, 1, 5, 0])) + run_engine(step_scan(detectors=[], params=[x_axis, 1, 5, 0])) From 2c4587a719d390092f79fc6e54adcd9f89310670 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 8 May 2026 15:21:04 +0000 Subject: [PATCH 10/26] Update comments --- src/dodal/plans/wrapped.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index 8d85520d8ef..fd2db2b28fb 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -11,7 +11,7 @@ from dodal.common import MsgGenerator from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator -"""This module wraps plan(s) from bluesky.plans so they are comptaible with blueapi. +"""This module wraps plan(s) from bluesky.plans so they are compatible with blueapi. Required decorators are installed on plan import. https://github.com/DiamondLightSource/blueapi/issues/474 @@ -228,7 +228,6 @@ def list_scan( metadata = metadata or {} metadata["shape"] = _make_list_scan_shape(params, grid=False) - # Not sure about this one yield from bp.list_scan(tuple(detectors), *tuple(params), md=metadata) # type: ignore From 647d095cc3c08c1753ba848f019909f9eb484ae1 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Mon, 11 May 2026 07:43:39 +0000 Subject: [PATCH 11/26] Update doc strings and error msgs --- src/dodal/plans/wrapped.py | 9 +++++---- tests/plans/test_wrapped.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py index fd2db2b28fb..ef17aaad961 100644 --- a/src/dodal/plans/wrapped.py +++ b/src/dodal/plans/wrapped.py @@ -423,7 +423,8 @@ def _make_step_scan_args_and_shape( current_list.append(param) else: raise ValueError( - f'Scan syntax only takes movables or numbers for params. You provided "{param}".' + "Scan syntax only takes movables or numbers as parameters. " + f'You provided "{param}".' ) step_scan_args: list[Movable | list[float]] = [] @@ -553,9 +554,9 @@ def step_grid_rscan( params: Annotated[ Sequence[Movable | float | int], Field( - description="List of tuples (device, parameter). For independent \ - trajectories, provide '[movable1, start1, stop1, step1, movable2, \ - start2, stop2, step2, ... , movableN, startN, stopN, stepN]'." + description="For independent trajectories, provide " + "'[movable1, start1, stop1, step1, movable2, start2, stop2, step2, ... , " + "movableN, startN, stopN, stepN]'." ), ], snake_axes: bool = True, # Currently specifying axes to snake is not supported diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py index 3d169c40fe5..c27c208a4ba 100644 --- a/tests/plans/test_wrapped.py +++ b/tests/plans/test_wrapped.py @@ -1010,7 +1010,7 @@ def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( ): with pytest.raises( ValueError, - match="Scan syntax only takes movables or numbers for params.", + match="Scan syntax only takes movables or numbers as parameters.", ): _make_step_scan_args_and_shape( [x_axis, 1, "3", 1, y_axis, 1, "4", 1], # type: ignore From b30a406df2fdea1a0613115dae4f7f730312e5f0 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 10:02:48 +0000 Subject: [PATCH 12/26] Changes so far --- src/dodal/plans/__init__.py | 33 - src/dodal/plans/scans/__init__.py | 33 + src/dodal/plans/scans/annotations.py | 78 ++ src/dodal/plans/{ => scans}/spec_path.py | 11 +- src/dodal/plans/scans/utils.py | 88 ++ src/dodal/plans/scans/validators.py | 62 ++ src/dodal/plans/scans/wrapped.py | 357 ++++++++ src/dodal/plans/wrapped.py | 578 ------------ system_tests/test_adsim.py | 2 +- tests/plans/scans/__init__.py | 0 tests/plans/scans/conftest.py | 56 ++ tests/plans/scans/test_utils.py | 159 ++++ tests/plans/scans/test_wrapped.py | 722 +++++++++++++++ tests/plans/test_compliance.py | 5 +- tests/plans/test_scanspec.py | 2 +- tests/plans/test_wrapped.py | 1033 ---------------------- 16 files changed, 1563 insertions(+), 1656 deletions(-) create mode 100644 src/dodal/plans/scans/__init__.py create mode 100644 src/dodal/plans/scans/annotations.py rename src/dodal/plans/{ => scans}/spec_path.py (87%) create mode 100644 src/dodal/plans/scans/utils.py create mode 100644 src/dodal/plans/scans/validators.py create mode 100644 src/dodal/plans/scans/wrapped.py delete mode 100644 src/dodal/plans/wrapped.py create mode 100644 tests/plans/scans/__init__.py create mode 100644 tests/plans/scans/conftest.py create mode 100644 tests/plans/scans/test_utils.py create mode 100644 tests/plans/scans/test_wrapped.py delete mode 100644 tests/plans/test_wrapped.py diff --git a/src/dodal/plans/__init__.py b/src/dodal/plans/__init__.py index 50b6a976229..e69de29bb2d 100644 --- a/src/dodal/plans/__init__.py +++ b/src/dodal/plans/__init__.py @@ -1,33 +0,0 @@ -from .spec_path import spec_scan -from .wrapped import ( - count, - list_grid_rscan, - list_grid_scan, - list_rscan, - list_scan, - num_grid_rscan, - num_grid_scan, - num_rscan, - num_scan, - step_grid_rscan, - step_grid_scan, - step_rscan, - step_scan, -) - -__all__ = [ - "count", - "list_grid_rscan", - "list_grid_scan", - "list_rscan", - "list_scan", - "num_grid_rscan", - "num_grid_scan", - "num_rscan", - "num_scan", - "spec_scan", - "step_grid_rscan", - "step_grid_scan", - "step_rscan", - "step_scan", -] diff --git a/src/dodal/plans/scans/__init__.py b/src/dodal/plans/scans/__init__.py new file mode 100644 index 00000000000..50b6a976229 --- /dev/null +++ b/src/dodal/plans/scans/__init__.py @@ -0,0 +1,33 @@ +from .spec_path import spec_scan +from .wrapped import ( + count, + list_grid_rscan, + list_grid_scan, + list_rscan, + list_scan, + num_grid_rscan, + num_grid_scan, + num_rscan, + num_scan, + step_grid_rscan, + step_grid_scan, + step_rscan, + step_scan, +) + +__all__ = [ + "count", + "list_grid_rscan", + "list_grid_scan", + "list_rscan", + "list_scan", + "num_grid_rscan", + "num_grid_scan", + "num_rscan", + "num_scan", + "spec_scan", + "step_grid_rscan", + "step_grid_scan", + "step_rscan", + "step_scan", +] diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py new file mode 100644 index 00000000000..d2dd2b64167 --- /dev/null +++ b/src/dodal/plans/scans/annotations.py @@ -0,0 +1,78 @@ +from collections.abc import Sequence +from typing import Annotated as A +from typing import Any, TypeVar + +from bluesky.protocols import Movable, Readable +from ophyd_async.core import AsyncReadable +from pydantic import BeforeValidator, Field + +from dodal.plans.scans.validators import trajectory_validator, validate_start_stop_step + +Number = float | int +T = TypeVar("T") + +DetectorsA = A[ + Sequence[Readable | AsyncReadable], + Field( + description="Set of readable devices, will take a reading at each point", + ), +] + +MovableStartStep = tuple[Movable[Number], Number, Number] + +MovableStartStepA = A[ + MovableStartStep, + Field( + description="Additional trajectories, each specified as a tuple of " + "(movable, start, step)." + ), +] + +MovableStartStop = tuple[Movable[Number], Number, Number] + +MovableStartStopA = A[ + MovableStartStop, + Field( + description="Additional trajectories, each specified as a tuple of " + "(movable, start, stop)." + ), +] + +MovableStartStopNum = tuple[Movable[Number], Number, Number, int] + +MovableStartStopNumA = A[ + MovableStartStopNum, + Field( + description="Additional trajectories, each specified as a tuple of " + "(movable, start, stop, num)." + ), +] + +MovableListOfPoints = tuple[Movable[Any], list[Any]] + +MovableListOfPointsA = A[ + MovableListOfPoints, + Field( + description="List of tuples (device, positions). For concurrent \ + trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ + [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'. Number \ + of points for each movable must be equal." + ), +] + +MovableStartStopStep = tuple[Movable[Number], Number, Number, Number] + + +MovableStartStopStepA = A[ + MovableStartStopStep, + Field( + description="Tuple containing (movable, start, stop, step) for a scan trajectory." + ), + BeforeValidator( + trajectory_validator( + length=4, + description="(movable, start, stop, step)", + validate=validate_start_stop_step, + ) + ), +] diff --git a/src/dodal/plans/spec_path.py b/src/dodal/plans/scans/spec_path.py similarity index 87% rename from src/dodal/plans/spec_path.py rename to src/dodal/plans/scans/spec_path.py index 6f7f5a47f0a..4b605198fb9 100644 --- a/src/dodal/plans/spec_path.py +++ b/src/dodal/plans/scans/spec_path.py @@ -3,25 +3,20 @@ from typing import Annotated, Any import bluesky.plans as bp -from bluesky.protocols import Movable, Readable +from bluesky.protocols import Movable from cycler import Cycler, cycler from pydantic import Field, validate_call from scanspec.specs import Spec from dodal.common import MsgGenerator from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator +from dodal.plans.scans.annotations import DetectorsA @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) def spec_scan( - detectors: Annotated[ - set[Readable], - Field( - description="Set of readable devices, will take a reading at each point, \ - in addition to any Movables in the Spec", - ), - ], + detectors: DetectorsA, spec: Annotated[ Spec[Movable], Field(description="ScanSpec modelling the path of the scan"), diff --git a/src/dodal/plans/scans/utils.py b/src/dodal/plans/scans/utils.py new file mode 100644 index 00000000000..53cbe4bf6f7 --- /dev/null +++ b/src/dodal/plans/scans/utils.py @@ -0,0 +1,88 @@ +from collections.abc import Iterable, Sequence +from decimal import Decimal + +import numpy as np + +from dodal.plans.scans.annotations import ( + MovableListOfPoints, + MovableStartStep, + MovableStartStopStep, + Number, + T, +) + + +def flatten(items: Iterable[Iterable[T]]) -> tuple[T, ...]: + """Flatten one level of nested iterables.""" + return tuple(item for group in items for item in group) + + +def _round_list_elements( + stepped_list: list[Number], params: list[Number] +) -> list[Number]: + decimals = [Decimal(str(param)) for param in params] + exponents = [d.as_tuple().exponent for d in decimals] + decimal_places = [-exponent for exponent in exponents] # type: ignore + max_decimal_places = max(decimal_places) + return np.round(stepped_list, decimals=max_decimal_places).tolist() + + +def _make_stepped_list_step(start: float, stop: float, step: float) -> list[Number]: + if abs(step) > abs(stop - start): + step = stop - start + step = abs(step) * np.sign(stop - start) + stepped_list = np.arange(start, stop, step).tolist() + if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.05): + stepped_list.append(stepped_list[-1] + step) + rounded_stepped_list = _round_list_elements( + stepped_list=stepped_list, params=[start, stop, step] + ) + return rounded_stepped_list + + +def _make_stepped_list_num(start: float, step: float, num: int) -> list[Number]: + if num == 0 or step == 0: + raise ValueError( + f"Number of points ({num}) and number of steps ({step}) cannot be zero." + ) + stepped_list = [start + (n * step) for n in range(num)] + rounded_stepped_list = _round_list_elements( + stepped_list=stepped_list, params=[start, step] + ) + return rounded_stepped_list + + +def make_step_scan_args_and_shape( + trajectory: MovableStartStopStep, extra_trajectories: Sequence[MovableStartStep] +) -> tuple[list[MovableListOfPoints], tuple[int, ...]]: + """Convert [x, (1, 5, 1), ...] to [x, [1, 2, 3, 4, 5], ...].""" + movable, start, stop, step = trajectory + + movable_values = _make_stepped_list_step(start, stop, step) + shape = [len(movable_values)] + step_scan_args: list[MovableListOfPoints] = [(movable, movable_values)] + + for et in extra_trajectories: + movable, start, step = et + # For a non-grid scan, subsequent axes have the same number + # of points as the first axis. + movable_values = _make_stepped_list_num(start, step, shape[0]) + step_scan_args.append((movable, movable_values)) + + return step_scan_args, tuple(shape) + + +def make_step_grid_scan_args_and_shape( + params: Sequence[MovableStartStopStep], +) -> tuple[list[MovableListOfPoints], tuple[int, ...]]: + """Convert [x, (1, 5, 1), ...] to [x, [1, 2, 3, 4, 5], ...].""" + step_scan_args: list[MovableListOfPoints] = [] + shape: list[int] = [] + + for trajectory in params: + movable, start, stop, step = trajectory + movable_values = _make_stepped_list_step(start, stop, step) + shape.append(len(movable_values)) + step_scan_args.append((movable, movable_values)) + + return step_scan_args, tuple(shape) diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py new file mode 100644 index 00000000000..bc93bbd983e --- /dev/null +++ b/src/dodal/plans/scans/validators.py @@ -0,0 +1,62 @@ +from collections.abc import Callable +from typing import Any + +from bluesky.protocols import HasName, Movable + + +def trajectory_validator( + *, + length: int, + description: str, + validate: Callable[[str, tuple[Any, ...], str], None] | None = None, +) -> Callable[[Any], Any]: + def validator(value: Any) -> Any: + if not isinstance(value, tuple): + raise ValueError(f"Trajectory must be a tuple of {description}.") + + if not value: + raise ValueError(f"Trajectory must contain {description}.") + + movable = value[0] + + if not isinstance(movable, Movable): + raise ValueError( + f"The first value in a trajectory must be Movable. Got {movable!r}." + ) + movable_name = movable.name if isinstance(movable, HasName) else repr(movable) + + formatted_value = (movable_name, *value[1:]) + + if len(value) != length: + raise ValueError( + f"Trajectory for {movable_name} must contain exactly " + f"{length} values: {description}. " + f"Got {len(value)} values: {formatted_value!r}" + ) + if validate is not None: + validate(movable_name, value, description) + + return value + + return validator + + +def validate_start_stop_step( + movable_name: str, + value: tuple[Any, ...], + description: str, +) -> None: + _, start, stop, step = value + + if step == 0: + raise ValueError( + f"Step size cannot be 0. " + f"Received ({movable_name}, {start}, {stop}, {step}) for " + f"{description}." + ) + + if start == stop: + raise ValueError( + f"Start and stop values cannot be the same. " + f"Received ({movable_name}, {start}, {stop}, {step}) for {description}." + ) diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py new file mode 100644 index 00000000000..490a038d8e9 --- /dev/null +++ b/src/dodal/plans/scans/wrapped.py @@ -0,0 +1,357 @@ +from collections.abc import Iterable, Sequence +from typing import Annotated as A +from typing import Any + +import bluesky.plans as bp +from bluesky.protocols import Movable +from pydantic import Field, NonNegativeFloat, validate_call + +from dodal.common import MsgGenerator +from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator +from dodal.plans.scans.annotations import ( + DetectorsA, + MovableListOfPoints, + MovableListOfPointsA, + MovableStartStep, + MovableStartStop, + MovableStartStopA, + MovableStartStopNum, + MovableStartStopNumA, + MovableStartStopStep, + MovableStartStopStepA, +) +from dodal.plans.scans.utils import ( + flatten, + make_step_grid_scan_args_and_shape, + make_step_scan_args_and_shape, +) + +"""This module wraps plan(s) from bluesky.plans so they are compatible with blueapi. +Required decorators are installed on plan import. +https://github.com/DiamondLightSource/blueapi/issues/474 + +Non-serialisable fields are ignored when they are optional. +https://github.com/DiamondLightSource/blueapi/issues/711 + +Using *args in plans is currently not supported. +https://github.com/DiamondLightSource/blueapi/issues/1450 + +We may also need other adjustments for UI purposes, e.g. + - Forcing uniqueness or orderedness of Readables. + - Limits and metadata (e.g. units). +""" + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def count( + detectors: DetectorsA, + num: A[int, Field(description="Number of frames to collect", ge=1)] = 1, + delay: A[ + NonNegativeFloat | Sequence[NonNegativeFloat], + Field( + description="Delay between readings: if tuple, len(delay) == num - 1 and \ + the delays are between each point, if value or None is the delay for every \ + gap", + json_schema_extra={"units": "s"}, + ), + ] = 0.0, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Reads from a number of devices. + + Wraps bluesky.plans.count(det, num, delay, md=metadata) exposing only serializable + parameters and metadata. + """ + if isinstance(delay, Sequence): + assert len(delay) == num - 1, ( + f"Number of delays given must be {num - 1}: was given {len(delay)}" + ) + metadata = metadata or {} + metadata["shape"] = (num,) + yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def num_scan( + detectors: DetectorsA, + trajectory: MovableStartStop, + *extra_axes: MovableStartStopA, + num: int, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent single or multi-motor trajector(y/ies). + + The scan is defined by number of points along scan trajector(y/ies). Wraps + bluesky.plans.scan(det, *args, num, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = (num,) + + yield from bp.scan( + detectors, *trajectory, *flatten(extra_axes), num=num, md=metadata + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def num_grid_scan( + detectors: DetectorsA, + trajectory: MovableStartStopNum, + *extra_trajectories: MovableStartStopNumA, + snake_axes: Iterable[Movable] | bool = False, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent multi-motor trajectories. + + The scan is defined by number of points along scan trajectories. Snakes all fast + axes by default (all axes but the first axis provided). Wraps + bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata). + """ + yield from bp.grid_scan( + detectors, + *trajectory, + *flatten(extra_trajectories), + snake_axes=snake_axes, + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def num_rscan( + detectors: DetectorsA, + trajectory: MovableStartStop, + *extra_trajectories: MovableStartStopA, + num: int, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent trajector(y/ies), relative to current position(s). + + The scan is defined by number of points along scan trajector(y/ies). Wraps + bluesky.plans.rel_scan(det, *args, num, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = (num,) + + yield from bp.rel_scan( + detectors, *trajectory, *flatten(extra_trajectories), num=num, md=metadata + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def num_grid_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopNum, + *extra_trajectories: MovableStartStopNumA, + snake_axes: list | bool = True, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent trajectories, relative to current positions. + + The scan is defined by number of points along scan trajectories. Snakes all fast + axes by default (all axes but the first axis provided). Wraps + bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata). + """ + yield from bp.rel_grid_scan( + detectors, + *trajectory, + *flatten(extra_trajectories), + snake_axes=snake_axes, + md=metadata, + ) + + +def _make_list_scan_shape( + params: Sequence[MovableListOfPoints], grid: bool +) -> tuple[int, ...]: + shape = [] + for param in params: + points = param[1] + # List arg must all be same size. If list missing or not same size, this will + # be validated by bp.list_scan. + dim = len(points) + shape.append(dim) + if not grid: + break + + return tuple(shape) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def list_scan( + detectors: DetectorsA, + trajectory: MovableListOfPoints, + *extra_trajectories: MovableListOfPointsA, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent single or multi-motor trajector(y/ies). + + The scan is defined by providing a list of points for each scan trajectory. + Wraps bluesky.plans.list_scan(det, *args, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = _make_list_scan_shape( + [trajectory, *extra_trajectories], grid=False + ) + # typing is wrong for list scan. + yield from bp.list_scan( + detectors, + *flatten([trajectory, *extra_trajectories]), # type: ignore + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def list_grid_scan( + detectors: DetectorsA, + trajectory: MovableListOfPoints, + *extra_trajectories: MovableListOfPointsA, + snake_axes: bool = False, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent trajectories. + + The scan is defined by providing a list of points for each scan trajectory. Snakes + all fast axes by default (all axes but the first axis provided). Wraps + bluesky.plans.list_grid_scan(det, *args, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = _make_list_scan_shape( + [trajectory, *extra_trajectories], grid=True + ) + yield from bp.list_grid_scan( + detectors, + *flatten([trajectory, *extra_trajectories]), + snake_axes=snake_axes, + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def list_rscan( + detectors: DetectorsA, + trajectory: MovableListOfPoints, + *extra_trajectories: MovableListOfPointsA, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent trajector(y/ies), relative to current position. + + The scan is defined by providing a list of points for each scan trajectory. + Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = _make_list_scan_shape( + [trajectory, *extra_trajectories], grid=False + ) + yield from bp.rel_list_scan( + detectors, *flatten([trajectory, *extra_trajectories]), md=metadata + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def list_grid_rscan( + detectors: DetectorsA, + trajectory: MovableListOfPoints, + *extra_trajectories: MovableListOfPointsA, + snake_axes: bool = True, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent trajectories, relative to current positions. + + The scan is defined by providing a list of points for each scan trajectory. Snakes + all fast axes by default (all axes but the first axis provided). Wraps + bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). + """ + metadata = metadata or {} + metadata["shape"] = _make_list_scan_shape( + [trajectory, *extra_trajectories], grid=True + ) + yield from bp.rel_list_grid_scan( + detectors, + *flatten([trajectory, *extra_trajectories]), + snake_axes=snake_axes, + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def step_scan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStep, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent trajectories with specified step size. + + Generates list(s) of points for each trajectory, used with + bluesky.plans.list_scan(det, *args, md=metadata). + """ + # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 + args, shape = make_step_scan_args_and_shape(trajectory, extra_trajectories) + metadata = metadata or {} + metadata["shape"] = shape + yield from bp.list_scan(detectors, *flatten(args), md=metadata) # type: ignore + + +@validate_call(config={"arbitrary_types_allowed": True}) +def step_grid_scan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStopStepA, + snake_axes: bool = True, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent trajectories with specified step size. + + Generates list(s) of points for each trajectory, used with + bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by + default (all axes but the first axis provided). + """ + # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 + args, shape = make_step_grid_scan_args_and_shape([trajectory, *extra_trajectories]) + metadata = metadata or {} + metadata["shape"] = shape + yield from bp.list_grid_scan( + detectors, *flatten(args), snake_axes=snake_axes, md=metadata + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def step_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopStep, + *extra_trajectories: MovableStartStep, + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan concurrent trajectories with specified step size, relative to position. + + Generates list(s) of points for each trajectory, used with + bluesky.plans.rel_list_scan(det, *args, md=metadata). + """ + # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 + args, shape = make_step_scan_args_and_shape(trajectory, extra_trajectories) + metadata = metadata or {} + metadata["shape"] = shape + yield from bp.rel_list_scan(detectors, *flatten(args), md=metadata) + + +@validate_call(config={"arbitrary_types_allowed": True}) +def step_grid_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopStep, + *extra_trajectories: MovableStartStopStepA, + snake_axes: bool = True, # Currently specifying axes to snake is not supported + metadata: dict[str, Any] | None = None, +) -> MsgGenerator: + """Scan independent trajectories with specified step size, relative to position. + + Generates list(s) of points for each trajectory, used with + bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by + default (all axes but the first axis provided). + """ + # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 + args, shape = make_step_grid_scan_args_and_shape([trajectory, *extra_trajectories]) + metadata = metadata or {} + metadata["shape"] = shape + yield from bp.rel_list_grid_scan( + detectors, *flatten(args), snake_axes=snake_axes, md=metadata + ) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py deleted file mode 100644 index ef17aaad961..00000000000 --- a/src/dodal/plans/wrapped.py +++ /dev/null @@ -1,578 +0,0 @@ -from collections.abc import Sequence -from decimal import Decimal -from typing import Annotated, Any, TypeVar - -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 - -from dodal.common import MsgGenerator -from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator - -"""This module wraps plan(s) from bluesky.plans so they are compatible with blueapi. -Required decorators are installed on plan import. -https://github.com/DiamondLightSource/blueapi/issues/474 - -Non-serialisable fields are ignored when they are optional. -https://github.com/DiamondLightSource/blueapi/issues/711 - -Using *args in plans is currently not supported. -https://github.com/DiamondLightSource/blueapi/issues/1450 - -We may also need other adjustments for UI purposes, e.g. - - Forcing uniqueness or orderedness of Readables. - - Limits and metadata (e.g. units). -""" - -T = TypeVar("T") - - -@attach_data_session_metadata_decorator() -@validate_call(config={"arbitrary_types_allowed": True}) -def count( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - min_length=1, - ), - ], - num: Annotated[int, Field(description="Number of frames to collect", ge=1)] = 1, - delay: Annotated[ - NonNegativeFloat | Sequence[NonNegativeFloat], - Field( - description="Delay between readings: if tuple, len(delay) == num - 1 and \ - the delays are between each point, if value or None is the delay for every \ - gap", - json_schema_extra={"units": "s"}, - ), - ] = 0.0, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Reads from a number of devices. - - Wraps bluesky.plans.count(det, num, delay, md=metadata) exposing only serializable - parameters and metadata. - """ - if isinstance(delay, Sequence): - assert len(delay) == num - 1, ( - f"Number of delays given must be {num - 1}: was given {len(delay)}" - ) - metadata = metadata or {} - metadata["shape"] = (num,) - yield from bp.count(tuple(detectors), num, delay=delay, md=metadata) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def num_scan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="List of tuples (device, parameter). For concurrent " - "trajectories, provide '[movable1, start1, stop1, movable2, start2, stop2, " - "... , movableN, startN, stopN]'." - ), - ], - num: int, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent single or multi-motor trajector(y/ies). - - The scan is defined by number of points along scan trajector(y/ies). Wraps - bluesky.plans.scan(det, *args, num, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = (num,) - - yield from bp.scan(tuple(detectors), *params, num=num, md=metadata) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def num_grid_scan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="List of tuples (device, parameter). For independent \ - trajectories, provide '[(movable1, [start1, stop1, num1]), (movable2, \ - [start2, stop2, num2]), ... , (movableN, [startN, stopN, numN])]'." - ), - ], - snake_axes: list | bool = True, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent multi-motor trajectories. - - The scan is defined by number of points along scan trajectories. Snakes all fast - axes by default (all axes but the first axis provided). Wraps - bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata). - """ - yield from bp.grid_scan( - tuple(detectors), *params, snake_axes=snake_axes, md=metadata - ) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def num_rscan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="List of tuples (device, parameter). For concurrent \ - trajectories, provide '[movable1, start1, stop1, movable2, start2, stop2, \ - ... , movableN, startN, stopN]'." - ), - ], - num: int | None = None, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent trajector(y/ies), relative to current position(s). - - The scan is defined by number of points along scan trajector(y/ies). Wraps - bluesky.plans.rel_scan(det, *args, num, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = (num,) - - yield from bp.rel_scan(tuple(detectors), *params, num=num, md=metadata) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def num_grid_rscan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="List of tuples (device, parameter). For independent \ - trajectories, provide '[(movable1, [start1, stop1, num1]), (movable2, \ - [start2, stop2, num2]), ... , (movableN, [startN, stopN, numN])]'." - ), - ], - snake_axes: list | bool = True, - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent trajectories, relative to current positions. - - The scan is defined by number of points along scan trajectories. Snakes all fast - axes by default (all axes but the first axis provided). Wraps - bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata). - """ - yield from bp.rel_grid_scan( - tuple(detectors), *params, snake_axes=snake_axes, md=metadata - ) - - -def _make_list_scan_shape( - params: Sequence[Movable | list[float | int]], grid: bool -) -> tuple[int, ...]: - shape = [] - for param in params: - # List arg must all be same size. If list missing or not same size, this will - # be validated by bp.list_scan. - if isinstance(param, list): - dim = len(param) - shape.append(dim) - if not grid: - break - - return tuple(shape) - - -@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", - ), - ], - params: Annotated[ - list[Movable | list[float | int]], - Field( - description="List of tuples (device, positions). For concurrent \ - trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ - [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'. Number \ - of points for each movable must be equal." - ), - ], - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent single or multi-motor trajector(y/ies). - - The scan is defined by providing a list of points for each scan trajectory. - Wraps bluesky.plans.list_scan(det, *args, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params, grid=False) - - yield from bp.list_scan(tuple(detectors), *tuple(params), md=metadata) # type: ignore - - -@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", - ), - ], - params: Annotated[ - Sequence[Movable | list[float | int]], - Field( - description="For independent trajectories, provide" - "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ..., " - "movableN, [point1, point2, ...]]'." - ), - ], - snake_axes: bool = True, # Currently specifying axes to snake is not supported - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent trajectories. - - The scan is defined by providing a list of points for each scan trajectory. Snakes - all fast axes by default (all axes but the first axis provided). Wraps - bluesky.plans.list_grid_scan(det, *args, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params, grid=True) - - yield from bp.list_grid_scan( - tuple(detectors), *params, snake_axes=snake_axes, md=metadata - ) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def list_rscan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | list[float | int]], - Field( - description="For concurrent trajectories, provide " - "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ..., " - "movableN, [point1, point2, ...]]'. Number \ - of points for each movable must be equal." - ), - ], - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent trajector(y/ies), relative to current position. - - The scan is defined by providing a list of points for each scan trajectory. - Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params, grid=False) - yield from bp.rel_list_scan(tuple(detectors), *params, md=metadata) - - -@validate_call(config={"arbitrary_types_allowed": True}) -def list_grid_rscan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - Sequence[Movable | list[float | int]], - Field( - description="For independent trajectories, provide " - "'[movable1, [point1, point2, ...], movable2, [point1, point2, ...], ... , " - "movableN, [point1, point2, ...]]'." - ), - ], - snake_axes: bool = True, # Currently specifying axes to snake is not supported - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent trajectories, relative to current positions. - - The scan is defined by providing a list of points for each scan trajectory. Snakes - all fast axes by default (all axes but the first axis provided). Wraps - bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). - """ - metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape(params, grid=True) - yield from bp.rel_list_grid_scan( - tuple(detectors), *params, snake_axes=snake_axes, md=metadata - ) - - -def _round_list_elements( - stepped_list: list[float | int], params: list[float | int] -) -> list[float | int]: - decimals = [Decimal(str(param)) for param in params] - exponents = [d.as_tuple().exponent for d in decimals] - decimal_places = [-exponent for exponent in exponents] # type: ignore - max_decimal_places = max(decimal_places) - return np.round(stepped_list, decimals=max_decimal_places).tolist() - - -def _make_stepped_list_step( - start: float, stop: float, step: float -) -> list[float | int]: - if start == stop: - raise ValueError( - f"Start ({start}) and stop ({stop}) values cannot be the same." - ) - if step == 0: - raise ValueError(f"Step size {step} cannot be zero.") - if abs(step) > abs(stop - start): - step = stop - start - step = abs(step) * np.sign(stop - start) - stepped_list = np.arange(start, stop, step).tolist() - if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.05): - stepped_list.append(stepped_list[-1] + step) - rounded_stepped_list = _round_list_elements( - stepped_list=stepped_list, params=[start, stop, step] - ) - return rounded_stepped_list - - -def _make_stepped_list_num(start: float, step: float, num: int) -> list[float | int]: - if num == 0 or step == 0: - raise ValueError( - f"Number of points ({num}) and number of steps ({step}) cannot be zero." - ) - stepped_list = [start + (n * step) for n in range(num)] - rounded_stepped_list = _round_list_elements( - stepped_list=stepped_list, params=[start, step] - ) - return rounded_stepped_list - - -def require( - value: object, - expected: type[T] | tuple[type[T], ...], - name: str, -) -> T: - expected_tuple = expected if isinstance(expected, tuple) else (expected,) - if not isinstance(value, expected_tuple): - allowed = ", ".join(t.__name__ for t in expected_tuple) - raise ValueError( - f"Parameter {name} must be one of type {allowed}, got {type(value).__name__}." - ) - return value # type: ignore[return-value] - - -def parse_full_axis( - values: Sequence[Movable | float | int], -) -> tuple[Movable, float, float, float]: - if len(values) != 4: - raise ValueError( - f"The axis must be movable, start, stop, step. You provided {values}" - ) - movable = require(values[0], Movable, "movable") - start = require(values[1], (int, float), "start") - stop = require(values[2], (int, float), "stop") - step = require(values[3], (int, float), "step") - return movable, start, stop, step - - -def parse_relative_axis( - values: Sequence[Movable | float | int], -) -> tuple[Movable, float, float]: - if len(values) != 3: - raise ValueError( - f"The axis must be movable, start, stop. You provided {', '.join(map(str, values))}" - ) - movable = require(values[0], Movable, "movable") - start = require(values[1], (int, float), "start") - step = require(values[2], (int, float), "step") - return movable, start, step - - -def _make_step_scan_args_and_shape( - params: Sequence[Movable | float | int], grid: bool -) -> tuple[list[Movable | list[float]], tuple[int, ...]]: - """Convert [x, 1, 4, 1, ...] to [x, [1, 2, 3, 4], ...].""" - list_of_movable_with_values: list[list[Movable | float | int]] = [] - current_list: list[Movable | float | int] = [] - for param in params: - if isinstance(param, Movable): - current_list = [param] - list_of_movable_with_values.append(current_list) - elif isinstance(param, (int, float)): - current_list.append(param) - else: - raise ValueError( - "Scan syntax only takes movables or numbers as parameters. " - f'You provided "{param}".' - ) - - step_scan_args: list[Movable | list[float]] = [] - shape = [] - first_axis = True - for movable_with_values in list_of_movable_with_values: - if first_axis or grid: - movable, start, stop, step = parse_full_axis(movable_with_values) - movable_values = _make_stepped_list_step(start, stop, step) - shape.append(len(movable_values)) - first_axis = False - else: - # If not a grid scan, expects start, stop for all other axes and use the - # first axis shape for the number of steps. - movable, start, step = parse_relative_axis(movable_with_values) - movable_values = _make_stepped_list_num(start, step, shape[0]) - step_scan_args.extend([movable, movable_values]) - - return step_scan_args, tuple(shape) - - -@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", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="For concurrent trajectories, provide " - "'[movable1, start1, stop1, step1, movable2, start2, step2, ... , " - "movableN, startN, stepN]'." - ), - ], - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent trajectories with specified step size. - - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_scan(det, *args, md=metadata). - """ - # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args_and_shape(params, grid=False) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.list_scan(tuple(detectors), *tuple(args), md=metadata) # type: ignore - - -@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", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="List of tuples (device, parameter). For independent \ - trajectories, provide '[movable1, start1, stop1, step1, movable2, start2, " - "stop2, step2, ... , movableN, startN, stopN, stepN]'." - ), - ], - snake_axes: bool = True, # Currently specifying axes to snake is not supported - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent trajectories with specified step size. - - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by - default (all axes but the first axis provided). - """ - # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args_and_shape(params, grid=True) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.list_grid_scan( - tuple(detectors), *args, snake_axes=snake_axes, md=metadata - ) - - -@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", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="For concurrent trajectories, provide " - "'[movable1, start1, stop1, step1, movable2, start2, step2, ... , " - "movableN, startN, stepN]'." - ), - ], - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan concurrent trajectories with specified step size, relative to position. - - Generates list(s) of points for each trajectory, used with - bluesky.plans.rel_list_scan(det, *args, md=metadata). - """ - # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args_and_shape(params, grid=False) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.rel_list_scan(tuple(detectors), *args, md=metadata) - - -@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", - ), - ], - params: Annotated[ - Sequence[Movable | float | int], - Field( - description="For independent trajectories, provide " - "'[movable1, start1, stop1, step1, movable2, start2, stop2, step2, ... , " - "movableN, startN, stopN, stepN]'." - ), - ], - snake_axes: bool = True, # Currently specifying axes to snake is not supported - metadata: dict[str, Any] | None = None, -) -> MsgGenerator: - """Scan independent trajectories with specified step size, relative to position. - - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by - default (all axes but the first axis provided). - """ - # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 - args, shape = _make_step_scan_args_and_shape(params, grid=True) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.rel_list_grid_scan( - tuple(detectors), *args, snake_axes=snake_axes, md=metadata - ) diff --git a/system_tests/test_adsim.py b/system_tests/test_adsim.py index 0e2323c99a2..003daa54516 100644 --- a/system_tests/test_adsim.py +++ b/system_tests/test_adsim.py @@ -23,7 +23,7 @@ from dodal.beamlines import adsim from dodal.devices.motors import XThetaStage -from dodal.plans import count +from dodal.plans.scans import count """System tests that can be run against the containerised IOCs from epics-containers: https://github.com/epics-containers/example-services diff --git a/tests/plans/scans/__init__.py b/tests/plans/scans/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/plans/scans/conftest.py b/tests/plans/scans/conftest.py new file mode 100644 index 00000000000..b54d93384ce --- /dev/null +++ b/tests/plans/scans/conftest.py @@ -0,0 +1,56 @@ +import pytest + +from dodal.plans.scans.annotations import ( + MovableListOfPoints, + MovableStartStep, + MovableStartStop, + MovableStartStopNum, + MovableStartStopStep, +) + + +@pytest.fixture +def trajectories_start_stop( + request: pytest.FixtureRequest, +) -> list[MovableStartStop]: + return [ + (request.getfixturevalue(axis), start, stop) + for axis, start, stop in request.param + ] + + +@pytest.fixture +def trajectories_start_stop_num( + request: pytest.FixtureRequest, +) -> list[MovableStartStopNum]: + return [ + (request.getfixturevalue(axis), start, stop, num) + for axis, start, stop, num in request.param + ] + + +@pytest.fixture +def trajectories_with_list( + request: pytest.FixtureRequest, +) -> list[MovableListOfPoints]: + return [(request.getfixturevalue(axis), points) for axis, points in request.param] + + +@pytest.fixture +def trajectories_start_step( + request: pytest.FixtureRequest, +) -> list[MovableStartStep]: + return [ + (request.getfixturevalue(axis), start, step) + for axis, start, step in request.param + ] + + +@pytest.fixture +def trajectories_start_stop_step( + request: pytest.FixtureRequest, +) -> list[MovableStartStopStep]: + return [ + (request.getfixturevalue(axis), start, stop, step) + for axis, start, stop, step in request.param + ] diff --git a/tests/plans/scans/test_utils.py b/tests/plans/scans/test_utils.py new file mode 100644 index 00000000000..6108e3eca35 --- /dev/null +++ b/tests/plans/scans/test_utils.py @@ -0,0 +1,159 @@ +# import re + +# import pytest + +# from dodal.devices.motors import Motor +# from dodal.plans.scans.annotations import Number +# from dodal.plans.scans.utils import ( +# _make_stepped_list_num, +# _make_stepped_list_step, +# _round_list_elements, +# make_step_grid_scan_args_and_shape, +# make_step_scan_args_and_shape, +# ) + + +# @pytest.mark.parametrize( +# "x_list, y_list, grid, final_shape, final_length", +# ( +# [[0, 10, 1], [0, 5], False, (11,), 4], +# [[0, 10, 1], [0, 5, 1], True, (11, 6), 4], +# ), +# ) +# def test_make_step_scan_args_and_shape( +# x_axis: Motor, +# x_list: list, +# y_axis: Motor, +# y_list: list, +# grid: bool, +# final_shape: list, +# final_length: int, +# ): +# args, shape = make_step_scan_args_and_shape( +# params=[x_axis, *x_list, y_axis, *y_list], grid=grid +# ) +# assert len(args) == final_length +# assert shape == final_shape + + +# def test_make_list_scan_args_fails_when_lists_are_different_lengths( +# x_axis: Motor, +# y_axis: Motor, +# ): +# with pytest.raises(ValueError): +# _make_step_scan_args_and_shape( +# params=[x_axis, 0, 1, 2, y_axis, 0, 1, 2, 3], grid=False +# ) + + +# @pytest.mark.parametrize( +# "stepped_list, params, rounded_element", +# ( +# [[0.1234, 1.1234, 2.1234], [0.123, 2.123, 1], 0.123], +# [[0.1234, 1.1234, 2.1234], [0.12, 2.12, 1], 0.12], +# [[0.1234, 1.1234, 2.1234], [0.1, 2.1, 1], 0.1], +# [[0.1234, 1.1234, 2.1234], [0, 2, 1], 0], +# ), +# ) +# def test_round_list_elements( +# stepped_list: list[float], params: list[float], rounded_element: float +# ): +# rounded_list = _round_list_elements(stepped_list, params) +# assert rounded_list[0] == rounded_element + + +# @pytest.mark.parametrize( +# "start, stop, step", +# ( +# [-1, 1, 0.1], +# [-2, 2, 0.2], +# [1, -1, -0.1], +# [2, -2, -0.2], +# [1, -1, 0.1], +# [2, -2, 0.2], +# ), +# ) +# def test_make_stepped_list_step(start: float, stop: float, step: float): +# stepped_list = _make_stepped_list_step(start, stop, step) +# stepped_list_length = len(stepped_list) +# assert stepped_list_length == 21 +# assert stepped_list[0] / stepped_list[-1] == -1 +# assert stepped_list[10] == 0 + + +# def test_make_stepped_list_step_with_large_step(): +# stepped_list = _make_stepped_list_step(0, 1, 5) +# stepped_list_length = len(stepped_list) +# assert stepped_list_length == 2 +# assert stepped_list[0] == 0 +# assert stepped_list[-1] == 1 + + +# @pytest.mark.parametrize("start, step", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) +# def test_make_stepped_list_num(start: float, step: float): +# stepped_list = _make_stepped_list_num(start, step, num=21) +# stepped_list_length = len(stepped_list) +# assert stepped_list_length == 21 +# assert stepped_list[0] / stepped_list[-1] == -1 +# assert stepped_list[10] == 0 + + +# def test_make_stepped_list_num_fails_when_num_is_zero(): +# start = stop = 1.1 +# with pytest.raises( +# ValueError, +# match=re.escape( +# f"Start ({start}) and stop ({stop}) values cannot be the same." +# ), +# ): +# _make_stepped_list_step(start=start, stop=stop, step=0.25) + + +# def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): +# with pytest.raises( +# ValueError, +# match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), +# ): +# _make_stepped_list_num(start=1, step=0, num=0) + + +# @pytest.mark.parametrize( +# "x_list, y_list, z_list, grid", +# ( +# [[0, 1], [0, 0.2], [0, 0.5], False], +# [[0, 1, 0.25], [0, 0.2], [0, 1, 0.2, 0.5], False], +# [[0, 1, 0.25], [0, 0.2], [0, 1, 0.5], True], +# [[0, 1, 0.25], [0, 1, 0.2], [0, 0.5], True], +# ), +# ) +# def test_make_step_scan_args_fails_when_given_incorrect_number_of_parameters( +# x_axis: Motor, +# x_list: list[Number], +# y_axis: Motor, +# y_list: list[Number], +# z_axis: Motor, +# z_list: list[Number], +# grid: bool, +# ): +# with pytest.raises(ValueError): +# make_step_scan_args_and_shape( +# params=[x_axis, *x_list, y_axis, *y_list, z_axis, *z_list], grid=grid +# ) + + +# def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( +# x_axis: Motor, +# y_axis: Motor, +# ): +# with pytest.raises( +# ValueError, +# match="Scan syntax only takes movables or numbers as parameters.", +# ): +# make_step_scan_args_and_shape( +# [x_axis, 1, "3", 1, y_axis, 1, "4", 1], # type: ignore +# grid=True, +# ) +# make_step_scan_args_and_shape( +# [x_axis, 1, "3", 1, y_axis, 1, "4"], # type: ignore +# grid=False, +# ) diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py new file mode 100644 index 00000000000..0a788669763 --- /dev/null +++ b/tests/plans/scans/test_wrapped.py @@ -0,0 +1,722 @@ +import math +import re +from collections.abc import Mapping, Sequence +from typing import cast + +import pytest +from bluesky.protocols import Readable +from bluesky.run_engine import RunEngine +from event_model.documents import ( + Event, + EventDescriptor, + RunStart, + RunStop, + StreamResource, +) +from ophyd_async.core import AsyncReadable, StandardDetector +from ophyd_async.testing import assert_emitted +from pydantic import ValidationError + +from dodal.devices.motors import Motor +from dodal.plans.scans.annotations import ( + MovableListOfPoints, + MovableStartStep, + MovableStartStop, + MovableStartStopNum, + MovableStartStopStep, + Number, +) +from dodal.plans.scans.wrapped import ( + count, + list_grid_rscan, + list_grid_scan, + list_rscan, + list_scan, + num_grid_rscan, + num_grid_scan, + num_rscan, + num_scan, + step_grid_rscan, + step_grid_scan, + step_rscan, + step_scan, +) + + +def assert_expected_shape( + run_engine_documents: Mapping[str, list[dict]], expected_shape: tuple[int, ...] +) -> None: + start = run_engine_documents["start"][0] + assert start["shape"] == expected_shape + + +def test_count_delay_validation(det: StandardDetector, run_engine: RunEngine): + args: dict[float | Sequence[float], str] = { # type: ignore + # List wrong length + (1,): "Number of delays given must be 2: was given 1", + (1, 2, 3): "Number of delays given must be 2: was given 3", + # Delay non-physical + # negative time + -1: "Input should be greater than or equal to 0", + (-1, 2): "Input should be greater than or equal to 0", + # # null time + None: "Input should be a valid number", + (None, 2): "Input should be a valid number", + # # NaN time + "foo": "Input should be a valid number", + ("foo", 2): "Input should be a valid number", + } + for delay, reason in args.items(): + with pytest.raises((ValidationError, AssertionError), match=reason): + run_engine(count([det], num=3, delay=delay)) + + +def test_count_detectors_validation(run_engine: RunEngine): + args: dict[str, Sequence[Readable | AsyncReadable]] = { + # No device to read + "1 validation error for count": set(), + # Not Readable + "Input should be an instance of Sequence": set("foo"), # type: ignore + } + for reason, dets in args.items(): + with pytest.raises(ValidationError, match=reason): + run_engine(count(dets)) + + +def test_count_num_validation(det: StandardDetector, run_engine: RunEngine): + args: dict[int, str] = { + -1: "Input should be greater than or equal to 1", + 0: "Input should be greater than or equal to 1", + "str": "Input should be a valid integer", # type: ignore + } + for num, reason in args.items(): + with pytest.raises(ValidationError, match=reason): + run_engine(count([det], num=num)) + + +@pytest.mark.parametrize("num, shape", ([1, (1,)], [3, (3,)])) +def test_count_plan_produces_expected_start_document( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, + num: int, + shape: tuple[int, ...], +): + run_engine(count([det], num=num)) + start = run_engine_documents.get("start") + assert start and len(start) == 1 + run_start = cast(RunStart, start[0]) + assert (hints := run_start.get("hints")) and ( + hints.get("dimensions") == [(("time",), "primary")] + ) + assert_expected_shape(run_engine_documents, shape) + + +@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) +def test_count_plan_produces_expected_stop_document( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, + num: int, + length: tuple[int, ...], +): + run_engine(count([det], num=num)) + stop = run_engine_documents.get("stop") + assert stop and len(stop) == 1 + run_stop = cast(RunStop, stop[0]) + assert run_stop.get("num_events") == {"primary": length} + assert run_stop.get("exit_status") == "success" + + +def test_count_plan_produces_expected_descriptor( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, +): + run_engine(count([det], num=1)) + desc = run_engine_documents.get("descriptor") + assert desc and len(desc) == 1 + event_desc = cast(EventDescriptor, desc[0]) + object_keys = event_desc.get("object_keys") + assert object_keys is not None and det.name in object_keys + assert event_desc.get("name") == "primary" + + +@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) +def test_count_plan_produces_expected_events( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, + num: int, + length: tuple[int, ...], +): + run_engine(count([det], num=num)) + event_docs = run_engine_documents.get("event") + assert event_docs and len(event_docs) == length + for i in range(len(event_docs)): + event = cast(Event, event_docs[i]) + assert not event.get("data") # empty data + assert event.get("seq_num") == i + 1 + + +@pytest.mark.parametrize("num", [1, 3]) +def test_count_plan_produces_expected_resources( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, + num: int, +): + run_engine(count([det], num=num)) + stream_resource_docs = run_engine_documents.get("stream_resource") + data_keys = [det.name, f"{det.name}-sum"] + assert stream_resource_docs and len(stream_resource_docs) == len(data_keys) + for i in range(len(stream_resource_docs)): + resource = cast(StreamResource, stream_resource_docs[i]) + assert resource.get("data_key") == data_keys[i] + + +@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) +def test_count_plan_produces_expected_datums( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + det: StandardDetector, + num: int, + length: tuple[int, ...], +): + run_engine(count([det], num=num)) + stream_datum = run_engine_documents.get("stream_datum") + data_keys = [det.name, f"{det.name}-sum"] + assert stream_datum and len(stream_datum) == len(data_keys) * length + + +def _assert_emitted( + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + num: int, + start: int = 1, + descriptor: int = 1, + stream_resource: int = 2, + stop: int = 1, +) -> None: + numbers = { + "start": start, + "descriptor": descriptor, + "event": num, + "stop": stop, + } + # If detector, add stream parts. + if len(detectors) > 0: + # Order matters + numbers = { + "start": start, + "descriptor": descriptor, + "stream_resource": stream_resource, + "stream_datum": num * stream_resource, + "event": num, + "stop": stop, + } + assert_emitted(run_engine_documents, **numbers) + + +@pytest.fixture(params=[0, 1], ids=["0 detector(s)", "1 detector(s)"]) +def detectors( + request: pytest.FixtureRequest, det: StandardDetector +) -> Sequence[StandardDetector]: + return [] if request.param == 0 else [det] + + +@pytest.mark.parametrize( + "trajectories_start_stop, num", + [ + ([("x_axis", 0.0, 2.2)], 5), + ([("x_axis", 1.1, -1.1)], 3), + ([("x_axis", -1.1, 1.1), ("y_axis", 2.2, -2.2)], 5), + ([("x_axis", 0, 1.1), ("y_axis", 2.2, 3.3)], 5), + ], + indirect=["trajectories_start_stop"], +) +def test_num_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop: list[MovableStartStop], + num: int, +): + run_engine( + num_scan( + detectors, trajectories_start_stop[0], *trajectories_start_stop[1:], num=num + ) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) + + +def test_num_scan_fails_when_given_wrong_number_of_params( + run_engine: RunEngine, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine(num_scan([], x_axis, -1, 1, (y_axis, 1, 5, 1), num=5)) # type: ignore + + +@pytest.mark.parametrize( + "trajectories_start_stop_num, snake_axes", + [ + ([("x_axis", -1.1, 1.1, 5)], True), + ([("x_axis", -1.1, 1.1, 5)], False), + ([("x_axis", 0, 1.1, 5), ("y_axis", 2.2, 3.3, 5)], True), + ([("x_axis", 0, 1.1, 5), ("y_axis", 2.2, 3.3, 5)], False), + ], + indirect=["trajectories_start_stop_num"], +) +def test_num_grid_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_num: list[MovableStartStopNum], + snake_axes: bool, +): + run_engine( + num_grid_scan( + detectors, + trajectories_start_stop_num[0], + *trajectories_start_stop_num[1:], + snake_axes=snake_axes, + ) + ) + expected_shape = tuple(num for _, _, _, num in trajectories_start_stop_num) + _assert_emitted(run_engine_documents, detectors, math.prod(expected_shape)) + assert_expected_shape(run_engine_documents, expected_shape) + + +# def test_num_grid_scan_fails_when_given_wrong_number_of_params( +# run_engine: RunEngine, +# x_axis: Motor, +# y_axis: Motor, +# ): +# with pytest.raises(ValueError): +# run_engine(num_grid_scan(detectors=[], params=[x_axis, 0, 1.1, 2, y_axis, 1.1])) + + +@pytest.mark.parametrize( + "x_start, x_stop, x_num, y_start, y_stop, y_num", + ([-1.1, 1.1, 5, 2.2, -2.2, 3], [0, 1.1, 3, 2.2, 3.3, 5]), +) +def test_num_scan_fails_when_asked_to_snake_slow_axis( + run_engine: RunEngine, + x_axis: Motor, + x_start: Number, + x_stop: Number, + x_num: int, + y_axis: Motor, + y_start: Number, + y_stop: Number, + y_num: int, +): + with pytest.raises(ValueError): + run_engine( + num_grid_scan( + [], + (x_axis, x_start, x_stop, x_num), + (y_axis, y_start, y_stop, y_num), + snake_axes=[x_axis], + ) + ) + + +@pytest.mark.parametrize( + "trajectories_start_stop, num", + [ + ([("x_axis", 0.0, 2.2)], 5), + ([("x_axis", 1.1, -1.1)], 3), + ([("x_axis", -1.1, 1.1), ("y_axis", 2.2, -2.2)], 6), + ([("x_axis", 0, 1.1), ("y_axis", 2.2, 3.3)], 5), + ], + indirect=["trajectories_start_stop"], +) +def test_num_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop: list[MovableStartStop], + num: int, +): + run_engine( + num_rscan( + detectors, trajectories_start_stop[0], *trajectories_start_stop[1:], num=num + ) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) + + +# @pytest.mark.parametrize( +# "x_list, y_list, num", ([[-1, 1], [2, 0], 0], [[-1, 1], [-1, 1], 3.5]) +# ) +# def test_num_rscan_fails_when_given_bad_info( +# run_engine: RunEngine, +# x_axis: Motor, +# x_list: list[float | int], +# y_axis: Motor, +# y_list: list[float | int], +# num: int, +# ): +# with pytest.raises(ValueError): +# run_engine( +# num_rscan( +# detectors=[], +# params=[x_axis, *x_list, y_axis, *y_list], +# num=num, +# ) +# ) + + +@pytest.mark.parametrize( + "trajectories_start_stop_num, snake_axes", + [ + ([("x_axis", -1.1, 1.1, 5)], True), + ([("x_axis", 0, 1.1, 5)], False), + ([("x_axis", -1.1, 1.1, 5), ("y_axis", 2.2, -2.2, 3)], True), + ([("x_axis", 0, 1.1, 5), ("y_axis", 2.2, 3.3, 5)], False), + ], + indirect=["trajectories_start_stop_num"], +) +def test_num_grid_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_num: list[MovableStartStopNum], + snake_axes: bool, +): + run_engine( + num_grid_rscan( + detectors, + trajectories_start_stop_num[0], + *trajectories_start_stop_num[1:], + snake_axes=snake_axes, + ) + ) + expected_shape = tuple(num for _, _, _, num in trajectories_start_stop_num) + _assert_emitted(run_engine_documents, detectors, math.prod(expected_shape)) + assert_expected_shape(run_engine_documents, expected_shape) + + +def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( + run_engine: RunEngine, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises(ValueError): + run_engine( + num_grid_rscan( + [], (x_axis, 1, 6, 10), (y_axis, -10, 0, 5), snake_axes=[x_axis] + ) + ) + + +@pytest.mark.parametrize( + "trajectories_with_list", + [ + [("x_axis", [0, 1, 2, 3])], + [("x_axis", [3, 2, 1]), ("y_axis", [1, 2, 3])], + [ + ("x_axis", [-1.1, -2.2, -3.3, -4.4, -5.5]), + ("y_axis", [1.1, 2.2, 3.3, 4.4, 5.5]), + ], + ], + indirect=True, +) +def test_list_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_with_list: list[MovableListOfPoints], +): + num = len(trajectories_with_list[0][1]) + run_engine( + list_scan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) + + +def test_list_scan_fails_with_differnt_list_lengths( + run_engine: RunEngine, x_axis: Motor, y_axis: Motor +): + with pytest.raises(ValueError): + run_engine(list_scan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) + + +@pytest.mark.parametrize( + "trajectories_with_list", + [ + [("x_axis", [0, 1, 2, 3])], + [("x_axis", [1.1, 2.2, 3.3])], + [("x_axis", [3, 2, 1]), ("y_axis", [1, 2, 3])], + [ + ("x_axis", [-1.1, -2.2, -3.3, -4.4, -5.5]), + ("y_axis", [1.1, 2.2, 3.3, 4.4, 5.5]), + ], + ], + indirect=True, +) +def test_list_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_with_list: list[MovableListOfPoints], +): + num = len(trajectories_with_list[0][1]) + run_engine( + list_rscan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, (num,)) + + +def test_list_rscan_fails_with_differnt_list_lengths( + run_engine: RunEngine, x_axis: Motor, y_axis: Motor +): + with pytest.raises(ValueError): + run_engine(list_rscan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) + + +@pytest.mark.parametrize( + "trajectories_with_list", + [ + [("x_axis", [-1.1, -2.2, -3.3, -4.4, -5.5])], + [("x_axis", [3, 2, 1]), ("y_axis", [1, 2, 3, 4])], + ], + indirect=True, +) +def test_list_grid_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_with_list: list[MovableListOfPoints], +): + shape = tuple(len(points) for _, points in trajectories_with_list) + num = math.prod(shape) + run_engine( + list_grid_scan( + detectors, trajectories_with_list[0], *trajectories_with_list[1:] + ) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, shape) + + +@pytest.mark.parametrize( + "trajectories_with_list", + [ + [("x_axis", [1.1, 2.2, 3.3, 4.4, 5.5])], + [("x_axis", [3, 2, 1]), ("y_axis", [1, 2, 3, 4])], + ], + indirect=True, +) +def test_list_grid_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_with_list: list[MovableListOfPoints], +): + shape = tuple(len(points) for _, points in trajectories_with_list) + num = math.prod(shape) + run_engine( + list_grid_rscan( + detectors, trajectories_with_list[0], *trajectories_with_list[1:] + ) + ) + _assert_emitted(run_engine_documents, detectors, num) + assert_expected_shape(run_engine_documents, shape) + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, trajectories_start_step, expected_num", + [ + ([("x_axis", 0, 1, 0.25)], [], 5), + ([("x_axis", 0, 1, 0.25)], [("y_axis", 0, 0.25)], 5), + ], + indirect=["trajectories_start_stop_step", "trajectories_start_step"], +) +def test_step_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_step: list[MovableStartStopStep], + trajectories_start_step: list[MovableStartStep], + expected_num: int, +): + run_engine( + step_scan( + detectors, + trajectories_start_stop_step[0], + *trajectories_start_step, + ) + ) + _assert_emitted(run_engine_documents, detectors, expected_num) + assert_expected_shape(run_engine_documents, (expected_num,)) + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, expected_shape, snake", + [ + ([("x_axis", 0, 1, 0.25)], (5,), True), + ([("x_axis", 0, 1, 0.25)], (5,), False), + ([("x_axis", 0, 10, 2.5), ("y_axis", 0, -10, -2.5)], (5, 5), True), + ([("x_axis", 0, 10, 2.5), ("y_axis", 0, -10, -2.5)], (5, 5), False), + ], + indirect=["trajectories_start_stop_step"], +) +def test_step_grid_scan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_step: list[MovableStartStopStep], + expected_shape: tuple[int, ...], + snake: bool, +): + run_engine( + step_grid_scan( + detectors, + trajectories_start_stop_step[0], + *trajectories_start_stop_step[1:], + snake_axes=snake, + ) + ) + _assert_emitted(run_engine_documents, detectors, math.prod(expected_shape)) + assert_expected_shape(run_engine_documents, expected_shape) + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, trajectories_start_step, expected_num", + [ + ([("x_axis", 0, 1, 0.25)], [], 5), + ([("x_axis", 0, 1, 0.25)], [("y_axis", 0, 0.25)], 5), + ], + indirect=["trajectories_start_stop_step", "trajectories_start_step"], +) +def test_step_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_step: list[MovableStartStopStep], + trajectories_start_step: list[MovableStartStep], + expected_num: int, +): + run_engine( + step_rscan( + detectors, + trajectories_start_stop_step[0], + *trajectories_start_step, + ) + ) + _assert_emitted(run_engine_documents, detectors, expected_num) + assert_expected_shape(run_engine_documents, (expected_num,)) + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, expected_shape, snake", + [ + ([("x_axis", 0, 1, 0.25)], (5,), True), + ([("x_axis", 0, 1, 0.25)], (5,), False), + ([("x_axis", 0, 10, 2.5), ("y_axis", 0, -10, -2.5)], (5, 5), True), + ([("x_axis", 0, 10, 2.5), ("y_axis", 0, -10, -2.5)], (5, 5), False), + ], + indirect=["trajectories_start_stop_step"], +) +def test_step_grid_rscan( + run_engine: RunEngine, + run_engine_documents: Mapping[str, list[dict]], + detectors: Sequence[StandardDetector], + trajectories_start_stop_step: list[MovableStartStopStep], + expected_shape: tuple[int, ...], + snake: bool, +): + run_engine( + step_grid_rscan( + detectors, + trajectories_start_stop_step[0], + *trajectories_start_stop_step[1:], + snake_axes=snake, + ) + ) + _assert_emitted(run_engine_documents, detectors, math.prod(expected_shape)) + assert_expected_shape(run_engine_documents, expected_shape) + + +def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( + run_engine: RunEngine, + x_axis: Motor, +): + with pytest.raises( + ValueError, + match=re.escape( + f"Trajectory for {x_axis.name} must contain exactly 4 values: " + "(movable, start, stop, step). Got 3 values: ('x_axis', 0, 1)" + ), + ): + run_engine(step_grid_scan([], (x_axis, 1, 5))) # type: ignore + + +def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( + run_engine: RunEngine, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises( + ValueError, + match=re.escape( + f"Trajectory for {y_axis.name} must contain exactly 4 values: " + "(movable, start, stop, step). Got 3 values: ('y_axis', 1, 2)" + ), + ): + run_engine(step_grid_scan([], (x_axis, 1, 5, 1), (y_axis, 1, 2))) # type: ignore + + +def test_step_scan_fails_with_step_size_zero( + run_engine: RunEngine, + x_axis: Motor, +): + start = 1 + stop = 5 + step = 0 + with pytest.raises( + ValueError, + match=re.escape( + f"Step size cannot be 0. " + f"Received ({x_axis.name}, {start}, {stop}, {step})" + " for (movable, start, stop, step)." + ), + ): + run_engine(step_scan([], (x_axis, start, stop, step))) + + +def test_step_scan_fails_with_start_and_stop_being_same_value( + run_engine: RunEngine, + x_axis: Motor, +): + start = stop = 0 + step = 5 + with pytest.raises( + ValueError, + match=re.escape( + f"Start and stop values cannot be the same. " + f"Received ({x_axis.name}, {start}, {stop}, {step}) " + "for (movable, start, stop, step)." + ), + ): + run_engine(step_scan([], (x_axis, start, stop, step))) + + +def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( + run_engine: RunEngine, + x_axis: Motor, + y_axis: Motor, +): + with pytest.raises( + ValueError, + match="The axis must be movable, start, stop.", + ): + run_engine(step_scan([], (x_axis, 0, 1, 0.1), (y_axis, 1, 5, 1))) # type: ignore diff --git a/tests/plans/test_compliance.py b/tests/plans/test_compliance.py index 2ab5d7ec65c..32d5285b713 100644 --- a/tests/plans/test_compliance.py +++ b/tests/plans/test_compliance.py @@ -5,8 +5,9 @@ from bluesky.utils import MsgGenerator -from dodal import plan_stubs, plans +from dodal import plan_stubs from dodal.common.types import PlanGenerator +from dodal.plans import scans """Bluesky distinguishes between `plans`: complete experimental proceedures, which open and close data collection runs, and which may be part of a larger plan that collect data @@ -63,7 +64,7 @@ def assert_metadata_requirements(plan: PlanGenerator, signature: inspect.Signatu def test_plans_comply(): - for plan in get_all_available_generators(plans): + for plan in get_all_available_generators(scans): signature = inspect.Signature.from_callable(plan) assert_hard_requirements(plan, signature) assert_metadata_requirements(plan, signature) diff --git a/tests/plans/test_scanspec.py b/tests/plans/test_scanspec.py index 55d1b5e8f28..869953eaf50 100644 --- a/tests/plans/test_scanspec.py +++ b/tests/plans/test_scanspec.py @@ -16,7 +16,7 @@ from ophyd_async.sim import SimMotor from scanspec.specs import Line -from dodal.plans import spec_scan +from dodal.plans.scans import spec_scan @pytest.fixture diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py deleted file mode 100644 index c27c208a4ba..00000000000 --- a/tests/plans/test_wrapped.py +++ /dev/null @@ -1,1033 +0,0 @@ -import re -from collections.abc import Mapping, Sequence -from typing import cast - -import pytest -from bluesky.protocols import Readable -from bluesky.run_engine import RunEngine -from event_model.documents import ( - Event, - EventDescriptor, - RunStart, - RunStop, - StreamResource, -) -from ophyd_async.core import ( - AsyncReadable, - StandardDetector, -) -from ophyd_async.testing import assert_emitted -from pydantic import ValidationError - -from dodal.devices.motors import Motor -from dodal.plans.wrapped import ( - _make_step_scan_args_and_shape, - _make_stepped_list_num, - _make_stepped_list_step, - _round_list_elements, - count, - list_grid_rscan, - list_grid_scan, - list_rscan, - list_scan, - num_grid_rscan, - num_grid_scan, - num_rscan, - num_scan, - require, - step_grid_rscan, - step_grid_scan, - step_rscan, - step_scan, -) - - -def assert_expected_shape( - run_engine_documents: Mapping[str, list[dict]], expected_shape: tuple[int, ...] -) -> None: - start = run_engine_documents["start"][0] - assert start["shape"] == expected_shape - - -def test_count_delay_validation(det: StandardDetector, run_engine: RunEngine): - args: dict[float | Sequence[float], str] = { # type: ignore - # List wrong length - (1,): "Number of delays given must be 2: was given 1", - (1, 2, 3): "Number of delays given must be 2: was given 3", - # Delay non-physical - # negative time - -1: "Input should be greater than or equal to 0", - (-1, 2): "Input should be greater than or equal to 0", - # # null time - None: "Input should be a valid number", - (None, 2): "Input should be a valid number", - # # NaN time - "foo": "Input should be a valid number", - ("foo", 2): "Input should be a valid number", - } - for delay, reason in args.items(): - with pytest.raises((ValidationError, AssertionError), match=reason): - run_engine(count([det], num=3, delay=delay)) - - -def test_count_detectors_validation(run_engine: RunEngine): - args: dict[str, Sequence[Readable | AsyncReadable]] = { - # No device to read - "1 validation error for count": set(), - # Not Readable - "Input should be an instance of Sequence": set("foo"), # type: ignore - } - for reason, dets in args.items(): - with pytest.raises(ValidationError, match=reason): - run_engine(count(dets)) - - -def test_count_num_validation(det: StandardDetector, run_engine: RunEngine): - args: dict[int, str] = { - -1: "Input should be greater than or equal to 1", - 0: "Input should be greater than or equal to 1", - "str": "Input should be a valid integer", # type: ignore - } - for num, reason in args.items(): - with pytest.raises(ValidationError, match=reason): - run_engine(count([det], num=num)) - - -@pytest.mark.parametrize("num, shape", ([1, (1,)], [3, (3,)])) -def test_count_plan_produces_expected_start_document( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, - num: int, - shape: tuple[int, ...], -): - run_engine(count([det], num=num)) - start = run_engine_documents.get("start") - assert start and len(start) == 1 - run_start = cast(RunStart, start[0]) - assert (hints := run_start.get("hints")) and ( - hints.get("dimensions") == [(("time",), "primary")] - ) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) -def test_count_plan_produces_expected_stop_document( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, - num: int, - length: tuple[int, ...], -): - run_engine(count([det], num=num)) - stop = run_engine_documents.get("stop") - assert stop and len(stop) == 1 - run_stop = cast(RunStop, stop[0]) - assert run_stop.get("num_events") == {"primary": length} - assert run_stop.get("exit_status") == "success" - - -def test_count_plan_produces_expected_descriptor( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, -): - run_engine(count([det], num=1)) - desc = run_engine_documents.get("descriptor") - assert desc and len(desc) == 1 - event_desc = cast(EventDescriptor, desc[0]) - object_keys = event_desc.get("object_keys") - assert object_keys is not None and det.name in object_keys - assert event_desc.get("name") == "primary" - - -@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) -def test_count_plan_produces_expected_events( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, - num: int, - length: tuple[int, ...], -): - run_engine(count([det], num=num)) - event_docs = run_engine_documents.get("event") - assert event_docs and len(event_docs) == length - for i in range(len(event_docs)): - event = cast(Event, event_docs[i]) - assert not event.get("data") # empty data - assert event.get("seq_num") == i + 1 - - -@pytest.mark.parametrize("num", [1, 3]) -def test_count_plan_produces_expected_resources( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, - num: int, -): - run_engine(count([det], num=num)) - stream_resource_docs = run_engine_documents.get("stream_resource") - data_keys = [det.name, f"{det.name}-sum"] - assert stream_resource_docs and len(stream_resource_docs) == len(data_keys) - for i in range(len(stream_resource_docs)): - resource = cast(StreamResource, stream_resource_docs[i]) - assert resource.get("data_key") == data_keys[i] - - -@pytest.mark.parametrize("num, length", ([1, 1], [3, 3])) -def test_count_plan_produces_expected_datums( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - det: StandardDetector, - num: int, - length: tuple[int, ...], -): - run_engine(count([det], num=num)) - stream_datum = run_engine_documents.get("stream_datum") - data_keys = [det.name, f"{det.name}-sum"] - assert stream_datum and len(stream_datum) == len(data_keys) * length - - -def test_count_with_no_detector_raise_error(run_engine: RunEngine): - with pytest.raises(ValidationError): - run_engine(count([])) - - -def _assert_emitted( - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - num: int, - start: int = 1, - descriptor: int = 1, - stream_resource: int = 2, - stop: int = 1, -) -> None: - numbers = { - "start": start, - "descriptor": descriptor, - "event": num, - "stop": stop, - } - # If detector, add stream parts. - if len(detectors) > 0: - # Order matters - numbers = { - "start": start, - "descriptor": descriptor, - "stream_resource": stream_resource, - "stream_datum": num * stream_resource, - "event": num, - "stop": stop, - } - assert_emitted(run_engine_documents, **numbers) - - -@pytest.fixture(params=[0, 1], ids=["0 detector(s)", "1 detector(s)"]) -def detectors( - request: pytest.FixtureRequest, det: StandardDetector -) -> Sequence[StandardDetector]: - return [] if request.param == 0 else [det] - - -@pytest.mark.parametrize("x_list, num", ([[0.0, 2.2], 5], [[1.1, -1.1], 3])) -def test_num_scan_with_one_axis( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - num: int, -): - run_engine(num_scan(detectors=detectors, params=[x_axis, *x_list], num=num)) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", ([[-1.1, 1.1], [2.2, -2.2], 5], [[0, 1.1], [2.2, 3.3], 5]) -) -def test_num_scan_with_two_axes( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, -): - run_engine( - num_scan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - num=num, - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -def test_num_scan_fails_when_given_wrong_number_of_params( - run_engine: RunEngine, x_axis: Motor -): - with pytest.raises(ValueError): - run_engine(num_scan(detectors=[], params=[x_axis, -1, 1, 5], num=5)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", - ([[-1, 1], [2, 0], 0], [[-1, 1], [-1, 1], 3.5], [[-1, 1], [-1, 1], -2]), -) -def test_num_scan_fails_when_given_bad_info( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, -): - with pytest.raises(ValueError): - run_engine( - num_scan( - detectors=[], - params=[x_axis, *x_list, y_axis, *y_list], - num=num, - ) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) -) -def test_num_grid_scan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: tuple[float, float, int], - y_axis: Motor, - y_list: tuple[float, float, int], -): - num = int(x_list[-1] * y_list[-1]) - run_engine( - num_grid_scan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) - - -@pytest.mark.parametrize( - "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) -) -def test_num_grid_scan_when_not_snaking( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: tuple[float, float, int], - y_axis: Motor, - y_list: tuple[float, float, int], -): - num = int(x_list[-1] * y_list[-1]) - run_engine( - num_grid_scan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=False, - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) - - -def test_num_grid_scan_fails_when_given_wrong_number_of_params( - run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, -): - with pytest.raises(ValueError): - run_engine(num_grid_scan(detectors=[], params=[x_axis, 0, 1.1, 2, y_axis, 1.1])) - - -@pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) -) -def test_num_scan_fails_when_asked_to_snake_slow_axis( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises(ValueError): - run_engine( - num_grid_scan( - detectors=[], - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=[x_axis], - ) - ) - - -@pytest.mark.parametrize("x_list, num", ([[0.0, 2.2], 5], [[1.1, -1.1], 3])) -def test_num_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - num: int, -): - run_engine(num_rscan(detectors=detectors, params=[x_axis, *x_list], num=num)) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", ([[-1.1, 1.1], [2.2, -2.2], 5], [[0, 1.1], [2.2, 3.3], 5]) -) -def test_num_rscan_with_two_axes( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, -): - run_engine( - num_rscan( - detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list], num=num - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", ([[-1, 1], [2, 0], 0], [[-1, 1], [-1, 1], 3.5]) -) -def test_num_rscan_fails_when_given_bad_info( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, -): - with pytest.raises(ValueError): - run_engine( - num_rscan( - detectors=[], - params=[x_axis, *x_list, y_axis, *y_list], - num=num, - ) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) -) -def test_num_grid_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: tuple[float, float, int], - y_axis: Motor, - y_list: tuple[float, float, int], -): - num = int(x_list[-1] * y_list[-1]) - run_engine( - num_grid_rscan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) - - -@pytest.mark.parametrize( - "x_list, y_list", ([(-1.1, 1.1, 5), (2.2, -2.2, 3)], [(0, 1.1, 5), (2.2, 3.3, 5)]) -) -def test_num_grid_rscan_when_not_snaking( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: tuple[float, float, int], - y_axis: Motor, - y_list: tuple[float, float, int], -): - num = int(x_list[-1] * y_list[-1]) - run_engine( - num_grid_rscan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=False, - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (x_list[2], y_list[2])) - - -@pytest.mark.parametrize( - "x_list, y_list", ([[-1.1, 1.1, 5], [2.2, -2.2, 3]], [[0, 1.1, 5], [2.2, 3.3, 5]]) -) -def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises(ValueError): - run_engine( - num_grid_rscan( - detectors=[], - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=[x_axis], - ) - ) - - -@pytest.mark.parametrize( - "x_list, y_list, grid, final_shape, final_length", - ( - [[0, 10, 1], [0, 5], False, (11,), 4], - [[0, 10, 1], [0, 5, 1], True, (11, 6), 4], - ), -) -def test_make_step_scan_args_and_shape( - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, - grid: bool, - final_shape: list, - final_length: int, -): - args, shape = _make_step_scan_args_and_shape( - params=[x_axis, *x_list, y_axis, *y_list], grid=grid - ) - assert len(args) == final_length - assert shape == final_shape - - -def test_make_list_scan_args_fails_when_lists_are_different_lengths( - x_axis: Motor, - y_axis: Motor, -): - with pytest.raises(ValueError): - _make_step_scan_args_and_shape( - params=[x_axis, 0, 1, 2, y_axis, 0, 1, 2, 3], grid=False - ) - - -@pytest.mark.parametrize("x_list", ([0, 1, 2, 3], [1.1, 2.2, 3.3])) -def test_list_scan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, -): - num = len(x_list) - run_engine(list_scan(detectors=detectors, params=[x_axis, x_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@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, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, -): - num = int(len(x_list)) - run_engine(list_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -def test_list_scan_fails_with_differnt_list_lengths( - run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, -): - with pytest.raises(ValueError): - run_engine( - list_scan( - detectors=[], - params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], - ) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", - ( - [[3, 2, 1], [1, 2, 3, 4]], - [[-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, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - num = int(len(x_list) * len(y_list)) - run_engine( - list_grid_scan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (len(x_list), len(y_list))) - - -@pytest.mark.parametrize("x_list", ([0, 1, 2, 3], [1.1, 2.2, 3.3])) -def test_list_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, -): - num = int(len(x_list)) - run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (len(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_rscan_with_two_axes( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, -): - num = int(len(x_list)) - - run_engine(list_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -def test_list_rscan_fails_with_differnt_list_lengths( - run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, -): - with pytest.raises(ValueError): - run_engine( - list_rscan( - detectors=[], - params=[x_axis, [1, 2, 3, 4, 5], y_axis, [1, 2, 3, 4]], - ) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", - ( - [[3, 2, 1], [1, 2, 3, 4]], - [[-1.1, -2.2, -3.3, -4.4, -5.5], [1.1, 2.2, 3.3, 4.4, 5.5]], - ), -) -def test_list_grid_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, -): - num = int(len(x_list) * len(y_list)) - - run_engine( - list_grid_rscan(detectors=detectors, params=[x_axis, x_list, y_axis, y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (len(x_list), len(y_list))) - - -@pytest.mark.parametrize( - "stepped_list, params, rounded_element", - ( - [[0.1234, 1.1234, 2.1234], [0.123, 2.123, 1], 0.123], - [[0.1234, 1.1234, 2.1234], [0.12, 2.12, 1], 0.12], - [[0.1234, 1.1234, 2.1234], [0.1, 2.1, 1], 0.1], - [[0.1234, 1.1234, 2.1234], [0, 2, 1], 0], - ), -) -def test_round_list_elements( - stepped_list: list[float], params: list[float], rounded_element: float -): - rounded_list = _round_list_elements(stepped_list, params) - assert rounded_list[0] == rounded_element - - -@pytest.mark.parametrize( - "start, stop, step", - ( - [-1, 1, 0.1], - [-2, 2, 0.2], - [1, -1, -0.1], - [2, -2, -0.2], - [1, -1, 0.1], - [2, -2, 0.2], - ), -) -def test_make_stepped_list_step(start: float, stop: float, step: float): - stepped_list = _make_stepped_list_step(start, stop, step) - stepped_list_length = len(stepped_list) - assert stepped_list_length == 21 - assert stepped_list[0] / stepped_list[-1] == -1 - assert stepped_list[10] == 0 - - -def test_make_stepped_list_step_with_large_step(): - stepped_list = _make_stepped_list_step(0, 1, 5) - stepped_list_length = len(stepped_list) - assert stepped_list_length == 2 - assert stepped_list[0] == 0 - assert stepped_list[-1] == 1 - - -@pytest.mark.parametrize("start, step", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) -def test_make_stepped_list_num(start: float, step: float): - stepped_list = _make_stepped_list_num(start, step, num=21) - stepped_list_length = len(stepped_list) - assert stepped_list_length == 21 - assert stepped_list[0] / stepped_list[-1] == -1 - assert stepped_list[10] == 0 - - -def test_make_stepped_list_num_fails_when_num_is_zero(): - start = stop = 1.1 - with pytest.raises( - ValueError, - match=re.escape( - f"Start ({start}) and stop ({stop}) values cannot be the same." - ), - ): - _make_stepped_list_step(start=start, stop=stop, step=0.25) - - -def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): - with pytest.raises( - ValueError, - match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), - ): - _make_stepped_list_num(start=1, step=0, num=0) - - -def test_require_raises_error_if_not_correct_type(): - with pytest.raises( - ValueError, match="Parameter test must be one of type str, got int." - ): - require(value=5, expected=str, name="test") - - -@pytest.mark.parametrize( - "x_list, y_list, z_list, grid", - ( - [[0, 1], [0, 0.2], [0, 0.5], False], - [[0, 1, 0.25], [0, 0.2], [0, 1, 0.2, 0.5], False], - [[0, 1, 0.25], [0, 0.2], [0, 1, 0.5], True], - [[0, 1, 0.25], [0, 1, 0.2], [0, 0.5], True], - ), -) -def test_make_step_scan_args_fails_when_given_incorrect_number_of_parameters( - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - z_axis: Motor, - z_list: list[float | int], - grid: bool, -): - with pytest.raises(ValueError): - _make_step_scan_args_and_shape( - params=[x_axis, *x_list, y_axis, *y_list, z_axis, *z_list], grid=grid - ) - - -@pytest.mark.parametrize( - "x_list, num", ([[0, 1, 0.1], 11], [[-1, 1, 0.1], 21], [[0, 10, 1], 11]) -) -def test_step_scan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - num, -): - run_engine(step_scan(detectors=detectors, params=[x_axis, *x_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", - ( - [[0, 1, 0.25], [0, 0.1], 5], - [[-1, 1, 0.25], [-1, 0.1], 9], - [[0, 10, 2.5], [0, 1], 5], - ), -) -def test_step_scan_with_multiple_axes( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num, -): - run_engine( - step_scan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, expected_num_x, y_list, expected_num_y, snake", - ( - [[0, 1, 0.25], 5, [0, 2, 0.5], 5, True], - [[0, 1, 0.25], 5, [0, 2, 0.5], 5, False], - [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, True], - [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, False], - [[0, 10, 2.5], 5, [0, -10, -2.5], 5, True], - [[0, 10, 2.5], 5, [0, -10, -2.5], 5, False], - ), -) -def test_step_grid_scan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - expected_num_x: int, - y_axis: Motor, - y_list: list[float | int], - expected_num_y: int, - snake: bool, -): - run_engine( - step_grid_scan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=snake, - ) - ) - _assert_emitted(run_engine_documents, detectors, expected_num_x * expected_num_y) - assert_expected_shape(run_engine_documents, (expected_num_x, expected_num_y)) - - -@pytest.mark.parametrize( - "x_list, y_list", ([[0, 1, 0.1], [0, 1, 0.1, 1]], [[0, 1, 0.1], [0]]) -) -def test_step_grid_scan_fails_when_given_incorrect_number_of_params( - run_engine: RunEngine, - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises(ValueError): - run_engine( - step_grid_scan( - detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list] - ) - ) - - -@pytest.mark.parametrize( - "x_list, num", - ([[0, 1, 0.1], 11], [[-1, 1, 0.1], 21], [[0, 10, 1], 11]), -) -def test_step_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - num: int, -): - run_engine(step_rscan(detectors=detectors, params=[x_axis, *x_list])) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, y_list, num", - ( - [[0, 1, 0.25], [0, 0.1], 5], - [[-1, 1, 0.25], [-1, 0.1], 9], - [[0, 10, 2.5], [0, 1], 5], - ), -) -def test_step_rscan_with_multiple_axes( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], - num: int, -): - run_engine( - step_rscan(detectors=detectors, params=[x_axis, *x_list, y_axis, *y_list]) - ) - _assert_emitted(run_engine_documents, detectors, num) - assert_expected_shape(run_engine_documents, (num,)) - - -@pytest.mark.parametrize( - "x_list, expected_num_x, y_list, expected_num_y, snake", - ( - [[0, 1, 0.25], 5, [0, 2, 0.5], 5, True], - [[0, 1, 0.25], 5, [0, 2, 0.5], 5, False], - [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, True], - [[-1, 1, 0.25], 9, [1, -1, -0.5], 5, False], - [[0, 10, 2.5], 5, [0, -10, -2.5], 5, True], - [[0, 10, 2.5], 5, [0, -10, -2.5], 5, False], - ), -) -def test_step_grid_rscan( - run_engine: RunEngine, - run_engine_documents: Mapping[str, list[dict]], - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list[float | int], - expected_num_x: int, - y_axis: Motor, - y_list: list[float | int], - expected_num_y: int, - snake: bool, -): - run_engine( - step_grid_rscan( - detectors=detectors, - params=[x_axis, *x_list, y_axis, *y_list], - snake_axes=snake, - ) - ) - _assert_emitted(run_engine_documents, detectors, expected_num_x * expected_num_y) - assert_expected_shape(run_engine_documents, (expected_num_x, expected_num_y)) - - -@pytest.mark.parametrize("x_list, y_list", ([[0, 1], [0, 1, 0.1]], [[0], [0, 1, 0.1]])) -def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axes( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises( - ValueError, - match="The axis must be movable, start, stop, step.", - ): - run_engine( - step_grid_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list]) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", ([[0, 1, 0.1], [0, 1, 0.1, 1]], [[0, 1, 0.1], [0]]) -) -def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_second_axes( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises( - ValueError, - match="The axis must be movable, start, stop, step.", - ): - run_engine( - step_grid_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list]) - ) - - -@pytest.mark.parametrize( - "x_list, y_list", ([[0, 1, 0.1], [0, 1, 0.1]], [[0, 1, 0.1], [0]]) -) -def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( - run_engine: RunEngine, - x_axis: Motor, - x_list: list[float | int], - y_axis: Motor, - y_list: list[float | int], -): - with pytest.raises( - ValueError, - match="The axis must be movable, start, stop.", - ): - run_engine(step_scan(detectors=[], params=[x_axis, *x_list, y_axis, *y_list])) - - -def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( - x_axis: Motor, - y_axis: Motor, -): - with pytest.raises( - ValueError, - match="Scan syntax only takes movables or numbers as parameters.", - ): - _make_step_scan_args_and_shape( - [x_axis, 1, "3", 1, y_axis, 1, "4", 1], # type: ignore - grid=True, - ) - _make_step_scan_args_and_shape( - [x_axis, 1, "3", 1, y_axis, 1, "4"], # type: ignore - grid=False, - ) - - -def test_step_scan_fails_with_step_size_zero( - run_engine: RunEngine, - x_axis: Motor, -): - with pytest.raises( - ValueError, - match="Step size 0 cannot be zero.", - ): - run_engine(step_scan(detectors=[], params=[x_axis, 1, 5, 0])) From f5f6b75d0088d6b3894090ca9762184c1b572209 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 10:18:14 +0000 Subject: [PATCH 13/26] Simply _round_list_elements and add doc strings to all uitls functions --- src/dodal/plans/scans/utils.py | 120 ++++++++++++++++++++++++++++----- 1 file changed, 104 insertions(+), 16 deletions(-) diff --git a/src/dodal/plans/scans/utils.py b/src/dodal/plans/scans/utils.py index 53cbe4bf6f7..63af7d5e07e 100644 --- a/src/dodal/plans/scans/utils.py +++ b/src/dodal/plans/scans/utils.py @@ -1,5 +1,6 @@ from collections.abc import Iterable, Sequence from decimal import Decimal +from typing import cast import numpy as np @@ -17,53 +18,123 @@ def flatten(items: Iterable[Iterable[T]]) -> tuple[T, ...]: return tuple(item for group in items for item in group) +def _decimal_places(value: Number) -> int: + """Return the number of decimal places represented by a numeric value. + + Uses the decimal representation of the value to avoid floating-point + representation artefacts when determining its precision. + """ + exponent = cast(int, Decimal(str(value)).as_tuple().exponent) + return max(0, -exponent) + + def _round_list_elements( - stepped_list: list[Number], params: list[Number] + values: list[Number], + params: list[Number], ) -> list[Number]: - decimals = [Decimal(str(param)) for param in params] - exponents = [d.as_tuple().exponent for d in decimals] - decimal_places = [-exponent for exponent in exponents] # type: ignore - max_decimal_places = max(decimal_places) - return np.round(stepped_list, decimals=max_decimal_places).tolist() + """Round values to the greatest decimal precision of the given parameters. + + This prevents floating-point arithmetic from producing values such as + ``0.30000000000000004`` when generating scan points. + + Args: + values: Values to round. + params: Input parameters whose decimal precision determines the + rounding precision. + """ + decimal_places = max(_decimal_places(param) for param in params) + return [round(value, decimal_places) for value in values] def _make_stepped_list_step(start: float, stop: float, step: float) -> list[Number]: + """Generate a list of points between start and stop using a step size. + + The step direction is adjusted to match the direction from start to stop. + If the step would exceed the total range, it is reduced to the range. + The stop value is included when the next step falls within 5% of the + requested step size. + + The resulting values are rounded to the greatest decimal precision of + the input parameters to avoid floating-point representation artefacts. + + Args: + start: Starting position. + stop: Final position. + step: Step size between consecutive positions. + + Returns: + A list of generated scan positions. + """ if abs(step) > abs(stop - start): step = stop - start step = abs(step) * np.sign(stop - start) stepped_list = np.arange(start, stop, step).tolist() if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.05): stepped_list.append(stepped_list[-1] + step) - rounded_stepped_list = _round_list_elements( - stepped_list=stepped_list, params=[start, stop, step] - ) + rounded_stepped_list = _round_list_elements(stepped_list, [start, stop, step]) return rounded_stepped_list def _make_stepped_list_num(start: float, step: float, num: int) -> list[Number]: + """Generate a list of points from a start position using a fixed count. + + Points are generated by repeatedly adding the step size to the starting + position. The resulting values are rounded to the greatest decimal + precision of the input parameters to avoid floating-point representation + artefacts. + + Args: + start: Starting position. + step: Step size between consecutive positions. + num: Number of points to generate. + + Returns: + A list containing ``num`` scan positions. + + Raises: + ValueError: If ``num`` or ``step`` is zero. + """ if num == 0 or step == 0: raise ValueError( f"Number of points ({num}) and number of steps ({step}) cannot be zero." ) stepped_list = [start + (n * step) for n in range(num)] - rounded_stepped_list = _round_list_elements( - stepped_list=stepped_list, params=[start, step] - ) + rounded_stepped_list = _round_list_elements(stepped_list, [start, step]) return rounded_stepped_list def make_step_scan_args_and_shape( trajectory: MovableStartStopStep, extra_trajectories: Sequence[MovableStartStep] ) -> tuple[list[MovableListOfPoints], tuple[int, ...]]: - """Convert [x, (1, 5, 1), ...] to [x, [1, 2, 3, 4, 5], ...].""" + """Generate list-scan arguments for a step scan. + + The first trajectory defines the scan range and number of points using + ``(movable, start, stop, step)``. Additional trajectories use + ``(movable, start, step)`` and are generated with the same number of + points as the first trajectory. + + This produces the list of [movable, [point1, point2, ...]] arguments required by + ``bluesky.plans.list_scan`` and the corresponding scan shape. + + Args: + trajectory: Primary trajectory defining the scan range and number of + points. + extra_trajectories: Additional trajectories. Each is generated with + the same number of points as the primary trajectory. + + Returns: + A tuple containing the generated ``(movable, points)`` arguments and + the scan shape. The shape contains a single dimension corresponding to the + number of points in the primary trajectory. + """ movable, start, stop, step = trajectory movable_values = _make_stepped_list_step(start, stop, step) shape = [len(movable_values)] step_scan_args: list[MovableListOfPoints] = [(movable, movable_values)] - for et in extra_trajectories: - movable, start, step = et + for extra_t in extra_trajectories: + movable, start, step = extra_t # For a non-grid scan, subsequent axes have the same number # of points as the first axis. movable_values = _make_stepped_list_num(start, step, shape[0]) @@ -75,7 +146,24 @@ def make_step_scan_args_and_shape( def make_step_grid_scan_args_and_shape( params: Sequence[MovableStartStopStep], ) -> tuple[list[MovableListOfPoints], tuple[int, ...]]: - """Convert [x, (1, 5, 1), ...] to [x, [1, 2, 3, 4, 5], ...].""" + """Generate list-grid-scan arguments for a stepped grid scan. + + Each trajectory defines an independent scan axis using + ``(movable, start, stop, step)``. The number of generated points for each + trajectory determines the corresponding dimension of the scan shape. + + This produces the list of [movable, [point1, point2, ...]] arguments required by + ``bluesky.plans.list_grid_scan`` and the corresponding multidimensional + scan shape. + + Args: + params: Trajectories defining the scan axes. Each trajectory consists + of a movable, start position, stop position, and step size. + + Returns: + A tuple containing a list of the generated movable with the list of points and + the scan shape, with one dimension for each trajectory. + """ step_scan_args: list[MovableListOfPoints] = [] shape: list[int] = [] From d6b910a54c5475ac41b43f72bd669c3433906293 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 12:13:55 +0000 Subject: [PATCH 14/26] Update validators to use templates --- src/dodal/plans/scans/annotations.py | 2 +- src/dodal/plans/scans/validators.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py index d2dd2b64167..d993d0bca0b 100644 --- a/src/dodal/plans/scans/annotations.py +++ b/src/dodal/plans/scans/annotations.py @@ -71,7 +71,7 @@ BeforeValidator( trajectory_validator( length=4, - description="(movable, start, stop, step)", + template="(movable, start, stop, step)", validate=validate_start_stop_step, ) ), diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py index bc93bbd983e..c6fd7094847 100644 --- a/src/dodal/plans/scans/validators.py +++ b/src/dodal/plans/scans/validators.py @@ -7,15 +7,15 @@ def trajectory_validator( *, length: int, - description: str, + template: str, validate: Callable[[str, tuple[Any, ...], str], None] | None = None, ) -> Callable[[Any], Any]: def validator(value: Any) -> Any: if not isinstance(value, tuple): - raise ValueError(f"Trajectory must be a tuple of {description}.") + raise ValueError(f"Trajectory must be a tuple of {template}.") if not value: - raise ValueError(f"Trajectory must contain {description}.") + raise ValueError(f"Trajectory must contain {template}.") movable = value[0] @@ -30,11 +30,11 @@ def validator(value: Any) -> Any: if len(value) != length: raise ValueError( f"Trajectory for {movable_name} must contain exactly " - f"{length} values: {description}. " + f"{length} values: {template}. " f"Got {len(value)} values: {formatted_value!r}" ) if validate is not None: - validate(movable_name, value, description) + validate(movable_name, value, template) return value @@ -44,7 +44,7 @@ def validator(value: Any) -> Any: def validate_start_stop_step( movable_name: str, value: tuple[Any, ...], - description: str, + template: str, ) -> None: _, start, stop, step = value @@ -52,11 +52,11 @@ def validate_start_stop_step( raise ValueError( f"Step size cannot be 0. " f"Received ({movable_name}, {start}, {stop}, {step}) for " - f"{description}." + f"{template}." ) if start == stop: raise ValueError( f"Start and stop values cannot be the same. " - f"Received ({movable_name}, {start}, {stop}, {step}) for {description}." + f"Received ({movable_name}, {start}, {stop}, {step}) for {template}." ) From 3f6b07f87a1f341c1d9350f84da1b0189d36528a Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 12:14:14 +0000 Subject: [PATCH 15/26] Add back utils tests --- tests/plans/scans/test_utils.py | 212 ++++++++++++++++-------------- tests/plans/scans/test_wrapped.py | 2 +- 2 files changed, 117 insertions(+), 97 deletions(-) diff --git a/tests/plans/scans/test_utils.py b/tests/plans/scans/test_utils.py index 6108e3eca35..fb0e54cd472 100644 --- a/tests/plans/scans/test_utils.py +++ b/tests/plans/scans/test_utils.py @@ -1,39 +1,56 @@ -# import re - -# import pytest - -# from dodal.devices.motors import Motor -# from dodal.plans.scans.annotations import Number -# from dodal.plans.scans.utils import ( -# _make_stepped_list_num, -# _make_stepped_list_step, -# _round_list_elements, -# make_step_grid_scan_args_and_shape, -# make_step_scan_args_and_shape, -# ) - - -# @pytest.mark.parametrize( -# "x_list, y_list, grid, final_shape, final_length", -# ( -# [[0, 10, 1], [0, 5], False, (11,), 4], -# [[0, 10, 1], [0, 5, 1], True, (11, 6), 4], -# ), -# ) -# def test_make_step_scan_args_and_shape( -# x_axis: Motor, -# x_list: list, -# y_axis: Motor, -# y_list: list, -# grid: bool, -# final_shape: list, -# final_length: int, -# ): -# args, shape = make_step_scan_args_and_shape( -# params=[x_axis, *x_list, y_axis, *y_list], grid=grid -# ) -# assert len(args) == final_length -# assert shape == final_shape +import re + +import pytest + +from dodal.plans.scans.annotations import MovableStartStep, MovableStartStopStep +from dodal.plans.scans.utils import ( + _make_stepped_list_num, + _make_stepped_list_step, + _round_list_elements, + make_step_grid_scan_args_and_shape, + make_step_scan_args_and_shape, +) + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, trajectories_start_stop, expected_shape, expected_length", + [ + ([("x_axis", 0, 10, 1)], [("y_axis", 0, 5)], (11,), 2), + ([("x_axis", 0, 10, 1)], [("y_axis", 0, 1)], (11,), 2), + ], + indirect=["trajectories_start_stop_step", "trajectories_start_stop"], +) +def test_make_step_scan_args_and_shape( + trajectories_start_stop_step: list[MovableStartStopStep], + trajectories_start_stop: list[MovableStartStep], + expected_shape: tuple[int, ...], + expected_length: int, +): + args, shape = make_step_scan_args_and_shape( + trajectory=trajectories_start_stop_step[0], + extra_trajectories=trajectories_start_stop, + ) + assert len(args) == expected_length + assert shape == expected_shape + + +@pytest.mark.parametrize( + "trajectories_start_stop_step, expected_shape, expected_length", + [ + ([("x_axis", 0, 10, 1), ("y_axis", 0, 5, 1)], (11, 6), 2), + ], + indirect=["trajectories_start_stop_step"], +) +def test_make_step_grid_scan_args_and_shape( + trajectories_start_stop_step: list[MovableStartStopStep], + expected_shape: tuple[int, ...], + expected_length: int, +): + args, shape = make_step_grid_scan_args_and_shape( + params=trajectories_start_stop_step + ) + assert len(args) == expected_length + assert shape == expected_shape # def test_make_list_scan_args_fails_when_lists_are_different_lengths( @@ -41,63 +58,64 @@ # y_axis: Motor, # ): # with pytest.raises(ValueError): -# _make_step_scan_args_and_shape( -# params=[x_axis, 0, 1, 2, y_axis, 0, 1, 2, 3], grid=False +# make_step_scan_args_and_shape( +# (x_axis, 0, 1, 2), (y_axis, 0, 1, 2, 3), # ) -# @pytest.mark.parametrize( -# "stepped_list, params, rounded_element", -# ( -# [[0.1234, 1.1234, 2.1234], [0.123, 2.123, 1], 0.123], -# [[0.1234, 1.1234, 2.1234], [0.12, 2.12, 1], 0.12], -# [[0.1234, 1.1234, 2.1234], [0.1, 2.1, 1], 0.1], -# [[0.1234, 1.1234, 2.1234], [0, 2, 1], 0], -# ), -# ) -# def test_round_list_elements( -# stepped_list: list[float], params: list[float], rounded_element: float -# ): -# rounded_list = _round_list_elements(stepped_list, params) -# assert rounded_list[0] == rounded_element - - -# @pytest.mark.parametrize( -# "start, stop, step", -# ( -# [-1, 1, 0.1], -# [-2, 2, 0.2], -# [1, -1, -0.1], -# [2, -2, -0.2], -# [1, -1, 0.1], -# [2, -2, 0.2], -# ), -# ) -# def test_make_stepped_list_step(start: float, stop: float, step: float): -# stepped_list = _make_stepped_list_step(start, stop, step) -# stepped_list_length = len(stepped_list) -# assert stepped_list_length == 21 -# assert stepped_list[0] / stepped_list[-1] == -1 -# assert stepped_list[10] == 0 - - -# def test_make_stepped_list_step_with_large_step(): -# stepped_list = _make_stepped_list_step(0, 1, 5) -# stepped_list_length = len(stepped_list) -# assert stepped_list_length == 2 -# assert stepped_list[0] == 0 -# assert stepped_list[-1] == 1 - - -# @pytest.mark.parametrize("start, step", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) -# def test_make_stepped_list_num(start: float, step: float): -# stepped_list = _make_stepped_list_num(start, step, num=21) -# stepped_list_length = len(stepped_list) -# assert stepped_list_length == 21 -# assert stepped_list[0] / stepped_list[-1] == -1 -# assert stepped_list[10] == 0 - - +@pytest.mark.parametrize( + "stepped_list, params, expected_rounded_element", + ( + [[0.1234, 1.1234, 2.1234], [0.123, 2.123, 1], 0.123], + [[0.1234, 1.1234, 2.1234], [0.12, 2.12, 1], 0.12], + [[0.1234, 1.1234, 2.1234], [0.1, 2.1, 1], 0.1], + [[0.1234, 1.1234, 2.1234], [0, 2, 1], 0], + ), +) +def test_round_list_elements( + stepped_list: list[float], params: list[float], expected_rounded_element: float +): + rounded_list = _round_list_elements(stepped_list, params) + assert rounded_list[0] == expected_rounded_element + + +@pytest.mark.parametrize( + "start, stop, step", + ( + [-1, 1, 0.1], + [-2, 2, 0.2], + [1, -1, -0.1], + [2, -2, -0.2], + [1, -1, 0.1], + [2, -2, 0.2], + ), +) +def test_make_stepped_list_step(start: float, stop: float, step: float): + stepped_list = _make_stepped_list_step(start, stop, step) + stepped_list_length = len(stepped_list) + assert stepped_list_length == 21 + assert stepped_list[0] / stepped_list[-1] == -1 + assert stepped_list[10] == 0 + + +def test_make_stepped_list_step_with_large_step(): + stepped_list = _make_stepped_list_step(0, 1, 5) + stepped_list_length = len(stepped_list) + assert stepped_list_length == 2 + assert stepped_list[0] == 0 + assert stepped_list[-1] == 1 + + +@pytest.mark.parametrize("start, step", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) +def test_make_stepped_list_num(start: float, step: float): + stepped_list = _make_stepped_list_num(start, step, num=21) + stepped_list_length = len(stepped_list) + assert stepped_list_length == 21 + assert stepped_list[0] / stepped_list[-1] == -1 + assert stepped_list[10] == 0 + + +# Is this needed? # def test_make_stepped_list_num_fails_when_num_is_zero(): # start = stop = 1.1 # with pytest.raises( @@ -109,14 +127,15 @@ # _make_stepped_list_step(start=start, stop=stop, step=0.25) -# def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): -# with pytest.raises( -# ValueError, -# match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), -# ): -# _make_stepped_list_num(start=1, step=0, num=0) +def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): + with pytest.raises( + ValueError, + match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), + ): + _make_stepped_list_num(start=1, step=0, num=0) +# Not needed, move to wrap level. # @pytest.mark.parametrize( # "x_list, y_list, z_list, grid", # ( @@ -141,6 +160,7 @@ # ) +# This needs to be moved to wrapped / validators # def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( # x_axis: Motor, # y_axis: Motor, diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index 0a788669763..bf12c0495db 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -654,7 +654,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( ValueError, match=re.escape( f"Trajectory for {x_axis.name} must contain exactly 4 values: " - "(movable, start, stop, step). Got 3 values: ('x_axis', 0, 1)" + "(movable, start, stop, step). Got 3 values: ('x_axis', 1, 5)" ), ): run_engine(step_grid_scan([], (x_axis, 1, 5))) # type: ignore From 8e5c0102e788767e7d8f5e73e676c4a36661f82f Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 13:29:33 +0000 Subject: [PATCH 16/26] Separate out types and annotations --- src/dodal/plans/scans/annotations.py | 56 +++++++++++++++------- src/dodal/plans/scans/types.py | 19 ++++++++ src/dodal/plans/scans/utils.py | 71 +++++++++++++++++++--------- src/dodal/plans/scans/validators.py | 51 +++++++------------- tests/plans/scans/test_utils.py | 26 ++++++---- tests/plans/scans/test_wrapped.py | 44 +++++++++-------- 6 files changed, 163 insertions(+), 104 deletions(-) create mode 100644 src/dodal/plans/scans/types.py diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py index d993d0bca0b..4097925144e 100644 --- a/src/dodal/plans/scans/annotations.py +++ b/src/dodal/plans/scans/annotations.py @@ -1,15 +1,18 @@ from collections.abc import Sequence from typing import Annotated as A -from typing import Any, TypeVar -from bluesky.protocols import Movable, Readable +from bluesky.protocols import Readable from ophyd_async.core import AsyncReadable from pydantic import BeforeValidator, Field -from dodal.plans.scans.validators import trajectory_validator, validate_start_stop_step - -Number = float | int -T = TypeVar("T") +from dodal.plans.scans.types import ( + MovableListOfPoints, + MovableStartStep, + MovableStartStop, + MovableStartStopNum, + MovableStartStopStep, +) +from dodal.plans.scans.validators import trajectory_validator DetectorsA = A[ Sequence[Readable | AsyncReadable], @@ -18,38 +21,51 @@ ), ] -MovableStartStep = tuple[Movable[Number], Number, Number] - MovableStartStepA = A[ MovableStartStep, Field( description="Additional trajectories, each specified as a tuple of " "(movable, start, step)." ), + BeforeValidator( + trajectory_validator( + length=3, + template="(movable, start, step)", + expected_type=MovableStartStop, + ) + ), ] -MovableStartStop = tuple[Movable[Number], Number, Number] - MovableStartStopA = A[ MovableStartStop, Field( description="Additional trajectories, each specified as a tuple of " "(movable, start, stop)." ), + BeforeValidator( + trajectory_validator( + length=3, + template="(movable, start, stop)", + expected_type=MovableStartStop, + ) + ), ] -MovableStartStopNum = tuple[Movable[Number], Number, Number, int] - MovableStartStopNumA = A[ MovableStartStopNum, Field( description="Additional trajectories, each specified as a tuple of " "(movable, start, stop, num)." ), + BeforeValidator( + trajectory_validator( + length=4, + template="(movable, start, stop, num)", + expected_type=MovableStartStopNum, + ) + ), ] -MovableListOfPoints = tuple[Movable[Any], list[Any]] - MovableListOfPointsA = A[ MovableListOfPoints, Field( @@ -58,11 +74,15 @@ [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'. Number \ of points for each movable must be equal." ), + BeforeValidator( + trajectory_validator( + length=2, + template="(movable, [point1, point2, ...])", + expected_type=MovableListOfPoints, + ) + ), ] -MovableStartStopStep = tuple[Movable[Number], Number, Number, Number] - - MovableStartStopStepA = A[ MovableStartStopStep, Field( @@ -72,7 +92,7 @@ trajectory_validator( length=4, template="(movable, start, stop, step)", - validate=validate_start_stop_step, + expected_type=MovableStartStopStep, ) ), ] diff --git a/src/dodal/plans/scans/types.py b/src/dodal/plans/scans/types.py new file mode 100644 index 00000000000..80d7e18ac3e --- /dev/null +++ b/src/dodal/plans/scans/types.py @@ -0,0 +1,19 @@ +from collections.abc import Sequence +from typing import Any + +from bluesky.protocols import Movable, Readable +from ophyd_async.core import AsyncReadable + +Number = float | int + +Detectors = Sequence[Readable | AsyncReadable] + +MovableStartStep = tuple[Movable[Number], Number, Number] + +MovableStartStop = tuple[Movable[Number], Number, Number] + +MovableStartStopNum = tuple[Movable[Number], Number, Number, int] + +MovableListOfPoints = tuple[Movable[Any], list[Any]] + +MovableStartStopStep = tuple[Movable[Number], Number, Number, Number] diff --git a/src/dodal/plans/scans/utils.py b/src/dodal/plans/scans/utils.py index 63af7d5e07e..be15c93262f 100644 --- a/src/dodal/plans/scans/utils.py +++ b/src/dodal/plans/scans/utils.py @@ -1,23 +1,30 @@ from collections.abc import Iterable, Sequence from decimal import Decimal -from typing import cast +from typing import TypeVar, cast import numpy as np +from bluesky.protocols import HasName -from dodal.plans.scans.annotations import ( +from dodal.plans.scans.types import ( MovableListOfPoints, MovableStartStep, + MovableStartStopNum, MovableStartStopStep, Number, - T, ) +T = TypeVar("T") + def flatten(items: Iterable[Iterable[T]]) -> tuple[T, ...]: """Flatten one level of nested iterables.""" return tuple(item for group in items for item in group) +def get_bluesky_obj_name(obj) -> str: + return obj.name if isinstance(obj, HasName) else repr(obj) + + def _decimal_places(value: Number) -> int: """Return the number of decimal places represented by a numeric value. @@ -46,7 +53,7 @@ def _round_list_elements( return [round(value, decimal_places) for value in values] -def _make_stepped_list_step(start: float, stop: float, step: float) -> list[Number]: +def _make_stepped_list_step(values: MovableStartStopStep) -> list[Number]: """Generate a list of points between start and stop using a step size. The step direction is adjusted to match the direction from start to stop. @@ -58,24 +65,40 @@ def _make_stepped_list_step(start: float, stop: float, step: float) -> list[Numb the input parameters to avoid floating-point representation artefacts. Args: - start: Starting position. - stop: Final position. - step: Step size between consecutive positions. + values: A tuple containing the movable, start position, stop position, + and step size. Returns: A list of generated scan positions. """ + movable, start, stop, step = values + + if step == 0: + raise ValueError( + f"Step size cannot be 0. " + "Expected (movable, start, stop, step). " + f"Received ({get_bluesky_obj_name(movable)}, {start}, {stop}, {step})." + ) + if start == stop: + raise ValueError( + f"Start and stop values cannot be the same. " + "Expected (movable, start, stop, step). " + f"Received ({get_bluesky_obj_name(movable)}, {start}, {stop}, {step})." + ) if abs(step) > abs(stop - start): step = stop - start + step = abs(step) * np.sign(stop - start) stepped_list = np.arange(start, stop, step).tolist() + if abs((stepped_list[-1] + step) - stop) <= abs(step * 0.05): stepped_list.append(stepped_list[-1] + step) + rounded_stepped_list = _round_list_elements(stepped_list, [start, stop, step]) return rounded_stepped_list -def _make_stepped_list_num(start: float, step: float, num: int) -> list[Number]: +def _make_stepped_list_num(values: MovableStartStopNum) -> list[Number]: """Generate a list of points from a start position using a fixed count. Points are generated by repeatedly adding the step size to the starting @@ -84,9 +107,8 @@ def _make_stepped_list_num(start: float, step: float, num: int) -> list[Number]: artefacts. Args: - start: Starting position. - step: Step size between consecutive positions. - num: Number of points to generate. + values: A tuple containing the movable, start position, stop position, + and number of points. Returns: A list containing ``num`` scan positions. @@ -94,9 +116,12 @@ def _make_stepped_list_num(start: float, step: float, num: int) -> list[Number]: Raises: ValueError: If ``num`` or ``step`` is zero. """ + movable, start, step, num = values if num == 0 or step == 0: raise ValueError( - f"Number of points ({num}) and number of steps ({step}) cannot be zero." + "Number of points and number of steps cannot be zero. " + "Expected (movable, start, step, num). " + f"Received ({get_bluesky_obj_name(movable), start, step, num}) " ) stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements(stepped_list, [start, step]) @@ -127,9 +152,8 @@ def make_step_scan_args_and_shape( the scan shape. The shape contains a single dimension corresponding to the number of points in the primary trajectory. """ - movable, start, stop, step = trajectory - - movable_values = _make_stepped_list_step(start, stop, step) + movable, _, _, _ = trajectory + movable_values = _make_stepped_list_step(trajectory) shape = [len(movable_values)] step_scan_args: list[MovableListOfPoints] = [(movable, movable_values)] @@ -137,7 +161,7 @@ def make_step_scan_args_and_shape( movable, start, step = extra_t # For a non-grid scan, subsequent axes have the same number # of points as the first axis. - movable_values = _make_stepped_list_num(start, step, shape[0]) + movable_values = _make_stepped_list_num((movable, start, step, shape[0])) step_scan_args.append((movable, movable_values)) return step_scan_args, tuple(shape) @@ -166,11 +190,14 @@ def make_step_grid_scan_args_and_shape( """ step_scan_args: list[MovableListOfPoints] = [] shape: list[int] = [] + try: + for trajectory in params: + movable, _, _, _ = trajectory + movable_values = _make_stepped_list_step(trajectory) + shape.append(len(movable_values)) + step_scan_args.append((movable, movable_values)) - for trajectory in params: - movable, start, stop, step = trajectory - movable_values = _make_stepped_list_step(start, stop, step) - shape.append(len(movable_values)) - step_scan_args.append((movable, movable_values)) + return step_scan_args, tuple(shape) - return step_scan_args, tuple(shape) + except Exception as e: + raise ValueError("Recieved input ") from e diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py index c6fd7094847..a92b264bc43 100644 --- a/src/dodal/plans/scans/validators.py +++ b/src/dodal/plans/scans/validators.py @@ -1,14 +1,16 @@ from collections.abc import Callable from typing import Any -from bluesky.protocols import HasName, Movable +from bluesky.protocols import Movable +from pydantic import TypeAdapter, ValidationError + +from dodal.plans.scans.utils import get_bluesky_obj_name def trajectory_validator( - *, length: int, template: str, - validate: Callable[[str, tuple[Any, ...], str], None] | None = None, + expected_type: Any, ) -> Callable[[Any], Any]: def validator(value: Any) -> Any: if not isinstance(value, tuple): @@ -21,42 +23,25 @@ def validator(value: Any) -> Any: if not isinstance(movable, Movable): raise ValueError( - f"The first value in a trajectory must be Movable. Got {movable!r}." + f"The first value in a trajectory must be Movable. Got {get_bluesky_obj_name(movable)!r}." ) - movable_name = movable.name if isinstance(movable, HasName) else repr(movable) - - formatted_value = (movable_name, *value[1:]) + formatted_values = (get_bluesky_obj_name(movable), *value[1:]) if len(value) != length: raise ValueError( - f"Trajectory for {movable_name} must contain exactly " - f"{length} values: {template}. " - f"Got {len(value)} values: {formatted_value!r}" + f"Trajectory must contain exactly {length} values. " + f"Expected {template}. Got {len(value)} values: {formatted_values!r}" ) - if validate is not None: - validate(movable_name, value, template) + try: + TypeAdapter( + expected_type, config={"arbitrary_types_allowed": True} + ).validate_python(value, strict=False) + except ValidationError as exc: + raise ValueError( + f"Trajectory has invalid types. Expected {template}. " + f"Received {formatted_values!r}." + ) from exc return value return validator - - -def validate_start_stop_step( - movable_name: str, - value: tuple[Any, ...], - template: str, -) -> None: - _, start, stop, step = value - - if step == 0: - raise ValueError( - f"Step size cannot be 0. " - f"Received ({movable_name}, {start}, {stop}, {step}) for " - f"{template}." - ) - - if start == stop: - raise ValueError( - f"Start and stop values cannot be the same. " - f"Received ({movable_name}, {start}, {stop}, {step}) for {template}." - ) diff --git a/tests/plans/scans/test_utils.py b/tests/plans/scans/test_utils.py index fb0e54cd472..6bbf6e14236 100644 --- a/tests/plans/scans/test_utils.py +++ b/tests/plans/scans/test_utils.py @@ -1,8 +1,9 @@ import re import pytest +from ophyd_async.sim import SimMotor -from dodal.plans.scans.annotations import MovableStartStep, MovableStartStopStep +from dodal.plans.scans.types import MovableStartStep, MovableStartStopStep from dodal.plans.scans.utils import ( _make_stepped_list_num, _make_stepped_list_step, @@ -90,16 +91,18 @@ def test_round_list_elements( [2, -2, 0.2], ), ) -def test_make_stepped_list_step(start: float, stop: float, step: float): - stepped_list = _make_stepped_list_step(start, stop, step) +def test_make_stepped_list_step( + x_axis: SimMotor, start: float, stop: float, step: float +): + stepped_list = _make_stepped_list_step((x_axis, start, stop, step)) stepped_list_length = len(stepped_list) assert stepped_list_length == 21 assert stepped_list[0] / stepped_list[-1] == -1 assert stepped_list[10] == 0 -def test_make_stepped_list_step_with_large_step(): - stepped_list = _make_stepped_list_step(0, 1, 5) +def test_make_stepped_list_step_with_large_step(x_axis: SimMotor): + stepped_list = _make_stepped_list_step((x_axis, 0, 1, 5)) stepped_list_length = len(stepped_list) assert stepped_list_length == 2 assert stepped_list[0] == 0 @@ -107,10 +110,11 @@ def test_make_stepped_list_step_with_large_step(): @pytest.mark.parametrize("start, step", ([-1, 0.1], [-2, 0.2], [1, -0.1], [2, -0.2])) -def test_make_stepped_list_num(start: float, step: float): - stepped_list = _make_stepped_list_num(start, step, num=21) +def test_make_stepped_list_num(x_axis: SimMotor, start: float, step: float): + num = 21 + stepped_list = _make_stepped_list_num((x_axis, start, step, num)) stepped_list_length = len(stepped_list) - assert stepped_list_length == 21 + assert stepped_list_length == num assert stepped_list[0] / stepped_list[-1] == -1 assert stepped_list[10] == 0 @@ -127,12 +131,14 @@ def test_make_stepped_list_num(start: float, step: float): # _make_stepped_list_step(start=start, stop=stop, step=0.25) -def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values(): +def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values( + x_axis: SimMotor, +): with pytest.raises( ValueError, match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), ): - _make_stepped_list_num(start=1, step=0, num=0) + _make_stepped_list_num((x_axis, 1, 0, 0)) # Not needed, move to wrap level. diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index bf12c0495db..541ce07a74b 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -14,11 +14,11 @@ StreamResource, ) from ophyd_async.core import AsyncReadable, StandardDetector +from ophyd_async.sim import SimMotor from ophyd_async.testing import assert_emitted from pydantic import ValidationError -from dodal.devices.motors import Motor -from dodal.plans.scans.annotations import ( +from dodal.plans.scans.types import ( MovableListOfPoints, MovableStartStep, MovableStartStop, @@ -253,8 +253,8 @@ def test_num_scan( def test_num_scan_fails_when_given_wrong_number_of_params( run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, + x_axis: SimMotor, + y_axis: SimMotor, ): with pytest.raises(ValueError): run_engine(num_scan([], x_axis, -1, 1, (y_axis, 1, 5, 1), num=5)) # type: ignore @@ -292,8 +292,8 @@ def test_num_grid_scan( # def test_num_grid_scan_fails_when_given_wrong_number_of_params( # run_engine: RunEngine, -# x_axis: Motor, -# y_axis: Motor, +# x_axis: SimMotor, +# y_axis: SimMotor, # ): # with pytest.raises(ValueError): # run_engine(num_grid_scan(detectors=[], params=[x_axis, 0, 1.1, 2, y_axis, 1.1])) @@ -305,11 +305,11 @@ def test_num_grid_scan( ) def test_num_scan_fails_when_asked_to_snake_slow_axis( run_engine: RunEngine, - x_axis: Motor, + x_axis: SimMotor, x_start: Number, x_stop: Number, x_num: int, - y_axis: Motor, + y_axis: SimMotor, y_start: Number, y_stop: Number, y_num: int, @@ -356,9 +356,9 @@ def test_num_rscan( # ) # def test_num_rscan_fails_when_given_bad_info( # run_engine: RunEngine, -# x_axis: Motor, +# x_axis: SimMotor, # x_list: list[float | int], -# y_axis: Motor, +# y_axis: SimMotor, # y_list: list[float | int], # num: int, # ): @@ -404,8 +404,8 @@ def test_num_grid_rscan( def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, + x_axis: SimMotor, + y_axis: SimMotor, ): with pytest.raises(ValueError): run_engine( @@ -442,7 +442,7 @@ def test_list_scan( def test_list_scan_fails_with_differnt_list_lengths( - run_engine: RunEngine, x_axis: Motor, y_axis: Motor + run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor ): with pytest.raises(ValueError): run_engine(list_scan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) @@ -476,7 +476,7 @@ def test_list_rscan( def test_list_rscan_fails_with_differnt_list_lengths( - run_engine: RunEngine, x_axis: Motor, y_axis: Motor + run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor ): with pytest.raises(ValueError): run_engine(list_rscan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) @@ -648,7 +648,7 @@ def test_step_grid_rscan( def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( run_engine: RunEngine, - x_axis: Motor, + x_axis: SimMotor, ): with pytest.raises( ValueError, @@ -662,8 +662,8 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, + x_axis: SimMotor, + y_axis: SimMotor, ): with pytest.raises( ValueError, @@ -677,7 +677,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( def test_step_scan_fails_with_step_size_zero( run_engine: RunEngine, - x_axis: Motor, + x_axis: SimMotor, ): start = 1 stop = 5 @@ -692,10 +692,12 @@ def test_step_scan_fails_with_step_size_zero( ): run_engine(step_scan([], (x_axis, start, stop, step))) + run_engine(step_scan([], (x_axis, start, stop, ["1"]))) # type: ignore + def test_step_scan_fails_with_start_and_stop_being_same_value( run_engine: RunEngine, - x_axis: Motor, + x_axis: SimMotor, ): start = stop = 0 step = 5 @@ -712,8 +714,8 @@ def test_step_scan_fails_with_start_and_stop_being_same_value( def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( run_engine: RunEngine, - x_axis: Motor, - y_axis: Motor, + x_axis: SimMotor, + y_axis: SimMotor, ): with pytest.raises( ValueError, From 32aa33ce58f38cd8909f7fe5aed45e4f4ea96b46 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 13:54:08 +0000 Subject: [PATCH 17/26] Tidy up --- src/dodal/plans/scans/annotations.py | 6 +-- src/dodal/plans/scans/spec_path.py | 2 +- src/dodal/plans/scans/utils.py | 22 +++++++-- src/dodal/plans/scans/wrapped.py | 69 +++++++++++++--------------- tests/plans/scans/test_utils.py | 28 ++++++----- tests/plans/test_scanspec.py | 2 +- 6 files changed, 72 insertions(+), 57 deletions(-) diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py index 4097925144e..1a2e8f1ade3 100644 --- a/src/dodal/plans/scans/annotations.py +++ b/src/dodal/plans/scans/annotations.py @@ -1,11 +1,9 @@ -from collections.abc import Sequence from typing import Annotated as A -from bluesky.protocols import Readable -from ophyd_async.core import AsyncReadable from pydantic import BeforeValidator, Field from dodal.plans.scans.types import ( + Detectors, MovableListOfPoints, MovableStartStep, MovableStartStop, @@ -15,7 +13,7 @@ from dodal.plans.scans.validators import trajectory_validator DetectorsA = A[ - Sequence[Readable | AsyncReadable], + Detectors, Field( description="Set of readable devices, will take a reading at each point", ), diff --git a/src/dodal/plans/scans/spec_path.py b/src/dodal/plans/scans/spec_path.py index 4b605198fb9..85617f7019f 100644 --- a/src/dodal/plans/scans/spec_path.py +++ b/src/dodal/plans/scans/spec_path.py @@ -39,7 +39,7 @@ def spec_scan( **(metadata or {}), } - yield from bp.scan_nd(tuple(detectors), _as_cycler(spec), md=_md) + yield from bp.scan_nd(detectors, _as_cycler(spec), md=_md) def _as_cycler(spec: Spec[Movable]) -> Cycler: diff --git a/src/dodal/plans/scans/utils.py b/src/dodal/plans/scans/utils.py index be15c93262f..9c417dfa852 100644 --- a/src/dodal/plans/scans/utils.py +++ b/src/dodal/plans/scans/utils.py @@ -25,6 +25,22 @@ def get_bluesky_obj_name(obj) -> str: return obj.name if isinstance(obj, HasName) else repr(obj) +def make_list_scan_shape( + params: Sequence[MovableListOfPoints], grid: bool +) -> tuple[int, ...]: + shape = [] + for param in params: + points = param[1] + # List arg must all be same size. If list missing or not same size, this will + # be validated by bp.list_scan. + dim = len(points) + shape.append(dim) + if not grid: + break + + return tuple(shape) + + def _decimal_places(value: Number) -> int: """Return the number of decimal places represented by a numeric value. @@ -114,14 +130,14 @@ def _make_stepped_list_num(values: MovableStartStopNum) -> list[Number]: A list containing ``num`` scan positions. Raises: - ValueError: If ``num`` or ``step`` is zero. + ValueError: If ``step`` or ``num`` is zero. """ movable, start, step, num = values if num == 0 or step == 0: raise ValueError( - "Number of points and number of steps cannot be zero. " + "Number of steps and number of points cannot be zero. " "Expected (movable, start, step, num). " - f"Received ({get_bluesky_obj_name(movable), start, step, num}) " + f"Received ({get_bluesky_obj_name(movable)}, {start}, {step}, {num})." ) stepped_list = [start + (n * step) for n in range(num)] rounded_stepped_list = _round_list_elements(stepped_list, [start, step]) diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py index 490a038d8e9..57ff07ff9f3 100644 --- a/src/dodal/plans/scans/wrapped.py +++ b/src/dodal/plans/scans/wrapped.py @@ -4,24 +4,22 @@ import bluesky.plans as bp from bluesky.protocols import Movable +from bluesky.utils import plan from pydantic import Field, NonNegativeFloat, validate_call from dodal.common import MsgGenerator from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator from dodal.plans.scans.annotations import ( DetectorsA, - MovableListOfPoints, MovableListOfPointsA, - MovableStartStep, - MovableStartStop, + MovableStartStepA, MovableStartStopA, - MovableStartStopNum, MovableStartStopNumA, - MovableStartStopStep, MovableStartStopStepA, ) from dodal.plans.scans.utils import ( flatten, + make_list_scan_shape, make_step_grid_scan_args_and_shape, make_step_scan_args_and_shape, ) @@ -44,6 +42,7 @@ @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) +@plan def count( detectors: DetectorsA, num: A[int, Field(description="Number of frames to collect", ge=1)] = 1, @@ -73,9 +72,10 @@ def count( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def num_scan( detectors: DetectorsA, - trajectory: MovableStartStop, + trajectory: MovableStartStopA, *extra_axes: MovableStartStopA, num: int, metadata: dict[str, Any] | None = None, @@ -94,9 +94,10 @@ def num_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def num_grid_scan( detectors: DetectorsA, - trajectory: MovableStartStopNum, + trajectory: MovableStartStopNumA, *extra_trajectories: MovableStartStopNumA, snake_axes: Iterable[Movable] | bool = False, metadata: dict[str, Any] | None = None, @@ -117,9 +118,10 @@ def num_grid_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def num_rscan( detectors: DetectorsA, - trajectory: MovableStartStop, + trajectory: MovableStartStopA, *extra_trajectories: MovableStartStopA, num: int, metadata: dict[str, Any] | None = None, @@ -138,9 +140,10 @@ def num_rscan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def num_grid_rscan( detectors: DetectorsA, - trajectory: MovableStartStopNum, + trajectory: MovableStartStopNumA, *extra_trajectories: MovableStartStopNumA, snake_axes: list | bool = True, metadata: dict[str, Any] | None = None, @@ -160,26 +163,11 @@ def num_grid_rscan( ) -def _make_list_scan_shape( - params: Sequence[MovableListOfPoints], grid: bool -) -> tuple[int, ...]: - shape = [] - for param in params: - points = param[1] - # List arg must all be same size. If list missing or not same size, this will - # be validated by bp.list_scan. - dim = len(points) - shape.append(dim) - if not grid: - break - - return tuple(shape) - - @validate_call(config={"arbitrary_types_allowed": True}) +@plan def list_scan( detectors: DetectorsA, - trajectory: MovableListOfPoints, + trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: @@ -189,7 +177,7 @@ def list_scan( Wraps bluesky.plans.list_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape( + metadata["shape"] = make_list_scan_shape( [trajectory, *extra_trajectories], grid=False ) # typing is wrong for list scan. @@ -201,9 +189,10 @@ def list_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def list_grid_scan( detectors: DetectorsA, - trajectory: MovableListOfPoints, + trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, snake_axes: bool = False, metadata: dict[str, Any] | None = None, @@ -215,7 +204,7 @@ def list_grid_scan( bluesky.plans.list_grid_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape( + metadata["shape"] = make_list_scan_shape( [trajectory, *extra_trajectories], grid=True ) yield from bp.list_grid_scan( @@ -227,9 +216,10 @@ def list_grid_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def list_rscan( detectors: DetectorsA, - trajectory: MovableListOfPoints, + trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: @@ -239,7 +229,7 @@ def list_rscan( Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape( + metadata["shape"] = make_list_scan_shape( [trajectory, *extra_trajectories], grid=False ) yield from bp.rel_list_scan( @@ -248,9 +238,10 @@ def list_rscan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def list_grid_rscan( detectors: DetectorsA, - trajectory: MovableListOfPoints, + trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, snake_axes: bool = True, metadata: dict[str, Any] | None = None, @@ -262,7 +253,7 @@ def list_grid_rscan( bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). """ metadata = metadata or {} - metadata["shape"] = _make_list_scan_shape( + metadata["shape"] = make_list_scan_shape( [trajectory, *extra_trajectories], grid=True ) yield from bp.rel_list_grid_scan( @@ -274,10 +265,11 @@ def list_grid_rscan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def step_scan( detectors: DetectorsA, trajectory: MovableStartStopStepA, - *extra_trajectories: MovableStartStep, + *extra_trajectories: MovableStartStepA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: """Scan concurrent trajectories with specified step size. @@ -293,6 +285,7 @@ def step_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def step_grid_scan( detectors: DetectorsA, trajectory: MovableStartStopStepA, @@ -316,10 +309,11 @@ def step_grid_scan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def step_rscan( detectors: DetectorsA, - trajectory: MovableStartStopStep, - *extra_trajectories: MovableStartStep, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStepA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: """Scan concurrent trajectories with specified step size, relative to position. @@ -335,9 +329,10 @@ def step_rscan( @validate_call(config={"arbitrary_types_allowed": True}) +@plan def step_grid_rscan( detectors: DetectorsA, - trajectory: MovableStartStopStep, + trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStopStepA, snake_axes: bool = True, # Currently specifying axes to snake is not supported metadata: dict[str, Any] | None = None, diff --git a/tests/plans/scans/test_utils.py b/tests/plans/scans/test_utils.py index 6bbf6e14236..0d918bd41e4 100644 --- a/tests/plans/scans/test_utils.py +++ b/tests/plans/scans/test_utils.py @@ -119,16 +119,18 @@ def test_make_stepped_list_num(x_axis: SimMotor, start: float, step: float): assert stepped_list[10] == 0 -# Is this needed? -# def test_make_stepped_list_num_fails_when_num_is_zero(): -# start = stop = 1.1 -# with pytest.raises( -# ValueError, -# match=re.escape( -# f"Start ({start}) and stop ({stop}) values cannot be the same." -# ), -# ): -# _make_stepped_list_step(start=start, stop=stop, step=0.25) +def test_make_stepped_list_num_fails_when_num_is_zero(x_axis: SimMotor): + start = stop = 1.1 + step = 0.25 + with pytest.raises( + ValueError, + match=re.escape( + f"Start and stop values cannot be the same. " + "Expected (movable, start, stop, step). " + f"Received (x_axis, {start}, {stop}, {step})." + ), + ): + _make_stepped_list_step((x_axis, start, stop, step)) def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values( @@ -136,7 +138,11 @@ def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values( ): with pytest.raises( ValueError, - match=re.escape("Number of points (0) and number of steps (0) cannot be zero."), + match=re.escape( + "Number of steps and number of points cannot be zero. " + "Expected (movable, start, step, num). " + "Received (x_axis, 1, 0, 0)." + ), ): _make_stepped_list_num((x_axis, 1, 0, 0)) diff --git a/tests/plans/test_scanspec.py b/tests/plans/test_scanspec.py index 869953eaf50..8764fc11adc 100644 --- a/tests/plans/test_scanspec.py +++ b/tests/plans/test_scanspec.py @@ -37,7 +37,7 @@ def documents_from_expected_shape( docs: dict[str, list[Document]] = {} run_engine( - spec_scan({det}, spec), # type: ignore + spec_scan([det], spec), lambda name, doc: docs.setdefault(name, []).append(doc), ) return docs From c482a9d7a2c49fea041baddae158d03c7c6a63c9 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 14:53:44 +0000 Subject: [PATCH 18/26] Update to be paramertised tests for complicance test, also remove no *args requirements --- tests/plans/test_compliance.py | 40 +++++++++++++++++++++------------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/tests/plans/test_compliance.py b/tests/plans/test_compliance.py index 32d5285b713..2f47e95a80d 100644 --- a/tests/plans/test_compliance.py +++ b/tests/plans/test_compliance.py @@ -3,6 +3,7 @@ from types import ModuleType from typing import Any, get_type_hints +import pytest from bluesky.utils import MsgGenerator from dodal import plan_stubs @@ -47,8 +48,8 @@ def assert_hard_requirements(plan: PlanGenerator, signature: inspect.Signature): assert plan.__doc__ is not None, f"'{plan.__name__}' has no docstring" for parameter in signature.parameters.values(): assert ( - parameter.kind is not parameter.VAR_POSITIONAL - and parameter.kind is not parameter.VAR_KEYWORD + # parameter.kind is not parameter.VAR_POSITIONAL # BlueAPI should support *args + parameter.kind is not parameter.VAR_KEYWORD ), f"'{plan.__name__}' has variadic arguments" @@ -63,16 +64,25 @@ def assert_metadata_requirements(plan: PlanGenerator, signature: inspect.Signatu assert metadata.default is None, f"'{plan.__name__}' metadata default is mutable" -def test_plans_comply(): - for plan in get_all_available_generators(scans): - signature = inspect.Signature.from_callable(plan) - assert_hard_requirements(plan, signature) - assert_metadata_requirements(plan, signature) - - -def test_stubs_comply(): - for stub in get_all_available_generators(plan_stubs): - signature = inspect.Signature.from_callable(stub) - assert_hard_requirements(stub, signature) - if "metadata" in signature.parameters: - assert_metadata_requirements(stub, signature) +@pytest.mark.parametrize( + "plan", + get_all_available_generators(scans), + ids=lambda plan: plan.__name__, +) +def test_plan_comply(plan): + signature = inspect.Signature.from_callable(plan) + assert_hard_requirements(plan, signature) + assert_metadata_requirements(plan, signature) + + +@pytest.mark.parametrize( + "stub", + get_all_available_generators(plan_stubs), + ids=lambda stub: stub.__name__, +) +def test_stub_comply(stub): + signature = inspect.Signature.from_callable(stub) + assert_hard_requirements(stub, signature) + + if "metadata" in signature.parameters: + assert_metadata_requirements(stub, signature) From 5ed89cf707edde5fd278331dd3f6427602be133a Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 14:54:04 +0000 Subject: [PATCH 19/26] Correct remaining tests plus add doc string --- src/dodal/plans/scans/validators.py | 33 +++++++++++++++++++++++++-- tests/plans/scans/test_wrapped.py | 35 +++++++++++++++-------------- 2 files changed, 49 insertions(+), 19 deletions(-) diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py index a92b264bc43..a76af4ce148 100644 --- a/src/dodal/plans/scans/validators.py +++ b/src/dodal/plans/scans/validators.py @@ -12,7 +12,35 @@ def trajectory_validator( template: str, expected_type: Any, ) -> Callable[[Any], Any]: + """Create a validator for a scan trajectory. + + The returned validator checks that the trajectory is a tuple with the + expected number of values, that its first value is a ``Movable``, and + that the complete tuple matches the supplied Pydantic type. + + Pydantic validation is performed with arbitrary types allowed so that + ``Movable`` protocol types can be validated as part of the trajectory. + Type validation is performed after the structural checks so that malformed + trajectories produce more useful error messages. + + Args: + length: Expected number of values in the trajectory tuple. + template: Human-readable description of the expected trajectory + structure, for example ``"(movable, start, stop, step)"``. + expected_type: Pydantic-compatible type describing the expected + trajectory, used to validate the types of the tuple elements. + + Returns: + A Pydantic-compatible validator function that validates a trajectory. + + Raises: + ValueError: If the value is not a tuple, is empty, contains an invalid + movable, has the wrong number of values, or contains values of + invalid types. + """ + def validator(value: Any) -> Any: + """Validate a single trajectory value.""" if not isinstance(value, tuple): raise ValueError(f"Trajectory must be a tuple of {template}.") @@ -23,14 +51,15 @@ def validator(value: Any) -> Any: if not isinstance(movable, Movable): raise ValueError( - f"The first value in a trajectory must be Movable. Got {get_bluesky_obj_name(movable)!r}." + "The first value in a trajectory must be Movable. " + f"Received {get_bluesky_obj_name(movable)!r}." ) formatted_values = (get_bluesky_obj_name(movable), *value[1:]) if len(value) != length: raise ValueError( f"Trajectory must contain exactly {length} values. " - f"Expected {template}. Got {len(value)} values: {formatted_values!r}" + f"Expected {template}. Received {len(value)} values: {formatted_values!r}" ) try: TypeAdapter( diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index 541ce07a74b..cb82277eb4b 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -653,8 +653,9 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( with pytest.raises( ValueError, match=re.escape( - f"Trajectory for {x_axis.name} must contain exactly 4 values: " - "(movable, start, stop, step). Got 3 values: ('x_axis', 1, 5)" + "Trajectory must contain exactly 4 values. " + "Expected (movable, start, stop, step). " + "Received 3 values: ('x_axis', 1, 5)" ), ): run_engine(step_grid_scan([], (x_axis, 1, 5))) # type: ignore @@ -668,8 +669,9 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( with pytest.raises( ValueError, match=re.escape( - f"Trajectory for {y_axis.name} must contain exactly 4 values: " - "(movable, start, stop, step). Got 3 values: ('y_axis', 1, 2)" + "Trajectory must contain exactly 4 values. " + "Expected (movable, start, stop, step). " + "Received 3 values: ('y_axis', 1, 2)" ), ): run_engine(step_grid_scan([], (x_axis, 1, 5, 1), (y_axis, 1, 2))) # type: ignore @@ -679,20 +681,15 @@ def test_step_scan_fails_with_step_size_zero( run_engine: RunEngine, x_axis: SimMotor, ): - start = 1 - stop = 5 - step = 0 with pytest.raises( ValueError, match=re.escape( - f"Step size cannot be 0. " - f"Received ({x_axis.name}, {start}, {stop}, {step})" - " for (movable, start, stop, step)." + "Step size cannot be 0. " + "Expected (movable, start, stop, step). " + "Received (x_axis, 1, 5, 0)" ), ): - run_engine(step_scan([], (x_axis, start, stop, step))) - - run_engine(step_scan([], (x_axis, start, stop, ["1"]))) # type: ignore + run_engine(step_scan([], (x_axis, 1, 5, 0))) def test_step_scan_fails_with_start_and_stop_being_same_value( @@ -704,9 +701,9 @@ def test_step_scan_fails_with_start_and_stop_being_same_value( with pytest.raises( ValueError, match=re.escape( - f"Start and stop values cannot be the same. " - f"Received ({x_axis.name}, {start}, {stop}, {step}) " - "for (movable, start, stop, step)." + "Start and stop values cannot be the same. " + "Expected (movable, start, stop, step). " + f"Received ({x_axis.name}, {start}, {stop}, {step})." ), ): run_engine(step_scan([], (x_axis, start, stop, step))) @@ -719,6 +716,10 @@ def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( ): with pytest.raises( ValueError, - match="The axis must be movable, start, stop.", + match=re.escape( + "Trajectory must contain exactly 3 values. " + "Expected (movable, start, step). " + "Received 4 values: ('y_axis', 1, 5, 1)" + ), ): run_engine(step_scan([], (x_axis, 0, 1, 0.1), (y_axis, 1, 5, 1))) # type: ignore From 99728ba2a6f1bbe4d71b38570ef45e1e57a2f99c Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 15:25:42 +0000 Subject: [PATCH 20/26] Add missing code coverage --- src/dodal/plans/scans/annotations.py | 2 +- src/dodal/plans/scans/utils.py | 18 +++++++---------- src/dodal/plans/scans/validators.py | 10 ++++------ tests/plans/scans/test_wrapped.py | 29 ++++++++++++++++++++++++++++ 4 files changed, 41 insertions(+), 18 deletions(-) diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py index 1a2e8f1ade3..95f37de6eec 100644 --- a/src/dodal/plans/scans/annotations.py +++ b/src/dodal/plans/scans/annotations.py @@ -29,7 +29,7 @@ trajectory_validator( length=3, template="(movable, start, step)", - expected_type=MovableStartStop, + expected_type=MovableStartStep, ) ), ] diff --git a/src/dodal/plans/scans/utils.py b/src/dodal/plans/scans/utils.py index 9c417dfa852..2f97e12236a 100644 --- a/src/dodal/plans/scans/utils.py +++ b/src/dodal/plans/scans/utils.py @@ -22,7 +22,7 @@ def flatten(items: Iterable[Iterable[T]]) -> tuple[T, ...]: def get_bluesky_obj_name(obj) -> str: - return obj.name if isinstance(obj, HasName) else repr(obj) + return obj.name if isinstance(obj, HasName) else str(obj) def make_list_scan_shape( @@ -206,14 +206,10 @@ def make_step_grid_scan_args_and_shape( """ step_scan_args: list[MovableListOfPoints] = [] shape: list[int] = [] - try: - for trajectory in params: - movable, _, _, _ = trajectory - movable_values = _make_stepped_list_step(trajectory) - shape.append(len(movable_values)) - step_scan_args.append((movable, movable_values)) - - return step_scan_args, tuple(shape) + for trajectory in params: + movable, _, _, _ = trajectory + movable_values = _make_stepped_list_step(trajectory) + shape.append(len(movable_values)) + step_scan_args.append((movable, movable_values)) - except Exception as e: - raise ValueError("Recieved input ") from e + return step_scan_args, tuple(shape) diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py index a76af4ce148..8618b836665 100644 --- a/src/dodal/plans/scans/validators.py +++ b/src/dodal/plans/scans/validators.py @@ -44,17 +44,15 @@ def validator(value: Any) -> Any: if not isinstance(value, tuple): raise ValueError(f"Trajectory must be a tuple of {template}.") - if not value: - raise ValueError(f"Trajectory must contain {template}.") - movable = value[0] + formatted_values = (get_bluesky_obj_name(movable), *value[1:]) if not isinstance(movable, Movable): raise ValueError( - "The first value in a trajectory must be Movable. " - f"Received {get_bluesky_obj_name(movable)!r}." + "The first value in a trajectory must implement the Movable protocol. " + f"{get_bluesky_obj_name(movable)} does not implement Movable. " + f"Received {formatted_values}." ) - formatted_values = (get_bluesky_obj_name(movable), *value[1:]) if len(value) != length: raise ValueError( diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index cb82277eb4b..b8489b0c5cf 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -723,3 +723,32 @@ def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( ), ): run_engine(step_scan([], (x_axis, 0, 1, 0.1), (y_axis, 1, 5, 1))) # type: ignore + + +def test_scan_fails_when_not_using_movable( + run_engine: RunEngine, + x_axis: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "The first value in a trajectory must implement the Movable protocol. " + "y_axis does not implement Movable. " + "Received ('y_axis', 1, 5, 1)." + ), + ): + run_engine(step_scan([], (x_axis, 0, 1, 0.1), ("y_axis", 1, 5, 1))) # type: ignore + + +def test_scan_fails_when_using_invalid_structure( + run_engine: RunEngine, + x_axis: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "Trajectory has invalid types. Expected (movable, start, stop, step). " + "Received ('x_axis', 0, 1, [0.1])." + ), + ): + run_engine(step_rscan([], (x_axis, 0, 1, [0.1]))) # type: ignore From 4bd9b64d5ca77d909ace9da0813602587be72074 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 15:34:25 +0000 Subject: [PATCH 21/26] Remove commented out code --- tests/plans/scans/test_utils.py | 54 ------------------------------- tests/plans/scans/test_wrapped.py | 30 ----------------- 2 files changed, 84 deletions(-) diff --git a/tests/plans/scans/test_utils.py b/tests/plans/scans/test_utils.py index 0d918bd41e4..1ac3ffd10c6 100644 --- a/tests/plans/scans/test_utils.py +++ b/tests/plans/scans/test_utils.py @@ -54,16 +54,6 @@ def test_make_step_grid_scan_args_and_shape( assert shape == expected_shape -# def test_make_list_scan_args_fails_when_lists_are_different_lengths( -# x_axis: Motor, -# y_axis: Motor, -# ): -# with pytest.raises(ValueError): -# make_step_scan_args_and_shape( -# (x_axis, 0, 1, 2), (y_axis, 0, 1, 2, 3), -# ) - - @pytest.mark.parametrize( "stepped_list, params, expected_rounded_element", ( @@ -145,47 +135,3 @@ def test_make_stepped_list_num_fails_when_given_equal_start_and_stop_values( ), ): _make_stepped_list_num((x_axis, 1, 0, 0)) - - -# Not needed, move to wrap level. -# @pytest.mark.parametrize( -# "x_list, y_list, z_list, grid", -# ( -# [[0, 1], [0, 0.2], [0, 0.5], False], -# [[0, 1, 0.25], [0, 0.2], [0, 1, 0.2, 0.5], False], -# [[0, 1, 0.25], [0, 0.2], [0, 1, 0.5], True], -# [[0, 1, 0.25], [0, 1, 0.2], [0, 0.5], True], -# ), -# ) -# def test_make_step_scan_args_fails_when_given_incorrect_number_of_parameters( -# x_axis: Motor, -# x_list: list[Number], -# y_axis: Motor, -# y_list: list[Number], -# z_axis: Motor, -# z_list: list[Number], -# grid: bool, -# ): -# with pytest.raises(ValueError): -# make_step_scan_args_and_shape( -# params=[x_axis, *x_list, y_axis, *y_list, z_axis, *z_list], grid=grid -# ) - - -# This needs to be moved to wrapped / validators -# def test_make_step_scan_args_and_shape_fails_with_invalid_type_args( -# x_axis: Motor, -# y_axis: Motor, -# ): -# with pytest.raises( -# ValueError, -# match="Scan syntax only takes movables or numbers as parameters.", -# ): -# make_step_scan_args_and_shape( -# [x_axis, 1, "3", 1, y_axis, 1, "4", 1], # type: ignore -# grid=True, -# ) -# make_step_scan_args_and_shape( -# [x_axis, 1, "3", 1, y_axis, 1, "4"], # type: ignore -# grid=False, -# ) diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index b8489b0c5cf..0666c2caf99 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -290,15 +290,6 @@ def test_num_grid_scan( assert_expected_shape(run_engine_documents, expected_shape) -# def test_num_grid_scan_fails_when_given_wrong_number_of_params( -# run_engine: RunEngine, -# x_axis: SimMotor, -# y_axis: SimMotor, -# ): -# with pytest.raises(ValueError): -# run_engine(num_grid_scan(detectors=[], params=[x_axis, 0, 1.1, 2, y_axis, 1.1])) - - @pytest.mark.parametrize( "x_start, x_stop, x_num, y_start, y_stop, y_num", ([-1.1, 1.1, 5, 2.2, -2.2, 3], [0, 1.1, 3, 2.2, 3.3, 5]), @@ -351,27 +342,6 @@ def test_num_rscan( assert_expected_shape(run_engine_documents, (num,)) -# @pytest.mark.parametrize( -# "x_list, y_list, num", ([[-1, 1], [2, 0], 0], [[-1, 1], [-1, 1], 3.5]) -# ) -# def test_num_rscan_fails_when_given_bad_info( -# run_engine: RunEngine, -# x_axis: SimMotor, -# x_list: list[float | int], -# y_axis: SimMotor, -# y_list: list[float | int], -# num: int, -# ): -# with pytest.raises(ValueError): -# run_engine( -# num_rscan( -# detectors=[], -# params=[x_axis, *x_list, y_axis, *y_list], -# num=num, -# ) -# ) - - @pytest.mark.parametrize( "trajectories_start_stop_num, snake_axes", [ From 9763df274f5a884b94f1e074ea63ba46838c9866 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 15:40:46 +0000 Subject: [PATCH 22/26] Reduce test_wrapped lines --- tests/plans/scans/test_wrapped.py | 82 +++++++++++++------------------ 1 file changed, 34 insertions(+), 48 deletions(-) diff --git a/tests/plans/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py index 0666c2caf99..88ff295f807 100644 --- a/tests/plans/scans/test_wrapped.py +++ b/tests/plans/scans/test_wrapped.py @@ -18,6 +18,7 @@ from ophyd_async.testing import assert_emitted from pydantic import ValidationError +from dodal.plans.scans import wrapped as sw from dodal.plans.scans.types import ( MovableListOfPoints, MovableStartStep, @@ -26,21 +27,6 @@ MovableStartStopStep, Number, ) -from dodal.plans.scans.wrapped import ( - count, - list_grid_rscan, - list_grid_scan, - list_rscan, - list_scan, - num_grid_rscan, - num_grid_scan, - num_rscan, - num_scan, - step_grid_rscan, - step_grid_scan, - step_rscan, - step_scan, -) def assert_expected_shape( @@ -68,7 +54,7 @@ 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(sw.count([det], num=3, delay=delay)) def test_count_detectors_validation(run_engine: RunEngine): @@ -80,7 +66,7 @@ def test_count_detectors_validation(run_engine: RunEngine): } for reason, dets in args.items(): with pytest.raises(ValidationError, match=reason): - run_engine(count(dets)) + run_engine(sw.count(dets)) def test_count_num_validation(det: StandardDetector, run_engine: RunEngine): @@ -91,7 +77,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(sw.count([det], num=num)) @pytest.mark.parametrize("num, shape", ([1, (1,)], [3, (3,)])) @@ -102,7 +88,7 @@ def test_count_plan_produces_expected_start_document( num: int, shape: tuple[int, ...], ): - run_engine(count([det], num=num)) + run_engine(sw.count([det], num=num)) start = run_engine_documents.get("start") assert start and len(start) == 1 run_start = cast(RunStart, start[0]) @@ -120,7 +106,7 @@ def test_count_plan_produces_expected_stop_document( num: int, length: tuple[int, ...], ): - run_engine(count([det], num=num)) + run_engine(sw.count([det], num=num)) stop = run_engine_documents.get("stop") assert stop and len(stop) == 1 run_stop = cast(RunStop, stop[0]) @@ -133,7 +119,7 @@ def test_count_plan_produces_expected_descriptor( run_engine_documents: Mapping[str, list[dict]], det: StandardDetector, ): - run_engine(count([det], num=1)) + run_engine(sw.count([det], num=1)) desc = run_engine_documents.get("descriptor") assert desc and len(desc) == 1 event_desc = cast(EventDescriptor, desc[0]) @@ -150,7 +136,7 @@ def test_count_plan_produces_expected_events( num: int, length: tuple[int, ...], ): - run_engine(count([det], num=num)) + run_engine(sw.count([det], num=num)) event_docs = run_engine_documents.get("event") assert event_docs and len(event_docs) == length for i in range(len(event_docs)): @@ -166,7 +152,7 @@ def test_count_plan_produces_expected_resources( det: StandardDetector, num: int, ): - run_engine(count([det], num=num)) + run_engine(sw.count([det], num=num)) stream_resource_docs = run_engine_documents.get("stream_resource") data_keys = [det.name, f"{det.name}-sum"] assert stream_resource_docs and len(stream_resource_docs) == len(data_keys) @@ -183,7 +169,7 @@ def test_count_plan_produces_expected_datums( num: int, length: tuple[int, ...], ): - run_engine(count([det], num=num)) + run_engine(sw.count([det], num=num)) stream_datum = run_engine_documents.get("stream_datum") data_keys = [det.name, f"{det.name}-sum"] assert stream_datum and len(stream_datum) == len(data_keys) * length @@ -243,7 +229,7 @@ def test_num_scan( num: int, ): run_engine( - num_scan( + sw.num_scan( detectors, trajectories_start_stop[0], *trajectories_start_stop[1:], num=num ) ) @@ -257,7 +243,7 @@ def test_num_scan_fails_when_given_wrong_number_of_params( y_axis: SimMotor, ): with pytest.raises(ValueError): - run_engine(num_scan([], x_axis, -1, 1, (y_axis, 1, 5, 1), num=5)) # type: ignore + run_engine(sw.num_scan([], x_axis, -1, 1, (y_axis, 1, 5, 1), num=5)) # type: ignore @pytest.mark.parametrize( @@ -278,7 +264,7 @@ def test_num_grid_scan( snake_axes: bool, ): run_engine( - num_grid_scan( + sw.num_grid_scan( detectors, trajectories_start_stop_num[0], *trajectories_start_stop_num[1:], @@ -307,7 +293,7 @@ def test_num_scan_fails_when_asked_to_snake_slow_axis( ): with pytest.raises(ValueError): run_engine( - num_grid_scan( + sw.num_grid_scan( [], (x_axis, x_start, x_stop, x_num), (y_axis, y_start, y_stop, y_num), @@ -334,7 +320,7 @@ def test_num_rscan( num: int, ): run_engine( - num_rscan( + sw.num_rscan( detectors, trajectories_start_stop[0], *trajectories_start_stop[1:], num=num ) ) @@ -360,7 +346,7 @@ def test_num_grid_rscan( snake_axes: bool, ): run_engine( - num_grid_rscan( + sw.num_grid_rscan( detectors, trajectories_start_stop_num[0], *trajectories_start_stop_num[1:], @@ -379,7 +365,7 @@ def test_num_grid_rscan_fails_when_asked_to_snake_slow_axis( ): with pytest.raises(ValueError): run_engine( - num_grid_rscan( + sw.num_grid_rscan( [], (x_axis, 1, 6, 10), (y_axis, -10, 0, 5), snake_axes=[x_axis] ) ) @@ -405,7 +391,7 @@ def test_list_scan( ): num = len(trajectories_with_list[0][1]) run_engine( - list_scan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) + sw.list_scan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) ) _assert_emitted(run_engine_documents, detectors, num) assert_expected_shape(run_engine_documents, (num,)) @@ -415,7 +401,7 @@ def test_list_scan_fails_with_differnt_list_lengths( run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor ): with pytest.raises(ValueError): - run_engine(list_scan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) + run_engine(sw.list_scan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) @pytest.mark.parametrize( @@ -439,7 +425,7 @@ def test_list_rscan( ): num = len(trajectories_with_list[0][1]) run_engine( - list_rscan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) + sw.list_rscan(detectors, trajectories_with_list[0], *trajectories_with_list[1:]) ) _assert_emitted(run_engine_documents, detectors, num) assert_expected_shape(run_engine_documents, (num,)) @@ -449,7 +435,7 @@ def test_list_rscan_fails_with_differnt_list_lengths( run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor ): with pytest.raises(ValueError): - run_engine(list_rscan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) + run_engine(sw.list_rscan([], (x_axis, [1, 2, 3, 4, 5]), (y_axis, [1, 2, 3, 4]))) @pytest.mark.parametrize( @@ -469,7 +455,7 @@ def test_list_grid_scan( shape = tuple(len(points) for _, points in trajectories_with_list) num = math.prod(shape) run_engine( - list_grid_scan( + sw.list_grid_scan( detectors, trajectories_with_list[0], *trajectories_with_list[1:] ) ) @@ -494,7 +480,7 @@ def test_list_grid_rscan( shape = tuple(len(points) for _, points in trajectories_with_list) num = math.prod(shape) run_engine( - list_grid_rscan( + sw.list_grid_rscan( detectors, trajectories_with_list[0], *trajectories_with_list[1:] ) ) @@ -519,7 +505,7 @@ def test_step_scan( expected_num: int, ): run_engine( - step_scan( + sw.step_scan( detectors, trajectories_start_stop_step[0], *trajectories_start_step, @@ -548,7 +534,7 @@ def test_step_grid_scan( snake: bool, ): run_engine( - step_grid_scan( + sw.step_grid_scan( detectors, trajectories_start_stop_step[0], *trajectories_start_stop_step[1:], @@ -576,7 +562,7 @@ def test_step_rscan( expected_num: int, ): run_engine( - step_rscan( + sw.step_rscan( detectors, trajectories_start_stop_step[0], *trajectories_start_step, @@ -605,7 +591,7 @@ def test_step_grid_rscan( snake: bool, ): run_engine( - step_grid_rscan( + sw.step_grid_rscan( detectors, trajectories_start_stop_step[0], *trajectories_start_stop_step[1:], @@ -628,7 +614,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_first_axis( "Received 3 values: ('x_axis', 1, 5)" ), ): - run_engine(step_grid_scan([], (x_axis, 1, 5))) # type: ignore + run_engine(sw.step_grid_scan([], (x_axis, 1, 5))) # type: ignore def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( @@ -644,7 +630,7 @@ def test_step_grid_scan_fails_when_given_wrong_number_of_args_for_other_axis( "Received 3 values: ('y_axis', 1, 2)" ), ): - run_engine(step_grid_scan([], (x_axis, 1, 5, 1), (y_axis, 1, 2))) # type: ignore + run_engine(sw.step_grid_scan([], (x_axis, 1, 5, 1), (y_axis, 1, 2))) # type: ignore def test_step_scan_fails_with_step_size_zero( @@ -659,7 +645,7 @@ def test_step_scan_fails_with_step_size_zero( "Received (x_axis, 1, 5, 0)" ), ): - run_engine(step_scan([], (x_axis, 1, 5, 0))) + run_engine(sw.step_scan([], (x_axis, 1, 5, 0))) def test_step_scan_fails_with_start_and_stop_being_same_value( @@ -676,7 +662,7 @@ def test_step_scan_fails_with_start_and_stop_being_same_value( f"Received ({x_axis.name}, {start}, {stop}, {step})." ), ): - run_engine(step_scan([], (x_axis, start, stop, step))) + run_engine(sw.step_scan([], (x_axis, start, stop, step))) def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( @@ -692,7 +678,7 @@ def test_step_scan_fails_when_given_wrong_number_of_args_for_second_axes( "Received 4 values: ('y_axis', 1, 5, 1)" ), ): - run_engine(step_scan([], (x_axis, 0, 1, 0.1), (y_axis, 1, 5, 1))) # type: ignore + run_engine(sw.step_scan([], (x_axis, 0, 1, 0.1), (y_axis, 1, 5, 1))) # type: ignore def test_scan_fails_when_not_using_movable( @@ -707,7 +693,7 @@ def test_scan_fails_when_not_using_movable( "Received ('y_axis', 1, 5, 1)." ), ): - run_engine(step_scan([], (x_axis, 0, 1, 0.1), ("y_axis", 1, 5, 1))) # type: ignore + run_engine(sw.step_scan([], (x_axis, 0, 1, 0.1), ("y_axis", 1, 5, 1))) # type: ignore def test_scan_fails_when_using_invalid_structure( @@ -721,4 +707,4 @@ def test_scan_fails_when_using_invalid_structure( "Received ('x_axis', 0, 1, [0.1])." ), ): - run_engine(step_rscan([], (x_axis, 0, 1, [0.1]))) # type: ignore + run_engine(sw.step_rscan([], (x_axis, 0, 1, [0.1]))) # type: ignore From 4ea39fe1d00b571e8f508cb3d2dd14c23a9560e7 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 16:02:25 +0000 Subject: [PATCH 23/26] Update doc strings for scan plans --- src/dodal/plans/scans/annotations.py | 18 +- src/dodal/plans/scans/wrapped.py | 377 +++++++++++++++++++++++---- 2 files changed, 334 insertions(+), 61 deletions(-) diff --git a/src/dodal/plans/scans/annotations.py b/src/dodal/plans/scans/annotations.py index 95f37de6eec..d6f215aa81e 100644 --- a/src/dodal/plans/scans/annotations.py +++ b/src/dodal/plans/scans/annotations.py @@ -22,8 +22,7 @@ MovableStartStepA = A[ MovableStartStep, Field( - description="Additional trajectories, each specified as a tuple of " - "(movable, start, step)." + description="Trajectory defined by a movable, start position, and step size." ), BeforeValidator( trajectory_validator( @@ -37,8 +36,7 @@ MovableStartStopA = A[ MovableStartStop, Field( - description="Additional trajectories, each specified as a tuple of " - "(movable, start, stop)." + description="Trajectory defined by a movable, start position, and stop position.", ), BeforeValidator( trajectory_validator( @@ -52,8 +50,8 @@ MovableStartStopNumA = A[ MovableStartStopNum, Field( - description="Additional trajectories, each specified as a tuple of " - "(movable, start, stop, num)." + description="Trajectory defined by a movable, start position, stop position, " + "and number of points." ), BeforeValidator( trajectory_validator( @@ -67,10 +65,7 @@ MovableListOfPointsA = A[ MovableListOfPoints, Field( - description="List of tuples (device, positions). For concurrent \ - trajectories, provide '[(movable1, [point1, point2, ...]), (movable2, \ - [point1, point2, ...]), ... , (movableN, [point1, point2, ...])]'. Number \ - of points for each movable must be equal." + description="Trajectory defined by a movable and a list of positions to move to." ), BeforeValidator( trajectory_validator( @@ -84,7 +79,8 @@ MovableStartStopStepA = A[ MovableStartStopStep, Field( - description="Tuple containing (movable, start, stop, step) for a scan trajectory." + description="Trajectory defined by a movable, start position, stop position, " + "and step size." ), BeforeValidator( trajectory_validator( diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py index 57ff07ff9f3..d8a6ce66ea3 100644 --- a/src/dodal/plans/scans/wrapped.py +++ b/src/dodal/plans/scans/wrapped.py @@ -57,10 +57,27 @@ def count( ] = 0.0, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Reads from a number of devices. + """Read from a number of devices. - Wraps bluesky.plans.count(det, num, delay, md=metadata) exposing only serializable - parameters and metadata. + Args: + detectors: Devices to trigger and read. + num: Number of readings to collect. + delay: Delay between readings in seconds. A single value applies to + every gap. A sequence specifies an individual delay for each gap + and must contain ``num - 1`` values. + metadata: Additional metadata to include in the run. + + Examples: + Collect 10 readings with a 1-second delay between each reading:: + + count([detector], num=10, delay=1.0) + + Use a different delay for each gap:: + + count([detector], num=3, delay=[0.5, 1.0]) + + Wraps: + ``bluesky.plans.count(det, num, delay, md=metadata)``. """ if isinstance(delay, Sequence): assert len(delay) == num - 1, ( @@ -76,20 +93,40 @@ def count( def num_scan( detectors: DetectorsA, trajectory: MovableStartStopA, - *extra_axes: MovableStartStopA, + *extra_trajectories: MovableStartStopA, num: int, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent single or multi-motor trajector(y/ies). + """Scan one or more motors over a specified range. + + The scan is defined by the number of points along each trajectory. + All trajectories are scanned concurrently. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, start position, + and stop position. + *extra_trajectories: Additional trajectories to scan concurrently. + num: Number of points in the scan. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from 0 to 10 in 11 points:: + + num_scan([detector], (motor, 0, 10), num=11) + + Scan two motors concurrently:: - The scan is defined by number of points along scan trajector(y/ies). Wraps - bluesky.plans.scan(det, *args, num, md=metadata). + num_scan([detector], (x_motor, 0, 10), (y_motor, 5, 15), num=11) + + Wraps: + ``bluesky.plans.scan(det, *args, num, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = (num,) yield from bp.scan( - detectors, *trajectory, *flatten(extra_axes), num=num, md=metadata + detectors, *trajectory, *flatten(extra_trajectories), num=num, md=metadata ) @@ -99,14 +136,35 @@ def num_grid_scan( detectors: DetectorsA, trajectory: MovableStartStopNumA, *extra_trajectories: MovableStartStopNumA, - snake_axes: Iterable[Movable] | bool = False, + snake_axes: Iterable[Movable] | bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: """Scan independent multi-motor trajectories. - The scan is defined by number of points along scan trajectories. Snakes all fast - axes by default (all axes but the first axis provided). Wraps - bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata). + Each trajectory is defined by a movable, start position, stop position, + and number of points. The trajectories are scanned independently to + produce a grid. By default, all axes except the first axis are snaked. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, start position, + stop position, and number of points. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Axes to snake, or ``True`` to snake all axes except the + first axis. ``False`` disables snaking. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from 0 to 10 using 11 points:: + + num_grid_scan([detector], (x_motor, 0, 10, 11)) + + Scan two motors over a 2D grid:: + + num_grid_scan([detector], (x_motor, 0, 10, 11), (y_motor, 0, 5, 6)) + + Wraps: + ``bluesky.plans.grid_scan(det, *args, snake_axes, md=metadata)``. """ yield from bp.grid_scan( detectors, @@ -126,14 +184,35 @@ def num_rscan( num: int, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent trajector(y/ies), relative to current position(s). + """Scan one or more motors relative to their current positions. + + Each trajectory defines a relative start and stop position. The scan is + performed using the specified number of points, with all trajectories + scanned concurrently. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, relative start + position, and relative stop position. + *extra_trajectories: Additional trajectories to scan concurrently. + num: Number of points in the scan. + metadata: Additional metadata to include in the run. - The scan is defined by number of points along scan trajector(y/ies). Wraps - bluesky.plans.rel_scan(det, *args, num, md=metadata). + Examples: + Scan one motor from its current position to 10 units above it, + using 11 points:: + + num_rscan([detector], (x_motor, 0, 10), num=11) + + Scan two motors concurrently relative to their current positions:: + + num_rscan([detector], (x_motor, 0, 10), (y_motor, -5, 5), num=11) + + Wraps: + ``bluesky.plans.rel_scan(det, *args, num, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = (num,) - yield from bp.rel_scan( detectors, *trajectory, *flatten(extra_trajectories), num=num, md=metadata ) @@ -148,11 +227,34 @@ def num_grid_rscan( snake_axes: list | bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan independent trajectories, relative to current positions. + """Scan independent trajectories relative to current positions. + + Each trajectory is defined by a movable, relative start position, relative + stop position, and number of points. The trajectories are scanned + independently to produce a grid. By default, all axes except the first + axis are snaked. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, relative start + position, relative stop position, and number of points. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Axes to snake, or ``True`` to snake all axes except the + first axis. ``False`` disables snaking. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from its current position to 10 units above it, + using 11 points:: + + num_grid_rscan([detector], (x_motor, 0, 10, 11)) + + Scan two motors over a 2D grid relative to their current positions:: - The scan is defined by number of points along scan trajectories. Snakes all fast - axes by default (all axes but the first axis provided). Wraps - bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata). + num_grid_rscan([detector], (x_motor, 0, 10, 11), (y_motor, -5, 5, 11)) + + Wraps: + ``bluesky.plans.rel_grid_scan(det, *args, snake_axes, md=metadata)``. """ yield from bp.rel_grid_scan( detectors, @@ -171,10 +273,31 @@ def list_scan( *extra_trajectories: MovableListOfPointsA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent single or multi-motor trajector(y/ies). + """Scan one or more motors through specified lists of positions. + + Each trajectory is defined by a movable and a list of positions. All + trajectories are scanned concurrently, with one point from each + trajectory used at each scan step. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable and a list of + positions. + *extra_trajectories: Additional trajectories to scan concurrently. + metadata: Additional metadata to include in the run. + + Examples: + Scan a motor through a list of positions:: + + list_scan([detector], (x_motor, [0, 1, 2, 3])) - The scan is defined by providing a list of points for each scan trajectory. - Wraps bluesky.plans.list_scan(det, *args, md=metadata). + Scan two motors concurrently through corresponding lists of + positions:: + + list_scan([detector], (x_motor, [0, 1, 2]), (y_motor, [10, 20, 30])) + + Wraps: + ``bluesky.plans.list_scan(det, *args, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = make_list_scan_shape( @@ -194,14 +317,35 @@ def list_grid_scan( detectors: DetectorsA, trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, - snake_axes: bool = False, + snake_axes: bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan independent trajectories. + """Scan independent trajectories through specified lists of positions. + + Each trajectory is defined by a movable and a list of positions. The + trajectories are scanned independently to produce a grid. By default, + snaking is disabled. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable and a list of + positions. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Whether to snake the fast axes. ``False`` disables + snaking. + metadata: Additional metadata to include in the run. - The scan is defined by providing a list of points for each scan trajectory. Snakes - all fast axes by default (all axes but the first axis provided). Wraps - bluesky.plans.list_grid_scan(det, *args, md=metadata). + Examples: + Scan one motor through a list of positions:: + + list_grid_scan([detector], (x_motor, [0, 1, 2, 3])) + + Scan two motors over a 2D grid:: + + list_grid_scan([detector], (x_motor, [0, 1, 2]), (y_motor, [10, 20, 30])) + + Wraps: + ``bluesky.plans.list_grid_scan(det, *args, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = make_list_scan_shape( @@ -223,10 +367,31 @@ def list_rscan( *extra_trajectories: MovableListOfPointsA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent trajector(y/ies), relative to current position. + """Scan one or more motors through relative positions. + + Each trajectory is defined by a movable and a list of positions relative + to the motor's current position. All trajectories are scanned concurrently, + with one point from each trajectory used at each scan step. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable and a list of + relative positions. + *extra_trajectories: Additional trajectories to scan concurrently. + metadata: Additional metadata to include in the run. - The scan is defined by providing a list of points for each scan trajectory. - Wraps bluesky.plans.rel_list_scan(det, *args, md=metadata). + Examples: + Scan a motor through relative positions:: + + list_rscan([detector], (x_motor, [0, 1, 2, 3])) + + Scan two motors concurrently through corresponding relative + positions:: + + list_rscan([detector], (x_motor, [0, 1, 2]), (y_motor, [-1, 0, 1])) + + Wraps: + ``bluesky.plans.rel_list_scan(det, *args, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = make_list_scan_shape( @@ -246,11 +411,32 @@ def list_grid_rscan( snake_axes: bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan independent trajectories, relative to current positions. + """Scan independent trajectories through relative positions. + + Each trajectory is defined by a movable and a list of positions relative + to its current position. The trajectories are scanned independently to + produce a grid. By default, all axes except the first axis are snaked. - The scan is defined by providing a list of points for each scan trajectory. Snakes - all fast axes by default (all axes but the first axis provided). Wraps - bluesky.plans.rel_list_grid_scan(det, *args, md=metadata). + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable and a list of + relative positions. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Whether to snake the fast axes. ``True`` enables snaking + and ``False`` disables it. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor through relative positions:: + + list_grid_rscan([detector], (x_motor, [0, 1, 2, 3])) + + Scan two motors over a 2D grid relative to their current positions:: + + list_grid_rscan([detector], (x_motor, [0, 1, 2]), (y_motor, [-1, 0, 1])) + + Wraps: + ``bluesky.plans.rel_list_grid_scan(det, *args, md=metadata)``. """ metadata = metadata or {} metadata["shape"] = make_list_scan_shape( @@ -272,10 +458,34 @@ def step_scan( *extra_trajectories: MovableStartStepA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent trajectories with specified step size. + """Scan one or more motors using specified step sizes. + + The primary trajectory is defined by a movable, start position, stop + position, and step size. Additional trajectories are defined by a + movable, start position, and step size and contain the same number of + points as the primary trajectory. All trajectories are scanned + concurrently. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, start position, + stop position, and step size. + *extra_trajectories: Additional trajectories to scan concurrently, + defined by a movable, start position, and step size. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from 0 to 10 in steps of 1:: + + step_scan([detector], (x_motor, 0, 10, 1)) + + Scan two motors concurrently, with the second motor starting at 5 + and using the same number of points as the primary trajectory:: - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_scan(det, *args, md=metadata). + step_scan([detector], (x_motor, 0, 10, 1), (y_motor, 5, 0.5)) + + Wraps: + ``bluesky.plans.list_scan(det, *args, md=metadata)``. """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 args, shape = make_step_scan_args_and_shape(trajectory, extra_trajectories) @@ -293,11 +503,32 @@ def step_grid_scan( snake_axes: bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan independent trajectories with specified step size. + """Scan independent trajectories using specified step sizes. + + Each trajectory is defined by a movable, start position, stop position, + and step size. The trajectories are scanned independently to produce a + grid. By default, all axes except the first axis are snaked. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, start position, + stop position, and step size. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Whether to snake the fast axes. ``True`` enables snaking + and ``False`` disables it. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from 0 to 10 in steps of 1:: + + step_grid_scan([detector], (x_motor, 0, 10, 1)) - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by - default (all axes but the first axis provided). + Scan two motors over a 2D grid:: + + step_grid_scan([detector], (x_motor, 0, 10, 1), (y_motor, 0, 5, 1)) + + Wraps: + ``bluesky.plans.list_grid_scan(det, *args, md=metadata)``. """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 args, shape = make_step_grid_scan_args_and_shape([trajectory, *extra_trajectories]) @@ -316,10 +547,34 @@ def step_rscan( *extra_trajectories: MovableStartStepA, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan concurrent trajectories with specified step size, relative to position. + """Scan one or more motors using relative step sizes. + + The primary trajectory is defined by a movable, relative start position, + relative stop position, and step size. Additional trajectories are defined + by a movable, relative start position, and step size and contain the same + number of points as the primary trajectory. All trajectories are scanned + concurrently. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, relative start + position, relative stop position, and step size. + *extra_trajectories: Additional trajectories to scan concurrently, + defined by a movable, relative start position, and step size. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from its current position to 10 units above it, + in steps of 1:: - Generates list(s) of points for each trajectory, used with - bluesky.plans.rel_list_scan(det, *args, md=metadata). + step_rscan([detector], (x_motor, 0, 10, 1)) + + Scan two motors concurrently using relative positions:: + + step_rscan([detector], (x_motor, 0, 10, 1), (y_motor, -5, 0.5)) + + Wraps: + ``bluesky.plans.rel_list_scan(det, *args, md=metadata)``. """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 args, shape = make_step_scan_args_and_shape(trajectory, extra_trajectories) @@ -334,14 +589,36 @@ def step_grid_rscan( detectors: DetectorsA, trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStopStepA, - snake_axes: bool = True, # Currently specifying axes to snake is not supported + snake_axes: bool = True, metadata: dict[str, Any] | None = None, ) -> MsgGenerator: - """Scan independent trajectories with specified step size, relative to position. + """Scan independent trajectories using relative step sizes. + + Each trajectory is defined by a movable, relative start position, relative + stop position, and step size. The trajectories are scanned independently + to produce a grid. By default, all axes except the first axis are snaked. + + Args: + detectors: Devices to trigger and read at each scan point. + trajectory: Primary trajectory defined by a movable, relative start + position, relative stop position, and step size. + *extra_trajectories: Additional trajectories to include in the grid. + snake_axes: Whether to snake the fast axes. ``True`` enables snaking + and ``False`` disables it. + metadata: Additional metadata to include in the run. + + Examples: + Scan one motor from its current position to 10 units above it, + in steps of 1:: + + step_grid_rscan([detector], (x_motor, 0, 10, 1)) + + Scan two motors over a 2D grid relative to their current positions:: + + step_grid_rscan([detector], (x_motor, 0, 10, 1), (y_motor, 0, 5, 1)) - Generates list(s) of points for each trajectory, used with - bluesky.plans.list_grid_scan(det, *args, md=metadata). Snakes all fast axes by - default (all axes but the first axis provided). + Wraps: + ``bluesky.plans.rel_list_grid_scan(det, *args, md=metadata)``. """ # TODO: move to using Linspace spec and spec_scan when stable and tested at v1.0 args, shape = make_step_grid_scan_args_and_shape([trajectory, *extra_trajectories]) From 0d65b0b02a6b7674f538cabda2286867873b5d20 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 16:09:48 +0000 Subject: [PATCH 24/26] Add missing plan typing for spec_path --- src/dodal/plans/scans/spec_path.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/dodal/plans/scans/spec_path.py b/src/dodal/plans/scans/spec_path.py index 85617f7019f..4a3edfec20f 100644 --- a/src/dodal/plans/scans/spec_path.py +++ b/src/dodal/plans/scans/spec_path.py @@ -4,6 +4,7 @@ import bluesky.plans as bp from bluesky.protocols import Movable +from bluesky.utils import plan from cycler import Cycler, cycler from pydantic import Field, validate_call from scanspec.specs import Spec @@ -15,6 +16,7 @@ @attach_data_session_metadata_decorator() @validate_call(config={"arbitrary_types_allowed": True}) +@plan def spec_scan( detectors: DetectorsA, spec: Annotated[ From f32aa86677a95a644ceb89b3af31f78c4b40eb4f Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 16:10:00 +0000 Subject: [PATCH 25/26] Assume blueapi supports *args --- src/dodal/plans/scans/wrapped.py | 3 --- tests/plans/test_compliance.py | 7 +++---- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py index d8a6ce66ea3..36f763a39d8 100644 --- a/src/dodal/plans/scans/wrapped.py +++ b/src/dodal/plans/scans/wrapped.py @@ -31,9 +31,6 @@ Non-serialisable fields are ignored when they are optional. https://github.com/DiamondLightSource/blueapi/issues/711 -Using *args in plans is currently not supported. -https://github.com/DiamondLightSource/blueapi/issues/1450 - We may also need other adjustments for UI purposes, e.g. - Forcing uniqueness or orderedness of Readables. - Limits and metadata (e.g. units). diff --git a/tests/plans/test_compliance.py b/tests/plans/test_compliance.py index 2f47e95a80d..6a390fc4b2c 100644 --- a/tests/plans/test_compliance.py +++ b/tests/plans/test_compliance.py @@ -47,10 +47,9 @@ def get_all_available_generators(mod: ModuleType) -> Iterable[PlanGenerator]: def assert_hard_requirements(plan: PlanGenerator, signature: inspect.Signature): assert plan.__doc__ is not None, f"'{plan.__name__}' has no docstring" for parameter in signature.parameters.values(): - assert ( - # parameter.kind is not parameter.VAR_POSITIONAL # BlueAPI should support *args - parameter.kind is not parameter.VAR_KEYWORD - ), f"'{plan.__name__}' has variadic arguments" + assert parameter.kind is not parameter.VAR_KEYWORD, ( + f"'{plan.__name__}' has variadic arguments" + ) def assert_metadata_requirements(plan: PlanGenerator, signature: inspect.Signature): From 057b47ddea4dcd1a906c3f8cfb2d569d73e125af Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Fri, 4 Sep 2026 16:19:22 +0000 Subject: [PATCH 26/26] Use CustomPlanMetadata --- src/dodal/plans/scans/spec_path.py | 6 +++--- src/dodal/plans/scans/wrapped.py | 29 ++++++++++++++--------------- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/src/dodal/plans/scans/spec_path.py b/src/dodal/plans/scans/spec_path.py index 4a3edfec20f..f1deb3078c2 100644 --- a/src/dodal/plans/scans/spec_path.py +++ b/src/dodal/plans/scans/spec_path.py @@ -1,10 +1,10 @@ import operator from functools import reduce -from typing import Annotated, Any +from typing import Annotated import bluesky.plans as bp from bluesky.protocols import Movable -from bluesky.utils import plan +from bluesky.utils import CustomPlanMetadata, plan from cycler import Cycler, cycler from pydantic import Field, validate_call from scanspec.specs import Spec @@ -23,7 +23,7 @@ def spec_scan( Spec[Movable], Field(description="ScanSpec modelling the path of the scan"), ], - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Generic plan for reading `detectors` at every point of a ScanSpec `Spec`. A `Spec` is an N-dimensional path. diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py index 36f763a39d8..e0960b2bda3 100644 --- a/src/dodal/plans/scans/wrapped.py +++ b/src/dodal/plans/scans/wrapped.py @@ -1,10 +1,9 @@ from collections.abc import Iterable, Sequence from typing import Annotated as A -from typing import Any import bluesky.plans as bp from bluesky.protocols import Movable -from bluesky.utils import plan +from bluesky.utils import CustomPlanMetadata, plan from pydantic import Field, NonNegativeFloat, validate_call from dodal.common import MsgGenerator @@ -52,7 +51,7 @@ def count( json_schema_extra={"units": "s"}, ), ] = 0.0, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Read from a number of devices. @@ -92,7 +91,7 @@ def num_scan( trajectory: MovableStartStopA, *extra_trajectories: MovableStartStopA, num: int, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors over a specified range. @@ -134,7 +133,7 @@ def num_grid_scan( trajectory: MovableStartStopNumA, *extra_trajectories: MovableStartStopNumA, snake_axes: Iterable[Movable] | bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent multi-motor trajectories. @@ -179,7 +178,7 @@ def num_rscan( trajectory: MovableStartStopA, *extra_trajectories: MovableStartStopA, num: int, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors relative to their current positions. @@ -222,7 +221,7 @@ def num_grid_rscan( trajectory: MovableStartStopNumA, *extra_trajectories: MovableStartStopNumA, snake_axes: list | bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent trajectories relative to current positions. @@ -268,7 +267,7 @@ def list_scan( detectors: DetectorsA, trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors through specified lists of positions. @@ -315,7 +314,7 @@ def list_grid_scan( trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, snake_axes: bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent trajectories through specified lists of positions. @@ -362,7 +361,7 @@ def list_rscan( detectors: DetectorsA, trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors through relative positions. @@ -406,7 +405,7 @@ def list_grid_rscan( trajectory: MovableListOfPointsA, *extra_trajectories: MovableListOfPointsA, snake_axes: bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent trajectories through relative positions. @@ -453,7 +452,7 @@ def step_scan( detectors: DetectorsA, trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStepA, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors using specified step sizes. @@ -498,7 +497,7 @@ def step_grid_scan( trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStopStepA, snake_axes: bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent trajectories using specified step sizes. @@ -542,7 +541,7 @@ def step_rscan( detectors: DetectorsA, trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStepA, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan one or more motors using relative step sizes. @@ -587,7 +586,7 @@ def step_grid_rscan( trajectory: MovableStartStopStepA, *extra_trajectories: MovableStartStopStepA, snake_axes: bool = True, - metadata: dict[str, Any] | None = None, + metadata: CustomPlanMetadata | None = None, ) -> MsgGenerator: """Scan independent trajectories using relative step sizes.