From 2e299b96ace58195de09fb673ab6a81116c69308 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Mon, 13 Jul 2026 16:12:39 +0200 Subject: [PATCH 1/3] Remove config models from lattice subpackage. --- pyaml/lattice/attribute_linker.py | 52 +++++------- pyaml/lattice/lattice_elements_linker.py | 29 +++---- pyaml/lattice/simulator.py | 104 +++++++++++++---------- tests/lattice/test_linkers.py | 13 ++- 4 files changed, 98 insertions(+), 100 deletions(-) diff --git a/pyaml/lattice/attribute_linker.py b/pyaml/lattice/attribute_linker.py index 793726f1..6fa61eff 100644 --- a/pyaml/lattice/attribute_linker.py +++ b/pyaml/lattice/attribute_linker.py @@ -1,5 +1,6 @@ +from dataclasses import dataclass + import at -from pydantic import ConfigDict from pyaml.common.element import Element from pyaml.lattice.lattice_elements_linker import ( @@ -11,34 +12,29 @@ PYAMLCLASS = "PyAtAttributeElementsLinker" -class ConfigModel(LinkerConfigModel): - """Base configuration model for linker definitions. - - This class defines the configuration structure used to instantiate - a specific linking strategy. Each concrete implementation of a - `LatticeElementsLinker` may define its own subclass extending this model - to include additional configuration parameters. +@dataclass +class PyAtAttributeConfigModel(LinkerConfigModel): + """Configuration model for ``PyAtAttributeElementsLinker``. - Attributes + Parameters ---------- - model_config : ConfigDict - Pydantic configuration allowing arbitrary field types and forbidding - unexpected extra keys. + attribute_name : str + Name of the PyAT element attribute used to identify matching + lattice elements. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") attribute_name: str class PyAtAttributeIdentifier(LinkerIdentifier): - """Abstract base class for identifiers used to match PyAML and PyAT elements. + """Identifier based on a PyAT element attribute. - The identifier acts as an intermediate representation between the PyAML - configuration and the PyAT lattice. Its exact structure depends on the - linking strategy (e.g., family name, element index, or user-defined tag). - - Subclasses should define the fields and logic necessary to represent - a unique reference to one or more PyAT elements. + Parameters + ---------- + attribute_name : str + Name of the PyAT attribute used for matching. + identifier + Expected value of the attribute. """ def __init__(self, attribute_name: str, identifier): @@ -50,19 +46,15 @@ def __repr__(self): class PyAtAttributeElementsLinker(LatticeElementsLinker): - """Abstract base class defining the interface for PyAT–PyAML element linking. - - Implementations of this class define how PyAML elements are matched - to PyAT elements based on a given linking strategy (e.g., by family name, - by index, or by a custom attribute). + """Link lattice elements using a specified PyAT element attribute. - Parameters - ---------- - config_model : ConfigModel - The configuration model for the linking strategy. + This linker associates PyAML elements with PyAT elements by comparing + the value of a configurable PyAT attribute against the identifier + extracted from the PyAML element. """ - def __init__(self, config_model: ConfigModel): + def __init__(self, attribute_name: str): + config_model = PyAtAttributeConfigModel(attribute_name) super().__init__(config_model) def get_element_identifier(self, element: Element) -> LinkerIdentifier: diff --git a/pyaml/lattice/lattice_elements_linker.py b/pyaml/lattice/lattice_elements_linker.py index de83cf33..aaed4b64 100644 --- a/pyaml/lattice/lattice_elements_linker.py +++ b/pyaml/lattice/lattice_elements_linker.py @@ -1,33 +1,27 @@ -from abc import ABCMeta, abstractmethod +from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Iterable import at from at import Lattice -from pydantic import BaseModel, ConfigDict from pyaml import PyAMLException from pyaml.common.element import Element -class LinkerConfigModel(BaseModel): +class LinkerConfigModel(ABC): """Base configuration model for linker definitions. This class defines the configuration structure used to instantiate a specific linking strategy. Each concrete implementation of a `LatticeElementsLinker` may define its own subclass extending this model to include additional configuration parameters. - - Attributes - ---------- - model_config : ConfigDict - Pydantic configuration allowing arbitrary field types and forbidding - unexpected extra keys. """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + pass -class LinkerIdentifier(metaclass=ABCMeta): +class LinkerIdentifier(ABC): """Abstract base class for identifiers used to match PyAML and PyAT elements. The identifier acts as an intermediate representation between the PyAML @@ -41,7 +35,7 @@ class LinkerIdentifier(metaclass=ABCMeta): pass -class LatticeElementsLinker(metaclass=ABCMeta): +class LatticeElementsLinker(ABC): """Abstract base class defining the interface for PyAT–PyAML element linking. Implementations of this class define how PyAML elements are matched @@ -50,6 +44,8 @@ class LatticeElementsLinker(metaclass=ABCMeta): Parameters ---------- + attribute_name: str + linker_config_model : LinkerConfigModel The configuration model for the linking strategy. @@ -61,7 +57,7 @@ class LatticeElementsLinker(metaclass=ABCMeta): def __init__(self, linker_config_model: LinkerConfigModel): self.linker_config_model = linker_config_model - self.lattice: Lattice = None + self.lattice: Lattice | None = None def set_lattice(self, lattice: Lattice): """ @@ -97,9 +93,10 @@ def get_element_identifier(self, element: Element) -> LinkerIdentifier: def _iter_matches(self, identifier: LinkerIdentifier) -> Iterable[at.Element]: """Yield all elements in the lattice whose matches the identifier.""" - for elem in self.lattice: - if self._test_at_element(identifier, elem): - yield elem + if self.lattice: + for elem in self.lattice: + if self._test_at_element(identifier, elem): + yield elem def get_at_elements(self, element_id: LinkerIdentifier | list[LinkerIdentifier]) -> list[at.Element]: """Return a list of PyAT elements matching the given identifiers. diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index e518731a..af72e0ee 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -1,11 +1,10 @@ from pathlib import Path import at -from pydantic import BaseModel, ConfigDict from ..bpm.bpm import BPM from ..common.abstract_aggregator import ScalarAggregator -from ..common.element import Element +from ..common.element import Element, __pyaml_repr__ from ..common.element_holder import ElementHolder from ..common.exception import PyAMLException from ..configuration import ROOT @@ -38,61 +37,66 @@ from ..rf.rf_transmitter import RFTransmitter from ..tuning_tools.measurement_tool import MeasurementTool from ..tuning_tools.tuning_tool import TuningTool -from .attribute_linker import ( - ConfigModel as PyAtAttrLinkerConfigModel, -) -from .attribute_linker import ( - PyAtAttributeElementsLinker, -) +from .attribute_linker import PyAtAttributeConfigModel, PyAtAttributeElementsLinker from .lattice_elements_linker import LatticeElementsLinker # Define the main class name for this module PYAMLCLASS = "Simulator" -class ConfigModel(BaseModel): - """ - Configuration model for Simulator - - Parameters - ---------- - name : str - Simulator name - lattice : str - AT lattice file - mat_key : str, optional - AT lattice ring name - linker : LatticeElementsLinker, optional - The linker configuration model - description : str , optional - Simulator description - """ - - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - - name: str - lattice: str - mat_key: str = None - linker: LatticeElementsLinker = None - description: str | None = None +class Simulator(ElementHolder): + """Simulator interface backed by a PyAT lattice. + The simulator loads a PyAT lattice from disk and attaches PyAML + elements to their corresponding PyAT elements. Once attached, the + resulting device objects expose read/write interfaces operating + directly on the simulation model. -class Simulator(ElementHolder): - """ - Class that implements access to AT simulator + Elements may be matched either using the default name-based lookup + or a custom :class:`LatticeElementsLinker`. """ - def __init__(self, cfg: ConfigModel): + def __init__( + self, + name: str, + lattice: str, + mat_key: str | None = None, + linker: LatticeElementsLinker | None = None, + description: str | None = None, + ): + """Create a simulator from a PyAT lattice. + + Parameters + ---------- + name : str + Name of the simulator. + lattice : str + Path to the PyAT lattice file, relative to the configured + PyAML root directory. + mat_key : str, optional + Variable name of the lattice when loading a MATLAB ``.mat`` + lattice file. + linker : LatticeElementsLinker, optional + Custom linker used to associate PyAML elements with PyAT + lattice elements. If omitted, elements are matched by name. + description : str, optional + Human-readable description of the simulator. + """ + super().__init__() - self._cfg = cfg - path: Path = ROOT.get() / cfg.lattice + self._name = name + self._lattice = lattice + self._mat_key = mat_key + self.description = description - if self._cfg.mat_key is None: + path: Path = ROOT.get() / self._lattice + + 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 = linker if self._linker: self._linker.set_lattice(self.ring) else: @@ -104,16 +108,24 @@ def __init__(self, cfg: ConfigModel): self._elements_indexing[e.FamName] = [e] def name(self) -> str: - return self._cfg.name + return self._name + + @property + def lattice(self) -> str: + return self._lattice def get_lattice(self) -> at.Lattice: return self.ring - def get_description(self) -> str: + @property + def mat_key(self) -> str | None: + return self._mat_key + + def get_description(self) -> str | None: """ 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 @@ -345,4 +357,4 @@ def get_at_elems(self, element: Element) -> list[at.Element]: return [elts[idx] for idx in indices] def __repr__(self): - return repr(self._cfg).replace("ConfigModel", self.__class__.__name__) + return __pyaml_repr__(self) diff --git a/tests/lattice/test_linkers.py b/tests/lattice/test_linkers.py index 9e8734bb..30f14e90 100644 --- a/tests/lattice/test_linkers.py +++ b/tests/lattice/test_linkers.py @@ -2,9 +2,6 @@ from pyaml import PyAMLException from pyaml.accelerator import Accelerator -from pyaml.lattice.attribute_linker import ( - ConfigModel as AttrConfigModel, -) from pyaml.lattice.attribute_linker import ( PyAtAttributeElementsLinker, PyAtAttributeIdentifier, @@ -40,7 +37,7 @@ def test_conf_with_linker(): def test_attribute_identifier_from_pyaml_name(lattice_with_custom_attr): """We bind to AT element attribute 'Tag'; identifier value comes from PyAML element .name""" - linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) + linker = PyAtAttributeElementsLinker(attribute_name="Tag") linker.set_lattice(lattice_with_custom_attr) pyaml_elem = DummyPyAMLElement(name="QF") # identifier="QF" ident = linker.get_element_identifier(pyaml_elem) @@ -50,7 +47,7 @@ def test_attribute_identifier_from_pyaml_name(lattice_with_custom_attr): def test_attribute_get_at_elements_all_matches(lattice_with_custom_attr): - linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) + linker = PyAtAttributeElementsLinker(attribute_name="Tag") linker.set_lattice(lattice_with_custom_attr) ident = PyAtAttributeIdentifier("Tag", "QF") matches = linker.get_at_elements(ident) @@ -60,7 +57,7 @@ def test_attribute_get_at_elements_all_matches(lattice_with_custom_attr): def test_attribute_get_at_element_first_match(lattice_with_custom_attr): - linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) + linker = PyAtAttributeElementsLinker(attribute_name="Tag") linker.set_lattice(lattice_with_custom_attr) ident = PyAtAttributeIdentifier("Tag", "QD") first = linker.get_at_element(ident) @@ -73,7 +70,7 @@ def test_attribute_get_at_element_first_match(lattice_with_custom_attr): def test_attribute_no_match_raises(lattice_with_custom_attr): - linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) + linker = PyAtAttributeElementsLinker(attribute_name="Tag") linker.set_lattice(lattice_with_custom_attr) ident = PyAtAttributeIdentifier("Tag", "ZZ") with pytest.raises(PyAMLException): @@ -83,7 +80,7 @@ def test_attribute_no_match_raises(lattice_with_custom_attr): def test_attribute_multiple_identifiers_accumulate(lattice_with_custom_attr): - linker = PyAtAttributeElementsLinker(AttrConfigModel(attribute_name="Tag")) + linker = PyAtAttributeElementsLinker(attribute_name="Tag") linker.set_lattice(lattice_with_custom_attr) ids = [PyAtAttributeIdentifier("Tag", "QF"), PyAtAttributeIdentifier("Tag", "QD")] res = linker.get_at_elements(ids) From c4073594450aabab911c100c35845f04bca61c12 Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Fri, 10 Jul 2026 18:58:28 +0200 Subject: [PATCH 2/3] Change ValidationMeta to inherit from ABCMeta to be compatible for classes that inherit from abstract classes. --- pyaml/validation/validation_models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyaml/validation/validation_models.py b/pyaml/validation/validation_models.py index afecfb36..0683d86d 100644 --- a/pyaml/validation/validation_models.py +++ b/pyaml/validation/validation_models.py @@ -2,6 +2,7 @@ import inspect import logging +from abc import ABCMeta from typing import Any from pydantic import BaseModel, ConfigDict, create_model @@ -23,7 +24,7 @@ class ValidationModel(PyAMLBaseModel): model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") -class ValidationMeta(type): +class ValidationMeta(ABCMeta): """ Metaclass that validates constructor arguments before object creation. From 61d050d91248c1469cd98646beef6f5ec2d9e0df Mon Sep 17 00:00:00 2001 From: Teresia Olsson Date: Mon, 13 Jul 2026 16:30:56 +0200 Subject: [PATCH 3/3] Add dynamic validation. --- pyaml/lattice/attribute_linker.py | 5 ++++- pyaml/lattice/simulator.py | 5 +++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pyaml/lattice/attribute_linker.py b/pyaml/lattice/attribute_linker.py index 6fa61eff..4babfc6f 100644 --- a/pyaml/lattice/attribute_linker.py +++ b/pyaml/lattice/attribute_linker.py @@ -9,6 +9,8 @@ LinkerIdentifier, ) +from ..validation import DynamicValidation, register_schema + PYAMLCLASS = "PyAtAttributeElementsLinker" @@ -45,7 +47,8 @@ def __repr__(self): return f"{self.attribute_name}={self.identifier}" -class PyAtAttributeElementsLinker(LatticeElementsLinker): +@register_schema +class PyAtAttributeElementsLinker(LatticeElementsLinker, DynamicValidation): """Link lattice elements using a specified PyAT element attribute. This linker associates PyAML elements with PyAT elements by comparing diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index af72e0ee..48855a79 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -37,14 +37,15 @@ from ..rf.rf_transmitter import RFTransmitter from ..tuning_tools.measurement_tool import MeasurementTool from ..tuning_tools.tuning_tool import TuningTool -from .attribute_linker import PyAtAttributeConfigModel, PyAtAttributeElementsLinker +from ..validation import DynamicValidation, register_schema from .lattice_elements_linker import LatticeElementsLinker # Define the main class name for this module PYAMLCLASS = "Simulator" -class Simulator(ElementHolder): +@register_schema +class Simulator(ElementHolder, DynamicValidation): """Simulator interface backed by a PyAT lattice. The simulator loads a PyAT lattice from disk and attaches PyAML