nbs-bl 0.2.0__py3-none-any.whl

nbs_bl/plans/suspenders.py

@@ -0,0 +1,87 @@
+from ..utils import merge_func
+from importlib.metadata import entry_points
+from bluesky.preprocessors import finalize_wrapper
+from bluesky.suspenders import SuspenderBase
+from bluesky import Msg
+from ..beamline import GLOBAL_BEAMLINE
+
+
+def get_suspender_entrypoints():
+    """
+    Get suspender entrypoints from beamline configuration.
+
+    Returns
+    -------
+    list
+        List of decorator functions loaded from entrypoints
+    """
+    config = GLOBAL_BEAMLINE.config
+    suspender_entrypoints = config.get("settings", {}).get("suspenders", [])
+    print(f"Suspender entrypoints: {suspender_entrypoints}")
+    suspenders = []
+    if suspender_entrypoints:
+        eps = entry_points()
+
+        for ep_name in suspender_entrypoints:
+            try:
+                # Look for entrypoint in nbs_bl.suspenders group
+                matches = eps.select(group="nbs_bl.suspenders", name=ep_name)
+                for match in matches:
+                    print(f"Loading suspender {match.name}")
+                    suspender = match.load()
+
+                    # Handle both single suspender and list of suspenders
+                    if isinstance(suspender, (list, tuple)):
+                        suspenders.extend(suspender)
+                    else:
+                        suspenders.append(suspender)
+            except Exception as e:
+                print(f"Failed to load suspender {ep_name}: {e}")
+
+    return suspenders
+
+
+def dynamic_suspenders(func):
+    # Load suspenders from entrypoints at decoration time
+    suspenders = get_suspender_entrypoints()
+
+    @merge_func(func)
+    def wrapper(*args, skip_suspenders=False, **kwargs):
+        """
+        Parameters
+        ----------
+        skip_suspenders : bool, optional
+            If True, do not install suspenders
+        """
+        # Check if suspenders should be skipped
+
+        if skip_suspenders or not suspenders:
+            return (yield from func(*args, **kwargs))
+        else:
+            suspender_list = []
+            for sus in suspenders:
+                if callable(sus) and not isinstance(sus, SuspenderBase):
+                    sus = sus()
+                if isinstance(sus, (list, tuple)):
+                    suspender_list.extend(sus)
+                else:
+                    suspender_list.append(sus)
+
+            def _install():
+                for susp in suspender_list:
+                    yield Msg("install_suspender", None, susp)
+
+            def _remove():
+                for susp in suspender_list:
+                    yield Msg("remove_suspender", None, susp)
+
+            @merge_func(func)
+            def _inner_plan(*args, **kwargs):
+                yield from _install()
+                return (yield from func(*args, **kwargs))
+
+            return (
+                yield from finalize_wrapper(_inner_plan(*args, **kwargs), _remove())
+            )
+
+    return wrapper

nbs_bl/plans/time_estimation.py

@@ -0,0 +1,168 @@
+import numpy as np
+from .scan_base import _make_gscan_points
+
+"""
+Time estimation functions for plan execution time calculation.
+
+Usage Examples:
+---------------
+1. Basic time estimation:
+   estimation_dict = {
+       "estimator": "generic_estimate",
+       "fixed": 10,
+       "overhead": 0.5,
+       "dwell": "dwell_time"
+   }
+   plan_args = {"dwell_time": 2.0, "points": 100}
+   time = generic_estimate("my_plan", plan_args, estimation_dict)
+
+2. With repeat functionality:
+   estimation_dict = {
+       "estimator": "generic_estimate", 
+       "fixed": 10,
+       "overhead": 0.5,
+       "dwell": "dwell_time",
+       "reset": 5.0  # 5 seconds between repeats
+   }
+   plan_args = {"dwell_time": 2.0, "points": 100, "repeat": 3}
+   time = generic_estimate("my_plan", plan_args, estimation_dict)
+   # Result: (base_time * 3) + (5.0 * 2) = base_time * 3 + 10.0
+"""
+
+
+def with_repeat(estimator_func):
+    """
+    Decorator to add repeat functionality to time estimators.
+
+    This decorator handles the "repeat" parameter in plan_args and "reset" parameter
+    in estimation_dict. It calculates the base plan time, multiplies by repeats,
+    and adds reset time for each repeat except the last one.
+
+    Parameters
+    ----------
+    estimator_func : callable
+        The original time estimator function
+
+    Returns
+    -------
+    callable
+        Wrapped estimator function that handles repeats
+    """
+
+    def wrapper(plan_name, plan_args, estimation_dict):
+        # Get repeat count and remove from plan_args to avoid confusion
+        repeat_count = plan_args.get("repeat", 1)
+
+        # Get reset time from estimation parameters
+        reset_time = estimation_dict.get("reset_time", 0)
+
+        # Calculate base plan time using original estimator
+        base_time = estimator_func(plan_name, plan_args, estimation_dict)
+
+        if base_time is None:
+            return None
+
+        # Calculate total time: base_time * repeats + reset_time * (repeats - 1)
+        total_time = base_time * repeat_count + reset_time * (repeat_count - 1)
+
+        return total_time
+
+    return wrapper
+
+
+def get_dwell(plan_args, estimation_dict):
+    dwell = estimation_dict.get("dwell", "dwell")
+    if isinstance(dwell, str):
+        return plan_args.get(dwell, 0)
+    else:
+        return dwell
+
+
+@with_repeat
+def generic_estimate(plan_name, plan_args, estimation_dict):
+    a = estimation_dict.get("fixed", 0)
+    b = estimation_dict.get("overhead", 0)
+    c = get_dwell(plan_args, estimation_dict)
+    if "points" in estimation_dict:
+        points = estimation_dict.get("points")
+        if isinstance(points, str):
+            if points == "args":
+                n = len(plan_args.get("args", []))
+            else:
+                n = plan_args.get(points, 0)
+        else:
+            n = points
+    else:
+        n = 0
+    return a + b * n + c * n
+
+
+@with_repeat
+def list_scan_estimate(plan_name, plan_args, estimation_dict):
+    a = estimation_dict.get("fixed", 0)
+    b = estimation_dict.get("overhead", 0)
+    c = get_dwell(plan_args, estimation_dict)
+
+    n = len(plan_args.get("args", ["", []])[1])
+    return a + b * n + c * n
+
+
+@with_repeat
+def grid_scan_estimate(plan_name, plan_args, estimation_dict):
+    a = estimation_dict.get("fixed", 0)
+    b = estimation_dict.get("overhead", 0)
+    c = get_dwell(plan_args, estimation_dict)
+
+    args = plan_args.get("args")
+    n_axes = len(args) // 4
+    n_points = 1
+    for i in range(n_axes):
+        n_points *= args[i * 4 + 3]
+    return a + b * n_points + c * n_points
+
+
+@with_repeat
+def list_grid_scan_estimate(plan_name, plan_args, estimation_dict):
+    a = estimation_dict.get("fixed", 0)
+    b = estimation_dict.get("overhead", 0)
+    c = get_dwell(plan_args, estimation_dict)
+
+    args = plan_args.get("args")
+    n_axes = len(args) // 2
+    n_points = 1
+    for i in range(n_axes):
+        n_points *= len(args[i * 2 + 1])
+    return a + b * n_points + c * n_points
+
+
+@with_repeat
+def fly_scan_estimate(plan_name, plan_args, estimation_dict):
+    if "args" in plan_args:
+        args = plan_args["args"]
+        if len(args) % 2 == 0:
+            args = args[1:]
+        start = args[0]
+        stop = args[1]
+        speed = args[2]
+    else:
+        start = plan_args.get("start", 0)
+        stop = plan_args.get("stop", 0)
+        speed = plan_args.get("speed", 1)
+    a = estimation_dict.get("fixed", 0)
+    return a + np.abs((stop - start) / speed)
+
+
+@with_repeat
+def gscan_estimate(plan_name, plan_args, estimation_dict):
+    if "region" in estimation_dict:
+        region = estimation_dict.get("region")
+    else:
+        region = plan_args.get("args")
+    if len(region) % 2 == 0:
+        region = region[1:]  # First argument is motor name
+    points = len(_make_gscan_points(*region))
+    a = estimation_dict.get("fixed", 0)
+    b = estimation_dict.get("overhead", 0)
+    c = get_dwell(plan_args, estimation_dict)
+
+    return a + b * points + c * points

nbs_bl/plans/xas.py

@@ -0,0 +1,123 @@
+# from ..settings import settings
+from ..plans.scans import nbs_gscan
+from ..utils import merge_func
+from ..plans.preprocessors import wrap_metadata
+from .plan_stubs import set_roi, clear_one_roi
+from ..help import _add_to_import_list, add_to_func_list, add_to_plan_time_dict
+from ..queueserver import GLOBAL_USER_STATUS
+from ..status import StatusDict
+from ..beamline import GLOBAL_BEAMLINE
+from .scan_base import _make_gscan_points
+
+# from ..settings import GLOBAL_SETTINGS as settings
+from os.path import join
+
+try:
+    import tomllib
+except ModuleNotFoundError:
+    import tomli as tomllib
+
+GLOBAL_XAS_PLANS = GLOBAL_USER_STATUS.request_status_dict("XAS_PLANS", use_redis=True)
+GLOBAL_XAS_PLANS.clear()
+
+
+def add_to_xas_list(f, key, **plan_info):
+    """
+    A function decorator that will add the plan to the built-in list
+    """
+    _add_to_import_list(f, "xas")
+    GLOBAL_XAS_PLANS[key] = {}
+    GLOBAL_XAS_PLANS[key].update(plan_info)
+    return f
+
+
+def _wrap_xas(element, edge):
+    def decorator(func):
+        return wrap_metadata({"element": element, "edge": edge, "scantype": "xas"})(
+            func
+        )
+
+    return decorator
+
+
+# Needs to have Element, Edge, Ref Element,
+def _xas_factory(energy_grid, element, edge, key):
+    @_wrap_xas(element, edge)
+    @wrap_metadata({"plan_name": key})
+    @merge_func(nbs_gscan, omit_params=["motor", "start", "stop", "step", "args"])
+    def inner(**kwargs):
+        """Parameters
+        ----------
+        repeat : int
+            Number of times to repeat the scan
+        **kwargs :
+            Arguments to be passed to tes_gscan
+
+        """
+        eref_sample = kwargs.pop("eref_sample", None)
+        if eref_sample is None:
+            eref_sample = element
+        yield from set_roi("pfy", energy_grid[0], energy_grid[-2])
+        yield from nbs_gscan(
+            GLOBAL_BEAMLINE.energy, *energy_grid, eref_sample=eref_sample, **kwargs
+        )
+        yield from clear_one_roi("pfy")
+
+    d = f"Perform an in-place xas scan for {element} with energy pattern {energy_grid} \n"
+    inner.__doc__ = d + inner.__doc__
+
+    inner.__qualname__ = key
+    inner.__name__ = key
+    inner._edge = edge
+    inner._short_doc = (
+        f"Do XAS for {element} from {energy_grid[0]} to {energy_grid[-2]}"
+    )
+    return inner
+
+
+@add_to_func_list
+def load_xas(filename):
+    """Load XAS plans from a TOML file and inject them into the IPython user namespace.
+
+    Parameters
+    ----------
+    filename : str
+        Path to the TOML file containing XAS plan definitions
+    """
+    try:
+        # Get IPython's user namespace
+        ip = get_ipython()
+        user_ns = ip.user_ns
+    except (NameError, AttributeError):
+        # Not running in IPython, just return the generated functions
+        user_ns = None
+
+    generated_plans = {}
+    with open(filename, "rb") as f:
+        regions = tomllib.load(f)
+        for key, value in regions.items():
+            name = value.get("name", key)
+            region = value.get("region")
+            element = value.get("element", "")
+            edge = value.get("edge", "")
+            xas_func = _xas_factory(region, element, edge, key)
+            add_to_xas_list(
+                xas_func, key, name=name, element=element, edge=edge, region=region
+            )
+            add_to_plan_time_dict(
+                xas_func,
+                "gscan_estimate",
+                fixed=5,
+                overhead=0.5,
+                dwell="dwell",
+                region=region,
+            )
+            # Store the function
+            generated_plans[key] = xas_func
+
+            # If we're in IPython, inject into user namespace
+            if user_ns is not None:
+                user_ns[key] = xas_func
+
+    # Return the generated plans dictionary in case it's needed
+    return generated_plans

nbs_bl/printing.py

@@ -0,0 +1,221 @@
+from IPython.utils.coloransi import TermColors as color
+from rich import print
+import sys
+import re
+from select import select
+
+
+def _strip_ansi(text):
+    """
+    Remove ANSI escape sequences from text.
+
+    Parameters
+    ----------
+    text : str
+        Text containing ANSI escape sequences
+
+    Returns
+    -------
+    str
+        Text with ANSI sequences removed
+    """
+    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+    return ansi_escape.sub("", text)
+
+
+def _ansilen(text):
+    """
+    Get visible length of text containing ANSI escape sequences.
+
+    Parameters
+    ----------
+    text : str
+        Text containing ANSI escape sequences
+
+    Returns
+    -------
+    int
+        Visual length of text
+    """
+    return len(_strip_ansi(text))
+
+
+def _wrap_text(text, width, subsequent_indent=""):
+    """
+    Wrap text while preserving ANSI escape sequences.
+
+    Parameters
+    ----------
+    text : str
+        Text to wrap
+    width : int
+        Maximum line width
+    subsequent_indent : str
+        Indentation for lines after the first
+
+    Returns
+    -------
+    str
+        Wrapped text with preserved ANSI sequences
+    """
+    words = text.split()
+    lines = []
+    current_line = []
+    current_length = 0
+
+    for word in words:
+        word_length = _ansilen(word)
+        if current_length + word_length + (1 if current_line else 0) <= width:
+            current_line.append(word)
+            current_length += word_length + (1 if current_line else 0)
+        else:
+            if current_line:
+                lines.append(" ".join(current_line))
+            current_line = [word]
+            current_length = word_length
+
+    if current_line:
+        lines.append(" ".join(current_line))
+
+    # Add indentation to subsequent lines
+    if subsequent_indent and len(lines) > 1:
+        lines[1:] = [subsequent_indent + line for line in lines[1:]]
+
+    return "\n".join(lines)
+
+
+def colored(text, tint="white", attrs=[]):
+    """
+    A simple wrapper around IPython's interface to TermColors
+    """
+    tint = tint.lower()
+    if "dark" in tint:
+        tint = tint[4:].lower()
+    elif "light" in tint:
+        tint = "bright_" + tint[5:].lower()
+    elif "blink" in tint:
+        tint = "bold " + tint[5:].lower()
+    elif "no" in tint:
+        tint = "reset"
+    else:
+        tint = tint.lower()
+    return "[{0}]{1}[/{0}]".format(tint, str(text))
+
+
+# def colored(text, tint="white", attrs=[]):
+#     """
+#     A simple wrapper around IPython's interface to TermColors
+#     """
+#     tint = tint.lower()
+#     if "dark" in tint:
+#         tint = "Dark" + tint[4:].capitalize()
+#     elif "light" in tint:
+#         tint = "Light" + tint[5:].capitalize()
+#     elif "blink" in tint:
+#         tint = "Blink" + tint[5:].capitalize()
+#     elif "no" in tint:
+#         tint = "Normal"
+#     else:
+#         tint = tint.capitalize()
+#     return "{0}{1}{2}".format(getattr(color, tint), str(text), color.Normal)
+
+def read_input(message, default, timeout, secs):
+    print(message)
+    rlist, _, _ = select([sys.stdin], [], [], secs)
+    if rlist:
+        s = sys.stdin.readline()
+        if ord(s[0]) == 13 or ord(s[0]) == 10:
+            return default
+        else:
+            return s.rstrip(chr(10) + chr(13))
+    else:
+        return timeout
+
+
+def run_report(thisfile):
+    """
+    Noisily proclaim to be importing a file of python code.
+    """
+    print(colored("Importing %s ..." % thisfile.split("/")[-1], "lightcyan"))
+
+
+def boxed_text(title, text, tint, width=82, shrink=False):
+    """
+    Put text in a lovely unicode block element box.
+
+    Parameters
+    ----------
+    title : str
+        Title to display at top of box
+    text : str or list
+        Text content to display in box
+    tint : str
+        Color to use for box borders
+    width : int
+        Maximum width of box
+    shrink : bool
+        Whether to shrink box to content width
+    """
+    if isinstance(text, str):
+        textlines = text.splitlines()
+    else:
+        textlines = text
+
+    if shrink:
+        width = min(width, max((_ansilen(line) for line in textlines)))
+
+    remainder = width - 5 - len(title)
+    ul = "\u250C"
+    ur = "\u2510"
+    ll = "\u2514"
+    lr = "\u2518"
+    bar = "\u2500"
+    strut = "\u2502"
+
+    print("")
+    print(colored("".join([ul, bar * 3, " ", title, " ", bar * remainder, ur]), tint))
+
+    for line1 in textlines:
+        lines = _wrap_text(line1, width, subsequent_indent=" " * 26)
+        for line in lines.splitlines():
+            line = line.rstrip()
+            add = " " * (width - _ansilen(line))
+            print("".join([colored(strut, tint), line, add, colored(strut, tint)]))
+
+    print(colored("".join([ll, bar * width, lr]), tint))
+
+
+def error_msg(text):
+    return colored(text, "lightred")
+
+
+def warning_msg(text):
+    return colored(text, "yellow")
+
+
+def url_msg(text):
+    return text
+
+
+def bold_msg(text):
+    return colored(text, "white")
+
+
+def verbosebold_msg(text):
+    return colored(text, "lightcyan")
+
+
+def list_msg(text):
+    return colored(text, "cyan")
+
+
+def disconnected_msg(text):
+    return colored(text, "purple")
+
+
+def info_msg(text):
+    return colored(text, "brown")
+
+
+def whisper(text):
+    return colored(text, "darkgray")

nbs_bl/qt/models/beamline.py

@@ -0,0 +1,11 @@
+from nbs_gui.models.beamline import GUIBeamlineModel
+
+
+class SSTBeamlineModel(GUIBeamlineModel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        try:
+            self.energy.obj.rotation_motor = self.primary_sampleholder.obj.r
+        except AttributeError as e:
+            print(f"Problem loading energy model: {e}")
+        print("Finished loading BeamlineModel")

nbs_bl/qt/models/energy.py

@@ -0,0 +1,53 @@
+from ..widgets.energy import EnergyControl, EnergyMonitor
+from nbs_gui.models import PseudoPositionerModel, PVPositionerModel, PVModel
+
+
+# Copied from ucal as an example
+class EnergyModel:
+    default_controller = EnergyControl
+    default_monitor = EnergyMonitor
+
+    def __init__(
+        self,
+        name,
+        obj,
+        group,
+        long_name,
+        **kwargs,
+    ):
+        print("Initializing Energy")
+        self.name = name
+        self.obj = obj
+        self.energy = PseudoPositionerModel(name, obj, group, name)
+        self.grating_motor = PVPositionerModel(
+            name=obj.monoen.gratingx.name,
+            obj=obj.monoen.gratingx,
+            group=group,
+            long_name=f"{name} Grating",
+        )
+        self.cff = PVModel(
+            obj.monoen.cff.name, obj.monoen.cff, group=group, long_name=f"{name} CFF"
+        )
+        self.real_motors = self.energy.real_motors
+        self.group = group
+        self.label = long_name
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+        print("Done Initializing Energy")
+
+    def iter_models(self):
+        """
+        Yield contained energy-related models for traversal.
+
+        Yields
+        ------
+        BaseModel
+            Contained models.
+        """
+        yield from (
+            self.energy,
+            self.grating_motor,
+            self.cff,
+        )
+        for motor in self.real_motors:
+            yield motor