nbs-bl 0.2.0__py3-none-any.whl

nbs_bl/plans/plan_stubs.py

@@ -0,0 +1,246 @@
+from ..beamline import GLOBAL_BEAMLINE
+from ..help import add_to_plan_list, add_to_condition_list
+import warnings
+from bluesky import Msg
+from bluesky.plan_stubs import rd, sleep, mv
+import time
+from typing import Optional
+from bluesky.protocols import Readable
+import typing
+from .conditions import is_signal_below, is_signal_equals, is_signal_above
+
+GLOBAL_EXPOSURE_TIME = 1.0
+
+
+def update_plan_status(status):
+    ret = yield Msg("update_plan_status", None, status)
+    return ret
+
+
+def clear_plan_status():
+    ret = yield Msg("clear_plan_status")
+    return ret
+
+
+def call_obj(obj, method, *args, **kwargs):
+    ret = yield Msg("call_obj", obj, *args, method=method, **kwargs)
+    return ret
+
+
+def sampleholder_move_sample_old(sampleholder, sample_id=None, **position):
+    """
+    Set and move a sample. Dangerous! Sample moves really need to go through
+    Bluesky's mv plan, or else they won't be properly pause-able
+    """
+    yield from call_obj(sampleholder, "move_sample", sample_id, **position)
+
+
+def sampleholder_move_sample(sampleholder, sample_id=None, **position):
+    """
+    Set and move a sample.
+    """
+    positions = yield from call_obj(
+        sampleholder, "get_sample_position", sample_id, **position
+    )
+    print(f"Moving {sampleholder.name} to {positions}")
+    yield from mv(sampleholder, positions)
+    print("Done Moving ", sampleholder.name)
+
+
+def sampleholder_set_sample(sampleholder, sample_id):
+    """
+    Set a sample without moving it
+    """
+    yield from call_obj(sampleholder, "set_sample", sample_id)
+
+
+@add_to_plan_list
+def set_exposure(time: Optional[float] = None, extra_dets=[]):
+    """Sets the exposure time for all active detectors"""
+    global GLOBAL_EXPOSURE_TIME
+    if time is not None:
+        GLOBAL_EXPOSURE_TIME = time
+
+    all_dets = GLOBAL_BEAMLINE.detectors.active + extra_dets
+    for d in all_dets:
+        try:
+            if hasattr(d, "set_exposure"):
+                yield from call_obj(d, "set_exposure", GLOBAL_EXPOSURE_TIME)
+        except RuntimeError as ex:
+            warnings.warn(repr(ex), RuntimeWarning)
+
+
+@add_to_plan_list
+def set_roi(label, llim, ulim):
+    for d in GLOBAL_BEAMLINE.detectors.active:
+        try:
+            if hasattr(d, "set_roi"):
+                yield from call_obj(d, "set_roi", label, llim, ulim)
+        except RuntimeError as ex:
+            warnings.warn(repr(ex), RuntimeWarning)
+
+
+@add_to_plan_list
+def clear_all_rois():
+    for d in GLOBAL_BEAMLINE.detectors.active:
+        try:
+            if hasattr(d, "clear_all_rois"):
+                yield from call_obj(d, "clear_all_rois")
+        except RuntimeError as ex:
+            warnings.warn(repr(ex), RuntimeWarning)
+
+
+@add_to_plan_list
+def clear_one_roi(label):
+    for d in GLOBAL_BEAMLINE.detectors.active:
+        try:
+            if hasattr(d, "clear_roi"):
+                yield from call_obj(d, "clear_roi", label)
+        except RuntimeError as ex:
+            warnings.warn(repr(ex), RuntimeWarning)
+
+
+@add_to_plan_list
+def wait_for_condition(
+    condition,
+    condition_args: typing.List = None,
+    condition_kwargs: typing.Dict = None,
+    timeout: Optional[float] = None,
+    sleep_time: float = 10,
+):
+    """
+    Wait for a condition function to return True.
+
+    Parameters
+    ----------
+    condition : str or callable
+        A condition function that returns a boolean.
+    condition_args : list, optional
+        Arguments for the condition function.
+    condition_kwargs : dict, optional
+        Keyword arguments for the condition function.
+    timeout : Optional[float], optional
+        Maximum time to wait in seconds. If None, wait indefinitely.
+    sleep_time : float, optional
+        Time to sleep between checks in seconds. Default is 10.
+
+    Returns
+    -------
+    bool
+        True when condition returns True.
+
+    Raises
+    ------
+    TimeoutError
+        If timeout is reached before condition is met.
+    """
+    if condition_args is None:
+        condition_args = []
+    if condition_kwargs is None:
+        condition_kwargs = {}
+
+    start_time = time.time()
+    while True:
+        if timeout is not None and (time.time() - start_time > timeout):
+            raise TimeoutError
+        result = yield from condition(*condition_args, **condition_kwargs)
+        if result:
+            return True
+        else:
+            yield from sleep(sleep_time)
+
+
+def wait_for_signal_below(
+    sig: Readable, val: float, timeout: Optional[float] = None, sleep_time: float = 10
+):
+    """
+    Wait for a readable object to go below a threshold value.
+
+    Parameters
+    ----------
+    sig : Readable
+        Any object that implements the Readable protocol (can be read with rd()).
+    val : float
+        The target value to wait for.
+    timeout : Optional[float], optional
+        Maximum time to wait in seconds. If None, wait indefinitely.
+    sleep_time : float, optional
+        Time to sleep between checks in seconds. Default is 10.
+
+    Returns
+    -------
+    bool
+        True when signal goes below threshold.
+
+    Raises
+    ------
+    TimeoutError
+        If timeout is reached before condition is met.
+    """
+    return wait_for_condition(
+        is_signal_below, [sig, val], timeout=timeout, sleep_time=sleep_time
+    )
+
+
+def wait_for_signal_equals(
+    sig: Readable, val: float, timeout: Optional[float] = None, sleep_time: float = 10
+):
+    """
+    Wait for a readable object to equal 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.
+    timeout : Optional[float], optional
+        Maximum time to wait in seconds. If None, wait indefinitely.
+    sleep_time : float, optional
+        Time to sleep between checks in seconds. Default is 10.
+
+    Returns
+    -------
+    bool
+        True when signal equals target value.
+
+    Raises
+    ------
+    TimeoutError
+        If timeout is reached before condition is met.
+    """
+    return wait_for_condition(
+        is_signal_equals, [sig, val], timeout=timeout, sleep_time=sleep_time
+    )
+
+
+def wait_for_signal_above(
+    sig: Readable, val: float, timeout: Optional[float] = None, sleep_time: float = 10
+):
+    """
+    Wait for a readable object to go above a threshold value.
+
+    Parameters
+    ----------
+    sig : Readable
+        Any object that implements the Readable protocol (can be read with rd()).
+    val : float
+        The target value to wait for.
+    timeout : Optional[float], optional
+        Maximum time to wait in seconds. If None, wait indefinitely.
+    sleep_time : float, optional
+        Time to sleep between checks in seconds. Default is 10.
+
+    Returns
+    -------
+    bool
+        True when signal goes above threshold.
+
+    Raises
+    ------
+    TimeoutError
+        If timeout is reached before condition is met.
+    """
+    return wait_for_condition(
+        is_signal_above, [sig, val], timeout=timeout, sleep_time=sleep_time
+    )

nbs_bl/plans/preprocessors.py

@@ -0,0 +1,160 @@
+from bluesky.plan_stubs import open_run, close_run
+from bluesky.utils import RunEngineControlException, make_decorator
+from bluesky.preprocessors import contingency_wrapper
+from bluesky.preprocessors import plan_mutator, ensure_generator, single_gen
+from bluesky import Msg
+from functools import wraps
+from ..utils import merge_func
+from copy import deepcopy
+import inspect
+import uuid
+from typing import Optional
+
+
+def sanitize(arg):
+    # storable_types = (int, float, str, bool)
+    if type(arg) == list:
+        _arg = []
+        for thing in arg:
+            try:
+                _arg.append(thing.name)
+            except Exception:
+                _arg.append(thing)
+        return _arg
+    else:
+        if hasattr(arg, "name"):
+            return arg.name
+        else:
+            return arg
+
+
+def plan_md_decorator(plan_function):
+    @wraps(plan_function)
+    def _inner(*args, md: Optional[dict] = None, plan_level=0, **kwargs):
+        plan_name = plan_function.__name__
+        md = md or {}
+        _md = {}
+        arguments = {}
+        s = inspect.signature(plan_function)
+        for p in s.parameters.values():
+            if p.default is not s.empty:
+                arguments[p.name] = p.default
+
+        for a, n in zip(args, plan_function.__code__.co_varnames):
+            arguments[n] = sanitize(a)
+        for k, v in kwargs.items():
+            arguments[k] = sanitize(v)
+        if "md" in arguments:
+            del arguments["md"]
+
+        _md["plan_history"] = []
+        if plan_level == 0:
+            _md["master_plan"] = plan_name
+            _md["batch_uid"] = str(uuid.uuid4())
+        _md.update(deepcopy(md))
+        _md["plan_history"].append(
+            {"plan_name": plan_name, "arguments": arguments, "plan_level": plan_level}
+        )
+        return plan_function(*args, md=_md, plan_level=plan_level + 1, **kwargs)
+
+    return _inner
+
+
+def wrap_metadata(param):
+    def decorator(func):
+        @merge_func(func)
+        def inner(*args, md: Optional[dict] = None, **kwargs):
+            md = md or {}
+            _md = {}
+            _md.update(param)
+            _md.update(md)
+            return (yield from func(*args, md=_md, **kwargs))
+
+        return inner
+
+    return decorator
+
+
+def run_return_wrapper(plan, *, md: Optional[dict] = None):
+    """Enclose in 'open_run' and 'close_run' messages.
+
+    Parameters
+    ----------
+    plan : iterable or iterator
+        a generator, list, or similar containing `Msg` objects
+    md : dict, optional
+        metadata to be passed into the 'open_run' message
+    """
+    yield from open_run(md)
+
+    def except_plan(e):
+        if isinstance(e, RunEngineControlException):
+            yield from close_run(exit_status=e.exit_status)
+        else:
+            yield from close_run(exit_status="fail", reason=str(e))
+
+    return (
+        yield from contingency_wrapper(
+            plan, except_plan=except_plan, else_plan=close_run
+        )
+    )
+
+
+run_return_decorator = make_decorator(run_return_wrapper)
+
+
+def plan_status_decorator(func):
+    """
+    Update plan status, and clear after plan completion
+
+    Parameters
+    ----------
+    plan : iterable or iterator
+        a generator, list, or similar containing `Msg` objects
+    status : str
+        the status to update to
+
+    Yields
+    ------
+    msg : Msg
+        messages from plan with 'update_plan_status', and 'clear_plan_status' messages inserted
+
+    """
+
+    @merge_func(func)
+    def _inner(*args, md: Optional[dict] = None, **kwargs):
+        # Get plan name from metadata if available
+        md = md or {}
+        plan_name = md.get("plan_name", func.__name__) if md else func.__name__
+        status = "Running " + plan_name
+        print(f"Plan status decorator: {status}")
+        update_msgs = [Msg("update_plan_status", None, status)]
+        clear_msgs = [Msg("clear_plan_status")]
+
+        def insert_after_open(msg):
+            if msg.command == "open_run":
+
+                def new_gen():
+                    yield from ensure_generator(update_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(clear_msgs)
+                    yield msg
+
+                return new_gen(), None
+            else:
+                return None, None
+
+        # Apply nested mutations.
+        plan1 = plan_mutator(func(*args, md=md, **kwargs), insert_after_open)
+        plan2 = plan_mutator(plan1, insert_before_close)
+        return (yield from plan2)
+
+    return _inner

nbs_bl/plans/scan_base.py

@@ -0,0 +1,58 @@
+def _make_gscan_points(*args, shift: float = 0):
+    """
+    Generate a sequence of energy scan points from a variable-length parameter list.
+
+    Parameters should be passed in the following format:
+        (estart1, delta1, estop1, delta2, estop2, ...)
+
+    - Each segment defines a range starting from the given estart and ending at the estop,
+      using the specified delta (step size).
+    - Delta values must be positive. The function automatically determines the direction
+      (increasing or decreasing) based on the estart and estop values.
+    - A single energy value can be passed to return a single-point list.
+    - Segments can be non-monotonic — i.e., you can mix increasing and decreasing ranges.
+    - An optional `shift` parameter can be used to apply a constant offset to all values.
+
+    Examples:
+        _make_gscan_points(250)
+            ➜ [250.0]
+
+        _make_gscan_points(250, 5, 260)
+            ➜ [250.0, 255.0, 260.0]
+
+        _make_gscan_points(264, 2, 260, 5, 250)
+            ➜ [264.0, 262.0, 260.0, 255.0, 250.0]
+
+    Raises:
+        TypeError if the number of arguments is incorrect
+        ValueError if any delta is zero
+    """
+    
+    if len(args) == 1:
+        return [float(args[0]) + shift]
+    if (len(args) - 1) % 2 != 0:
+        raise TypeError(
+            "gscan received an even number of arguments. Either a stop or a step-size is missing.  Expected format: (estart1, delta1, estop1, delta2, estop2, ...)"
+        )
+    points = []
+    for i in range(1, len(args) - 1, 2):
+        estart = float(args[i - 1]) + shift
+        delta = abs(args[i]) ## Ensures delta is positive in case users input negative delta for reverse energy list.
+        estop = float(args[i + 1]) + shift
+
+        if delta == 0:
+            raise ValueError("Step size (delta) cannot be zero.")
+
+        step = delta if estop > estart else -delta
+
+        if not points or points[-1] != estart:
+            points.append(estart)
+
+        next_point = estart + step
+        while (step > 0 and next_point < estop - step / 2.0) or (step < 0 and next_point > estop - step / 2.0):
+            points.append(next_point)
+            next_point += step
+
+        points.append(estop)
+
+    return points