Include Strength and IntegratedStrength in magnet elements

pyat/at/lattice/elements/magnet_elements.py

@@ -16,7 +16,7 @@
 
 import warnings
 from collections.abc import Generator
-from typing import Any
+from typing import Any, Sequence
 
 import numpy as np
 
@@ -34,7 +34,7 @@ class ThinMultipole(Element):
     """Thin multipole element"""
 
     _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ["PolynomA", "PolynomB"]
-    _conversions = dict(Element._conversions, K=float, H=float)
+    _conversions = dict(Element._conversions, K=float, H=float, O=float)
     _stacklevel = 4  # Stacklevel for warnings
 
     def __init__(self, family_name: str, poly_a, poly_b, **kwargs):
@@ -105,6 +105,13 @@ def check_polynom(keyname, arg):
                         and (kvalue.size < 3 or argvalue[2] != kvalue[2])
                     ):
                         raise seterr(family_name, "PolynomB", kvalue, "h", argvalue[2])
+                elif issubclass(self.__class__, Octupole):
+                    if (
+                        keyname == "PolynomB"
+                        and argvalue[3] != 0.0
+                        and (kvalue.size < 4 or argvalue[3] != kvalue[3])
+                    ):
+                        raise seterr(family_name, "PolynomB", kvalue, "o", argvalue[3]) 
                 elif np.any(argvalue) and not np.array_equiv(kvalue, argvalue):
                     raise seterr(family_name, keyname, kvalue, argname, argvalue)
                 else:
@@ -139,6 +146,7 @@ def check_strength(keyname, index):
         prmpolb = check_polynom("PolynomB", poly_b)
         check_strength("K", 1)
         check_strength("H", 2)
+        check_strength("O", 3)
         # Determine the length and order of PolynomA and PolynomB
         len_a, ord_a = getpol(prmpola)
         len_b, ord_b = getpol(prmpolb)
@@ -168,29 +176,80 @@ def __setattr__(self, key, value):
                 raise ValueError(f"MaxOrder must be smaller than {lmax}")
         super().__setattr__(key, value)
 
+    def _get_order(self):
+        order = getattr(self, 'DefaultOrder', None)
+        if order is None:
+            msg = (f"Default Order not defined for class "
+                   f"{self.__class__.__name__}, please access the strength "
+                   f"with PolynomA/B or use predefined magnet class")
+            raise AtError(msg)
+        return order
+
+    def _get_strength(self, order, integrated=False):
+        try:
+            k = self.PolynomB[order]
+        except IndexError:
+            return 0.0
+        if integrated:
+            l = getattr(self,'Length', 1.0)
+            k = k*l if l > 0.0 else k   
+        return k
+
+    def _set_strength(self, value, order, integrated=False):
+        if integrated:
+            l = getattr(self,'Length', 1.0)    
+            value = value/l if l > 0.0 else value
+        self.PolynomB[order] = value    
+
+
     # noinspection PyPep8Naming
     @property
     def K(self) -> float:
         """Focusing strength [mˆ-2]"""
-        arr = self.PolynomB
-        return 0.0 if len(arr) < 2 else float(arr[1])
+        return self._get_strength(1)
 
     # noinspection PyPep8Naming
     @K.setter
     def K(self, strength: float):
-        self.PolynomB[1] = strength
+        self._set_strength(strength, 1)
 
     # noinspection PyPep8Naming
     @property
     def H(self) -> float:
         """Sextupolar strength [mˆ-3]"""
-        arr = self.PolynomB
-        return 0.0 if len(arr) < 3 else float(arr[2])
+        return self._get_strength(2)
 
     # noinspection PyPep8Naming
     @H.setter
     def H(self, strength):
-        self.PolynomB[2] = strength
+        self._set_strength(strength, 2)
+
+        # noinspection PyPep8Naming
+    @property
+    def O(self) -> float:
+        """Octupolar strength [mˆ-4]"""
+        return self._get_strength(3)
+
+    # noinspection PyPep8Naming
+    @O.setter
+    def O(self, strength):
+        self._set_strength(strength, 3)
+
+    @property
+    def Strength(self): 
+        return self._get_strength(self._get_order())
+
+    @Strength.setter
+    def Strength(self, strength):
+        self._set_strength(strength, self._get_order())
+
+    @property
+    def IntegratedStrength(self):
+        return self._get_strength(self._get_order(), True) 
+
+    @IntegratedStrength.setter
+    def IntegratedStrength(self, strength):
+        self._set_strength(strength, self._get_order(), True)  
 
 
 class Multipole(_Radiative, LongElement, ThinMultipole):
@@ -231,7 +290,7 @@ def is_compatible(self, other) -> bool:
             return True
         else:
             return False
-
+        
 
 class Dipole(Radiative, Multipole):
     """Dipole element"""
@@ -436,12 +495,38 @@ def items(self) -> Generator[tuple[str, Any], None, None]:
 
 
 class Octupole(Multipole):
-    """Octupole element, with no changes from multipole at present"""
+    """Octupole element"""
 
-    _BUILD_ATTRIBUTES = Multipole._BUILD_ATTRIBUTES
+    _BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ["O"]
+    _stacklevel = 7  # Stacklevel for warnings
 
     DefaultOrder = 3
 
+    def __init__(self, family_name: str, length: float, o: float = 0.0, **kwargs):
+        """
+        Args:
+            family_name:    Name of the element
+            length:         Element length [m]
+            o:              Octupolar strength [mˆ-4]
+
+        Keyword Arguments:
+            PolynomB:           straight multipoles
+            PolynomA:           skew multipoles
+            MaxOrder:           Number of desired multipoles
+            NumIntSteps=10:     Number of integration steps
+            KickAngle:          Correction deviation angles (H, V)
+            FieldScaling:       Scaling factor applied to the magnetic field
+              (*PolynomA* and *PolynomB*)
+
+        Default PassMethod: ``StrMPoleSymplectic4Pass``
+        """
+        kwargs.setdefault("PassMethod", "StrMPoleSymplectic4Pass")
+        super().__init__(family_name, length, [], [0.0, 0.0, 0.0, o], **kwargs)
+
+    def items(self) -> Generator[tuple[str, Any], None, None]:
+        yield from super().items()
+        yield "O", self.O
+
 
 class Corrector(LongElement):
     """Corrector element"""

pyat/at/load/elegant.py

@@ -183,7 +183,7 @@ class KOCT(_ElegantElement):
     def to_at(self, l, k3=0.0, **params):  # noqa: E741
         poly_b = [0.0, 0.0, 0.0, k3 / 6.0]
         poly_a = [0.0, 0.0, 0.0, 0.0]
-        return [elt.Octupole(self.name, l, poly_a, poly_b, **params)]
+        return [elt.Octupole(self.name, l, k3/6.0, **params)]
 
     @classmethod
     def from_at(cls, kwargs):

pyat/at/load/madx.py

@@ -410,9 +410,11 @@ def from_at(cls, kwargs):
 class octupole(_MadElement):
     @mad_element
     def to_at(self, l, k3=0.0, k3s=0.0, **params):  # noqa: E741
-        poly_b = [0.0, 0.0, 0.0, k3 / 6.0]
-        poly_a = [0.0, 0.0, 0.0, k3s / 6.0]
-        return [elt.Octupole(self.name, l, poly_a, poly_b, **self.meval(params))]
+        atparams = {}
+        k3s = float(k3s)
+        if k3s !=0.0:
+            atparams["PolynomA"] = [0.0, 0.0, 0.0, k3s / 6.0]
+        return [elt.Octupole(self.name, l, k3/6.0, *atparams, **self.meval(params))]
 
     @classmethod
     def from_at(cls, kwargs):

pyat/test/test_basic_elements.py

@@ -119,9 +119,27 @@ def test_sextupole():
 
 
 def test_octupole():
-    o = elements.Octupole('octupole', 1.0, [], [0.0, 0.0, 0.0, 0.0])
+    o = elements.Octupole('octupole', 1.0)
     assert o.MaxOrder == 3
     assert len(o.PolynomA) == 4
+    assert o.O == 0.0
+    o = elements.Octupole('octupole', 1.0, -0.5)
+    assert o.MaxOrder == 3
+    assert len(o.PolynomA) == 4
+    assert o.O == -0.5
+    o = elements.Octupole('octupole', 1.0, PolynomB=[0.0, 0.0, 0.0, 0.005, 0.0])
+    assert o.MaxOrder == 3
+    assert len(o.PolynomA) == 5
+    assert o.O == 0.005
+    o = elements.Octupole('octupole', 1.0, PolynomB=[0.0, 0.0, 0.0, 0.005, 0.001])
+    assert o.MaxOrder == 4
+    assert len(o.PolynomA) == 5
+    assert o.O == 0.005
+    o = elements.Octupole('octupole', 1.0, PolynomB=[0.0, 0.0, 0.5, 0.005, 0.001],
+                           MaxOrder=3)
+    assert o.MaxOrder == 3
+    assert len(o.PolynomA) == 5
+    assert o.O == 0.005
 
 
 def test_thinmultipole():
@@ -146,7 +164,6 @@ def test_multipole():
     [
         (elements.Multipole, ["Multi", 1.0]),
         (elements.ThinMultipole, ["ThinMulti"]),
-        (elements.Octupole, ["Oct", 1.0]),
     ],
 )
 def test_PolynomA_strength_prioritisation(element_type, args):
@@ -177,7 +194,6 @@ def test_PolynomA_strength_prioritisation(element_type, args):
     [
         (elements.Multipole, ["Multi", 1.0]),
         (elements.ThinMultipole, ["ThinMulti"]),
-        (elements.Octupole, ["Oct", 1.0]),
     ],
 )
 def test_PolynomB_strength_prioritisation(element_type, args):
@@ -208,10 +224,9 @@ def test_PolynomB_strength_prioritisation(element_type, args):
     [
         (elements.Multipole, ["Multi", 1.0]),
         (elements.ThinMultipole, ["ThinMulti"]),
-        (elements.Octupole, ["Oct", 1.0]),
     ],
 )
-def test_K_and_H_strength_prioritisation(element_type, args):
+def test_K_and_H_and_O_strength_prioritisation(element_type, args):
     # Warning and no change when K is zero
     with pytest.warns(lattice.AtWarning):
         elem = element_type(*args, [], [0.0, 1.0], K=0.0)
@@ -220,6 +235,10 @@ def test_K_and_H_strength_prioritisation(element_type, args):
     with pytest.warns(lattice.AtWarning):
         elem = element_type(*args, [], [0.0, 0.0, 2.0], H=0.0)
     assert elem.PolynomB[2] == 2.0
+        # Warning and no change when O is zero
+    with pytest.warns(lattice.AtWarning):
+        elem = element_type(*args, [], [0.0, 0.0, 0.0, 3.0], O=0.0)
+    assert elem.PolynomB[3] == 3.0
     # Warning and no change when poly_b[1] and K are the same and non-zero
     with pytest.warns(lattice.AtWarning):
         elem = element_type(*args, [], [0.0, 1.0], K=1.0)
@@ -228,18 +247,28 @@ def test_K_and_H_strength_prioritisation(element_type, args):
     with pytest.warns(lattice.AtWarning):
         elem = element_type(*args, [], [0.0, 0.0, 2.0], H=2.0)
     assert elem.PolynomB[2] == 2.0
+        # Warning and no change when poly_b[3] and O are the same and non-zero
+    with pytest.warns(lattice.AtWarning):
+        elem = element_type(*args, [], [0.0, 0.0, 0.0, 3.0], O=3.0)
+    assert elem.PolynomB[3] == 3.0
     # Error when K is non-zero and different to poly_b[1], even if poly_b[1] is 0.0
     with pytest.raises(lattice.AtError):
         elem = element_type(*args, [], [0.0, 0.0], K=2.0)
     # Error when H is non-zero and different to poly_b[2], even if poly_b[2] is 0.0
     with pytest.raises(lattice.AtError):
         elem = element_type(*args, [], [0.0, 0.0, 0.0], H=2.0)
+    # Error when O is non-zero and different to poly_b[3], even if poly_b[3] is 0.0
+    with pytest.raises(lattice.AtError):
+        elem = element_type(*args, [], [0.0, 0.0 ,0.0, 0.0], O=3.0)
     # Error when len(poly_b)<2 and K is specified
     with pytest.raises(lattice.AtError):
         elem = element_type(*args, [], [0.0], K=1.0)
     # Error when len(poly_b)<3 and H is specified
     with pytest.raises(lattice.AtError):
         elem = element_type(*args, [], [0.0, 0.0], H=1.0)
+     # Error when len(poly_b)<4 and O is specified
+    with pytest.raises(lattice.AtError):
+        elem = element_type(*args, [], [0.0, 0.0, 0.0], O=1.0)       
 
 
 @pytest.mark.parametrize(

pyat/test/test_load_save.py

@@ -20,7 +20,8 @@ def simple_hmba(hmba_lattice):
     # Replace Multipoles by Octupoles
     for id in ring.get_uint32_index("OF*"):
         q = ring[id]
-        ring[id] = elt.Octupole(q.FamName, q.Length, q.PolynomA, q.PolynomB)
+        ring[id] = elt.Octupole(q.FamName, q.Length, PolynomA=q.PolynomA,
+                                PolynomB=q.PolynomB)
     # Disable useless e;ements
     for id in ring.get_uint32_index("SH*"):
         q = ring[id]