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, freeze: bool = True) -> Generator[tuple[str, Any], None, None]:
yield from super().items(freeze=freeze)
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, freeze: bool = True) -> Generator[tuple[str, Any], None, None]:
yield from super().items(freeze=freeze)
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]
