diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index 5016bf2b..e75b282d 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING from .. import PyAMLException -from ..common.element import ElementConfigModel +from ..validation import DynamicValidation, register_schema from .chromaticity_monitor import ChomaticityMonitor from .response_matrix_data import ResponseMatrixData from .tuning_tool import TuningTool @@ -10,10 +10,6 @@ if TYPE_CHECKING: from ..arrays.magnet_array import MagnetArray -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier import logging import time @@ -25,55 +21,47 @@ PYAMLCLASS = "Chromaticity" -class ConfigModel(ElementConfigModel): - """ - Configuration model for Tune - - Parameters - ---------- - sextu_array_name : str - Array name of sextu used to adjust the chromaticity - chromaticty_monitor_name : str - Name of the diagnostic pyaml device for measuring the chromaticity - response_matrix : str | ResponseMatrixData - filename or data of the chromaticity response matrix - """ - - sextu_array_name: str - chromaticty_monitor_name: str - response_matrix: str | ResponseMatrixData - - -class Chromaticity(TuningTool): +@register_schema +class Chromaticity(TuningTool, DynamicValidation): """ Class providing chromaticity adjustment tool """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, name: str, sextu_array_name: str, chromaticity_monitor_name: str, response_matrix: str | ResponseMatrixData + ): """ - Construct a chromaticity adjustment object. + Initialize a chromaticity adjustment tool. Parameters ---------- - cfg : ConfigModel - Configuration for the tune adjustment. - """ - super().__init__(cfg.name) - self._cfg = cfg + name : str + Name of the tuning tool. + sextu_array_name : str + Name of the sextupole array used to adjust the chromaticity. + chromaticity_monitor_name : str + Name of the chromaticity monitor used for readback. + response_matrix : str | ResponseMatrixData + Chromaticity response matrix or path to a saved response matrix file. + """ + super().__init__(name) + self.sextu_array_name = sextu_array_name + self._chromaticity_monitor_name = chromaticity_monitor_name + self.response_matrix_file = response_matrix self._response_matrix = None self._correctionmat = None # If the configuration response matrix is a filename, load it - if type(cfg.response_matrix) is str: + if type(self.response_matrix_file) is str: try: - cfg.response_matrix = ResponseMatrixData.load(cfg.response_matrix) + self._response_matrix = ResponseMatrixData.load(self.response_matrix_file) except Exception as e: logger.warning(f"{str(e)}") - cfg.response_matrix = None + self._response_matrix = None # Invert matrix - if cfg.response_matrix: - self._response_matrix = np.array(cfg.response_matrix._cfg.matrix) + if self._response_matrix: + self._response_matrix = np.array(self._response_matrix._cfg.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -84,7 +72,7 @@ def response_matrix(self) -> ResponseMatrixData | None: """ Return the response matrix if it has been loaded None otherwise """ - return self._cfg.response_matrix + return self._response_matrix def load(self, load_path: Path): """ @@ -95,19 +83,19 @@ def load(self, load_path: Path): load_path : Path Filename of the :class:`~.ResponseMatrixData` to load """ - self._cfg.response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._cfg.response_matrix._cfg.matrix) + self._response_matrix = ResponseMatrixData.load(load_path) + self._response_matrix = np.array(self._response_matrix._cfg.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property def _cm(self) -> "ChomaticityMonitor": self.check_peer() - return self.peer.get_chromaticity_monitor(self._cfg.chromaticty_monitor_name) + return self.peer.get_chromaticity_monitor(self._chromaticity_monitor_name) @property def _sextu(self) -> "MagnetArray": self.check_peer() - return self.peer.get_magnets(self._cfg.sextu_array_name) + return self.peer.get_magnets(self.sextu_array_name) def get(self): """ @@ -117,7 +105,7 @@ def get(self): def readback(self): """ - Launch a chromaticty scan and returns the measured chromaticity. + Launch a chromaticity scan and returns the measured chromaticity. """ self._cm.measure() return self._cm.chromaticity.get() diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index 4c4ea1cc..2d4c24d5 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -1,64 +1,22 @@ -from ..common.abstract import ReadFloatArray -from ..common.constants import Action -from ..common.element import ElementConfigModel -from ..common.exception import PyAMLException -from ..tuning_tools.measurement_tool import MeasurementTool, MeasurementToolConfigModel - -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier import logging +from collections.abc import Callable from time import sleep import matplotlib.pyplot as plt import numpy as np -from pydantic import ConfigDict + +from ..common.abstract import ReadFloatArray +from ..common.constants import Action +from ..common.element import __pyaml_repr__ +from ..common.exception import PyAMLException +from ..tuning_tools.measurement_tool import MeasurementTool +from ..validation import DynamicValidation, register_schema logger = logging.getLogger(__name__) PYAMLCLASS = "ChomaticityMonitor" -class ConfigModel(MeasurementToolConfigModel): - """ - Configuration model for Chromaticity Monitor. - - Parameters - ---------- - betatron_tune_name : str - Name of the diagnostic pyaml device for measuring the tune - rf_plant_name : str - Name of main RF frequency plant - bpm_array_name : str,optional - Name of main BPM array used for dispersion fit - e_delta : float, optional - Default variation of relative energy during chromaticity measurement: - f0 - f0 * E_delta * alphac < f_RF < f0 + f0 * E_delta * alphac, - by default 0.001 - max_e_delta : float, optional - Maximum authorized variation of relative energy during chromaticity - measurement, by default 0.004 - fit_order : int, optional - Chomaticity fitting order, by default 1 - fit_disp_order : int, optional - Dispersion fitting order, by default 1 - fit_dispersion : bool, optional - Dispersion fitting, by default False - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - betatron_tune_name: str - rf_plant_name: str - bpm_array_name: str | None = None - e_delta: float = 0.001 - max_e_delta: float = 0.004 - fit_order: int = 1 - fit_disp_order: int = 1 - fit_dispersion: bool = False - - class RChromaDispArray(ReadFloatArray): """ Class providing read access to chromaticity or dispersion. @@ -81,27 +39,89 @@ def unit(self) -> str: return self.unit -class ChomaticityMonitor(MeasurementTool): +@register_schema +class ChomaticityMonitor(MeasurementTool, DynamicValidation): """ Class providing access to a chromaticity monitor of a physical or simulated lattice. The monitor provides horizontal and vertical chromaticity measurements. """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + betatron_tune_name: str, + rf_plant_name: str, + bpm_array_name: str | None = None, + e_delta: float = 0.001, + max_e_delta: float = 0.004, + fit_order: int = 1, + fit_disp_order: int = 1, + fit_dispersion: bool = False, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): """ - Construct a ChomaticityMonitor. + Initialize a chromaticity monitor. + + The monitor performs chromaticity measurements by varying the RF + frequency, acquiring tune measurements from a betatron tune monitor, + and optionally fitting the machine dispersion using BPM orbit data. Parameters ---------- - cfg : ConfigModel - Configuration for the ChromaticityMonitor, including betatron - tune monitor, RF plant, and defaults parameters. + name : str + Name of the chromaticity monitor. + betatron_tune_name : str + Name of the betatron tune monitor used to measure the horizontal + and vertical tunes. + rf_plant_name : str + Name of the RF plant used to vary the RF frequency. + bpm_array_name : str, optional + Name of the BPM array used for dispersion measurements. Required + only when dispersion fitting is enabled. + e_delta : float, optional + Default relative momentum deviation used during the measurement. + max_e_delta : float, optional + Maximum permitted relative momentum deviation. + fit_order : int, optional + Polynomial order used to fit the chromaticity. + fit_disp_order : int, optional + Polynomial order used to fit the dispersion. + fit_dispersion : bool, optional + Whether to fit the machine dispersion in addition to the + chromaticity. + n_step : int, optional + Number of RF frequency steps. + sleep_between_step : float, optional + Delay in seconds after changing the RF frequency. + n_avg_meas : int, optional + Number of tune (and orbit) measurements to average at each RF + frequency. + sleep_between_meas : float, optional + Delay in seconds between consecutive measurements during + averaging. """ - super().__init__(cfg.name) - self._cfg = cfg + + super().__init__(name) + + self.betatron_tune_name = betatron_tune_name + self.rf_plant_name = rf_plant_name + self.bpm_array_name = bpm_array_name + self.e_delta = e_delta + self.max_e_delta = max_e_delta + self.fit_order = fit_order + self.fit_disp_order = fit_disp_order + self.fit_dispersion = fit_dispersion + 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 + self._chromaticity = RChromaDispArray(self, "chromaticity", "1") - self._dipsersion = RChromaDispArray(self, "dispersion", "m") + self._dispersion = RChromaDispArray(self, "dispersion", "m") self._alphac = None @property @@ -129,22 +149,22 @@ def dispersion(self) -> ReadFloatArray: ReadFloatArray Array of dispersion values [[dx, dy],[d'x, d'y],...] """ - return self._dipsersion + return self._dispersion def measure( self, - n_step: int = None, - alphac: float = None, - e_delta: float = None, - max_e_delta: float = None, - n_avg_meas: int = None, - sleep_between_meas: float = None, - sleep_between_step: float = None, - fit_order: int = None, - fit_disp_order: int = None, + n_step: int | None = None, + alphac: float | None = None, + e_delta: float | None = None, + max_e_delta: float | None = None, + n_avg_meas: int | None = None, + sleep_between_meas: float | None = None, + sleep_between_step: float | None = None, + fit_order: int | None = None, + fit_disp_order: int | None = None, fit_dispersion: bool | None = None, - do_plot: bool = None, - callback: callable = None, + do_plot: bool | None = None, + callback: Callable | None = None, ): """ Main function for chromaticity measurment. @@ -202,16 +222,16 @@ def measure( dtune:np.array # The tune variation (on Action.RESTORE) """ - n_step = n_step if n_step is not None else self._cfg.n_step + n_step = n_step if n_step is not None else self.n_step alphac = alphac if alphac is not None else self._alphac - e_delta = e_delta if e_delta is not None else self._cfg.e_delta - max_e_delta = max_e_delta if max_e_delta is not None else self._cfg.max_e_delta - n_avg_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_between_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas - sleep_between_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - fit_order = fit_order if fit_order is not None else self._cfg.fit_order - fit_disp_order = fit_disp_order if fit_disp_order is not None else self._cfg.fit_disp_order - fit_dispersion = fit_dispersion if fit_dispersion is not None else self._cfg.fit_dispersion + e_delta = e_delta if e_delta is not None else self.e_delta + max_e_delta = max_e_delta if max_e_delta is not None else self.max_e_delta + n_avg_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + sleep_between_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas + sleep_between_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + fit_order = fit_order if fit_order is not None else self.fit_order + fit_disp_order = fit_disp_order if fit_disp_order is not None else self.fit_disp_order + fit_dispersion = fit_dispersion if fit_dispersion is not None else self.fit_dispersion if abs(e_delta) > abs(max_e_delta): logger.warning(f"e_delta={e_delta} is greater than max_e_delta={max_e_delta}") @@ -224,14 +244,14 @@ def measure( # Get devices self.check_peer() - tm = self.peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) - rf = self.peer.get_rf_plant(self._cfg.rf_plant_name) + tm = self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + rf = self.peer.get_rf_plant(self.rf_plant_name) bpms = None n_bpm = 0 orbit = None - if fit_dispersion and fit_disp_order is not None and self._cfg.bpm_array_name is not None: + if fit_dispersion and fit_disp_order is not None and self.bpm_array_name is not None: # For dispersion fit - bpms = self.peer.get_bpms(self._cfg.bpm_array_name) + bpms = self.peer.get_bpms(self.bpm_array_name) n_bpm = len(bpms) f0 = rf.frequency.get() @@ -365,3 +385,10 @@ def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False): fig.tight_layout() plt.show() + + def _after_attach(self): + self._chromaticity = RChromaDispArray(self, "chromaticity", "1") + self._dispersion = RChromaDispArray(self, "dispersion", "m") + + def __repr__(self): + return __pyaml_repr__(self, exclude=["chromaticity", "dispersion"]) diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index 8b995fbf..ab0e23d7 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -3,11 +3,10 @@ from typing import Callable, Optional import numpy as np -from pydantic import ConfigDict from ..common.constants import Action -from ..common.element import ElementConfigModel -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel logger = logging.getLogger(__name__) @@ -15,31 +14,78 @@ PYAMLCLASS = "ChromaticityResponseMatrix" -class ConfigModel(MeasurementToolConfigModel): +@register_schema +class ChromaticityResponseMatrix(MeasurementTool, DynamicValidation): """ - Configuration model for Tune response matrix + Measure the chromaticity response matrix of a lattice. + + This tool perturbs each sextupole in a sextupole array and measures the + resulting chromaticity variation using a chromaticity monitor. The measured + slopes are assembled into a response matrix with chromaticity components as + the observables and sextupoles as the variables. Parameters ---------- + name : str + Name of the response matrix measurement tool. sextu_array_name : str - Array name of sextupole used to adjust the chromaticity + Name of the sextupole array to excite. chromaticity_name : str - Name of the diagnostic chromaticy monitor + Name of the chromaticity monitor used to measure the response. sextu_delta : float - Delta strength used to get the response matrix + Default sextupole excitation applied during the measurement. + n_step : int, optional + Default number of excitation steps used to fit the response. + sleep_between_step : float, optional + Default delay in seconds after changing the sextupole strength. + n_avg_meas : int, optional + Default number of chromaticity measurements to average at each step. + sleep_between_meas : float, optional + Default delay in seconds between averaged measurements. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - sextu_array_name: str - chromaticity_name: str - sextu_delta: float + def __init__( + self, + name: str, + sextu_array_name: str, + chromaticity_name: str, + sextu_delta: float, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): + """ + Initialize the chromaticity response matrix measurement tool. + Parameters + ---------- + name : str + Name of the response matrix measurement tool. + sextu_array_name : str + Name of the sextupole array to excite. + chromaticity_name : str + Name of the chromaticity monitor used to measure the response. + sextu_delta : float + Default sextupole excitation applied during the measurement. + n_step : int, optional + Default number of excitation steps used to fit the response. + sleep_between_step : float, optional + Default delay in seconds after changing the sextupole strength. + n_avg_meas : int, optional + Default number of chromaticity measurements to average at each step. + sleep_between_meas : float, optional + Default delay in seconds between averaged measurements. + """ -class ChromaticityResponseMatrix(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + super().__init__(name) + self.sextu_array_name = sextu_array_name + self.chromaticity_name = chromaticity_name + self.sextu_delta = sextu_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 self.aborted = False def measure( @@ -117,8 +163,8 @@ def callback(action: Action, data:dict): """ # Get devices self.check_peer() - sextus = self._peer.get_magnets(self._cfg.sextu_array_name) - cm = self._peer.get_chromaticity_monitor(self._cfg.chromaticity_name) + sextus = self._peer.get_magnets(self.sextu_array_name) + cm = self._peer.get_chromaticity_monitor(self.chromaticity_name) self._register_callback(callback) self._init_measure("pyaml.tuning_tools.response_matrix_data") @@ -131,11 +177,11 @@ def callback(action: Action, data:dict): return False initial_chroma = cm.chromaticity.get() - delta = sextu_delta if sextu_delta is not None else self._cfg.sextu_delta - nb_step = n_step if n_step is not None else self._cfg.n_step - 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 + delta = sextu_delta if sextu_delta is not None else self.sextu_delta + nb_step = n_step if n_step is not None else self.n_step + 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 err = None aborted = False diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index edec89e4..c1cbeba6 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -1,3 +1,4 @@ +import copy import logging from abc import ABCMeta, abstractmethod from pathlib import Path @@ -166,10 +167,14 @@ def _register_callback(self, callback: Callable): self._callback = callback def attach(self, peer: "ElementHolder") -> Self: - """ - Create a new reference to attach this measurement tool object to a simulator - or a control system. - """ - obj = self.__class__(self._cfg) + if hasattr(self, "_cfg"): + obj = self.__class__(self._cfg) + else: + obj = copy.copy(self) + obj._after_attach() obj._peer = peer return obj + + def _after_attach(self) -> None: + """Hook for subclasses to rebind internal references after attach.""" + pass diff --git a/pyaml/tuning_tools/tuning_tool.py b/pyaml/tuning_tools/tuning_tool.py index 8f043bc4..f24f42c9 100644 --- a/pyaml/tuning_tools/tuning_tool.py +++ b/pyaml/tuning_tools/tuning_tool.py @@ -1,3 +1,4 @@ +import copy from typing import Self from ..common.element import Element @@ -17,6 +18,14 @@ def attach(self, peer: "ElementHolder") -> Self: Create a new reference to attach this tuning tool object to a simulator or a control system. """ - obj = self.__class__(self._cfg) + if hasattr(self, "_cfg"): + obj = self.__class__(self._cfg) + else: + obj = copy.copy(self) + obj._after_attach() obj._peer = peer return obj + + def _after_attach(self) -> None: + """Hook for subclasses to rebind internal references after attach.""" + pass diff --git a/tests/config/EBSOrbit.yaml b/tests/config/EBSOrbit.yaml index 9b66318a..ce29b43f 100644 --- a/tests/config/EBSOrbit.yaml +++ b/tests/config/EBSOrbit.yaml @@ -53,7 +53,7 @@ devices: sextu_delta: 1e-6 - type: pyaml.tuning_tools.chromaticity name: DEFAULT_CHROMATICITY_CORRECTION - chromaticty_monitor_name: CHROMATICITY_MONITOR + chromaticity_monitor_name: CHROMATICITY_MONITOR sextu_array_name: Sext response_matrix: ${path:ideal_chroma_resp.json} - type: pyaml.tuning_tools.orbit