diff --git a/pyaml/common/element.py b/pyaml/common/element.py index 7314c8d1..cdf83239 100644 --- a/pyaml/common/element.py +++ b/pyaml/common/element.py @@ -232,7 +232,9 @@ class ElementConfigModel(BaseModel): - [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. + of lattice element is used for indexing. ``name`` refers to the + ``FamName`` of the lattice elements, or to the attribute configured + in the simulator linker when one is used. """ model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") @@ -260,7 +262,9 @@ class Element: | ``[name]#start_idx..end_idx`` : Lattice 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. + of lattice element is used for indexing. ``name`` refers to the + ``FamName`` of the lattice elements, or to the attribute configured + in the simulator linker when one is used. description : str or None, optional Human-readable element description. diff --git a/pyaml/lattice/attribute_linker.py b/pyaml/lattice/attribute_linker.py index affa5bcf..053ef10e 100644 --- a/pyaml/lattice/attribute_linker.py +++ b/pyaml/lattice/attribute_linker.py @@ -109,7 +109,23 @@ 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 self.get_identifier_for_name(element.get_name()) + + def get_identifier_for_name(self, name: str) -> LinkerIdentifier: + """ + Get the identifier for a name used in an element ``lattice_names`` selector. + + Parameters + ---------- + name : str + Expected value of the configured attribute. + + Returns + ------- + LinkerIdentifier + The identifier matching the lattice elements whose attribute equals ``name`` + """ + return PyAtAttributeIdentifier(self.linker_config_model.attribute_name, name) def _test_at_element(self, identifier: PyAtAttributeIdentifier, element: at.Element) -> bool: """ diff --git a/pyaml/lattice/lattice_elements_linker.py b/pyaml/lattice/lattice_elements_linker.py index ef11e559..87703667 100644 --- a/pyaml/lattice/lattice_elements_linker.py +++ b/pyaml/lattice/lattice_elements_linker.py @@ -70,6 +70,8 @@ class LatticeElementsLinker(ABC): Set the lattice for element linking. get_element_identifier(element) Get the identifier for linking an element. + get_identifier_for_name(name) + Get the identifier for a name given in an element ``lattice_names`` selector. get_at_elements(element_id) Return a list of PyAT elements matching the given identifiers. get_at_element(element_id) @@ -130,6 +132,31 @@ def get_element_identifier(self, element: Element) -> LinkerIdentifier: """ pass + def get_identifier_for_name(self, name: str) -> LinkerIdentifier: + """ + Get the identifier for a name used in an element ``lattice_names`` selector. + + Called when a PyAML element selects its lattice elements explicitly + (``lattice_names: list(a,b)``, ``a@0,1`` or ``a#0..2``): each ``name`` of the + selector is resolved through the linker instead of the element name. + + Parameters + ---------- + name : str + Name given in the ``lattice_names`` selector. + + Returns + ------- + LinkerIdentifier + The identifier matching the lattice element(s) designated by ``name``. + + Raises + ------ + NotImplementedError + If the linker does not support ``lattice_names`` selectors. + """ + raise NotImplementedError(f"{self.__class__.__name__} does not support 'lattice_names' selectors") + def _iter_matches(self, identifier: LinkerIdentifier) -> Iterable[at.Element]: """ Yield all lattice elements matching ``identifier``. diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 071d66a8..10bd537f 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -387,6 +387,12 @@ def get_at_elems(self, element: Element) -> list[at.Element]: """ Resolve a PyAML element to matching PyAT lattice elements. + The ``lattice_names`` selector of the element is honoured with or without a + linker. With a linker, the names of the selector are resolved through + ``linker.get_identifier_for_name`` (e.g. the configured attribute for + ``pyaml.lattice.attribute_linker``) instead of ``FamName``; without a + selector the linker uses the element name. + Parameters ---------- element : Element @@ -398,11 +404,28 @@ def get_at_elems(self, element: Element) -> list[at.Element]: Matching Accelerator Toolbox lattice elements. """ if self._linker: - 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._lattice}") - return element_list + if element.get_lattice_names() is None: + # By name + return self._linker.get_at_elements(self._linker.get_element_identifier(element)) + # By list + nameList = self.get_names(element) + if nameList is not None: + identifiers = [self._linker.get_identifier_for_name(name) for name in nameList] + return self._linker.get_at_elements(identifiers) + # By name or indices + name, indices = self.get_indices(element) + if name is None: + # Direct indexing in the ring + return [self.ring[idx] for idx in indices] + elts = self._linker.get_at_elements(self._linker.get_identifier_for_name(name)) + if indices is None: + return elts + try: + return [elts[idx] for idx in indices] + except IndexError as err: + strErr = f"{element.get_name()}: Invalid lattice_names index " + strErr += f"for {element.get_lattice_names()}, {len(elts)} element(s) found for {name}" + raise PyAMLException(strErr) from err else: # By list nameList = self.get_names(element) diff --git a/tests/config/EBSNames-linker.yaml b/tests/config/EBSNames-linker.yaml new file mode 100644 index 00000000..98ae50b2 --- /dev/null +++ b/tests/config/EBSNames-linker.yaml @@ -0,0 +1,95 @@ +type: pyaml.accelerator +facility: ESRF +machine: sr +energy: 6e9 +simulators: + - type: pyaml.lattice.simulator + lattice: sr/lattices/ebs-original-names.mat + name: design + linker: + type: pyaml.lattice.attribute_linker + attribute_name: FamName # equivalent to the default linker +controls: + - type: tango.pyaml.controlsystem + tango_host: ebs-simu-3:10000 + name: live + catalog: catalogs/ebs_catalogs.yaml +data_folder: /data/store +devices: +- type: pyaml.magnet.quadrupole + name: QF1E + model: + type: pyaml.magnet.linear_model + calibration_factor: 1.00054 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c04-e/current +- type: pyaml.magnet.quadrupole + name: QF1E-ALL + lattice_names: list(QF1E,QF1I) + model: + type: pyaml.magnet.linear_model + calibration_factor: 1.00054 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c04-e/current +- type: pyaml.magnet.quadrupole + name: QF1E-C05 + lattice_names: QF1E@1 + model: + type: pyaml.magnet.linear_model + calibration_factor: 0.996841 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c05-a/current +- type: pyaml.magnet.quadrupole + name: QF1E-C04-C05-C06 + lattice_names: QF1E@0,1,2 + model: + type: pyaml.magnet.linear_model + calibration_factor: 1.00191 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c05-e/current +- type: pyaml.magnet.quadrupole + name: QF1E-C04-C05-C06-2 + lattice_names: "@140,290,424" + model: + type: pyaml.magnet.linear_model + calibration_factor: 1.00191 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c05-e/current +- type: pyaml.magnet.quadrupole + name: QF1E-C04-C05-C06-3 + lattice_names: QF1E#0..3 + model: + type: pyaml.magnet.linear_model + calibration_factor: 1.00191 + crosstalk: 1.0 + curve: + type: pyaml.magnet.csvcurve + file: sr/magnet_models/QF1_strength.csv + unit: 1/m + hardware_unit: A + powerconverter: srmag/vps-qf1/c05-e/current diff --git a/tests/lattice/test_linkers.py b/tests/lattice/test_linkers.py index 75585da7..957c2cae 100644 --- a/tests/lattice/test_linkers.py +++ b/tests/lattice/test_linkers.py @@ -1,3 +1,4 @@ +import at import pytest from pyaml import PyAMLException @@ -6,6 +7,7 @@ PyAtAttributeElementsLinker, PyAtAttributeIdentifier, ) +from pyaml.lattice.simulator import Simulator # ----------------------- # Dummy PyAML Element @@ -13,14 +15,18 @@ class DummyPyAMLElement: - """Minimal stand-in for a PyAML Element: only provides .name.""" + """Minimal stand-in for a PyAML Element: only provides .name and .lattice_names.""" - def __init__(self, name: str): + def __init__(self, name: str, lattice_names: str | None = None): self._name = name + self._lattice_names = lattice_names def get_name(self) -> str: return self._name + def get_lattice_names(self) -> str | None: + return self._lattice_names + def test_conf_with_linker(): sr: Accelerator = Accelerator.load("tests/config/sr-attribute-linker.yaml") @@ -94,90 +100,27 @@ def check_index(ring, elts, indices): assert ring.index(e) == indices[idx] -def test_various_naming_addressing(): - sr = Accelerator.load("tests/config/EBSNames.yaml", ignore_external=True) +# Ring indices of the QF1E / QF1I quadrupoles of sr/lattices/ebs-original-names.mat +QF1E_INDICES = [ + 140, 290, 424, 576, 712, 848, 982, 1116, 1250, 1384, 1525, 1662, 1803, 1937, 2078, 2212, + 2348, 2482, 2616, 2762, 2898, 3046, 3182, 3316, 3452, 3593, 3732, 3873, 4009, 4150, 4286, +] # fmt: skip +QF1E_ALL_INDICES = QF1E_INDICES + [4430] + + +@pytest.mark.parametrize("config", ["tests/config/EBSNames.yaml", "tests/config/EBSNames-linker.yaml"]) +def test_various_naming_addressing(config): + """lattice_names selectors give the same elements with and without a linker.""" + sr = Accelerator.load(config, ignore_external=True) ring = sr.design.get_lattice() elts = sr.design.magnet.get("QF1E").strength._elements assert len(elts) == 31 - check_index( - ring, - elts, - [ - 140, - 290, - 424, - 576, - 712, - 848, - 982, - 1116, - 1250, - 1384, - 1525, - 1662, - 1803, - 1937, - 2078, - 2212, - 2348, - 2482, - 2616, - 2762, - 2898, - 3046, - 3182, - 3316, - 3452, - 3593, - 3732, - 3873, - 4009, - 4150, - 4286, - ], - ) + check_index(ring, elts, QF1E_INDICES) elts = sr.design.magnet.get("QF1E-ALL").strength._elements assert len(elts) == 32 - check_index( - ring, - elts, - [ - 140, - 290, - 424, - 576, - 712, - 848, - 982, - 1116, - 1250, - 1384, - 1525, - 1662, - 1803, - 1937, - 2078, - 2212, - 2348, - 2482, - 2616, - 2762, - 2898, - 3046, - 3182, - 3316, - 3452, - 3593, - 3732, - 3873, - 4009, - 4150, - 4286, - 4430, - ], - ) + check_index(ring, elts, QF1E_ALL_INDICES) elts = sr.design.magnet.get("QF1E-C05").strength._elements assert len(elts) == 1 @@ -194,3 +137,60 @@ def test_various_naming_addressing(): elts = sr.design.magnet.get("QF1E-C04-C05-C06-3").strength._elements assert len(elts) == 3 check_index(ring, elts, [140, 290, 424]) + + +# ----------------------- +# lattice_names selectors with a linker +# ----------------------- + + +@pytest.fixture +def simulator_with_tag_linker(lattice_with_custom_attr, tmp_path) -> Simulator: + """Simulator on the 'Tag' lattice: D1(Tag=D1) QF(Tag=QF) QF2(Tag=QF) QD(Tag=QD).""" + lattice_file = tmp_path / "tag_lattice.m" + at.save_m(lattice_with_custom_attr, str(lattice_file)) + return Simulator(name="design", lattice=str(lattice_file), linker=PyAtAttributeElementsLinker(attribute_name="Tag")) + + +def _tags(elements): + return [getattr(e, "Tag", None) for e in elements] + + +def test_linker_without_lattice_names_uses_element_name(simulator_with_tag_linker): + elts = simulator_with_tag_linker.get_at_elems(DummyPyAMLElement("QF")) + assert _tags(elts) == ["QF", "QF"] + + +def test_linker_lattice_names_list(simulator_with_tag_linker): + sim = simulator_with_tag_linker + assert _tags(sim.get_at_elems(DummyPyAMLElement("anything", "list(QF)"))) == ["QF", "QF"] + assert _tags(sim.get_at_elems(DummyPyAMLElement("anything", "list(QF,QD)"))) == ["QF", "QF", "QD"] + + +def test_linker_lattice_names_indices(simulator_with_tag_linker): + sim = simulator_with_tag_linker + ring = sim.ring + assert sim.get_at_elems(DummyPyAMLElement("x", "QF@1")) == [ring[2]] + assert sim.get_at_elems(DummyPyAMLElement("x", "QF@0,1")) == [ring[1], ring[2]] + assert sim.get_at_elems(DummyPyAMLElement("x", "QF#0..2")) == [ring[1], ring[2]] + # Empty name: direct indexing in the whole ring + assert sim.get_at_elems(DummyPyAMLElement("x", "@3")) == [ring[3]] + assert sim.get_at_elems(DummyPyAMLElement("x", "#1..3")) == [ring[1], ring[2]] + + +def test_linker_lattice_names_share_lattice_element(simulator_with_tag_linker): + """Two PyAML elements with different names can drive the same PyAT element.""" + sim = simulator_with_tag_linker + main = sim.get_at_elems(DummyPyAMLElement("QD")) + coil = sim.get_at_elems(DummyPyAMLElement("QD_H", "list(QD)")) + assert main == coil == [sim.ring[3]] + + +def test_linker_lattice_names_errors(simulator_with_tag_linker): + sim = simulator_with_tag_linker + with pytest.raises(PyAMLException): + sim.get_at_elems(DummyPyAMLElement("x", "list(ZZ)")) + with pytest.raises(PyAMLException): + sim.get_at_elems(DummyPyAMLElement("x", "ZZ@0")) + with pytest.raises(PyAMLException): + sim.get_at_elems(DummyPyAMLElement("x", "QF@5"))