backup from ndxchronus

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