nbs-bl 0.2.0__py3-none-any.whl

nbs_bl/plans/scan_decorators.py

@@ -0,0 +1,524 @@
+from functools import wraps
+from ..beamline import GLOBAL_BEAMLINE
+from ..detectors import (
+    activate_detector,
+    deactivate_detector,
+    activate_detector_set,
+)
+from ..utils import merge_func
+from .plan_stubs import (
+    set_exposure,
+    sampleholder_set_sample,
+    sampleholder_move_sample,
+    update_plan_status,
+    clear_plan_status,
+)
+from bluesky.plan_stubs import mv, trigger_and_read, declare_stream
+from bluesky.utils import separate_devices
+from bluesky.preprocessors import stage_wrapper, plan_mutator, set_run_key_wrapper
+from .preprocessors import wrap_metadata, plan_status_decorator
+from .suspenders import dynamic_suspenders
+from .groups import repeat
+
+# from ..settings import settings
+from typing import Optional
+from importlib.metadata import entry_points
+from functools import reduce
+from datetime import datetime
+from bluesky_queueserver import parameter_annotation_decorator
+
+
+def wrap_scantype(scantype):
+    def decorator(func):
+        return wrap_metadata({"scantype": scantype})(func)
+
+    return decorator
+
+
+def _beamline_setup(func):
+    blconf = GLOBAL_BEAMLINE.config.get("configuration", {})
+    if blconf.get("has_slits", False):
+        func = _slit_setup(func)
+    if blconf.get("has_motorized_samples", False):
+        func = _sample_setup_with_move(func)
+    else:
+        func = _sample_setup_no_move(func)
+    func = _eref_setup(func)
+    func = _energy_setup(func)
+    return func
+
+
+def _eref_setup(func):
+    @merge_func(func)
+    def _inner(
+        *args, eref_sample: Optional[str] = None, md: Optional[dict] = None, **kwargs
+    ):
+        """
+        Parameters
+        ----------
+        eref_sample : str, optional
+            The energy reference sample. If given, the selected reference sample is set
+        """
+        md = md or {}
+        _md = {}
+        blconf = GLOBAL_BEAMLINE.config.get("configuration", {})
+        if eref_sample is not None and blconf.get("has_motorized_eref", False):
+            yield from update_plan_status("moving energy reference sample")
+            yield from sampleholder_move_sample(
+                GLOBAL_BEAMLINE.reference_sampleholder, eref_sample
+            )
+
+            _md.update({"reference_sample": eref_sample})
+        _md.update(md)
+        return (yield from func(*args, md=_md, **kwargs))
+
+    return _inner
+
+
+def _sample_setup_with_move(func):
+    @merge_func(func)
+    def _inner(
+        *args,
+        sample: Optional[str] = None,
+        sample_position: Optional[dict] = {},
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        sample : str, optional
+            The sample id. If given, the selected sample metadata is set
+        sample_position: dict, optional
+            A dictionary of positions relative to the sample center. Parameters not given will be assumed to
+            be the default for the sampleholder (typically moving the sample into the beam at a typical angle)
+        """
+        if sample is not None:
+            yield from update_plan_status("moving_sample")
+            yield from sampleholder_move_sample(
+                GLOBAL_BEAMLINE.primary_sampleholder, sample, **sample_position
+            )
+        return (yield from func(*args, **kwargs))
+
+    return _inner
+
+
+def _sample_setup_no_move(func):
+    @merge_func(func)
+    def _inner(*args, sample=None, **kwargs):
+        """
+        Parameters
+        ----------
+        sample : str, optional
+            The sample id. If given, the selected sample metadata is set, but the sample is not moved
+        """
+        if sample is not None:
+            yield from sampleholder_set_sample(
+                GLOBAL_BEAMLINE.primary_sampleholder, sample
+            )
+        return (yield from func(*args, **kwargs))
+
+    return _inner
+
+
+def _slit_setup(func):
+    @merge_func(func)
+    def _inner(*args, eslit: Optional[float] = None, **kwargs):
+        """
+        Parameters
+        ----------
+        eslit : float, optional
+            If not None, will set the beamline exit slit prior to the plan start.
+        """
+        if eslit is not None:
+            yield from update_plan_status("setting exit slit")
+            yield from mv(GLOBAL_BEAMLINE.slits, eslit)
+        return (yield from func(*args, **kwargs))
+
+    return _inner
+
+
+def _energy_setup(func):
+    @merge_func(func)
+    def _inner(
+        *args,
+        energy: Optional[float] = None,
+        polarization: Optional[float] = None,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        energy : float, optional
+            If not None, will set the beamline energy prior to the plan start.
+        """
+        print("Energy Setup Decorator")
+
+        if energy is not None:
+            print("Setting up energy")
+            yield from update_plan_status("setting up energy")
+            yield from mv(GLOBAL_BEAMLINE.energy, energy)
+        if polarization is not None and hasattr(GLOBAL_BEAMLINE, "polarization"):
+            yield from update_plan_status("setting up polarization")
+            yield from mv(GLOBAL_BEAMLINE.polarization, polarization)
+        return (yield from func(*args, **kwargs))
+
+    return _inner
+
+
+def staged_baseline_wrapper(plan, devices, name="staged_baseline"):
+    """
+    Preprocessor that records a baseline of all `devices` after `open_run`
+
+    The readings are designated for a separate event stream named 'baseline' by
+    default.
+
+    Parameters
+    ----------
+    plan : iterable or iterator
+        a generator, list, or similar containing `Msg` objects
+    devices : collection
+        collection of Devices to read
+        If None, the plan passes through unchanged.
+    name : string, optional
+        name for event stream; by default, 'baseline'
+
+    Yields
+    ------
+    msg : Msg
+        messages from plan, with 'set' messages inserted
+    """
+
+    def head():
+        yield from declare_stream(*devices, name=name)
+        yield from trigger_and_read(devices, name=name)
+
+    def tail():
+        yield from trigger_and_read(devices, name=name)
+
+    def insert_baseline(msg):
+        if msg.command == "open_run":
+            if msg.run is not None:
+                return None, set_run_key_wrapper(head(), msg.run)
+            else:
+                return None, head()
+
+        elif msg.command == "close_run":
+            if msg.run is not None:
+                return set_run_key_wrapper(tail(), msg.run), None
+            else:
+                return tail(), None
+
+        return None, None
+
+    if not devices:
+        # no-op
+        return (yield from plan)
+    else:
+        return (yield from plan_mutator(plan, insert_baseline))
+
+
+def _nbs_setup_detectors_with_baseline(func):
+    @parameter_annotation_decorator({"parameters": {"dwell": {"min": 0}}})
+    @merge_func(func, ["detectors"])
+    def _inner(*args, extra_dets=[], dwell: Optional[float] = None, **kwargs):
+        """
+        Parameters
+        ----------
+        extra_dets : list, optional
+            A list of extra detectors to be activated for the scan, by default [].
+        dwell : float, optional
+            The exposure time in seconds for all detectors, by default None.
+            If None, do not set any exposure time, and assume that detectors are already set
+        """
+
+        # for det in extra_dets:
+        #    activate_detector(det)
+        print("Detector Setup Decorator with Baseline")
+        if dwell is not None:
+            yield from set_exposure(dwell, extra_dets=extra_dets)
+
+        all_dets = separate_devices(GLOBAL_BEAMLINE.detectors.active + extra_dets)
+        # We need to stage only the baseline devices that are not already in the list of detectors
+        # In order to avoid double staging
+        if hasattr(GLOBAL_BEAMLINE, "staged_baseline"):
+            staged_baseline_devices = separate_devices(
+                GLOBAL_BEAMLINE.staged_baseline.values()
+            )
+            devices_to_stage = [
+                dev for dev in staged_baseline_devices if dev not in all_dets
+            ]
+        else:
+            staged_baseline_devices = []
+            devices_to_stage = []
+
+        def inner():
+            yield from func(all_dets, *args, **kwargs)
+
+        ret = yield from stage_wrapper(
+            staged_baseline_wrapper(inner(), staged_baseline_devices), devices_to_stage
+        )
+
+        # for det in extra_dets:
+        #    deactivate_detector(det)
+
+        return ret
+
+    return _inner
+
+
+def _nbs_setup_detectors(func):
+    @parameter_annotation_decorator({"parameters": {"dwell": {"min": 0}}})
+    @merge_func(func, ["detectors"])
+    def _inner(*args, extra_dets=[], dwell: Optional[float] = None, **kwargs):
+        """
+        Parameters
+        ----------
+        extra_dets : list, optional
+            A list of extra detectors to be activated for the scan, by default [].
+        dwell : float, optional
+            The exposure time in seconds for all detectors, by default None.
+            If None, do not set any exposure time, and assume that detectors are already set
+        """
+
+        # for det in extra_dets:
+        #    activate_detector(det)
+        print("Detector Setup Decorator")
+        if dwell is not None:
+            yield from set_exposure(dwell, extra_dets=extra_dets)
+        all_dets = GLOBAL_BEAMLINE.detectors.active + extra_dets
+        ret = yield from func(all_dets, *args, **kwargs)
+
+        # for det in extra_dets:
+        #    deactivate_detector(det)
+
+        return ret
+
+    return _inner
+
+
+def _nbs_add_plot_md(func):
+    @merge_func(func)
+    def _inner(*args, md: Optional[dict] = None, plot_detectors: list = None, **kwargs):
+        md = md or {}
+        plot_hints = {}
+        if plot_detectors is not None:
+            activate_detector_set(plot_detectors)
+        plot_hints = GLOBAL_BEAMLINE.detectors.get_plot_hints()
+        _md = {"plot_hints": plot_hints}
+        _md.update(md)
+        return (yield from func(*args, md=_md, **kwargs))
+
+    return _inner
+
+
+def _nbs_add_sample_md(func):
+    @merge_func(func)
+    def _inner(*args, md: Optional[dict] = None, **kwargs):
+        """
+        Sample information is automatically added to the run md
+        """
+        md = md or {}
+        if hasattr(GLOBAL_BEAMLINE, "current_sample"):
+            _md = {
+                "sample_name": GLOBAL_BEAMLINE.current_sample.get("name", ""),
+                "sample_id": GLOBAL_BEAMLINE.current_sample.get("sample_id", ""),
+                "sample_desc": GLOBAL_BEAMLINE.current_sample.get("description", ""),
+                "sample_set": GLOBAL_BEAMLINE.current_sample.get("group", ""),
+                "sample_info": {
+                    k: v
+                    for k, v in GLOBAL_BEAMLINE.current_sample.items()
+                    if k not in ["name", "sample_id", "description", "group"]
+                },
+            }
+            _md.update(md)
+            return (yield from func(*args, md=_md, **kwargs))
+        else:
+            return (yield from func(*args, md=md, **kwargs))
+
+    return _inner
+
+
+def _nbs_add_comment(func):
+    @merge_func(func)
+    def _inner(
+        *args,
+        md: Optional[dict] = None,
+        comment: Optional[str] = None,
+        group_name: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Parameters
+        ----------
+        comment : str, optional
+            A comment that will be added into the run metadata. If not provided, no comment will be added.
+        group_name : str, optional
+            A group name label that will be added into the run metadata.
+        """
+        md = md or {}
+        if comment is not None:
+            _md = {"comment": comment}
+        else:
+            _md = {}
+        if group_name is not None:
+            _md["group_name"] = group_name
+        _md.update(md)
+        return (yield from func(*args, md=_md, **kwargs))
+
+    return _inner
+
+
+def nbs_base_scan_decorator(func):
+    @repeat
+    @plan_status_decorator
+    @_nbs_add_plan_args
+    @_beamline_setup
+    @_nbs_setup_detectors
+    @_nbs_add_sample_md
+    @_nbs_add_plot_md
+    @_nbs_add_comment
+    @merge_func(func)
+    def _inner(*args, **kwargs):
+        return (yield from func(*args, **kwargs))
+
+    return _inner
+
+
+def _nbs_add_plan_args(func):
+    @merge_func(func)
+    def _inner(*args, md: Optional[dict] = None, **kwargs):
+        md = md or {}
+        _md = {}
+        _md.update(md)
+        _md.update(
+            {
+                "plan_passed_name": func.__name__,
+                # Don't have a way to normalize args and kwargs, copy.deepcopy is dangerous
+                # "plan_passed_args": args,
+                # "plan_passed_kwargs": kwargs,
+                "time_human": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            }
+        )
+        return (yield from func(*args, md=_md, **kwargs))
+
+    return _inner
+
+
+"""
+    This was too clever, and was a mistake
+def nbs_add_bl_prefix(func):
+    base_name = func.__name__
+    plan_name = f"{settings.beamline_prefix}_{base_name}"
+    func.__name__ = plan_name
+    return func
+"""
+
+
+def wrap_plan_name(func):
+    return wrap_metadata({"plan_name": func.__name__})(func)
+
+
+def nbs_builtin_scan_wrapper(func):
+    """
+    Designed to wrap bluesky built-in scans to produce an nbs version
+    """
+
+    @wrap_plan_name
+    @nbs_base_scan_decorator
+    @merge_func(func)
+    def _inner(*args, **kwargs):
+        print("Running ", func.__name__)
+        return (yield from func(*args, **kwargs))
+
+    # _inner = wrap_metadata({"plan_name": plan_name})(nbs_base_scan_decorator(func))
+
+    d = f"""Modifies {func.__name__} to automatically fill
+dets with global active beamline detectors.
+Other detectors may be added on the fly via extra_dets
+---------------------------------------------------------
+"""
+
+    _inner.__doc__ = d + _inner.__doc__
+    return _inner
+
+
+def get_decorator_entrypoints():
+    """
+    Get decorator entrypoints from beamline configuration.
+
+    Returns
+    -------
+    list
+        List of decorator functions loaded from entrypoints
+    """
+    config = GLOBAL_BEAMLINE.config
+    decorator_entrypoints = config.get("settings", {}).get("plan_decorators", [])
+
+    decorators = []
+    if decorator_entrypoints:
+        eps = entry_points()
+        for ep_name in decorator_entrypoints:
+            try:
+                # Look for entrypoint in nbs_bl.plan_decorators group
+                matches = eps.select(group="nbs_bl.plan_decorators", name=ep_name)
+                for match in matches:
+                    decorator = match.load()
+                    decorators.append(decorator)
+            except Exception as e:
+                print(f"Failed to load decorator {ep_name}: {e}")
+
+    return decorators
+
+
+def dynamic_scan_wrapper(func, func_name=None):
+    """
+    A flexible wrapper for Bluesky scans that incorporates base NBS functionality
+    and additional decorators from entrypoints.
+
+    Parameters
+    ----------
+    func : callable
+        The scan function to wrap
+
+    Returns
+    -------
+    callable
+        Wrapped scan function with all decorators applied
+    """
+    # Get base decorators
+    base_decorators = [
+        merge_func(func, use_func_name=False),
+        dynamic_suspenders,
+        nbs_base_scan_decorator,
+    ]
+
+    # Get additional decorators from entrypoints
+    additional_decorators = get_decorator_entrypoints()
+
+    # Combine all decorators
+    all_decorators = base_decorators + additional_decorators + [wrap_plan_name]
+
+    func_name = func_name or func.__name__
+
+    # Apply decorators in order
+    def _inner(*args, **kwargs):
+        print("Running ", func_name)
+        return (yield from func(*args, **kwargs))
+
+    if func_name is not None:
+        _inner.__name__ = func_name
+    # Apply all decorators in sequence
+    wrapped = reduce(lambda f, dec: dec(f), all_decorators, _inner)
+
+    # Add documentation
+    d = f"""Wraps {func.__name__} with additional arguments and default detectors.
+
+Other detectors may be added on the fly via extra_dets
+
+
+Additional decorators loaded:
+{[d.__name__ for d in additional_decorators]}
+"""
+    wrapped.__doc__ = d + (wrapped.__doc__ or "")
+
+    return wrapped

nbs_bl/plans/scans.py

@@ -0,0 +1,145 @@
+from .scan_decorators import dynamic_scan_wrapper
+from .scan_base import _make_gscan_points
+from .flyscan_base import fly_scan
+from ..help import add_to_scan_list, add_to_plan_time_dict
+from ..utils import merge_func
+from ..beamline import GLOBAL_BEAMLINE as bl
+from .plan_stubs import update_plan_status
+
+import bluesky.plans as bp
+from bluesky.plan_stubs import mv
+
+_scan_list = [
+    bp.count,
+    bp.scan,
+    bp.rel_scan,
+    bp.list_scan,
+    bp.rel_list_scan,
+    bp.list_grid_scan,
+    bp.rel_list_grid_scan,
+    bp.log_scan,
+    bp.rel_log_scan,
+    bp.grid_scan,
+    bp.rel_grid_scan,
+    bp.scan_nd,
+    bp.spiral,
+    bp.spiral_fermat,
+    bp.spiral_square,
+    bp.rel_spiral,
+    bp.rel_spiral_fermat,
+    bp.rel_spiral_square,
+    fly_scan,
+]
+
+for _scan in _scan_list:
+    _fixedname = f"nbs_{_scan.__name__}"
+    _newscan = dynamic_scan_wrapper(_scan, func_name=_fixedname)
+    # _newscan.__name__ = _fixedname
+    globals()[_fixedname] = _newscan
+    add_to_scan_list(_newscan)
+
+add_to_plan_time_dict(
+    nbs_count, "generic_estimate", fixed=0, overhead=0.05, dwell="dwell", points="num"
+)
+for _scan in [
+    bp.scan,
+    nbs_scan,
+    bp.rel_scan,
+    nbs_rel_scan,
+    bp.log_scan,
+    nbs_log_scan,
+    bp.rel_log_scan,
+    nbs_rel_log_scan,
+]:
+    add_to_plan_time_dict(
+        _scan, "generic_estimate", fixed=5, overhead=0.5, dwell="dwell", points="num"
+    )
+for _scan in [bp.list_scan, nbs_list_scan, bp.rel_list_scan, nbs_rel_list_scan]:
+    add_to_plan_time_dict(
+        _scan, "list_scan_estimate", fixed=5, overhead=0.5, dwell="dwell"
+    )
+for _scan in [bp.grid_scan, nbs_grid_scan, bp.rel_grid_scan, nbs_rel_grid_scan]:
+    add_to_plan_time_dict(
+        _scan, "grid_scan_estimate", fixed=5, overhead=0.5, dwell="dwell"
+    )
+
+for _scan in [
+    bp.list_grid_scan,
+    nbs_list_grid_scan,
+    bp.rel_list_grid_scan,
+    nbs_rel_list_grid_scan,
+]:
+    add_to_plan_time_dict(
+        _scan, "list_grid_scan_estimate", fixed=5, overhead=0.5, dwell="dwell"
+    )
+
+
+@add_to_scan_list
+@merge_func(
+    nbs_list_scan,
+    omit_params=["points"],
+    exclude_wrapper_args=False,
+    use_func_name=False,
+)
+def nbs_gscan(motor, start: float, *args, extra_dets=[], shift: float = 0, **kwargs):
+    """A variable step scan of a motor and the default detectors.
+
+    Other detectors may be added via extra_dets
+
+    Parameters
+    ----------
+    motor : Ophyd Device
+        The motor object to scan
+    start : float
+        Starting position for the scan
+    step : float, optional
+        Step size for first region
+    stop : float, optional
+        First stopping position
+    step2 : float, optional
+        Step size for second region. Additional stop/step pairs can be provided
+    stop2 : float, optional
+        Second stopping position. Additional stop/step pairs can be provided
+    *args : float, optional
+        Additional stop/step pairs can be provided
+    extra_dets : list
+        A list of detectors to add for just this scan
+    shift : float
+        A value to shift all start/stop positions by
+
+    Examples
+    --------
+    # Scan from 270 to 280 in steps of 1
+    nbs_gscan(energy, 270, 1, 280)
+
+    # Scan from 270 to 280 in steps of 1, then from 280 to 290 in steps of 0.25
+    nbs_gscan(energy, 270, 1, 280, 0.25, 290)
+
+    # Scan with multiple regions of different step sizes
+    nbs_gscan(energy, 270, 1, 280, 0.5, 285, 0.1, 290)
+    """
+    points = _make_gscan_points(start, *args, shift=shift)
+    # Move motor to start position first
+    yield from update_plan_status(
+        f"moving motor {motor.name} to start position {points[0]}"
+    )
+    yield from mv(motor, points[0])
+    return (yield from nbs_list_scan(motor, points, extra_dets=extra_dets, **kwargs))
+
+
+@add_to_scan_list
+@merge_func(nbs_gscan, use_func_name=False, omit_params=["motor"])
+def nbs_energy_scan(*args, **kwargs):
+    """A variable step scan of the energy and the default detectors.
+
+    The energy motor is automatically included.Other detectors may be added via extra_dets.
+    """
+    motor = bl.energy
+    return (yield from nbs_gscan(motor, *args, **kwargs))
+
+
+for _scan in [nbs_gscan, nbs_energy_scan]:
+    add_to_plan_time_dict(_scan, "gscan_estimate", fixed=5, overhead=0.5, dwell="dwell")
+
+for _scan in [fly_scan, nbs_fly_scan]:
+    add_to_plan_time_dict(_scan, "fly_scan_estimate", fixed=5)