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..d6f215aa81e --- /dev/null +++ b/src/dodal/plans/scans/annotations.py @@ -0,0 +1,92 @@ +from typing import Annotated as A + +from pydantic import BeforeValidator, Field + +from dodal.plans.scans.types import ( + Detectors, + MovableListOfPoints, + MovableStartStep, + MovableStartStop, + MovableStartStopNum, + MovableStartStopStep, +) +from dodal.plans.scans.validators import trajectory_validator + +DetectorsA = A[ + Detectors, + Field( + description="Set of readable devices, will take a reading at each point", + ), +] + +MovableStartStepA = A[ + MovableStartStep, + Field( + description="Trajectory defined by a movable, start position, and step size." + ), + BeforeValidator( + trajectory_validator( + length=3, + template="(movable, start, step)", + expected_type=MovableStartStep, + ) + ), +] + +MovableStartStopA = A[ + MovableStartStop, + Field( + description="Trajectory defined by a movable, start position, and stop position.", + ), + BeforeValidator( + trajectory_validator( + length=3, + template="(movable, start, stop)", + expected_type=MovableStartStop, + ) + ), +] + +MovableStartStopNumA = A[ + MovableStartStopNum, + Field( + description="Trajectory defined by a movable, start position, stop position, " + "and number of points." + ), + BeforeValidator( + trajectory_validator( + length=4, + template="(movable, start, stop, num)", + expected_type=MovableStartStopNum, + ) + ), +] + +MovableListOfPointsA = A[ + MovableListOfPoints, + Field( + description="Trajectory defined by a movable and a list of positions to move to." + ), + BeforeValidator( + trajectory_validator( + length=2, + template="(movable, [point1, point2, ...])", + expected_type=MovableListOfPoints, + ) + ), +] + +MovableStartStopStepA = A[ + MovableStartStopStep, + Field( + description="Trajectory defined by a movable, start position, stop position, " + "and step size." + ), + BeforeValidator( + trajectory_validator( + length=4, + template="(movable, start, stop, step)", + expected_type=MovableStartStopStep, + ) + ), +] diff --git a/src/dodal/plans/spec_path.py b/src/dodal/plans/scans/spec_path.py similarity index 79% rename from src/dodal/plans/spec_path.py rename to src/dodal/plans/scans/spec_path.py index 6f7f5a47f0a..f1deb3078c2 100644 --- a/src/dodal/plans/spec_path.py +++ b/src/dodal/plans/scans/spec_path.py @@ -1,32 +1,29 @@ 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, Readable +from bluesky.protocols import Movable +from bluesky.utils import CustomPlanMetadata, plan 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}) +@plan 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"), ], - 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. @@ -44,7 +41,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/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 new file mode 100644 index 00000000000..2f97e12236a --- /dev/null +++ b/src/dodal/plans/scans/utils.py @@ -0,0 +1,215 @@ +from collections.abc import Iterable, Sequence +from decimal import Decimal +from typing import TypeVar, cast + +import numpy as np +from bluesky.protocols import HasName + +from dodal.plans.scans.types import ( + MovableListOfPoints, + MovableStartStep, + MovableStartStopNum, + MovableStartStopStep, + Number, +) + +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 str(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. + + 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( + values: list[Number], + params: list[Number], +) -> list[Number]: + """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(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. + 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: + 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(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 + position. The resulting values are rounded to the greatest decimal + precision of the input parameters to avoid floating-point representation + artefacts. + + Args: + values: A tuple containing the movable, start position, stop position, + and number of points. + + Returns: + A list containing ``num`` scan positions. + + Raises: + ValueError: If ``step`` or ``num`` is zero. + """ + movable, start, step, num = values + if num == 0 or step == 0: + raise ValueError( + "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})." + ) + stepped_list = [start + (n * step) for n in range(num)] + 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, ...]]: + """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, _, _, _ = trajectory + movable_values = _make_stepped_list_step(trajectory) + shape = [len(movable_values)] + step_scan_args: list[MovableListOfPoints] = [(movable, movable_values)] + + 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((movable, 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, ...]]: + """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] = [] + 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) diff --git a/src/dodal/plans/scans/validators.py b/src/dodal/plans/scans/validators.py new file mode 100644 index 00000000000..8618b836665 --- /dev/null +++ b/src/dodal/plans/scans/validators.py @@ -0,0 +1,74 @@ +from collections.abc import Callable +from typing import Any + +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, + 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}.") + + 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 implement the Movable protocol. " + f"{get_bluesky_obj_name(movable)} does not implement Movable. " + f"Received {formatted_values}." + ) + + if len(value) != length: + raise ValueError( + f"Trajectory must contain exactly {length} values. " + f"Expected {template}. Received {len(value)} values: {formatted_values!r}" + ) + 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 diff --git a/src/dodal/plans/scans/wrapped.py b/src/dodal/plans/scans/wrapped.py new file mode 100644 index 00000000000..e0960b2bda3 --- /dev/null +++ b/src/dodal/plans/scans/wrapped.py @@ -0,0 +1,625 @@ +from collections.abc import Iterable, Sequence +from typing import Annotated as A + +import bluesky.plans as bp +from bluesky.protocols import Movable +from bluesky.utils import CustomPlanMetadata, 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, + MovableListOfPointsA, + MovableStartStepA, + MovableStartStopA, + MovableStartStopNumA, + 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, +) + +"""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 + +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}) +@plan +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: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """Read from a number of devices. + + 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, ( + 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}) +@plan +def num_scan( + detectors: DetectorsA, + trajectory: MovableStartStopA, + *extra_trajectories: MovableStartStopA, + num: int, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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:: + + 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_trajectories), num=num, md=metadata + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +@plan +def num_grid_scan( + detectors: DetectorsA, + trajectory: MovableStartStopNumA, + *extra_trajectories: MovableStartStopNumA, + snake_axes: Iterable[Movable] | bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """Scan independent multi-motor trajectories. + + 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, + *trajectory, + *flatten(extra_trajectories), + snake_axes=snake_axes, + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +@plan +def num_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopA, + *extra_trajectories: MovableStartStopA, + num: int, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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. + + 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 + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +@plan +def num_grid_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopNumA, + *extra_trajectories: MovableStartStopNumA, + snake_axes: list | bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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:: + + 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, + *trajectory, + *flatten(extra_trajectories), + snake_axes=snake_axes, + md=metadata, + ) + + +@validate_call(config={"arbitrary_types_allowed": True}) +@plan +def list_scan( + detectors: DetectorsA, + trajectory: MovableListOfPointsA, + *extra_trajectories: MovableListOfPointsA, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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])) + + 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( + [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}) +@plan +def list_grid_scan( + detectors: DetectorsA, + trajectory: MovableListOfPointsA, + *extra_trajectories: MovableListOfPointsA, + snake_axes: bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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. + + 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( + [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}) +@plan +def list_rscan( + detectors: DetectorsA, + trajectory: MovableListOfPointsA, + *extra_trajectories: MovableListOfPointsA, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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. + + 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( + [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}) +@plan +def list_grid_rscan( + detectors: DetectorsA, + trajectory: MovableListOfPointsA, + *extra_trajectories: MovableListOfPointsA, + snake_axes: bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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. + + 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( + [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}) +@plan +def step_scan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStepA, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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:: + + 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) + 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}) +@plan +def step_grid_scan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStopStepA, + snake_axes: bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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)) + + 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]) + 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}) +@plan +def step_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStepA, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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:: + + 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) + metadata = metadata or {} + metadata["shape"] = shape + yield from bp.rel_list_scan(detectors, *flatten(args), md=metadata) + + +@validate_call(config={"arbitrary_types_allowed": True}) +@plan +def step_grid_rscan( + detectors: DetectorsA, + trajectory: MovableStartStopStepA, + *extra_trajectories: MovableStartStopStepA, + snake_axes: bool = True, + metadata: CustomPlanMetadata | None = None, +) -> MsgGenerator: + """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)) + + 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]) + 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 1ea01906db8..00000000000 --- a/src/dodal/plans/wrapped.py +++ /dev/null @@ -1,577 +0,0 @@ -from collections.abc import Sequence -from decimal import Decimal -from typing import Annotated, Any - -import bluesky.plans as bp -import numpy as np -from bluesky.protocols import Movable, Readable -from ophyd_async.core import AsyncReadable -from pydantic import Field, NonNegativeFloat, validate_call - -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 -https://github.com/DiamondLightSource/blueapi/issues/474 - -Non-serialisable fields are ignored when they are optional -https://github.com/DiamondLightSource/blueapi/issues/711 - -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: 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) - - -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 - - -@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[ - list[tuple[Movable, list[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). - """ - # 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 - - yield from bp.scan(tuple(detectors), *args, 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[ - list[tuple[Movable, list[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). - """ - # 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) - - -@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[ - list[tuple[Movable, list[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). - """ - # 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 - - yield from bp.rel_scan(tuple(detectors), *args, 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[ - list[tuple[Movable, list[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). - """ - # 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 - ) - - -def _make_list_scan_args(params: list[tuple[Movable, list[float | int]]], grid: bool): - shape = [] - args = [] - 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 - - -@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[tuple[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). - """ - args, shape = _make_list_scan_args(params=params, grid=False) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.list_scan(tuple(detectors), *args, md=metadata) - - -@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[ - list[tuple[Movable, list[float | int]]], - Field( - description="List of tuples (device, positions). 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). - """ - args, shape = _make_list_scan_args(params=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 list_rscan( - detectors: Annotated[ - Sequence[Readable | AsyncReadable], - Field( - description="Set of readable devices, will take a reading at each point", - ), - ], - params: Annotated[ - list[tuple[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 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). - """ - 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) - - -@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[ - list[tuple[Movable, list[float | int]]], - Field( - description="List of tuples (device, positions). 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). - """ - args, shape = _make_list_scan_args(params=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 - ) - - -def _round_list_elements(stepped_list, params) -> list[float]: - 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: - if start == stop: - raise ValueError( - f"Start ({start}) and stop ({stop}) values cannot be the same." - ) - 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, step, num) -> list: - 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( - params: list[tuple[Movable, list[float | int]]], grid: bool -) -> tuple[list[Any], list[float]]: - args = [] - 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 - - -@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[ - list[tuple[Movable, list[float | int]]], - Field( - description="List of tuples (device, parameter). 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(params, grid=False) - metadata = metadata or {} - metadata["shape"] = shape - - yield from bp.list_scan(tuple(detectors), *args, md=metadata) - - -@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[ - list[tuple[Movable, list[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(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[ - list[tuple[Movable, list[float | int]]], - Field( - description="List of tuples (device, parameter). 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(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[ - list[tuple[Movable, list[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, 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(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..1ac3ffd10c6 --- /dev/null +++ b/tests/plans/scans/test_utils.py @@ -0,0 +1,137 @@ +import re + +import pytest +from ophyd_async.sim import SimMotor + +from dodal.plans.scans.types 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 + + +@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( + 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(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 + 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(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 == num + assert stepped_list[0] / stepped_list[-1] == -1 + assert stepped_list[10] == 0 + + +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( + x_axis: SimMotor, +): + with pytest.raises( + ValueError, + 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/scans/test_wrapped.py b/tests/plans/scans/test_wrapped.py new file mode 100644 index 00000000000..88ff295f807 --- /dev/null +++ b/tests/plans/scans/test_wrapped.py @@ -0,0 +1,710 @@ +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.sim import SimMotor +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, + MovableStartStop, + MovableStartStopNum, + MovableStartStopStep, + Number, +) + + +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(sw.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(sw.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(sw.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(sw.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(sw.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(sw.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(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)): + 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(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) + 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(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 + + +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( + sw.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: SimMotor, + y_axis: SimMotor, +): + with pytest.raises(ValueError): + run_engine(sw.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( + sw.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) + + +@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: SimMotor, + x_start: Number, + x_stop: Number, + x_num: int, + y_axis: SimMotor, + y_start: Number, + y_stop: Number, + y_num: int, +): + with pytest.raises(ValueError): + run_engine( + sw.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( + sw.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( + "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( + sw.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: SimMotor, + y_axis: SimMotor, +): + with pytest.raises(ValueError): + run_engine( + sw.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( + 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,)) + + +def test_list_scan_fails_with_differnt_list_lengths( + run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor +): + with pytest.raises(ValueError): + run_engine(sw.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( + 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,)) + + +def test_list_rscan_fails_with_differnt_list_lengths( + run_engine: RunEngine, x_axis: SimMotor, y_axis: SimMotor +): + with pytest.raises(ValueError): + run_engine(sw.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( + sw.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( + sw.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( + sw.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( + sw.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( + sw.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( + sw.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: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "Trajectory must contain exactly 4 values. " + "Expected (movable, start, stop, step). " + "Received 3 values: ('x_axis', 1, 5)" + ), + ): + 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( + run_engine: RunEngine, + x_axis: SimMotor, + y_axis: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "Trajectory must contain exactly 4 values. " + "Expected (movable, start, stop, step). " + "Received 3 values: ('y_axis', 1, 2)" + ), + ): + 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( + run_engine: RunEngine, + x_axis: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "Step size cannot be 0. " + "Expected (movable, start, stop, step). " + "Received (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( + run_engine: RunEngine, + x_axis: SimMotor, +): + start = stop = 0 + step = 5 + with pytest.raises( + ValueError, + match=re.escape( + "Start and stop values cannot be the same. " + "Expected (movable, start, stop, step). " + f"Received ({x_axis.name}, {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( + run_engine: RunEngine, + x_axis: SimMotor, + y_axis: SimMotor, +): + with pytest.raises( + ValueError, + match=re.escape( + "Trajectory must contain exactly 3 values. " + "Expected (movable, start, step). " + "Received 4 values: ('y_axis', 1, 5, 1)" + ), + ): + 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( + 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(sw.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(sw.step_rscan([], (x_axis, 0, 1, [0.1]))) # type: ignore diff --git a/tests/plans/test_compliance.py b/tests/plans/test_compliance.py index 2ab5d7ec65c..6a390fc4b2c 100644 --- a/tests/plans/test_compliance.py +++ b/tests/plans/test_compliance.py @@ -3,10 +3,12 @@ from types import ModuleType from typing import Any, get_type_hints +import pytest 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 @@ -45,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 - and 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): @@ -62,16 +63,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(plans): - 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) diff --git a/tests/plans/test_scanspec.py b/tests/plans/test_scanspec.py index 55d1b5e8f28..8764fc11adc 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 @@ -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 diff --git a/tests/plans/test_wrapped.py b/tests/plans/test_wrapped.py deleted file mode 100644 index 43bf8b21ec3..00000000000 --- a/tests/plans/test_wrapped.py +++ /dev/null @@ -1,1024 +0,0 @@ -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_list_scan_args, - _make_num_scan_args, - _make_step_scan_args, - _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, - step_grid_rscan, - step_grid_scan, - step_rscan, - step_scan, -) - - -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)) - print(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 run_start.get("shape") == shape - assert (hints := run_start.get("hints")) and ( - hints.get("dimensions") == [(("time",), "primary")] - ) - - -@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([])) - - -@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], - 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) - - -@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) - - -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)) - - -@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, - detectors: Sequence[StandardDetector], - 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=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: list[float | int], - y_axis: Motor, - y_list: list[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) - - -@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: list[float | int], - y_axis: Motor, - y_list: list[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) - - -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])] - ) - ) - - -@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, - 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( - num_grid_scan( - detectors=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) - - -@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) - - -@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, - detectors: Sequence[StandardDetector], - 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=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: list[float | int], - y_axis: Motor, - y_list: list[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) - - -@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: list[float | int], - y_axis: Motor, - y_list: list[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) - - -@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, - 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( - num_grid_rscan( - detectors=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, 1, 2], [3, 4, 5], False, [3], 4], [[0, 1, 2], [3, 4, 5, 6], True, [3, 4], 4]), -) -def test_make_list_scan_args( - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, - grid: bool, - final_shape: list, - final_length: int, -): - args, shape = _make_list_scan_args( - 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_list_scan_args( - 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 = int(len(x_list)) - - run_engine(list_scan(detectors=detectors, params=[(x_axis, x_list)])) - _assert_emitted(run_engine_documents, detectors, 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) - - -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, - 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, - y_axis: Motor, - y_list: list, -): - 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) - - -@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) - - -@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) - - -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, - 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) - - -@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_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) - - -@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( - 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 - - -@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, - y_axis: Motor, - y_list: list, - z_axis: Motor, - z_list: list, - 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 - ) - - -@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, - num, -): - run_engine(step_scan(detectors=detectors, params=[(x_axis, x_list)])) - _assert_emitted(run_engine_documents, detectors, 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, - y_axis: Motor, - y_list: list, - num, -): - run_engine( - step_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], - [[0, 10, 2.5], [0, -10, -2.5], 25], - ), -) -def test_step_grid_scan( - 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, -): - 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, - y_axis: Motor, - y_list: list, - num, -): - run_engine( - step_grid_scan( - detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], - snake_axes=False, - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - - -@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, - y_axis: Motor, - y_list: list, -): - 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, - num, -): - run_engine(step_rscan(detectors=detectors, params=[(x_axis, x_list)])) - _assert_emitted(run_engine_documents, detectors, 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, - y_axis: Motor, - y_list: list, - num, -): - run_engine( - step_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], - [[0, 10, 2.5], [0, -10, -2.5], 25], - ), -) -def test_step_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, -): - 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, - y_axis: Motor, - y_list: list, - num, -): - run_engine( - step_grid_rscan( - detectors=detectors, - params=[(x_axis, x_list), (y_axis, y_list)], - snake_axes=False, - ) - ) - _assert_emitted(run_engine_documents, detectors, num) - - -@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( - run_engine: RunEngine, - detectors: Sequence[StandardDetector], - x_axis: Motor, - x_list: list, - y_axis: Motor, - y_list: list, -): - with pytest.raises(ValueError): - run_engine( - step_grid_rscan( - detectors=detectors, params=[(x_axis, x_list), (y_axis, y_list)] - ) - )