diff --git a/pyaml/common/element_holder.py b/pyaml/common/element_holder.py index a06a0693..e5900068 100644 --- a/pyaml/common/element_holder.py +++ b/pyaml/common/element_holder.py @@ -21,7 +21,7 @@ from ..magnet.serialized_magnet import SerializedMagnets from ..rf.rf_plant import RFPlant from ..rf.rf_transmitter import RFTransmitter -from ..tuning_tools.chromaticity_monitor import ChomaticityMonitor +from ..tuning_tools.chromaticity_monitor import ChromaticityMonitor from .element import Element if TYPE_CHECKING: @@ -293,7 +293,7 @@ def add_tool(self, tool: Element): # ---- Chromaticity ------------------------------------------------- - def get_chromaticity_monitor(self, name: str) -> ChomaticityMonitor: + def get_chromaticity_monitor(self, name: str) -> ChromaticityMonitor: obj = self.__get("Chomaticity monitor", name, self.__TUNING_TOOLS) return obj diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index 5016bf2b..0afbfb0e 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -2,18 +2,14 @@ from typing import TYPE_CHECKING from .. import PyAMLException -from ..common.element import ElementConfigModel -from .chromaticity_monitor import ChomaticityMonitor +from ..validation import DynamicValidation, register_schema +from .chromaticity_monitor import ChromaticityMonitor from .response_matrix_data import ResponseMatrixData from .tuning_tool import TuningTool 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.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.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property - def _cm(self) -> "ChomaticityMonitor": + def _cm(self) -> "ChromaticityMonitor": 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..fc0e02d7 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -1,62 +1,20 @@ -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 - -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 - """ +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 - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") +logger = logging.getLogger(__name__) - 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 +PYAMLCLASS = "ChromaticityMonitor" class RChromaDispArray(ReadFloatArray): @@ -65,7 +23,7 @@ class RChromaDispArray(ReadFloatArray): Returns arrays of shape (fit_order,2) or None """ - def __init__(self, parent: "ChomaticityMonitor", name: str, unit: str): + def __init__(self, parent: "ChromaticityMonitor", name: str, unit: str): self._parent = parent self._name = name self._unit = unit @@ -81,27 +39,89 @@ def unit(self) -> str: return self.unit -class ChomaticityMonitor(MeasurementTool): +@register_schema +class ChromaticityMonitor(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..c8305bb7 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -1,45 +1,92 @@ import logging import time +from dataclasses import asdict 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 .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool +from .response_matrix_data import ResponseMatrixData logger = logging.getLogger(__name__) 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 +164,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 +178,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 @@ -206,12 +253,12 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - mat = ResponseMatrixDataConfigModel( + mat = ResponseMatrixData( matrix=chromamat.T.tolist(), variable_names=sextus.names(), observable_names=[cm.get_name() + ".x", cm.get_name() + ".y"], ) - self.latest_measurement.update(mat.model_dump()) + self.latest_measurement.update(asdict(mat)) self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data" return True diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index d62dbefe..176d3627 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -1,14 +1,12 @@ import logging -from typing import Callable, Optional, Self +from typing import Callable, Optional -from pydantic import ConfigDict from pySC.apps import measure_dispersion from pySC.apps.codes import DispersionCode from ..common.constants import Action -from ..common.element import ElementConfigModel -from ..common.element_holder import ElementHolder from ..external.pySC_interface import pySCInterface +from ..validation import DynamicValidation, register_schema from .measurement_tool import MeasurementTool logger = logging.getLogger(__name__) @@ -16,35 +14,41 @@ PYAMLCLASS = "Dispersion" -class ConfigModel(ElementConfigModel): - """ - Configuration model for dispersion measurement +@register_schema +class Dispersion(MeasurementTool, DynamicValidation): + """Measure beam dispersion by changing the RF frequency. + + The measurement uses a :class:`pySCInterface` to change the frequency of + an RF plant and acquire orbit data from a BPM array. Progress is reported + through the callback mechanism provided by :class:`MeasurementTool`. Parameters ---------- + name : str + Name of the dispersion measurement tool. bpm_array_name : str - BPM array name + Name of the BPM array used to measure the orbit. rf_plant_name : str - RF plant name + Name of the RF plant whose frequency is varied. frequency_delta : float - Frequency delta for measurement - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - bpm_array_name: str - rf_plant_name: str - frequency_delta: float + RF-frequency change applied during the measurement. + Attributes + ---------- + bpm_array_name : str + Name of the BPM array used for the measurement. + rf_plant_name : str + Name of the RF plant used for the measurement. + frequency_delta : float + RF-frequency change applied during the measurement. + """ -class Dispersion(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + def __init__(self, name: str, bpm_array_name: str, rf_plant_name: str, frequency_delta: float): + super().__init__(name) - self.bpm_array_name = cfg.bpm_array_name - self.rf_plant_name = cfg.rf_plant_name - self.frequency_delta = cfg.frequency_delta + self.bpm_array_name = bpm_array_name + self.rf_plant_name = rf_plant_name + self.frequency_delta = frequency_delta def measure( self, diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index edec89e4..c1e7beba 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 @@ -15,34 +16,6 @@ logger = logging.getLogger(__name__) -class MeasurementToolConfigModel(ElementConfigModel): - """ - Measurement tool configuration model - - Parameters - ---------- - n_step: int, optional - Number of measurement step [-delta/n_step..delta/n_step] - Default 1 - sleep_between_step: float, optional - Default sleep time after an actuator excitation - Default: 0 - n_avg_meas : int, optional - Default number of measurement per step used for averaging - Default 1 - sleep_between_meas: float, optional - Default sleep time between two measurments - Default: 0 - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - n_step: Optional[int] = 1 - sleep_between_step: Optional[float] = 0 - n_avg_meas: Optional[int] = 1 - sleep_between_meas: Optional[float] = 0 - - class MeasurementTool(Element, metaclass=ABCMeta): """ Base class for measurement tool such as reponse matrix measurement or other scans. @@ -166,10 +139,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/orbit.py b/pyaml/tuning_tools/orbit.py index 6baa1f88..67151e77 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -1,25 +1,17 @@ import logging +from dataclasses import asdict from pathlib import Path -from typing import TYPE_CHECKING, Literal, Optional, Union - -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier +from typing import Literal, Optional, Union import numpy as np -from pydantic import ConfigDict - -if TYPE_CHECKING: - from ..common.element_holder import ElementHolder from pySC import ResponseMatrix as pySC_ResponseMatrix from pySC.apps import orbit_correction from ..arrays.magnet_array import MagnetArray -from ..common.element import Element, ElementConfigModel from ..common.exception import PyAMLException from ..external.pySC_interface import pySCInterface from ..rf.rf_plant import RFPlant +from ..validation import DynamicValidation, register_schema from .orbit_response_matrix_data import OrbitResponseMatrixData from .tuning_tool import TuningTool @@ -29,57 +21,56 @@ PYAMLCLASS = "Orbit" -class ConfigModel(ElementConfigModel): - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - bpm_array_name: str - hcorr_array_name: str - vcorr_array_name: str - rf_plant_name: Optional[str] = None - singular_values: Optional[int] = None - singular_values_H: Optional[int] = None - singular_values_V: Optional[int] = None - virtual_target: float = 0 - response_matrix: Union[str, OrbitResponseMatrixData] - +@register_schema +class Orbit(TuningTool, DynamicValidation): + def __init__( + self, + name: str, + bpm_array_name: str, + hcorr_array_name: str, + vcorr_array_name: str, + response_matrix: Union[str, OrbitResponseMatrixData], + rf_plant_name: Optional[str] = None, + singular_values: Optional[int] = None, + singular_values_H: Optional[int] = None, + singular_values_V: Optional[int] = None, + virtual_target: float = 0, + ): + super().__init__(name) -class Orbit(TuningTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg - self.bpm_array_name = cfg.bpm_array_name - self.hcorr_array_name = cfg.hcorr_array_name - self.vcorr_array_name = cfg.vcorr_array_name + self.bpm_array_name = bpm_array_name + self.hcorr_array_name = hcorr_array_name + self.vcorr_array_name = vcorr_array_name self._pySC_response_matrix = None + self.rf_plant_name = rf_plant_name + self.virtual_target = virtual_target - self.virtual_target = cfg.virtual_target - - if cfg.singular_values is None: - if cfg.singular_values_H is None or cfg.singular_values_V is None: + if singular_values is None: + if singular_values_H is None or singular_values_V is None: raise PyAMLException( "Either `singular_values` or `singular_values_H` and `singular_values_V` must be provided." ) - self.singular_values_H = cfg.singular_values_H - self.singular_values_V = cfg.singular_values_V + self.singular_values_H = singular_values_H + self.singular_values_V = singular_values_V else: - if cfg.singular_values_H is not None or cfg.singular_values_V is not None: + if singular_values_H is not None or singular_values_V is not None: raise PyAMLException( "Either `singular_values` or `singular_values_H` and `singular_values_V` must be provided, not both." ) - self.singular_values_H = cfg.singular_values - self.singular_values_V = cfg.singular_values + self.singular_values_H = singular_values + self.singular_values_V = singular_values # If the configuration response matrix is a filename, load it - if type(cfg.response_matrix) is str: + if type(response_matrix) is str: try: - cfg.response_matrix = OrbitResponseMatrixData.load(cfg.response_matrix) + self._response_matrix = OrbitResponseMatrixData.load(response_matrix) except Exception as e: - logger.warning(f"Loading {cfg.response_matrix} failed {str(e)}") - cfg.response_matrix = None + logger.warning(f"Loading {response_matrix} failed {str(e)}") + self._response_matrix = None # Converts to self._pySC_response_matrix - if cfg.response_matrix: - self._set_response_matrix(cfg.response_matrix) + if self._response_matrix: + self._set_response_matrix(self._response_matrix) self._hcorr: MagnetArray = None self._vcorr: MagnetArray = None @@ -95,16 +86,16 @@ def load(self, load_path: Path): load_path : Path Filename of the :class:`~.OrbitResponseMatrixData` to load """ - self._cfg.response_matrix = OrbitResponseMatrixData.load(load_path) - self._set_response_matrix(self._cfg.response_matrix) + self._response_matrix = OrbitResponseMatrixData.load(load_path) + self._set_response_matrix(self.response_matrix) def _set_response_matrix(self, mat): - m = mat._cfg.model_dump() + m = asdict(mat) m["input_names"] = m.pop("variable_names") m["output_names"] = m.pop("observable_names") m["input_planes"] = m.pop("variable_planes") m["output_planes"] = m.pop("observable_planes") - self._cfg.response_matrix = mat + self._response_matrix = mat self._pySC_response_matrix = pySC_ResponseMatrix.model_validate(m) @property @@ -112,7 +103,7 @@ def response_matrix(self) -> OrbitResponseMatrixData | None: """ Return the response matrix if it has been loaded None otherwise """ - return self._cfg.response_matrix + return self._response_matrix def correct( self, @@ -312,11 +303,11 @@ def get_rf_weight(self) -> float: return self._pySC_response_matrix.rf_weight def post_init(self): - self._hcorr = self.peer.get_magnets(self._cfg.hcorr_array_name) - self._vcorr = self.peer.get_magnets(self._cfg.vcorr_array_name) + self._hcorr = self.peer.get_magnets(self.hcorr_array_name) + self._vcorr = self.peer.get_magnets(self.vcorr_array_name) hvElts = [] hvElts.extend(self._hcorr) hvElts.extend(self._vcorr) self._hvcorr = MagnetArray("", hvElts) - if self._cfg.rf_plant_name is not None: - self._rf_plant = self.peer.get_rf_plant(self._cfg.rf_plant_name) + if self.rf_plant_name is not None: + self._rf_plant = self.peer.get_rf_plant(self.rf_plant_name) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 68d26fba..cf90bd7c 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -1,55 +1,102 @@ import logging -from pathlib import Path -from typing import Callable, List, Optional, Self +from dataclasses import asdict +from typing import Callable, List, Optional import pySC -from pydantic import ConfigDict from pySC.apps import measure_ORM from pySC.apps.codes import ResponseCode from ..common.constants import Action from ..external.pySC_interface import pySCInterface -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .orbit_response_matrix_data import ConfigModel as OrbitResponseMatrixDataConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool +from .orbit_response_matrix_data import OrbitResponseMatrixData logger = logging.getLogger(__name__) PYAMLCLASS = "OrbitResponseMatrix" -class ConfigModel(MeasurementToolConfigModel): - """ - Configuration model for orbit response matrix measurement +@register_schema +class OrbitResponseMatrix(MeasurementTool, DynamicValidation): + """Measure an orbit response matrix using BPMs and orbit correctors. + + The orbit response matrix describes the change in measured beam position + produced by a change in corrector strength. This measurement tool uses + horizontal and vertical corrector arrays together with a BPM array to + perform the measurement through the pySC interface. + + After a successful measurement, the result is converted to + :class:`OrbitResponseMatrixData` and stored in ``latest_measurement``. + The stored data includes the response matrix, corrector and BPM names, + and the plane associated with each variable and observable. Parameters ---------- + name : str + Name of the measurement tool. bpm_array_name : str - BPM array name + Name of the BPM array used to measure the orbit. hcorr_array_name : str - Horizontal corrector array name + Name of the horizontal corrector array. vcorr_array_name : str - Vertical corrector array name + Name of the vertical corrector array. corrector_delta : float - Corrector delta for measurement + Change in corrector strength applied during the measurement. + n_step : int, optional + Number of strength steps used for each corrector. The default is 1. + sleep_between_step : float, optional + Time in seconds to wait after changing a corrector strength. The + default is 0. + n_avg_meas : int, optional + Number of orbit measurements averaged at each corrector setting. The + default is 1. + sleep_between_meas : float, optional + Time in seconds to wait between orbit measurements used for averaging. + The default is 0. + + Attributes + ---------- + bpm_array_name : str + Name of the configured BPM array. + hcorr_array_name : str + Name of the configured horizontal corrector array. + vcorr_array_name : str + Name of the configured vertical corrector array. + corrector_delta : float + Corrector-strength change used for the measurement. + n_step : int + Configured number of corrector-strength steps. + sleep_between_step : float + Configured delay between corrector-strength changes. + n_avg_meas : int + Configured number of orbit measurements to average. + sleep_between_meas : float + Configured delay between averaged orbit measurements. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - bpm_array_name: str - hcorr_array_name: str - vcorr_array_name: str - corrector_delta: float - - -class OrbitResponseMatrix(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + def __init__( + self, + name: str, + bpm_array_name: str, + hcorr_array_name: str, + vcorr_array_name: str, + corrector_delta: float, + n_step: Optional[int] = 1, + sleep_between_step: Optional[float] = 0, + n_avg_meas: Optional[int] = 1, + sleep_between_meas: Optional[float] = 0, + ): + super().__init__(name) - self.bpm_array_name = cfg.bpm_array_name - self.hcorr_array_name = cfg.hcorr_array_name - self.vcorr_array_name = cfg.vcorr_array_name - self.corrector_delta = cfg.corrector_delta + self.bpm_array_name = bpm_array_name + self.hcorr_array_name = hcorr_array_name + self.vcorr_array_name = vcorr_array_name + self.corrector_delta = corrector_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, @@ -91,9 +138,9 @@ def measure( reading. If the callback returns false, then the process is aborted. """ - 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( @@ -155,11 +202,11 @@ def measure( return False orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump()) - self.latest_measurement.update(orm_data.model_dump()) + self.latest_measurement.update(asdict(orm_data)) return True - def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel: + def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixData: # all metadata is discarded here. Should we keep something? element_holder = self._peer @@ -187,5 +234,5 @@ def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataC "observable_planes": observable_planes, } - orm_data = OrbitResponseMatrixDataConfigModel(**orm_data_model) + orm_data = OrbitResponseMatrixData(**orm_data_model) return orm_data diff --git a/pyaml/tuning_tools/orbit_response_matrix_data.py b/pyaml/tuning_tools/orbit_response_matrix_data.py index 6baa9d10..3993b414 100644 --- a/pyaml/tuning_tools/orbit_response_matrix_data.py +++ b/pyaml/tuning_tools/orbit_response_matrix_data.py @@ -1,34 +1,41 @@ +from dataclasses import dataclass from typing import Optional -from .response_matrix_data import ConfigModel as ReponseMatrixDataConfigModel +from ..validation import DynamicValidation, register_schema from .response_matrix_data import ResponseMatrixData PYAMLCLASS = "OrbitResponseMatrixData" -class ConfigModel(ReponseMatrixDataConfigModel): - """ - Configuration model for orbit response matrix +@register_schema +@dataclass +class OrbitResponseMatrixData(ResponseMatrixData, DynamicValidation): + """Store orbit response matrix data and related metadata. + + In addition to the response matrix, variable names, and observable names + provided by :class:`ResponseMatrixData`, this class stores the RF response + and the plane associated with each variable and observable. Parameters ---------- - rf_response : list[float], optional - RF response data - variable_names : list[str], optional - Vaiable plane names, basically the plane of the actuators - observable_names : list[str], optional - Observable plane names, basically the plane of measurements + matrix : list[list[float]] + Orbit response matrix. Each row corresponds to an observable and each + column corresponds to a variable. + variable_names : list[str] or None + Names of the response-matrix variables, typically orbit correctors. + observable_names : list[str] + Names of the response-matrix observables, typically beam position + monitors. + rf_response : list[float] or None, optional + Orbit response to an RF-frequency change. + variable_planes : list[str] or None, optional + Plane associated with each response-matrix variable, typically the + plane of the corresponding actuator. + observable_planes : list[str] or None, optional + Plane associated with each response-matrix observable, typically the + plane of the corresponding measurement. """ rf_response: Optional[list[float]] = None variable_planes: Optional[list[str]] = None observable_planes: Optional[list[str]] = None - - -class OrbitResponseMatrixData(ResponseMatrixData): - """ - Orbit response matrix loader - """ - - def __init__(self, cfg: ConfigModel): - self._cfg = cfg diff --git a/pyaml/tuning_tools/response_matrix_data.py b/pyaml/tuning_tools/response_matrix_data.py index 075e407f..323682c9 100644 --- a/pyaml/tuning_tools/response_matrix_data.py +++ b/pyaml/tuning_tools/response_matrix_data.py @@ -1,48 +1,57 @@ +from dataclasses import dataclass from pathlib import Path -from typing import Optional - -from pydantic import BaseModel, ConfigDict from pyaml.common.element import __pyaml_repr__ from pyaml.configuration.factory import Factory from .. import PyAMLException from ..configuration.fileloader import load +from ..validation import DynamicValidation, register_schema PYAMLCLASS = "ResponseMatrixData" -class ConfigModel(BaseModel): - """ - Base configuration model for response matrix +@register_schema +@dataclass +class ResponseMatrixData(DynamicValidation): + """Response matrix data and its associated variable and observable names. Parameters ---------- matrix : list[list[float]] - Response matrix data (rows for observable, cols for variables) - variable_names : list[str], optional - Variable names, basically the actuators - observables_names : list[str] - Observable names, basically the measurements + Response matrix values. Each row corresponds to an observable and each + column corresponds to a variable. + variable_names : list[str] or None + Names of the variables represented by the matrix columns, typically + actuators. May be ``None`` if the names are unavailable. + observable_names : list[str] + Names of the observables represented by the matrix rows, typically + measured quantities. """ matrix: list[list[float]] - variable_names: Optional[list[str]] observable_names: list[str] - - -class ResponseMatrixData(object): - """ - Generic response matrix loader - """ - - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + variable_names: list[str] | None = None + type: str | None = None @staticmethod def load(filename: str) -> "ResponseMatrixData": - """ - Load a reponse matrix from a configuration file + """Load response matrix data from a configuration file. + + Parameters + ---------- + filename : str + Path to the response matrix configuration file. + + Returns + ------- + ResponseMatrixData + Response matrix data constructed from the configuration file. + + Raises + ------ + PyAMLException + If the specified file does not exist. """ path = Path(filename) if path.exists(): @@ -50,6 +59,3 @@ def load(filename: str) -> "ResponseMatrixData": return Factory.build(config_dict, ignore_external=False) else: raise PyAMLException(f"{filename}: file not found") - - def __repr__(self): - return __pyaml_repr__(self) diff --git a/pyaml/tuning_tools/tune.py b/pyaml/tuning_tools/tune.py index 382092bf..3426a8e6 100644 --- a/pyaml/tuning_tools/tune.py +++ b/pyaml/tuning_tools/tune.py @@ -5,13 +5,7 @@ import numpy as np -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier - from .. import PyAMLException -from ..common.element import ElementConfigModel from .response_matrix_data import ResponseMatrixData from .tuning_tool import TuningTool @@ -19,62 +13,77 @@ from ..arrays.magnet_array import MagnetArray from ..diagnostics.tune_monitor import BetatronTuneMonitor +from ..validation import DynamicValidation, register_schema + logger = logging.getLogger(__name__) # Define the main class name for this module PYAMLCLASS = "Tune" -class ConfigModel(ElementConfigModel): - """ - Configuration model for Tune +@register_schema +class Tune(TuningTool, DynamicValidation): + """Adjust the horizontal and vertical betatron tunes. + + The tune correction is calculated from a response matrix describing the + change in horizontal and vertical tune produced by changes in quadrupole + strength. The pseudo-inverse of this matrix is used to convert a requested + tune change into the corresponding quadrupole-strength changes. + + The response matrix may be supplied directly as a + :class:`ResponseMatrixData` instance or loaded from a file. The tune is + measured using the configured betatron tune monitor, while corrections are + applied through the configured quadrupole array. Parameters ---------- + name : str + Name of the tuning tool. quad_array_name : str - Array name of quad used to adjust the tune + Name of the quadrupole array used to adjust the tune. betatron_tune_name : str - Name of the diagnostic pyaml device for measuring the tune - quad_delta : float - Delta strength used to get the response matrix - - """ - - quad_array_name: str - betatron_tune_name: str - response_matrix: str | ResponseMatrixData - - -class Tune(TuningTool): - """ - Class providing tune adjustment tool + Name of the betatron tune monitor used to measure the horizontal and + vertical tunes. + response_matrix : str or ResponseMatrixData + Tune response matrix or path to a file containing the response matrix. + The matrix is expected to have one row for each tune plane and one + column for each quadrupole. + + Attributes + ---------- + quad_array_name : str + Name of the configured quadrupole array. + betatron_tune_name : str + Name of the configured betatron tune monitor. + response_matrix : ResponseMatrixData or None + Loaded tune response matrix. """ - def __init__(self, cfg: ConfigModel): - """ - Construct a Tune adjustment object. - - Parameters - ---------- - cfg : ConfigModel - Configuration for the tune adjustment. - """ - super().__init__(cfg.name) - self._cfg = cfg - self._response_matrix = None + def __init__( + self, + name: str, + quad_array_name: str, + betatron_tune_name: str, + response_matrix: str | ResponseMatrixData, + ): + super().__init__(name) + + self.quad_array_name = quad_array_name + self.betatron_tune_name = betatron_tune_name + self._response_matrix = response_matrix 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) is str: try: - cfg.response_matrix = ResponseMatrixData.load(cfg.response_matrix) + self._response_matrix = ResponseMatrixData.load(self._response_matrix) 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.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -90,8 +99,8 @@ def load(self, 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.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property @@ -99,17 +108,17 @@ 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 @property def _tm(self) -> "BetatronTuneMonitor": self.check_peer() - return self.peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) + return self.peer.get_betatron_tune_monitor(self.betatron_tune_name) @property def _quads(self) -> "MagnetArray": self.check_peer() - return self.peer.get_magnets(self._cfg.quad_array_name) + return self.peer.get_magnets(self.quad_array_name) def get(self): """ diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 15f432dc..67151270 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -1,4 +1,5 @@ import logging +from dataclasses import asdict from time import sleep from typing import Callable, Optional @@ -6,39 +7,104 @@ from pydantic import ConfigDict from ..common.constants import Action -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool +from .response_matrix_data import ResponseMatrixData logger = logging.getLogger(__name__) PYAMLCLASS = "TuneResponseMatrix" -class ConfigModel(MeasurementToolConfigModel): - """ - Configuration model for Tune response matrix +@register_schema +class TuneResponseMatrix(MeasurementTool, DynamicValidation): + """Measure the response of the betatron tune to quadrupole-strength changes. + + The tune response matrix describes the change in horizontal and vertical + betatron tune produced by changes in quadrupole strength. Each quadrupole + in the configured magnet array is varied over a range of strengths, and + the resulting tune values are measured and averaged. + + For measurements using more than one strength step, a linear fit is used + to determine the tune response to each quadrupole. The resulting matrix + has one row for each tune plane and one column for each quadrupole. + + After a successful measurement, the result is stored in + :attr:`MeasurementTool.latest_measurement` as a + :class:`ResponseMatrixData` representation. Parameters ---------- + name : str + Name of the measurement tool. quad_array_name : str - Array name of quad used to adjust the tune + Name of the quadrupole array used for the measurement. betatron_tune_name : str - Name of the diagnostic pyaml device for measuring the tune + Name of the betatron tune monitor used to measure the horizontal and + vertical tunes. quad_delta : float - Delta strength used to get the response matrix + Maximum positive and negative quadrupole-strength change applied during + the measurement. + n_step : int, optional + Number of quadrupole-strength settings used for each quadrupole. The + settings are distributed linearly from ``-quad_delta`` to + ``quad_delta``. The default is 1. + sleep_between_step : float, optional + Time in seconds to wait after changing a quadrupole strength and before + measuring the tune. The default is 0. + n_avg_meas : int, optional + Number of tune measurements averaged at each strength setting. The + default is 1. + sleep_between_meas : float, optional + Time in seconds to wait between tune measurements used for averaging. + The default is 0. + + Attributes + ---------- + quad_array_name : str + Name of the configured quadrupole array. + betatron_tune_name : str + Name of the configured betatron tune monitor. + quad_delta : float + Configured quadrupole-strength change. + n_step : int + Configured number of strength settings. + sleep_between_step : float + Configured delay after each strength change. + n_avg_meas : int + Configured number of tune measurements to average. + sleep_between_meas : float + Configured delay between averaged tune measurements. + + Notes + ----- + The generated response matrix has shape ``(2, n_quadrupoles)``. The first + row contains the horizontal tune response and the second row contains the + vertical tune response. + + Quadrupole strengths are restored after each individual scan and again when + the measurement exits because of an error or interruption. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - quad_array_name: str - betatron_tune_name: str - quad_delta: float - - -class TuneResponseMatrix(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + def __init__( + self, + name: str, + quad_array_name: str, + betatron_tune_name: str, + quad_delta: float, + n_step: Optional[int] = 1, + sleep_between_step: Optional[float] = 0, + n_avg_meas: Optional[int] = 1, + sleep_between_meas: Optional[float] = 0, + ): + super().__init__(name) + self.quad_array_name = quad_array_name + self.betatron_tune_name = betatron_tune_name + self.quad_delta = quad_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, @@ -114,16 +180,16 @@ def callback(action: Action, data:dict): """ # Get devices self.check_peer() - quads = self._peer.get_magnets(self._cfg.quad_array_name) - tm = self._peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) + quads = self._peer.get_magnets(self.quad_array_name) + tm = self._peer.get_betatron_tune_monitor(self.betatron_tune_name) tunemat = np.zeros((len(quads), 2)) initial_tune = tm.tune.get() - delta = quad_delta if quad_delta is not None else self._cfg.quad_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 = quad_delta if quad_delta is not None else self.quad_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 self._register_callback(callback) self._init_measure("pyaml.tuning_tools.response_matrix_data") @@ -193,12 +259,12 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - mat = ResponseMatrixDataConfigModel( + mat = ResponseMatrixData( matrix=tunemat.T.tolist(), variable_names=quads.names(), observable_names=[tm.get_name() + ".x", tm.get_name() + ".y"], ) - self.latest_measurement.update(mat.model_dump()) + self.latest_measurement.update(asdict(mat)) self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data" return True 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/pyaml/validation/schema_builder.py b/pyaml/validation/schema_builder.py index 652a30fb..ea334aa5 100644 --- a/pyaml/validation/schema_builder.py +++ b/pyaml/validation/schema_builder.py @@ -107,7 +107,7 @@ def _configuration_schema_from_basemodel( if not isinstance(validation_model, type) or not issubclass(validation_model, BaseModel): raise TypeError("validation_model must be a subclass of pydantic.BaseModel.") - fields: dict[str, tuple[object, object]] = {} + fields: dict[str, Any] = {} for field_name, field_info in validation_model.model_fields.items(): if field_name in RESERVED_CONFIGURATION_FIELDS: @@ -125,7 +125,7 @@ def _configuration_schema_from_basemodel( ) -def _field_definition_from_field_info(field_name: str, field_info: FieldInfo) -> tuple[object, object]: +def _field_definition_from_field_info(field_name: str, field_info: FieldInfo) -> tuple[Any, Any]: """ Convert a Pydantic field definition into a ``create_model`` field tuple. """ @@ -152,7 +152,7 @@ def _field_definition_from_field_info(field_name: str, field_info: FieldInfo) -> return annotation, default -def _resolve_annotation(annotation: object) -> object: +def _resolve_annotation(annotation: Any) -> Any: """ Resolve an annotation into a schema-friendly type. @@ -225,12 +225,12 @@ def _resolve_annotation(annotation: object) -> object: raise TypeError(f"Unsupported generic annotation: {annotation!r}") -def _field_kwargs(field_name: str, field_info: object) -> dict[str, object]: +def _field_kwargs(field_name: str, field_info: Any) -> dict[str, Any]: """ Collect supported field metadata for ``pydantic.create_model``. """ - kwargs: dict[str, object] = {} + kwargs: dict[str, Any] = {} description = getattr(field_info, "description", None) if description is not None: @@ -266,7 +266,7 @@ def _configuration_schema_from_constructor(cls: type) -> type[ConfigurationSchem if not isinstance(cls, type): raise TypeError("cls must be a class.") - fields = _fields_from_constructor_signature( + fields: dict[str, Any] = _fields_from_constructor_signature( cls, expand_arbitrary_types=True, ) @@ -279,7 +279,7 @@ def _configuration_schema_from_constructor(cls: type) -> type[ConfigurationSchem ) -def _fields_from_constructor_signature(cls: type, expand_arbitrary_types: bool = False) -> dict[str, tuple[object, object]]: +def _fields_from_constructor_signature(cls: type, expand_arbitrary_types: bool = False) -> dict[str, tuple[Any, Any]]: """ Extract field definitions from a class constructor signature. diff --git a/pyaml/validation/validation_models.py b/pyaml/validation/validation_models.py index 0683d86d..4bbe7b29 100644 --- a/pyaml/validation/validation_models.py +++ b/pyaml/validation/validation_models.py @@ -3,7 +3,7 @@ import inspect import logging from abc import ABCMeta -from typing import Any +from typing import Any, ClassVar from pydantic import BaseModel, ConfigDict, create_model @@ -94,19 +94,41 @@ def __call__(cls, *args: Any, **kwargs: Any): return super().__call__(**validated.model_dump()) +class ValidationModelDescriptor: + """Provide a lazily generated validation model on the class.""" + + def __get__( + self, + instance: object | None, + owner: type["DynamicValidation"], + ) -> type[ValidationModel]: + model = owner.__dict__.get("_validation_model") + + if model is None: + model = owner._build_validation_model() + owner._validation_model = model + + return model + + class DynamicValidation(metaclass=ValidationMeta): - """ - Base class for automatic constructor argument validation. + """Base class for automatic constructor argument validation. - When a subclass is defined, a validation model is generated - automatically from its constructor signature and assigned to - ``validation_model``. The generated model is then used to validate - constructor arguments before object creation. + When a subclass is defined, a validation model is generated from either + its explicitly declared constructor or its directly declared class + annotations. + + Class annotations are used when no constructor has yet been defined. This + supports classes decorated with ``@dataclass``, because their generated + constructor is not available while ``__init_subclass__`` is running. Subclasses must not define ``validation_model`` manually. """ - validation_model: type[ValidationModel] | None = None + _validation_model: ClassVar[type[ValidationModel] | None] = None + + # Lazy construction of the validation class + validation_model: ClassVar[ValidationModelDescriptor] = ValidationModelDescriptor() def __init_subclass__(cls, **kwargs): """ @@ -120,31 +142,33 @@ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - if getattr(cls, "validation_model", None) is not None: + # Check if validation model already exists + # This only checks for the current class and not parent classes + if "validation_model" in cls.__dict__: raise TypeError(f"{cls.__name__} may not define validation_model manually.") - cls.validation_model = cls._build_validation_model() + # Set the _validation_model to None since should be built using lazy construction + cls._validation_model = None @classmethod def _build_validation_model(cls) -> type[ValidationModel]: - """ - Generate a validation model from the constructor signature. + """Generate a validation model from the class definition. - The generated model contains one field for each parameter in the - subclass's ``__init__`` method, excluding ``self``, ``*args``, and - ``**kwargs``. Field types are obtained from the constructor's type - annotations and default values are preserved. + For classes with an explicitly defined ``__init__``, fields are + extracted from the constructor signature. Otherwise, fields are + extracted from annotations declared directly on the class. The latter + supports dataclasses before their generated constructor is available. Returns ------- type[ValidationModel] - A dynamically generated subclass of :class:`ValidationModel` - representing the constructor arguments accepted by the subclass. + Dynamically generated validation model representing the arguments + accepted by the class. """ logger.debug("Building validation model for %s.", f"{cls.__module__}.{cls.__name__}") - fields = _fields_from_constructor_signature(cls, expand_arbitrary_types=False) + fields: dict[str, Any] = _fields_from_constructor_signature(cls, expand_arbitrary_types=False) model = create_model(f"{cls.__name__}ValidationModel", **fields, __base__=ValidationModel) 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