diff --git a/pyaml/magnet/cfm_magnet.py b/pyaml/magnet/cfm_magnet.py index 29ac96c5..489017c3 100644 --- a/pyaml/magnet/cfm_magnet.py +++ b/pyaml/magnet/cfm_magnet.py @@ -2,11 +2,12 @@ from ..common import abstract from ..common.abstract import RWMapper -from ..common.element import Element, ElementConfigModel, __pyaml_repr__ +from ..common.element import Element, __pyaml_repr__ from ..common.exception import PyAMLException from ..configuration.factory import ELEMENT_REGISTRY +from ..validation import DynamicValidation, register_schema from .hcorrector import HCorrector -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet from .model import MagnetModel from .octupole import Octupole from .quadrupole import Quadrupole @@ -31,39 +32,64 @@ PYAMLCLASS = "CombinedFunctionMagnet" -class ConfigModel(ElementConfigModel): - mapping: list[list[str]] - """Name mapping for multipoles - (i.e. [[B0,C01A-H],[A0,C01A-H],[B2,C01A-S]])""" - model: MagnetModel | None = None - """Object in charge of converting magnet strenghts to currents""" - - -class CombinedFunctionMagnet(Element): - """CombinedFunctionMagnet class""" - - def __init__(self, cfg: ConfigModel, peer=None): - super().__init__(cfg.name) - self._cfg = cfg - self.model = cfg.model +@register_schema +class CombinedFunctionMagnet(Element, DynamicValidation): + """ + Combined function magnet made up of several virtual single-function magnets. + + This class represents a magnet whose effect is described by multiple multipole + components, such as a corrector, quadrupole, sextupole, or octupole family. + Each entry in ``mapping`` creates a virtual magnet backed by the same + underlying magnet model, and the virtual magnets are exposed as individual + elements while still belonging to the same combined-function object. + + Parameters + ---------- + name : str + Name of the combined-function magnet. + mapping : list[list[str]] + List of ``[multipole, magnet_name]`` pairs. The first entry selects the + virtual magnet type, and the second entry gives the name of the virtual + magnet. + model : MagnetModel | None, optional + Magnet model used to convert strengths to hardware values and vice versa. + description : str | None, optional + Human-readable description of the magnet. + peer : object, optional + Control-system or simulator peer used when attaching the magnet. + + Raises + ------ + PyAMLException + If the mapping is invalid, if an unsupported multipole is requested, or + if the model does not provide the required multipole information. + """ + + def __init__( + self, name: str, mapping: list[list[str]], model: MagnetModel | None = None, description: str | None = None, peer=None + ): + super().__init__(name, None, description) + + self._mapping = mapping + self.model = model self.__virtuals: list[Magnet] = [] - self.__strengths: abstract.ReadWriteFloatArray = None - self.__hardwares: abstract.ReadWriteFloatArray = None + self.__strengths: abstract.ReadWriteFloatArray | None = None + self.__hardwares: abstract.ReadWriteFloatArray | None = None if peer is None: # Configuration part - if self.model is not None and not hasattr(self.model._cfg, "multipoles"): - raise PyAMLException(f"{cfg.name} model: mutipolesfield required for combined function magnet") + if self.model is not None and not hasattr(self.model, "multipoles"): + raise PyAMLException(f"{name} model: mutipoles field required for combined function magnet") idx = 0 self.polynoms = [] - for _idx, m in enumerate(cfg.mapping): + for _idx, m in enumerate(self._mapping): # Check mapping validity if len(m) != 2: raise PyAMLException("Invalid CombinedFunctionMagnet mapping for {m}") if m[0] not in _fmap: raise PyAMLException(m[0] + " not implemented for combined function magnet") - if m[0] not in self.model._cfg.multipoles: + if m[0] not in self.model.multipoles: raise PyAMLException(m[0] + " not found in underlying magnet model") self.polynoms.append(_fmap[m[0]].polynom) # Create the virtual magnet for the correspoding multipole @@ -81,16 +107,16 @@ def get_model_name(self) -> str: """ Returns the model name of this magnet """ - return self._cfg.name + return self._name def __create_virutal_manget(self, name: str, idx: int) -> Magnet: args = {"name": name, "model": self.model} - mVirtual: Magnet = _fmap[idx](MagnetConfigModel(**args)) + mVirtual: Magnet = _fmap[idx](**args) mVirtual.set_model_name(self.get_name()) return mVirtual def nb_multipole(self) -> int: - return len(self._cfg.mapping) + return len(self._mapping) def attach( self, @@ -100,13 +126,13 @@ def attach( ) -> list[Magnet]: l = [] # Attached the CombinedFunctionMagnet itself - nCFM = CombinedFunctionMagnet(self._cfg, peer) + nCFM = CombinedFunctionMagnet(self._name, self._mapping, self.model, self._description, peer) nCFM.__strengths = strengths nCFM.__hardwares = hardwares l.append(nCFM) # Construct a single function magnet for each multipole # of this combined function magnet - for idx, _m in enumerate(self._cfg.mapping): + for idx, _m in enumerate(self._mapping): strength = RWMapper(strengths, idx) hardware = RWMapper(hardwares, idx) if self.model.has_hardware() else None l.append(self.__virtuals[idx].attach(peer, strength, hardware)) diff --git a/pyaml/magnet/corrector.py b/pyaml/magnet/corrector.py index fc9de129..40cc4ad1 100644 --- a/pyaml/magnet/corrector.py +++ b/pyaml/magnet/corrector.py @@ -8,7 +8,7 @@ class RWCorrectorAngle(abstract.ReadWriteFloatScalar): """ Set the angle of a horizontal or vertical corrector. KickAngle sign convention is defined the a global PyAML constant - (see pyaml.common.constant.HORIZONATL_KICK_SIGN). + (see pyaml.common.constant.HORIZONTAL_KICK_SIGN). To change the convention, you have execute the code below prior to everything: import pyaml.common.constants pyaml.common.constants.HORIZONTAL_KICK_SIGN = -1.0 diff --git a/pyaml/magnet/csvcurve.py b/pyaml/magnet/csvcurve.py index 475b5211..df61eb14 100644 --- a/pyaml/magnet/csvcurve.py +++ b/pyaml/magnet/csvcurve.py @@ -1,51 +1,75 @@ -from pathlib import Path - import numpy as np -from pydantic import BaseModel, ConfigDict +from numpy.typing import NDArray +from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException from ..configuration.fileloader import ROOT +from ..validation import DynamicValidation, register_schema from .curve import Curve # Define the main class name for this module PYAMLCLASS = "CSVCurve" -class ConfigModel(BaseModel): +@register_schema +class CSVCurve(Curve, DynamicValidation): """ - Configuration model for CSV curve + Curve loaded from a CSV file. - Parameters - ---------- - file : str - CSV file that contains the curve (n rows, 2 columns) - """ + This class reads a CSV file containing a two-column numeric dataset + representing ``(x, y)`` coordinate pairs. The file is loaded during + initialization and stored internally as a NumPy array. - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + The CSV file must: - file: str + - contain exactly two columns, + - use commas as the delimiter, + - contain values that can be converted to ``float``. + Parameters + ---------- + file : str + Path to the CSV file. Relative paths are resolved using the + project's configured root directory. -class CSVCurve(Curve): - """ - Class for load CSV (x,y) curve + Attributes + ---------- + file : str + Path to the CSV file provided during initialization. + + Raises + ------ + PyAMLException + If the file cannot be parsed as a numeric CSV or if the loaded + data does not have shape ``(n, 2)``. + + Notes + ----- + The loaded curve is stored internally as a NumPy array with shape + ``(n, 2)``, where the first column contains x-values and the second + column contains y-values. The data can be accessed using + :meth:`get_curve`. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, file: str): + self._file = file # Load CSV curve - path = ROOT.expand_path(cfg.file) + path = ROOT.expand_path(self._file) try: self._curve = np.genfromtxt(path, delimiter=",", dtype=float, loose=False) except ValueError as e: - raise PyAMLException(f"CSVCurve(file='{cfg.file}',dtype=float): {str(e)}") from None + raise PyAMLException(f"CSVCurve(file='{self._file}',dtype=float): {str(e)}") from None _s = np.shape(self._curve) if len(_s) != 2 or _s[1] != 2: - raise PyAMLException(f"CSVCurve(file='{cfg.file}',dtype=float):wrong shape (2,2) expected but got {str(_s)}") + raise PyAMLException(f"CSVCurve(file='{self._file}',dtype=float):wrong shape (2,2) expected but got {str(_s)}") + + @property + def file(self): + return self._file - def get_curve(self) -> np.array: + def get_curve(self) -> NDArray[np.float64]: """ Get the curve data. @@ -57,4 +81,4 @@ def get_curve(self) -> np.array: return self._curve def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/pyaml/magnet/csvmatrix.py b/pyaml/magnet/csvmatrix.py index 6726f944..a912ec99 100644 --- a/pyaml/magnet/csvmatrix.py +++ b/pyaml/magnet/csvmatrix.py @@ -1,46 +1,51 @@ -from pathlib import Path - import numpy as np -from pydantic import BaseModel, ConfigDict +from numpy.typing import NDArray +from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException from ..configuration.fileloader import ROOT +from ..validation import DynamicValidation, register_schema from .matrix import Matrix # Define the main class name for this module PYAMLCLASS = "CSVMatrix" -class ConfigModel(BaseModel): +@register_schema +class CSVMatrix(Matrix, DynamicValidation): """ - Configuration model for CSV matrix + Matrix loaded from a CSV file. + + This class reads a CSV file containing numeric values and stores the + resulting matrix as a NumPy array. Parameters ---------- file : str - CSV file that contains the matrix - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - file: str + Path to the CSV file. Relative paths are resolved using the + project's configured root directory. - -class CSVMatrix(Matrix): - """ - Class for loading CSV matrix + Raises + ------ + PyAMLException + If the CSV file cannot be parsed as a numeric array. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, file: str): + self._file = file + # Load CSV matrix - path = ROOT.expand_path(cfg.file) + path = ROOT.expand_path(self._file) try: self._mat = np.genfromtxt(path, delimiter=",", dtype=float, loose=False) except ValueError as e: - raise PyAMLException(f"CSVMatrix(file='{cfg.file}',dtype=float): {str(e)}") from None + raise PyAMLException(f"CSVMatrix(file='{self._file}',dtype=float): {str(e)}") from None + + @property + def file(self): + return self._file - def get_matrix(self) -> np.array: + def get_matrix(self) -> NDArray[np.float64]: """ Get the matrix data. @@ -52,4 +57,4 @@ def get_matrix(self) -> np.array: return self._mat def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/pyaml/magnet/hcorrector.py b/pyaml/magnet/hcorrector.py index 964c4e29..4868ba4a 100644 --- a/pyaml/magnet/hcorrector.py +++ b/pyaml/magnet/hcorrector.py @@ -1,30 +1,28 @@ +import copy +from typing import Self + from ..common import abstract from ..common.constants import HORIZONTAL_KICK_SIGN from ..lattice.polynom_info import PolynomInfo +from ..validation import DynamicValidation, register_schema from .corrector import RWCorrectorAngle -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "HCorrector" -class ConfigModel(MagnetConfigModel): - """Configuration model for Horizontal Corrector magnet.""" - - ... - - -class HCorrector(Magnet): +@register_schema +class HCorrector(Magnet, DynamicValidation): """Horizontal Corrector class""" polynom = PolynomInfo("PolynomB", 0, HORIZONTAL_KICK_SIGN) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) self.__angle = RWCorrectorAngle(self) @property @@ -33,3 +31,17 @@ def angle(self) -> abstract.ReadWriteFloatScalar: Set the kick angle. """ return self.__angle + + def attach( + self, + peer, + strength: abstract.ReadWriteFloatScalar, + hardware: abstract.ReadWriteFloatScalar, + ) -> Self: + """ + Create a new reference to attach this magnet to a simulator + or a control systemand. + """ + obj = super().attach(peer, strength, hardware) + obj.__angle = RWCorrectorAngle(obj) + return obj diff --git a/pyaml/magnet/identity_cfm_model.py b/pyaml/magnet/identity_cfm_model.py index 33d8e05d..98d3846c 100644 --- a/pyaml/magnet/identity_cfm_model.py +++ b/pyaml/magnet/identity_cfm_model.py @@ -1,64 +1,71 @@ import numpy as np -from pydantic import BaseModel, ConfigDict from .. import PyAMLException from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess +from ..validation import DynamicValidation, register_schema from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "IdentityCFMagnetModel" -class ConfigModel(BaseModel): +@register_schema +class IdentityCFMagnetModel(MagnetModel, DynamicValidation): """ - Configuration model for identity combined function magnet model + Identity combined-function magnet model. - Parameters - ---------- - multipoles : list[str] - List of supported functions: A0, B0, A1, B1, etc (i.e. [B0, A1, B2]) - powerconverters : list[DeviceAccess], optional - Power converter devices to apply current - physics : list[DeviceAccess], optional - Magnet devices to apply strength - units : list[str] - List of strength units (i.e. ['rad', 'm-1', 'm-2']) - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - multipoles: list[str] - powerconverters: list[str | None] | None = None - physics: list[str | None] | None = None - units: list[str] + This magnet model maps strengths and hardware values directly to the + underlying devices without applying any conversion. It is intended for cases + where the physics values and hardware values are identical, so the model acts + as an identity transform for all supported multipoles. + Parameters -class IdentityCFMagnetModel(MagnetModel): - """ - Class that map values to underlying devices without conversion + multipoles : list[str] + List of supported multipoles, for example ["B0", "A1", "B2"]. + powerconverters : list[str | None] | None, optional + Names of the power converter devices used for hardware access. + physics : list[str | None] | None, optional + Names of the physics devices used for strength access. + units : list[str] | None, optional + List of units for the supported multipoles. + + Raises + + PyAMLException + If both physics and powerconverters are missing, if both are + provided at the same time, or if the configuration lengths do not match. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + multipoles: list[str], + powerconverters: list[str | None] | None = None, + physics: list[str | None] | None = None, + units: list[str] | None = None, + ): + self.multipoles = multipoles + self._powerconverters = powerconverters + self._physics = physics + self.units = units # Check config - self.__nbFunction: int = len(cfg.multipoles) + self.__nbFunction: int = len(multipoles) - if cfg.physics is None and cfg.powerconverters is None: + if self._physics is None and self._powerconverters is None: raise PyAMLException("Invalid IdentityCFMagnetModel configuration,physics or powerconverters device required") - if cfg.physics is not None and cfg.powerconverters is not None: + if self._physics is not None and self._powerconverters is not None: raise PyAMLException( "Invalid IdentityCFMagnetModel configuration,physics or powerconverters device required but not both" ) - if cfg.physics: - self.__devices = cfg.physics + if self._physics: + self.__devices = self._physics else: - self.__devices = cfg.powerconverters + self.__devices = self._powerconverters self.__nbDev: int = len(self.__devices) - self.__check_len(cfg.units, "units", self.__nbFunction) + self.__check_len(self.units, "units", self.__nbFunction) def __check_len(self, obj, name, expected_len): lgth = len(obj) @@ -74,10 +81,10 @@ def compute_strengths(self, currents: np.array) -> np.array: return currents def get_strength_units(self) -> list[str]: - return self._cfg.units + return self.units def get_hardware_units(self) -> list[str]: - return self._cfg.units + return self.units def get_device_names(self) -> list[str | None]: return self.__devices @@ -86,10 +93,10 @@ def set_magnet_rigidity(self, brho: np.double): pass def has_physics(self) -> bool: - return self._cfg.physics is not None + return self._physics is not None def has_hardware(self) -> bool: - return self._cfg.powerconverters is not None + return self._powerconverters is not None def __repr__(self): return __pyaml_repr__(self) diff --git a/pyaml/magnet/identity_model.py b/pyaml/magnet/identity_model.py index d8e6689c..64270f67 100644 --- a/pyaml/magnet/identity_model.py +++ b/pyaml/magnet/identity_model.py @@ -1,54 +1,63 @@ import numpy as np -from pydantic import BaseModel, ConfigDict from .. import PyAMLException from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "IdentityMagnetModel" -class ConfigModel(BaseModel): +class IdentityMagnetModel(MagnetModel): """ - Configuration model for identity magnet model + Identity magnet model. + + This model maps magnet strengths directly to hardware values without applying + any conversion. It is intended for cases where the physics value and the + hardware value are identical, so the model behaves as an identity transform. Parameters ---------- - powerconverter : DeviceAccess, optional - Power converter device to apply current - physics : DeviceAccess, optional - Magnet device to apply strength - unit : str - Unit of the strength (i.e. 1/m or m-1) - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - powerconverter: str | None = None - physics: str | None = None - unit: str - - -class IdentityMagnetModel(MagnetModel): - """ - Class that map value to underlying device without conversion + powerconverter : str | None, optional + Name of the power converter device used to apply current. + physics : str | None, optional + Name of the physics device used to apply strength. + unit : str | None, optional + Unit of the magnet strength and hardware value, for example ``"1/m"`` or + ``"m-1"``. + + Raises + ------ + PyAMLException + If both ``physics`` and ``powerconverter`` are missing, or if both are + provided at the same time. + + Notes + ----- + The model does not perform any numerical conversion: strengths are returned + as hardware values and hardware values are returned as strengths. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg - self.__unit = cfg.unit - if cfg.physics is None and cfg.powerconverter is None: + def __init__( + self, + powerconverter: str | None = None, + physics: str | None = None, + unit: str | None = None, + ): + self._physics = physics + self._powerconverter = powerconverter + self._unit = unit + + if self._physics is None and self._powerconverter is None: raise PyAMLException("Invalid IdentityMagnetModel configuration,physics or powerconverter device required") - if cfg.physics is not None and cfg.powerconverter is not None: + if self._physics is not None and self._powerconverter is not None: raise PyAMLException( "Invalid IdentityMagnetModel configuration,physics or powerconverter device required but not both" ) - if cfg.physics: - self.__device = cfg.physics + if self._physics: + self.__device = self._physics else: - self.__device = cfg.powerconverter + self.__device = self._powerconverter def compute_hardware_values(self, strengths: np.array) -> np.array: return strengths @@ -57,10 +66,10 @@ def compute_strengths(self, currents: np.array) -> np.array: return currents def get_strength_units(self) -> list[str]: - return [self.__unit] + return [self._unit] def get_hardware_units(self) -> list[str]: - return [self.__unit] + return [self._unit] def get_device_names(self) -> list[str | None]: return [self.__device] @@ -69,10 +78,10 @@ def set_magnet_rigidity(self, brho: np.double): pass def has_physics(self) -> bool: - return self._cfg.physics is not None + return self._physics is not None def has_hardware(self) -> bool: - return self._cfg.powerconverter is not None + return self._powerconverter is not None def __repr__(self): return __pyaml_repr__(self) diff --git a/pyaml/magnet/inline_curve.py b/pyaml/magnet/inline_curve.py index 1ae6e2a3..064d3fe7 100644 --- a/pyaml/magnet/inline_curve.py +++ b/pyaml/magnet/inline_curve.py @@ -1,45 +1,52 @@ -from pathlib import Path - import numpy as np -from pydantic import BaseModel, ConfigDict +from numpy.typing import NDArray +from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException +from ..validation import DynamicValidation, register_schema from .curve import Curve # Define the main class name for this module PYAMLCLASS = "InlineCurve" -class ConfigModel(BaseModel): +@register_schema +class InlineCurve(Curve, DynamicValidation): """ - Configuration model for inline curve + Curve defined directly from an in-memory matrix of (x, y) points. + + This class stores curve data provided as a list of rows, where each row must + contain exactly two values: the x-coordinate and the y-coordinate. The data is + converted to a NumPy array and validated to ensure it has shape ``(n, 2)``. Parameters ---------- mat : list[list[float]] - Curve data (n rows, 2 columns) + Curve data as a two-column matrix. Each row represents one point in the + curve, with ``mat[i][0]`` being the x-value and ``mat[i][1]`` being the + y-value. + + Raises + ------ + PyAMLException + If the provided matrix does not have two columns. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - mat: list[list[float]] - + def __init__(self, mat: list[list[float]]): + self._mat = mat -class InlineCurve(Curve): - """ - Class for load CSV (x,y) curve - """ - - def __init__(self, cfg: ConfigModel): - self._cfg = cfg # Load the curve - self._curve = np.array(self._cfg.mat) + self._curve = np.array(self._mat) _s = np.shape(self._curve) if len(_s) != 2 or _s[1] != 2: - raise PyAMLException(f"InlineCurve(mat='{cfg.mat}',dtype=float): wrong shape (2,2) expected but got {str(_s)}") + raise PyAMLException(f"InlineCurve(mat='{self._mat}',dtype=float): wrong shape (2,2) expected but got {str(_s)}") + + @property + def mat(self): + return self._mat - def get_curve(self) -> np.array: + def get_curve(self) -> NDArray[np.float64]: """ Get the curve data. @@ -51,4 +58,4 @@ def get_curve(self) -> np.array: return self._curve def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/pyaml/magnet/inline_matrix.py b/pyaml/magnet/inline_matrix.py index 25d03c78..534fb14b 100644 --- a/pyaml/magnet/inline_matrix.py +++ b/pyaml/magnet/inline_matrix.py @@ -1,38 +1,39 @@ import numpy as np -from pydantic import BaseModel, ConfigDict +from numpy.typing import NDArray +from ..common.element import __pyaml_repr__ +from ..validation import DynamicValidation, register_schema from .matrix import Matrix # Define the main class name for this module PYAMLCLASS = "InlineMatrix" -class ConfigModel(BaseModel): +@register_schema +class InlineMatrix(Matrix, DynamicValidation): """ - Configuration model for inline matrix + Matrix defined directly from in-memory data. Parameters ---------- mat : list[list[float]] - The matrix - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - mat: list[list[float]] - + Matrix data given as a nested list of numbers. -class InlineMatrix(Matrix): - """ - Class for loading CSV matrix + Attributes + ---------- + _mat : np.ndarray + Internal NumPy representation of the matrix. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, mat: list[list[float]]): # Load the matrix - self._mat = np.array(self._cfg.mat) + self._mat = np.array(mat) + + @property + def mat(self): + return self._mat - def get_matrix(self) -> np.array: + def get_matrix(self) -> NDArray[np.float64]: """ Get the matrix data. @@ -44,4 +45,4 @@ def get_matrix(self) -> np.array: return self._mat def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/pyaml/magnet/linear_cfm_model.py b/pyaml/magnet/linear_cfm_model.py index 32c0cc80..3dede9aa 100644 --- a/pyaml/magnet/linear_cfm_model.py +++ b/pyaml/magnet/linear_cfm_model.py @@ -1,9 +1,8 @@ import numpy as np -from pydantic import BaseModel, ConfigDict from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException -from ..control.deviceaccess import DeviceAccess +from ..validation import DynamicValidation, register_schema from .curve import Curve from .matrix import Matrix from .model import MagnetModel @@ -12,101 +11,113 @@ PYAMLCLASS = "LinearCFMagnetModel" -class ConfigModel(BaseModel): +@register_schema +class LinearCFMagnetModel(MagnetModel, DynamicValidation): """ - Configuration model for linear combined function magnet model + Linear combined-function magnet model. + + This model converts between magnet strengths and hardware currents using + precomputed excitation curves, optional calibration and pseudo-current + corrections, and an optional coupling matrix to separate or combine multipole + responses. It is intended for combined-function magnets where several + multipoles share a common hardware representation. Parameters - ---------- + multipoles : list[str] - List of supported functions: A0, B0, A1, B1, etc (i.e. [B0, A1, B2]) + Names of the supported multipoles, for example ["B0", "A1", "B2"]. curves : list[Curve] - Excitation curves, 1 curve per function - calibration_factors : list[float], optional - Correction factor applied to curves, 1 factor per function. - Default: ones - calibration_offsets : list[float], optional - Correction offset applied to curves, 1 offset per function. - Default: zeros - pseudo_factors : list[float], optional - Factors applied to 'pseudo currents', 1 factor per function. - Default: ones - pseudo_offsets : list[float], optional - Offsets applied to 'pseudo currents', 1 offset per function. - Default: zeros - powerconverters : list[str] - List of power converter devices to apply currents (can be different - from number of functions) - matrix : Matrix, optional - n x m matrix (n rows for n functions, m columns for m currents) - to handle multipoles separation. Default: Identity - units : list[str] - List of strength units (i.e. ['rad', 'm-1', 'm-2']) + Excitation curves, one per multipole. + powerconverters : list[str | None] + Names of the power converter devices associated with the hardware currents. hardware_units : list[str] - List of hardware units (i.e. 'A' , 'V') - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - multipoles: list[str] - curves: list[Curve] - calibration_factors: list[float] = None - calibration_offsets: list[float] = None - pseudo_factors: list[float] = None - pseudo_offsets: list[float] = None - powerconverters: list[str | None] - matrix: Matrix = None - units: list[str] - hardware_units: list[str] - - -class LinearCFMagnetModel(MagnetModel): - """ - Class providing a simple linear model for combined function magnets. A matrix - can handle separation of multipoles. A pseudo current is a linear combination - of power supply currents associated to a single function. + Units of the hardware variables, one per power converter. + calibration_factors : list[float] | None, optional + Multiplicative correction factors applied to the excitation curves. + Defaults to ones. + calibration_offsets : list[float] | None, optional + Additive correction offsets applied to the excitation curves. + Defaults to zeros. + pseudo_factors : list[float] | None, optional + Multiplicative factors applied to pseudo currents. + Defaults to ones. + pseudo_offsets : list[float] | None, optional + Additive offsets applied to pseudo currents. + Defaults to zeros. + matrix : Matrix | None, optional + Coupling matrix mapping power-supply currents to pseudo currents. + Defaults to the identity matrix. + units : list[str] | None, optional + Strength units, one per multipole. + + Raises + + PyAMLException + If the list lengths do not match, if the matrix has the wrong shape, or + if the configuration is otherwise inconsistent. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + multipoles: list[str], + curves: list[Curve], + powerconverters: list[str | None], + hardware_units: list[str], + calibration_factors: list[float] | None = None, + calibration_offsets: list[float] | None = None, + pseudo_factors: list[float] | None = None, + pseudo_offsets: list[float] | None = None, + matrix: Matrix | None = None, + units: list[str] | None = None, + ): + self.multipoles = multipoles + self._curves = curves + self._powerconverters = powerconverters + self.hardware_units = hardware_units + self._calibration_factors = calibration_factors + self._calibration_offsets = calibration_offsets + self._pseudo_factors = pseudo_factors + self._pseudo_offsets = pseudo_offsets + self._matrix = matrix + self.units = units self._brho = np.nan # Check config - self.__nbFunction: int = len(cfg.multipoles) - self.__nbPS: int = len(cfg.powerconverters) + self.__nbFunction: int = len(self.multipoles) + self.__nbPS: int = len(self._powerconverters) - if cfg.calibration_factors is None: + if self._calibration_factors is None: self.__calibration_factors = np.ones(self.__nbFunction) else: - self.__calibration_factors = cfg.calibration_factors + self.__calibration_factors = self._calibration_factors - if cfg.calibration_offsets is None: + if self._calibration_offsets is None: self.__calibration_offsets = np.zeros(self.__nbFunction) else: - self.__calibration_offsets = cfg.calibration_offsets + self.__calibration_offsets = self._calibration_offsets - if cfg.pseudo_factors is None: + if self._pseudo_factors is None: self.__pf = np.ones(self.__nbFunction) else: - self.__pf = cfg.pseudo_factors + self.__pf = self._pseudo_factors - if cfg.pseudo_offsets is None: + if self._pseudo_offsets is None: self.__po = np.zeros(self.__nbFunction) else: - self.__po = cfg.pseudo_factors + self.__po = self._pseudo_factors self.__check_len(self.__calibration_factors, "calibration_factors", self.__nbFunction) self.__check_len(self.__calibration_offsets, "calibration_offsets", self.__nbFunction) self.__check_len(self.__pf, "pseudo_factors", self.__nbFunction) self.__check_len(self.__po, "pseudo_offsets", self.__nbFunction) - self.__check_len(cfg.units, "units", self.__nbFunction) - self.__check_len(cfg.hardware_units, "hardware_units", self.__nbPS) - self.__check_len(cfg.curves, "curves", self.__nbFunction) + self.__check_len(self.units, "units", self.__nbFunction) + self.__check_len(self.hardware_units, "hardware_units", self.__nbPS) + self.__check_len(self._curves, "curves", self.__nbFunction) - if cfg.matrix is None: + if self._matrix is None: self.__matrix = np.identity(self.__nbFunction) else: - self.__matrix = cfg.matrix.get_matrix() + self.__matrix = self._matrix.get_matrix() _s = np.shape(self.__matrix) @@ -119,7 +130,7 @@ def __init__(self, cfg: ConfigModel): self.__rcurves = [] # Apply factor and offset - for idx, c in enumerate(cfg.curves): + for idx, c in enumerate(self._curves): self.__curves.append(c.get_curve()) self.__curves[idx][:, 1] *= self.__calibration_factors[idx] self.__curves[idx][:, 1] += self.__calibration_offsets[idx] @@ -150,13 +161,13 @@ def compute_strengths(self, currents: np.array) -> np.array: return _strength def get_strength_units(self) -> list[str]: - return self._cfg.units + return self.units def get_hardware_units(self) -> list[str]: - return self._cfg.hardware_units + return self.hardware_units def get_device_names(self) -> list[str | None]: - return self._cfg.powerconverters + return self._powerconverters def set_magnet_rigidity(self, brho: np.double): self._brho = brho diff --git a/pyaml/magnet/linear_model.py b/pyaml/magnet/linear_model.py index 3b9864e9..0f91a81f 100644 --- a/pyaml/magnet/linear_model.py +++ b/pyaml/magnet/linear_model.py @@ -1,7 +1,7 @@ import numpy as np -from pydantic import BaseModel, ConfigDict from ..common.element import __pyaml_repr__ +from ..validation import DynamicValidation, register_schema from .curve import Curve from .model import MagnetModel @@ -9,60 +9,66 @@ PYAMLCLASS = "LinearMagnetModel" -class ConfigModel(BaseModel): +@register_schema +class LinearMagnetModel(MagnetModel, DynamicValidation): """ - Linear magnet model. + Linear magnet model for a single magnet function. + + This model converts between magnet strengths and hardware currents using a + single excitation curve or, if no curve is provided, a linear scaling with an + optional offset. It is intended for single-function magnets such as correctors + or other elements that use one current channel. Parameters ---------- - curve : Curve or None, optional - Curve object used for interpolation. By default, - identity curve is used. - powerconverter : str or None, optional - Power converter device to apply currrent + unit : str + Unit of the magnet strength, for example ``"1/m"`` or ``"m-1"``. + hardware_unit : str + Unit of the hardware value, for example ``"A"`` or ``"V"``. + curve : Curve | None, optional + Excitation curve used for interpolation. If omitted, a linear conversion + is used instead. + powerconverter : str | None, optional + Name of the power converter device used to apply current. calibration_factor : float, optional - Correction factor applied to the curve. Default: 1.0 + Multiplicative correction applied to the curve or linear scaling. + Default is ``1.0``. calibration_offset : float, optional - Correction offset applied to the curve. Default: 0.0 + Additive correction applied to the curve or linear scaling. + Default is ``0.0``. crosstalk : float, optional - Crosstalk factor. Default: 1.0 - unit : str - Unit of the strength (i.e. 1/m or m-1) - hardware_unit : str - Hardware units (i.e. 'A' , 'V') - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - curve: Curve | None = None - powerconverter: str | None - calibration_factor: float = 1.0 - calibration_offset: float = 0.0 - crosstalk: float = 1.0 - unit: str - hardware_unit: str - - -class LinearMagnetModel(MagnetModel): - """ - Class that handle manget current/strength conversion using - linear interpolation for a single function magnet + Crosstalk factor applied together with the calibration factor. + Default is ``1.0``. + + Notes + ----- + If a curve is provided, the model interpolates between strength and current + values using the curve and its inverse. If no curve is provided, the model + uses a simple linear relation with the stored scaling factor and offset. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg - if self._cfg.curve: - self.__curve = cfg.curve.get_curve() - self.__curve[:, 1] = self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + cfg.calibration_offset + def __init__( + self, + unit: str, + hardware_unit: str, + curve: Curve | None = None, + powerconverter: str | None = None, + calibration_factor: float = 1.0, + calibration_offset: float = 0.0, + crosstalk: float = 1.0, + ): + if curve: + self.__curve = curve.get_curve() + self.__curve[:, 1] = self.__curve[:, 1] * calibration_factor * crosstalk + calibration_offset self.__rcurve = Curve.inverse(self.__curve) else: self.__curve = None self.__rcurve = None - self.__g = cfg.calibration_factor * cfg.crosstalk - self.__o = cfg.calibration_offset - self.__strength_unit = cfg.unit - self.__hardware_unit = cfg.hardware_unit - self.__ps = cfg.powerconverter + self.__g = calibration_factor * crosstalk + self.__o = calibration_offset + self.__strength_unit = unit + self.__hardware_unit = hardware_unit + self.__ps = powerconverter self.__brho = np.nan def compute_hardware_values(self, strengths: np.array) -> np.array: diff --git a/pyaml/magnet/linear_serialized_model.py b/pyaml/magnet/linear_serialized_model.py index bb2eb6f4..b1d88ca5 100644 --- a/pyaml/magnet/linear_serialized_model.py +++ b/pyaml/magnet/linear_serialized_model.py @@ -1,57 +1,17 @@ import numpy as np -from pydantic import BaseModel, ConfigDict from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException -from ..control.deviceaccess import DeviceAccess +from ..validation import DynamicValidation, register_schema from .curve import Curve -from .inline_curve import ConfigModel as InlineCurveModel from .inline_curve import InlineCurve -from .linear_model import ConfigModel as LinearConfigModel from .linear_model import LinearMagnetModel -from .matrix import Matrix from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "LinearSerializedMagnetModel" -class ConfigModel(BaseModel): - """ - Configuration model for linear serialized magnet model - - Parameters - ---------- - curves : Curve or list[Curve] - Excitation curves, 1 curve for all or 1 curve per magnet - calibration_factors : float or list[float], optional - Correction factor applied to curves, 1 factor for all or 1 factor per magnet. - Default: ones - calibration_offsets : float or list[float], optional - Correction offset applied to curves, 1 offset for all or 1 offset per magnet. - Default: zeros - crosstalk : float or list[float], optional - Crosstalk factors. Default: 1.0 - powerconverter : str - The hardware can be a single power supply or a list of power supplies. - If a list is provided, the same value will be affected to all of them - unit : str - Strength unit: rad, m-1, m-2 - hardware_unit : str - Hardware unit (i.e. 'A' , 'V') - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - curves: Curve | list[Curve] - calibration_factors: float | list[float] = None - calibration_offsets: float | list[float] = None - crosstalk: float | list[float] = 1.0 - powerconverter: str | None - unit: str - hardware_unit: str - - def _get_length(elem) -> int: if elem is None: return 0 @@ -82,47 +42,98 @@ def _check_len(obj, name, expected_length): ) -class LinearSerializedMagnetModel(MagnetModel): +@register_schema +class LinearSerializedMagnetModel(MagnetModel, DynamicValidation): """ - Class providing a simple linear model for combined function magnets. A matrix - can handle separation of multipoles. A pseudo current is a linear combination - of power supply currents associated to a single function. + Linear model for a serialized combined-function magnet. + + This model wraps one or more :class:`LinearMagnetModel` instances and + represents a combined-function magnet as a set of sub-models that share a + common hardware readout. The excitation curve may be provided either as a + single :class:`Curve` instance, which is replicated for all magnets, or as a + list of curves with one entry per magnet. + + The model supports per-magnet calibration factors, calibration offsets, and + crosstalk values. Strengths are converted to a single hardware current by + evaluating the individual sub-models and averaging their corresponding + hardware values. Conversely, a hardware current is mapped back to a + strength value for each sub-model. + + Parameters + ---------- + curves : Curve or list[Curve] + Excitation curve(s) used by the magnet model. A single curve is reused + for all magnets, while a list must contain one curve per magnet. + calibration_factors : float or list[float], optional + Multiplicative correction factor applied to each curve. + If omitted, defaults to 1.0 for all magnets. + calibration_offsets : float or list[float], optional + Additive correction offset applied to each curve. + If omitted, defaults to 0.0 for all magnets. + crosstalk : float or list[float], optional + Crosstalk factor for each magnet. If omitted, defaults to 1.0. + powerconverter : str, optional + Name of the associated power converter or power converter group. + unit : str, optional + Strength unit, such as ``rad``, ``m-1`` or ``m-2``. + hardware_unit : str, optional + Hardware unit, typically ``A`` or ``V``. + + Notes + ----- + The number of magnets is inferred from the longest list among the supplied + configuration values. Scalars are expanded to match that length. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + curves: Curve | list[Curve], + calibration_factors: float | list[float] | None = None, + calibration_offsets: float | list[float] | None = None, + crosstalk: float | list[float] = 1.0, + powerconverter: str | None = None, + unit: str | None = None, + hardware_unit: str | None = None, + ): self.__brho = np.nan + self._curves = curves + self._calibration_factors = calibration_factors + self._calibration_offsets = calibration_offsets + self._crosstalk = crosstalk + self._powerconverter = powerconverter + self._unit = unit + self._hardware_unit = hardware_unit # Check config - self.__nbMagnets: int = _get_max_length(cfg.curves, cfg.calibration_factors, cfg.calibration_offsets, cfg.crosstalk) + self.__nbMagnets: int = _get_max_length(curves, calibration_factors, calibration_offsets, crosstalk) self.__calibration_factors = np.ones(self.__nbMagnets) self.__calibration_offsets = np.ones(self.__nbMagnets) self.__crosstalk = np.ones(self.__nbMagnets) - self.__curves = _to_list_of_length(self._cfg.curves, self.__nbMagnets) + self.__curves = _to_list_of_length(curves, self.__nbMagnets) self.__sub_models: list[LinearMagnetModel] = [] def __initialize(self): - if self._cfg.calibration_factors is None: + if self._calibration_factors is None: self.__calibration_factors = np.ones(self.__nbMagnets) else: - self.__calibration_factors = _to_list_of_length(self._cfg.calibration_factors, self.__nbMagnets) + self.__calibration_factors = _to_list_of_length(self._calibration_factors, self.__nbMagnets) - if self._cfg.calibration_offsets is None: + if self._calibration_offsets is None: self.__calibration_offsets = np.zeros(self.__nbMagnets) else: - self.__calibration_offsets = _to_list_of_length(self._cfg.calibration_offsets, self.__nbMagnets) + self.__calibration_offsets = _to_list_of_length(self._calibration_offsets, self.__nbMagnets) - if self._cfg.crosstalk is None: + if self._crosstalk is None: self.__crosstalk = np.zeros(self.__nbMagnets) else: - self.__crosstalk = _to_list_of_length(self._cfg.crosstalk, self.__nbMagnets) - self.__curves = _to_list_of_length(self._cfg.curves, self.__nbMagnets) - if isinstance(self._cfg.curves, list): - self.__curves = self._cfg.curves + self.__crosstalk = _to_list_of_length(self._crosstalk, self.__nbMagnets) + self.__curves = _to_list_of_length(self._curves, self.__nbMagnets) + if isinstance(self._curves, list): + self.__curves = self._curves else: self.__curves: list[Curve] = [] for _ in range(self.__nbMagnets): - curve = InlineCurve(InlineCurveModel(mat=self._cfg.curves.get_curve())) + curve = InlineCurve(mat=self._curves.get_curve()) self.__curves.append(curve) _check_len(self.__calibration_factors, "calibration_factors", self.__nbMagnets) @@ -132,16 +143,16 @@ def __initialize(self): self.__sub_models: list[LinearMagnetModel] = [] for magnet_idx in range(self.__nbMagnets): - sub_model = LinearConfigModel( + sub_model = dict( curve=self.__curves[magnet_idx], calibration_factor=self.__calibration_factors[magnet_idx], calibration_offset=self.__calibration_offsets[magnet_idx], crosstalk=self.__crosstalk[magnet_idx], - powerconverter=self._cfg.powerconverter, - unit=self._cfg.unit, - hardware_unit=self._cfg.hardware_unit, + powerconverter=self._powerconverter, + unit=self._unit, + hardware_unit=self._hardware_unit, ) - self.__sub_models.append(LinearMagnetModel(sub_model)) + self.__sub_models.append(LinearMagnetModel(**sub_model)) def set_number_of_magnets(self, nb_magnets: int): self.__nbMagnets = nb_magnets @@ -159,13 +170,13 @@ def compute_strengths(self, currents: np.array) -> np.array: return np.array([model.compute_strengths([current])[0] for model in self.__sub_models]) def get_strength_units(self) -> list[str]: - return [self._cfg.unit] * self.__nbMagnets + return [self._unit] * self.__nbMagnets def get_hardware_units(self) -> list[str]: - return [self._cfg.hardware_unit] + return [self._hardware_unit] def get_device_names(self) -> list[str | None]: - return [self._cfg.powerconverter] + return [self._powerconverter] def set_magnet_rigidity(self, brho: np.double): self.__brho = brho diff --git a/pyaml/magnet/magnet.py b/pyaml/magnet/magnet.py index 2dca34c9..ae22deb1 100644 --- a/pyaml/magnet/magnet.py +++ b/pyaml/magnet/magnet.py @@ -1,36 +1,23 @@ +import copy +from typing import Self + +import numpy as np from scipy.constants import speed_of_light from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel +from ..common.element import Element from .model import MagnetModel -try: - from typing import Self # Python 3.11+ -except ImportError: - from typing_extensions import Self # Python 3.10 and earlier -import numpy as np - - -class MagnetConfigModel(ElementConfigModel): - """ - Configuration model for magnet elements. - - Attributes - ---------- - model : MagnetModel or None, optional - Object in charge of converting magnet strengths to power supply values - """ - - model: MagnetModel | None = None - class Magnet(Element): """ Class providing access to one magnet of a physical or simulated lattice """ - def __init__(self, name: str, model: MagnetModel = None): + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): """ Construct a magnet @@ -41,7 +28,7 @@ def __init__(self, name: str, model: MagnetModel = None): model : MagnetModel Magnet model in charge of computing coil(s) current """ - super().__init__(name) + super().__init__(name, lattice_names, description) self.__model = model self.__strength: abstract.ReadWriteFloatScalar = None self.__hardware: abstract.ReadWriteFloatScalar = None @@ -85,7 +72,7 @@ def attach( Create a new reference to attach this magnet to a simulator or a control systemand. """ - obj = self.__class__(self._cfg) + obj = copy.copy(self) obj.__modelName = self.__modelName obj.__strength = strength obj.__hardware = hardware diff --git a/pyaml/magnet/octupole.py b/pyaml/magnet/octupole.py index d9710a09..e41c88e0 100644 --- a/pyaml/magnet/octupole.py +++ b/pyaml/magnet/octupole.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "Octupole" -class ConfigModel(MagnetConfigModel): - """Configuration model for Octupole magnet.""" - - ... - - -class Octupole(Magnet): +@register_schema +class Octupole(Magnet, DynamicValidation): """Octupole class""" polynom = PolynomInfo("PolynomB", 3) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/quadrupole.py b/pyaml/magnet/quadrupole.py index 3c720ab6..6ddb252f 100644 --- a/pyaml/magnet/quadrupole.py +++ b/pyaml/magnet/quadrupole.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "Quadrupole" -class ConfigModel(MagnetConfigModel): - """Configuration model for Quadrupole magnet.""" - - ... - - -class Quadrupole(Magnet): +@register_schema +class Quadrupole(Magnet, DynamicValidation): """Quadrupole class""" polynom = PolynomInfo("PolynomB", 1) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/serialized_magnet.py b/pyaml/magnet/serialized_magnet.py index 3a0bf0fc..4aa72c20 100644 --- a/pyaml/magnet/serialized_magnet.py +++ b/pyaml/magnet/serialized_magnet.py @@ -3,30 +3,25 @@ from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel, __pyaml_repr__ +from ..common.element import Element, __pyaml_repr__ from ..configuration.factory import ELEMENT_REGISTRY -from ..control.deviceaccess import DeviceAccess +from ..validation import DynamicValidation, register_schema from .function_mapping import function_map -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "SerializedMagnets" -class ConfigModel(ElementConfigModel): - function: str - """List of magnets""" - elements: list[str] | str - """List of magnets""" - model: MagnetModel | None = None - """Object in charge of converting magnet strengths to currents""" - - class ReadWriteSerializedStrengths(abstract.ReadWriteFloatScalar): - def __init__(self, cfg: ConfigModel, elements: list[abstract.ReadWriteFloatScalar]): + def __init__( + self, + elements: list[abstract.ReadWriteFloatScalar], + model: MagnetModel | None = None, + ): self.elements = elements - self._cfg = cfg + self.model = model def get(self) -> float: return sum([elem.get() for elem in self.elements]) @@ -38,10 +33,10 @@ def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") def unit(self) -> str: - return self._cfg.model.get_strength_units()[0] + return self.model.get_strength_units()[0] def get_model(self) -> MagnetModel: - return self._cfg.model + return self.model def get_elements(self): return self.elements @@ -51,49 +46,83 @@ def set_magnet_rigidity(self, brho: np.double): class ReadWriteSerializedHardwares(ReadWriteSerializedStrengths): - def __init__(self, cfg: ConfigModel, elements: list[abstract.ReadWriteFloatScalar]): - super().__init__(cfg, elements) + def __init__( + self, + elements: list[abstract.ReadWriteFloatScalar], + model: MagnetModel | None = None, + ): + super().__init__(elements, model) def unit(self) -> str: - return self._cfg.model.get_hardware_units()[0] + return self.model.get_hardware_units()[0] def set_magnet_rigidity(self, brho: np.double): [element.set_magnet_rigidity(brho) for element in self.elements] -class SerializedMagnets(Element): +@register_schema +class SerializedMagnets(Element, DynamicValidation): """ - Class managing serialized magnets: a set of magnet with the same set point. - The set point is usually managed by only one power supply but it can be covered by several ones. - If several power supplies + Serialized group of magnets that share the same set point. + This class represents a set of magnets that are controlled together as a + single logical device. The serialized magnets may be driven by one power + supply or by several power supplies, but they share a common physics or + hardware set point through a combined read/write interface. Parameters ---------- - cfg : ConfigModel - Configuration object TODO: to describe + name : str + Name of the serialized magnet group. + function : str + Magnet function identifier used to select the concrete virtual magnet type. + elements : list[str] | str + Names of the individual magnets in the group. + model : MagnetModel | None, optional + Magnet model used to convert between strengths and hardware values. + description : str | None, optional + Human-readable description of the serialized magnet group. + peer : object, optional + Control-system or simulator peer used when attaching the magnet group. Raises ------ - pyaml.PyAMLException - In case of wrong initialization + PyAMLException + If the requested function is not implemented or if the configuration is + invalid. + + Notes + ----- + The serialized group stores a virtual magnet for each underlying element. When + attached, each virtual magnet is bound to the same peer and the group exposes + aggregate strength and hardware accessors. """ - def __init__(self, cfg: ConfigModel, peer=None): - super().__init__(cfg.name) - self._cfg = cfg - self.model = cfg.model + def __init__( + self, + name: str, + function: str, + elements: list[str] | str, + model: MagnetModel | None = None, + description: str | None = None, + peer=None, + ): + super().__init__(name, None, description) + + self.function = function + self.model = model + self.polynom = None self.__strengths = None self.__hardwares = None self.__virtuals: list[Magnet] = [] - self.__elements = cfg.elements if isinstance(cfg.elements, list) else [cfg.elements] + self.__elements = elements if isinstance(elements, list) else [elements] self.model.set_number_of_magnets(len(self.__elements)) if peer is None: # Configuration part - self.polynom = function_map[self._cfg.function].polynom - if self._cfg.function not in function_map: - raise PyAMLException(self._cfg.function + " not implemented for serialized magnet") + self.polynom = function_map[self.function].polynom + if self.function not in function_map: + raise PyAMLException(self.function + " not implemented for serialized magnet") for element in self.__elements: # Check mapping validity # Create the virtual magnet for the corresponding magnet @@ -107,7 +136,7 @@ def __init__(self, cfg: ConfigModel, peer=None): def __create_virtual_magnet(self, name: str) -> Magnet: args = {"name": name, "model": self.model} - virtual: Magnet = function_map[self._cfg.function](MagnetConfigModel(**args)) + virtual: Magnet = function_map[self.function](**args) virtual.set_model_name(self.get_name()) return virtual @@ -124,9 +153,9 @@ def attach( hardwares: list[abstract.ReadWriteFloatScalar], ) -> list[Magnet]: l = [] - n_ser_mag = SerializedMagnets(self._cfg, peer) - n_ser_mag.__strengths = ReadWriteSerializedStrengths(self._cfg, strengths) - n_ser_mag.__hardwares = ReadWriteSerializedHardwares(self._cfg, hardwares) + n_ser_mag = SerializedMagnets(self._name, self.function, self.__elements, self.model, self.description, peer) + n_ser_mag.__strengths = ReadWriteSerializedStrengths(strengths, self.model) + n_ser_mag.__hardwares = ReadWriteSerializedHardwares(hardwares, self.model) l.append(n_ser_mag) # Construct a single magnet for each magnet. sub_magnets: list[Magnet] = [] diff --git a/pyaml/magnet/sextupole.py b/pyaml/magnet/sextupole.py index e8218ca3..2cf97ab9 100644 --- a/pyaml/magnet/sextupole.py +++ b/pyaml/magnet/sextupole.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "Sextupole" -class ConfigModel(MagnetConfigModel): - """Configuration model for Sextupole magnet.""" - - ... - - -class Sextupole(Magnet): +@register_schema +class Sextupole(Magnet, DynamicValidation): """Sextupole class""" polynom = PolynomInfo("PolynomB", 2) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/skewoctu.py b/pyaml/magnet/skewoctu.py index 109edd78..e25e2979 100644 --- a/pyaml/magnet/skewoctu.py +++ b/pyaml/magnet/skewoctu.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "SkewOctu" -class ConfigModel(MagnetConfigModel): - """Configuration model for SkewOctu magnet.""" - - ... - - -class SkewOctu(Magnet): +@register_schema +class SkewOctu(Magnet, DynamicValidation): """SkewOctu class""" polynom = PolynomInfo("PolynomA", 3) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/skewquad.py b/pyaml/magnet/skewquad.py index f79e2c30..505ac03b 100644 --- a/pyaml/magnet/skewquad.py +++ b/pyaml/magnet/skewquad.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "SkewQuad" -class ConfigModel(MagnetConfigModel): - """Configuration model for SkewQuad magnet.""" - - ... - - -class SkewQuad(Magnet): +@register_schema +class SkewQuad(Magnet, DynamicValidation): """SkewQuad class""" polynom = PolynomInfo("PolynomA", 1) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/skewsext.py b/pyaml/magnet/skewsext.py index 640d996e..ad68cbd9 100644 --- a/pyaml/magnet/skewsext.py +++ b/pyaml/magnet/skewsext.py @@ -1,24 +1,19 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from ..validation import DynamicValidation, register_schema +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "SkewSext" -class ConfigModel(MagnetConfigModel): - """Configuration model for SkewSext magnet.""" - - ... - - -class SkewSext(Magnet): +@register_schema +class SkewSext(Magnet, DynamicValidation): """SkewSext class""" polynom = PolynomInfo("PolynomA", 2) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) diff --git a/pyaml/magnet/spline_model.py b/pyaml/magnet/spline_model.py index 69fd18aa..e9fed917 100644 --- a/pyaml/magnet/spline_model.py +++ b/pyaml/magnet/spline_model.py @@ -3,7 +3,7 @@ from scipy.interpolate import make_smoothing_spline from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess +from ..validation import DynamicValidation, register_schema from .curve import Curve from .model import MagnetModel @@ -11,60 +11,67 @@ PYAMLCLASS = "SplineMagnetModel" -class ConfigModel(BaseModel): +@register_schema +class SplineMagnetModel(MagnetModel, DynamicValidation): """ - Configuration model for spline magnet model + Magnet model that converts between strength and hardware current using + spline interpolation. + + The model represents a single-function magnet and builds two smoothing + splines from the supplied excitation curve: + + - a forward spline for converting hardware current to magnet strength + - an inverse spline for converting magnet strength to hardware current + + The input curve is first scaled by ``calibration_factor`` and ``crosstalk`` + and shifted by ``calibration_offset`` before the splines are created. Parameters ---------- curve : Curve - Curve object used for interpolation - powerconverter : DeviceAccess, optional - Power converter device to apply current + Excitation curve used for interpolation. + powerconverter : str | None, optional + Name of the associated power converter device. calibration_factor : float, optional - Correction factor applied to the curve. Default: 1.0 + Multiplicative correction applied to the curve. Default is ``1.0``. calibration_offset : float, optional - Correction offset applied to the curve. Default: 0.0 + Additive correction applied to the curve. Default is ``0.0``. crosstalk : float, optional - Crosstalk factor. Default: 1.0 - unit : str - Unit of the strength (i.e. 1/m or m-1) - hardware_unit : str - Hardware units (i.e. 'A' , 'V') + Crosstalk factor applied to the curve. Default is ``1.0``. + unit : str | None, optional + Strength unit, such as ``m-1`` or ``m-2``. + hardware_unit : str | None, optional + Hardware unit, such as ``A`` or ``V``. alpha : float, optional - Regularization parameter (alpha >= 0). alpha = 0 means the interpolation - passes through all the points of the curve. Default: 0.0 - """ + Smoothing parameter passed to :func:`scipy.interpolate.make_smoothing_spline`. + ``alpha = 0`` gives exact interpolation through the data points. - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - curve: Curve - powerconverter: str | None - calibration_factor: float = 1.0 - calibration_offset: float = 0.0 - crosstalk: float = 1.0 - unit: str - hardware_unit: str - alpha: float = 0.0 - - -class SplineMagnetModel(MagnetModel): - """ - Class that handle manget current/strength conversion using - spline interpolation for a single function magnet + Notes + ----- + The magnet rigidity ``brho`` must be set with :meth:`set_magnet_rigidity` + before using the conversion methods. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg - self.__curve = cfg.curve.get_curve() - self.__curve[:, 1] = self.__curve[:, 1] * cfg.calibration_factor * cfg.crosstalk + cfg.calibration_offset + def __init__( + self, + curve: Curve, + powerconverter: str | None = None, + calibration_factor: float = 1.0, + calibration_offset: float = 0.0, + crosstalk: float = 1.0, + unit: str | None = None, + hardware_unit: str | None = None, + alpha: float = 0.0, + ): + self.__curve = curve.get_curve() + self.__curve[:, 1] = self.__curve[:, 1] * calibration_factor * crosstalk + calibration_offset rcurve = Curve.inverse(self.__curve) - self.__strength_unit = cfg.unit - self.__hardware_unit = cfg.hardware_unit + self.__strength_unit = unit + self.__hardware_unit = hardware_unit self.__brho = np.nan - self.__ps = cfg.powerconverter - self.__spl = make_smoothing_spline(self.__curve[:, 0], self.__curve[:, 1], lam=cfg.alpha) - self.__rspl = make_smoothing_spline(rcurve[:, 0], rcurve[:, 1], lam=cfg.alpha) + self.__ps = powerconverter + self.__spl = make_smoothing_spline(self.__curve[:, 0], self.__curve[:, 1], lam=alpha) + self.__rspl = make_smoothing_spline(rcurve[:, 0], rcurve[:, 1], lam=alpha) def compute_hardware_values(self, strengths: np.array) -> np.array: _current = self.__rspl(strengths[0] * self.__brho) diff --git a/pyaml/magnet/vcorrector.py b/pyaml/magnet/vcorrector.py index cf3a0d08..2a4d9a91 100644 --- a/pyaml/magnet/vcorrector.py +++ b/pyaml/magnet/vcorrector.py @@ -1,29 +1,27 @@ +import copy +from typing import Self + from ..common import abstract from ..lattice.polynom_info import PolynomInfo +from ..validation import DynamicValidation, register_schema from .corrector import RWCorrectorAngle -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet +from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "VCorrector" -class ConfigModel(MagnetConfigModel): - """Configuration model for Vertical Corrector magnet.""" - - ... - - -class VCorrector(Magnet): +@register_schema +class VCorrector(Magnet, DynamicValidation): """Vertical Corrector class""" polynom = PolynomInfo("PolynomA", 0) - def __init__(self, cfg: ConfigModel): - super().__init__( - cfg.name, - cfg.model if hasattr(cfg, "model") else None, - ) - self._cfg = cfg + def __init__( + self, name: str, model: MagnetModel | None = None, lattice_names: str | None = None, description: str | None = None + ): + super().__init__(name, model, lattice_names, description) self.__angle = RWCorrectorAngle(self) @property @@ -32,3 +30,17 @@ def angle(self) -> abstract.ReadWriteFloatScalar: Set the kick angle. """ return self.__angle + + def attach( + self, + peer, + strength: abstract.ReadWriteFloatScalar, + hardware: abstract.ReadWriteFloatScalar, + ) -> Self: + """ + Create a new reference to attach this magnet to a simulator + or a control systemand. + """ + obj = super().attach(peer, strength, hardware) + obj.__angle = RWCorrectorAngle(obj) + return obj diff --git a/pyaml/validation/configuration_models.py b/pyaml/validation/configuration_models.py index dc9bc104..b77f732f 100644 --- a/pyaml/validation/configuration_models.py +++ b/pyaml/validation/configuration_models.py @@ -2,9 +2,10 @@ import importlib import logging -from typing import ClassVar +from typing import Any, ClassVar -from pydantic import AliasChoices, BaseModel, ConfigDict, Field +from pydantic import AliasChoices, BaseModel, ConfigDict, Field, GetCoreSchemaHandler +from pydantic_core import CoreSchema, core_schema logger = logging.getLogger(__name__) @@ -30,11 +31,16 @@ def model_dump_json(self, **kwargs): class ConfigurationSchema(PyAMLBaseModel): """ - Base model for validating externally supplied configuration data. + Base model for configuration schemas. + + Each configuration schema defines the expected input for constructing a + specific object. The required ``class`` field specifies the fully + qualified class path of the object to construct. - Each configuration schema defines the expected input for configuring a - specific object and includes a ``class`` field containing its fully - qualified class path. + Notes + ----- + Virtual subclasses may be registered to allow compatible schema types to + be accepted during validation. """ model_config = ConfigDict(validate_by_name=True, validate_by_alias=True, arbitrary_types_allowed=False, extra="forbid") @@ -44,6 +50,130 @@ class ConfigurationSchema(PyAMLBaseModel): alias="class", ) + # Add implementation which mimics required behaviour of virtual subclasses + # Real virtual subclasses are not compatible with Pydantic + + _virtual_subclasses: ClassVar[set[type["ConfigurationSchema"]]] + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + cls._virtual_subclasses = set() + + @classmethod + def virtual_subclasses(cls) -> set[type["ConfigurationSchema"]]: + """ + Return the registered virtual subclasses. + + Returns + ------- + set[type[ConfigurationSchema]] + A copy of the set containing all virtual subclasses registered + for this schema class. + """ + return set(cls._virtual_subclasses) + + @classmethod + def register_virtual_subclass(cls, subclass: type["ConfigurationSchema"]) -> None: + """ + Register a virtual subclass. + + Parameters + ---------- + subclass : type[ConfigurationSchema] + Subclass to register as a virtual subclass of this schema. + + Raises + ------ + TypeError + If ``subclass`` is not a ``ConfigurationSchema`` subclass. + """ + + if not isinstance(subclass, type) or not issubclass(subclass, ConfigurationSchema): + raise TypeError( + f"Cannot register {subclass!r} as a virtual subclass of {cls.__name__}: " + "it is not a ConfigurationSchema subclass." + ) + + cls._virtual_subclasses.add(subclass) + + @classmethod + def is_virtual_subclass_of(cls, superclass: type["ConfigurationSchema"]) -> bool: + """ + Check whether this class is a subclass of another schema. + + Parameters + ---------- + superclass : type[ConfigurationSchema] + Schema class to compare against. + + Returns + ------- + bool + ``True`` if ``cls`` is ``superclass`` or a registered virtual + subclass of it, otherwise ``False``. + + Raises + ------ + TypeError + If ``superclass`` is not a ``ConfigurationSchema`` subclass. + """ + + if not isinstance(superclass, type) or not issubclass(superclass, ConfigurationSchema): + raise TypeError(f"{superclass!r} is not a ConfigurationSchema.") + + checked: set[type["ConfigurationSchema"]] = set() + + def subclass_exists(current_cls: type["ConfigurationSchema"]) -> bool: + if current_cls in checked: + return False + checked.add(current_cls) + + if current_cls is cls: + return True + + for child in current_cls._virtual_subclasses: + if subclass_exists(child): + return True + + return False + + return subclass_exists(superclass) + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + """ + Customize Pydantic core schema generation. + + Parameters + ---------- + source : Any + Source type passed by Pydantic during schema generation. + handler : GetCoreSchemaHandler + Pydantic schema handler used to generate the default schema. + + Returns + ------- + CoreSchema + Wrapped core schema that accepts registered virtual subclasses. + """ + + # Get the normal schema + schema = handler(source) + + def validate(value, nxt): + # Already the expected schema type + if isinstance(value, cls): + return value + + # A registered virtual subclass instance should also be accepted + if isinstance(value, ConfigurationSchema) and type(value).is_virtual_subclass_of(cls): + return value + + # Otherwise let Pydantic validate normally + return nxt(value) + + return core_schema.no_info_wrap_validator_function(validate, schema) + class ModuleConfigurationSchema(PyAMLBaseModel): """ @@ -60,7 +190,7 @@ class ModuleConfigurationSchema(PyAMLBaseModel): module_path: str = Field( description="Fully qualified module path.", - alias=AliasChoices(*MODULE_PATH_ALIASES), + validation_alias=AliasChoices(*MODULE_PATH_ALIASES), ) def to_configuration(self) -> ConfigurationSchema: diff --git a/pyaml/validation/generator.py b/pyaml/validation/generator.py index d890b046..0f9d8aca 100644 --- a/pyaml/validation/generator.py +++ b/pyaml/validation/generator.py @@ -175,6 +175,7 @@ def model_schema(self, schema: core_schema.ModelSchema) -> dict[str, Any]: """ base_schema = super().model_schema(schema) + model_cls = schema.get("cls") logging.debug(f"Base schema is extracted from {model_cls}.") @@ -198,7 +199,12 @@ def model_schema(self, schema: core_schema.ModelSchema) -> dict[str, Any]: { schema_cls for _, schema_cls in self._registry.items() - if isinstance(schema_cls, type) and issubclass(schema_cls, model_cls) and schema_cls is not model_cls + if ( + isinstance(schema_cls, type) + and issubclass(schema_cls, ConfigurationSchema) + and schema_cls is not model_cls + and (issubclass(schema_cls, model_cls) or schema_cls.is_virtual_subclass_of(model_cls)) + ) }, key=lambda cls: f"{cls.__module__}.{cls.__name__}", ) diff --git a/pyaml/validation/schema_builder.py b/pyaml/validation/schema_builder.py index 652a30fb..a2ed5759 100644 --- a/pyaml/validation/schema_builder.py +++ b/pyaml/validation/schema_builder.py @@ -44,6 +44,49 @@ ) +def generate_class_path(source: type) -> str: + """ + Generate the fully qualified class path for a type. + + Parameters + ---------- + source : type + Class or type object to convert into a fully qualified path. + + Returns + ------- + str + Fully qualified class path in the form ``module.name``. + """ + return f"{source.__module__}.{source.__name__}" + + +def _extract_source_bases(source: type) -> list[type]: + """ + Extract relevant base classes from a type's method resolution order. + + Parameters + ---------- + source : type + Class whose non-framework base classes should be collected. + + Returns + ------- + list[type] + List of base classes from the class hierarchy, excluding + ``object``, :class:`pydantic.BaseModel`, and + :class:`ConfigurationSchema`. + """ + bases: list[type] = [] + + for base in source.__mro__[1:]: + if base in (object, BaseModel, ConfigurationSchema): + continue + bases.append(base) + + return bases + + def generate_configuration_schema(source: type) -> type[ConfigurationSchema]: """ Generate a configuration schema for a class or Pydantic model. @@ -61,7 +104,7 @@ def generate_configuration_schema(source: type) -> type[ConfigurationSchema]: raise TypeError("Source must be a class.") registry = SchemaRegistry() - class_path = f"{source.__module__}.{source.__name__}" + class_path = generate_class_path(source) existing = registry.get(class_path) if existing is not None: @@ -82,6 +125,25 @@ def generate_configuration_schema(source: type) -> type[ConfigurationSchema]: logger.debug("Register schema for %s.", class_path) registry.register(class_path, schema) + # Keep information about the inheritance relations + bases = _extract_source_bases(source) + base_schemas: list[type[ConfigurationSchema]] = [] + + for base in bases: + base_class_path = generate_class_path(base) + base_schema = registry.get(base_class_path) + + if base_schema is None: + base_schema = generate_configuration_schema(base) + logger.debug("Generate and register schema for base class %s.", base) + + base_schemas.append(base_schema) + + # Register virtual subclass links + for base_schema in base_schemas: + logger.debug("Register virtual subclass %s under %s.", schema, base_schema) + base_schema.register_virtual_subclass(schema) + return schema diff --git a/tests/configuration/test_curve.py b/tests/magnet/test_curve.py similarity index 82% rename from tests/configuration/test_curve.py rename to tests/magnet/test_curve.py index 68848a7d..cf8597d1 100644 --- a/tests/configuration/test_curve.py +++ b/tests/magnet/test_curve.py @@ -1,13 +1,12 @@ import numpy as np from pyaml.configuration import ROOT -from pyaml.magnet.csvcurve import ConfigModel, CSVCurve +from pyaml.magnet.csvcurve import CSVCurve from pyaml.magnet.curve import Curve def curve_test(file: str, current: float, strength: float): - curveConfig = ConfigModel(file=file) - curve = CSVCurve(curveConfig) + curve = CSVCurve(file=file) curveData = curve.get_curve() icurveData = Curve.inverse(curveData) x1 = np.interp(current, curveData[:, 0], curveData[:, 1]) diff --git a/tests/validation/test_generator.py b/tests/validation/test_generator.py index e32521f9..97dfc34d 100644 --- a/tests/validation/test_generator.py +++ b/tests/validation/test_generator.py @@ -42,6 +42,10 @@ class ChildSchemaB(ParentSchema): b: str = "x" +class VirtualChildSchema(ConfigurationSchema): + b: str = "x" + + class ContainerSchema(ConfigurationSchema): model: ChildSchemaA | None = Field( default=None, @@ -125,10 +129,37 @@ def test_generate_replaces_parent_schema_with_registered_subclasses( schema = SchemaGenerator.generate("pkg.module.Parent") - child_refs = {item["$ref"] for item in schema.get("anyOf", [])} + anyof = schema.get("anyOf", []) + refs = {item["$ref"] for item in anyof if "$ref" in item} + if refs: + assert "#/$defs/ChildSchemaA" in refs + assert "#/$defs/ChildSchemaB" in refs + else: + class_paths = {item["properties"]["class"]["const"] for item in anyof if "properties" in item} + assert "pkg.module.ChildA" in class_paths + assert "pkg.module.ChildB" in class_paths - assert "#/$defs/ChildSchemaA" in child_refs - assert "#/$defs/ChildSchemaB" in child_refs + +def test_generate_includes_real_and_virtual_subclasses( + registry: SchemaRegistry, +): + registry.register("pkg.module.Parent", ParentSchema) + registry.register("pkg.module.ChildA", ChildSchemaA) + registry.register("pkg.module.VirtualChild", VirtualChildSchema) + + ParentSchema.register_virtual_subclass(VirtualChildSchema) + + schema = SchemaGenerator.generate("pkg.module.Parent") + + anyof = schema.get("anyOf", []) + refs = {item["$ref"] for item in anyof if "$ref" in item} + if refs: + assert "#/$defs/ChildSchemaA" in refs + assert "#/$defs/VirtualChildSchema" in refs + else: + class_paths = {item["properties"]["class"]["const"] for item in anyof if "properties" in item} + assert "pkg.module.ChildA" in class_paths + assert "pkg.module.VirtualChild" in class_paths def test_model_schema_preserves_metadata_from_parent_schema( @@ -138,12 +169,10 @@ def test_model_schema_preserves_metadata_from_parent_schema( registry.register("pkg.module.ChildA", ChildSchemaA) registry.register("pkg.module.ChildB", ChildSchemaB) - generator = RegistryJsonSchema() - base_schema = ParentSchema.__pydantic_core_schema__ - - schema = generator.model_schema(base_schema) + schema = ParentSchema.model_json_schema(schema_generator=RegistryJsonSchema) assert schema["title"] == "ParentSchema" + assert schema["description"] == "Parent schema used to test inheritance." # ========================================================== diff --git a/tests/validation/test_models.py b/tests/validation/test_models.py index 45ddbc06..0c362716 100644 --- a/tests/validation/test_models.py +++ b/tests/validation/test_models.py @@ -97,6 +97,89 @@ def test_configuration_schema_dump_uses_alias_when_requested(): assert dumped == {"class": "pkg.module.Class"} +def test_configuration_schema_registers_virtual_subclass(): + class Curve(ConfigurationSchema): + pass + + class Curve1(ConfigurationSchema): + pass + + Curve.register_virtual_subclass(Curve1) + + assert Curve1 in Curve.virtual_subclasses() + assert Curve1.is_virtual_subclass_of(Curve) + + +def test_configuration_schema_is_virtual_subclass_of_is_transitive(): + class Curve(ConfigurationSchema): + pass + + class Curve1(ConfigurationSchema): + pass + + class Curve2(ConfigurationSchema): + pass + + Curve.register_virtual_subclass(Curve1) + Curve1.register_virtual_subclass(Curve2) + + assert Curve1.is_virtual_subclass_of(Curve) + assert Curve2.is_virtual_subclass_of(Curve) + assert Curve2.is_virtual_subclass_of(Curve1) + assert Curve.is_virtual_subclass_of(Curve) + assert not Curve.is_virtual_subclass_of(Curve1) + assert not Curve1.is_virtual_subclass_of(Curve2) + + +def test_configuration_schema_virtual_subclasses_are_per_class(): + class Curve(ConfigurationSchema): + pass + + class Curve1(ConfigurationSchema): + pass + + class Curve2(ConfigurationSchema): + pass + + Curve.register_virtual_subclass(Curve1) + Curve1.register_virtual_subclass(Curve2) + + assert Curve.virtual_subclasses() == {Curve1} + assert Curve1.virtual_subclasses() == {Curve2} + assert Curve2.virtual_subclasses() == set() + + +def test_configuration_schema_virtual_subclasses_returns_copy(): + class Curve(ConfigurationSchema): + pass + + class Curve1(ConfigurationSchema): + pass + + Curve.register_virtual_subclass(Curve1) + + subclasses = Curve.virtual_subclasses() + subclasses.clear() + + assert Curve.virtual_subclasses() == {Curve1} + + +def test_configuration_schema_register_virtual_subclass_rejects_non_schema(): + class Curve(ConfigurationSchema): + pass + + with pytest.raises(TypeError, match="is not a ConfigurationSchema subclass"): + Curve.register_virtual_subclass(int) # type: ignore[arg-type] + + +def test_configuration_schema_is_virtual_subclass_of_rejects_non_schema(): + class Curve(ConfigurationSchema): + pass + + with pytest.raises(TypeError, match="is not a ConfigurationSchema"): + Curve.is_virtual_subclass_of(int) # type: ignore[arg-type] + + # ========================================================== # ValidationModel # ========================================================== diff --git a/tests/validation/test_schema_builder.py b/tests/validation/test_schema_builder.py index 8c84dd13..3f6abefe 100644 --- a/tests/validation/test_schema_builder.py +++ b/tests/validation/test_schema_builder.py @@ -103,6 +103,69 @@ class ValidationModel(BaseModel): ) +def test_generate_configuration_schema_registers_parent_chain(): + class Curve: + pass + + class Curve1(Curve): + pass + + class Curve2(Curve1): + pass + + curve2_schema = generate_configuration_schema(Curve2) + + registry = SchemaRegistry() + curve_path = f"{Curve.__module__}.{Curve.__name__}" + curve1_path = f"{Curve1.__module__}.{Curve1.__name__}" + curve2_path = f"{Curve2.__module__}.{Curve2.__name__}" + + curve_schema = registry.get(curve_path) + curve1_schema = registry.get(curve1_path) + curve2_schema_from_registry = registry.get(curve2_path) + + assert curve_schema is not None + assert curve1_schema is not None + assert curve2_schema_from_registry is curve2_schema + + assert curve1_schema.is_virtual_subclass_of(curve_schema) + assert curve2_schema.is_virtual_subclass_of(curve1_schema) + assert curve2_schema.is_virtual_subclass_of(curve_schema) + assert not curve_schema.is_virtual_subclass_of(curve1_schema) + + +def test_generate_configuration_schema_registers_multiple_inheritance(): + class Curve: + pass + + class Monitor: + pass + + class Curve1(Curve): + pass + + class Curve2(Curve1, Monitor): + pass + + curve2_schema = generate_configuration_schema(Curve2) + + registry = SchemaRegistry() + curve_schema = registry.get(f"{Curve.__module__}.{Curve.__name__}") + monitor_schema = registry.get(f"{Monitor.__module__}.{Monitor.__name__}") + curve1_schema = registry.get(f"{Curve1.__module__}.{Curve1.__name__}") + curve2_schema_from_registry = registry.get(f"{Curve2.__module__}.{Curve2.__name__}") + + assert curve_schema is not None + assert monitor_schema is not None + assert curve1_schema is not None + assert curve2_schema_from_registry is curve2_schema + + assert curve1_schema.is_virtual_subclass_of(curve_schema) + assert curve2_schema.is_virtual_subclass_of(curve1_schema) + assert curve2_schema.is_virtual_subclass_of(curve_schema) + assert curve2_schema.is_virtual_subclass_of(monitor_schema) + + def test_fields_from_constructor_signature(): class Example: def __init__(self, x: int, y: str = "a", *args, **kwargs):