Source code for at.lattice.elements.magnet_elements

r""":py:class:`.Element` classes for magnets.

.. _at-field-expansion:

Expansion of the magnetic field in AT
-------------------------------------

The multipole magnets are expressed according to the following expression for the
multipole expansion of the magnetic field:

.. math:: \frac{B_y + iB_x}{B\rho} = \sum_{n=0}^{MaxOrder}(b_n+ ia_n)(x+iy)^n

where :math:`n` is the multipole order (0 for dipole, 1 for quadrupole…).

The :math:`b_n` coefficients describe the "normal" magnetic field (mid-plane symmetry)
and are given in the `PolynomB` sequence.

The :math:`a_n` coefficients describe the "skew" magnetic field and are given in the
`PolynomA` sequence.

.. note::

   This field expansion differs from the one used in MAD or other programs. See `PALS
   <https://pals-project.readthedocs.io/en/latest/element-parameters.html#magneticmultipolep-magnetic-multipole-parameters>`_
   for a definition of the MAD/PALS field expansion. Practically, we have:

   .. math:: \begin{eqnarray} Kn_n &= n!\:b_n \\ Ks_n &= n!\:a_n \end{eqnarray}

"""

from __future__ import annotations

__all__ = [
    "Bend",
    "Corrector",
    "Dipole",
    "Multipole",
    "Octupole",
    "Quadrupole",
    "Sextupole",
    "ThinMultipole",
    "Wiggler",
]

import warnings
from collections.abc import Generator, Callable
from typing import Any
import contextlib
from warnings import warn
from math import tan, atan

import numpy as np

from ..exceptions import AtError, AtWarning
from .conversions import _float, _array
from .abstract_elements import Radiative, _Radiative
from .element_object import Element
from .basic_elements import LongElement

# AtWarning from this module should always be issued (not only on the first occurrence)
warnings.filterwarnings("always", category=AtWarning, module=__name__)


def _warn(doc: str) -> str:
    """Add a notice on AT / PALS field expansion."""
    return "\n".join(
        (
            doc,
            "",
            "The AT field expansion differs from the MAD expansion by a factor n!. "
            "See :ref:`here <at-field-expansion>` for the definition of the AT field "
            "expansion and `PALS <https://pals-project.readthedocs.io/en/latest/element"
            "-parameters.html#magneticmultipolep-magnetic-multipole-parameters>`_ "
            "for the MAD/PALS field expansion.",
        )
    )


def _k_property(attrname: str, order: int, doc: str) -> property:
    """Return a property for a field component."""

    def kget(mult: ThinMultipole) -> float:
        if order > mult.MaxOrder:
            return 0.0
        else:
            return float(getattr(mult, attrname)[order])

    def kset(mult: ThinMultipole, value: float) -> None:
        getattr(mult, attrname)[order] = value

    return property(kget, kset, doc=doc)


def _s_property(attrname: str, doc: str) -> property:
    """Return a property for the strength of the main field component."""

    def kget(mult: ThinMultipole) -> float:
        order = getattr(mult, "DefaultOrder", None)
        if order is None:
            msg = "Undefined order. You need to set a 'DefaultOrder' attribute."
            raise AttributeError(msg)
        elif order > mult.MaxOrder:
            return 0.0
        else:
            return float(getattr(mult, attrname)[order])

    def kset(mult: ThinMultipole, value: float) -> None:
        order = getattr(mult, "DefaultOrder", None)
        if order is None:
            msg = "Undefined order. You need to set a 'DefaultOrder' attribute."
            raise AttributeError(msg)
        getattr(mult, attrname)[order] = value

    return property(kget, kset, doc=doc)


class _LinkedArray(np.ndarray):
    """Linked array.

    A numpy array triggering an action on any modification of its elements.
    """

    setitem: Callable[[int, Any], None]

    def __setitem__(self, key, value):
        # normal action
        super().__setitem__(key, value)
        with contextlib.suppress(AttributeError):
            # external action, if one is defined
            self.setitem(key, value)

    def __repr__(self) -> str:
        # Simulate a standard ndarray
        return repr(self.view(np.ndarray))


[docs] class ThinMultipole(Element): """Thin multipole element.""" # Class attributes _BUILD_ATTRIBUTES = [*Element._BUILD_ATTRIBUTES, "PolynomA", "PolynomB"] _conversions = dict(Element._conversions, K=float, H=float) _stacklevel = 4 # Stacklevel for warnings # Instance attributes PolynomA: np.ndarray #: Integrated skew strengths PolynomB: np.ndarray #: Integrated normal strengths MaxOrder: int #: Highest multipole order def __init__(self, family_name: str, poly_a, poly_b, **kwargs): """ Args: family_name: Name of the element poly_a: Array of integrated skew multipole components poly_b: Array of integrated normal multipole components. Keyword arguments: MaxOrder: Number of desired multipoles. Default: highest index of non-zero polynomial coefficients FieldScaling: Scaling factor applied to the magnetic field (*PolynomA* and *PolynomB*) Default PassMethod: ``ThinMPolePass`` """ def make_same_length(arr1, arr2): lmax = max(len(arr1), len(arr2)) arr1 = np.pad(arr1, (0, lmax - len(arr1)), "constant") arr2 = np.pad(arr2, (0, lmax - len(arr2)), "constant") return arr1, arr2 def seterr(name, kname, kval, aname, aval): mess = ( f"Element {name}: Conflicting element data, {kname!r} ({kval}) " f"in kwargs does not match positional argument {aname!r} ({aval})." ) return AtError(mess) def setwarn(name, aname, kname): mess = ( f"Element {name}: Duplicate element data, both positional " f"argument {aname!r} and {kname!r} in kwargs passed." ) warnings.warn(AtWarning(mess), stacklevel=self._stacklevel) def check_polynom(keyname, arg): argvalue = self._conversions[keyname](arg) argname = f"poly_{keyname[-1].lower()}" if keyname in kwargs: kvalue = self._conversions[keyname](kwargs.pop(keyname)) kvalue, argvalue = make_same_length(kvalue, argvalue) if issubclass(self.__class__, (Dipole, Quadrupole)): if ( keyname == "PolynomB" and argvalue[1] != 0.0 and (kvalue.size < 2 or argvalue[1] != kvalue[1]) ): raise seterr(family_name, "PolynomB", kvalue, "k", argvalue[1]) elif issubclass(self.__class__, Sextupole): if ( keyname == "PolynomB" and argvalue[2] != 0.0 and (kvalue.size < 3 or argvalue[2] != kvalue[2]) ): raise seterr(family_name, "PolynomB", kvalue, "h", argvalue[2]) elif np.any(argvalue) and not np.array_equiv(kvalue, argvalue): raise seterr(family_name, keyname, kvalue, argname, argvalue) else: setwarn(family_name, argname, keyname) return kvalue else: return argvalue def check_strength(keyname, index): if keyname in kwargs: k = self._conversions[keyname](kwargs.pop(keyname)) vname = f"poly_b[{index}]" if len(prmpolb) > index: if k != 0.0 and k != prmpolb[index]: raise seterr(family_name, "K", k, vname, prmpolb[index]) else: setwarn(family_name, "poly_b", "K") else: raise seterr(family_name, "K", k, vname, 0.0) # To avoid potentially unintended behaviour, due to the constructor being # passed multiple definitions of the same thing, we do the following: # - If it is present in kwargs, 'PolynomA' takes priority over 'poly_a'. # - If it is present in kwargs, 'PolynomB' takes priority over 'poly_b' which # in-turn takes priority over 'K' and 'H' if they are present in kwargs. # - Whenever this happens, we either raise an error or give a warning. If the # duplicate definitions contain contradictory non-zero data then we raise an # error, otherwise we give a warning. # Check kwargs and poly_a & poly_b for compatibility and convert to ParamArray prmpola = check_polynom("PolynomA", poly_a) prmpolb = check_polynom("PolynomB", poly_b) check_strength("K", 1) check_strength("H", 2) # Determine the length and order of PolynomA and PolynomB len_a = len(prmpola) len_b = len(prmpolb) deforder = max(getattr(self, "DefaultOrder", 0), len_a - 1, len_b - 1) # Remove MaxOrder maxorder = int(kwargs.pop("MaxOrder", deforder)) kwargs.setdefault("PassMethod", "ThinMPolePass") super().__init__(family_name, **kwargs) # Set MaxOrder while PolynomA and PolynomB are not set yet super().__setattr__("MaxOrder", maxorder) # Adjust polynom lengths and set them len_ab = max(maxorder, deforder) + 1 self.PolynomA = np.pad(prmpola, (0, len_ab - len_a)) self.PolynomB = np.pad(prmpolb, (0, len_ab - len_b)) def __setattr__(self, key, value): """Check the compatibility of MaxOrder, PolynomA and PolynomB.""" def maxord(poly): nonzero = np.flatnonzero(poly != 0.0) return len(poly), nonzero[-1] if len(nonzero) > 0 else -1 def ck(attrname): poly = getattr(self, attrname) return maxord(poly) polys = ("PolynomA", "PolynomB") if key in polys: lmin = self.MaxOrder lenp, ordp = maxord(np.asarray(value)) if not lenp > lmin: msg = f"Length of {key} must be larger than {lmin}" raise ValueError(msg) if ordp > lmin: msg = f"Some values of {key} are truncated by MaxOrder={lmin}" warn(AtWarning(msg), stacklevel=2) elif key == "MaxOrder": intval = int(value) lens, ords = zip(*(ck(k) for k in polys), strict=True) if intval >= (lmin := min(lens)): msg = f"MaxOrder must be smaller than {lmin}" raise ValueError(msg) if intval < max(ords): msg = f"Some values are truncated by MaxOrder={intval}" warn(AtWarning(msg), stacklevel=2) super().__setattr__(key, 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]) # noinspection PyPep8Naming @K.setter def K(self, strength: float): self.PolynomB[1] = strength # 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]) # noinspection PyPep8Naming @H.setter def H(self, strength): self.PolynomB[2] = strength @property def IntegratedPolynomA(self) -> np.ndarray: """Integrated skew strengths.""" return self.PolynomA @IntegratedPolynomA.setter def IntegratedPolynomA(self, value) -> None: self.PolynomA = value @property def IntegratedPolynomB(self) -> np.ndarray: """Integrated normal strengths.""" return self.PolynomB @IntegratedPolynomB.setter def IntegratedPolynomB(self, value) -> None: self.PolynomB = value Kn0L = _k_property( "IntegratedPolynomB", 0, "Integrated normal dipolar strength.\n\n" "*Kn0L* is the opposite of the horizontal momentum kick: " r":math:`\Delta p_x = -\mathrm{Kn0L}`", ) Ks0L = _k_property( "IntegratedPolynomA", 0, "Integrated skew dipolar strength.\n\n" "*Ks0L* is the vertical momentum kick: " r":math:`\Delta p_y = \mathrm{Ks0L}`", ) Kn1L = _k_property( "IntegratedPolynomB", 1, "Integrated normal quadrupolar strength [m^-1].", ) Ks1L = _k_property( "IntegratedPolynomA", 1, "Integrated skew quadrupolar strength [m^-1].", ) Kn2L = _k_property( "IntegratedPolynomB", 2, _warn("Integrated normal sextupolar strength [m^-2]."), ) Ks2L = _k_property( "IntegratedPolynomA", 2, _warn("Integrated skew sextupolar strength [m^-2]."), ) Kn3L = _k_property( "IntegratedPolynomB", 3, _warn("Integrated normal octupolar strength [m^-3]."), ) Ks3L = _k_property( "IntegratedPolynomA", 3, _warn("Integrated skew octupolar strength [m^-3]."), ) IntegratedStrength = _s_property( "IntegratedPolynomB", "Integrated strength of the main field component.", )
[docs] class Multipole(_Radiative, LongElement, ThinMultipole): """Multipole element.""" # Class attributes _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "PolynomA", "PolynomB"] _stacklevel = 6 # Stacklevel for warnings # Instance attributes PolynomA: np.ndarray #: Skew strengths PolynomB: np.ndarray #: Normal strengths MaxOrder: int #: Highest multipole order NumIntSteps: int #: Number of integration slices def __init__(self, family_name: str, length: float, poly_a, poly_b, **kwargs): """ Args: family_name: Name of the element length: Element length [m] poly_a: Array of skew multipole components poly_b: Array of normal multipole components. Keyword arguments: MaxOrder: Number of desired multipoles. Default: highest index of non-zero polynomial coefficients NumIntSteps: Number of integration steps (default: 10) KickAngle: Correction deviation angles (H, V) FieldScaling: Scaling factor applied to the magnetic field (*PolynomA* and *PolynomB*) Default PassMethod: ``StrMPoleSymplectic4Pass`` """ kwargs.setdefault("PassMethod", "StrMPoleSymplectic4Pass") kwargs.setdefault("NumIntSteps", 10) super().__init__(family_name, length, poly_a, poly_b, **kwargs)
[docs] def is_compatible(self, other) -> bool: if super().is_compatible(other) and self.MaxOrder == other.MaxOrder: for i in range(self.MaxOrder + 1): if self.PolynomB[i] != other.PolynomB[i]: return False if self.PolynomA[i] != other.PolynomA[i]: return False return True else: return False
def _setter(self, attrname: str) -> Callable[[int, Any], None]: def setr(key, value): getattr(self, attrname)[key] = np.array(value) / self.Length return setr @property def IntegratedPolynomA(self) -> np.ndarray: pa = self.PolynomA ipa = _LinkedArray(pa.shape, dtype=pa.dtype, buffer=self.Length * pa) ipa.setitem = self._setter("PolynomA") return ipa @IntegratedPolynomA.setter def IntegratedPolynomA(self, value): self.PolynomA = np.array(value) / self.Length @property def IntegratedPolynomB(self) -> np.ndarray: pb = self.PolynomB ipb = _LinkedArray(pb.shape, dtype=pb.dtype, buffer=self.Length * pb) ipb.setitem = self._setter("PolynomB") return ipb @IntegratedPolynomB.setter def IntegratedPolynomB(self, value) -> None: self.PolynomB = np.array(value) / self.Length Strength = _s_property("PolynomB", "Strength of the main field component.") Kn0 = _k_property("PolynomB", 0, "Normal dipolar strength [m^-1].") Ks0 = _k_property("PolynomA", 0, "Skew dipolar strength [m^-1].") Kn1 = _k_property("PolynomB", 1, "Normal quadrupolar strength [m^-2].") Ks1 = _k_property("PolynomA", 1, "Skew quadrupolar strength [m^-2].") Kn2 = _k_property("PolynomB", 2, _warn("Normal sextupolar m^-3].")) Ks2 = _k_property("PolynomA", 2, _warn("Skew sextupolar m^-3].")) Kn3 = _k_property("PolynomB", 3, _warn("Normal octupolar strength [m^-4].")) Ks3 = _k_property("PolynomA", 3, _warn("Skew octupolar strength m^-4]."))
[docs] class Dipole(Radiative, Multipole): """Dipole element.""" _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "BendingAngle", "K"] _conversions = dict( Multipole._conversions, EntranceAngle=float, ExitAngle=float, FringeInt1=float, FringeInt2=float, FringeQuadEntrance=int, FringeQuadExit=int, FringeBendEntrance=int, FringeBendExit=int, ) _file_classname = "Bend" _stacklevel = 7 # Stacklevel for warnings _entrance_fields = [ *Multipole._entrance_fields, "EntranceAngle", "FringeInt1", "FringeBendEntrance", "FringeQuadEntrance", ] _exit_fields = [ *Multipole._exit_fields, "ExitAngle", "FringeInt2", "FringeBendExit", "FringeQuadExit", ] DefaultOrder = 1 # Instance attributes BendingAngle: float EntranceAngle: float ExitAngle: float FringeBendEntrance: int FringeBendExit: int FringeQuadEntrance: int FringeQuadExit: int FullGap: float def __init__( self, family_name: str, length: float, bending_angle: float = 0.0, k: float = 0.0, **kwargs, ): """ Args: family_name: Name of the element length: Element length [m] bending_angle: Bending angle [rd] k: Focusing strength [m^-2]. Keyword arguments: EntranceAngle=0.0: entrance angle ExitAngle=0.0: exit angle PolynomB: normal multipoles PolynomA: skew multipoles MaxOrder=0: Number of desired multipoles NumIntSteps=10: Number of integration steps FullGap: Magnet full gap FringeInt1: Extension of the entrance fringe field FringeInt2: Extension of the exit fringe field FringeBendEntrance: 1: legacy version Brown First Order (default) 2: SOLEIL close to second order of Brown 3: THOMX FringeBendExit: See *FringeBendEntrance* FringeQuadEntrance: 0: no fringe field effect (default) 1: Lee-Whiting's thin lens limit formula 2: elegant-like FringeQuadExit: See *FringeQuadEntrance* fringeIntM0: Integrals for FringeQuad method 2 fringeIntP0: KickAngle: Correction deviation angles (H, V) FieldScaling: Scaling factor applied to the magnetic field Available PassMethods: :ref:`BndMPoleSymplectic4Pass`, :ref:`BendLinearPass`, :ref:`ExactSectorBendPass`, :ref:`ExactRectangularBendPass`, :ref:`ExactRectBendPass`, BndStrMPoleSymplectic4Pass Default PassMethod: :ref:`BndMPoleSymplectic4Pass` """ kwargs.setdefault("BendingAngle", bending_angle) kwargs.setdefault("EntranceAngle", 0.0) kwargs.setdefault("ExitAngle", 0.0) kwargs.setdefault("PassMethod", "BndMPoleSymplectic4Pass") super().__init__(family_name, length, [], [0.0, k], **kwargs)
[docs] def items(self) -> Generator[tuple[str, Any], None, None]: yield from super().items() yield "K", self.K
def _part(self, fr, sumfr): pp = super()._part(fr, sumfr) pp.BendingAngle = fr / sumfr * self.BendingAngle pp.EntranceAngle = 0.0 pp.ExitAngle = 0.0 return pp
[docs] def is_compatible(self, other) -> bool: def invrho(dip: Dipole): return dip.BendingAngle / dip.Length return ( super().is_compatible(other) and self.ExitAngle == -other.EntranceAngle and abs(invrho(self) - invrho(other)) <= 1.0e-6 )
[docs] def merge(self, other) -> None: super().merge(other) # noinspection PyAttributeOutsideInit self.ExitAngle = other.ExitAngle self.BendingAngle += other.BendingAngle
# Bend is a synonym of Dipole. Bend = Dipole
[docs] class Quadrupole(Radiative, Multipole): """Quadrupole element.""" _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "K"] _conversions = dict( Multipole._conversions, FringeQuadEntrance=int, FringeQuadExit=int ) _stacklevel = 7 # Stacklevel for warnings _entrance_fields = [*Multipole._entrance_fields, "FringeQuadEntrance"] _exit_fields = [*Multipole._exit_fields, "FringeQuadExit"] DefaultOrder = 1 # Instance attributes FringeQuadEntrance: int FringeQuadExit: int def __init__(self, family_name: str, length: float, k: float = 0.0, **kwargs): """ Args: family_name: Name of the element length: Element length [m] k: Focusing strength [m^-2]. Keyword Arguments: PolynomB: normal multipoles PolynomA: skew multipoles MaxOrder=1: Number of desired multipoles NumIntSteps=10: Number of integration steps FringeQuadEntrance: 0: no fringe field effect (default) 1: Lee-Whiting's thin lens limit formula 2: elegant-like FringeQuadExit: See ``FringeQuadEntrance`` fringeIntM0: Integrals for FringeQuad method 2 fringeIntP0: 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, k], **kwargs)
[docs] def items(self) -> Generator[tuple[str, Any], None, None]: yield from super().items() yield "K", self.K
[docs] class Sextupole(Multipole): """Sextupole element.""" _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "H"] _stacklevel = 7 # Stacklevel for warnings DefaultOrder = 2 def __init__(self, family_name: str, length: float, h: float = 0.0, **kwargs): """ Args: family_name: Name of the element length: Element length [m] h: Sextupolar strength [m^-3]. Keyword Arguments: PolynomB: normal 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, h], **kwargs)
[docs] class Octupole(Multipole): """Octupole element, with no changes from multipole at present.""" _BUILD_ATTRIBUTES = Multipole._BUILD_ATTRIBUTES DefaultOrder = 3
[docs] class Corrector(LongElement): """Corrector element.""" # Class attributes _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "KickAngle"] # Instance attributes KickAngle: np.ndarray #: (H, V) deviation angles def __init__(self, family_name: str, length: float, kick_angle, **kwargs): """ Args: family_name: Name of the element length: Element length [m] KickAngle: Correction deviation angles (H, V). Keyword Args: FieldScaling: Scaling factor applied to the magnetic field (*KickAngle*) Default PassMethod: ``CorrectorPass`` """ kwargs.setdefault("PassMethod", "CorrectorPass") super().__init__(family_name, length, KickAngle=kick_angle, **kwargs) @property def Kn0L(self) -> float: r"""Opposite of the horizontal momentum kick - :math:`\Delta p_x = -\mathrm{Kn0L}`.""" return -tan(self.KickAngle[0]) @Kn0L.setter def Kn0L(self, value: float) -> None: self.KickAngle[0] = -atan(value) @property def Ks0L(self) -> float: r"""Vertical momentum kick - :math:`\Delta p_y = \mathrm{Ks0L}`.""" return tan(self.KickAngle[1]) @Ks0L.setter def Ks0L(self, value: float) -> None: self.KickAngle[1] = atan(value)
[docs] class Wiggler(Radiative, LongElement): """Wiggler element. See atwiggler.m """ _BUILD_ATTRIBUTES = [*LongElement._BUILD_ATTRIBUTES, "Lw", "Bmax"] _conversions = dict( Element._conversions, Lw=_float, Bmax=_float, Energy=_float, Bx=lambda v: _array(v, (6, -1)), By=lambda v: _array(v, (6, -1)), Nstep=int, Nmeth=int, NHharm=int, NVharm=int, ) # Instance attributes Lw: float Bmax: float Energy: float Bx: np.ndarray By: np.ndarray Nstep: int Nmeth: int NHharm: int NVharm: int # noinspection PyPep8Naming def __init__( self, family_name: str, length: float, wiggle_period: float, b_max: float, energy: float = 0.0, *, Nstep: int = 5, Nmeth: int = 4, By=(1, 1, 0, 1, 1, 0), Bx=(), **kwargs, ): """ Args: length: total length of the wiggler wiggle_period: length must be a multiple of this b_max: peak wiggler field [Tesla] energy: kept for backwards compatibility but always ignored Nstep: number of integration steps. Nmeth: symplectic integration order: 2 or 4 Bx: harmonics for horizontal wiggler: (6, nHharm) array-like object By: harmonics for vertical wiggler (6, nHharm) array-like object. Default PassMethod: ``GWigSymplecticPass`` """ kwargs.setdefault("PassMethod", "GWigSymplecticPass") n_wiggles = length / wiggle_period if abs(round(n_wiggles) - n_wiggles) > 1e-6: msg = ( "Wiggler: length / wiggle_period is not an " f"integer. ({length}/{wiggle_period}={n_wiggles})" ) raise ValueError(msg) if energy: warnings.warn( AtWarning( f"Wiggler: positional argument energy={energy} has been ignored as " "lattice.Energy is always used instead." ), stacklevel=2, ) super().__init__( family_name, length, Lw=wiggle_period, Bmax=b_max, Nstep=Nstep, Nmeth=Nmeth, By=By, Bx=Bx, **kwargs, ) for i, b in enumerate(self.By.T): dk = abs(b[3] ** 2 - b[4] ** 2 - b[2] ** 2) / abs(b[4]) if dk > 1e-6: msg = f"Wiggler(H): kx^2 + kz^2 -ky^2 !=0, i = {i}" raise ValueError(msg) for i, b in enumerate(self.Bx.T): dk = abs(b[2] ** 2 - b[4] ** 2 - b[3] ** 2) / abs(b[4]) if dk > 1e-6: msg = f"Wiggler(V): ky^2 + kz^2 -kx^2 !=0, i = {i}" raise ValueError(msg) self.NHharm = self.By.shape[1] self.NVharm = self.Bx.shape[1]
[docs] def divide(self, frac) -> list[Element]: # A wiggler is indivisible return [self]