diff --git a/src/ibex_bluesky_core/plans/muons/dae_magnet_scan.py b/src/ibex_bluesky_core/plans/muons/dae_magnet_scan.py index f079e59..1c5d7e3 100644 --- a/src/ibex_bluesky_core/plans/muons/dae_magnet_scan.py +++ b/src/ibex_bluesky_core/plans/muons/dae_magnet_scan.py @@ -4,10 +4,12 @@ from collections.abc import Generator import bluesky.plans as bp +import bluesky.plan_stubs as bps import matplotlib import matplotlib.pyplot as plt from bluesky.utils import Msg from ophyd_async.plan_stubs import ensure_connected +from ophyd_async.core import Device from ibex_bluesky_core.callbacks import ISISCallbacks from ibex_bluesky_core.devices import get_pv_prefix @@ -23,33 +25,37 @@ from ibex_bluesky_core.run_engine import get_run_engine from ibex_bluesky_core.callbacks.fitting.fitting_utils import Gaussian +import matplotlib.pyplot as plt +from ibex_bluesky_core.plan_stubs import call_sync +try: + from itertools import batched +except ImportError: + from itertools import islice + def batched(iterable, n, *, strict=False): + if n < 1: + raise ValueError('n must be at least one') + iterator = iter(iterable) + while batch := tuple(islice(iterator, n)): + if strict and len(batch) != n: + raise ValueError('batched(): incomplete batch') + yield batch default_prefix = get_pv_prefix() + def dae_magnet_plan( - block, - start, - stop, + *args, num, periods=True, frames=500, save_run=True, - magnet_tolerance=0.01, - magnet_settle_time=1, - prefix = default_prefix + prefix = default_prefix, + confident: bool = False, ) -> Generator[Msg, None, None]: """Scan a DAE against a magnet.""" - - def check_within_tolerance(setpoint: float, actual: float) -> bool: - return setpoint - magnet_tolerance <= actual <= setpoint + magnet_tolerance - - magnet = block_rw_rbv( - float, - block, - write_config=BlockWriteConfig( - settle_time_s=magnet_settle_time, set_success_func=check_within_tolerance - ), - ) + + plt.close("all") + plt.show() if periods: controller = PeriodPerPointController(save_run=save_run) @@ -67,7 +73,13 @@ def check_within_tolerance(setpoint: float, actual: float) -> bool: reducer=reducer, ) - yield from ensure_connected(magnet, dae, force_reconnect=True) + to_connect = [dae] + to_measure = [] + for item in args: + if isinstance(item, Device): + to_connect.append(item) + to_measure.append(item.name) + yield from ensure_connected(*to_connect, force_reconnect=True) if periods: yield from set_num_periods(dae, num) @@ -76,14 +88,28 @@ def check_within_tolerance(setpoint: float, actual: float) -> bool: icc = ISISCallbacks( y=reducer.intensity.name, - x=magnet.name, + x=args[0].name, yerr=reducer.intensity_stddev.name, fit=Gaussian().fit(), + measured_fields=to_measure[1:], ) @icc def _inner() -> Generator[Msg, None, None]: - yield from bp.scan([dae], magnet, start, stop, num=num) + yield from bp.scan([dae], *args, num=num) yield from _inner() print(icc.live_fit.result.fit_report()) + fitted_value = icc.live_fit.result.params["x0"].value + if args[1] <= abs(fitted_value) <= args[2]: + if not confident: + ans = yield from call_sync(input, f"set {args[0].name} to {fitted_value}? y/n") + else: + ans = "y" + if ans.lower().startswith("y"): + fitted_fraction = (fitted_value - args[1]) / (args[2] - args[1]) + for device, start, stop in batched(args, n=3): + value = start + fitted_fraction * (stop - start) + print(f"Setting {device.name} to {value}") + yield from bps.mv(device, abs(value)) + return icc.live_fit.result