Skip to content

Commit

Permalink
Move plans from deprecated module
Browse files Browse the repository at this point in the history
  • Loading branch information
DiamondJoseph committed Oct 24, 2024
1 parent 5778b5f commit 32dd6e1
Show file tree
Hide file tree
Showing 8 changed files with 576 additions and 0 deletions.
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ dependencies = [
"aiohttp",
"redis",
"deepdiff",
"scanspec>=0.7.3",
]

dynamic = ["version"]
Expand All @@ -47,6 +48,7 @@ dev = [
# Commented out due to dependency version conflict with pydantic 1.x
# "copier",
"myst-parser",
"ophyd_async[sim]",
"pipdeptree",
"pre-commit",
"psutil",
Expand Down
151 changes: 151 additions & 0 deletions src/dodal/plan_stubs/wrapped.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
import itertools
from collections.abc import Mapping
from typing import Annotated, Any

import bluesky.plan_stubs as bps
from bluesky.protocols import Movable
from bluesky.utils import MsgGenerator

"""
Wrappers for Bluesky built-in plan stubs with type hinting
"""

Group = Annotated[str, "String identifier used by 'wait' or stubs that await"]


# https://github.com/bluesky/bluesky/issues/1821
def set_absolute(
movable: Movable, value: Any, group: Group | None = None, wait: bool = False
) -> MsgGenerator:
"""
Set a device, wrapper for `bp.abs_set`.
Args:
movable (Movable): The device to set
value (T): The new value
group (Group | None, optional): The message group to associate with the
setting, for sequencing. Defaults to None.
wait (bool, optional): The group should wait until all setting is complete
(e.g. a motor has finished moving). Defaults to False.
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""
return (yield from bps.abs_set(movable, value, group=group, wait=wait))


# https://github.com/bluesky/bluesky/issues/1821
def set_relative(
movable: Movable, value: Any, group: Group | None = None, wait: bool = False
) -> MsgGenerator:
"""
Change a device, wrapper for `bp.rel_set`.
Args:
movable (Movable): The device to set
value (T): The new value
group (Group | None, optional): The message group to associate with the
setting, for sequencing. Defaults to None.
wait (bool, optional): The group should wait until all setting is complete
(e.g. a motor has finished moving). Defaults to False.
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""

return (yield from bps.rel_set(movable, value, group=group, wait=wait))


# https://github.com/bluesky/bluesky/issues/1821
def move(moves: Mapping[Movable, Any], group: Group | None = None) -> MsgGenerator:
"""
Move a device, wrapper for `bp.mv`.
Args:
moves (Mapping[Movable, T]): Mapping of Movables to target positions
group (Group | None, optional): The message group to associate with the
setting, for sequencing. Defaults to None.
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""

return (
# https://github.com/bluesky/bluesky/issues/1809
yield from bps.mv(*itertools.chain.from_iterable(moves.items()), group=group) # type: ignore
)


# https://github.com/bluesky/bluesky/issues/1821
def move_relative(
moves: Mapping[Movable, Any], group: Group | None = None
) -> MsgGenerator:
"""
Move a device relative to its current position, wrapper for `bp.mvr`.
Args:
moves (Mapping[Movable, T]): Mapping of Movables to target deltas
group (Group | None, optional): The message group to associate with the
setting, for sequencing. Defaults to None.
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""

return (
# https://github.com/bluesky/bluesky/issues/1809
yield from bps.mvr(*itertools.chain.from_iterable(moves.items()), group=group) # type: ignore
)


def sleep(time: float) -> MsgGenerator:
"""
Suspend all action for a given time, wrapper for `bp.sleep`
Args:
time (float): Time to wait in seconds
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""

return (yield from bps.sleep(time))


def wait(
group: Group | None = None,
timeout: float | None = None,
) -> MsgGenerator:
"""
Wait for a group status to complete, wrapper for `bp.wait`
Args:
group (Group | None, optional): The name of the group to wait for, defaults
to None, in which case waits for all
groups that have not yet been awaited.
timeout (float | None, default=None): a timeout in seconds
Returns:
MsgGenerator: Plan
Yields:
Iterator[MsgGenerator]: Bluesky messages
"""

return (yield from bps.wait(group, timeout=timeout))
4 changes: 4 additions & 0 deletions src/dodal/plans/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
from .scanspec import spec_scan
from .wrapped import count

__all__ = ["count", "spec_scan"]
65 changes: 65 additions & 0 deletions src/dodal/plans/scanspec.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import operator
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.data_session 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: dict[str, Any] | None = None,
) -> MsgGenerator:
"""Generic plan for reading `detectors` at every point of a ScanSpec `Spec`.
A `Spec` is an N-dimensional path.
"""
_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(tuple(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()))
42 changes: 42 additions & 0 deletions src/dodal/plans/wrapped.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
from typing import Annotated, Any

import bluesky.plans as bp
from bluesky.protocols import Readable
from pydantic import Field, NonNegativeFloat, validate_call

from dodal.common import MsgGenerator
from dodal.plan_stubs.data_session 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[
NonNegativeFloat | list[NonNegativeFloat],
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"},
),
] = 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, list):
assert (
delays := len(delay)
) == num - 1, f"Number of delays given must be {num - 1}: was given {delays}"
metadata = metadata or {}
metadata["shape"] = (num,)
yield from bp.count(tuple(detectors), num, delay=delay, md=metadata)
Loading

0 comments on commit 32dd6e1

Please sign in to comment.