diff --git a/pyaml/accelerator.py b/pyaml/accelerator.py index 2af928af..dc273e2c 100644 --- a/pyaml/accelerator.py +++ b/pyaml/accelerator.py @@ -4,23 +4,20 @@ from pydantic import BaseModel, ConfigDict, Field -from .arrays.array import ArrayConfig -from .common.element import Element +from .arrays.array import Array, ArraySchema +from .common.element import Element, ElementSchema from .common.element_holder import ElementHolder from .common.exception import PyAMLConfigException from .configuration import ConfigurationManager, UnsupportedConfigurationRootError from .configuration.factory import Factory -from .control.controlsystem import ControlSystem -from .lattice.simulator import Simulator +from .control.controlsystem import ControlSystem, ControlSystemSchema +from .lattice.simulator import Simulator, SimulatorSchema from .yellow_pages import YellowPages -# Define the main class name for this module -PYAMLCLASS = "Accelerator" - -class ConfigModel(BaseModel): +class AcceleratorSchema(BaseModel): """ - Configuration model for Accelerator + Schema for Accelerator Parameters ---------- @@ -50,70 +47,86 @@ class ConfigModel(BaseModel): Element list """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") facility: str machine: str energy: float alphac: float | None = None harmonic_number: int | None = None - controls: list[ControlSystem] = None - simulators: list[Simulator] = None + controls: list[ControlSystemSchema] = None + simulators: list[SimulatorSchema] = None data_folder: str description: str | None = None - arrays: list[ArrayConfig] = Field(default=None, repr=False) - devices: list[Element] = Field(repr=False) + arrays: list[ArraySchema] = Field(default=None, repr=False) + devices: list[ElementSchema] = Field(default=None, repr=False) class Accelerator(object): """PyAML top level class""" - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + facility: str, + machine: str, + energy: float, + data_folder: str, + alphac: float | None = None, + harmonic_number: int | None = None, + controls: list[ControlSystem] = None, + simulators: list[Simulator] = None, + description: str | None = None, + arrays: list[Array] | None = None, + devices: list[Element] | None = None, + ): + self.facility = facility + self.machine = machine + __design = None __live = None + self._controls: dict[str, ElementHolder] = {} self._simulators: dict[str, ElementHolder] = {} - if cfg.controls is not None: - for c in cfg.controls: + if controls is not None: + for c in controls: if c.name() == "live": self.__live = c else: # Add as dynamic attribute setattr(self, c.name(), c) - c.fill_device(cfg.devices) + c.fill_device(devices) c._peer = self self._controls[c.name()] = c - if cfg.simulators is not None: - for s in cfg.simulators: + if simulators is not None: + for s in simulators: if s.name() == "design": self.__design = s else: # Add as dynamic attribute setattr(self, s.name(), s) - s.fill_device(cfg.devices) + s.fill_device(devices) s._peer = self self._simulators[s.name()] = s - if cfg.arrays is not None: - for a in cfg.arrays: - if cfg.simulators is not None: - for s in cfg.simulators: + if arrays is not None: + for a in arrays: + if simulators is not None: + for s in simulators: a.fill_array(s) - if cfg.controls is not None: - for c in cfg.controls: + if controls is not None: + for c in controls: a.fill_array(c) - if cfg.energy is not None: - self.set_energy(cfg.energy) + if energy is not None: + self.set_energy(energy) - if cfg.alphac is not None: - self.set_mcf(cfg.alphac) + if alphac is not None: + self.set_mcf(alphac) - if cfg.harmonic_number is not None: - self.set_harmonic_number(cfg.harmonic_number) + if harmonic_number is not None: + self.set_harmonic_number(harmonic_number) self._yellow_pages = YellowPages(self) @@ -121,12 +134,12 @@ def __init__(self, cfg: ConfigModel): def _set_properties(self, method: str, value): # Sets global property - if self._cfg.simulators is not None: - for s in self._cfg.simulators: + if self._simulators is not None: + for s in self._simulators: m = getattr(s, method) m(value) - if self._cfg.controls is not None: - for c in self._cfg.controls: + if self._controls is not None: + for c in self._controls: m = getattr(c, method) m(value) @@ -180,24 +193,24 @@ def add_device(self, config: dict, ignore_external=False): "Invalid device type, Element or sub classes of Element expected " + f"but got {dev.__class__.__name__}" ) - self._cfg.devices.append(dev) - if self._cfg.controls is not None: - for c in self._cfg.controls: + self._devices.append(dev) + if self._controls is not None: + for c in self._controls: c.fill_device([dev]) - if self._cfg.simulators is not None: - for s in self._cfg.simulators: + if self._simulators is not None: + for s in self._simulators: s.fill_device([dev]) def post_init(self): """ Method triggered after all initialisations are done """ - if self._cfg.simulators is not None: - for s in self._cfg.simulators: + if self._simulators is not None: + for s in self._simulators: s.post_init() - if self._cfg.controls is not None: - for c in self._cfg.controls: + if self._controls is not None: + for c in self._controls: c.post_init() def get_description(self) -> str: diff --git a/pyaml/arrays/array.py b/pyaml/arrays/array.py index e3b59b73..4d0e2381 100644 --- a/pyaml/arrays/array.py +++ b/pyaml/arrays/array.py @@ -1,5 +1,5 @@ """ -Array configuration +Base classes for arrays """ from pydantic import BaseModel, ConfigDict @@ -9,33 +9,36 @@ from ..common.element_holder import ElementHolder -class ArrayConfigModel(BaseModel): +class ArraySchema(BaseModel): """ - Base class for configuration of array of :py:class:`~pyaml.arrays.element.Element`, - :py:class:`~pyaml.arrays.bpm.BPM`, :py:class:`~pyaml.arrays.magnet.Magnet` or - :py:class:`~pyaml.arrays.cfm_magnet.CombinedFunctionMagnet`. + Schema for configuration of array. Parameters ---------- name : str Name of the array elements : list[str] - List of pyaml element names + List of element names """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") name: str elements: list[str] -class ArrayConfig(object): +class Array: """ Base class that implements configuration for access to arrays (families) """ - def __init__(self, cfg: ArrayConfigModel): - self._cfg = cfg + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -56,7 +59,7 @@ def fill_array(self, holder: ElementHolder): """ raise PyAMLException("Array.fill_array() is not subclassed") - def __repr__(self): - # ArrayConfigModel is a super class - # ConfigModel is expected from sub classes - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + # def __repr__(self): + # # ArrayConfigModel is a super class + # # ConfigModel is expected from sub classes + # return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/arrays/bpm.py b/pyaml/arrays/bpm.py index 0883a378..d78e6e5f 100644 --- a/pyaml/arrays/bpm.py +++ b/pyaml/arrays/bpm.py @@ -1,15 +1,8 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from .array import Array -# Define the main class name for this module -PYAMLCLASS = "BPM" - -class ConfigModel(ArrayConfigModel): - """Configuration model for :py:class:`.BPMArray`.""" - - -class BPM(ArrayConfig): +class BPM(Array): """ :py:class:`.BPMArray` configuration. @@ -27,21 +20,15 @@ class BPM(ArrayConfig): - BPM_C04-02 - BPM_C04-03 - A :py:class:`.BPMArray` configuration can also be created by code using the - following example: - - .. code-block:: python - - >>> from pyaml.arrays.bpm import BPM,ConfigModel as BPMArrayConfigModel - >>> bpm_cfg = BPM(BPMArrayConfigModel( - name="BPM", - elements=["BPM_C04-01","BPM_C04-02","BPM_C04-03"] - )) - """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -66,4 +53,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate the :py:class:`.BPMArray` with. """ - holder.fill_bpm_array(self._cfg.name, self._cfg.elements) + holder.fill_bpm_array(self._name, self._elements) diff --git a/pyaml/arrays/cfm_magnet.py b/pyaml/arrays/cfm_magnet.py index f69ce658..15e4e10e 100644 --- a/pyaml/arrays/cfm_magnet.py +++ b/pyaml/arrays/cfm_magnet.py @@ -1,35 +1,20 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from .array import Array -# Define the main class name for this module -PYAMLCLASS = "CombinedFunctionMagnet" - -class ConfigModel(ArrayConfigModel): - """Configuration model for Combined Function Magnet array.""" - - ... - - -class CombinedFunctionMagnet(ArrayConfig): +class CombinedFunctionMagnet(Array): """ - Combined function magnet array confirguration - - Example - ------- - - A magnet array configuration can also be created by code using - the following example:: + Combined function magnet array confirguration. - from pyaml.arrays.cfm_magnet import CombinedFunctionMagnet - from pyaml.arrays.cfm_magnet import ConfigModel as CFMagnetConfigModel - magArray = CombinedFunctionMagnet( - CFMagnetConfigModel(name="myCFM", elements=["mag1","mag2"]) - ) """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -40,4 +25,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with combined function magnet array """ - holder.fill_cfm_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_cfm_magnet_array(self._name, self._elements) diff --git a/pyaml/arrays/element.py b/pyaml/arrays/element.py index 6a40d027..039f4572 100644 --- a/pyaml/arrays/element.py +++ b/pyaml/arrays/element.py @@ -1,36 +1,20 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from .array import Array -# Define the main class name for this module -PYAMLCLASS = "Element" - -class ConfigModel(ArrayConfigModel): - """Configuration model for :py:class:`.ElementArray`.""" - - -class Element(ArrayConfig): +class Element(Array): """ :py:class:`.ElementArray` configuration. - Example - ------- - - An element array configuration can also be created by code using - the following example: - - .. code-block:: python - - from pyaml.arrays.element import Element,ConfigModel as ElementArrayConfigModel - elt_cfg = Element( - ElementArrayConfigModel(name="MyArray", elements=["BPM_C04-01","SH1A-C04-H"]) - ) - - """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -54,4 +38,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with element array """ - holder.fill_element_array(self._cfg.name, self._cfg.elements) + holder.fill_element_array(self._name, self._elements) diff --git a/pyaml/arrays/magnet.py b/pyaml/arrays/magnet.py index d65282cd..083a8a94 100644 --- a/pyaml/arrays/magnet.py +++ b/pyaml/arrays/magnet.py @@ -1,34 +1,19 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from .array import Array -# Define the main class name for this module -PYAMLCLASS = "Magnet" - -class ConfigModel(ArrayConfigModel): - """Configuration model for Magnet array.""" - - ... - - -class Magnet(ArrayConfig): +class Magnet(Array): """ - Magnet array confirguration - - Example - ------- - - A magnet array configuration can also be created by code using - the following example:: - - from pyaml.arrays.magnet import Magnet,ConfigModel as MagnetArrayConfigModel - magArray = Magnet( - MagnetArrayConfigModel(name="MyMags", elements=["mag1","mag2"]) - ) + Magnet array confirguration. """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -39,4 +24,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with magnet array """ - holder.fill_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_magnet_array(self._name, self._elements) diff --git a/pyaml/arrays/serialized_magnet.py b/pyaml/arrays/serialized_magnet.py index fd412403..86e60ecb 100644 --- a/pyaml/arrays/serialized_magnet.py +++ b/pyaml/arrays/serialized_magnet.py @@ -1,35 +1,20 @@ from ..common.element_holder import ElementHolder -from .array import ArrayConfig, ArrayConfigModel +from .array import Array -# Define the main class name for this module -PYAMLCLASS = "SerializedMagnets" - -class ConfigModel(ArrayConfigModel): - """Configuration model for Serialized Magnets array.""" - - ... - - -class SerializedMagnets(ArrayConfig): +class SerializedMagnets(Array): """ Serialized magnets array configuration - Example - ------- - - A magnet array configuration can also be created by code using - the following example:: - - from pyaml.arrays.serialized_magnet import SerializedMagnets - from pyaml.arrays.serialized_magnet import ConfigModel as SerializedMagnetConfigModel - magArray = SerializedMagnets( - SerializedMagnetConfigModel(name="mySerializedMagnets", elements=["mag1","mag2"]) - ) """ - def __init__(self, cfg: ArrayConfigModel): - super().__init__(cfg) + def __init__( + self, + name: str, + elements: list[str], + ): + self._name = name + self._elements = elements def fill_array(self, holder: ElementHolder): """ @@ -40,4 +25,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with serialized magnet array """ - holder.fill_serialized_magnet_array(self._cfg.name, self._cfg.elements) + holder.fill_serialized_magnet_array(self._name, self._elements) diff --git a/pyaml/bpm/bpm.py b/pyaml/bpm/bpm.py index 1a14f4b9..4e734020 100644 --- a/pyaml/bpm/bpm.py +++ b/pyaml/bpm/bpm.py @@ -1,5 +1,5 @@ -from ..bpm.bpm_model import BPMModel -from ..common.element import Element, ElementConfigModel +from ..bpm.bpm_model import BPMModel, BPMModelSchema +from ..common.element import Element, ElementSchema from ..common.exception import PyAMLException from ..lattice.abstract_impl import RBpmArray, RWBpmOffsetArray, RWBpmTiltScalar @@ -8,10 +8,8 @@ except ImportError: from typing_extensions import Self # Python 3.10 and earlier -PYAMLCLASS = "BPM" - -class ConfigModel(ElementConfigModel): +class BPMSchema(ElementSchema): """ Configuration model for BPM element. @@ -21,7 +19,7 @@ class ConfigModel(ElementConfigModel): Object in charge of BPM modeling """ - model: BPMModel | None = None + model: BPMModelSchema | None = None class BPM(Element): @@ -29,7 +27,11 @@ class BPM(Element): Class providing access to one BPM of a physical or simulated lattice """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + model: BPMModel | None = None, + ): """ Construct a BPM @@ -41,10 +43,8 @@ def __init__(self, cfg: ConfigModel): BPM model in charge of computing beam position """ - super().__init__(cfg.name) - - self.__model = cfg.model if hasattr(cfg, "model") else None - self._cfg = cfg + super().__init__(name) + self.__model = model self.__positions = None self.__offset = None self.__tilt = None @@ -146,7 +146,7 @@ def attach( """ # Attach positions, offset and tilt attributes and returns a new # reference - obj = self.__class__(self._cfg) + obj = self.__class__(self.name, self.__model) obj.__model = self.__model obj.__positions = positions obj.__offset = offset diff --git a/pyaml/bpm/bpm_model.py b/pyaml/bpm/bpm_model.py index e00c2aa3..efefa1e5 100644 --- a/pyaml/bpm/bpm_model.py +++ b/pyaml/bpm/bpm_model.py @@ -2,10 +2,14 @@ import numpy as np from numpy.typing import NDArray +from pydantic import BaseModel from ..control.deviceaccess import DeviceAccess +class BPMModelSchema(BaseModel): ... + + class BPMModel(metaclass=ABCMeta): """ Abstract class providing interface to accessing BPM positions, offsets, diff --git a/pyaml/bpm/bpm_simple_model.py b/pyaml/bpm/bpm_simple_model.py index 582a6025..79e38c20 100644 --- a/pyaml/bpm/bpm_simple_model.py +++ b/pyaml/bpm/bpm_simple_model.py @@ -5,21 +5,19 @@ from pyaml.bpm.bpm_model import BPMModel from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema +from .bpm_model import BPMModelSchema -# Define the main class name for this module -PYAMLCLASS = "BPMSimpleModel" - -class ConfigModel(BaseModel): +class BPMSimpleModelSchema(BPMModelSchema): """ Configuration model for BPM simple model Parameters ---------- - x_pos : DeviceAccess, optional + x_pos : DeviceAccessSchema, optional Horizontal position device - y_pos : DeviceAccess, optional + y_pos : DeviceAccessSchema, optional Vertical position device x_pos_index : int, optional Index in the array when specified, otherwise scalar @@ -29,10 +27,10 @@ class ConfigModel(BaseModel): value is expected """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") - x_pos: DeviceAccess | None - y_pos: DeviceAccess | None + x_pos: DeviceAccessSchema | None + y_pos: DeviceAccessSchema | None x_pos_index: int | None = None y_pos_index: int | None = None @@ -43,10 +41,17 @@ class BPMSimpleModel(BPMModel): offset values. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg - self.__x_pos = cfg.x_pos - self.__y_pos = cfg.y_pos + def __init__( + self, + x_pos: DeviceAccess | None, + y_pos: DeviceAccess | None, + x_pos_index: int | None = None, + y_pos_index: int | None = None, + ): + self.__x_pos = x_pos + self.__y_pos = y_pos + self.__x_pos_index = x_pos_index + self.__y_pos_index = y_pos_index def get_pos_devices(self) -> list[DeviceAccess | None]: """ @@ -92,7 +97,7 @@ def x_pos_index(self) -> int | None: int Index in the array, None for a scalar value """ - return self._cfg.x_pos_index + return self.__x_pos_index def y_pos_index(self) -> int | None: """ @@ -105,7 +110,8 @@ def y_pos_index(self) -> int | None: int Index in the array, None for a scalar value """ - return self._cfg.y_pos_index + return self.__.y_pos_index + - def __repr__(self): - return __pyaml_repr__(self) +# def __repr__(self): +# return __pyaml_repr__(self) diff --git a/pyaml/bpm/bpm_tiltoffset_model.py b/pyaml/bpm/bpm_tiltoffset_model.py index 6bccbd8c..e79c3593 100644 --- a/pyaml/bpm/bpm_tiltoffset_model.py +++ b/pyaml/bpm/bpm_tiltoffset_model.py @@ -3,18 +3,15 @@ from pydantic import BaseModel, ConfigDict from pyaml.bpm.bpm_model import BPMModel -from pyaml.bpm.bpm_simple_model import BPMSimpleModel +from pyaml.bpm.bpm_simple_model import BPMSimpleModel, BPMSimpleModelSchema from ..common.element import __pyaml_repr__ from ..control.deviceaccess import DeviceAccess -# Define the main class name for this module -PYAMLCLASS = "BPMTiltOffsetModel" - # TODO: Implepement indexed offset and tilt -class ConfigModel(BaseModel): +class BPMTiltOffsetModelSchema(BPMSimpleModelSchema): """ Configuration model for BPM with tilt and offset @@ -32,15 +29,11 @@ class ConfigModel(BaseModel): BPM tilt device """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") - x_pos: DeviceAccess | None - y_pos: DeviceAccess | None - x_pos_index: int | None = None - y_pos_index: int | None = None - x_offset: DeviceAccess | None - y_offset: DeviceAccess | None - tilt: DeviceAccess | None + x_offset: DeviceAccess | None = None + y_offset: DeviceAccess | None = None + tilt: DeviceAccess | None = None class BPMTiltOffsetModel(BPMSimpleModel): @@ -49,13 +42,21 @@ class BPMTiltOffsetModel(BPMSimpleModel): offset values. """ - def __init__(self, cfg: ConfigModel): - super().__init__(cfg) - self.__x_pos = cfg.x_pos - self.__y_pos = cfg.y_pos - self.__x_offset = cfg.x_offset - self.__y_offset = cfg.y_offset - self.__tilt = cfg.tilt + def __init__( + self, + x_pos: DeviceAccess | None, + y_pos: DeviceAccess | None, + x_pos_index: int | None = None, + y_pos_index: int | None = None, + x_offset: DeviceAccess | None = None, + y_offset: DeviceAccess | None = None, + tilt: DeviceAccess | None = None, + ): + super().__init__(x_pos=x_pos, y_pos=y_pos, x_pos_index=x_pos_index, y_pos_index=y_pos_index) + + self.__x_offset = x_offset + self.__y_offset = y_offset + self.__tilt = tilt def get_pos_devices(self) -> list[DeviceAccess | None]: """ @@ -90,5 +91,6 @@ def get_offset_devices(self) -> list[DeviceAccess | None]: """ return [self.__x_offset, self.__y_offset] - def __repr__(self): - return __pyaml_repr__(self) + +# def __repr__(self): +# return __pyaml_repr__(self) diff --git a/pyaml/common/element.py b/pyaml/common/element.py index f323f25a..862813c7 100644 --- a/pyaml/common/element.py +++ b/pyaml/common/element.py @@ -1,47 +1,75 @@ +import warnings from typing import TYPE_CHECKING -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, field_validator from .exception import PyAMLException +from .utils import __pyaml_repr__ if TYPE_CHECKING: from ..common.element_holder import ElementHolder -def __pyaml_repr__(obj): +class ElementSchema(BaseModel): """ - Returns a string representation of a pyaml object - """ - if hasattr(obj, "_cfg"): - if isinstance(obj, Element): - return repr(obj._cfg).replace( - "ConfigModel(", - obj.__class__.__name__ + "(peer='" + obj.get_peer_name() + "', ", - ) - else: - # no peer - return repr(obj._cfg).replace("ConfigModel", obj.__class__.__name__) - else: - # Object is not yet fully constructed - if isinstance(obj, Element): - return f"{obj.__class__.__name__}: {obj.get_name()}" - else: - return f"{obj.__class__.__name__}" + Base schema for element configuration. + Parameters + ---------- + name : str + The name of the element. + description : str, optional + Description of the element. + lattice_names : str or None, optional + The name(s) of the associated element(s) in the lattice. By default, + the element name is used. lattice_name accept the following + syntax: + - list(name,[name]) : Element names + - [name]@idx[,idx] : Element indices in the subset formed by name. + - [name]#start_idx..end_idx : Element range in the subset formed by name. + In the above syntax, if the name is not specficied, the whole set + of lattice element is used for indexing. + """ -class ElementConfigModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: str = Field(description="Name of the element") + description: str | None = Field(default=None, description="Description of the element.") + lattice_names: str | None = Field( + default=None, description="The name(s) of the associated element(s) in the lattice." + ) + + # Validate the syntax for the lattice_names field + # This validation is currently very basic and can be improved + @field_validator("lattice_names") + @classmethod + def validate_lattice_names(cls, v): + if v is None: + return v + + # Example: very simplified checks + if v.startswith("list(") and v.endswith(")"): + return v + if "@" in v: + return v + if "#" in v and ".." in v: + return v + raise ValueError("Invalid lattice_names syntax") + + +class Element: """ - Base class for element configuration. + Class providing access to one element of a physical or simulated lattice Parameters ---------- name : str - The name of the PyAML element. + The name of the element. description : str, optional Description of the element. lattice_names : str or None, optional The name(s) of the associated element(s) in the lattice. By default, - the PyAML element name is used. lattice_name accept the following + the element name is used. lattice_name accept the following syntax: - list(name,[name]) : Element names - [name]@idx[,idx] : Element indices in the subset formed by name. @@ -50,46 +78,37 @@ class ElementConfigModel(BaseModel): of lattice element is used for indexing. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - name: str - description: str | None = None - lattice_names: str | None = None + def __init__(self, name: str, description: str | None = None, lattice_names: str | None = None): + self._name: str = name + self.description = description + # If no lattice names are given put it to the name of the element + if lattice_names: + self._lattice_names = lattice_names + else: + self._lattice_names = self._name -class Element(object): - """ - Class providing access to one element of a physical or simulated lattice + self._peer: ElementHolder | None = None # Peer: ControlSystem, Simulator - Attributes: - name: str - The unique name identifying the element in the configuration file - """ + @property + def name(self) -> str: + return self._name - def __init__(self, name: str): - self._name: str = name - self._peer: "ElementHolder" = None # Peer: ControlSystem, Simulator + # TODO: implement name setter -> this requires checking so the name is unique def get_name(self) -> str: - """ - Returns the name of the element - """ - return self._name + warnings.warn( + "get_name() is deprecated; use .name instead", + DeprecationWarning, + stacklevel=2, + ) + return self.name - def get_lattice_names(self) -> str: - """ - Returns the name of associated lattice element(s) - """ - if not hasattr(self, "_cfg"): - return self._name - else: - return self._cfg.lattice_names + @property + def lattice_names(self) -> str: + return self._lattice_names - def get_description(self) -> str: - """ - Returns the description of the element - """ - return self._cfg.description + # TODO: implement lattice_names setter -> this requires validation of the format def set_energy(self, E: float): """ @@ -115,7 +134,7 @@ def check_peer(self): to a simulator or to a control system """ if self._peer is None: - raise PyAMLException(f"{str(self)} is not attachedto a control system or the a simulator") + raise PyAMLException(f"{str(self)} is not attached to a control system or a simulator.") @property def peer(self) -> "ElementHolder": diff --git a/pyaml/common/utils.py b/pyaml/common/utils.py new file mode 100644 index 00000000..8ee5f88a --- /dev/null +++ b/pyaml/common/utils.py @@ -0,0 +1,26 @@ +"""Module for utility functions.""" + + +def __pyaml_repr__(obj): + """ + Returns a string representation of a pyaml object + """ + + attrs = {} + + # Instance attributes + for k, v in obj.__dict__.items(): + # Exclude private attributes + if not k.startswith("_"): + attrs[k] = v + + # Properties + for name, attr in vars(type(obj)).items(): + if isinstance(attr, property): + try: + attrs[name] = getattr(obj, name) + except Exception as e: + attrs[name] = f"" + + parts = ", ".join(f"{k}={v!r}" for k, v in attrs.items()) + return f"{obj.__class__.__name__}({parts})" diff --git a/pyaml/configuration/csvcurve.py b/pyaml/configuration/csvcurve.py index fadb0a45..5562bfe3 100644 --- a/pyaml/configuration/csvcurve.py +++ b/pyaml/configuration/csvcurve.py @@ -5,13 +5,10 @@ from ..common.exception import PyAMLException from ..configuration.fileloader import get_path -from .curve import Curve +from .curve import Curve, CurveSchema -# Define the main class name for this module -PYAMLCLASS = "CSVCurve" - -class ConfigModel(BaseModel): +class CSVCurveSchema(CurveSchema): """ Configuration model for CSV curve @@ -21,7 +18,7 @@ class ConfigModel(BaseModel): CSV file that contains the curve (n rows, 2 columns) """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") file: str @@ -31,23 +28,20 @@ class CSVCurve(Curve): Class for load CSV (x,y) curve """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, file: str): + self._file = file # Load CSV curve - path = get_path(cfg.file) + path = get_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):" - f"wrong shape (2,2) expected but got {str(_s)}" + f"CSVCurve(file='{self._file}',dtype=float):wrong shape (2,2) expected but got {str(_s)}" ) def get_curve(self) -> np.array: @@ -61,5 +55,6 @@ def get_curve(self) -> np.array: """ return self._curve - def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + +# def __repr__(self): +# return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/csvmatrix.py b/pyaml/configuration/csvmatrix.py index daed97f8..55788aab 100644 --- a/pyaml/configuration/csvmatrix.py +++ b/pyaml/configuration/csvmatrix.py @@ -5,13 +5,10 @@ from ..common.exception import PyAMLException from ..configuration.fileloader import get_path -from .matrix import Matrix +from .matrix import Matrix, MatrixSchema -# Define the main class name for this module -PYAMLCLASS = "CSVMatrix" - -class ConfigModel(BaseModel): +class CSVMatrixSchema(MatrixSchema): """ Configuration model for CSV matrix @@ -21,7 +18,7 @@ class ConfigModel(BaseModel): CSV file that contains the matrix """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") file: str @@ -31,16 +28,15 @@ class CSVMatrix(Matrix): Class for loading CSV matrix """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, file: str): + self._file = file + # Load CSV matrix - path = get_path(cfg.file) + path = get_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 def get_matrix(self) -> np.array: """ @@ -53,5 +49,6 @@ def get_matrix(self) -> np.array: """ return self._mat - def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + +# def __repr__(self): +# return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/curve.py b/pyaml/configuration/curve.py index df73ec01..1c4254b8 100644 --- a/pyaml/configuration/curve.py +++ b/pyaml/configuration/curve.py @@ -1,6 +1,11 @@ from abc import ABCMeta, abstractmethod import numpy as np +from pydantic import BaseModel, ConfigDict + + +class CurveSchema(BaseModel): + model_config = ConfigDict(extra="forbid") class Curve(metaclass=ABCMeta): diff --git a/pyaml/configuration/inline_curve.py b/pyaml/configuration/inline_curve.py index e92071ab..ba4d5536 100644 --- a/pyaml/configuration/inline_curve.py +++ b/pyaml/configuration/inline_curve.py @@ -5,13 +5,10 @@ from ..common.exception import PyAMLException from ..configuration import get_root_folder -from .curve import Curve +from .curve import Curve, CurveSchema -# Define the main class name for this module -PYAMLCLASS = "InlineCurve" - -class ConfigModel(BaseModel): +class InlineCurveSchema(CurveSchema): """ Configuration model for inline curve @@ -21,7 +18,7 @@ class ConfigModel(BaseModel): Curve data (n rows, 2 columns) """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") mat: list[list[float]] @@ -31,15 +28,16 @@ class InlineCurve(Curve): Class for load CSV (x,y) curve """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, mat: list[list[float]]): + self._mat = mat + # 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)}" + f"InlineCurve(mat='{self._mat}',dtype=float): wrong shape (2,2) expected but got {str(_s)}" ) def get_curve(self) -> np.array: @@ -53,5 +51,6 @@ def get_curve(self) -> np.array: """ return self._curve - def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + +# def __repr__(self): +# return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/inline_matrix.py b/pyaml/configuration/inline_matrix.py index 25d03c78..493fab4a 100644 --- a/pyaml/configuration/inline_matrix.py +++ b/pyaml/configuration/inline_matrix.py @@ -1,13 +1,10 @@ import numpy as np from pydantic import BaseModel, ConfigDict -from .matrix import Matrix +from .matrix import Matrix, MatrixSchema -# Define the main class name for this module -PYAMLCLASS = "InlineMatrix" - -class ConfigModel(BaseModel): +class InlineMatrixSchema(MatrixSchema): """ Configuration model for inline matrix @@ -17,7 +14,7 @@ class ConfigModel(BaseModel): The matrix """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") mat: list[list[float]] @@ -27,10 +24,11 @@ class InlineMatrix(Matrix): Class for loading CSV matrix """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__(self, mat: list[list[float]]): + self._mat = mat + # Load the matrix - self._mat = np.array(self._cfg.mat) + self._matrix = np.array(self._mat) def get_matrix(self) -> np.array: """ @@ -41,7 +39,8 @@ def get_matrix(self) -> np.array: np.array Matrix data as a numpy array """ - return self._mat + return self._matrix + - def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) +# def __repr__(self): +# return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/configuration/matrix.py b/pyaml/configuration/matrix.py index dc5b9ac7..6b2509b8 100644 --- a/pyaml/configuration/matrix.py +++ b/pyaml/configuration/matrix.py @@ -1,6 +1,11 @@ from abc import ABCMeta, abstractmethod from numpy import array +from pydantic import BaseModel, ConfigDict + + +class MatrixSchema(BaseModel): + model_config = ConfigDict(extra="forbid") class Matrix(metaclass=ABCMeta): diff --git a/pyaml/configuration/schema_registry.py b/pyaml/configuration/schema_registry.py new file mode 100644 index 00000000..437d1d06 --- /dev/null +++ b/pyaml/configuration/schema_registry.py @@ -0,0 +1,123 @@ +# config/schema_registry.py +from enum import Enum +from typing import Any, Type + +# from typing import Literal +from pydantic import BaseModel + +from ..accelerator import AcceleratorSchema +from ..arrays.array import ArraySchema +from ..common.element import ElementSchema +from ..control.controlsystem import ControlSystemSchema +from ..lattice.simulator import SimulatorSchema + + +class BaseSchema(Enum): + ACCELERATOR = AcceleratorSchema + CONTROLSYSTEM = ControlSystemSchema + SIMULATOR = SimulatorSchema + ARRAY = ArraySchema + DEVICE = ElementSchema + + +class SchemaRegistry: + """ + Singleton registry for schemas. + The registry maps the class type to which + schema that should be used for validation. + + The registry is divided into namespaces + based on the base schemas. + """ + + _instance: "SchemaRegistry | None" = None + + # Schemas follow the convention: + # namespace -> class_str -> validation model class + _schemas: dict[BaseSchema, dict[type[Any], Type[BaseModel]]] + + def __new__(cls) -> "SchemaRegistry": + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._schemas = {baseschema: {} for baseschema in BaseSchema} + return cls._instance + + def register(self, class_: type[Any], schema: Type[BaseModel], baseschema: BaseSchema | None = None) -> None: + if baseschema is None: + baseschema = get_baseschema(schema) + + if baseschema is None: + raise TypeError( + f"Could not infer a BaseSchema from {schema.__name__}. A baseschema must be explicitly given." + ) + + elif not isinstance(baseschema, BaseSchema): + raise TypeError(f"{baseschema!r} is not a valid BaseSchema.") + + # Check if already registered + if self.contains(class_, baseschema): + raise ValueError(f"{class_.__name__} is already registered under {baseschema.name}.") + + self._schemas[baseschema][class_] = schema + + def contents( + self, + ) -> dict[BaseSchema, dict[type[Any], Type[BaseModel]]]: + return {baseschema: dict(schemas) for baseschema, schemas in self._schemas.items()} + + def get( + self, + class_: type[Any], + baseschema: BaseSchema, + ) -> Type[BaseModel]: + _validate_baseschema(baseschema) + + try: + return self._schemas[baseschema][class_] + + except KeyError: + raise KeyError(f"{class_.__name__} is not registered under {baseschema.name}.") from None + + def __getitem__( + self, + item: tuple[type[Any], BaseSchema], + ) -> Type[BaseModel]: + class_, baseschema = item + + return self.get(class_, baseschema) + + def contains( + self, + class_: type[Any], + baseschema: BaseSchema, + ) -> bool: + _validate_baseschema(baseschema) + + return class_ in self._schemas[baseschema] + + def __contains__( + self, + item: tuple[type[Any], BaseSchema], + ) -> bool: + class_, baseschema = item + + return self.contains(class_, baseschema) + + +def get_baseschema(schema: Type[BaseModel]) -> BaseSchema | None: + matches = [baseschema for baseschema in BaseSchema if issubclass(schema, baseschema.value)] + + if len(matches) > 1: + raise TypeError(f"{schema.__name__} matches multiple base schemas: {[m.name for m in matches]}") + + if not matches: + return None + + return matches[0] + + +def _validate_baseschema( + baseschema: BaseSchema, +) -> None: + if not isinstance(baseschema, BaseSchema): + raise TypeError(f"{baseschema!r} is not a valid BaseSchema.") diff --git a/pyaml/control/controlsystem.py b/pyaml/control/controlsystem.py index 33f57f84..1fd9c6b9 100644 --- a/pyaml/control/controlsystem.py +++ b/pyaml/control/controlsystem.py @@ -1,6 +1,8 @@ from abc import ABCMeta, abstractmethod from typing import Tuple +from pydantic import BaseModel + from ..bpm.bpm import BPM from ..bpm.bpm_model import BPMModel from ..common.abstract import RWMapper @@ -38,6 +40,13 @@ from .deviceaccess import DeviceAccess +class ControlSystemSchema(BaseModel): + name: str + debug_level: str = None + scalar_aggregator: str | None = None + vector_aggregator: str | None = None + + class ControlSystem(ElementHolder, metaclass=ABCMeta): """ Abstract class providing access to a control system float variable diff --git a/pyaml/control/deviceaccess.py b/pyaml/control/deviceaccess.py index ec37cd4f..51c06398 100644 --- a/pyaml/control/deviceaccess.py +++ b/pyaml/control/deviceaccess.py @@ -1,8 +1,14 @@ from abc import ABCMeta, abstractmethod +from pydantic import BaseModel, ConfigDict + # TODO: correctly type value +class DeviceAccessSchema(BaseModel): + model_config = ConfigDict(extra="forbid") + + class DeviceAccess(metaclass=ABCMeta): """ Abstract class providing access to a control system variable diff --git a/pyaml/diagnostics/chromaticity_monitor.py b/pyaml/diagnostics/chromaticity_monitor.py index 550749ed..c02b7705 100644 --- a/pyaml/diagnostics/chromaticity_monitor.py +++ b/pyaml/diagnostics/chromaticity_monitor.py @@ -2,7 +2,7 @@ from ..common.constants import Action from ..common.element import ElementConfigModel from ..common.exception import PyAMLException -from ..tuning_tools.measurement_tool import MeasurementTool, MeasurementToolConfigModel +from ..tuning_tools.measurement_tool import MeasurementTool, MeasurementToolSchema try: from typing import Self # Python 3.11+ @@ -17,10 +17,8 @@ logger = logging.getLogger(__name__) -PYAMLCLASS = "ChomaticityMonitor" - -class ConfigModel(MeasurementToolConfigModel): +class ChomaticityMonitorSchema(MeasurementToolSchema): """ Configuration model for Chromaticity Monitor. @@ -47,7 +45,7 @@ class ConfigModel(MeasurementToolConfigModel): Dispersion fitting, by default False """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") betatron_tune_name: str rf_plant_name: str @@ -88,20 +86,35 @@ class ChomaticityMonitor(MeasurementTool): 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, + ): """ Construct a ChomaticityMonitor. - Parameters - ---------- - cfg : ConfigModel - Configuration for the ChromaticityMonitor, including betatron - tune monitor, RF plant, and defaults parameters. """ - 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._chromaticity = RChromaDispArray(self, "chromaticity", "1") - self._dipsersion = RChromaDispArray(self, "dispersion", "m") + self._dispersion = RChromaDispArray(self, "dispersion", "m") self._alphac = None @property @@ -129,7 +142,7 @@ def dispersion(self) -> ReadFloatArray: ReadFloatArray Array of dispersion values [[dx, dy],[d'x, d'y],...] """ - return self._dipsersion + return self._dispersion def measure( self, @@ -202,19 +215,19 @@ 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}") + logger.warning(f"e_delta={self._e_delta} is greater than max_e_delta={self._max_e_delta}") if alphac is None: raise PyAMLException("Moment compaction factor is not defined") @@ -224,14 +237,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() diff --git a/pyaml/diagnostics/tune_monitor.py b/pyaml/diagnostics/tune_monitor.py index 1a534c1c..35b0a376 100644 --- a/pyaml/diagnostics/tune_monitor.py +++ b/pyaml/diagnostics/tune_monitor.py @@ -1,6 +1,6 @@ from ..common.abstract import ReadFloatArray -from ..common.element import Element, ElementConfigModel -from ..control.deviceaccess import DeviceAccess +from ..common.element import Element, ElementSchema +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .atune_monitor import ABetatronTuneMonitor try: @@ -10,25 +10,23 @@ import numpy as np from pydantic import ConfigDict -PYAMLCLASS = "BetatronTuneMonitor" - -class ConfigModel(ElementConfigModel): +class BetatronTuneMonitorSchema(ElementSchema): """ - Configuration model for BetatronTuneMonitor + Configuration schema for BetatronTuneMonitor Parameters ---------- - tune_h : DeviceAccess, optional + tune_h : DeviceAccessSchema, optional Horizontal betatron tune device - tune_v : DeviceAccess, optional + tune_v : DeviceAccessSchema, optional Vertical betatron tune device """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") - tune_h: DeviceAccess | None - tune_v: DeviceAccess | None + tune_h: DeviceAccessSchema | None = None + tune_v: DeviceAccessSchema | None = None rf_plant_name: str | None = None @@ -39,24 +37,30 @@ class BetatronTuneMonitor(Element, ABetatronTuneMonitor): The monitor provides horizontal and vertical betatron tune measurements. """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + description: str | None = None, + lattice_names: str | None = None, + tune_h: DeviceAccess | None = None, + tune_v: DeviceAccess | None = None, + rf_plant_name: str | None = None, + ): """ Construct a BetatronTuneMonitor - - Parameters - ---------- - cfg : ConfigModel - Configuration for the BetatronTuneMonitor, including - device access for horizontal and vertical tunes. """ - super().__init__(cfg.name) - self._cfg = cfg - self.__tune = None - self._h = None + super().__init__(name, description, lattice_names) + + self._tune_h = tune_h + self._tune_v = tune_v + self.rf_plant_name = rf_plant_name + + self._tune = None + self._harmonic_number = None def set_harmonic(self, h: int): - self._h = float(h) + self._harmonic_number = float(h) @property def tune(self) -> ReadFloatArray: @@ -87,13 +91,13 @@ def __init__(self, parent: BetatronTuneMonitor): self.parent = parent def get(self) -> np.array: - h = self.parent._h - rf_name = self.parent._cfg.rf_plant_name - if h is not None and rf_name is not None: + harmonic_number = self.parent._harmonic_number + rf_name = self.parent.rf_plant_name + if harmonic_number is not None and rf_name is not None: tune = self.parent.tune.get() rf = self.parent.peer.get_rf_plant(rf_name) freq = rf.frequency.get() - return tune * freq / h + return tune * freq / harmonic_number def unit(self) -> str: return "Hz" @@ -117,7 +121,9 @@ def attach(self, peer, betatron_tune: ReadFloatArray) -> Self: Self A new attached instance of TuneMonitor """ - obj = self.__class__(self._cfg) - obj.__tune = betatron_tune + obj = self.__class__( + self.name, self._description, self._lattice_names, self._tune_h, self._tune_v, self.rf_plant_name + ) + obj._tune = betatron_tune obj._peer = peer return obj diff --git a/pyaml/lattice/attribute_linker.py b/pyaml/lattice/attribute_linker.py index 820b8c97..9ab454e0 100644 --- a/pyaml/lattice/attribute_linker.py +++ b/pyaml/lattice/attribute_linker.py @@ -4,14 +4,12 @@ from pyaml.common.element import Element from pyaml.lattice.lattice_elements_linker import ( LatticeElementsLinker, - LinkerConfigModel, LinkerIdentifier, + LinkerSchema, ) -PYAMLCLASS = "PyAtAttributeElementsLinker" - -class ConfigModel(LinkerConfigModel): +class PyAtAttributeIdentifierSchema(LinkerSchema): """Base configuration model for linker definitions. This class defines the configuration structure used to instantiate @@ -26,7 +24,7 @@ class ConfigModel(LinkerConfigModel): unexpected extra keys. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") attribute_name: str @@ -62,8 +60,9 @@ class PyAtAttributeElementsLinker(LatticeElementsLinker): The configuration model for the linking strategy. """ - def __init__(self, config_model: ConfigModel): - super().__init__(config_model) + def __init__(self, attribute_name: str): + super().__init__(self) + self._attribute_name = attribute_name def get_element_identifier(self, element: Element) -> LinkerIdentifier: """ @@ -79,12 +78,8 @@ def get_element_identifier(self, element: Element) -> LinkerIdentifier: LinkerIdentifier The identifier for linking the element """ - return PyAtAttributeIdentifier( - self.linker_config_model.attribute_name, element.get_name() - ) + return PyAtAttributeIdentifier(self.self._attribute_name, element.get_name()) - def _test_at_element( - self, identifier: PyAtAttributeIdentifier, element: at.Element - ) -> bool: + def _test_at_element(self, identifier: PyAtAttributeIdentifier, element: at.Element) -> bool: attr_value = getattr(element, identifier.attribute_name, None) return attr_value == identifier.identifier diff --git a/pyaml/lattice/lattice_elements_linker.py b/pyaml/lattice/lattice_elements_linker.py index 68b40da1..9778e161 100644 --- a/pyaml/lattice/lattice_elements_linker.py +++ b/pyaml/lattice/lattice_elements_linker.py @@ -9,7 +9,7 @@ from pyaml.common.element import Element -class LinkerConfigModel(BaseModel): +class LatticeElementsLinkerSchema(BaseModel): """Base configuration model for linker definitions. This class defines the configuration structure used to instantiate @@ -24,7 +24,7 @@ class LinkerConfigModel(BaseModel): unexpected extra keys. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") class LinkerIdentifier(metaclass=ABCMeta): @@ -48,19 +48,13 @@ class LatticeElementsLinker(metaclass=ABCMeta): to PyAT elements based on a given linking strategy (e.g., by family name, by index, or by a custom attribute). - Parameters - ---------- - linker_config_model : LinkerConfigModel - The configuration model for the linking strategy. - Attributes ---------- lattice : at.Lattice Reference to the PyAT lattice handled by this linker. """ - def __init__(self, linker_config_model: LinkerConfigModel): - self.linker_config_model = linker_config_model + def __init__(self): self.lattice: Lattice = None def set_lattice(self, lattice: Lattice): @@ -75,9 +69,7 @@ def set_lattice(self, lattice: Lattice): self.lattice = lattice @abstractmethod - def _test_at_element( - self, identifier: LinkerIdentifier, element: at.Element - ) -> bool: + def _test_at_element(self, identifier: LinkerIdentifier, element: at.Element) -> bool: pass @abstractmethod @@ -103,9 +95,7 @@ def _iter_matches(self, identifier: LinkerIdentifier) -> Iterable[at.Element]: if self._test_at_element(identifier, elem): yield elem - def get_at_elements( - self, element_id: LinkerIdentifier | list[LinkerIdentifier] - ) -> list[at.Element]: + def get_at_elements(self, element_id: LinkerIdentifier | list[LinkerIdentifier]) -> list[at.Element]: """Return a list of PyAT elements matching the given identifiers. This method should resolve one or multiple PyAML identifiers @@ -139,8 +129,7 @@ def get_at_elements( if not results: raise PyAMLException( - f"No PyAT elements found for identifier(s): " - f"{', '.join(i.__repr__() for i in identifiers)}" + f"No PyAT elements found for identifier(s): {', '.join(i.__repr__() for i in identifiers)}" ) return results @@ -164,6 +153,4 @@ def get_at_element(self, element_id: LinkerIdentifier) -> at.Element: """ for elem in self._iter_matches(element_id): return elem - raise PyAMLException( - f"No PyAT element found for FamName: {element_id.__repr__()}" - ) + raise PyAMLException(f"No PyAT element found for FamName: {element_id.__repr__()}") diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index db2f4d13..5f117434 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -43,13 +43,10 @@ from .attribute_linker import ( PyAtAttributeElementsLinker, ) -from .lattice_elements_linker import LatticeElementsLinker +from .lattice_elements_linker import LatticeElementsLinker, LatticeElementsLinkerSchema -# Define the main class name for this module -PYAMLCLASS = "Simulator" - -class ConfigModel(BaseModel): +class SimulatorSchema(BaseModel): """ Configuration model for Simulator @@ -61,18 +58,18 @@ class ConfigModel(BaseModel): AT lattice file mat_key : str, optional AT lattice ring name - linker : LatticeElementsLinker, optional + linker : LatticeElementsLinkerSchema, optional The linker configuration model description : str , optional Simulator description """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") name: str lattice: str mat_key: str = None - linker: LatticeElementsLinker = None + linker: LatticeElementsLinkerSchema = None description: str | None = None @@ -81,17 +78,29 @@ class Simulator(ElementHolder): Class that implements access to AT simulator """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + lattice: str, + mat_key: str = None, + linker: LatticeElementsLinker = None, + description: str | None = None, + ): super().__init__() - self._cfg = cfg - path: Path = get_root_folder() / cfg.lattice + self._name = name + self._lattice = lattice + self._mat_key = mat_key + self._linker = linker + self._description = description + + path: Path = get_root_folder() / self._lattice - if self._cfg.mat_key is None: + if self._mat_key is None: self.ring = at.load_lattice(path) else: - self.ring = at.load_lattice(path, mat_key=f"{self._cfg.mat_key}") + self.ring = at.load_lattice(path, mat_key=f"{self._mat_key}") - self._linker = cfg.linker + self._linker = self._linker if self._linker: self._linker.set_lattice(self.ring) else: @@ -103,7 +112,7 @@ def __init__(self, cfg: ConfigModel): self._elements_indexing[e.FamName] = [e] def name(self) -> str: - return self._cfg.name + return self._name def get_lattice(self) -> at.Lattice: return self.ring @@ -112,7 +121,7 @@ def get_description(self) -> str: """ Returns the description of the accelerator """ - return self._cfg.description + return self._description def create_magnet_strength_aggregator(self, magnets: list[Magnet]) -> ScalarAggregator: # No magnet aggregator for simulator @@ -202,13 +211,13 @@ def fill_device(self, elements: list[Element]): self.add_bpm(e) elif isinstance(e, RFPlant): - if e._cfg.transmitters: + if e._transmitters: cavs: list[at.Element] = [] harmonics: list[float] = [] attachedTrans: list[RFTransmitter] = [] - for t in e._cfg.transmitters: + for t in e._transmitters: cavsPerTrans: list[at.Element] = [] - for c in t._cfg.cavities: + for c in t._cavities: # Expect unique name for cavities cav = self.get_at_elems(Element(c)) if len(cav) > 1: @@ -218,7 +227,7 @@ def fill_device(self, elements: list[Element]): if len(cav) == 0: raise PyAMLException(f"RF transmitter {t.get_name()}, No cavity found") cavsPerTrans.append(cav[0]) - harmonics.append(t._cfg.harmonic) + harmonics.append(t._harmonic) voltage = RWRFVoltageScalar(cavsPerTrans) phase = RWRFPhaseScalar(cavsPerTrans) nt = t.attach(self, voltage, phase) @@ -306,7 +315,7 @@ def get_at_elems(self, element: Element) -> list[at.Element]: identifier = self._linker.get_element_identifier(element) element_list = self._linker.get_at_elements(identifier) if not element_list: - raise PyAMLException(f"{identifier} not found in lattice:{self._cfg.lattice}") + raise PyAMLException(f"{identifier} not found in lattice:{self._lattice}") return element_list else: # By list @@ -316,7 +325,7 @@ def get_at_elems(self, element: Element) -> list[at.Element]: names = [] for name in nameList: if name not in self._elements_indexing: - raise PyAMLException(f"{name} not found in lattice:{self._cfg.lattice}") + raise PyAMLException(f"{name} not found in lattice:{self._lattice}") elts = self._elements_indexing[name] names.extend(elts) return names @@ -329,12 +338,13 @@ def get_at_elems(self, element: Element) -> list[at.Element]: return [self.ring[idx] for idx in indices] else: if name not in self._elements_indexing: - raise PyAMLException(f"{name} not found in lattice:{self._cfg.lattice}") + raise PyAMLException(f"{name} not found in lattice:{self._lattice}") elts = self._elements_indexing[name] if indices is None: return elts else: return [elts[idx] for idx in indices] - def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + +# def __repr__(self): +# return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) diff --git a/pyaml/magnet/cfm_magnet.py b/pyaml/magnet/cfm_magnet.py index 9249a6c8..14473a22 100644 --- a/pyaml/magnet/cfm_magnet.py +++ b/pyaml/magnet/cfm_magnet.py @@ -2,12 +2,12 @@ from ..common import abstract from ..common.abstract import RWMapper -from ..common.element import Element, ElementConfigModel, __pyaml_repr__ +from ..common.element import Element, ElementSchema, __pyaml_repr__ from ..common.exception import PyAMLException from ..configuration import Factory from .hcorrector import HCorrector -from .magnet import Magnet, MagnetConfigModel -from .model import MagnetModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel, MagnetModelSchema from .octupole import Octupole from .quadrupole import Quadrupole from .sextupole import Sextupole @@ -27,54 +27,44 @@ "A3": SkewOctu, } -# Define the main class name for this module -PYAMLCLASS = "CombinedFunctionMagnet" - -class ConfigModel(ElementConfigModel): +class CombinedFunctionMagnetSchema(ElementSchema): mapping: list[list[str]] """Name mapping for multipoles (i.e. [[B0,C01A-H],[A0,C01A-H],[B2,C01A-S]])""" - model: MagnetModel | None = None + model: MagnetModelSchema | 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 + def __init__(self, name: str, mapping: list[list[str]], model: MagnetModel | None = None, peer=None): + super().__init__(name) + self._mapping = mapping + self.model = model self.__virtuals: list[Magnet] = [] self.__strengths: abstract.ReadWriteFloatArray = None self.__hardwares: abstract.ReadWriteFloatArray = 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: mutipoles" - "field required for combined function magnet" - ) + if self.model is not None and not hasattr(self.model, "_multipoles"): + raise PyAMLException(f"{self._name} model: mutipolesfield 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}" - ) + 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: + raise PyAMLException(m[0] + " not implemented for combined function magnet") + 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 - vm = self.__create_virutal_manget(m[1], m[0]) + vm = self.__create_virtual_magnet(m[1], m[0]) self.__virtuals.append(vm) # Register the virtual element in the factory to have # a coherent factory and improve error reporting @@ -88,16 +78,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: + def __create_vitual_magnet(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, @@ -107,13 +97,13 @@ def attach( ) -> list[Magnet]: l = [] # Attached the CombinedFunctionMagnet itself - nCFM = CombinedFunctionMagnet(self._cfg, peer) + nCFM = CombinedFunctionMagnet(self._name, self._mapping, self._model, self._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)) @@ -127,9 +117,7 @@ def strengths(self) -> abstract.ReadWriteFloatScalar: """ self.check_peer() if self.__strengths is None: - raise PyAMLException( - f"{str(self)} has no model that supports physics units" - ) + raise PyAMLException(f"{str(self)} has no model that supports physics units") return self.__strengths @property @@ -140,9 +128,7 @@ def hardwares(self) -> abstract.ReadWriteFloatScalar: """ self.check_peer() if self.__hardwares is None: - raise PyAMLException( - f"{str(self)} has no model that supports hardware units" - ) + raise PyAMLException(f"{str(self)} has no model that supports hardware units") return self.__hardwares def set_energy(self, E: float): diff --git a/pyaml/magnet/hcorrector.py b/pyaml/magnet/hcorrector.py index 964c4e29..deea1f40 100644 --- a/pyaml/magnet/hcorrector.py +++ b/pyaml/magnet/hcorrector.py @@ -2,13 +2,11 @@ from ..common.constants import HORIZONTAL_KICK_SIGN from ..lattice.polynom_info import PolynomInfo from .corrector import RWCorrectorAngle -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "HCorrector" - -class ConfigModel(MagnetConfigModel): +class HCorrectorSchema(MagnetSchema): """Configuration model for Horizontal Corrector magnet.""" ... @@ -19,12 +17,10 @@ class HCorrector(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) self.__angle = RWCorrectorAngle(self) @property diff --git a/pyaml/magnet/identity_cfm_model.py b/pyaml/magnet/identity_cfm_model.py index 1b4a6bd3..d881c6c7 100644 --- a/pyaml/magnet/identity_cfm_model.py +++ b/pyaml/magnet/identity_cfm_model.py @@ -3,14 +3,14 @@ from .. import PyAMLException from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "IdentityCFMagnetModel" -class ConfigModel(BaseModel): +class IdentityCFMagnetModelSchema(BaseModel): """ Configuration model for identity combined function magnet model @@ -26,12 +26,12 @@ class ConfigModel(BaseModel): List of strength units (i.e. ['rad', 'm-1', 'm-2']) """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") multipoles: list[str] - powerconverters: list[DeviceAccess | None] | None = None - physics: list[DeviceAccess | None] | None = None - units: list[str] + powerconverters: list[DeviceAccessSchema | None] | None = None + physics: list[DeviceAccessSchema | None] | None = None + units: list[str] = "" class IdentityCFMagnetModel(MagnetModel): @@ -39,37 +39,43 @@ class IdentityCFMagnetModel(MagnetModel): Class that map values to underlying devices without conversion """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + multipoles: list[str], + powerconverters: list[DeviceAccess | None] | None = None, + physics: list[DeviceAccess | None] | None = None, + units: list[str] = "", + ): + self._multipoles = multipoles + self._powerconverters = powerconverters + self._physics = physics + self._units = units # Check config - self.__nbFunction: int = len(cfg.multipoles) + self.__nbFunction: int = len(self._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" + "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" + "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) if lgth != expected_len: raise PyAMLException( - f"{name} does not have the expected " - f"number of items ({expected_len} items expected but got {lgth})" + f"{name} does not have the expected number of items ({expected_len} items expected but got {lgth})" ) def compute_hardware_values(self, strengths: np.array) -> np.array: @@ -79,10 +85,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_devices(self) -> list[DeviceAccess | None]: return self.__devices @@ -91,10 +97,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 132fff1f..1965cff6 100644 --- a/pyaml/magnet/identity_model.py +++ b/pyaml/magnet/identity_model.py @@ -3,14 +3,11 @@ from .. import PyAMLException from ..common.element import __pyaml_repr__ -from ..control.deviceaccess import DeviceAccess +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "IdentityMagnetModel" - -class ConfigModel(BaseModel): +class IdentityMagnetModelSchema(BaseModel): """ Configuration model for identity magnet model @@ -24,11 +21,11 @@ class ConfigModel(BaseModel): Unit of the strength (i.e. 1/m or m-1) """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") - powerconverter: DeviceAccess | None = None - physics: DeviceAccess | None = None - unit: str + powerconverter: DeviceAccessSchema | None = None + physics: DeviceAccessSchema | None = None + unit: str = "" class IdentityMagnetModel(MagnetModel): @@ -36,23 +33,26 @@ class IdentityMagnetModel(MagnetModel): Class that map value to underlying device without conversion """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg - self.__unit = cfg.unit - if cfg.physics is None and cfg.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: + def __init__( + self, + powerconverter: DeviceAccess | None = None, + physics: DeviceAccess | None = None, + unit: str = "", + ): + self._powerconverter = powerconverter + self._physics = physics + self.__unit = unit + + if self._physics is None and self._powerconverter is None: + raise PyAMLException("Invalid IdentityMagnetModel configuration,physics or powerconverter device required") + 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" + "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 @@ -73,10 +73,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/linear_cfm_model.py b/pyaml/magnet/linear_cfm_model.py index 6ddc5421..e2ac62e0 100644 --- a/pyaml/magnet/linear_cfm_model.py +++ b/pyaml/magnet/linear_cfm_model.py @@ -1,18 +1,18 @@ import numpy as np from pydantic import BaseModel, ConfigDict -from ..common.element import __pyaml_repr__ +# from ..common.element import __pyaml_repr__ from ..common.exception import PyAMLException from ..configuration.curve import Curve from ..configuration.matrix import Matrix -from ..control.deviceaccess import DeviceAccess +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .model import MagnetModel # Define the main class name for this module PYAMLCLASS = "LinearCFMagnetModel" -class ConfigModel(BaseModel): +class LinearCFMagnetModelSchema(BaseModel): """ Configuration model for linear combined function magnet model @@ -44,7 +44,7 @@ class ConfigModel(BaseModel): List of strength units (i.e. ['rad', 'm-1', 'm-2']) """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") multipoles: list[str] curves: list[Curve] @@ -52,7 +52,7 @@ class ConfigModel(BaseModel): calibration_offsets: list[float] = None pseudo_factors: list[float] = None pseudo_offsets: list[float] = None - powerconverters: list[DeviceAccess | None] + powerconverters: list[DeviceAccessSchema | None] matrix: Matrix = None units: list[str] @@ -64,63 +64,77 @@ class LinearCFMagnetModel(MagnetModel): of power supply currents associated to a single function. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + multipoles: list[str], + curves: list[Curve], + powerconverters: list[DeviceAccess | None], + units: list[str], + calibration_factors: list[float] = None, + calibration_offsets: list[float] = None, + pseudo_factors: list[float] = None, + pseudo_offsets: list[float] = None, + matrix: Matrix = None, + ): + self._multipoles = (multipoles,) + self._curves = (curves,) + self._powerconverters = (powerconverters,) + self.units = (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._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.__check_len( - self.__calibration_factors, "calibration_factors", self.__nbFunction - ) - self.__check_len( - self.__calibration_offsets, "calibration_offsets", self.__nbFunction - ) + 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.curves, "curves", self.__nbFunction) + self.__check_len(self._units, "units", self.__nbFunction) + 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) if len(_s) != 2 or _s[0] != self.__nbFunction or _s[1] != self.__nbPS: raise PyAMLException( - "matrix wrong dimension " - f"({self.__nbFunction}x{self.__nbPS} expected but got {_s[0]}x{_s[1]})" + f"matrix wrong dimension ({self.__nbFunction}x{self.__nbPS} expected but got {_s[0]}x{_s[1]})" ) self.__curves = [] 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] @@ -133,18 +147,13 @@ def __check_len(self, obj, name, expected_len): lgth = len(obj) if lgth != expected_len: raise PyAMLException( - f"{name} does not have the expected " - f"number of items ({expected_len} items expected but got {lgth})" + f"{name} does not have the expected number of items ({expected_len} items expected but got {lgth})" ) def compute_hardware_values(self, strengths: np.array) -> np.array: _pI = np.zeros(self.__nbFunction) for idx, c in enumerate(self.__rcurves): - _pI[idx] = ( - self.__pf[idx] - * np.interp(strengths[idx] * self._brho, c[:, 0], c[:, 1]) - + self.__po[idx] - ) + _pI[idx] = self.__pf[idx] * np.interp(strengths[idx] * self._brho, c[:, 0], c[:, 1]) + self.__po[idx] _currents = np.matmul(self.__inv, _pI) return _currents @@ -152,30 +161,24 @@ def compute_strengths(self, currents: np.array) -> np.array: _strength = np.zeros(self.__nbFunction) _pI = np.matmul(self.__matrix, currents) for idx, c in enumerate(self.__curves): - _strength[idx] = ( - np.interp( - (_pI[idx] - self.__po[idx]) / self.__pf[idx], c[:, 0], c[:, 1] - ) - / self._brho - ) + _strength[idx] = np.interp((_pI[idx] - self.__po[idx]) / self.__pf[idx], c[:, 0], c[:, 1]) / self._brho return _strength def get_strength_units(self) -> list[str]: - return self._cfg.units + return self._units def get_hardware_units(self) -> list[str]: - return np.array([p.unit() for p in self._cfg.powerconverters]) + return np.array([p.unit() for p in self._powerconverters]) def get_devices(self) -> list[DeviceAccess]: - return self._cfg.powerconverters + return self._powerconverters def set_magnet_rigidity(self, brho: np.double): self._brho = brho def has_hardware(self) -> bool: - return (self.__nbPS == self.__nbFunction) and np.allclose( - self.__matrix, np.eye(self.__nbFunction) - ) + return (self.__nbPS == self.__nbFunction) and np.allclose(self.__matrix, np.eye(self.__nbFunction)) + - def __repr__(self): - return __pyaml_repr__(self) +# def __repr__(self): +# return __pyaml_repr__(self) diff --git a/pyaml/magnet/linear_model.py b/pyaml/magnet/linear_model.py index 44e182bc..e9172530 100644 --- a/pyaml/magnet/linear_model.py +++ b/pyaml/magnet/linear_model.py @@ -1,16 +1,12 @@ import numpy as np -from pydantic import BaseModel, ConfigDict -from ..common.element import __pyaml_repr__ -from ..configuration.curve import Curve -from ..control.deviceaccess import DeviceAccess -from .model import MagnetModel +# from ..common.element import __pyaml_repr__ +from ..configuration.curve import Curve, CurveSchema +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema +from .model import MagnetModel, MagnetModelSchema -# Define the main class name for this module -PYAMLCLASS = "LinearMagnetModel" - -class ConfigModel(BaseModel): +class LinearMagnetModelSchema(MagnetModelSchema): """ Linear magnet model. @@ -32,14 +28,12 @@ class ConfigModel(BaseModel): """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - curve: Curve | None = None - powerconverter: DeviceAccess | None + curve: CurveSchema | None = None + powerconverter: DeviceAccessSchema | None = None calibration_factor: float = 1.0 calibration_offset: float = 0.0 crosstalk: float = 1.0 - unit: str + unit: str = "" class LinearMagnetModel(MagnetModel): @@ -48,40 +42,41 @@ class LinearMagnetModel(MagnetModel): linear interpolation for a single function magnet """ - 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 - ) - self.__rcurve = Curve.inverse(self.__curve) + def __init__( + self, + curve: Curve | None = None, + powerconverter: DeviceAccess | None = None, + calibration_factor: float = 1.0, + calibration_offset: float = 0.0, + crosstalk: float = 1.0, + unit: str = "", + ): + 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.powerconverter.unit() + self.__g = calibration_factor * crosstalk + self.__o = calibration_offset + + self.__strength_unit = unit + if powerconverter: + self.__hardware_unit = powerconverter.unit() self.__brho = np.nan - self.__ps = cfg.powerconverter + self.__ps = powerconverter def compute_hardware_values(self, strengths: np.array) -> np.array: if self.__rcurve is not None: - _current = np.interp( - strengths[0] * self.__brho, self.__rcurve[:, 0], self.__rcurve[:, 1] - ) + _current = np.interp(strengths[0] * self.__brho, self.__rcurve[:, 0], self.__rcurve[:, 1]) else: _current = (strengths[0] * self.__brho) / self.__g + self.__o return np.array([_current]) def compute_strengths(self, currents: np.array) -> np.array: if self.__curve is not None: - _strength = ( - np.interp(currents[0], self.__curve[:, 0], self.__curve[:, 1]) - / self.__brho - ) + _strength = np.interp(currents[0], self.__curve[:, 0], self.__curve[:, 1]) / self.__brho else: _strength = ((currents[0] - self.__o) * self.__g) / self.__brho return np.array([_strength]) @@ -98,5 +93,6 @@ def get_devices(self) -> list[DeviceAccess]: def set_magnet_rigidity(self, brho: np.double): self.__brho = brho - def __repr__(self): - return __pyaml_repr__(self) + +# def __repr__(self): +# return __pyaml_repr__(self) diff --git a/pyaml/magnet/linear_serialized_model.py b/pyaml/magnet/linear_serialized_model.py index cfae4bb3..eaf7efe0 100644 --- a/pyaml/magnet/linear_serialized_model.py +++ b/pyaml/magnet/linear_serialized_model.py @@ -7,16 +7,13 @@ from ..configuration.inline_curve import ConfigModel as InlineCurveModel from ..configuration.inline_curve import InlineCurve from ..configuration.matrix import Matrix -from ..control.deviceaccess import DeviceAccess +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .linear_model import ConfigModel as LinearConfigModel from .linear_model import LinearMagnetModel from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "LinearSerializedMagnetModel" - -class ConfigModel(BaseModel): +class LinearSerializedMagnetModelSchema(BaseModel): """ Configuration model for linear serialized magnet model @@ -39,13 +36,13 @@ class ConfigModel(BaseModel): Strength unit: rad, m-1, m-2 """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(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: DeviceAccess + powerconverter: DeviceAccessSchema unit: str @@ -75,8 +72,7 @@ def _check_len(obj, name, expected_length): length = len(obj) if length != expected_length: raise PyAMLException( - f"{name} does not have the expected " - f"number of items ({expected_length} items expected but got {length})" + f"{name} does not have the expected number of items ({expected_length} items expected but got {length})" ) @@ -87,46 +83,54 @@ class LinearSerializedMagnetModel(MagnetModel): of power supply currents associated to a single function. """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + curves: Curve | list[Curve], + powerconverter: DeviceAccess, + unit: str, + calibration_factors: float | list[float] = None, + calibration_offsets: float | list[float] = None, + crosstalk: float | list[float] = 1.0, + ): + self._curves = curves + self._powerconverter = powerconverter + self._calibration_factors = calibration_factors + self._calibration_offsets = calibration_offsets + self._unit = unit self.__brho = np.nan # Check config self.__nbMagnets: int = _get_max_length( - cfg.curves, cfg.calibration_factors, cfg.calibration_offsets, cfg.crosstalk + self._curves, self._calibration_factors, self._calibration_offsets, self._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(InlineCurveModel(mat=self._curves.get_curve())) self.__curves.append(curve) _check_len(self.__calibration_factors, "calibration_factors", self.__nbMagnets) @@ -141,8 +145,8 @@ def __initialize(self): 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, + powerconverter=self._powerconverter, + unit=self._unit, ) self.__sub_models.append(LinearMagnetModel(sub_model)) @@ -157,26 +161,23 @@ def compute_hardware_values(self, strengths: np.array) -> np.array: return np.array( [ model.compute_hardware_values([strength]) - for strength, model in zip(strengths, self.__sub_models) + for strength, model in zip(strengths, self.__sub_models, strict=True) ] ) def compute_strengths(self, currents: np.array) -> np.array: return np.array( - [ - model.compute_strengths([current]) - for current, model in zip(currents, self.__sub_models) - ] + [model.compute_strengths([current]) for current, model in zip(currents, self.__sub_models, strict=True)] ) def get_strength_units(self) -> list[str]: - return self._cfg.units + return self._units def get_hardware_units(self) -> list[str]: - return [p.unit() for p in self._cfg.__sub_models] + return [p.unit() for p in self.__sub_models] def get_devices(self) -> list[DeviceAccess]: - return self._cfg.powerconverters + return self._powerconverters def set_magnet_rigidity(self, brho: np.double): self.__brho = brho diff --git a/pyaml/magnet/magnet.py b/pyaml/magnet/magnet.py index d98b5ca7..d748ae34 100644 --- a/pyaml/magnet/magnet.py +++ b/pyaml/magnet/magnet.py @@ -1,9 +1,10 @@ +from pydantic import Field from scipy.constants import speed_of_light from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel -from .model import MagnetModel +from ..common.element import Element, ElementSchema +from .model import MagnetModel, MagnetModelSchema try: from typing import Self # Python 3.11+ @@ -12,17 +13,19 @@ import numpy as np -class MagnetConfigModel(ElementConfigModel): +class MagnetSchema(ElementSchema): """ Configuration model for magnet elements. Attributes ---------- - model : MagnetModel or None, optional - Object in charge of converting magnet strengths to power supply values + model : MagnetModelSchema or None, optional + Schema for object in charge of converting magnet strengths to power supply values """ - model: MagnetModel | None = None + model: MagnetModelSchema | None = Field( + default=None, description="Schema for object in charge of converting magnet strengths to power supply values" + ) class Magnet(Element): @@ -30,18 +33,31 @@ 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): """ Construct a magnet Parameters ---------- - name : str + name: str Element name + description : str, optional + Description of the element. + lattice_names : str or None, optional + The name(s) of the associated element(s) in the lattice. By default, + the PyAML element name is used. lattice_name accept the following + syntax: + - list(name,[name]) : Element names + - [name]@idx[,idx] : Element indices in the subset formed by name. + - [name]#start_idx..end_idx : Element range in the subset formed by name. + In the above syntax, if the name is not specficied, the whole set + of lattice element is used for indexing. model : MagnetModel Magnet model in charge of computing coil(s) current """ - super().__init__(name) + super().__init__(name, description, lattice_names) self.__model = model self.__strength: abstract.ReadWriteFloatScalar = None self.__hardware: abstract.ReadWriteFloatScalar = None @@ -85,7 +101,7 @@ def attach( Create a new reference to attach this magnet to a simulator or a control systemand. """ - obj = self.__class__(self._cfg) + obj = self.__class__(self._name, self._description, self._lattice_names, self.__model) obj.__modelName = self.__modelName obj.__strength = strength obj.__hardware = hardware diff --git a/pyaml/magnet/model.py b/pyaml/magnet/model.py index cce2bfdc..ea86e74c 100644 --- a/pyaml/magnet/model.py +++ b/pyaml/magnet/model.py @@ -2,10 +2,15 @@ import numpy as np import numpy.typing as npt +from pydantic import BaseModel, ConfigDict from ..control.deviceaccess import DeviceAccess +class MagnetModelSchema(BaseModel): + model_config = ConfigDict(extra="forbid") + + class MagnetModel(metaclass=ABCMeta): """ Abstract class providing strength to coil current conversion @@ -13,9 +18,7 @@ class MagnetModel(metaclass=ABCMeta): """ @abstractmethod - def compute_hardware_values( - self, strengths: npt.NDArray[np.float64] - ) -> npt.NDArray[np.float64]: + def compute_hardware_values(self, strengths: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: """ Compute hardware value(s) from magnet strength(s) @@ -33,9 +36,7 @@ def compute_hardware_values( pass @abstractmethod - def compute_strengths( - self, hardware_values: npt.NDArray[np.float64] - ) -> npt.NDArray[np.float64]: + def compute_strengths(self, hardware_values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: """ Compute magnet strength(s) from hardware value(s) diff --git a/pyaml/magnet/octupole.py b/pyaml/magnet/octupole.py index d9710a09..1ed3101e 100644 --- a/pyaml/magnet/octupole.py +++ b/pyaml/magnet/octupole.py @@ -1,12 +1,10 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "Octupole" - -class ConfigModel(MagnetConfigModel): - """Configuration model for Octupole magnet.""" +class OctupoleSchema(MagnetSchema): + """Configuration model for Sextupole magnet.""" ... @@ -16,9 +14,7 @@ class Octupole(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/quadrupole.py b/pyaml/magnet/quadrupole.py index 3c720ab6..df47b50a 100644 --- a/pyaml/magnet/quadrupole.py +++ b/pyaml/magnet/quadrupole.py @@ -1,12 +1,10 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "Quadrupole" - -class ConfigModel(MagnetConfigModel): - """Configuration model for Quadrupole magnet.""" +class QuadrupoleSchema(MagnetSchema): + """Schema for Quadrupole magnet.""" ... @@ -16,9 +14,7 @@ class Quadrupole(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/serialized_magnet.py b/pyaml/magnet/serialized_magnet.py index 10613c04..29fb635d 100644 --- a/pyaml/magnet/serialized_magnet.py +++ b/pyaml/magnet/serialized_magnet.py @@ -3,18 +3,15 @@ from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel, __pyaml_repr__ +from ..common.element import Element, ElementSchema, __pyaml_repr__ from ..configuration import Factory from ..control.deviceaccess import DeviceAccess from .function_mapping import function_map -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "SerializedMagnets" - -class ConfigModel(ElementConfigModel): +class SerializedMagnetsSchema(ElementSchema): function: str """List of magnets""" elements: list[str] | str @@ -24,9 +21,15 @@ class ConfigModel(ElementConfigModel): class ReadWriteSerializedStrengths(abstract.ReadWriteFloatScalar): - def __init__(self, cfg: ConfigModel, elements: list[abstract.ReadWriteFloatScalar]): - self.elements = elements - self._cfg = cfg + def __init__( + self, + function: str, + elements: list[abstract.ReadWriteFloatScalar], + model: MagnetModel | None = None, + ): + self._function = function + self._elements = elements + self._model = model def get(self) -> float: return sum([elem.get() for elem in self.elements]) @@ -38,10 +41,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,11 +54,16 @@ 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, + function: str, + elements: list[abstract.ReadWriteFloatScalar], + model: MagnetModel | None = None, + ): + super().__init__(function, 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] @@ -67,33 +75,30 @@ class SerializedMagnets(Element): The set point is usually managed by only one power supply but it can be covered by several ones. If several power supplies - - Parameters - ---------- - cfg : ConfigModel - Configuration object TODO: to describe - - Raises - ------ - pyaml.PyAMLException - In case of wrong initialization """ - 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[abstract.ReadWriteFloatScalar], + model: MagnetModel | None = None, + peer=None, + ): + super().__init__(name) + 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 +112,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 +129,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, peer) + n_ser_mag.__strengths = ReadWriteSerializedStrengths(self._function, self.__elements, self.model, strengths) + n_ser_mag.__hardwares = ReadWriteSerializedHardwares(self._function, self.__elements, self.model, hardwares) l.append(n_ser_mag) # Construct a single magnet for each magnet. sub_magnets: list[Magnet] = [] @@ -174,4 +179,4 @@ def get_devices(self) -> list[DeviceAccess]: if isinstance(self.model.powerconverter, list): return self.model.powerconverter else: - return [self._cfg.powerconverter] + return [self._powerconverter] diff --git a/pyaml/magnet/sextupole.py b/pyaml/magnet/sextupole.py index e8218ca3..6bcff439 100644 --- a/pyaml/magnet/sextupole.py +++ b/pyaml/magnet/sextupole.py @@ -1,11 +1,9 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "Sextupole" - -class ConfigModel(MagnetConfigModel): +class SextupoleSchema(MagnetSchema): """Configuration model for Sextupole magnet.""" ... @@ -16,9 +14,7 @@ class Sextupole(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/skewoctu.py b/pyaml/magnet/skewoctu.py index 109edd78..f957362d 100644 --- a/pyaml/magnet/skewoctu.py +++ b/pyaml/magnet/skewoctu.py @@ -1,11 +1,9 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "SkewOctu" - -class ConfigModel(MagnetConfigModel): +class SkewOctuSchema(MagnetSchema): """Configuration model for SkewOctu magnet.""" ... @@ -16,9 +14,7 @@ class SkewOctu(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/skewquad.py b/pyaml/magnet/skewquad.py index f79e2c30..d942e03e 100644 --- a/pyaml/magnet/skewquad.py +++ b/pyaml/magnet/skewquad.py @@ -1,11 +1,9 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "SkewQuad" - -class ConfigModel(MagnetConfigModel): +class SkewQuadSchema(MagnetSchema): """Configuration model for SkewQuad magnet.""" ... @@ -16,9 +14,7 @@ class SkewQuad(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/skewsext.py b/pyaml/magnet/skewsext.py index 640d996e..675bca60 100644 --- a/pyaml/magnet/skewsext.py +++ b/pyaml/magnet/skewsext.py @@ -1,11 +1,9 @@ from ..lattice.polynom_info import PolynomInfo -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "SkewSext" - -class ConfigModel(MagnetConfigModel): +class SkewSextSchema(MagnetSchema): """Configuration model for SkewSext magnet.""" ... @@ -16,9 +14,7 @@ class SkewSext(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) diff --git a/pyaml/magnet/spline_model.py b/pyaml/magnet/spline_model.py index 2f589753..d30ff728 100644 --- a/pyaml/magnet/spline_model.py +++ b/pyaml/magnet/spline_model.py @@ -3,15 +3,12 @@ from scipy.interpolate import make_smoothing_spline from ..common.element import __pyaml_repr__ -from ..configuration.curve import Curve -from ..control.deviceaccess import DeviceAccess +from ..configuration.curve import Curve, CurveSchema +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "SplineMagnetModel" - -class ConfigModel(BaseModel): +class SplineMagnetModelSchema(BaseModel): """ Configuration model for spline magnet model @@ -34,10 +31,10 @@ class ConfigModel(BaseModel): passes through all the points of the curve. Default: 0.0 """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") - curve: Curve - powerconverter: DeviceAccess | None + curve: CurveSchema + powerconverter: DeviceAccessSchema | None calibration_factor: float = 1.0 calibration_offset: float = 0.0 crosstalk: float = 1.0 @@ -51,22 +48,25 @@ class SplineMagnetModel(MagnetModel): spline interpolation for a single function magnet """ - 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, + unit: str, + powerconverter: DeviceAccessSchema | None, + calibration_factor: float = 1.0, + calibration_offset: float = 0.0, + crosstalk: float = 1.0, + 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.powerconverter.unit() + self.__strength_unit = unit + self.__hardware_unit = powerconverter.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..f43f0389 100644 --- a/pyaml/magnet/vcorrector.py +++ b/pyaml/magnet/vcorrector.py @@ -1,13 +1,11 @@ from ..common import abstract from ..lattice.polynom_info import PolynomInfo from .corrector import RWCorrectorAngle -from .magnet import Magnet, MagnetConfigModel +from .magnet import Magnet, MagnetSchema +from .model import MagnetModel -# Define the main class name for this module -PYAMLCLASS = "VCorrector" - -class ConfigModel(MagnetConfigModel): +class VCorrectorSchema(MagnetSchema): """Configuration model for Vertical Corrector magnet.""" ... @@ -18,12 +16,10 @@ class VCorrector(Magnet): 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, description: str | None = None, lattice_names: str | None = None, model: MagnetModel = None + ): + super().__init__(name, description, lattice_names, model) self.__angle = RWCorrectorAngle(self) @property diff --git a/pyaml/rf/rf_plant.py b/pyaml/rf/rf_plant.py index 5c5bfe53..5fea06a9 100644 --- a/pyaml/rf/rf_plant.py +++ b/pyaml/rf/rf_plant.py @@ -8,19 +8,16 @@ from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel -from ..control.deviceaccess import DeviceAccess -from .rf_transmitter import RFTransmitter +from ..common.element import Element, ElementSchema +from ..control.deviceaccess import DeviceAccess, DeviceAccessSchema +from .rf_transmitter import RFTransmitter, RFTransmitterSchema -# Define the main class name for this module -PYAMLCLASS = "RFPlant" - -class ConfigModel(ElementConfigModel): - masterclock: DeviceAccess | None = None +class RFPlantSchema(ElementSchema): + masterclock: DeviceAccessSchema | None = None """Device to apply main RF frequency""" - transmitters: list[RFTransmitter] | None = None - """List of RF trasnmitters""" + transmitters: list[RFTransmitterSchema] | None = None + """List of RF transmitters""" class RFPlant(Element): @@ -28,23 +25,31 @@ class RFPlant(Element): Main RF object """ - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg - self.__frequency = None - self.__voltage = None + def __init__( + self, + name: str, + description: str | None = None, + lattice_names: str | None = None, + masterclock: DeviceAccess | None = None, + transmitters: list[RFTransmitter] | None = None, + ): + super().__init__(name, description, lattice_names) + self._masterclock = masterclock + self._transmitters = transmitters + self._frequency = None + self._voltage = None @property def frequency(self) -> abstract.ReadWriteFloatScalar: - if self.__frequency is None: + if self._frequency is None: raise PyAMLException(f"{str(self)} has no masterclock device defined") - return self.__frequency + return self._frequency @property def voltage(self) -> abstract.ReadWriteFloatScalar: - if self.__voltage is None: - raise PyAMLException(f"{str(self)} has no trasmitter device defined") - return self.__voltage + if self._voltage is None: + raise PyAMLException(f"{str(self)} has no transmitter device defined") + return self._voltage def attach( self, @@ -53,9 +58,9 @@ def attach( voltage: abstract.ReadWriteFloatScalar, ) -> Self: # Attach frequency attribute and returns a new reference - obj = self.__class__(self._cfg) - obj.__frequency = frequency - obj.__voltage = voltage + obj = self.__class__(self._name, self._description, self._lattice_names, self._masterclock, self._transmitters) + obj._frequency = frequency + obj._voltage = voltage obj._peer = peer return obj @@ -70,25 +75,25 @@ def __init__(self, transmitters: list[RFTransmitter]): transmitters : list[RFTransmitter] List of attached transmitters """ - self.__trans = transmitters + self._trans = transmitters def get(self) -> float: sum = 0 # Count only fundamental harmonic - for t in self.__trans: - if t._cfg.harmonic == 1.0: + for t in self._trans: + if t.harmonic == 1.0: sum += t.voltage.get() return sum def set(self, value: float): # Assume that sum of transmitter (fundamental harmonic) distribution is 1 - for t in self.__trans: - if t._cfg.harmonic == 1.0: - v = value * t._cfg.distribution + for t in self._trans: + if t.harmonic == 1.0: + v = value * t.distribution t.voltage.set(v) def set_and_wait(self, value: float): raise NotImplementedError("Not implemented yet.") def unit(self) -> str: - return self.__trans[0]._cfg.phase.unit() + return self._trans[0].phase.unit() diff --git a/pyaml/rf/rf_transmitter.py b/pyaml/rf/rf_transmitter.py index 013dbbf7..b474cb98 100644 --- a/pyaml/rf/rf_transmitter.py +++ b/pyaml/rf/rf_transmitter.py @@ -8,14 +8,11 @@ from .. import PyAMLException from ..common import abstract -from ..common.element import Element, ElementConfigModel +from ..common.element import Element, ElementSchema from ..control.deviceaccess import DeviceAccess -# Define the main class name for this module -PYAMLCLASS = "RFTransmitter" - -class ConfigModel(ElementConfigModel): +class RFTransmitterSchema(ElementSchema): """ Configuration model for RF Transmitter. @@ -34,7 +31,7 @@ class ConfigModel(ElementConfigModel): by default 1.0 """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") voltage: DeviceAccess | None = None phase: DeviceAccess | None = None @@ -48,11 +45,26 @@ class RFTransmitter(Element): Class that handle a RF transmitter """ - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg - self.__voltage = None - self.__phase = None + def __init__( + self, + name: str, + cavities: list[str], + description: str | None = None, + lattice_names: str | None = None, + voltage: DeviceAccess | None = None, + phase: DeviceAccess | None = None, + harmonic: float = 1.0, + distribution: float = 1.0, + ): + super().__init__(name) + self._cavities = cavities + self._voltage = voltage + self._phase = phase + self.harmonic = harmonic + self.distribution = distribution + + self._voltage = None + self._phase = None @property def voltage(self) -> abstract.ReadWriteFloatScalar: @@ -69,11 +81,9 @@ def voltage(self) -> abstract.ReadWriteFloatScalar: PyAMLException If transmitter is unattached or has no voltage device defined """ - if self.__voltage is None: - raise PyAMLException( - f"{str(self)} is unattached or has no voltage device defined" - ) - return self.__voltage + if self._voltage is None: + raise PyAMLException(f"{str(self)} is unattached or has no voltage device defined") + return self._voltage @property def phase(self) -> abstract.ReadWriteFloatScalar: @@ -90,11 +100,9 @@ def phase(self) -> abstract.ReadWriteFloatScalar: PyAMLException If transmitter is unattached or has no phase device defined """ - if self.__phase is None: - raise PyAMLException( - f"{str(self)} is unattached or has no phase device defined" - ) - return self.__phase + if self._phase is None: + raise PyAMLException(f"{str(self)} is unattached or has no phase device defined") + return self._phase def attach( self, @@ -120,8 +128,17 @@ def attach( A new attached instance of RFTransmitter """ # Attach voltage and phase attribute and returns a new reference - obj = self.__class__(self._cfg) - obj.__voltage = voltage - obj.__phase = phase + obj = self.__class__( + self._name, + self._cavities, + self.description, + self.lattice_names, + self.voltage, + self.phase, + self.harmonic, + self.distribution, + ) + obj._voltage = voltage + obj._phase = phase obj._peer = peer return obj diff --git a/pyaml/tuning_tools/chromaticity.py b/pyaml/tuning_tools/chromaticity.py index 27a4d892..280ad7c7 100644 --- a/pyaml/tuning_tools/chromaticity.py +++ b/pyaml/tuning_tools/chromaticity.py @@ -2,9 +2,9 @@ from typing import TYPE_CHECKING from .. import PyAMLException -from ..common.element import ElementConfigModel +from ..common.element import ElementSchema from ..diagnostics.chromaticity_monitor import ChomaticityMonitor -from .response_matrix_data import ResponseMatrixData +from .response_matrix_data import ResponseMatrixData, ResponseMatrixDataSchema from .tuning_tool import TuningTool if TYPE_CHECKING: @@ -21,11 +21,8 @@ logger = logging.getLogger(__name__) -# Define the main class name for this module -PYAMLCLASS = "Chromaticity" - -class ConfigModel(ElementConfigModel): +class ChromaticitySchema(ElementSchema): """ Configuration model for Tune @@ -41,7 +38,7 @@ class ConfigModel(ElementConfigModel): sextu_array_name: str chromaticty_monitor_name: str - response_matrix: str | ResponseMatrixData + response_matrix: str | ResponseMatrixDataSchema class Chromaticity(TuningTool): @@ -49,31 +46,34 @@ class Chromaticity(TuningTool): Class providing chromaticity adjustment tool """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name, + sextu_array_name: str, + chromaticty_monitor_name: str, + response_matrix: str | ResponseMatrixData, + ): """ Construct a chromaticity adjustment object. - Parameters - ---------- - cfg : ConfigModel - Configuration for the tune adjustment. """ - super().__init__(cfg.name) - self._cfg = cfg + super().__init__(name) + self._sextu_array_name = sextu_array_name + self._chromaticty_monitor_name = chromaticty_monitor_name 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(response_matrix) is str: try: - cfg.response_matrix = ResponseMatrixData.load(cfg.response_matrix) + response_matrix = ResponseMatrixData.load(response_matrix) except Exception as e: logger.warning(f"{str(e)}") - cfg.response_matrix = None + response_matrix = None # Invert matrix - if cfg.response_matrix: - self._response_matrix = np.array(cfg.response_matrix._cfg.matrix) + if response_matrix: + self._response_matrix = np.array(response_matrix._matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -84,7 +84,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 +95,19 @@ def load(self, load_path: Path): load_path : Path Filename of the :class:`~.ResponseMatrixData` to load """ - self._cfg.response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._cfg.response_matrix._cfg.matrix) + self._response_matrix = ResponseMatrixData.load(load_path) + self._response_matrix = np.array(self._response_matrix._matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property def _cm(self) -> "ChomaticityMonitor": self.check_peer() - return self.peer.get_chromaticity_monitor(self._cfg.chromaticty_monitor_name) + return self.peer.get_chromaticity_monitor(self._chromaticty_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): """ diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index 8b995fbf..8b5ca2cb 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -7,15 +7,13 @@ from ..common.constants import Action from ..common.element import ElementConfigModel -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from .measurement_tool import MeasurementTool, MeasurementToolSchema +from .response_matrix_data import ResponseMatrixDataSchema logger = logging.getLogger(__name__) -PYAMLCLASS = "ChromaticityResponseMatrix" - -class ConfigModel(MeasurementToolConfigModel): +class ChromaticityResponseMatrixSchema(MeasurementToolSchema): """ Configuration model for Tune response matrix @@ -29,7 +27,7 @@ class ConfigModel(MeasurementToolConfigModel): Delta strength used to get the response matrix """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") sextu_array_name: str chromaticity_name: str @@ -37,9 +35,17 @@ class ConfigModel(MeasurementToolConfigModel): class ChromaticityResponseMatrix(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg + def __init__( + self, + name: str, + sextu_array_name: str, + chromaticity_name: str, + sextu_delta: float, + ): + super().__init__(name) + self._sextu_array_name = (sextu_array_name,) + self._chromaticity_name = (chromaticity_name,) + self._sextu_delta = (sextu_delta,) self.aborted = False def measure( @@ -117,8 +123,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 +137,11 @@ def callback(action: Action, data:dict): return False initial_chroma = cm.chromaticity.get() - delta = sextu_delta if sextu_delta is not None else self._cfg.sextu_delta - nb_step = n_step if n_step is not None else self._cfg.n_step - nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas + delta = sextu_delta if sextu_delta is not None else self._sextu_delta + nb_step = n_step if n_step is not None else self._n_step + nb_meas = n_avg_meas if n_avg_meas is not None else self._n_avg_meas + sleep_step = sleep_between_step if sleep_between_step is not None else self._.sleep_between_step + sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._sleep_between_meas err = None aborted = False @@ -206,7 +212,7 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - mat = ResponseMatrixDataConfigModel( + mat = ResponseMatrixDataSchema( matrix=chromamat.T.tolist(), variable_names=sextus.names(), observable_names=[cm.get_name() + ".x", cm.get_name() + ".y"], diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index 3f094eec..4354671d 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -6,17 +6,15 @@ from pySC.apps.codes import DispersionCode from ..common.constants import Action -from ..common.element import ElementConfigModel +from ..common.element import ElementSchema from ..common.element_holder import ElementHolder from ..external.pySC_interface import pySCInterface from .measurement_tool import MeasurementTool logger = logging.getLogger(__name__) -PYAMLCLASS = "Dispersion" - -class ConfigModel(ElementConfigModel): +class DispersionSchema(ElementSchema): """ Configuration model for dispersion measurement @@ -30,7 +28,7 @@ class ConfigModel(ElementConfigModel): Frequency delta for measurement """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") bpm_array_name: str rf_plant_name: str @@ -38,13 +36,17 @@ class ConfigModel(ElementConfigModel): class Dispersion(MeasurementTool): - def __init__(self, cfg: ConfigModel): - super().__init__(cfg.name) - self._cfg = cfg - - self.bpm_array_name = cfg.bpm_array_name - self.rf_plant_name = cfg.rf_plant_name - self.frequency_delta = cfg.frequency_delta + def __init__( + self, + name: str, + bpm_array_name: str, + rf_plant_name: str, + frequency_delta: float, + ): + super().__init__(name) + self.bpm_array_name = bpm_array_name + self.rf_plant_name = rf_plant_name + self.frequency_delta = frequency_delta def measure( self, diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index e2e40955..fac10eb7 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -6,7 +6,7 @@ from pydantic import ConfigDict from ..common.constants import Action -from ..common.element import Element, ElementConfigModel +from ..common.element import Element, ElementSchema from ..common.exception import PyAMLException if TYPE_CHECKING: @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) -class MeasurementToolConfigModel(ElementConfigModel): +class MeasurementToolSchema(ElementSchema): """ Measurement tool configuration model @@ -35,7 +35,7 @@ class MeasurementToolConfigModel(ElementConfigModel): Default: 0 """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") n_step: Optional[int] = 1 sleep_between_step: Optional[float] = 0 diff --git a/pyaml/tuning_tools/orbit.py b/pyaml/tuning_tools/orbit.py index 2dff6a18..2e84b305 100644 --- a/pyaml/tuning_tools/orbit.py +++ b/pyaml/tuning_tools/orbit.py @@ -16,21 +16,19 @@ from pySC.apps import orbit_correction from ..arrays.magnet_array import MagnetArray -from ..common.element import Element, ElementConfigModel +from ..common.element import Element, ElementSchema from ..common.exception import PyAMLException from ..external.pySC_interface import pySCInterface from ..rf.rf_plant import RFPlant -from .orbit_response_matrix_data import OrbitResponseMatrixData +from .orbit_response_matrix_data import OrbitResponseMatrixData, OrbitResponseMatrixDataSchema from .tuning_tool import TuningTool logger = logging.getLogger(__name__) logging.getLogger("pyaml.external.pySC").setLevel(logging.WARNING) -PYAMLCLASS = "Orbit" - -class ConfigModel(ElementConfigModel): - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") +class OrbitSchema(ElementSchema): + model_config = ConfigDict(extra="forbid") bpm_array_name: str hcorr_array_name: str @@ -40,47 +38,58 @@ class ConfigModel(ElementConfigModel): singular_values_H: Optional[int] = None singular_values_V: Optional[int] = None virtual_target: float = 0 - response_matrix: Union[str, OrbitResponseMatrixData] + response_matrix: Union[str, OrbitResponseMatrixDataSchema] 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 + def __init__( + self, + name, + 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) + 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.virtual_target = cfg.virtual_target + self.virtual_target = 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) + 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)}") + response_matrix = None # Converts to self._pySC_response_matrix - if cfg.response_matrix: - self._set_response_matrix(cfg.response_matrix) + if response_matrix: + self._set_response_matrix(response_matrix) self._hcorr: MagnetArray = None self._vcorr: MagnetArray = None @@ -96,16 +105,17 @@ 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) - - def _set_response_matrix(self, mat): - m = mat._cfg.model_dump() - m["input_names"] = m.pop("variable_names") - m["output_names"] = m.pop("observable_names") - m["input_planes"] = m.pop("variable_planes") - m["output_planes"] = m.pop("observable_planes") - self._cfg.response_matrix = mat + self._response_matrix = OrbitResponseMatrixData.load(load_path) + self._set_response_matrix(self._response_matrix) + + def _set_response_matrix(self, mat: OrbitResponseMatrixData): + m = mat + m["input_names"] = m._variable_names + m["output_names"] = m._observable_names + m["input_planes"] = m._variable_planes + m["output_planes"] = m._observable_planes + + self._response_matrix = mat self._pySC_response_matrix = pySC_ResponseMatrix.model_validate(m) @property @@ -113,7 +123,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, @@ -313,11 +323,11 @@ def get_rf_weight(self) -> float: return self._pySC_response_matrix.rf_weight def post_init(self): - self._hcorr = self.peer.get_magnets(self._cfg.hcorr_array_name) - self._vcorr = self.peer.get_magnets(self._cfg.vcorr_array_name) + self._hcorr = self.peer.get_magnets(self._hcorr_array_name) + self._vcorr = self.peer.get_magnets(self._vcorr_array_name) hvElts = [] hvElts.extend(self._hcorr) hvElts.extend(self._vcorr) self._hvcorr = MagnetArray("", hvElts) - if self._cfg.rf_plant_name is not None: - self._rf_plant = self.peer.get_rf_plant(self._cfg.rf_plant_name) + if self._rf_plant_name is not None: + self._rf_plant = self.peer.get_rf_plant(self._rf_plant_name) diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 99c8d2fd..83db9b92 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -10,15 +10,13 @@ from ..common.constants import Action from ..common.element import ElementConfigModel from ..external.pySC_interface import pySCInterface -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .orbit_response_matrix_data import ConfigModel as OrbitResponseMatrixDataConfigModel +from .measurement_tool import MeasurementTool, MeasurementToolSchema +from .orbit_response_matrix_data import OrbitResponseMatrixDataSchema logger = logging.getLogger(__name__) -PYAMLCLASS = "OrbitResponseMatrix" - -class ConfigModel(MeasurementToolConfigModel): +class OrbitResponseMatrixSchema(MeasurementToolSchema): """ Configuration model for orbit response matrix measurement @@ -34,7 +32,7 @@ class ConfigModel(MeasurementToolConfigModel): Corrector delta for measurement """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") bpm_array_name: str hcorr_array_name: str @@ -43,14 +41,20 @@ class ConfigModel(MeasurementToolConfigModel): 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, + ): + 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 def measure( self, @@ -92,9 +96,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( @@ -161,7 +165,7 @@ def measure( self.latest_measurement.update(orm_data.model_dump()) self.latest_measurement["type"] = "pyaml.tuning_tools.orbit_response_matrix_data" - def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel: + def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataSchema: # all metadata is discarded here. Should we keep something? element_holder = self._peer @@ -189,5 +193,5 @@ def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataC "observable_planes": observable_planes, } - orm_data = OrbitResponseMatrixDataConfigModel(**orm_data_model) + orm_data = OrbitResponseMatrixDataSchema(**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..f4918f03 100644 --- a/pyaml/tuning_tools/orbit_response_matrix_data.py +++ b/pyaml/tuning_tools/orbit_response_matrix_data.py @@ -1,12 +1,10 @@ from typing import Optional -from .response_matrix_data import ConfigModel as ReponseMatrixDataConfigModel +from .response_matrix_data import ConfigModel as ReponseMatrixDataSchema from .response_matrix_data import ResponseMatrixData -PYAMLCLASS = "OrbitResponseMatrixData" - -class ConfigModel(ReponseMatrixDataConfigModel): +class OrbitResponseMatrixDataSchema(ReponseMatrixDataSchema): """ Configuration model for orbit response matrix @@ -30,5 +28,16 @@ class OrbitResponseMatrixData(ResponseMatrixData): Orbit response matrix loader """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + matrix: list[list[float]], + variable_names: Optional[list[str]], + observable_names: list[str], + rf_response: Optional[list[float]] = None, + variable_planes: Optional[list[str]] = None, + observable_planes: Optional[list[str]] = None, + ): + super.__init__(matrix, variable_names, observable_names) + self._rf_response = rf_response + self._variable_planes = variable_planes + self._observable_planes = observable_planes diff --git a/pyaml/tuning_tools/response_matrix_data.py b/pyaml/tuning_tools/response_matrix_data.py index 4cfd4800..a2fb2a91 100644 --- a/pyaml/tuning_tools/response_matrix_data.py +++ b/pyaml/tuning_tools/response_matrix_data.py @@ -9,10 +9,8 @@ from .. import PyAMLException from ..configuration.fileloader import load -PYAMLCLASS = "ResponseMatrixData" - -class ConfigModel(BaseModel): +class ResponseMatrixDataSchema(BaseModel): """ Base configuration model for response matrix @@ -36,8 +34,15 @@ class ResponseMatrixData(object): Generic response matrix loader """ - def __init__(self, cfg: ConfigModel): - self._cfg = cfg + def __init__( + self, + matrix: list[list[float]], + variable_names: Optional[list[str]], + observable_names: list[str], + ): + self._matrix = matrix + self._variable_names = variable_names + self._observable_names = observable_names @staticmethod def load(filename: str) -> "ResponseMatrixData": diff --git a/pyaml/tuning_tools/tune.py b/pyaml/tuning_tools/tune.py index 382092bf..24b33c3b 100644 --- a/pyaml/tuning_tools/tune.py +++ b/pyaml/tuning_tools/tune.py @@ -12,7 +12,7 @@ from .. import PyAMLException from ..common.element import ElementConfigModel -from .response_matrix_data import ResponseMatrixData +from .response_matrix_data import ResponseMatrixData, ResponseMatrixDataSchema from .tuning_tool import TuningTool if TYPE_CHECKING: @@ -21,11 +21,8 @@ logger = logging.getLogger(__name__) -# Define the main class name for this module -PYAMLCLASS = "Tune" - -class ConfigModel(ElementConfigModel): +class TuneSchema(ElementConfigModel): """ Configuration model for Tune @@ -42,7 +39,7 @@ class ConfigModel(ElementConfigModel): quad_array_name: str betatron_tune_name: str - response_matrix: str | ResponseMatrixData + response_matrix: str | ResponseMatrixDataSchema class Tune(TuningTool): @@ -50,31 +47,34 @@ class Tune(TuningTool): Class providing tune adjustment tool """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + quad_array_name: str, + betatron_tune_name: str, + response_matrix: str | ResponseMatrixData, + ): """ Construct a Tune adjustment object. - Parameters - ---------- - cfg : ConfigModel - Configuration for the tune adjustment. """ - super().__init__(cfg.name) - self._cfg = cfg + super().__init__(name) + self._quad_array_name = quad_array_name + self._betatron_tune_name = betatron_tune_name 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(response_matrix) is str: try: - cfg.response_matrix = ResponseMatrixData.load(cfg.response_matrix) + response_matrix = ResponseMatrixData.load(response_matrix) except Exception as e: logger.warning(f"{str(e)}") - cfg.response_matrix = None + response_matrix = None # Invert matrix - if cfg.response_matrix: - self._response_matrix = np.array(cfg.response_matrix._cfg.matrix) + if response_matrix: + self._response_matrix = np.array(response_matrix._matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) # TODO: Initialise first setpoint @@ -90,8 +90,8 @@ def load(self, load_path: Path): Filename of the :class:`~.ResponseMatrixData` to load """ - self._cfg.response_matrix = ResponseMatrixData.load(load_path) - self._response_matrix = np.array(self._cfg.response_matrix._cfg.matrix) + self._response_matrix = ResponseMatrixData.load(load_path) + self._response_matrix = np.array(self._response_matrix._matrix) self._correctionmat = np.linalg.pinv(self._response_matrix) @property @@ -99,17 +99,17 @@ def response_matrix(self) -> ResponseMatrixData | None: """ Return the response matrix if it has been loaded None otherwise """ - return self._cfg.response_matrix + return self._response_matrix @property def _tm(self) -> "BetatronTuneMonitor": self.check_peer() - return self.peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) + return self.peer.get_betatron_tune_monitor(self._betatron_tune_name) @property def _quads(self) -> "MagnetArray": self.check_peer() - return self.peer.get_magnets(self._cfg.quad_array_name) + return self.peer.get_magnets(self._quad_array_name) def get(self): """ diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 15f432dc..1fa1c946 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -6,15 +6,13 @@ from pydantic import ConfigDict from ..common.constants import Action -from .measurement_tool import MeasurementTool, MeasurementToolConfigModel -from .response_matrix_data import ConfigModel as ResponseMatrixDataConfigModel +from .measurement_tool import MeasurementTool, MeasurementToolSchema +from .response_matrix_data import ResponseMatrixDataSchema logger = logging.getLogger(__name__) -PYAMLCLASS = "TuneResponseMatrix" - -class ConfigModel(MeasurementToolConfigModel): +class TuneResponseMatrixSchema(MeasurementToolSchema): """ Configuration model for Tune response matrix @@ -28,7 +26,7 @@ class ConfigModel(MeasurementToolConfigModel): Delta strength used to get the response matrix """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + model_config = ConfigDict(extra="forbid") quad_array_name: str betatron_tune_name: str @@ -36,9 +34,17 @@ class ConfigModel(MeasurementToolConfigModel): 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, + ): + super().__init__(name) + self._quad_array_name = quad_array_name + self._betatron_tune_name = betatron_tune_name + self._quad_delta = quad_delta def measure( self, @@ -114,16 +120,16 @@ def callback(action: Action, data:dict): """ # Get devices self.check_peer() - quads = self._peer.get_magnets(self._cfg.quad_array_name) - tm = self._peer.get_betatron_tune_monitor(self._cfg.betatron_tune_name) + quads = self._peer.get_magnets(self._quad_array_name) + tm = self._peer.get_betatron_tune_monitor(self._betatron_tune_name) tunemat = np.zeros((len(quads), 2)) initial_tune = tm.tune.get() - delta = quad_delta if quad_delta is not None else self._cfg.quad_delta - nb_step = n_step if n_step is not None else self._cfg.n_step - nb_meas = n_avg_meas if n_avg_meas is not None else self._cfg.n_avg_meas - sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step - sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas + delta = quad_delta if quad_delta is not None else self._quad_delta + nb_step = n_step if n_step is not None else self._n_step + nb_meas = n_avg_meas if n_avg_meas is not None else self._n_avg_meas + sleep_step = sleep_between_step if sleep_between_step is not None else self._sleep_between_step + sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._sleep_between_meas self._register_callback(callback) self._init_measure("pyaml.tuning_tools.response_matrix_data") @@ -193,7 +199,7 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - mat = ResponseMatrixDataConfigModel( + mat = ResponseMatrixDataSchema( matrix=tunemat.T.tolist(), variable_names=quads.names(), observable_names=[tm.get_name() + ".x", tm.get_name() + ".y"],