diff --git a/pyaml/common/constants.py b/pyaml/common/constants.py index 4d4ee5ac..dc9a49c8 100644 --- a/pyaml/common/constants.py +++ b/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." diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index b9e60493..8f6c122e 100644 --- a/pyaml/external/pySC_interface.py +++ b/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 diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index feb8bccd..c7758fe4 100644 --- a/pyaml/tuning_tools/bba.py +++ b/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 diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index 81fccb11..a7c42662 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/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 diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index 61468323..c566c0e0 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/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: diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 14bd96d2..a85b70ed 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/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