diff --git a/pyaml/common/element_holder.py b/pyaml/common/element_holder.py index db6b5215..8db92878 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: @@ -294,7 +294,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 e75b282d..0afbfb0e 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -3,7 +3,7 @@ from .. import PyAMLException from ..validation import DynamicValidation, register_schema -from .chromaticity_monitor import ChomaticityMonitor +from .chromaticity_monitor import ChromaticityMonitor from .response_matrix_data import ResponseMatrixData from .tuning_tool import TuningTool @@ -61,7 +61,7 @@ def __init__( # Invert matrix if self._response_matrix: - self._response_matrix = np.array(self._response_matrix._cfg.matrix) + self._response_matrix = np.array(self._response_matrix.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -84,11 +84,11 @@ def load(self, load_path: Path): Filename of the :class:`~.ResponseMatrixData` to load """ self._response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._response_matrix._cfg.matrix) + 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._chromaticity_monitor_name) diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index 2d4c24d5..fc0e02d7 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -14,7 +14,7 @@ logger = logging.getLogger(__name__) -PYAMLCLASS = "ChomaticityMonitor" +PYAMLCLASS = "ChromaticityMonitor" class RChromaDispArray(ReadFloatArray): @@ -23,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 @@ -40,7 +40,7 @@ def unit(self) -> str: @register_schema -class ChomaticityMonitor(MeasurementTool, DynamicValidation): +class ChromaticityMonitor(MeasurementTool, DynamicValidation): """ Class providing access to a chromaticity monitor of a physical or simulated lattice. The monitor provides diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index ab0e23d7..c8305bb7 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -1,5 +1,6 @@ import logging import time +from dataclasses import asdict from typing import Callable, Optional import numpy as np @@ -7,7 +8,7 @@ from ..common.constants import Action from ..validation import DynamicValidation, register_schema from .measurement_tool import MeasurementTool -from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from .response_matrix_data import ResponseMatrixData logger = logging.getLogger(__name__) @@ -252,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 c1cbeba6..0d2cf448 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -17,31 +17,12 @@ 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 + n_step: int = 10 + sleep_between_step: float = 0 + n_avg_meas: int = 1 + sleep_between_meas: float = 0 class MeasurementTool(Element, metaclass=ABCMeta): 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 db267725..2fa17fda 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/validation/schema_builder.py b/pyaml/validation/schema_builder.py index a2ed5759..edd95d5f 100644 --- a/pyaml/validation/schema_builder.py +++ b/pyaml/validation/schema_builder.py @@ -169,7 +169,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: @@ -187,7 +187,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. """ @@ -214,7 +214,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. @@ -287,12 +287,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: @@ -328,7 +328,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, ) @@ -341,7 +341,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 1cd9e8bc..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) + # 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)