From 4e8ad25353cab046327651c7c1ba54df5175d6e6 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Mon, 13 Jul 2026 18:30:55 +0200 Subject: [PATCH 01/10] Changes to remove config models from chromaticity. --- pyaml/tuning_tools/chromaticity.py | 74 +++---- pyaml/tuning_tools/chromaticity_monitor.py | 195 ++++++++++-------- .../chromaticity_response_matrix.py | 94 ++++++--- pyaml/tuning_tools/measurement_tool.py | 15 +- pyaml/tuning_tools/tuning_tool.py | 11 +- tests/config/EBSOrbit.yaml | 2 +- 6 files changed, 233 insertions(+), 158 deletions(-) diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index 5016bf2b..e75b282d 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING from .. import PyAMLException -from ..common.element import ElementConfigModel +from ..validation import DynamicValidation, register_schema from .chromaticity_monitor import ChomaticityMonitor from .response_matrix_data import ResponseMatrixData from .tuning_tool import TuningTool @@ -10,10 +10,6 @@ if TYPE_CHECKING: from ..arrays.magnet_array import MagnetArray -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier import logging import time @@ -25,55 +21,47 @@ PYAMLCLASS = "Chromaticity" -class ConfigModel(ElementConfigModel): - """ - Configuration model for Tune - - Parameters - ---------- - sextu_array_name : str - Array name of sextu used to adjust the chromaticity - chromaticty_monitor_name : str - Name of the diagnostic pyaml device for measuring the chromaticity - response_matrix : str | ResponseMatrixData - filename or data of the chromaticity response matrix - """ - - sextu_array_name: str - chromaticty_monitor_name: str - response_matrix: str | ResponseMatrixData - - -class Chromaticity(TuningTool): +@register_schema +class Chromaticity(TuningTool, DynamicValidation): """ Class providing chromaticity adjustment tool """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, name: str, sextu_array_name: str, chromaticity_monitor_name: str, response_matrix: str | ResponseMatrixData + ): """ - Construct a chromaticity adjustment object. + Initialize a chromaticity adjustment tool. Parameters ---------- - cfg : ConfigModel - Configuration for the tune adjustment. - """ - super().__init__(cfg.name) - self._cfg = cfg + name : str + Name of the tuning tool. + sextu_array_name : str + Name of the sextupole array used to adjust the chromaticity. + chromaticity_monitor_name : str + Name of the chromaticity monitor used for readback. + response_matrix : str | ResponseMatrixData + Chromaticity response matrix or path to a saved response matrix file. + """ + super().__init__(name) + self.sextu_array_name = sextu_array_name + self._chromaticity_monitor_name = chromaticity_monitor_name + self.response_matrix_file = response_matrix self._response_matrix = None self._correctionmat = None # If the configuration response matrix is a filename, load it - if type(cfg.response_matrix) is str: + if type(self.response_matrix_file) is str: try: - cfg.response_matrix = ResponseMatrixData.load(cfg.response_matrix) + self._response_matrix = ResponseMatrixData.load(self.response_matrix_file) except Exception as e: logger.warning(f"{str(e)}") - cfg.response_matrix = None + self._response_matrix = None # Invert matrix - if cfg.response_matrix: - self._response_matrix = np.array(cfg.response_matrix._cfg.matrix) + if self._response_matrix: + self._response_matrix = np.array(self._response_matrix._cfg.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -84,7 +72,7 @@ def response_matrix(self) -> ResponseMatrixData | None: """ Return the response matrix if it has been loaded None otherwise """ - return self._cfg.response_matrix + return self._response_matrix def load(self, load_path: Path): """ @@ -95,19 +83,19 @@ def load(self, load_path: Path): load_path : Path Filename of the :class:`~.ResponseMatrixData` to load """ - self._cfg.response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._cfg.response_matrix._cfg.matrix) + self._response_matrix = ResponseMatrixData.load(load_path) + self._response_matrix = np.array(self._response_matrix._cfg.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property def _cm(self) -> "ChomaticityMonitor": self.check_peer() - return self.peer.get_chromaticity_monitor(self._cfg.chromaticty_monitor_name) + return self.peer.get_chromaticity_monitor(self._chromaticity_monitor_name) @property def _sextu(self) -> "MagnetArray": self.check_peer() - return self.peer.get_magnets(self._cfg.sextu_array_name) + return self.peer.get_magnets(self.sextu_array_name) def get(self): """ @@ -117,7 +105,7 @@ def get(self): def readback(self): """ - Launch a chromaticty scan and returns the measured chromaticity. + Launch a chromaticity scan and returns the measured chromaticity. """ self._cm.measure() return self._cm.chromaticity.get() diff --git a/pyaml/tuning_tools/chromaticity_monitor.py b/pyaml/tuning_tools/chromaticity_monitor.py index 4c4ea1cc..2d4c24d5 100644 --- a/pyaml/tuning_tools/chromaticity_monitor.py +++ b/pyaml/tuning_tools/chromaticity_monitor.py @@ -1,64 +1,22 @@ -from ..common.abstract import ReadFloatArray -from ..common.constants import Action -from ..common.element import ElementConfigModel -from ..common.exception import PyAMLException -from ..tuning_tools.measurement_tool import MeasurementTool, MeasurementToolConfigModel - -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier import logging +from collections.abc import Callable from time import sleep import matplotlib.pyplot as plt import numpy as np -from pydantic import ConfigDict + +from ..common.abstract import ReadFloatArray +from ..common.constants import Action +from ..common.element import __pyaml_repr__ +from ..common.exception import PyAMLException +from ..tuning_tools.measurement_tool import MeasurementTool +from ..validation import DynamicValidation, register_schema logger = logging.getLogger(__name__) PYAMLCLASS = "ChomaticityMonitor" -class ConfigModel(MeasurementToolConfigModel): - """ - Configuration model for Chromaticity Monitor. - - Parameters - ---------- - betatron_tune_name : str - Name of the diagnostic pyaml device for measuring the tune - rf_plant_name : str - Name of main RF frequency plant - bpm_array_name : str,optional - Name of main BPM array used for dispersion fit - e_delta : float, optional - Default variation of relative energy during chromaticity measurement: - f0 - f0 * E_delta * alphac < f_RF < f0 + f0 * E_delta * alphac, - by default 0.001 - max_e_delta : float, optional - Maximum authorized variation of relative energy during chromaticity - measurement, by default 0.004 - fit_order : int, optional - Chomaticity fitting order, by default 1 - fit_disp_order : int, optional - Dispersion fitting order, by default 1 - fit_dispersion : bool, optional - Dispersion fitting, by default False - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - betatron_tune_name: str - rf_plant_name: str - bpm_array_name: str | None = None - e_delta: float = 0.001 - max_e_delta: float = 0.004 - fit_order: int = 1 - fit_disp_order: int = 1 - fit_dispersion: bool = False - - class RChromaDispArray(ReadFloatArray): """ Class providing read access to chromaticity or dispersion. @@ -81,27 +39,89 @@ def unit(self) -> str: return self.unit -class ChomaticityMonitor(MeasurementTool): +@register_schema +class ChomaticityMonitor(MeasurementTool, DynamicValidation): """ Class providing access to a chromaticity monitor of a physical or simulated lattice. The monitor provides horizontal and vertical chromaticity measurements. """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + betatron_tune_name: str, + rf_plant_name: str, + bpm_array_name: str | None = None, + e_delta: float = 0.001, + max_e_delta: float = 0.004, + fit_order: int = 1, + fit_disp_order: int = 1, + fit_dispersion: bool = False, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): """ - Construct a ChomaticityMonitor. + Initialize a chromaticity monitor. + + The monitor performs chromaticity measurements by varying the RF + frequency, acquiring tune measurements from a betatron tune monitor, + and optionally fitting the machine dispersion using BPM orbit data. Parameters ---------- - cfg : ConfigModel - Configuration for the ChromaticityMonitor, including betatron - tune monitor, RF plant, and defaults parameters. + name : str + Name of the chromaticity monitor. + betatron_tune_name : str + Name of the betatron tune monitor used to measure the horizontal + and vertical tunes. + rf_plant_name : str + Name of the RF plant used to vary the RF frequency. + bpm_array_name : str, optional + Name of the BPM array used for dispersion measurements. Required + only when dispersion fitting is enabled. + e_delta : float, optional + Default relative momentum deviation used during the measurement. + max_e_delta : float, optional + Maximum permitted relative momentum deviation. + fit_order : int, optional + Polynomial order used to fit the chromaticity. + fit_disp_order : int, optional + Polynomial order used to fit the dispersion. + fit_dispersion : bool, optional + Whether to fit the machine dispersion in addition to the + chromaticity. + n_step : int, optional + Number of RF frequency steps. + sleep_between_step : float, optional + Delay in seconds after changing the RF frequency. + n_avg_meas : int, optional + Number of tune (and orbit) measurements to average at each RF + frequency. + sleep_between_meas : float, optional + Delay in seconds between consecutive measurements during + averaging. """ - super().__init__(cfg.name) - self._cfg = cfg + + super().__init__(name) + + self.betatron_tune_name = betatron_tune_name + self.rf_plant_name = rf_plant_name + self.bpm_array_name = bpm_array_name + self.e_delta = e_delta + self.max_e_delta = max_e_delta + self.fit_order = fit_order + self.fit_disp_order = fit_disp_order + self.fit_dispersion = fit_dispersion + self.n_step = n_step + self.sleep_between_step = sleep_between_step + self.n_avg_meas = n_avg_meas + self.sleep_between_meas = sleep_between_meas + self._chromaticity = RChromaDispArray(self, "chromaticity", "1") - self._dipsersion = RChromaDispArray(self, "dispersion", "m") + self._dispersion = RChromaDispArray(self, "dispersion", "m") self._alphac = None @property @@ -129,22 +149,22 @@ def dispersion(self) -> ReadFloatArray: ReadFloatArray Array of dispersion values [[dx, dy],[d'x, d'y],...] """ - return self._dipsersion + return self._dispersion def measure( self, - n_step: int = None, - alphac: float = None, - e_delta: float = None, - max_e_delta: float = None, - n_avg_meas: int = None, - sleep_between_meas: float = None, - sleep_between_step: float = None, - fit_order: int = None, - fit_disp_order: int = None, + n_step: int | None = None, + alphac: float | None = None, + e_delta: float | None = None, + max_e_delta: float | None = None, + n_avg_meas: int | None = None, + sleep_between_meas: float | None = None, + sleep_between_step: float | None = None, + fit_order: int | None = None, + fit_disp_order: int | None = None, fit_dispersion: bool | None = None, - do_plot: bool = None, - callback: callable = None, + do_plot: bool | None = None, + callback: Callable | None = None, ): """ Main function for chromaticity measurment. @@ -202,16 +222,16 @@ def measure( dtune:np.array # The tune variation (on Action.RESTORE) """ - n_step = n_step if n_step is not None else self._cfg.n_step + n_step = n_step if n_step is not None else self.n_step alphac = alphac if alphac is not None else self._alphac - e_delta = e_delta if e_delta is not None else self._cfg.e_delta - max_e_delta = max_e_delta if max_e_delta is not None else self._cfg.max_e_delta - n_avg_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_between_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas - sleep_between_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - fit_order = fit_order if fit_order is not None else self._cfg.fit_order - fit_disp_order = fit_disp_order if fit_disp_order is not None else self._cfg.fit_disp_order - fit_dispersion = fit_dispersion if fit_dispersion is not None else self._cfg.fit_dispersion + e_delta = e_delta if e_delta is not None else self.e_delta + max_e_delta = max_e_delta if max_e_delta is not None else self.max_e_delta + n_avg_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + sleep_between_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas + sleep_between_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + fit_order = fit_order if fit_order is not None else self.fit_order + fit_disp_order = fit_disp_order if fit_disp_order is not None else self.fit_disp_order + fit_dispersion = fit_dispersion if fit_dispersion is not None else self.fit_dispersion if abs(e_delta) > abs(max_e_delta): logger.warning(f"e_delta={e_delta} is greater than max_e_delta={max_e_delta}") @@ -224,14 +244,14 @@ def measure( # Get devices self.check_peer() - tm = self.peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) - rf = self.peer.get_rf_plant(self._cfg.rf_plant_name) + tm = self.peer.get_betatron_tune_monitor(self.betatron_tune_name) + rf = self.peer.get_rf_plant(self.rf_plant_name) bpms = None n_bpm = 0 orbit = None - if fit_dispersion and fit_disp_order is not None and self._cfg.bpm_array_name is not None: + if fit_dispersion and fit_disp_order is not None and self.bpm_array_name is not None: # For dispersion fit - bpms = self.peer.get_bpms(self._cfg.bpm_array_name) + bpms = self.peer.get_bpms(self.bpm_array_name) n_bpm = len(bpms) f0 = rf.frequency.get() @@ -365,3 +385,10 @@ def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False): fig.tight_layout() plt.show() + + def _after_attach(self): + self._chromaticity = RChromaDispArray(self, "chromaticity", "1") + self._dispersion = RChromaDispArray(self, "dispersion", "m") + + def __repr__(self): + return __pyaml_repr__(self, exclude=["chromaticity", "dispersion"]) diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index 8b995fbf..ab0e23d7 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -3,11 +3,10 @@ from typing import Callable, Optional import numpy as np -from pydantic import ConfigDict from ..common.constants import Action -from ..common.element import ElementConfigModel -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel logger = logging.getLogger(__name__) @@ -15,31 +14,78 @@ PYAMLCLASS = "ChromaticityResponseMatrix" -class ConfigModel(MeasurementToolConfigModel): +@register_schema +class ChromaticityResponseMatrix(MeasurementTool, DynamicValidation): """ - Configuration model for Tune response matrix + Measure the chromaticity response matrix of a lattice. + + This tool perturbs each sextupole in a sextupole array and measures the + resulting chromaticity variation using a chromaticity monitor. The measured + slopes are assembled into a response matrix with chromaticity components as + the observables and sextupoles as the variables. Parameters ---------- + name : str + Name of the response matrix measurement tool. sextu_array_name : str - Array name of sextupole used to adjust the chromaticity + Name of the sextupole array to excite. chromaticity_name : str - Name of the diagnostic chromaticy monitor + Name of the chromaticity monitor used to measure the response. sextu_delta : float - Delta strength used to get the response matrix + Default sextupole excitation applied during the measurement. + n_step : int, optional + Default number of excitation steps used to fit the response. + sleep_between_step : float, optional + Default delay in seconds after changing the sextupole strength. + n_avg_meas : int, optional + Default number of chromaticity measurements to average at each step. + sleep_between_meas : float, optional + Default delay in seconds between averaged measurements. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - sextu_array_name: str - chromaticity_name: str - sextu_delta: float + def __init__( + self, + name: str, + sextu_array_name: str, + chromaticity_name: str, + sextu_delta: float, + n_step: int = 1, + sleep_between_step: float = 0, + n_avg_meas: int = 1, + sleep_between_meas: float = 0, + ): + """ + Initialize the chromaticity response matrix measurement tool. + Parameters + ---------- + name : str + Name of the response matrix measurement tool. + sextu_array_name : str + Name of the sextupole array to excite. + chromaticity_name : str + Name of the chromaticity monitor used to measure the response. + sextu_delta : float + Default sextupole excitation applied during the measurement. + n_step : int, optional + Default number of excitation steps used to fit the response. + sleep_between_step : float, optional + Default delay in seconds after changing the sextupole strength. + n_avg_meas : int, optional + Default number of chromaticity measurements to average at each step. + sleep_between_meas : float, optional + Default delay in seconds between averaged measurements. + """ -class ChromaticityResponseMatrix(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + super().__init__(name) + self.sextu_array_name = sextu_array_name + self.chromaticity_name = chromaticity_name + self.sextu_delta = sextu_delta + self.n_step = n_step + self.sleep_between_step = sleep_between_step + self.n_avg_meas = n_avg_meas + self.sleep_between_meas = sleep_between_meas self.aborted = False def measure( @@ -117,8 +163,8 @@ def callback(action: Action, data:dict): """ # Get devices self.check_peer() - sextus = self._peer.get_magnets(self._cfg.sextu_array_name) - cm = self._peer.get_chromaticity_monitor(self._cfg.chromaticity_name) + sextus = self._peer.get_magnets(self.sextu_array_name) + cm = self._peer.get_chromaticity_monitor(self.chromaticity_name) self._register_callback(callback) self._init_measure("pyaml.tuning_tools.response_matrix_data") @@ -131,11 +177,11 @@ def callback(action: Action, data:dict): return False initial_chroma = cm.chromaticity.get() - delta = sextu_delta if sextu_delta is not None else self._cfg.sextu_delta - nb_step = n_step if n_step is not None else self._cfg.n_step - nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas + delta = sextu_delta if sextu_delta is not None else self.sextu_delta + nb_step = n_step if n_step is not None else self.n_step + nb_meas = n_avg_meas if n_avg_meas is not None else self.n_avg_meas + sleep_step = sleep_between_step if sleep_between_step is not None else self.sleep_between_step + sleep_meas = sleep_between_meas if sleep_between_meas is not None else self.sleep_between_meas err = None aborted = False diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index edec89e4..c1cbeba6 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -1,3 +1,4 @@ +import copy import logging from abc import ABCMeta, abstractmethod from pathlib import Path @@ -166,10 +167,14 @@ def _register_callback(self, callback: Callable): self._callback = callback def attach(self, peer: "ElementHolder") -> Self: - """ - Create a new reference to attach this measurement tool object to a simulator - or a control system. - """ - obj = self.__class__(self._cfg) + if hasattr(self, "_cfg"): + obj = self.__class__(self._cfg) + else: + obj = copy.copy(self) + obj._after_attach() obj._peer = peer return obj + + def _after_attach(self) -> None: + """Hook for subclasses to rebind internal references after attach.""" + pass diff --git a/pyaml/tuning_tools/tuning_tool.py b/pyaml/tuning_tools/tuning_tool.py index 8f043bc4..f24f42c9 100644 --- a/pyaml/tuning_tools/tuning_tool.py +++ b/pyaml/tuning_tools/tuning_tool.py @@ -1,3 +1,4 @@ +import copy from typing import Self from ..common.element import Element @@ -17,6 +18,14 @@ def attach(self, peer: "ElementHolder") -> Self: Create a new reference to attach this tuning tool object to a simulator or a control system. """ - obj = self.__class__(self._cfg) + if hasattr(self, "_cfg"): + obj = self.__class__(self._cfg) + else: + obj = copy.copy(self) + obj._after_attach() obj._peer = peer return obj + + def _after_attach(self) -> None: + """Hook for subclasses to rebind internal references after attach.""" + pass diff --git a/tests/config/EBSOrbit.yaml b/tests/config/EBSOrbit.yaml index 9b66318a..ce29b43f 100644 --- a/tests/config/EBSOrbit.yaml +++ b/tests/config/EBSOrbit.yaml @@ -53,7 +53,7 @@ devices: sextu_delta: 1e-6 - type: pyaml.tuning_tools.chromaticity name: DEFAULT_CHROMATICITY_CORRECTION - chromaticty_monitor_name: CHROMATICITY_MONITOR + chromaticity_monitor_name: CHROMATICITY_MONITOR sextu_array_name: Sext response_matrix: ${path:ideal_chroma_resp.json} - type: pyaml.tuning_tools.orbit From 6e8904cc40d1bdd095dc35c970362f77cafad563 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 16:50:47 +0200 Subject: [PATCH 02/10] Remove config model from dispersion. --- pyaml/tuning_tools/dispersion.py | 52 +++++++++++++++++--------------- 1 file changed, 28 insertions(+), 24 deletions(-) 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, From 0c8a95ec246a4aaa2531d237423ba467561815ce Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 18:20:15 +0200 Subject: [PATCH 03/10] Updates to the schema builder for dynamic validation to also work for dataclasses. --- pyaml/validation/schema_builder.py | 118 ++++++++++++++++++++++---- pyaml/validation/validation_models.py | 35 ++++---- 2 files changed, 119 insertions(+), 34 deletions(-) diff --git a/pyaml/validation/schema_builder.py b/pyaml/validation/schema_builder.py index 652a30fb..66679cd1 100644 --- a/pyaml/validation/schema_builder.py +++ b/pyaml/validation/schema_builder.py @@ -3,6 +3,8 @@ import inspect import logging import types +from dataclasses import MISSING +from dataclasses import Field as DataclassField from datetime import date, datetime, time, timedelta from decimal import Decimal from enum import Enum @@ -77,7 +79,7 @@ def generate_configuration_schema(source: type) -> type[ConfigurationSchema]: source.__module__, ) else: - schema = _configuration_schema_from_constructor(source) + schema = _configuration_schema_from_class(source) logger.debug("Register schema for %s.", class_path) registry.register(class_path, schema) @@ -107,7 +109,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 +127,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 +154,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 +227,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: @@ -255,21 +257,18 @@ def _field_kwargs(field_name: str, field_info: object) -> dict[str, object]: return kwargs -def _configuration_schema_from_constructor(cls: type) -> type[ConfigurationSchema]: - """ - Generate a configuration schema from a class constructor signature. +def _configuration_schema_from_class(cls: type) -> type[ConfigurationSchema]: + """Generate a configuration schema from a class definition. - The resulting schema contains one field for each constructor parameter, - with annotations and default values preserved. + Fields are extracted from an explicitly declared constructor when one + exists. Otherwise, fields are extracted from annotations declared directly + on the class, which supports dataclasses before their generated + constructors are available. """ - if not isinstance(cls, type): raise TypeError("cls must be a class.") - fields = _fields_from_constructor_signature( - cls, - expand_arbitrary_types=True, - ) + fields: dict[str, Any] = _fields_from_class_definition(cls, expand_arbitrary_types=True) return create_model( _generate_schema_name(cls), @@ -279,7 +278,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. @@ -322,3 +321,88 @@ def _fields_from_constructor_signature(cls: type, expand_arbitrary_types: bool = fields[name] = (annotation, default) return fields + + +def _fields_from_class_annotations(cls: type, expand_arbitrary_types: bool = False) -> dict[str, tuple[Any, Any]]: + """Extract field definitions from annotations declared on a class. + + This is primarily used for classes that will be transformed by + ``@dataclass``. During ``__init_subclass__``, the dataclass constructor + has not yet been generated, but the class annotations and defaults are + already available. + + Parameters + ---------- + cls + Class whose directly declared annotations should be inspected. + expand_arbitrary_types + If true, unsupported annotations are expanded recursively into + schema-friendly types. + + Returns + ------- + dict[str, tuple[Any, Any]] + Field definitions suitable for ``pydantic.create_model``. + """ + declared_annotations = cls.__dict__.get("__annotations__", {}) + type_hints = get_type_hints(cls, include_extras=True) + + fields: dict[str, tuple[Any, Any]] = {} + + for name in declared_annotations: + if name in RESERVED_CONFIGURATION_FIELDS: + raise ValueError(f"{cls.__name__} defines reserved field {name!r}, which is owned by ConfigurationSchema.") + + annotation = type_hints.get(name, Any) + + if expand_arbitrary_types: + annotation = _resolve_annotation(annotation) + + default = cls.__dict__.get(name, ...) + + # Support dataclasses.field(...), although @dataclass has not yet + # processed it when this function is called from __init_subclass__. + if isinstance(default, DataclassField): + if default.default is not MISSING: + default = default.default + elif default.default_factory is not MISSING: + default = Field(default_factory=default.default_factory) + else: + default = ... + + fields[name] = (annotation, default) + + return fields + + +def _fields_from_class_definition( + cls: type, + expand_arbitrary_types: bool = False, +) -> dict[str, tuple[Any, Any]]: + """Extract validation fields from a class definition. + + Fields are read from an explicitly defined constructor when one exists. + Otherwise, fields are read from annotations declared directly on the + class. The annotation fallback supports classes that are about to be + transformed by ``@dataclass``. + """ + if "__init__" in cls.__dict__: + return _fields_from_constructor_signature( + cls, + expand_arbitrary_types=expand_arbitrary_types, + ) + + declared_annotations = cls.__dict__.get("__annotations__", {}) + + if declared_annotations: + return _fields_from_class_annotations( + cls, + expand_arbitrary_types=expand_arbitrary_types, + ) + + # No constructor or locally declared annotations. Fall back to the + # effective inherited constructor. + return _fields_from_constructor_signature( + cls, + expand_arbitrary_types=expand_arbitrary_types, + ) diff --git a/pyaml/validation/validation_models.py b/pyaml/validation/validation_models.py index 0683d86d..f1e04a9a 100644 --- a/pyaml/validation/validation_models.py +++ b/pyaml/validation/validation_models.py @@ -8,7 +8,7 @@ from pydantic import BaseModel, ConfigDict, create_model from .configuration_models import PyAMLBaseModel -from .schema_builder import _fields_from_constructor_signature +from .schema_builder import _fields_from_class_definition logger = logging.getLogger(__name__) @@ -95,13 +95,15 @@ def __call__(cls, *args: Any, **kwargs: Any): 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. """ @@ -120,31 +122,30 @@ def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - if getattr(cls, "validation_model", None) is not None: + if "validation_model" in cls.__dict__: raise TypeError(f"{cls.__name__} may not define validation_model manually.") cls.validation_model = cls._build_validation_model() @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_class_definition(cls, expand_arbitrary_types=False) model = create_model(f"{cls.__name__}ValidationModel", **fields, __base__=ValidationModel) From 3942bd6d6a507289d8b43009ffe3b05a6eeb99a4 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 19:53:42 +0200 Subject: [PATCH 04/10] Initial changes for removing config model from response matrix data. --- pyaml/common/element_holder.py | 4 +- pyaml/tuning_tools/chromaticity.py | 8 +-- pyaml/tuning_tools/chromaticity_monitor.py | 6 +- .../chromaticity_response_matrix.py | 7 ++- pyaml/tuning_tools/orbit.py | 13 +---- pyaml/tuning_tools/orbit_response_matrix.py | 9 +-- .../orbit_response_matrix_data.py | 45 ++++++++------ pyaml/tuning_tools/response_matrix_data.py | 58 ++++++++++--------- pyaml/tuning_tools/tune.py | 2 +- pyaml/tuning_tools/tune_response_matrix.py | 7 ++- 10 files changed, 84 insertions(+), 75 deletions(-) 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 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/orbit.py b/pyaml/tuning_tools/orbit.py index 6baa1f88..83425498 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -1,17 +1,10 @@ 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 @@ -99,7 +92,7 @@ def load(self, load_path: Path): self._set_response_matrix(self._cfg.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") diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 68d26fba..4055d9a6 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -1,4 +1,5 @@ import logging +from dataclasses import asdict from pathlib import Path from typing import Callable, List, Optional, Self @@ -10,7 +11,7 @@ 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 .orbit_response_matrix_data import OrbitResponseMatrixData logger = logging.getLogger(__name__) @@ -155,11 +156,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 +188,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..8a5fd804 100644 --- a/pyaml/tuning_tools/tune.py +++ b/pyaml/tuning_tools/tune.py @@ -91,7 +91,7 @@ def load(self, load_path: Path): """ self._cfg.response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._cfg.response_matrix._cfg.matrix) + self._response_matrix = np.array(self._cfg.response_matrix.matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 15f432dc..04227663 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 @@ -7,7 +8,7 @@ from ..common.constants import Action from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from .response_matrix_data import ResponseMatrixData logger = logging.getLogger(__name__) @@ -193,12 +194,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 From c5f89a82ee941387a8adef0fb9016eb5b4e44e4f Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:04:02 +0200 Subject: [PATCH 05/10] Remove config model from orbit response matrix measurement. --- pyaml/tuning_tools/orbit_response_matrix.py | 106 ++++++++++++++------ 1 file changed, 76 insertions(+), 30 deletions(-) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 4055d9a6..cf90bd7c 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -1,16 +1,15 @@ import logging from dataclasses import asdict -from pathlib import Path -from typing import Callable, List, Optional, Self +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 ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool from .orbit_response_matrix_data import OrbitResponseMatrixData logger = logging.getLogger(__name__) @@ -18,39 +17,86 @@ 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, @@ -92,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( From ecf5b5a26471e27d4ece82fb6a2963d8f2854968 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:12:03 +0200 Subject: [PATCH 06/10] Remove config model from orbit. --- pyaml/tuning_tools/orbit.py | 88 ++++++++++++++++++------------------- 1 file changed, 43 insertions(+), 45 deletions(-) diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index 83425498..67151e77 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -4,15 +4,14 @@ from typing import Literal, Optional, Union import numpy as np -from pydantic import ConfigDict 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 @@ -22,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 @@ -88,8 +86,8 @@ 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 = asdict(mat) @@ -97,7 +95,7 @@ def _set_response_matrix(self, mat): 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 @@ -105,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, @@ -305,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) From 281b5292e34adbe6d1f785556a281aaec4b02001 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:17:51 +0200 Subject: [PATCH 07/10] Remove config model from tune response matrix. --- pyaml/tuning_tools/tune_response_matrix.py | 115 ++++++++++++++++----- 1 file changed, 90 insertions(+), 25 deletions(-) diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 04227663..67151270 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -7,7 +7,8 @@ from pydantic import ConfigDict from ..common.constants import Action -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..validation import DynamicValidation, register_schema +from .measurement_tool import MeasurementTool from .response_matrix_data import ResponseMatrixData logger = logging.getLogger(__name__) @@ -15,31 +16,95 @@ 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, @@ -115,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") From ac5335c646f901023f2a7104f98706eeb81ee89f Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:27:40 +0200 Subject: [PATCH 08/10] Remove config models for tune. --- pyaml/tuning_tools/tune.py | 101 ++++++++++++++++++++----------------- 1 file changed, 55 insertions(+), 46 deletions(-) diff --git a/pyaml/tuning_tools/tune.py b/pyaml/tuning_tools/tune.py index 8a5fd804..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.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): """ From 5056f10df61e8af3f70e94733076ffdcddef8dfa Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:28:18 +0200 Subject: [PATCH 09/10] Remove config model for measurement tool. --- pyaml/tuning_tools/measurement_tool.py | 28 -------------------------- 1 file changed, 28 deletions(-) diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index c1cbeba6..c1e7beba 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -16,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. From 4e3c1b46f5efdeab89f1e2ba2e8ba1562b84cfb3 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Tue, 14 Jul 2026 20:59:35 +0200 Subject: [PATCH 10/10] Change to laze creation of the validation model in DynamicValidation to simplify model generation for dataclasses. --- pyaml/validation/schema_builder.py | 106 +++----------------------- pyaml/validation/validation_models.py | 33 ++++++-- 2 files changed, 39 insertions(+), 100 deletions(-) diff --git a/pyaml/validation/schema_builder.py b/pyaml/validation/schema_builder.py index 66679cd1..ea334aa5 100644 --- a/pyaml/validation/schema_builder.py +++ b/pyaml/validation/schema_builder.py @@ -3,8 +3,6 @@ import inspect import logging import types -from dataclasses import MISSING -from dataclasses import Field as DataclassField from datetime import date, datetime, time, timedelta from decimal import Decimal from enum import Enum @@ -79,7 +77,7 @@ def generate_configuration_schema(source: type) -> type[ConfigurationSchema]: source.__module__, ) else: - schema = _configuration_schema_from_class(source) + schema = _configuration_schema_from_constructor(source) logger.debug("Register schema for %s.", class_path) registry.register(class_path, schema) @@ -257,18 +255,21 @@ def _field_kwargs(field_name: str, field_info: Any) -> dict[str, Any]: return kwargs -def _configuration_schema_from_class(cls: type) -> type[ConfigurationSchema]: - """Generate a configuration schema from a class definition. +def _configuration_schema_from_constructor(cls: type) -> type[ConfigurationSchema]: + """ + Generate a configuration schema from a class constructor signature. - Fields are extracted from an explicitly declared constructor when one - exists. Otherwise, fields are extracted from annotations declared directly - on the class, which supports dataclasses before their generated - constructors are available. + The resulting schema contains one field for each constructor parameter, + with annotations and default values preserved. """ + if not isinstance(cls, type): raise TypeError("cls must be a class.") - fields: dict[str, Any] = _fields_from_class_definition(cls, expand_arbitrary_types=True) + fields: dict[str, Any] = _fields_from_constructor_signature( + cls, + expand_arbitrary_types=True, + ) return create_model( _generate_schema_name(cls), @@ -321,88 +322,3 @@ def _fields_from_constructor_signature(cls: type, expand_arbitrary_types: bool = fields[name] = (annotation, default) return fields - - -def _fields_from_class_annotations(cls: type, expand_arbitrary_types: bool = False) -> dict[str, tuple[Any, Any]]: - """Extract field definitions from annotations declared on a class. - - This is primarily used for classes that will be transformed by - ``@dataclass``. During ``__init_subclass__``, the dataclass constructor - has not yet been generated, but the class annotations and defaults are - already available. - - Parameters - ---------- - cls - Class whose directly declared annotations should be inspected. - expand_arbitrary_types - If true, unsupported annotations are expanded recursively into - schema-friendly types. - - Returns - ------- - dict[str, tuple[Any, Any]] - Field definitions suitable for ``pydantic.create_model``. - """ - declared_annotations = cls.__dict__.get("__annotations__", {}) - type_hints = get_type_hints(cls, include_extras=True) - - fields: dict[str, tuple[Any, Any]] = {} - - for name in declared_annotations: - if name in RESERVED_CONFIGURATION_FIELDS: - raise ValueError(f"{cls.__name__} defines reserved field {name!r}, which is owned by ConfigurationSchema.") - - annotation = type_hints.get(name, Any) - - if expand_arbitrary_types: - annotation = _resolve_annotation(annotation) - - default = cls.__dict__.get(name, ...) - - # Support dataclasses.field(...), although @dataclass has not yet - # processed it when this function is called from __init_subclass__. - if isinstance(default, DataclassField): - if default.default is not MISSING: - default = default.default - elif default.default_factory is not MISSING: - default = Field(default_factory=default.default_factory) - else: - default = ... - - fields[name] = (annotation, default) - - return fields - - -def _fields_from_class_definition( - cls: type, - expand_arbitrary_types: bool = False, -) -> dict[str, tuple[Any, Any]]: - """Extract validation fields from a class definition. - - Fields are read from an explicitly defined constructor when one exists. - Otherwise, fields are read from annotations declared directly on the - class. The annotation fallback supports classes that are about to be - transformed by ``@dataclass``. - """ - if "__init__" in cls.__dict__: - return _fields_from_constructor_signature( - cls, - expand_arbitrary_types=expand_arbitrary_types, - ) - - declared_annotations = cls.__dict__.get("__annotations__", {}) - - if declared_annotations: - return _fields_from_class_annotations( - cls, - expand_arbitrary_types=expand_arbitrary_types, - ) - - # No constructor or locally declared annotations. Fall back to the - # effective inherited constructor. - return _fields_from_constructor_signature( - cls, - expand_arbitrary_types=expand_arbitrary_types, - ) diff --git a/pyaml/validation/validation_models.py b/pyaml/validation/validation_models.py index f1e04a9a..4bbe7b29 100644 --- a/pyaml/validation/validation_models.py +++ b/pyaml/validation/validation_models.py @@ -3,12 +3,12 @@ import inspect import logging from abc import ABCMeta -from typing import Any +from typing import Any, ClassVar from pydantic import BaseModel, ConfigDict, create_model from .configuration_models import PyAMLBaseModel -from .schema_builder import _fields_from_class_definition +from .schema_builder import _fields_from_constructor_signature logger = logging.getLogger(__name__) @@ -94,6 +94,23 @@ 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. @@ -108,7 +125,10 @@ class DynamicValidation(metaclass=ValidationMeta): 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): """ @@ -122,10 +142,13 @@ 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]: @@ -145,7 +168,7 @@ def _build_validation_model(cls) -> type[ValidationModel]: logger.debug("Building validation model for %s.", f"{cls.__module__}.{cls.__name__}") - fields: dict[str, Any] = _fields_from_class_definition(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)