nbs-bl 0.2.0__py3-none-any.whl

nbs_bl/plans/maximizers.py

@@ -0,0 +1,423 @@
+from bluesky_live.bluesky_run import BlueskyRun, DocumentCache
+import bluesky.preprocessors as bpp
+from .preprocessors import run_return_decorator
+from bluesky.plan_stubs import mv, mvr, trigger_and_read
+import numpy as np
+
+from nbs_bl.plans.flyscan_base import fly_scan
+from .scan_decorators import wrap_plan_name
+import pandas as pd
+from typing import List
+
+
+def process_fly_scan(run, signal_names=None, time_offsets=None):
+    if time_offsets is None:
+        time_offsets = {}
+
+    df = pd.DataFrame()
+    for stream_name in run:
+        if "monitor" not in stream_name:
+            print(stream_name)
+            continue
+        column_name = stream_name.replace("_monitor", "")
+        newdf = pd.DataFrame(
+            {
+                "time": run[stream_name].read()["time"].to_numpy(),
+                column_name: run[stream_name].read()[column_name].to_numpy(),
+            }
+        ).set_index("time")
+        newdf.index += time_offsets.get(stream_name, 0.0)
+        df = pd.concat((df, newdf))
+    non_monitor_signals = [
+        signal for signal in signal_names if signal not in df.columns
+    ]
+    print("Non-monitor signals: ", non_monitor_signals)
+    primary = run.primary.read()
+    for signal_name in non_monitor_signals:
+        if signal_name in primary:
+            newdf = pd.DataFrame(
+                {
+                    "time": primary["time"].to_numpy(),
+                    signal_name: primary[signal_name].to_numpy(),
+                }
+            ).set_index("time")
+            newdf.index += time_offsets.get(signal_name, 0.0)
+            df = pd.concat((df, newdf))
+    # df = df[~df.index.duplicated(keep='first')].sort_index().interpolate(method='index').ffill().bfill()
+
+    df = (
+        df.groupby("time")
+        .mean()
+        .sort_index()
+        .interpolate(method="index")
+        .ffill()
+        .bfill()
+    )
+
+    return df
+
+
+def find_optimum_motor_pos(
+    run,
+    motor_name: str,
+    signal_names: List[str],
+    time_offsets=None,
+    invert=False,
+):
+    """
+    df = process_fly_scan(run, [motor_name] + signal_names, time_offsets)
+    max_signal_dict = {}
+    for signal_name in signal_names:
+        idx = df[signal_name].idxmax() if not invert else df[signal_name].idxmin()
+        max_signal_dict[signal_name] = {}
+        max_signal_dict[signal_name]["time"] = idx
+        max_signal_dict[signal_name][motor_name] = df[motor_name][idx]
+        max_signal_dict[signal_name]["value"] = df[signal_name][idx]
+    return max_signal_dict
+    """
+    table = run.primary.read()
+    max_info = {}
+    for detname in signal_names:
+        if invert:
+            idx = int(table[detname].argmin())
+            print(f"Minimum found at step {idx} for detector {detname}")
+        else:
+            idx = int(table[detname].argmax())
+            print(f"Maximum found at step {idx} for detector {detname}")
+        max_info[detname] = {"idx": idx, "value": table[detname][idx]}
+        max_val = float(table[motor_name][idx])
+        max_info[detname][motor_name] = max_val
+    return max_info
+
+
+@wrap_plan_name
+def fly_max(
+    detectors,
+    motor,
+    *args,
+    max_channel=None,
+    invert=False,
+    end_on_max=True,
+    md=None,
+    period: float = 1,
+    time_offsets: dict = None,
+    **kwargs,
+):
+    r"""
+    plan: tune a motor to the maximum of signal(motor)
+
+    Initially, traverse the range from start to stop with
+    the number of points specified.  Repeat with progressively
+    smaller step size until the minimum step size is reached.
+    Rescans will be centered on the signal maximum
+    with original scan range reduced by ``step_factor``.
+
+    Set ``snake=True`` if your positions are reproducible
+    moving from either direction.  This will not
+    decrease the number of traversals required to reach convergence.
+    Snake motion reduces the total time spent on motion
+    to reset the positioner.  For some positioners, such as
+    those with hysteresis, snake scanning may not be appropriate.
+    For such positioners, always approach the positions from the
+    same direction.
+
+    Note:  if there are multiple maxima, this function may find a smaller one
+    unless the initial number of steps is large enough.
+
+    Parameters
+    ----------
+    detectors : Signal
+        list of 'readable' objects
+    motor : object
+        any 'settable' object (motor, temp controller, etc.)
+    start : float
+        start of range
+    stop : float
+        end of range, note: start < stop
+    velocities : list of floats
+        list of speeds to set motor to during run.
+    max_channel : list of strings
+        detector fields whose output is to maximize. If not given, the first detector is used.
+        (the first will be maximized, but secondardy maxes will be recorded during the scans for the first -
+        if the maxima are not in the same range this will not be useful)
+    md : dict, optional
+        metadata
+    period: float, optional
+        Detector integration time in seconds.
+    **kwargs : dict, optional
+        additional arguments to pass to fly_scan
+
+    """
+    if max_channel is None:
+        max_channel = [detectors[0].name]
+    _md = {
+        "maximizer_args": {
+            "plan": "fly_scan",
+            "max_channel": max_channel,
+            "end_on_max": end_on_max,
+            "invert": invert,
+        },
+        "hints": {},
+    }
+    _md.update(md or {})
+
+    dc = DocumentCache()
+
+    @bpp.subs_decorator(dc)
+    def inner_maximizer():
+        yield from fly_scan(detectors, motor, *args, md=_md, period=period, **kwargs)
+        run = BlueskyRun(dc)
+        motor_name = motor.name
+        max_info = {}
+        max_info = find_optimum_motor_pos(run, motor_name, max_channel, invert=invert)
+        if end_on_max:
+            motor_pos = max_info[max_channel[0]][motor_name]
+            print(f"going to position {motor_pos} for {motor_name}")
+            yield from mv(motor, motor_pos)
+        return max_info
+
+    return (yield from inner_maximizer())
+
+
+@wrap_plan_name
+def find_max(
+    plan,
+    dets,
+    *args,
+    max_channel=None,
+    invert=False,
+    end_on_max=True,
+    md=None,
+    **kwargs,
+):
+    """
+    invert turns find_max into find_min
+    """
+    dc = DocumentCache()
+
+    md = md or {}
+
+    _md = {
+        "maximizer_args": {
+            "plan": plan.__name__,
+            "max_channel": max_channel,
+            "invert": invert,
+            "end_on_max": end_on_max,
+        },
+    }
+    _md.update(md)
+
+    @bpp.subs_decorator(dc)
+    def inner_maximizer():
+        yield from plan(dets, *args, md=_md, **kwargs)
+        run = BlueskyRun(dc)
+        table = run.primary.read()
+        motor_names = run.metadata["start"]["motors"]
+        motors = [m for m in args if getattr(m, "name", None) in motor_names]
+        if max_channel is None:
+            detname = dets[0].name
+        else:
+            detname = max_channel
+        if invert:
+            max_idx = int(table[detname].argmin())
+            print(f"Minimum found at step {max_idx} for detector {detname}")
+        else:
+            max_idx = int(table[detname].argmax())
+            print(f"Maximum found at step {max_idx} for detector {detname}")
+        ret = []
+        for m in motors:
+            max_val = float(table[m.name][max_idx])
+            print(f"setting {m.name} to {max_val}")
+            ret.append([m, max_val])
+        if end_on_max:
+            print("going to found motor positions")
+            flat_list = [item for sublist in ret for item in sublist]
+            yield from mv(*flat_list)
+        return ret
+
+    return (yield from inner_maximizer())
+
+
+def find_min(plan, dets, *args):
+    return (yield from find_max(plan, dets, *args, invert=True))
+
+
+def find_max_deriv(plan, dets, *args, max_channel=None):
+    dc = DocumentCache()
+
+    @bpp.subs_decorator(dc)
+    def inner_maximizer():
+        yield from plan(dets, *args)
+        run = BlueskyRun(dc)
+        table = run.primary.read()
+        motor_names = run.metadata["start"]["motors"]
+        motors = [m for m in args if getattr(m, "name", None) in motor_names]
+        if len(motors) > 1:
+            print(
+                "Derivative with multiple motors unsupported, \
+            taking first motor"
+            )
+
+        if max_channel is None:
+            detname = dets[0].name
+        else:
+            detname = max_channel
+        motname = motors[0].name
+        max_idx = np.argmax(np.abs(np.gradient(table[detname], table[motname])))
+        print(f"Maximum derivative found at step {max_idx} for detector {detname}")
+        ret = []
+        for m in motors:
+            max_val = float(table[m.name][max_idx])
+            print(f"setting {m.name} to {max_val}")
+            ret.append([m, max_val])
+            yield from mv(m, max_val)
+        return ret
+
+    return (yield from inner_maximizer())
+
+
+def find_halfmax(plan, dets, *args, max_channel=None, **kwargs):
+    """
+    For a plan and detector that goes from low to high, find
+    the motor value where the detector is half of the maximum
+    value
+    """
+    dc = DocumentCache()
+
+    @bpp.subs_decorator(dc)
+    def inner_maximizer():
+        yield from plan(dets, *args, **kwargs)
+        run = BlueskyRun(dc)
+        table = run.primary.read()
+        motor_names = run.metadata["start"]["motors"]
+        motors = [m for m in args if getattr(m, "name", None) in motor_names]
+        if max_channel is None:
+            detname = dets[0].name
+        else:
+            detname = max_channel
+        max_val = float(table[detname].max())
+        halftable = table[detname] - max_val / 2.0
+        half_idx = 0
+        for n, v in enumerate(halftable.data):
+            if v > 0:
+                half_idx = n
+                break
+        ret = []
+        for m in motors:
+            mot_val = float(table[m.name][half_idx])
+            print(f"setting {m.name} to {mot_val}")
+            ret.append([m, mot_val])
+            yield from mv(m, mot_val)
+        return ret
+
+    return (yield from inner_maximizer())
+
+
+def halfmax_adaptive(
+    dets, motor, step: float = 5, precision: float = 1, maxct=None, max_channel=None
+):
+    if max_channel is None:
+        detname = dets[0].name
+    else:
+        detname = max_channel
+
+    def ct():
+        ret = yield from trigger_and_read(dets)
+        return ret[detname]["value"]
+
+    @bpp.stage_decorator(dets)
+    @run_return_decorator()
+    def halfmax_inner(step, maxct=None):
+        if maxct is None:
+            maxct = yield from ct()
+        current = maxct
+        while np.abs(step) > precision / 2.0:
+            yield from mvr(motor, -1 * step)
+            current = yield from ct()
+            while current > maxct / 2.0:
+                yield from mvr(motor, step)
+                current = yield from ct()
+            step = step / 2.0
+            if np.abs(step) > precision / 2.0:
+                print(f"{detname} halfmax at {motor.name}:{motor.position}")
+                print(f"Value: {current}, Max: {maxct}, reducing step to {step}")
+
+        print(f"{detname} halfmax at {motor.name}:{motor.position}")
+        print(f"Value: {current}, Max: {maxct}, " f"precision: {precision} reached")
+        return motor.position
+
+    return (yield from halfmax_inner(step, maxct))
+
+
+def threshold_adaptive(
+    dets, motor, threshold, step: float = 2, limit: int = 15, max_channel=None
+):
+    """
+    Attempt to get a detector over a threshold by moving a motor
+    """
+
+    if max_channel is None:
+        detname = dets[0].name
+    else:
+        detname = max_channel
+
+    def ct():
+        ret = yield from trigger_and_read(dets)
+        return ret[detname]["value"]
+
+    @bpp.stage_decorator(dets)
+    @run_return_decorator()
+    def inner_threshold():
+        pos = motor.position
+        maxpos = pos
+
+        n = 0
+        current = yield from ct()
+
+        if current > threshold:
+            mincurrent = current
+            minpos = motor.position
+            print(
+                f"Starting above threshold of {detname}, try to get below {threshold} by moving {motor.name} with starting position {pos} and step -{step}"
+            )
+            while current > threshold and n < limit:
+                yield from mvr(motor, -1 * step)
+                current = yield from ct()
+                if current < mincurrent:
+                    mincurrent = current
+                    minpos = motor.position
+                n += 1
+            if current < threshold:
+                pass
+            else:
+                raise ValueError(
+                    f"Detector {detname} did not fall below {threshold} after"
+                    f" {limit} moves of {motor.name} with -{step} step "
+                    f"size. Minimum value was {mincurrent} at {minpos}."
+                    f"Check if {motor.name} is going the right direction,"
+                    f" and if {detname} is on."
+                )
+
+        print(
+            f"Searching for threshold value {threshold} of {detname} for {motor.name} with starting position {pos} and step {step}"
+        )
+        maxcurrent = current
+        while current < threshold and n < limit:
+            yield from mvr(motor, step)
+            current = yield from ct()
+            if current > maxcurrent:
+                maxcurrent = current
+                maxpos = motor.position
+            n += 1
+        if current > threshold:
+            return motor.position
+        else:
+            raise ValueError(
+                f"Detector {detname} did not exceed {threshold} after"
+                f" {limit} moves of {motor.name} with {step} step "
+                f"size. Maximum value was {maxcurrent} at {maxpos}."
+                f"Check if {motor.name} is going the right direction,"
+                f" and if {detname} is on."
+            )
+
+    return (yield from inner_threshold())

nbs_bl/plans/metaplans.py

@@ -0,0 +1,179 @@
+import time
+import typing
+from ..help import add_to_plan_list
+
+from bluesky_queueserver import parameter_annotation_decorator
+
+
+@add_to_plan_list
+@parameter_annotation_decorator(
+    {
+        "parameters": {
+            "plans": {
+                "annotation": "typing.List[__PLAN__]",
+                "description": "List of plan names or callables to run in sequence.",
+            }
+        }
+    }
+)
+def repeat_plan_sequence_for_duration(
+    plans,
+    plan_args_list: typing.List[typing.List],
+    plan_kwargs_list: typing.List[typing.Dict],
+    duration: float,
+):
+    """
+    Repeat a sequence of plans for a specified duration.
+
+    Parameters
+    ----------
+    plans : list of str or callables
+        List of plan names (or callables) to run in sequence.
+    plan_args_list : list of list
+        List of argument lists, one for each plan.
+    plan_kwargs_list : list of dict
+        List of kwargs dicts, one for each plan.
+    duration : float
+        Total duration to repeat the sequence (seconds).
+
+    Yields
+    ------
+    Msg
+        Bluesky messages from the repeated plans.
+    """
+    start_time = time.time()
+    n_plans = len(plans)
+    idx = 0
+    while (time.time() - start_time) < duration:
+        plan = plans[idx % n_plans]
+        args = plan_args_list[idx % n_plans]
+        kwargs = plan_kwargs_list[idx % n_plans]
+        yield from plan(*args, **kwargs)
+        idx += 1
+
+
+@add_to_plan_list
+@parameter_annotation_decorator(
+    {
+        "parameters": {
+            "plans": {
+                "annotation": "typing.List[__PLAN__]",
+                "description": "List of plan names or callables to run in sequence.",
+            },
+            "condition": {
+                "annotation": "__PLAN__",
+                "description": "A plan that returns a boolean; controls loop continuation.",
+            },
+        }
+    }
+)
+def repeat_plan_sequence_while_condition(
+    plans,
+    plan_args_list: typing.List[typing.List],
+    plan_kwargs_list: typing.List[typing.Dict],
+    condition,
+    condition_args: typing.List = None,
+    condition_kwargs: typing.Dict = None,
+):
+    """
+    Repeat a sequence of plans while a condition plan returns True.
+
+    Parameters
+    ----------
+    plans : list of str or callables
+        List of plan names (or callables) to run in sequence.
+    plan_args_list : list of list
+        List of argument lists, one for each plan.
+    plan_kwargs_list : list of dict
+        List of kwargs dicts, one for each plan.
+    condition : str or callable
+        A plan that returns a boolean; controls loop continuation.
+    condition_args : list, optional
+        Arguments for the condition plan.
+    condition_kwargs : dict, optional
+        Keyword arguments for the condition plan.
+
+    Yields
+    ------
+    Msg
+        Bluesky messages from the repeated plans.
+    """
+    n_plans = len(plans)
+    idx = 0
+    if condition_args is None:
+        condition_args = []
+    if condition_kwargs is None:
+        condition_kwargs = {}
+    while True:
+        val = yield from condition(*condition_args, **condition_kwargs)
+        if not val:
+            break
+        plan = plans[idx % n_plans]
+        args = plan_args_list[idx % n_plans]
+        kwargs = plan_kwargs_list[idx % n_plans]
+        yield from plan(*args, **kwargs)
+        idx += 1
+
+
+@add_to_plan_list
+@parameter_annotation_decorator(
+    {
+        "parameters": {
+            "plans": {
+                "annotation": "typing.List[__PLAN__]",
+                "description": "List of plan names or callables to run in sequence.",
+            },
+            "condition": {
+                "annotation": "__PLAN__",
+                "description": "A plan that returns a boolean; controls loop continuation.",
+            },
+        }
+    }
+)
+def repeat_plan_sequence_until_condition(
+    plans,
+    plan_args_list: typing.List[typing.List],
+    plan_kwargs_list: typing.List[typing.Dict],
+    condition,
+    condition_args: typing.List = None,
+    condition_kwargs: typing.Dict = None,
+):
+    """
+    Repeat a sequence of plans until a condition plan returns True.
+
+    Parameters
+    ----------
+    plans : list of str or callables
+        List of plan names (or callables) to run in sequence.
+    plan_args_list : list of list
+        List of argument lists, one for each plan.
+    plan_kwargs_list : list of dict
+        List of kwargs dicts, one for each plan.
+    condition : str or callable
+        A plan that returns a boolean; controls loop continuation.
+    condition_args : list, optional
+        Arguments for the condition plan.
+    condition_kwargs : dict, optional
+        Keyword arguments for the condition plan.
+
+    Yields
+    ------
+    Msg
+        Bluesky messages from the repeated plans.
+    """
+    n_plans = len(plans)
+    idx = 0
+    if condition_args is None:
+        condition_args = []
+    if condition_kwargs is None:
+        condition_kwargs = {}
+    while True:
+        val = yield from condition(*condition_args, **condition_kwargs)
+        if val:
+            break
+        plan = plans[idx % n_plans]
+        args = plan_args_list[idx % n_plans]
+        kwargs = plan_kwargs_list[idx % n_plans]
+        yield from plan(*args, **kwargs)
+        idx += 1
+