diff --git a/pyaml/tuning_tools/bba.py b/pyaml/tuning_tools/bba.py index d1a71050..227df103 100644 --- a/pyaml/tuning_tools/bba.py +++ b/pyaml/tuning_tools/bba.py @@ -4,7 +4,6 @@ import matplotlib.pyplot as plt import numpy as np import pySC -from pydantic import ConfigDict from pySC.apps import measure_bba from pySC.apps.bba import BBAAnalysis from pySC.apps.codes import BBACode @@ -12,58 +11,93 @@ from ..common.constants import Action from ..common.exception import PyAMLException from ..external.pySC_interface import pySCInterface -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool logger = logging.getLogger(__name__) PYAMLCLASS = "BBA" -class ConfigModel(MeasurementToolConfigModel): +@register_schema +class BBA(MeasurementTool, DynamicValidation): """ - Configuration model for Beam Based Alignment. - BBA finds the magnetic center of a quad (zero crossing). + Beam-based alignment measurement tool. + + This tool determines the magnetic center of a quadrupole by varying its + strength while applying controlled horizontal and vertical orbit offsets. + The quadrupole center is identified from the corresponding zero crossings + in the BPM response. Parameters ---------- + name : str + Name of the measurement tool. bpm_array_name : str - BPM array name (orbit) + Name of the BPM array used to measure the orbit. bpm_name : str - BPM to be corrected (close to the quad) + Name of the BPM located near the quadrupole whose center is measured. hcorr_name : str - Horizontal corrector used to make a deviation in the quad + Name of the horizontal corrector used to create horizontal orbit + offsets at the quadrupole. vcorr_name : str - Vertical corrector used to make a deviation in the quad + Name of the vertical corrector used to create vertical orbit offsets + at the quadrupole. quad_name : str - Quadrupole used to find the center + Name of the quadrupole to align. hcorr_delta : float - Horizontal corrector delta strength + Change in horizontal corrector strength used for each horizontal + orbit-offset step. vcorr_delta : float - Vertical corrector delta strength + Change in vertical corrector strength used for each vertical + orbit-offset step. hquad_delta : float - Quadrupole delta strength (for h search) + Change in quadrupole strength used during the horizontal alignment + measurement. vquad_delta : float - Quadrupole delta strength (for v search) - + Change in quadrupole strength used during the vertical alignment + measurement. + n_step : int, default=1 + Number of orbit-offset steps to perform in each plane. + sleep_between_step : float, default=0 + Time in seconds to wait after changing an orbit offset. + n_avg_meas : int, default=1 + Number of BPM measurements to average at each step. + sleep_between_meas : float, default=0 + Time in seconds to wait between individual BPM measurements. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - bpm_array_name: str - bpm_name: str - hcorr_name: str - vcorr_name: str - quad_name: str - hcorr_delta: float - vcorr_delta: float - hquad_delta: float - vquad_delta: float - - -class BBA(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + def __init__( + self, + name: str, + bpm_array_name: str, + bpm_name: str, + hcorr_name: str, + vcorr_name: str, + quad_name: str, + hcorr_delta: float, + vcorr_delta: float, + hquad_delta: float, + vquad_delta: float, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): + super().__init__(name) + self.bpm_array_name = bpm_array_name + self.bpm_name = bpm_name + self.hcorr_name = hcorr_name + self.vcorr_name = vcorr_name + self.quad_name = quad_name + self.hcorr_delta = hcorr_delta + self.vcorr_delta = vcorr_delta + self.hquad_delta = hquad_delta + self.vquad_delta = vquad_delta + self.n_step = n_step + self.sleep_between_step = sleep_between_step + self.n_avg_meas = n_avg_meas + self.sleep_between_meas = sleep_between_meas def measure( self, @@ -111,29 +145,29 @@ def measure( plane: str, optional Plane to perform ("H" or "V", None => both plane) """ - nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas + nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas element_holder = self._peer interface = pySCInterface( element_holder=element_holder, - bpm_array_name=self._cfg.bpm_array_name, + bpm_array_name=self.bpm_array_name, ) interface.set_wait_time = sleep_step interface.read_wait_time = sleep_meas - bpms_names = element_holder.bpms.get(self._cfg.bpm_array_name).names() + bpms_names = element_holder.bpms.get(self.bpm_array_name).names() bba_pySC_config = { - "number": bpms_names.index(self._cfg.bpm_name), - "QUAD": self._cfg.quad_name, - "HCORR": self._cfg.hcorr_name, - "VCORR": self._cfg.vcorr_name, - "HCORR_delta": self._cfg.hcorr_delta, - "QUAD_dk_H": self._cfg.hquad_delta, - "VCORR_delta": self._cfg.vcorr_delta, - "QUAD_dk_V": self._cfg.vquad_delta, + "number": bpms_names.index(self.bpm_name), + "QUAD": self.quad_name, + "HCORR": self.hcorr_name, + "VCORR": self.vcorr_name, + "HCORR_delta": self.hcorr_delta, + "QUAD_dk_H": self.hquad_delta, + "VCORR_delta": self.vcorr_delta, + "QUAD_dk_V": self.vquad_delta, "magnet_type": "normal_quadrupole", } @@ -142,10 +176,10 @@ def measure( generator = measure_bba( interface=interface, - bpm_name=self._cfg.bpm_name, + bpm_name=self.bpm_name, config=bba_pySC_config, shots_per_orbit=nb_meas, - n_corr_steps=self._cfg.n_step, + n_corr_steps=self.n_step, bipolar=False, skip_save=True, plane=plane, @@ -266,7 +300,7 @@ def plot_data(self, plane: str): axes["A"].plot(xx[final_mask], yy[final_mask], ".", c="C0") axes["A"].plot(xx[~final_mask], yy[~final_mask], ".", c="C1") - axes["A"].set_xlabel(f"BPM position [μm]\n{self._cfg.bpm_name} offset = {offset * 1e6:.3f} [μm]") + axes["A"].set_xlabel(f"BPM position [μm]\n{self.bpm_name} offset = {offset * 1e6:.3f} [μm]") axes["A"].set_ylabel("Modulation [μm]") axes["A"].grid() @@ -285,6 +319,6 @@ def plot_data(self, plane: str): axes["C"].grid() fig.tight_layout() - fig.canvas.manager.set_window_title(f"{plane} BBA {self._cfg.bpm_name}") + fig.canvas.manager.set_window_title(f"{plane} BBA {self.bpm_name}") plt.show() diff --git a/pyaml/tuning_tools/bba2.py b/pyaml/tuning_tools/bba2.py index 084c4900..8678740c 100644 --- a/pyaml/tuning_tools/bba2.py +++ b/pyaml/tuning_tools/bba2.py @@ -4,63 +4,17 @@ import matplotlib.pyplot as plt import numpy as np -from pydantic import ConfigDict from ..common.constants import Action from ..common.exception import PyAMLException -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool logger = logging.getLogger(__name__) PYAMLCLASS = "BBA2" -class ConfigModel(MeasurementToolConfigModel): - """ - Configuration model for Beam Based Alignment. - BBA finds the magnetic center of a quad (zero crossing). - - Parameters - ---------- - bpm_array_name : str - BPM array name (orbit) - bpm_name : str - BPM to be corrected (close to the quad) - hcorr_name : str - Horizontal corrector used to make a deviation in the quad - vcorr_name : str - Vertical corrector used to make a deviation in the quad - quad_name : str - Quadrupole used to find the center - tune_correction_name: str - Tune tuning tool - hcorr_delta : float - Horizontal corrector delta strength - vcorr_delta : float - Vertical corrector delta strength - quad_delta : float - Quadrupole delta strength - bipolar_delta : bool - Perform scan at initial_quad_strength +/- quad delta - minicyle_sleep_time : float - Time to wait for quad mini cycle - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - bpm_array_name: str - bpm_name: str - hcorr_name: str - vcorr_name: str - quad_name: str - tune_correction_name: str - hcorr_delta: float - vcorr_delta: float - quad_delta: float - bipolar_delta: bool = False - minicyle_sleep_time: float = 5 - - class BBAData: def __init__(self): self.k = [] # Fitted kick @@ -87,10 +41,95 @@ def update_offset(self, x, error, fit): self.lastfit = fit -class BBA2(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg +@register_schema +class BBA2(MeasurementTool, DynamicValidation): + """ + Beam-based alignment tool with tune correction. + + This tool determines the magnetic center of a quadrupole from the BPM + response to controlled horizontal and vertical orbit offsets. During the + measurement, the quadrupole strength is varied and tune correction is used + to compensate for the resulting tune change. + + Parameters + ---------- + name : str + Name of the measurement tool. + bpm_array_name : str + Name of the BPM array used to measure the orbit. + bpm_name : str + Name of the BPM located near the quadrupole whose center is measured. + hcorr_name : str + Name of the horizontal corrector used to create horizontal orbit + offsets at the quadrupole. + vcorr_name : str + Name of the vertical corrector used to create vertical orbit offsets + at the quadrupole. + quad_name : str + Name of the quadrupole to align. + tune_correction_name : str + Name of the tune-correction tool used to compensate for tune changes + caused by varying the quadrupole strength. + hcorr_delta : float + Change in horizontal corrector strength used for each horizontal + orbit-offset step. + vcorr_delta : float + Change in vertical corrector strength used for each vertical + orbit-offset step. + quad_delta : float + Change in quadrupole strength used for the alignment measurement. + bipolar_delta : bool, default=False + If `True`, vary the quadrupole strength both above and below its + initial value by `quad_delta`. If `False`, apply the change in one + direction only. + minicycle_sleep_time : float, default=5 + Time in seconds to wait for the quadrupole minicycle to complete after + changing its strength. + n_step : int, default=1 + Number of orbit-offset steps to perform in each plane. + sleep_between_step : float, default=0 + Time in seconds to wait after changing an orbit offset. + n_avg_meas : int, default=1 + Number of BPM measurements to average at each step. + sleep_between_meas : float, default=0 + Time in seconds to wait between individual BPM measurements. + """ + + def __init__( + self, + name: str, + bpm_array_name: str, + bpm_name: str, + hcorr_name: str, + vcorr_name: str, + quad_name: str, + tune_correction_name: str, + hcorr_delta: float, + vcorr_delta: float, + quad_delta: float, + bipolar_delta: bool = False, + minicycle_sleep_time: float = 5, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): + super().__init__(name) + self.bpm_array_name = bpm_array_name + self.bpm_name = bpm_name + self.hcorr_name = hcorr_name + self.vcorr_name = vcorr_name + self.quad_name = quad_name + self.tune_correction_name = tune_correction_name + self.hcorr_delta = hcorr_delta + self.vcorr_delta = vcorr_delta + self.quad_delta = quad_delta + self.bipolar_delta = bipolar_delta + self.minicycle_sleep_time = minicycle_sleep_time + self.n_step = n_step + self.sleep_between_step = sleep_between_step + self.n_avg_meas = n_avg_meas + self.sleep_between_meas = sleep_between_meas @staticmethod def _x_intercept(x, k, n): @@ -331,17 +370,17 @@ def measure( plane: str, optional Plane to perform ("H" or "V", None => both plane) """ - self._nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - self._sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - self._sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas + self._nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + self._sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + self._sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas # Device handles self.check_peer() - self._h_steer = self.peer.magnet.get(self._cfg.hcorr_name) - self._v_steer = self.peer.magnet.get(self._cfg.vcorr_name) - self._quad = self.peer.magnet.get(self._cfg.quad_name) - self._bpms = self.peer.bpms.get(self._cfg.bpm_array_name) - self._bpmi = self._bpms.names().index(self._cfg.bpm_name) + self._h_steer = self.peer.magnet.get(self.hcorr_name) + self._v_steer = self.peer.magnet.get(self.vcorr_name) + self._quad = self.peer.magnet.get(self.quad_name) + self._bpms = self.peer.bpms.get(self.bpm_array_name) + self._bpmi = self._bpms.names().index(self.bpm_name) # Initial values self._initial_k0 = [self._h_steer.strength.get(), self._v_steer.strength.get()] @@ -349,26 +388,26 @@ def measure( self._quad_polarity = np.sign(self._initial_k1) if self._initial_k1 != 0 else 1 self._ref_ios, fx, fy = self._init_responses( - self._cfg.tune_correction_name, - self._cfg.bpm_array_name, - self._cfg.quad_name, - self._cfg.hcorr_name, - self._cfg.vcorr_name, + self.tune_correction_name, + self.bpm_array_name, + self.quad_name, + self.hcorr_name, + self.vcorr_name, self._bpmi, ) - dk0h = self._cfg.hcorr_delta if plane is None or plane == "H" else 0 - dk0v = self._cfg.vcorr_delta if plane is None or plane == "V" else 0 - dk1 = self._cfg.quad_delta - bidelta = self._cfg.bipolar_delta + dk0h = self.hcorr_delta if plane is None or plane == "H" else 0 + dk0v = self.vcorr_delta if plane is None or plane == "V" else 0 + dk1 = self.quad_delta + bidelta = self.bipolar_delta doH = dk0h != 0 doV = dk0v != 0 aborted = False err = None - logger.debug(f"Initial H corrector {self._cfg.hcorr_name} value: {self._initial_k0[0]} rad") - logger.debug(f"Initial V corrector {self._cfg.vcorr_name} value: {self._initial_k0[1]} rad") - logger.debug(f"Initial quad {self._cfg.quad_name} value: {self._initial_k1} m-1") + logger.debug(f"Initial H corrector {self.hcorr_name} value: {self._initial_k0[0]} rad") + logger.debug(f"Initial V corrector {self.vcorr_name} value: {self._initial_k0[1]} rad") + logger.debug(f"Initial quad {self.quad_name} value: {self._initial_k1} m-1") try: self._register_callback(callback) @@ -379,16 +418,16 @@ def measure( Y = BBAData() # Mini cycle - if self._cfg.minicyle_sleep_time > 0: - logger.debug(f"Quad mini cycling {self._cfg.quad_name}") + if self.minicycle_sleep_time > 0: + logger.debug(f"Quad mini cycling {self.quad_name}") _str = self._initial_k1 + dk1 * self._quad_polarity self._quad.strength.set(_str) self.send_callback(Action.APPLY, {"step": -1, "magnet": self._quad.name, "strength": _str}) - time.sleep(self._cfg.minicyle_sleep_time) + time.sleep(self.minicycle_sleep_time) _str = self._initial_k1 self._quad.strength.set(_str) self.send_callback(Action.APPLY, {"step": -1, "magnet": self._quad.name, "strength": _str}) - time.sleep(self._cfg.minicyle_sleep_time) + time.sleep(self.minicycle_sleep_time) opt_found = False self._step = 0 @@ -541,6 +580,6 @@ def plot_data(self): self.plot_plane_data(ax, "VData") fig.tight_layout() - fig.canvas.manager.set_window_title(f"BBA {self._cfg.bpm_name}") + fig.canvas.manager.set_window_title(f"BBA {self.bpm_name}") plt.show() diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index 002ef7b4..92a59d05 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -2,12 +2,10 @@ import logging from abc import ABCMeta, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING, Callable, Optional, Self - -from pydantic import ConfigDict +from typing import TYPE_CHECKING, Callable, Self from ..common.constants import Action -from ..common.element import Element, ElementConfigModel +from ..common.element import Element from ..common.exception import PyAMLException if TYPE_CHECKING: @@ -16,15 +14,6 @@ logger = logging.getLogger(__name__) -class MeasurementToolConfigModel(ElementConfigModel): - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - n_step: int = 10 - sleep_between_step: float = 0 - n_avg_meas: int = 1 - sleep_between_meas: float = 0 - - class MeasurementTool(Element, metaclass=ABCMeta): """ Base class for measurement tool such as reponse matrix measurement or other scans. @@ -148,11 +137,8 @@ def _register_callback(self, callback: Callable): self._callback = callback def attach(self, peer: "ElementHolder") -> Self: - if hasattr(self, "_cfg"): - obj = self.__class__(self._cfg) - else: - obj = copy.copy(self) - obj._after_attach() + obj = copy.copy(self) + obj._after_attach() obj._peer = peer return obj diff --git a/tests/tuning_tools/test_bba.py b/tests/tuning_tools/test_bba.py index aaeb6a75..330c8118 100644 --- a/tests/tuning_tools/test_bba.py +++ b/tests/tuning_tools/test_bba.py @@ -19,7 +19,7 @@ def test_bba(): # BBA (model dependant method) bba = SR.get_bba("BBA2-BPM_C04-04") - bba._cfg.minicyle_sleep_time = 0 + bba.minicycle_sleep_time = 0 bba.measure() assert np.abs(bba.h_offset() - 20e-6) < 1e-6 assert np.abs(bba.v_offset() + 15e-6) < 1e-6