30 modular linking between pyat and pyaml elements

pyaml/lattice/attribute_linker.py

@@ -0,0 +1,68 @@
+import at
+from pydantic import ConfigDict
+
+from pyaml.lattice.element import Element
+from pyaml.lattice.lattice_elements_linker import LinkerIdentifier, LinkerConfigModel, LatticeElementsLinker
+
+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.
+
+    Attributes
+    ----------
+    model_config : ConfigDict
+        Pydantic configuration allowing arbitrary field types and forbidding
+        unexpected extra keys.
+    """
+    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.
+
+    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.
+    """
+
+    def __init__(self, attribute_name:str, identifier):
+        self.attribute_name = attribute_name
+        self.identifier = identifier
+
+    def __repr__(self):
+        return f"{self.attribute_name}={self.identifier}"
+
+
+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).
+
+    Parameters
+    ----------
+    config_model : ConfigModel
+        The configuration model for the linking strategy.
+    """
+
+    def __init__(self, config_model:ConfigModel):
+        super().__init__(config_model)
+
+    def get_element_identifier(self, element: Element) -> LinkerIdentifier:
+        return PyAtAttributeIdentifier(self.linker_config_model.attribute_name, element.name)
+
+    def _test_at_element(self, identifier: PyAtAttributeIdentifier, element: at.Element) -> bool:
+        attr_value = getattr(element, identifier.attribute_name, None)
+        return attr_value == identifier.identifier

pyaml/lattice/lattice_elements_linker.py

@@ -0,0 +1,140 @@
+from abc import ABCMeta, abstractmethod
+from typing import Iterable
+
+import at
+from at import Lattice
+from pydantic import BaseModel, ConfigDict
+
+from pyaml import PyAMLException
+from pyaml.lattice.element import Element
+
+
+class LinkerConfigModel(BaseModel):
+    """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")
+
+
+class LinkerIdentifier(metaclass=ABCMeta):
+    """Abstract base class for identifiers used to match PyAML and PyAT elements.
+
+    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.
+    """
+    pass
+
+
+class LatticeElementsLinker(metaclass=ABCMeta):
+    """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).
+
+    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
+        self.lattice:Lattice = None
+
+    def set_lattice(self, lattice:Lattice):
+        self.lattice = lattice
+
+    @abstractmethod
+    def _test_at_element(self, identifier: LinkerIdentifier, element:at.Element) -> bool:
+        pass
+
+    @abstractmethod
+    def get_element_identifier(self, element:Element) -> LinkerIdentifier:
+        pass
+
+    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
+
+    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
+        into their corresponding PyAT elements according to the specific
+        linking strategy implemented.
+
+        Parameters
+        ----------
+        element_id : LinkerIdentifier or list of LinkerIdentifier
+            One or several identifiers describing which PyAT elements
+            to retrieve.
+
+        Returns
+        -------
+        list of at.Element
+            The list of matching PyAT elements found in the lattice.
+
+        Raises
+        ------
+        PyAMLException
+            If no element matches the given identifier(s).
+        """
+        if isinstance(element_id, LinkerIdentifier):
+            identifiers = [element_id]
+        else:
+            identifiers = element_id
+
+        results: list[at.Element] = []
+        for ident in identifiers:
+            results.extend(self._iter_matches(ident))
+
+        if not results:
+            raise PyAMLException(
+                f"No PyAT elements found for identifier(s): "
+                f"{', '.join(i.__repr__() for i in identifiers)}"
+            )
+        return results
+
+    def get_at_element(self, element_id:LinkerIdentifier) -> at.Element:
+        """Return a single PyAT element matching the given identifier.
+
+        Parameters
+        ----------
+        element_id : LinkerIdentifier
+            Identifier describing the PyAT element to retrieve.
+
+        Returns
+        -------
+        at.Element
+            The PyAT element matching the identifier.
+
+        Raises
+        ------
+        PyAMLException
+            If no element matches the identifier.
+        """
+        for elem in self._iter_matches(element_id):
+            return elem
+        raise PyAMLException(f"No PyAT element found for FamName: {element_id.__repr__()}")

pyaml/lattice/simulator.py

@@ -1,5 +1,8 @@
 from pydantic import BaseModel,ConfigDict
 import at
+
+from .attribute_linker import PyAtAttributeElementsLinker, ConfigModel as PyAtAttrLinkerConfigModel
+from .lattice_elements_linker import LatticeElementsLinker
 from ..configuration import get_root_folder
 from .element import Element
 from pathlib import Path
@@ -22,6 +25,8 @@ class ConfigModel(BaseModel):
     """AT lattice file"""
     mat_key: str = None
     """AT lattice ring name"""
+    linker: LatticeElementsLinker = None
+    """The linker configuration model"""
 
 class Simulator(ElementHolder):
     """
@@ -31,13 +36,16 @@ class Simulator(ElementHolder):
     def __init__(self, cfg: ConfigModel):
         super().__init__()
         self._cfg = cfg
+        self._linker = cfg.linker if cfg.linker else PyAtAttributeElementsLinker(PyAtAttrLinkerConfigModel(attribute_name="FamName"))
         path:Path = get_root_folder() / cfg.lattice
 
         if(self._cfg.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._linker.set_lattice(self.ring)
+
     def name(self) -> str:
        return self._cfg.name
 
@@ -54,23 +62,24 @@ def fill_device(self,elements:list[Element]):
        for e in elements:
           # Need conversion to physics unit to work with simulator
           if isinstance(e,Magnet):
-            current = RWHardwareScalar(self.get_at_elems(e.name),e.polynom,e.model) if e.model.has_physics() else None
-            strength = RWStrengthScalar(self.get_at_elems(e.name),e.polynom,e.model) if e.model.has_physics() else None
+            current = RWHardwareScalar(self.get_at_elems(e),e.polynom,e.model) if e.model.has_physics() else None
+            strength = RWStrengthScalar(self.get_at_elems(e),e.polynom,e.model) if e.model.has_physics() else None
             # Create a unique ref for this simulator
             m = e.attach(strength,current)
             self.add_magnet(str(m),m)
           elif isinstance(e,CombinedFunctionMagnet):
             self.add_magnet(str(e),e)
-            currents = RWHardwareArray(self.get_at_elems(e.name),e.polynoms,e.model) if e.model.has_physics() else None
-            strengths = RWStrengthArray(self.get_at_elems(e.name),e.polynoms,e.model) if e.model.has_physics() else None
+            currents = RWHardwareArray(self.get_at_elems(e),e.polynoms,e.model) if e.model.has_physics() else None
+            strengths = RWStrengthArray(self.get_at_elems(e),e.polynoms,e.model) if e.model.has_physics() else None
             # Create unique refs of each function for this simulator
             ms = e.attach(strengths,currents)
             for m in ms:
               self.add_magnet(str(m),m)
               self.add_magnet(str(m),m)
 
-    def get_at_elems(self,elementName:str) -> list[at.Element]:
-       elementList = [e for e in self.ring if e.FamName == elementName]
-       if not elementList:
-          raise Exception(f"{elementName} not found in lattice:{self._cfg.lattice}")
-       return elementList
+    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 Exception(f"{identifier} not found in lattice:{self._cfg.lattice}")
+       return element_list

tests/config/sr-attribute-linker.yaml

@@ -0,0 +1,28 @@
+type: pyaml.pyaml
+instruments:
+  - type: pyaml.instrument
+    name: sr
+    energy: 6e9
+    simulators:
+      - type: pyaml.lattice.simulator
+        lattice: sr/lattices/ebs.mat
+        name: design
+        linker:
+          type: pyaml.lattice.attribute_linker
+          attribute_name: FamName # equivalent to the default linker
+    data_folder: /data/store
+    arrays:
+      - type: pyaml.arrays.hcorrector
+        name: HCORR
+        elements:
+          - SH1A-C01-H
+          - SH1A-C02-H
+      - type: pyaml.arrays.vcorrector
+        name: VCORR
+        elements:
+          - SH1A-C01-V
+          - SH1A-C02-V
+    devices:
+      - sr/quadrupoles/QF1AC01.yaml
+      - sr/correctors/SH1AC01.yaml
+      - sr/correctors/SH1AC02.yaml

tests/conftest.py

@@ -1,4 +1,6 @@
 import types
+
+import at
 import pytest
 import subprocess
 import sys
@@ -165,3 +167,27 @@ def register_mock_strategy():
     Factory.register_strategy(strategy)
     yield
     Factory.remove_strategy(strategy)
+
+
+# -----------------------
+# Linkers fixtures
+# -----------------------
+
+
+@pytest.fixture
+def lattice_with_famnames() -> at.Lattice:
+    """Lattice with duplicate FamName to test multi-match and first-element behavior."""
+    qf1 = at.elements.Quadrupole('QF_1', 0.2); qf1.FamName = 'QF'
+    qf2 = at.elements.Quadrupole('QF_2', 0.25); qf2.FamName = 'QF'
+    qd1 = at.elements.Quadrupole('QD_1', 0.3); qd1.FamName = 'QD'
+    return at.Lattice([qf1, qf2, qd1], energy=3e9)
+
+
+@pytest.fixture
+def lattice_with_custom_attr() -> at.Lattice:
+    """Lattice where a custom attribute (e.g., 'Tag') is set on elements."""
+    d1 = at.elements.Drift('D1', 1.0);         setattr(d1, "Tag", "D1")
+    qf = at.elements.Quadrupole('QF', 0.2);    setattr(qf, "Tag", "QF")
+    qf2 = at.elements.Quadrupole('QF2', 0.2);  setattr(qf2, "Tag", "QF")
+    qd = at.elements.Quadrupole('QD', 0.3);    setattr(qd, "Tag", "QD")
+    return at.Lattice([d1, qf, qf2, qd], energy=3e9)