Average multiple shots (in Orbit correction and ORM), Action for callback before each shot.

pyaml/common/constants.py

@@ -15,7 +15,7 @@
 
 class Action(Enum):
     """
-    Identify callback points in measurement-tool workflows.
+    Identify callback points in measurement and orbit-correction workflows.
 
     Attributes
     ----------
@@ -27,6 +27,8 @@ class Action(Enum):
         Callback triggered immediately after taking a measurement.
     INIT : int
         Triggered once before the measurement starts
+    BEFORE_BPM_READING : int
+        Triggered before each BPM acquisition and its configured read delay.
     """
 
     APPLY = 0
@@ -37,3 +39,5 @@ class Action(Enum):
     "Triggered immediately after taking a measurement."
     INIT = 3
     "Triggered once before the measurement starts"
+    BEFORE_BPM_READING = 4
+    "Triggered before each BPM acquisition and its configured read delay."

pyaml/external/pySC_interface.py

@@ -6,12 +6,14 @@
 """
 
 import time
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Callable, Optional, Tuple
 
 import numpy as np
 
 if TYPE_CHECKING:
+    from ..common.element import Element
     from ..common.holders.element_holder import ElementHolder
+from ..common.constants import Action
 from ..common.exception import PyAMLException
 
 
@@ -27,6 +29,20 @@ class pySCInterface:
         Name of the BPM array used for orbit readback.
     rf_plant_name : Optional[str]
         Optional RF plant name, required for RF-related pySC operations.
+    callback : Callable, optional
+        Called as ``callback(action, data)`` with ``Action.BEFORE_BPM_READING``
+        before each BPM acquisition.  Return ``True`` to continue;
+        a false value raises ``KeyboardInterrupt``.
+    orbits_to_average : int, default 1
+        Positive number of BPM acquisitions to average per ``get_orbit()`` call.
+    tool : Element, optional
+        Owning tool passed as callback data ``source``. Defaults to this interface
+        when used without an owning tool.
+
+    Raises
+    ------
+    ValueError
+        If ``orbits_to_average`` is not a positive integer.
 
     Methods
     -------
@@ -44,16 +60,23 @@ class pySCInterface:
 
     set_wait_time: float = 0
     read_wait_time: float = 0
+    orbits_to_average: int = 1
 
     def __init__(
         self,
         element_holder: "ElementHolder",
         bpm_array_name: str,
         rf_plant_name: Optional[str] = None,
+        callback: Optional[Callable] = None,
+        orbits_to_average: int = 1,
+        tool: Optional["Element"] = None,
     ):
         """
         Initialize a pySC adapter from a PyAML element holder.
         """
+        if not isinstance(orbits_to_average, (int, np.integer)) or orbits_to_average < 1:
+            raise ValueError("orbits_to_average must be a positive integer.")
+
         self.element_holder = element_holder
 
         self.bpm_array = element_holder.diagnostic.bpms.get(bpm_array_name)
@@ -64,11 +87,40 @@ def __init__(
         else:
             self.rf_plant = None
 
+        self._callback = callback
+        if tool is None:
+            self.tool = self
+        else:
+            self.tool = tool
+        self.orbits_to_average = orbits_to_average
+
+    @property
+    def peer(self) -> "ElementHolder":
+        return self.element_holder
+
+    def attached_to(self) -> str:
+        """Return the peer type and name, or ``"None"`` when unattached."""
+        return "None" if self.peer is None else f"{self.peer.__class__.__name__}:{self.peer.name()}"
+
     def get_orbit(self) -> Tuple[np.array, np.array]:
-        """Return horizontal and vertical orbit readings from the BPM array."""
+        """Return mean horizontal and vertical orbits over ``orbits_to_average`` reads.
+
+        Invoke ``Action.BEFORE_BPM_READING`` before each acquisition and its
+        ``read_wait_time`` delay, including each sample used in the average.
+        """
+        self.send_callback(Action.BEFORE_BPM_READING, {})
         time.sleep(self.read_wait_time)
         positions = self.bpm_array.positions.get()
-        return positions[:, 0], positions[:, 1]
+        x = positions[:, 0] / self.orbits_to_average
+        y = positions[:, 1] / self.orbits_to_average
+
+        for _ in range(1, self.orbits_to_average):
+            self.send_callback(Action.BEFORE_BPM_READING, {})
+            time.sleep(self.read_wait_time)
+            positions = self.bpm_array.positions.get()
+            x += positions[:, 0] / self.orbits_to_average
+            y += positions[:, 1] / self.orbits_to_average
+        return x, y
 
     def get(self, name: str) -> float:
         """
@@ -133,3 +185,49 @@ def set_rf_main_frequency(self, value: float) -> None:
         self.rf_plant.frequency.set(value)
         time.sleep(self.set_wait_time)
         return
+
+    # TODO: this function is copied from MeasurementTool class, to be generalized?
+    def send_callback(self, action: Action, cb_data: dict, raiseException: bool = True):
+        """
+        Notify the caller about measurement progress.
+
+        If the registered callback returns ``False``, the scan is aborted by
+        raising ``KeyboardInterrupt`` when ``raiseException`` is ``True``.
+        Callback data contains the attached ``mode`` and the owning tool as
+        ``source`` (or this interface when no tool was provided).
+        Callback example:
+
+        .. code-block:: python
+
+            def callback(action: Action, data: dict):
+                print(f"{action}, data:{data}")
+                return True
+
+            # Correct the orbit using the above callback
+            sr.design.tool.orbit.correct(callback=callback)
+
+        Parameters
+        ----------
+        action : Action
+          Measurement lifecycle event.
+
+        cb_data : dict
+          Mutable event data passed to the callback.
+        raiseException : bool, optional
+          Whether a callback rejection should raise ``KeyboardInterrupt``.
+
+        Returns
+        -------
+        bool
+            Callback result, or ``True`` when no callback is registered.
+        """
+        ok = True
+        if self._callback is not None:
+            # Add source and peer
+            cb_data["mode"] = self.tool.attached_to()
+            cb_data["source"] = self.tool
+            ok = self._callback(action, cb_data)
+        if not ok and raiseException:
+            # Abort, same as ctrl+C
+            raise KeyboardInterrupt
+        return ok

pyaml/tuning_tools/bba.py

@@ -244,6 +244,7 @@ def measure(
         interface = pySCInterface(
             element_holder=element_holder,
             bpm_array_name=self.bpm_array_name,
+            tool=self,
         )
         interface.set_wait_time = sleep_step
         interface.read_wait_time = sleep_meas

pyaml/tuning_tools/dispersion.py

@@ -116,8 +116,9 @@ def measure(
         set_waiting_time : float
             Delay in seconds after changing the RF frequency.
         callback : Optional[Callable]
-            Optional callback invoked after applying, measuring, and restoring
-            settings.
+            Optional callback invoked before each BPM reading
+            (``Action.BEFORE_BPM_READING``) and after applying, measuring, and
+            restoring settings.
 
         Returns
         -------
@@ -130,6 +131,8 @@ def measure(
             element_holder=element_holder,
             bpm_array_name=self.bpm_array_name,
             rf_plant_name=self.rf_plant_name,
+            callback=callback,
+            tool=self,
         )
         interface.set_wait_time = set_waiting_time
 

pyaml/tuning_tools/orbit.py

@@ -8,7 +8,7 @@
 import logging
 from dataclasses import asdict
 from pathlib import Path
-from typing import Literal, Optional, Union
+from typing import Callable, Literal, Optional, Union
 
 import numpy as np
 from pySC import ResponseMatrix as pySC_ResponseMatrix
@@ -260,6 +260,8 @@ def correct(
         reference: Optional[np.ndarray] = None,
         rf: bool = False,
         virtual_target: Optional[float] = None,
+        orbits_to_average: int = 1,
+        callback: Optional[Callable] = None,
     ):
         """
         Perform orbit correction using the configured response matrix and corrector
@@ -299,6 +301,22 @@ def correct(
         virtual_target : float, optional
             Target sum of horizontal corrector changes, in radians. Defaults to
             the configured value.
+        orbits_to_average : int, default 1
+            Positive number of BPM acquisitions to average for each orbit read
+            used to calculate the correction.
+        callback : Callable, optional
+            Called as ``callback(action, data)`` with ``Action.BEFORE_BPM_READING``
+            before each BPM acquisition, including each sample in an average.
+            Data includes the attached ``mode`` and this tool as ``source``.
+            Return ``True`` to continue; a false value raises ``KeyboardInterrupt``
+            and aborts correction before applying trims.
+
+        Raises
+        ------
+        ValueError
+            If ``orbits_to_average`` is not a positive integer.
+        KeyboardInterrupt
+            If the callback requests an abort.
         """
 
         if self._pySC_response_matrix is None:
@@ -307,6 +325,9 @@ def correct(
         interface = pySCInterface(
             element_holder=self.peer,
             bpm_array_name=self.bpm_array_name,
+            callback=callback,
+            tool=self,
+            orbits_to_average=orbits_to_average,
         )
 
         if singular_values_h is not None:

pyaml/tuning_tools/orbit_response_matrix.py

@@ -206,6 +206,8 @@ def measure(
         interface = pySCInterface(
             element_holder=element_holder,
             bpm_array_name=self.bpm_array_name,
+            callback=callback,
+            tool=self,
         )
         # TODO handle sleep_meas
         interface.set_wait_time = sleep_step