diff --git a/docs/reference/api.md b/docs/reference/api.md index 0c92e9e065..a4b8292144 100644 --- a/docs/reference/api.md +++ b/docs/reference/api.md @@ -9,6 +9,7 @@ dodal.common dodal.devices dodal.plans + dodal.plan_stubs .. automodule:: dodal diff --git a/pyproject.toml b/pyproject.toml index 18155c1580..816285631d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ "aiofiles", "aiohttp", "redis", + "scanspec>=0.7.3", ] dynamic = ["version"] diff --git a/src/dodal/plan_stubs/__init__.py b/src/dodal/plan_stubs/__init__.py new file mode 100644 index 0000000000..6db5b3b480 --- /dev/null +++ b/src/dodal/plan_stubs/__init__.py @@ -0,0 +1,26 @@ +from .check_topup import check_topup_and_wait_if_necessary, wait_for_topup_complete +from .data_session import ( + DATA_SESSION, + attach_data_session_metadata_decorator, + attach_data_session_metadata_wrapper, +) +from .motor_utils import ( + MoveTooLarge, + check_and_cache_values, + home_and_reset_decorator, + home_and_reset_wrapper, + move_and_reset_wrapper, +) + +__all__ = [ + "DATA_SESSION", + "attach_data_session_metadata_wrapper", + "check_and_cache_values", + "check_topup_and_wait_if_necessary", + "wait_for_topup_complete", + "attach_data_session_metadata_decorator", + "MoveTooLarge", + "home_and_reset_decorator", + "home_and_reset_wrapper", + "move_and_reset_wrapper", +] diff --git a/src/dodal/plans/check_topup.py b/src/dodal/plan_stubs/check_topup.py similarity index 100% rename from src/dodal/plans/check_topup.py rename to src/dodal/plan_stubs/check_topup.py diff --git a/src/dodal/plans/data_session_metadata.py b/src/dodal/plan_stubs/data_session.py similarity index 100% rename from src/dodal/plans/data_session_metadata.py rename to src/dodal/plan_stubs/data_session.py diff --git a/src/dodal/plans/motor_util_plans.py b/src/dodal/plan_stubs/motor_utils.py similarity index 98% rename from src/dodal/plans/motor_util_plans.py rename to src/dodal/plan_stubs/motor_utils.py index 80d0e93611..5033f1c7fd 100644 --- a/src/dodal/plans/motor_util_plans.py +++ b/src/dodal/plan_stubs/motor_utils.py @@ -23,7 +23,7 @@ def __init__( super().__init__(*args) -def _check_and_cache_values( +def check_and_cache_values( devices_and_positions: dict[AnyDevice, float], smallest_move: float, maximum_move: float, @@ -87,7 +87,7 @@ def move_and_reset_wrapper( on. If false it is left up to the caller to wait on them. Defaults to True. """ - initial_positions = yield from _check_and_cache_values( + initial_positions = yield from check_and_cache_values( device_and_positions, smallest_move, maximum_move ) diff --git a/src/dodal/plans/__init__.py b/src/dodal/plans/__init__.py new file mode 100644 index 0000000000..fb40245969 --- /dev/null +++ b/src/dodal/plans/__init__.py @@ -0,0 +1,4 @@ +from .scanspec import spec_scan +from .wrapped import count + +__all__ = ["count", "spec_scan"] diff --git a/src/dodal/plans/scanspec.py b/src/dodal/plans/scanspec.py new file mode 100644 index 0000000000..adb5fea97f --- /dev/null +++ b/src/dodal/plans/scanspec.py @@ -0,0 +1,64 @@ +import operator +from collections.abc import Mapping +from functools import reduce +from typing import Annotated, Any + +import bluesky.plans as bp +from bluesky.protocols import Movable, Readable +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 import attach_data_session_metadata_decorator + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def spec_scan( + detectors: Annotated[ + set[Readable], + Field( + description="Set of readable devices, will take a reading at each point, \ + in addition to any Movables in the Spec", + ), + ], + spec: Annotated[ + Spec[Movable], + Field(description="ScanSpec modelling the path of the scan"), + ], + metadata: Mapping[str, Any] | None = None, +) -> MsgGenerator: + """Generic plan for reading `detectors` at every point of a ScanSpec `spec`.""" + _md = { + "plan_args": { + "detectors": {det.name for det in detectors}, + "spec": repr(spec), + }, + "plan_name": "spec_scan", + "shape": spec.shape(), + **(metadata or {}), + } + + yield from bp.scan_nd(detectors, _as_cycler(spec), md=_md) + + +def _as_cycler( + spec: Spec[Movable], # type: ignore +) -> Cycler: + """ + Convert a scanspec to a cycler for compatibility with legacy Bluesky plans such as + `bp.scan_nd`. Use the midpoints of the scanspec since cyclers are normally used + for software triggered scans. + + Args: + spec: A scanspec + + Returns: + Cycler: A new cycler + """ + + midpoints = spec.frames().midpoints + # Need to "add" the cyclers for all the axes together. The code below is + # effectively: cycler(motor1, [...]) + cycler(motor2, [...]) + ... + return reduce(operator.add, (cycler(*args) for args in midpoints.items())) diff --git a/src/dodal/plans/wrapped.py b/src/dodal/plans/wrapped.py new file mode 100644 index 0000000000..61c0ecec3d --- /dev/null +++ b/src/dodal/plans/wrapped.py @@ -0,0 +1,38 @@ +from collections.abc import Mapping +from typing import Annotated, Any + +import bluesky.plans as bp +from bluesky.protocols import Readable +from pydantic import Field, PositiveFloat, validate_call + +from dodal.common import MsgGenerator +from dodal.plan_stubs import attach_data_session_metadata_decorator + + +@attach_data_session_metadata_decorator() +@validate_call(config={"arbitrary_types_allowed": True}) +def count( + detectors: Annotated[ + set[Readable], + 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[ + PositiveFloat | list[PositiveFloat] | None, + Field( + description="Delay between readings: if list, 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"}, + ), + ] = None, + metadata: Mapping[str, Any] | None = None, +) -> MsgGenerator: + if isinstance(delay, list): + assert ( + delays := len(delay) + ) == num - 1, f"Number of delays given must be {num - 1}: was given {delays} " + yield from bp.count(detectors, num, delay=delay, md=metadata or {}) diff --git a/tests/plans/test_motor_util_plans.py b/tests/plan_stubs/test_motor_util_plans.py similarity index 96% rename from tests/plans/test_motor_util_plans.py rename to tests/plan_stubs/test_motor_util_plans.py index dc4a761ca2..b30cc9caf6 100644 --- a/tests/plans/test_motor_util_plans.py +++ b/tests/plan_stubs/test_motor_util_plans.py @@ -14,9 +14,9 @@ from ophyd_async.epics.motor import Motor from dodal.devices.util.test_utils import patch_motor -from dodal.plans.motor_util_plans import ( +from dodal.plan_stubs import ( MoveTooLarge, - _check_and_cache_values, + check_and_cache_values, home_and_reset_wrapper, ) @@ -59,7 +59,7 @@ def my_device(RE): "device_type", [DeviceWithOnlyMotors, DeviceWithNoMotors, DeviceWithSomeMotors], ) -@patch("dodal.plans.motor_util_plans.move_and_reset_wrapper") +@patch("dodal.plan_stubs.motor_utils.move_and_reset_wrapper") def test_given_types_of_device_when_home_and_reset_wrapper_called_then_motors_and_zeros_passed_to_move_and_reset_wrapper( patch_move_and_reset, device_type, RE ): @@ -80,7 +80,7 @@ def test_given_a_device_when_check_and_cache_values_then_motor_values_returned( set_mock_value(motor.user_readback, i * 100) motors_and_positions: dict[Motor, float] = RE( - _check_and_cache_values( + check_and_cache_values( {motor_obj: 0.0 for motor_obj in my_device.motors}, 0, 1000 ) ).plan_result # type: ignore @@ -109,7 +109,7 @@ def test_given_a_device_with_a_too_large_move_when_check_and_cache_values_then_e motors_and_positions = {motor_obj: new_position for motor_obj in my_device.motors} with pytest.raises(MoveTooLarge) as e: - RE(_check_and_cache_values(motors_and_positions, 0, max)) + RE(check_and_cache_values(motors_and_positions, 0, max)) assert e.value.axis == my_device.y assert e.value.maximum_move == max @@ -136,7 +136,7 @@ def test_given_a_device_where_one_move_too_small_when_check_and_cache_values_the } motors_and_positions: dict[Motor, float] = RE( - _check_and_cache_values(motors_and_new_positions, min, 1000) + check_and_cache_values(motors_and_new_positions, min, 1000) ).plan_result # type: ignore cached_positions = motors_and_positions.values() @@ -156,7 +156,7 @@ def test_given_a_device_where_all_moves_too_small_when_check_and_cache_values_th motors_and_new_positions = {motor_obj: 0.0 for motor_obj in my_device.motors} motors_and_positions: dict[Motor, float] = RE( - _check_and_cache_values(motors_and_new_positions, 40, 1000) + check_and_cache_values(motors_and_new_positions, 40, 1000) ).plan_result # type: ignore cached_positions = motors_and_positions.values() diff --git a/tests/plans/test_topup_plan.py b/tests/plan_stubs/test_topup_plan.py similarity index 91% rename from tests/plans/test_topup_plan.py rename to tests/plan_stubs/test_topup_plan.py index 2b9fcce512..658b448ebe 100644 --- a/tests/plans/test_topup_plan.py +++ b/tests/plan_stubs/test_topup_plan.py @@ -7,7 +7,7 @@ from dodal.beamlines import i03 from dodal.devices.synchrotron import Synchrotron, SynchrotronMode -from dodal.plans.check_topup import ( +from dodal.plan_stubs import ( check_topup_and_wait_if_necessary, wait_for_topup_complete, ) @@ -18,8 +18,8 @@ def synchrotron(RE) -> Synchrotron: return i03.synchrotron(fake_with_ophyd_sim=True) -@patch("dodal.plans.check_topup.wait_for_topup_complete") -@patch("dodal.plans.check_topup.bps.sleep") +@patch("dodal.plan_stubs.wait_for_topup_complete") +@patch("bluesky.plan_stubs.sleep") def test_when_topup_before_end_of_collection_wait( fake_sleep: MagicMock, fake_wait: MagicMock, synchrotron: Synchrotron, RE: RunEngine ): @@ -37,8 +37,8 @@ def test_when_topup_before_end_of_collection_wait( fake_sleep.assert_called_once_with(61.0) -@patch("dodal.plans.check_topup.bps.rd") -@patch("dodal.plans.check_topup.bps.sleep") +@patch("bluesky.plan_stubs.rd") +@patch("bluesky.plan_stubs.sleep") def test_wait_for_topup_complete( fake_sleep: MagicMock, fake_rd: MagicMock, synchrotron: Synchrotron, RE: RunEngine ): @@ -59,8 +59,8 @@ def fake_generator(value): fake_sleep.assert_called_with(0.1) -@patch("dodal.plans.check_topup.bps.sleep") -@patch("dodal.plans.check_topup.bps.null") +@patch("bluesky.plan_stubs.sleep") +@patch("bluesky.plan_stubs.null") def test_no_waiting_if_decay_mode( fake_null: MagicMock, fake_sleep: MagicMock, synchrotron: Synchrotron, RE: RunEngine ): @@ -77,7 +77,7 @@ def test_no_waiting_if_decay_mode( assert fake_sleep.call_count == 0 -@patch("dodal.plans.check_topup.bps.null") +@patch("bluesky.plan_stubs.null") def test_no_waiting_when_mode_does_not_allow_gating( fake_null: MagicMock, synchrotron: Synchrotron, RE: RunEngine ): @@ -120,7 +120,7 @@ def test_no_waiting_when_mode_does_not_allow_gating( (29, 39, 35, 1, 0, "topup_long_delay.txt"), ], ) -@patch("dodal.plans.check_topup.bps.sleep") +@patch("bluesky.plan_stubs.sleep") def test_topup_not_allowed_when_exceeds_threshold_percentage_of_topup_time( mock_sleep, RE: RunEngine, diff --git a/tests/preprocessors/test_filesystem_metadata.py b/tests/preprocessors/test_filesystem_metadata.py index 6904f3bb1c..5d4694d647 100644 --- a/tests/preprocessors/test_filesystem_metadata.py +++ b/tests/preprocessors/test_filesystem_metadata.py @@ -25,7 +25,7 @@ LocalDirectoryServiceClient, StaticVisitPathProvider, ) -from dodal.plans.data_session_metadata import ( +from dodal.plan_stubs import ( DATA_SESSION, attach_data_session_metadata_wrapper, )