nbs-bl 0.2.0__py3-none-any.whl

nbs_bl/plans/batches.py

@@ -0,0 +1,174 @@
+"""
+Tools to manage "batches" of runs.
+
+The goal is to provide tools to easily group more than one run into a larger
+unit that can be atomically retried for analysis.
+
+"""
+
+import uuid
+from functools import partial
+from itertools import count
+
+from bluesky import Msg
+from bluesky.plan_stubs import open_run, mv, trigger_and_read
+from bluesky.preprocessors import set_run_key_wrapper, subs_wrapper, msg_mutator
+from ophyd import Device, Signal, Component as Cpt
+
+# do not leak imports or helpers
+__all__ = ["setup_batch"]
+
+
+class RunMd(Device):
+    """A helper synthetic device to read per-run batch data from."""
+
+    uid = Cpt(Signal, value="", kind="hinted")
+    comment = Cpt(Signal, value="", kind="normal")
+    index = Cpt(Signal, value=0, kind="hinted")
+
+
+def setup_batch(batch_md, *, comment_function=None):
+    """
+    Set up a "batch" run.
+
+    This will create an additional run, on top of any wrapped runs that
+    includes *batch_md* flatted into the start document, a key `'purpose'` with
+    the value `"batch header"` and a key `'batch_uid'` with a generated uid..
+
+    The primary event stream of this run will include the keys: `'step_uid'`,
+    `'step_comment'` and `'step_index'` extracted from the "batched" runs.
+
+    Each wrapped run will have the key `'batch_md'` with the *batch_md* as the
+    value, `'batch_uid'` with the generated uid as the value and
+    `'batch_index'` with the running count of runs in this batch (starting from
+    0).  If the start documents already contain any of these keys the user
+    values will be respected (but do this at your own risk).
+
+    Parameters
+    ----------
+    batch_md : dict[str, Any]
+        Needs to be insertable to a start document.
+
+    comment_function : Optional[Callable[Start, str]]
+        A function to extarct a string comment from a start document.  If
+        this raises it will kill the scan.
+
+        If not specific defaults to `f"step {index}"`
+
+    Yields
+    ------
+    msg : Msg
+        To open a run for the "header" run.
+
+    Returns
+    -------
+    add_to_batch : GeneratorFunction[plan] -> Any
+        This wraps the inner plan in the batch.
+
+        What ever the inner plan returns (if anything) will be returned by the
+        wrapper.
+
+    close_batch : Callable -> None
+        Yield from this plan to close the batch (emit a stop document to
+
+        Only run this once!
+
+    Examples
+    --------
+    Typical usage::
+
+       def batch(batch_md, *, N=5, comment_function=None):
+           add_to_batch, close_batch = yield from setup_batch(
+               batch_md, comment_function=comment_function
+           )
+           for j in range(N):
+               yield from add_to_batch(inner_plan())
+           yield from close_batch()
+
+    """
+    # do not mutate the input!
+    batch_md = dict(batch_md)
+    batch_md.pop("batch_uid", None)
+    md = RunMd(name="step")
+    run_index = count()
+    batch_uid = str(uuid.uuid4())
+
+    srk_wrapper = partial(set_run_key_wrapper, run=f"batch_leader-{batch_uid}")
+
+    yield from srk_wrapper(
+        open_run(md={**batch_md, "purpose": "batch header", "batch_uid": batch_uid})
+    )
+
+    def enrich_metadata(msg):
+        if msg.command != "open_run":
+            return msg
+        # TODO maybe force these?
+        msg.kwargs.setdefault("batch_md", batch_md)
+        msg.kwargs.setdefault("batch_index", next(run_index))
+        msg.kwargs.setdefault("batch_uid", batch_uid)
+        return msg
+
+    def add_to_batch(inner):
+        """
+        Wrap a plan to be included in the batch.
+
+        This function is bound to the batch that created it via closures.
+
+        Parameters
+        ----------
+        inner : plan
+            The plan to wrap.  This may create any number of runs.
+        """
+        starts = []
+        ret = yield from subs_wrapper(
+            msg_mutator(inner, enrich_metadata),
+            {"start": [lambda name, doc: starts.append(doc)]},
+        )
+        for start in starts:
+            j = start["batch_index"]
+            comment = (
+                comment_function(start) if comment_function is not None else f"step {j}"
+            )
+            yield from mv(
+                *(md.uid, start["uid"]),
+                *(md.index, j),
+                *(md.comment, comment),
+            )
+
+            yield from srk_wrapper(trigger_and_read([md]))
+        # return what ever the wrapped plan returned
+        return ret
+
+    def close_batch(exit_status=None, reason=None):
+        """
+        Close the "header" run.
+
+        This function is bound to the batch that created it via closures.
+
+        .. warning ::
+
+            Only run this once!
+
+        Parameters
+        ----------
+        exit_status : {None, 'success', 'abort', 'fail'}
+            The exit status to report in the Stop document
+        reason : str, optional
+            Long-form description of why the run ended
+
+        Yields
+        ------
+        msg : Msg
+            Msg('close_run')
+
+        """
+        return (
+            yield Msg(
+                "close_run",
+                exit_status=exit_status,
+                reason=reason,
+                run=f"batch_leader-{batch_uid}",
+            )
+        )
+
+    return add_to_batch, close_batch

nbs_bl/plans/conditions.py

@@ -0,0 +1,77 @@
+from ..help import add_to_condition_list
+from bluesky.plan_stubs import rd
+from bluesky.protocols import Readable
+
+
+@add_to_condition_list
+def is_signal_below(sig: Readable, val: float):
+    """
+    Check if a readable object is below a threshold value.
+
+    Parameters
+    ----------
+    sig : Readable
+        Any object that implements the Readable protocol (can be read with rd()).
+    val : float
+        The threshold value to wait for.
+
+    Returns
+    -------
+    bool
+        True when signal is below threshold.
+
+    Raises
+    ------
+    """
+    reading = yield from rd(sig)
+    return reading < val
+
+
+@add_to_condition_list
+def is_signal_equals(sig: Readable, val: float):
+    """
+    Check if a readable object equals a target value.
+
+    Parameters
+    ----------
+    sig : Readable
+        Any object that implements the Readable protocol (can be read with rd()).
+    val : float
+        The target value to wait for.
+
+    Returns
+    -------
+    bool
+        True when signal equals target.
+
+    Raises
+    ------
+    """
+    reading = yield from rd(sig)
+    return reading == val
+
+
+@add_to_condition_list
+def is_signal_above(sig: Readable, val: float):
+    """
+    Check if a readable object is above a threshold value.
+
+    Parameters
+    ----------
+    sig : Readable
+        Any object that implements the Readable protocol (can be read with rd()).
+    val : float
+        The threshold value to wait for.
+
+    Returns
+    -------
+    bool
+        True when signal is above threshold.
+
+    Raises
+    ------
+    TimeoutError
+        If timeout is reached before condition is met.
+    """
+    reading = yield from rd(sig)
+    return reading > val

nbs_bl/plans/flyscan_base.py

@@ -0,0 +1,180 @@
+from bluesky.protocols import Readable, Flyable
+from bluesky.utils import get_hinted_fields
+import bluesky.preprocessors as bpp
+from bluesky.plan_stubs import trigger_and_read
+from warnings import warn
+from .plan_stubs import call_obj
+
+from bluesky.utils import Msg, ensure_generator, short_uid as _short_uid, single_gen
+from bluesky.preprocessors import plan_mutator
+from typing import Optional
+
+
+def flystream_during_wrapper(plan, flyers):
+    """
+    Kickoff and collect "flyer" (asynchronously collect) objects during runs.
+    This is a preprocessor that insert messages immediately after a run is
+    opened and before it is closed.
+    Parameters
+    ----------
+    plan : iterable or iterator
+        a generator, list, or similar containing `Msg` objects
+    flyers : collection
+        objects that support the flyer interface
+    Yields
+    ------
+    msg : Msg
+        messages from plan with 'kickoff', 'wait' and 'collect' messages
+        inserted
+    See Also
+    --------
+    :func:`bluesky.plans.fly`
+    """
+    grp1 = _short_uid("flyers-kickoff")
+    grp2 = _short_uid("flyers-complete")
+    kickoff_msgs = [Msg("kickoff", flyer, group=grp1) for flyer in flyers]
+    complete_msgs = [Msg("complete", flyer, group=grp2) for flyer in flyers]
+    collect_msgs = [Msg("collect", flyer) for flyer in flyers]
+    if flyers:
+        # If there are any flyers, insert a 'wait' Msg after kickoff, complete
+        kickoff_msgs += [Msg("wait", None, group=grp1)]
+        complete_msgs += [Msg("wait", None, group=grp2)]
+
+    def insert_after_open(msg):
+        if msg.command == "open_run":
+
+            def new_gen():
+                yield from ensure_generator(kickoff_msgs)
+
+            return single_gen(msg), new_gen()
+        else:
+            return None, None
+
+    def insert_before_close(msg):
+        if msg.command == "close_run":
+
+            def new_gen():
+                yield from ensure_generator(complete_msgs)
+                yield from ensure_generator(collect_msgs)
+                yield msg
+
+            return new_gen(), None
+        else:
+            return None, None
+
+    # Apply nested mutations.
+    plan1 = plan_mutator(plan, insert_after_open)
+    plan2 = plan_mutator(plan1, insert_before_close)
+    return (yield from plan2)
+
+
+def fly_scan(
+    detectors,
+    motor,
+    start,
+    stop,
+    *args,
+    md: Optional[dict] = None,
+    period: Optional[float] = None,
+    stream: bool = True,
+    **kwargs,
+):
+    """
+    Perform a fly scan over the specified motor range.
+
+    Parameters
+    ----------
+    detectors : list
+        List of detectors to use for the scan
+    motor : ophyd.Device
+        Motor to be scanned
+    start : float
+        Starting position of the scan
+    stop : float
+        Ending position of the scan
+    *args : float, optional
+        Additional scan parameters in groups of 3: start, stop, speed.
+        For example:
+        start1, stop1, speed1[, start2, stop2, speed2, ...]
+        This allows for multiple trajectory segments in a single scan
+    md : dict, optional
+        Metadata dictionary to be included with the scan
+    period : float, optional
+        Time period between data points. If None, uses detector's default period
+    stream : bool, optional
+        If True, continuously stream data from detectors during the scan.
+        If False, collect data only at specified points. Default is True
+
+    Returns
+    -------
+    uid : str
+        Unique identifier for the scan
+
+    Notes
+    -----
+    When stream=True, detectors will continuously collect data during motor movement,
+    providing higher time resolution but potentially more data volume.
+    When stream=False, data is collected only at specific points, reducing data volume
+    but potentially missing intermediate states.
+
+    Examples
+    --------
+    # Simple scan with one trajectory, default motor speed
+    >>> fly_scan([det], motor, 0, 10)
+
+    # Multi-segment scan with different speeds
+    >>> fly_scan([det], motor, 0, 10, 2, 10, 20, 5)
+    # This will scan from 0->10 at speed 2, then 10->20 at speed 5
+    """
+
+    md = md or {}
+
+    flyers = [d for d in detectors + [motor] if isinstance(d, Flyable)]
+    readers = [d for d in detectors + [motor] if isinstance(d, Readable)]
+
+    _md = {
+        "detectors": [det.name for det in readers],
+        "motors": [motor.name],
+        "plan_args": {
+            "detectors": list(map(repr, detectors)),
+            "args": [repr(motor), start, stop] + [a for a in args],
+        },
+        "plan_name": "fly_scan",
+        "hints": {},
+    }
+    _md.update(md or {})
+
+    x_fields = get_hinted_fields(motor)
+    default_dimensions = [(x_fields, "primary")]
+
+    default_hints = {}
+    if len(x_fields) > 0:
+        default_hints.update(dimensions=default_dimensions)
+
+    _md["hints"] = default_hints
+    _md["hints"].update(md.get("hints", {}) or {})
+
+    if period is not None:
+        for d in readers:
+            try:
+                if hasattr(d, "set_exposure"):
+                    yield from call_obj(d, "set_exposure", period)
+            except RuntimeError as ex:
+                warn(repr(ex), RuntimeWarning)
+
+    yield from call_obj(motor, "preflight", start, stop, *args, **kwargs)
+
+    @bpp.stage_decorator(readers)
+    @bpp.run_decorator(md=_md)
+    def inner_flyscan():
+        status = yield from call_obj(motor, "fly")
+
+        while not status.done:
+            yield from trigger_and_read(readers)
+
+        yield from call_obj(motor, "land")
+
+    if stream:
+        return (yield from flystream_during_wrapper(inner_flyscan(), flyers))
+    else:
+        return (yield from inner_flyscan())

nbs_bl/plans/groups.py

@@ -0,0 +1,55 @@
+import uuid
+from functools import wraps
+from bluesky.preprocessors import inject_md_wrapper
+from .preprocessors import wrap_metadata, merge_func
+
+
+def repeat(func):
+    @merge_func(func)
+    def inner(*args, repeat: int = 1, **kwargs):
+        """
+        Parameters
+        ----------
+        repeat : int, optional
+            The number of times to repeat the entire scan
+        """
+        if repeat > 1:
+            repeat_uid = str(uuid.uuid4())
+            return_list = []
+            for i in range(repeat):
+                repeat_md = {"repeat": {"uid": repeat_uid, "len": repeat, "index": i}}
+                r = yield from wrap_metadata(repeat_md)(func)(*args, **kwargs)
+                return_list.append(r)
+            return return_list
+        else:
+            return (yield from func(*args, **kwargs))
+
+    return inner
+
+
+def group(groupname):
+    def decorator(func):
+        @merge_func(func)
+        def inner(*args, **kwargs):
+            md = {"group_md": {"uid": str(uuid.uuid4()), "name": groupname}}
+            return (yield from inject_md_wrapper(func(*args, **kwargs), md))
+
+        return inner
+
+    return decorator
+
+
+def simple_1d_sequence_factory(length, label, device=None):
+    sq = {"shape": length, "label": label, "uid": str(uuid.uuid4())}
+    if device is not None:
+        sq["device_name"] = device.name
+    index = iter(range(length))
+
+    def add_to_sequence(plan, value):
+        return (
+            yield from inject_md_wrapper(
+                plan, {"sequence": {**sq, "index": next(index), "value": value}}
+            )
+        )
+
+    return add_to_sequence