New serialized magnet holder API

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)

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]):

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)

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)

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)
 

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)
 

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)

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