From 9b44c57769f7c163580938d17d0c537c9810d9e5 Mon Sep 17 00:00:00 2001 From: PONS Date: Fri, 14 Aug 2026 15:36:21 +0200 Subject: [PATCH] New serialized magnet holder API --- pyaml/arrays/serialized_magnet.py | 2 +- pyaml/common/element_holder.py | 35 ++++++----------- pyaml/common/serialized_magnet_holder.py | 39 +++++++++++++++++++ pyaml/common/serialized_magnets_holder.py | 47 +++++++++++++++++++++++ pyaml/control/controlsystem.py | 2 +- pyaml/lattice/simulator.py | 2 +- tests/arrays/test_arrays.py | 2 +- tests/magnet/test_serialized_magnets.py | 10 ++--- 8 files changed, 107 insertions(+), 32 deletions(-) create mode 100644 pyaml/common/serialized_magnet_holder.py create mode 100644 pyaml/common/serialized_magnets_holder.py diff --git a/pyaml/arrays/serialized_magnet.py b/pyaml/arrays/serialized_magnet.py index e055b3f13..d6a41af6d 100644 --- a/pyaml/arrays/serialized_magnet.py +++ b/pyaml/arrays/serialized_magnet.py @@ -36,4 +36,4 @@ def fill_array(self, holder: ElementHolder): holder : ElementHolder The element holder to populate with serialized magnet array """ - holder.fill_serialized_magnet_array(self._name, self._elements) + holder.serialized_magnets.add(self._name, self._elements) diff --git a/pyaml/common/element_holder.py b/pyaml/common/element_holder.py index 471d6333b..847d41672 100644 --- a/pyaml/common/element_holder.py +++ b/pyaml/common/element_holder.py @@ -25,6 +25,8 @@ from .element import Element from .magnet_holder import MagnetHolder from .magnets_holder import MagnetsHolder +from .serialized_magnet_holder import SerializedMagnetHolder +from .serialized_magnets_holder import SerializedMagnetsHolder if TYPE_CHECKING: from ..accelerator import Accelerator @@ -75,6 +77,8 @@ def __init__(self): # Sub holders self._magnet_holder = MagnetHolder(self) self._magnets_holder = MagnetsHolder(self) + self._serialized_magnet_holder = SerializedMagnetHolder(self) + self._serialized_magnets_holder = SerializedMagnetsHolder(self) @property def peer(self) -> "Accelerator": @@ -91,6 +95,14 @@ def magnet(self) -> MagnetHolder: def magnets(self) -> MagnetsHolder: return self._magnets_holder + @property + def serialized_magnet(self) -> SerializedMagnetHolder: + return self._serialized_magnet_holder + + @property + def serialized_magnets(self) -> SerializedMagnetsHolder: + return self._serialized_magnets_holder + def post_init(self): """ Method triggered after all initialisations are done @@ -216,29 +228,6 @@ def get_cfm_magnets(self, name: str) -> CombinedFunctionMagnetArray: def get_all_cfm_magnets(self) -> list[CombinedFunctionMagnet]: return [value for key, value in self._CFM_MAGNETS.items()] - # Serialized magnets - - def fill_serialized_magnet_array(self, arrayName: str, elementNames: list[str]): - self._fill_array( - arrayName, - elementNames, - self.get_serialized_magnet, - SerializedMagnetsArray, - self._SERIALIZED_MAGNETS_ARRAYS, - ) - - def get_serialized_magnet(self, name: str) -> Magnet: - return self._get("SerializedMagnets", name, self._SERIALIZED_MAGNETS) - - def add_serialized_magnet(self, m: Magnet): - self._add(self._SERIALIZED_MAGNETS, m) - - def get_serialized_magnets(self, name: str) -> SerializedMagnetsArray: - return self._get("SerializedMagnets array", name, self._SERIALIZED_MAGNETS_ARRAYS) - - def get_all_serialized_magnets(self) -> list[SerializedMagnets]: - return [value for key, value in self._SERIALIZED_MAGNETS.items()] - # BPMs def fill_bpm_array(self, arrayName: str, elementNames: list[str]): diff --git a/pyaml/common/serialized_magnet_holder.py b/pyaml/common/serialized_magnet_holder.py new file mode 100644 index 000000000..c778270e1 --- /dev/null +++ b/pyaml/common/serialized_magnet_holder.py @@ -0,0 +1,39 @@ +from typing import TYPE_CHECKING + +from ..magnet.serialized_magnet import SerializedMagnets + +if TYPE_CHECKING: + from .element_holder import ElementHolder + + +class SerializedMagnetHolder: + def __init__(self, peer: "ElementHolder"): + self._peer = peer + + def all(self) -> list[SerializedMagnets]: + """ + Returns all serialized magnets as a list + """ + return [value for key, value in self._peer._SERIALIZED_MAGNETS.items()] + + def get(self, name: str) -> SerializedMagnets: + """ + Returns the specified magnet + + Parameters + ---------- + name : str + Name of the magnet + """ + return self._peer._get("Serialized magnet", name, self._peer._SERIALIZED_MAGNETS) + + def add(self, m: SerializedMagnets): + """ + Adds the specified magnet to the holder + + Parameters + ---------- + m : Magnet + Magnet to be added + """ + self._peer._add(self._peer._SERIALIZED_MAGNETS, m) diff --git a/pyaml/common/serialized_magnets_holder.py b/pyaml/common/serialized_magnets_holder.py new file mode 100644 index 000000000..44b160021 --- /dev/null +++ b/pyaml/common/serialized_magnets_holder.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ..arrays.serialized_magnet_array import SerializedMagnetsArray + +if TYPE_CHECKING: + from .element_holder import ElementHolder + + +class SerializedMagnetsHolder: + def __init__(self, peer: "ElementHolder"): + self._peer = peer + + def get(self, name: str = None) -> SerializedMagnetsArray: + """ + Returns the specified serialized magnet array or all serialized magnets if no name specified + + Parameters + ---------- + name : str + Name of the serialized magnet array + """ + if name is None: + return SerializedMagnetsArray("", self._peer.serialized_magnet.all()) + else: + return self._peer._get("serialized Magnet array", name, self._peer._SERIALIZED_MAGNETS_ARRAYS) + + def add(self, arrayName: str, elementNames: list[str]): + """ + Adds the specified serialied magnet array to the holder + + Parameters + ---------- + arrayName : str + Array name + elementNames : list[str] + List of magnet names + """ + self._peer._fill_array( + arrayName, + elementNames, + self._peer.serialized_magnet.get, + SerializedMagnetsArray, + self._peer._SERIALIZED_MAGNETS_ARRAYS, + ) + + def __getitem__(self, key): + return self.get().__getitem__(key) diff --git a/pyaml/control/controlsystem.py b/pyaml/control/controlsystem.py index 85798d983..eb7df8909 100644 --- a/pyaml/control/controlsystem.py +++ b/pyaml/control/controlsystem.py @@ -163,7 +163,7 @@ def fill_device(self, elements: list[Element]): currents.append(current) strengths.append(strength) ms = e.attach(self, strengths, currents) - self.add_serialized_magnet(ms[0]) + self.serialized_magnet.add(ms[0]) for m in ms[1:]: self.magnet.add(m) diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 339bb8578..d67f656ca 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -188,7 +188,7 @@ def fill_device(self, elements: list[Element]): linked_currents.append(current) linked_strengths.append(strength) ms = e.attach(self, linked_strengths, linked_currents) - self.add_serialized_magnet(ms[0]) + self.serialized_magnet.add(ms[0]) for m in ms[1:]: self.magnet.add(m) diff --git a/tests/arrays/test_arrays.py b/tests/arrays/test_arrays.py index cc8cff7aa..f1466dc68 100644 --- a/tests/arrays/test_arrays.py +++ b/tests/arrays/test_arrays.py @@ -249,7 +249,7 @@ def test_arrays(install_test_package): ) def test_serialized_magnets_arrays(sr_file): sr: Accelerator = Accelerator.load(sr_file, include_locations=False, ignore_external=True) - the_serie = sr.design.get_serialized_magnets("series") + the_serie = sr.design.serialized_magnets.get("series") strength = the_serie.strengths.get() assert len(strength) == 1 print(strength) diff --git a/tests/magnet/test_serialized_magnets.py b/tests/magnet/test_serialized_magnets.py index 130379e3d..1e95ca706 100644 --- a/tests/magnet/test_serialized_magnets.py +++ b/tests/magnet/test_serialized_magnets.py @@ -59,7 +59,7 @@ def test_magnet_modification(sr_file): print(sr.yellow_pages) - sm: SerializedMagnets = sr.design.get_serialized_magnet("mySeriesOfMagnets") + sm: SerializedMagnets = sr.design.serialized_magnet.get("mySeriesOfMagnets") element_names = sm._SerializedMagnets__elements lattice = sr.design.get_lattice() @@ -110,11 +110,11 @@ def test_tune(sr_file): sr = Accelerator.load(sr_file, include_locations=False, ignore_external=True) sr.design.get_lattice().disable_6d() - m = sr.design.get_serialized_magnet("QF1A") + m = sr.design.serialized_magnet.get("QF1A") print(f"m.strength.get()={m.strength.get()}") assert len(m.get_magnets()) == m.get_nb_magnets() - quadForTuneDesign = sr.design.get_serialized_magnets("QForTune") + quadForTuneDesign = sr.design.serialized_magnets.get("QForTune") tune_monitor = sr.design.get_betatron_tune_monitor("BETATRON_TUNE") # Build tune response matrix tunemat = np.zeros((len(quadForTuneDesign), 2)) @@ -162,7 +162,7 @@ def test_tune(sr_file): ) def test_get_device_names(sr_file): sr: Accelerator = Accelerator.load(sr_file, include_locations=False, ignore_external=True) - sm: SerializedMagnets = sr.design.get_serialized_magnet("QF1A") + sm: SerializedMagnets = sr.design.serialized_magnet.get("QF1A") device_names = sm.get_device_names() @@ -192,7 +192,7 @@ def get_strengths_from_lattice(sr: Accelerator, sm: SerializedMagnets) -> list: ) def test_strength_computation(sr_file): sr: Accelerator = Accelerator.load(sr_file, include_locations=False, ignore_external=True) - sm: SerializedMagnets = sr.design.get_serialized_magnet("QF1A") + sm: SerializedMagnets = sr.design.serialized_magnet.get("QF1A") assert sm.get_nb_magnets() == 31 sm.strength.set(24.0) assert abs(sm.strength.get() - 24.0) < 1e-3