VariableMultipole compatibility between matlab AT and pyat

atmat/lattice/element_creation/atvariablemultipole.m

@@ -18,17 +18,17 @@
 %    FREQUENCYB     Frequency of SINE excitation for PolynomB
 %    PHASEA         Phase of SINE excitation for PolynomA
 %    PHASEB         Phase of SINE excitation for PolynomB
-%	 MAXORDER       Order of the multipole for a scalar amplitude
+%    MAXORDER       Order of the multipole for a scalar amplitude
 %    SEED           Input seed for the random number generator
 %    FUNCA          ARBITRARY excitation turn-by-turn kick list for PolynomA
 %    FUNCB          ARBITRARY excitation turn-by-turn kick list for PolynomB
 %    PERIODIC       If true (default) the user input kick list is repeated
 %    RAMPS          Vector (t0, t1, t2, t3) in turn number to define the ramping of the excitation
 %                   * t<t0: excitation amlpitude is zero
-%                   * t0<t<t1: exciation amplitude is linearly ramped up
-%                   * t1<t<t2: exciation amplitude is constant             
-%                   * t2<t<t3: exciation amplitude is linearly ramped down
-%                   * t3<t: exciation amplitude is zero 
+%                   * t0<t<t1: excitation amplitude is linearly ramped up
+%                   * t1<t<t2: excitation amplitude is constant
+%                   * t2<t<t3: excitation amplitude is linearly ramped down
+%                   * t3<t:    excitation amplitude is zero
 %
 %  OUTPUTS
 %  1. ELEM - Structure with the AT element
@@ -108,15 +108,21 @@ function setsine(rsrc, ab)
     function rsrc = setmaxorder(rsrc)
         if isfield(rsrc,'AmplitudeA')
             mxa=find(abs(rsrc.AmplitudeA)>0,1,'last');
+            if isempty(mxa)
+                mxa=1;
+            end
         else
             mxa=0;
         end
         if isfield(rsrc,'AmplitudeB')
             mxb=find(abs(rsrc.AmplitudeB)>0,1,'last');
+            if isempty(mxb)
+                mxb=1;
+            end
         else
             mxb=0;
         end
-        mxab=max([mxa,mxb,rsrc.MaxOrder-1]);
+        mxab=max([mxa,mxb,rsrc.MaxOrder+1]);
         rsrc.MaxOrder=mxab-1;
         if isfield(rsrc,'AmplitudeA')
             rsrc.AmplitudeA(mxa+1:mxab)=0;

pyat/at/lattice/variable_elements.py

@@ -1,64 +1,85 @@
-import numpy
+from __future__ import annotations
+
+from datetime import datetime
+from enum import IntEnum
+
+import numpy as np
+
 from .elements import Element, _array
 from .utils import AtError
-from typing import Optional, Union
-from enum import IntEnum
-from datetime import datetime
 
 
 class ACMode(IntEnum):
-    """Class to define the excitation types"""
+    """Class to define the excitation types."""
+
     SINE = 0
     WHITENOISE = 1
     ARBITRARY = 2
 
 
 class VariableMultipole(Element):
-    """Class to generate an AT variable thin multipole element
-    """
+    """Class to generate an AT variable thin multipole element."""
+
     _BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES
-    _conversions = dict(Element._conversions,
-                        Mode=int,
-                        AmplitudeA=_array, AmplitudeB=_array,
-                        FrequencyA=float, FrequencyB=float,
-                        PhaseA=float, PhaseB=float,
-                        Seed=int, NSamplesA=int, NSamplesB=int,
-                        FuncA=_array, FuncB=_array, Ramps=_array,
-                        Periodic=bool)
-
-    def __init__(self, family_name, AmplitudeA=None, AmplitudeB=None,
-                 mode=ACMode.SINE, **kwargs):
+    _conversions = dict(
+        Element._conversions,
+        Mode=int,
+        AmplitudeA=_array,
+        AmplitudeB=_array,
+        FrequencyA=float,
+        FrequencyB=float,
+        PhaseA=float,
+        PhaseB=float,
+        Seed=int,
+        NSamplesA=int,
+        NSamplesB=int,
+        FuncA=_array,
+        FuncB=_array,
+        Ramps=_array,
+        Periodic=bool,
+    )
+
+    def __init__(self, family_name: str, **kwargs: dict[str, any]):
         # noinspection PyUnresolvedReferences
         r"""
+        Create VariableMultipole.
+
         Parameters:
             family_name(str):    Element name
-
         Keyword Arguments:
-            AmplitudeA(list,float): Amplitude of the excitation for PolynomA. Default None
-            AmplitudeB(list,float): Amplitude of the excitation for PolynomB. Default None
+            AmplitudeA(list,float): Amplitude of the excitation for PolynomA.
+                Default None
+            AmplitudeB(list,float): Amplitude of the excitation for PolynomB.
+                Default None
             mode(ACMode): defines the evaluation grid. Default ACMode.SINE
-
               * :py:attr:`.ACMode.SINE`: sine function
               * :py:attr:`.ACMode.WHITENOISE`: gaussian white noise
-              * :py:attr:`.GridMode.ARBITRARY`: user defined turn-by-turn kick list
+              * :py:attr:`.ACMode.ARBITRARY`: user defined turn-by-turn kick list
             FrequencyA(float): Frequency of the sine excitation for PolynomA
             FrequencyB(float): Frequency of the sine excitation for PolynomB
             PhaseA(float): Phase of the sine excitation for PolynomA. Default 0
             PhaseB(float): Phase of the sine excitation for PolynomB. Default 0
             MaxOrder(int): Order of the multipole for scalar amplitude. Default 0
+                It is overwritten with the length of Amplitude(A,B)
             Seed(int): Seed of the random number generator for white
                        noise excitation. Default datetime.now()
             FuncA(list): User defined tbt kick list for PolynomA
             FuncB(list): User defined tbt kick list for PolynomB
             Periodic(bool): If True (default) the user defined kick is repeated
-            Ramps(list): Vector (t0, t1, t2, t3) in turn number to define the ramping 
-                         of the excitation
+            Ramps(list): Vector (t0, t1, t2, t3) in turn number to define
+                        the ramping of the excitation
 
               * ``t<t0``: excitation amplitude is zero
-              * ``t0<t<t1``: exciation amplitude is linearly ramped up
-              * ``t1<t<t2``: exciation amplitude is constant             
-              * ``t2<t<t3``: exciation amplitude is linearly ramped down
-              * ``t3<t``: exciation amplitude is zero           
+              * ``t0<t<t1``: excitation amplitude is linearly ramped up
+              * ``t1<t<t2``: excitation amplitude is constant
+              * ``t2<t<t3``: excitation amplitude is linearly ramped down
+              * ``t3<t``: excitation amplitude is zero
+        Raises:
+            AtError if none of AmplitudeA or AmplitudeB is passed.
+            AtError if ramp is not vector of length 4 when using Ramps
+            AtError when Frequency is not defined if using Mode ``ACMode.SINE``
+            AtError when Funct is not defined if using Mode ``ACMode.ARBITRARY``
+
 
         Examples:
 
@@ -68,74 +89,103 @@ def __init__(self, family_name, AmplitudeA=None, AmplitudeB=None,
 
         .. note::
 
-            * For all excitation modes all least one amplitude has to be provided.
+            * If no parameters are given it will be initialized as IdentityPass.
+            * For all excitation modes at least one amplitudes (A or B) has
+              to be provided.
               The default excitation is ``ACMode.SINE``
             * For ``mode=ACMode.SINE`` the ``Frequency(A,B)`` corresponding to the
-              ``Amplitude(A,B)`` has to be provided
+              ``Amplitude(A,B)`` has to be provided.
             * For ``mode=ACMode.ARBITRARY`` the ``Func(A,B)`` corresponding to the
-              ``Amplitude(A,B)`` has to be provided
+              ``Amplitude(A,B)`` has to be provided.
         """
-        kwargs.setdefault('PassMethod', 'VariableThinMPolePass')
-        self.MaxOrder = kwargs.pop('MaxOrder', 0)
-        self.Periodic = kwargs.pop('Periodic', True)
-        self.Mode = int(mode)
-        if AmplitudeA is None and AmplitudeB is None:
-            raise AtError('Please provide at least one amplitude for A or B')
-        AmplitudeB = self._set_params(AmplitudeB, mode, 'B', **kwargs)
-        AmplitudeA = self._set_params(AmplitudeA, mode, 'A', **kwargs)
-        self._setmaxorder(AmplitudeA, AmplitudeB)
-        if mode == ACMode.WHITENOISE:
-            self.Seed = kwargs.pop('Seed', datetime.now().timestamp())
-        self.PolynomA = numpy.zeros(self.MaxOrder+1)
-        self.PolynomB = numpy.zeros(self.MaxOrder+1)
-        ramps = kwargs.pop('Ramps', None)
-        if ramps is not None:
-            assert len(ramps)==4, \
-                'Ramps has to be a vector with 4 elements'
-            self.Ramps = ramps
-        super(VariableMultipole, self).__init__(family_name, **kwargs)
-
-    def _setmaxorder(self, ampa, ampb):
+        if len(kwargs) > 0:
+            self.FamName = family_name
+            if "AmplitudeA" not in kwargs and "AmplitudeB" not in kwargs:
+                raise AtError("Please provide at least one amplitude for A or B")
+            # start setting up Amplitudes and modes
+            # fist modes are called differently
+            modepyatinput = kwargs.pop("mode", ACMode.SINE)
+            modefromdict = kwargs.pop("Mode", None)
+            if modefromdict is not None:
+                mode = modefromdict  # matlab class
+            else:
+                mode = modepyatinput
+            self.Mode = int(mode)
+            # MaxOrder is later overwritten by lenth of amplitudes
+            self.MaxOrder = kwargs.get("MaxOrder", 0)
+            amplitudea = None
+            amplitudeb = None
+            if "AmplitudeA" in kwargs:
+                amplitudea = kwargs.pop("AmplitudeA", None)
+                amplitudea = self._set_params(amplitudea, mode, "A", **kwargs)
+            if "AmplitudeB" in kwargs:
+                amplitudeb = kwargs.pop("AmplitudeB", None)
+                amplitudeb = self._set_params(amplitudeb, mode, "B", **kwargs)
+            if amplitudea is not None:
+                self.AmplitudeA = amplitudea
+            if amplitudeb is not None:
+                self.AmplitudeB = amplitudeb
+            # end setting up Amplitudes and modes
+            kwargs.setdefault("PassMethod", "VariableThinMPolePass")
+            # this overwrites MaxOrder
+            self._setmaxorder(amplitudea, amplitudeb)
+            if mode == ACMode.WHITENOISE and "Seed" not in kwargs:
+                kwargs.update({"Seed": datetime.now().timestamp()})
+            self.Periodic = kwargs.pop("Periodic", True)
+            self.PolynomA = kwargs.pop("PolynomA", np.zeros(self.MaxOrder + 1))
+            self.PolynomB = kwargs.pop("PolynomB", np.zeros(self.MaxOrder + 1))
+            # check ramps
+            ramps = kwargs.pop("Ramps", None)
+            if ramps is not None:
+                if len(ramps) != 4:
+                    raise AtError("Ramps has to be a vector with 4 elements")
+                self.Ramps = ramps
+        # fill in super class
+        super().__init__(family_name, **kwargs)
+
+    def _setmaxorder(self, ampa: np.ndarray, ampb: np.ndarray):
         mxa, mxb = 0, 0
         if ampa is not None:
-            mxa = numpy.max(numpy.nonzero(ampa))
+            mxa = np.max(np.append(np.nonzero(ampa), 0))
         if ampb is not None:
-            mxb = numpy.max(numpy.nonzero(ampb))
+            mxb = np.max(np.append(np.nonzero(ampb), 0))
         self.MaxOrder = max(mxa, mxb)
         if ampa is not None:
             delta = self.MaxOrder - len(ampa)
             if delta > 0:
-                ampa = numpy.pad(ampa, (0, delta))
-            setattr(self, 'AmplitudeA', ampa)
+                ampa = np.pad(ampa, (0, delta))
+            self.AmplitudeA = ampa
         if ampb is not None:
             delta = self.MaxOrder + 1 - len(ampb)
             if delta > 0:
-                ampb = numpy.pad(ampb, (0, delta))
-            setattr(self, 'AmplitudeB', ampb)
+                ampb = np.pad(ampb, (0, delta))
+            self.AmplitudeB = ampb
 
-    def _set_params(self, amplitude, mode, ab, **kwargs):
+    def _set_params(
+        self, amplitude: int or str, mode, a_b: str, **kwargs: dict[str, any]
+    ):
         if amplitude is not None:
-            if numpy.isscalar(amplitude):
-                amp = numpy.zeros(self.MaxOrder)
-                amplitude = numpy.append(amp, amplitude)
+            if np.isscalar(amplitude):
+                amp = np.zeros(self.MaxOrder)
+                amplitude = np.append(amp, amplitude)
             if mode == ACMode.SINE:
-                self._set_sine(ab, **kwargs)
+                self._set_sine(a_b, **kwargs)
             if mode == ACMode.ARBITRARY:
-                self._set_arb(ab, **kwargs)
+                self._set_arb(a_b, **kwargs)
         return amplitude
 
-    def _set_sine(self, ab, **kwargs):
-        frequency = kwargs.pop('Frequency' + ab, None)
-        phase = kwargs.pop('Phase' + ab, 0)
-        assert frequency is not None, \
-            'Please provide a value for Frequency' + ab
-        setattr(self, 'Frequency' + ab, frequency)
-        setattr(self, 'Phase' + ab, phase)
+    def _set_sine(self, a_b: str, **kwargs: dict[str, any]):
+        frequency = kwargs.pop("Frequency" + a_b, None)
+        phase = kwargs.pop("Phase" + a_b, 0)
+        if frequency is None:
+            raise AtError("Please provide a value for Frequency" + a_b)
+        setattr(self, "Frequency" + a_b, frequency)
+        setattr(self, "Phase" + a_b, phase)
 
-    def _set_arb(self, ab, **kwargs):
-        func = kwargs.pop('Func' + ab, None)
+    def _set_arb(self, a_b: str, **kwargs: dict[str, any]):
+        func = kwargs.pop("Func" + a_b, None)
         nsamp = len(func)
-        assert func is not None, \
-            'Please provide a value for Func' + ab
-        setattr(self, 'Func' + ab, func)
-        setattr(self, 'NSamples' + ab, nsamp)
+        if func is None:
+            raise AtError("Please provide a value for Func" + a_b)
+        setattr(self, "Func" + a_b, func)
+        setattr(self, "NSamples" + a_b, nsamp)