From d8a9af2d8066565e490916103ea584199d6cbe1a Mon Sep 17 00:00:00 2001 From: gubaidulinvadim Date: Thu, 15 Oct 2026 10:48:01 +0200 Subject: [PATCH] Added Accelerator.add() utility function to add an element (= attach element to all available holders) --- pyaml/accelerator.py | 87 +++++++++++-- tests/test_accelerator_add.py | 223 ++++++++++++++++++++++++++++++++++ 2 files changed, 301 insertions(+), 9 deletions(-) create mode 100644 tests/test_accelerator_add.py diff --git a/pyaml/accelerator.py b/pyaml/accelerator.py index 0bc27301..9915e221 100644 --- a/pyaml/accelerator.py +++ b/pyaml/accelerator.py @@ -6,7 +6,7 @@ from .arrays.array import ArrayConfig from .common.element import Element, __pyaml_repr__ -from .common.exception import PyAMLConfigException +from .common.exception import PyAMLConfigException, PyAMLException from .common.holders.element_holder import ElementHolder from .configuration import ConfigurationManager, UnsupportedConfigurationRootError from .configuration.factory import Factory @@ -73,6 +73,8 @@ class Accelerator(DynamicValidation): Set the number of bucket. add_device(config, ignore_external=False) Dynamically add a device to this accelerator. + add(element) + Bind and register an element or tool in every mode. post_init() Method triggered after all initialisations are done get_description() @@ -208,6 +210,7 @@ def set_energy(self, E: float): Energy value to set in eV """ self._set_properties("_set_energy", E) + self._energy = E def set_mcf(self, alphac: float): """ @@ -219,6 +222,7 @@ def set_mcf(self, alphac: float): Moment compaction factor """ self._set_properties("_set_mcf", alphac) + self._alphac = alphac def set_harmonic_number(self, h: int): """ @@ -230,6 +234,78 @@ def set_harmonic_number(self, h: int): Number of bucket """ self._set_properties("_set_harmonic", h) + self._harmonic_number = h + + def add(self, element: Element) -> None: + """ + Bind and register an element or tool in every mode. + + Use the same backend binding as configured devices, including any + backend-specific mode restrictions. Newly registered instances receive + the current accelerator properties and run their post-initialization + hooks. Existing elements are not reinitialized. + + Parameters + ---------- + element : Element + Element to add, including tuning and measurement tools, magnets, + BPMs, RF plants, and diagnostics. + + Raises + ------ + TypeError + If the supplied object is not an Element. + PyAMLException + If the element's name is already registered in any mode, or + backend binding fails. + + Notes + ----- + Binding normally creates a separate attached instance in each mode, + leaving the supplied element reusable. Access these instances through + their mode's holders or by name, for example ``sr.design["ALTUNE"]``. + + Examples + -------- + .. code-block:: python + + tune = Tune( + "ALTUNE", + quad_array_name="QCORR", + response_matrix="respmat.json", + betatron_tune_name="BETATRON_TUNE2", + ) + sr.add(tune) + print(sr.design.tool.get("ALTUNE").tune_monitor) + print(sr.live.tool.get("ALTUNE").tune_monitor) + """ + if not isinstance(element, Element): + raise TypeError(f"Accelerator.add() expects an Element, got {type(element).__name__}") + + holders = (*self._simulators.values(), *self._controls.values()) + for holder in holders: + if element.name in holder._ALL: + raise PyAMLException(f"Duplicate element {element.name} in mode {holder.name()}") + + added: list[Element] = [] + for holder in holders: + existing_names = set(holder._ALL) + holder.fill_device([element]) + added.extend(instance for name, instance in holder._ALL.items() if name not in existing_names) + + for instance in added: + if self._energy is not None: + instance.set_energy(self._energy) + if self._alphac is not None: + instance.set_mcf(self._alphac) + if self._harmonic_number is not None: + instance.set_harmonic(self._harmonic_number) + for instance in added: + instance.post_init() + + if self._devices is None: + self._devices = [] + self._devices.append(element) def add_device(self, config: dict, ignore_external=False): """ @@ -250,14 +326,7 @@ 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._devices.append(dev) - if self._controls is not None: - for c in self._controls: - c.fill_device([dev]) - - if self._simulators is not None: - for s in self._simulators: - s.fill_device([dev]) + self.add(dev) def post_init(self): """ diff --git a/tests/test_accelerator_add.py b/tests/test_accelerator_add.py new file mode 100644 index 00000000..7c72a141 --- /dev/null +++ b/tests/test_accelerator_add.py @@ -0,0 +1,223 @@ +import copy + +import pytest + +from pyaml.accelerator import Accelerator +from pyaml.common.element import Element +from pyaml.common.exception import PyAMLException +from pyaml.diagnostics.tune_monitor import BetatronTuneMonitor +from pyaml.magnet.identity_model import IdentityMagnetModel +from pyaml.magnet.quadrupole import Quadrupole +from pyaml.tuning_tools.orbit import Orbit +from pyaml.tuning_tools.tune import Tune +from pyaml.tuning_tools.tune_response_matrix import TuneResponseMatrix +from pyaml.tuning_tools.tuning_tool import TuningTool + + +class TrackingElement(Element): + def __init__(self, name): + super().__init__(name) + self.energy = None + self.alphac = None + self.harmonic_number = None + self.post_init_count = 0 + + def _fill_device(self, holder): + instance = copy.copy(self) + instance._peer = holder + holder.add_element(instance) + + def set_energy(self, energy): + self.energy = energy + + def set_mcf(self, alphac): + self.alphac = alphac + + def set_harmonic(self, harmonic_number): + self.harmonic_number = harmonic_number + + def post_init(self): + self.check_peer() + assert self.energy is not None + self.post_init_count += 1 + + +@pytest.fixture +def accelerator_factory(simulator_fragment_factory): + def build(simulators=("design", "simulation"), controls=("live", "control")): + return Accelerator.from_dict( + { + "type": "pyaml.accelerator", + "facility": "TEST", + "machine": "sr", + "energy": 6e9, + "alphac": 1e-4, + "harmonic_number": 992, + "simulators": [simulator_fragment_factory(name) for name in simulators], + "controls": [{"type": "tango.pyaml.controlsystem", "name": name} for name in controls], + "devices": [], + } + ) + + return build + + +@pytest.mark.parametrize( + ("simulators", "controls"), + [ + (("design", "simulation"), ("live", "control")), + (("design", "simulation"), ()), + ((), ("live", "control")), + ((), ()), + (("shared",), ("shared",)), + ], +) +def test_add_binds_and_initializes_only_new_instances(accelerator_factory, simulators, controls): + accelerator = accelerator_factory(simulators, controls) + holders = [*accelerator.simulators().values(), *accelerator.controls().values()] + existing = TrackingElement("EXISTING") + accelerator.add(existing) + element = TrackingElement("NEW_ELEMENT") + + assert accelerator.add(element) is None + + instances = [] + for holder in holders: + attached = holder[element.name] + instances.append(attached) + assert attached is not element + assert attached.peer is holder + assert attached.peer.peer is accelerator + assert (attached.energy, attached.alphac, attached.harmonic_number) == (6e9, 1e-4, 992) + assert attached.post_init_count == 1 + assert holder[existing.name].post_init_count == 1 + assert len({id(attached) for attached in instances}) == len(holders) + assert element.peer is None + assert element.energy is None + assert element.post_init_count == 0 + assert accelerator._devices == [existing, element] + + +def test_add_uses_current_accelerator_properties(accelerator_factory): + accelerator = accelerator_factory() + accelerator.set_energy(3e9) + accelerator.set_mcf(2e-4) + accelerator.set_harmonic_number(400) + + accelerator.add(TrackingElement("NEW_ELEMENT")) + + for holder in accelerator.modes().values(): + attached = holder["NEW_ELEMENT"] + assert (attached.energy, attached.alphac, attached.harmonic_number) == (3e9, 2e-4, 400) + + +@pytest.mark.parametrize("tool_type", [Tune, TuneResponseMatrix]) +def test_add_binds_tune_tools_to_each_modes_devices(tool_type, config_root_path): + accelerator = Accelerator.load("tests/config/EBSTune.yaml") + kwargs = {"name": "ALTUNE", "quad_array_name": "QForTune", "betatron_tune_name": "BETATRON_TUNE"} + if tool_type is Tune: + kwargs["response_matrix"] = str(config_root_path / "tune_response.json") + else: + kwargs["quad_delta"] = 1e-4 + tool = tool_type(**kwargs) + + accelerator.add(tool) + + for holder in accelerator.modes().values(): + attached = holder.tool.get("ALTUNE") + assert isinstance(attached, tool_type) + assert attached is holder["ALTUNE"] + assert attached.peer is holder + assert attached.tune_monitor is holder.diagnostic.get("BETATRON_TUNE") + assert attached.quadrupoles is holder.magnets.get("QForTune") + assert tool.peer is None + assert accelerator.design.tool.get("ALTUNE") is not accelerator.live.tool.get("ALTUNE") + assert accelerator.yellow_pages.availability("ALTUNE") == {"design", "live"} + + +def test_add_initializes_orbit_handles(config_root_path): + accelerator = Accelerator.load("tests/config/EBSOrbit.yaml") + tool = Orbit( + name="ALTORBIT", + bpm_array_name="BPM", + hcorr_array_name="HCorr", + vcorr_array_name="VCorr", + response_matrix=str(config_root_path / "ideal_orm_disp.json"), + rf_plant_name="DEFAULT_RF_PLANT", + singular_values=1, + ) + + accelerator.add(tool) + + for holder in accelerator.modes().values(): + attached = holder.tool.get(tool.name) + assert attached.correctors.names() == holder.magnets.get("HCorr").names() + holder.magnets.get("VCorr").names() + assert all(corrector.peer is holder for corrector in attached.correctors) + assert attached.rf_plant is holder.rf.get("DEFAULT_RF_PLANT") + assert tool.peer is None + + +def test_add_registers_diagnostics_in_every_mode(accelerator_factory): + accelerator = accelerator_factory() + monitor = BetatronTuneMonitor("NEW_TUNE") + + accelerator.add(monitor) + + for holder in accelerator.modes().values(): + attached = holder.diagnostic.get(monitor.name) + assert attached is holder[monitor.name] + assert attached.peer is holder + assert attached is not monitor + assert attached._h == 992 + assert monitor.peer is None + assert accelerator.yellow_pages.availability(monitor.name) == set(accelerator.modes()) + + +def test_add_binds_magnet_strength_to_the_lattice(accelerator_factory): + accelerator = accelerator_factory(controls=()) + magnet = Quadrupole( + "NEW_QUAD", + lattice_names="QF1A-C01@0", + model=IdentityMagnetModel(physics="QUAD_STRENGTH", unit="1/m"), + ) + + accelerator.add(magnet) + accelerator.design.magnet.get(magnet.name).strength.set(0.5) + + assert accelerator.design.magnet.get(magnet.name).strength.get() == pytest.approx(0.5) + assert accelerator.simulation.magnet.get(magnet.name).strength.get() != pytest.approx(0.5) + assert magnet.peer is None + + +@pytest.mark.parametrize("existing_element", [TuningTool("DUPLICATE"), Element("DUPLICATE")]) +def test_add_rejects_duplicate_names_before_registering_any_instances(accelerator_factory, existing_element): + accelerator = accelerator_factory() + accelerator.control.add_element(existing_element) + + with pytest.raises(PyAMLException, match="Duplicate element DUPLICATE in mode control"): + accelerator.add(TrackingElement("DUPLICATE")) + + assert accelerator.control["DUPLICATE"] is existing_element + for holder in (accelerator.design, accelerator.simulation, accelerator.live): + assert "DUPLICATE" not in holder.get().names() + assert accelerator._devices == [] + + +@pytest.mark.parametrize("invalid_element", [None, {}, object()]) +def test_add_rejects_non_elements(accelerator_factory, invalid_element): + accelerator = accelerator_factory(simulators=(), controls=()) + + with pytest.raises(TypeError, match="Accelerator.add\\(\\) expects an Element"): + accelerator.add(invalid_element) + + +def test_add_device_uses_the_same_binding_and_initialization(accelerator_factory): + accelerator = accelerator_factory() + + accelerator.add_device({"type": "pyaml.diagnostics.tune_monitor", "name": "NEW_TUNE"}) + + for holder in accelerator.modes().values(): + monitor = holder.diagnostic.get("NEW_TUNE") + assert monitor.peer is holder + assert monitor._h == 992 + assert accelerator._devices[0].name == "NEW_TUNE"