Move util functions from i22 branch of ophyd_async that may be generally useful

pyproject.toml

@@ -15,6 +15,7 @@ description = "Common Diamond specific Bluesky plans and functions"
 dependencies = [
  "blueapi",
  "ophyd",
+ "ophyd_async @ git+https://github.com/bluesky/ophyd-async.git@add-detector-logic",
  "scanspec"
 ] # Add project dependencies here, e.g. ["click", "numpy"]
 dynamic = ["version"]

src/dls_bluesky_core/core/__init__.py

@@ -1,6 +1,14 @@
-from .types import MsgGenerator, PlanGenerator
+from .coordination import group_uuid
+from .maths import in_micros, step_to_num
+from .scanspecs import get_duration
+from .types import MsgGenerator, PlanGenerator, ScanAxis
 
 __all__ = [
+ "get_duration",
+ "group_uuid",
+ "in_micros",
  "MsgGenerator",
  "PlanGenerator",
+ "ScanAxis",
+ "step_to_num",
 ]

src/dls_bluesky_core/core/coordination.py

@@ -0,0 +1,14 @@
+import uuid
+
+
+def group_uuid(name: str) -> str:
+ """
+ Returns a unique but human-readable string, to assist debugging orchestrated groups.
+
+ Args:
+ name (str): A human readable name
+
+ Returns:
+ readable_uid (str): name appended with a unique string
+ """
+ return f"{name}-{str(uuid.uuid4())[:6]}"

src/dls_bluesky_core/core/maths.py

@@ -0,0 +1,47 @@
+from typing import Tuple
+
+import numpy as np
+
+
+def step_to_num(start: float, stop: float, step: float) -> Tuple[float, float, int]:
+ """
+ Standard handling for converting from start, stop, step to start, stop, num
+ Forces step to be same direction as length
+ Includes a final point if it is within 1% of the end point (consistent with GDA)
+
+ Args:
+ start (float):
+ Start of length, will be returned unchanged
+ stop (float):
+ End of length, if length/step does not divide cleanly will be returned
+ extended up to 1% of step, or else truncated.
+ step (float):
+ Length of a step along the line formed from start to stop.
+ If stop < start, will be coerced to be backwards.
+
+ Returns:
+ start, truncated_stop, num = Tuple[float, float, int]
+ start will be returned unchanged
+ truncated_stop = start + num * step
+ num is the maximal number of steps that could fit into the length.
+
+ """
+ # Make step be the right direction
+ step = abs(step) if stop >= start else -abs(step)
+ # If stop is within 1% of a step then include it
+ steps = int((stop - start) / step + 0.01)
+ return start, start + steps * step, steps + 1 # include 1st point
+
+
+def in_micros(t: float) -> int:
+ """
+ Converts between units of T and units of microT.
+ For example, from seconds to microseconds.
+
+ Args:
+ t (float): A time in seconds, or other measurement in units of U
+ Returns:
+ t (int): A time in microseconds, rounded up to the nearest whole microsecond,
+ or other measurement in units of microU, rounded up to the nearest whole microU.
+ """
+ return int(np.ceil(t * 1e6))

src/dls_bluesky_core/core/scanspecs.py

@@ -0,0 +1,32 @@
+from typing import List, Optional
+
+import numpy as np
+from scanspec.core import Frames
+from scanspec.specs import DURATION
+
+
+def get_duration(frames: List[Frames]) -> Optional[float]:
+ """
+ Returns the duration of a number of ScanSpec frames, if known and consistent.
+
+ Args:
+ frames (List[Frames]): A number of Frame objects
+
+ Raises:
+ ValueError: If any frames do not have a Duration defined.
+
+ Returns:
+ duration (float): if all frames have a consistent duration
+ None: otherwise
+
+ """
+ for fs in frames:
+ if DURATION in fs.axes():
+ durations = fs.midpoints[DURATION]
+ first_duration = durations[0]
+ if np.all(durations == first_duration):
+ # Constant duration, return it
+ return first_duration
+ else:
+ return None
+ raise ValueError("Duration not specified in Spec")

src/dls_bluesky_core/core/types.py

@@ -1,8 +1,11 @@
-from typing import Any, Callable, Generator
+from typing import Any, Callable, Generator, Union
 
 from bluesky import Msg
+from ophyd_async.core import Device
+from scanspec.specs import DURATION
 
 # 'A true "plan", usually the output of a generator function'
 MsgGenerator = Generator[Msg, Any, None]
 # 'A function that generates a plan'
 PlanGenerator = Callable[..., MsgGenerator]
+ScanAxis = Union[Device, DURATION]

src/dls_bluesky_core/plans/scanspec.py

@@ -1,13 +1,18 @@
 import operator
 from functools import reduce
-from typing import Any, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional, Sequence
 
+import bluesky.plan_stubs as bps
 import bluesky.plans as bp
+import bluesky.preprocessors as bpp
 from bluesky.protocols import Movable, Readable
 from cycler import Cycler, cycler
-from scanspec.specs import Spec
+from ophyd_async.core import Device
+from ophyd_async.core.flyer import ScanSpecFlyable
+from scanspec.specs import Repeat, Spec
 
-from dls_bluesky_core.core import MsgGenerator
+from ..core import MsgGenerator
+from ..stubs.flyables import fly_and_collect
 
 """
 Plans related to the use of the `ScanSpec https://github.com/dls-controls/scanspec`
@@ -57,6 +62,78 @@ def scan(
  yield from bp.scan_nd(detectors, cycler, md=_md)
 
 
+def scanspec_fly(
+ flyer: ScanSpecFlyable,
+ hw_spec: Spec[ScanSpecFlyable],
+ sw_dets: Optional[Sequence[Device]] = None,
+ sw_spec: Optional[Spec[Device]] = None,
+ # TODO: How to allow in Blueapi REST call?
+ # TODO: Allow defining processors for @BeforeScan/@AfterScan equivalent logic?
+ # setup_detectors: Optional[Collection[Msg]] = None,
+ # flush_period: float = 0.5, # TODO: Should we allow overriding default?
+ metadata: Optional[Dict[str, Any]] = None,
+) -> MsgGenerator:
+ """
+ TODO
+
+ Args:
+ flyer (ScanSpecFlyable):
+ A Flyable that traces the path of hw_spec Spec at every point of sw_spec.
+ hw_spec (Spec[Device]):
+ The 'inner scan' performed at each point of the outer scan.
+ sw_dets (Optional[Sequence[Device]]):
+ Any detectors to be triggered at every point of an outer scan.
+ sw_spec (Optional[Spec[Device]]):
+ The 'outer scan' trajectory to be followed by any non-flying motors.
+ Defaults to a one-shot no-op Spec.
+ # setup_detectors (Iterator[Msg]):
+ Any Msgs required to set up the detectors # TODO: Make a pre-processor?
+ # flush_period (float): # TODO: Allow non-default?
+ Timeout for calls to complete when the flyer is kicked off.
+ metadata (Optional[Dict[str, Any]]):
+ Key-value metadata to include in exported data, defaults to None.
+
+ Returns:
+ MsgGenerator: _description_
+
+ Yields:
+ Iterator[MsgGenerator]: _description_
+ """
+ sw_dets = sw_dets or []
+ sw_spec = sw_spec or Repeat()
+ detectors: Sequence[Device] = [flyer] + sw_dets
+ plan_args = {
+ "flyer": repr(flyer),
+ "hw_spec": repr(hw_spec),
+ "sw_dets": [repr(det) for det in sw_dets],
+ "sw_spec": repr(sw_spec),
+ }
+
+ _md = {
+ "plan_args": plan_args,
+ "detectors": [det.name for det in detectors],
+ "hints": {}, # TODO: Dimension hints from ScanSpec?
+ }
+ _md.update(metadata or {})
+
+ @bpp.stage_decorator(detectors)
+ @bpp.run_decorator(md=_md)
+ def hw_scanspec_fly() -> MsgGenerator:
+ # yield from setup_detectors
+ yield from bps.declare_stream(*sw_dets, name="sw")
+ yield from bps.declare_stream(flyer, name="hw")
+ for point in sw_spec.midpoints():
+ # Move flyer to start too
+ point[flyer] = hw_spec
+ # TODO: need to make pos_cache optional in this func
+ yield from bps.move_per_step(point, dict())
+ yield from bps.trigger_and_read(sw_dets)
+ yield from bps.checkpoint()
+ yield from fly_and_collect(flyer, checkpoint_every_collect=True)
+
+ return (yield from hw_scanspec_fly())
+
+
 def _scanspec_to_cycler(spec: Spec[str], axes: Mapping[str, Movable]) -> Cycler:
  """
  Convert a scanspec to a cycler for compatibility with legacy Bluesky plans such as

src/dls_bluesky_core/stubs/flyables.py

@@ -0,0 +1,23 @@
+import bluesky.plan_stubs as bps
+from bluesky.protocols import Flyable
+
+from dls_bluesky_core.core import MsgGenerator, group_uuid
+
+
+def fly_and_collect(
+ flyer: Flyable, flush_period: float = 0.5, checkpoint_every_collect: bool = False
+) -> MsgGenerator:
+ yield from bps.kickoff(flyer)
+ complete_group = group_uuid("complete")
+ yield from bps.complete(flyer, group=complete_group)
+ done = False
+ while not done:
+ try:
+ yield from bps.wait(group=complete_group, timeout=flush_period)
+ except TimeoutError:
+ pass
+ else:
+ done = True
+ yield from bps.collect(flyer, stream=True, return_payload=False)
+ if checkpoint_every_collect:
+ yield from bps.checkpoint()

tests/core/test_funcs.py

@@ -0,0 +1,68 @@
+import uuid
+from typing import Optional
+
+import pytest
+
+from dls_bluesky_core.core import group_uuid, in_micros, step_to_num
+
+
+@pytest.mark.parametrize(
+ "s,us",
+ [
+ (4.000_001, 4_000_001),
+ (4.999_999, 4_999_999),
+ (4, 4_000_000),
+ (-4.000_001, -4_000_001),
+ (-4.999_999, -4_999_999),
+ (-4, -4_000_000),
+ (4.000_000_1, 4_000_001),
+ (4.999_999_9, 5_000_000),
+ (0.1, 100_000),
+ (0.000_000_1, 1),
+ (-4.000_000_5, -4_000_000),
+ (-4.999_999_9, -4_999_999),
+ (-4.05, -4_050_000),
+ ],
+)
+def test_in_micros(s: float, us: int):
+ assert in_micros(s) == us
+
+
+@pytest.mark.parametrize(
+ "start,stop,step,expected_num,truncated_stop",
+ [
+ (0, 0, 1, 1, None), # start=stop, 1 point at start
+ (0, 0.5, 1, 1, 0), # step>length, 1 point at start
+ (0, 1, 1, 2, None), # stop=start+step, point at start & stop
+ (0, 0.99, 1, 2, 1), # stop >= start + 0.99*step, included
+ (0, 0.98, 1, 1, 0), # stop < start + 0.99*step, not included
+ (0, 1.01, 1, 2, 1), # stop >= start + 0.99*step, included
+ (0, 1.75, 0.25, 8, 1.75),
+ (0, 0, -1, 1, None), # start=stop, 1 point at start
+ (0, 0.5, -1, 1, 0), # abs(step)>length, 1 point at start
+ (0, -1, 1, 2, None), # stop=start+-abs(step), point at start & stop
+ (0, -0.99, 1, 2, -1), # stop >= start + 0.99*-abs(step), included
+ (0, -0.98, 1, 1, 0), # stop < start + 0.99*-abs(step), not included
+ (0, -1.01, 1, 2, -1), # stop >= start + 0.99*-abs(step), included
+ (0, -1.75, 0.25, 8, -1.75),
+ ],
+)
+def test_step_to_num(
+ start: float,
+ stop: float,
+ step: float,
+ expected_num: int,
+ truncated_stop: Optional[float],
+):
+ truncated_stop = stop if truncated_stop is None else truncated_stop
+ actual_start, actual_stop, num = step_to_num(start, stop, step)
+ assert actual_start == start
+ assert actual_stop == truncated_stop
+ assert num == expected_num
+
+
+@pytest.mark.parametrize("group", ["foo", "bar", "baz", str(uuid.uuid4())])
+def test_group_uid(group: str):
+ gid = group_uuid(group)
+ assert gid.startswith(f"{group}-")
+ assert not gid.endswith(f"{group}-")

tests/plans/test_compliance.py

@@ -9,7 +9,7 @@
 
 def is_bluesky_plan_generator(func: Any) -> bool:
  try:
- return get_type_hints(func).get("return") is MsgGenerator
+ return get_type_hints(func).get("return") == MsgGenerator
  except TypeError:
  # get_type_hints fails on some objects (such as Union or Optional)
  return False
@@ -45,7 +45,7 @@ def assert_metadata_requirements(plan: PlanGenerator, signature: inspect.Signatu
  ), f"'{plan.__name__}' does not allow metadata"
  metadata = signature.parameters["metadata"]
  assert (
- metadata.annotation is Optional[Mapping[str, Any]]
+ metadata.annotation == Optional[Mapping[str, Any]]
  and metadata.default is not inspect.Parameter.empty
  ), f"'{plan.__name__}' metadata is not optional"
  assert metadata.default is None, f"'{plan.__name__}' metadata default is mutable"