From 82b64029208f1546c6135d647272863b6b5193bc Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 16:10:06 +0200 Subject: [PATCH 01/11] add averaging in orbit correct, and a callback whenever bpms are read during orbit correction, new action for bpm reading --- pyaml/common/constants.py | 2 + pyaml/external/pySC_interface.py | 65 +++++++++++++++++++++++++++++++- pyaml/tuning_tools/orbit.py | 6 ++- 3 files changed, 70 insertions(+), 3 deletions(-) diff --git a/pyaml/common/constants.py b/pyaml/common/constants.py index 4d4ee5ac..fbbd4dd3 100644 --- a/pyaml/common/constants.py +++ b/pyaml/common/constants.py @@ -37,3 +37,5 @@ class Action(Enum): "Triggered immediately after taking a measurement." INIT = 3 "Triggered once before the measurement starts" + BEFORE_BPM_READING = 4 + "Triggered immediately before taking a measurement." diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index b9e60493..ebb5d461 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -6,12 +6,13 @@ """ 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.holders.element_holder import ElementHolder +from ..common.constants import Action from ..common.exception import PyAMLException @@ -44,12 +45,14 @@ 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, ): """ Initialize a pySC adapter from a PyAML element holder. @@ -64,11 +67,23 @@ def __init__( else: self.rf_plant = None + self._callback = callback + def get_orbit(self) -> Tuple[np.array, np.array]: """Return horizontal and vertical orbit readings from the BPM array.""" + self.send_callback(Action.BEFORE_MEASURE, {}) time.sleep(self.read_wait_time) positions = self.bpm_array.positions.get() - return positions[:, 0], positions[:, 1] + x = positions[:, 0] + y = positions[:, 1] + + 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 +148,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``. + The measurement tool adds its mode and source to ``cb_data`` before + invoking the callback. + Callback example: + + .. code-block:: python + + def callback(action: Action, data: dict): + print(f"{action}, data:{data}") + return True + + # Measure a tune response matrix using the above callback + sr.design.tool.trm.measure(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"] = f"{self.attached_to()}" + cb_data["source"] = self + 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/orbit.py b/pyaml/tuning_tools/orbit.py index 61468323..fd3f7cec 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 @@ -307,6 +309,8 @@ def correct( interface = pySCInterface( element_holder=self.peer, bpm_array_name=self.bpm_array_name, + callback=callback, + orbits_to_average=orbits_to_average, ) if singular_values_h is not None: From b94b10458a39d9e160992cf8dd17d6b0e5237162 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 16:24:16 +0200 Subject: [PATCH 02/11] fix Action and add callback for ORM bpm reading --- pyaml/external/pySC_interface.py | 2 +- pyaml/tuning_tools/orbit_response_matrix.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index ebb5d461..65e2a1f2 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -71,7 +71,7 @@ def __init__( def get_orbit(self) -> Tuple[np.array, np.array]: """Return horizontal and vertical orbit readings from the BPM array.""" - self.send_callback(Action.BEFORE_MEASURE, {}) + self.send_callback(Action.BEFORE_BPM_READING, {}) time.sleep(self.read_wait_time) positions = self.bpm_array.positions.get() x = positions[:, 0] diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 14bd96d2..44ae7642 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -206,6 +206,7 @@ def measure( interface = pySCInterface( element_holder=element_holder, bpm_array_name=self.bpm_array_name, + callback=callback, ) # TODO handle sleep_meas interface.set_wait_time = sleep_step From 0342bbd1d9b2a233c5d3cb18e42d843b8cea9ba3 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 16:27:53 +0200 Subject: [PATCH 03/11] fix orbit averages arguement --- pyaml/external/pySC_interface.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index 65e2a1f2..a5517a07 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -53,6 +53,7 @@ def __init__( bpm_array_name: str, rf_plant_name: Optional[str] = None, callback: Optional[Callable] = None, + orbits_to_average: int = 1, ): """ Initialize a pySC adapter from a PyAML element holder. @@ -68,6 +69,7 @@ def __init__( self.rf_plant = None self._callback = callback + self.orbits_to_average = orbits_to_average def get_orbit(self) -> Tuple[np.array, np.array]: """Return horizontal and vertical orbit readings from the BPM array.""" From 5283c17c36c1b57aa6d6a0df2f9b5d2dd27159b4 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 16:43:15 +0200 Subject: [PATCH 04/11] fix callback sending --- pyaml/external/pySC_interface.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index a5517a07..f22a7f72 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -189,7 +189,8 @@ def callback(action: Action, data: dict): ok = True if self._callback is not None: # Add source and peer - cb_data["mode"] = f"{self.attached_to()}" + # cb_data["mode"] = f"{self.attached_to()}" + cb_data["mode"] = "" cb_data["source"] = self ok = self._callback(action, cb_data) if not ok and raiseException: From ebfc7addd354c2d049bd075cbe463d5c18905b2b Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 16:48:17 +0200 Subject: [PATCH 05/11] fix peer --- pyaml/external/pySC_interface.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index f22a7f72..6fc242c1 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -71,6 +71,10 @@ def __init__( self._callback = callback self.orbits_to_average = orbits_to_average + @property + def peer(self) -> "ElementHolder": + return self.element_holder + def get_orbit(self) -> Tuple[np.array, np.array]: """Return horizontal and vertical orbit readings from the BPM array.""" self.send_callback(Action.BEFORE_BPM_READING, {}) From cc48caee8a44a08e8edcff055cbd8bdd29bfb055 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 17:35:13 +0200 Subject: [PATCH 06/11] fix averaging bug --- pyaml/external/pySC_interface.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index 6fc242c1..06dcf648 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -80,8 +80,8 @@ def get_orbit(self) -> Tuple[np.array, np.array]: self.send_callback(Action.BEFORE_BPM_READING, {}) time.sleep(self.read_wait_time) positions = self.bpm_array.positions.get() - x = positions[:, 0] - y = 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, {}) From 2085345c9f9d0ae1b34df05f93291ecf029fbb08 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 17:39:34 +0200 Subject: [PATCH 07/11] mode in callback data --- pyaml/external/pySC_interface.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index 06dcf648..b9557bbf 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -75,6 +75,10 @@ def __init__( 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.""" self.send_callback(Action.BEFORE_BPM_READING, {}) @@ -193,8 +197,7 @@ def callback(action: Action, data: dict): ok = True if self._callback is not None: # Add source and peer - # cb_data["mode"] = f"{self.attached_to()}" - cb_data["mode"] = "" + cb_data["mode"] = self.attached_to() cb_data["source"] = self ok = self._callback(action, cb_data) if not ok and raiseException: From 8800a8781a813e5ba863f3933b9cbdef02a1b220 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 18:03:34 +0200 Subject: [PATCH 08/11] check that orbits_to_average is valid --- pyaml/external/pySC_interface.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index b9557bbf..004dd6f2 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -58,6 +58,9 @@ def __init__( """ 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) From a26abbfe465f4eede525b59365c31513a3842a76 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 18:46:37 +0200 Subject: [PATCH 09/11] add callback to dispersion --- pyaml/tuning_tools/dispersion.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index 81fccb11..4f9fd1fb 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,7 @@ def measure( element_holder=element_holder, bpm_array_name=self.bpm_array_name, rf_plant_name=self.rf_plant_name, + callback=callback, ) interface.set_wait_time = set_waiting_time From e00566f7cf2e64b279e7f8910e2bce6139158b77 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 18:54:49 +0200 Subject: [PATCH 10/11] modify docstrings --- pyaml/common/constants.py | 6 ++++-- pyaml/external/pySC_interface.py | 21 ++++++++++++++++++--- pyaml/tuning_tools/orbit.py | 15 +++++++++++++++ 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/pyaml/common/constants.py b/pyaml/common/constants.py index fbbd4dd3..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 @@ -38,4 +40,4 @@ class Action(Enum): INIT = 3 "Triggered once before the measurement starts" BEFORE_BPM_READING = 4 - "Triggered immediately before taking a measurement." + "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 004dd6f2..1f0039c2 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -28,6 +28,17 @@ 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. + + Raises + ------ + ValueError + If ``orbits_to_average`` is not a positive integer. Methods ------- @@ -83,7 +94,11 @@ def attached_to(self) -> str: 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() @@ -179,8 +194,8 @@ def callback(action: Action, data: dict): print(f"{action}, data:{data}") return True - # Measure a tune response matrix using the above callback - sr.design.tool.trm.measure(callback=callback) + # Correct the orbit using the above callback + sr.design.tool.orbit.correct(callback=callback) Parameters ---------- diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index fd3f7cec..d7965c2f 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -301,6 +301,21 @@ 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. + 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: From 1deef285eb7d07192402bfa249fa5142cf336a66 Mon Sep 17 00:00:00 2001 From: kparasch Date: Fri, 23 Oct 2026 19:06:50 +0200 Subject: [PATCH 11/11] pass tool instead of pySC interface in the callback of pySC --- pyaml/external/pySC_interface.py | 17 +++++++++++++---- pyaml/tuning_tools/bba.py | 1 + pyaml/tuning_tools/dispersion.py | 1 + pyaml/tuning_tools/orbit.py | 2 ++ pyaml/tuning_tools/orbit_response_matrix.py | 1 + 5 files changed, 18 insertions(+), 4 deletions(-) diff --git a/pyaml/external/pySC_interface.py b/pyaml/external/pySC_interface.py index 1f0039c2..8f6c122e 100644 --- a/pyaml/external/pySC_interface.py +++ b/pyaml/external/pySC_interface.py @@ -11,6 +11,7 @@ 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 @@ -34,6 +35,9 @@ class pySCInterface: 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 ------ @@ -65,6 +69,7 @@ def __init__( 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. @@ -83,6 +88,10 @@ def __init__( 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 @@ -184,8 +193,8 @@ def send_callback(self, action: Action, cb_data: dict, raiseException: bool = Tr If the registered callback returns ``False``, the scan is aborted by raising ``KeyboardInterrupt`` when ``raiseException`` is ``True``. - The measurement tool adds its mode and source to ``cb_data`` before - invoking the callback. + 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 @@ -215,8 +224,8 @@ def callback(action: Action, data: dict): ok = True if self._callback is not None: # Add source and peer - cb_data["mode"] = self.attached_to() - cb_data["source"] = self + 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 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 4f9fd1fb..a7c42662 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -132,6 +132,7 @@ def measure( 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 d7965c2f..c566c0e0 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -307,6 +307,7 @@ def correct( 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. @@ -325,6 +326,7 @@ def correct( element_holder=self.peer, bpm_array_name=self.bpm_array_name, callback=callback, + tool=self, orbits_to_average=orbits_to_average, ) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 44ae7642..a85b70ed 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -207,6 +207,7 @@ def measure( 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