"""
Module to define common elements used in AT.
Each element has a default PassMethod attribute for which it should have the
appropriate attributes. If a different PassMethod is set, it is the caller's
responsibility to ensure that the appropriate attributes are present.
"""
from __future__ import annotations
import abc
import re
from abc import ABC
from collections.abc import Generator, Iterable
from copy import copy, deepcopy
from typing import Any, Optional
import numpy
# noinspection PyProtectedMember
from .variables import _nop
def _array(value, shape=(-1,), dtype=numpy.float64):
# Ensure proper ordering(F) and alignment(A) for "C" access in integrators
return numpy.require(value, dtype=dtype, requirements=['F', 'A']).reshape(
shape, order='F')
def _array66(value):
return _array(value, shape=(6, 6))
def _float(value) -> float:
return float(value)
def _int(value, vmin: Optional[int] = None, vmax: Optional[int] = None) -> int:
intv = int(value)
if vmin is not None and intv < vmin:
raise ValueError(f"Value must be greater of equal to {vmin}")
if vmax is not None and intv > vmax:
raise ValueError(f"Value must be smaller of equal to {vmax}")
return intv
[docs]
class LongtMotion(ABC):
"""Abstract Base class for all Element classes whose instances may modify
the particle momentum
Allows to identify elements potentially inducing longitudinal motion.
Subclasses of :py:class:`LongtMotion` must provide two methods for
enabling longitudinal motion:
* ``_get_longt_motion(self)`` must return the activation state,
* ``set_longt_motion(self, enable, new_pass=None, copy=False, **kwargs)``
must enable or disable longitudinal motion.
"""
@abc.abstractmethod
def _get_longt_motion(self):
return False
# noinspection PyShadowingNames
[docs]
@abc.abstractmethod
def set_longt_motion(self, enable, new_pass=None, copy=False, **kwargs):
"""Enable/Disable longitudinal motion
Parameters:
enable: :py:obj:`True`: for enabling, :py:obj:`False` for
disabling
new_pass: New PassMethod:
* :py:obj:`None`: makes no change,
* ``'auto'``: Uses the default conversion,
* Anything else is used as the new PassMethod.
copy: If True, returns a modified copy of the element,
otherwise modifies the element in-place
"""
# noinspection PyUnresolvedReferences
if new_pass is None or new_pass == self.PassMethod:
return self if copy else None
if copy:
newelem = deepcopy(self)
newelem.PassMethod = new_pass
return newelem
# noinspection PyAttributeOutsideInit
self.PassMethod = new_pass
# noinspection PyUnresolvedReferences
class _DictLongtMotion(LongtMotion):
# noinspection PyShadowingNames
"""Mixin class for elements implementing a 'default_pass' class attribute
:py:class:`DictLongtMotion` provides:
* a :py:meth:`set_longt_motion` method setting the PassMethod according
to the ``default_pass`` dictionary.
* a :py:obj:`.longt_motion` property set to :py:obj:`True` when the
PassMethod is ``default_pass[True]``
The class must have a ``default_pass`` class attribute, a dictionary
such that:
* ``default_pass[False]`` is the PassMethod when radiation is turned
OFF,
* ``default_pass[True]`` is the default PassMethod when radiation is
turned ON.
The :py:class:`DictLongtMotion` class must be set as the first base class.
Example:
>>> class QuantumDiffusion(_DictLongtMotion, Element):
...
... default_pass = {False: 'IdentityPass', True: 'QuantDiffPass'}
Defines a class such that :py:meth:`set_longt_motion` will select
``'IdentityPass'`` or ``'IdentityPass'``.
"""
def _get_longt_motion(self):
return self.PassMethod != self.default_pass[False]
# noinspection PyShadowingNames
def set_longt_motion(self, enable, new_pass=None, **kwargs):
if new_pass == 'auto':
new_pass = self.default_pass[enable]
return super().set_longt_motion(enable, new_pass=new_pass, **kwargs)
# noinspection PyUnresolvedReferences
class _Radiative(LongtMotion):
# noinspection PyShadowingNames
r"""Mixin class for radiating elements
:py:class:`_Radiative` implements the mechanism for converting the pass
methods of radiating elements. It provides:
* a :py:meth:`set_longt_motion` method setting the PassMethod
according to the following rule:
* ``enable == True``: replace "\*Pass" by "\*RadPass"
* ``enable == False``: replace "\*RadPass" by "\*Pass"
* a :py:obj:`.longt_motion` property set to true when the PassMethod
ends with "RadPass"
The :py:class:`_Radiative` class must be set as the first base class.
Example:
>>> class Multipole(_Radiative, LongElement, ThinMultipole):
Defines a class where :py:meth:`set_longt_motion` will convert the
PassMethod according to the \*Pass or \*RadPass suffix.
"""
def _get_longt_motion(self):
return self.PassMethod.endswith(('RadPass', 'QuantPass'))
def _autopass(self, enable):
if enable:
root = self.PassMethod.replace('QuantPass', 'Pass').\
replace('RadPass', 'Pass')
return ''.join((root[:-4], 'RadPass'))
elif self.longt_motion:
root = self.PassMethod.replace('QuantPass', 'Pass').\
replace('RadPass', 'Pass')
return root
else:
return None
# noinspection PyTypeChecker,PyShadowingNames
def set_longt_motion(self, enable, new_pass=None, copy=False, **kwargs):
if new_pass == 'auto':
new_pass = self._autopass(enable)
if new_pass is None or new_pass == self.PassMethod:
return self if copy else None
if enable:
def setpass(el):
el.PassMethod = new_pass
el.Energy = kwargs['energy']
else:
def setpass(el):
el.PassMethod = new_pass
try:
del el.Energy
except AttributeError:
pass
if copy:
newelem = deepcopy(self)
setpass(newelem)
return newelem
setpass(self)
[docs]
class Radiative(_Radiative):
# noinspection PyShadowingNames
r"""Mixin class for default radiating elements (:py:class:`.Dipole`,
:py:class:`.Quadrupole`, :py:class:`.Wiggler`)
:py:class:`Radiative` is a base class for the subset of radiative elements
considered as the ones to be turned on by default: :py:class:`.Dipole`,
:py:class:`.Quadrupole` and :py:class:`.Wiggler`, excluding the higher
order multipoles.
:py:class:`Radiative` inherits from :py:class:`_Radiative` and does not
add any new functionality. Its purpose is to identify the default set of
radiating elements.
Example:
>>> class Dipole(Radiative, Multipole):
Defines a class belonging to the default radiating elements. It
converts the PassMethod according to the "\*Pass" or "\*RadPass"
suffix.
"""
pass
[docs]
class Collective(_DictLongtMotion):
"""Mixin class for elements representing collective effects
Derived classes will automatically set the
:py:attr:`~Element.is_collective` property when the element is active.
The class must have a ``default_pass`` class attribute, a dictionary such
that:
* ``default_pass[False]`` is the PassMethod when collective effects
are turned OFF,
* ``default_pass[True]`` is the default PassMethod when collective effects
are turned ON.
The :py:class:`Collective` class must be set as the first base class.
Example:
>>> class WakeElement(Collective, Element):
...
... default_pass = {False: 'IdentityPass', True: 'WakeFieldPass'}
Defines a class where the :py:attr:`~Element.is_collective` property is
handled
"""
def _get_collective(self):
# noinspection PyUnresolvedReferences
return self.PassMethod != self.default_pass[False]
[docs]
@abc.abstractmethod
def clear_history(self):
pass
[docs]
class Element(object):
"""Base class for AT elements"""
_BUILD_ATTRIBUTES = ['FamName']
_conversions = dict(FamName=str, PassMethod=str, Length=_float,
R1=_array66, R2=_array66,
T1=lambda v: _array(v, (6,)),
T2=lambda v: _array(v, (6,)),
RApertures=lambda v: _array(v, (4,)),
EApertures=lambda v: _array(v, (2,)),
KickAngle=lambda v: _array(v, (2,)),
PolynomB=_array, PolynomA=_array,
BendingAngle=_float,
MaxOrder=_int, NumIntSteps=lambda v: _int(v, vmin=0),
Energy=_float,
)
_entrance_fields = ['T1', 'R1']
_exit_fields = ['T2', 'R2']
_no_swap = _entrance_fields + _exit_fields
def __init__(self, family_name: str, **kwargs):
"""
Parameters:
family_name: Name of the element
All keywords will be set as attributes of the element
"""
self.FamName = family_name
self.Length = kwargs.pop('Length', 0.0)
self.PassMethod = kwargs.pop('PassMethod', 'IdentityPass')
self.update(kwargs)
def __setattr__(self, key, value):
try:
value = self._conversions.get(key, _nop)(value)
except Exception as exc:
exc.args = ('In element {0}, parameter {1}: {2}'.format(
self.FamName, key, exc),)
raise
else:
super(Element, self).__setattr__(key, value)
def __str__(self):
return "\n".join(
[self.__class__.__name__ + ":"]
+ [f"{k:>14}: {v!s}" for k, v in self.items()]
)
def __repr__(self):
clsname, args, kwargs = self.definition
keywords = [f"{arg!r}" for arg in args]
keywords += [f"{k}={v!r}" for k, v in kwargs.items()]
args = re.sub(r"\n\s*", " ", ", ".join(keywords))
return f"{clsname}({args})"
[docs]
def equals(self, other) -> bool:
"""Whether an element is equivalent to another.
This implementation was found to be too slow for the generic
__eq__ method when comparing lattices.
"""
return repr(self) == repr(other)
[docs]
def divide(self, frac) -> list[Element]:
"""split the element in len(frac) pieces whose length
is frac[i]*self.Length
Parameters:
frac: length of each slice expressed as a fraction of the
initial length. ``sum(frac)`` may differ from 1.
Returns:
elem_list: a list of elements equivalent to the original.
Example:
>>> Drift('dr', 0.5).divide([0.2, 0.6, 0.2])
[Drift('dr', 0.1), Drift('dr', 0.3), Drift('dr', 0.1)]
"""
# Bx default, the element is indivisible
return [self]
[docs]
def swap_faces(self, copy=False):
"""Swap the faces of an element, alignment errors are ignored"""
def swapattr(element, attro, attri):
val = getattr(element, attri)
delattr(element, attri)
return attro, val
if copy:
el = self.copy()
else:
el = self
# Remove and swap entrance and exit attributes
fin = dict(swapattr(el, kout, kin)
for kin, kout in zip(el._entrance_fields,
el._exit_fields)
if kin in vars(el)
and kin not in el._no_swap)
fout = dict(swapattr(el, kin, kout)
for kin, kout in zip(el._entrance_fields,
el._exit_fields)
if kout in vars(el)
and kout not in el._no_swap)
# Apply swapped entrance and exit attributes
for key, value in fin.items():
setattr(el, key, value)
for key, value in fout.items():
setattr(el, key, value)
return el if copy else None
[docs]
def update(self, *args, **kwargs):
"""
update(**kwargs)
update(mapping, **kwargs)
update(iterable, **kwargs)
Update the element attributes with the given arguments
"""
attrs = dict(*args, **kwargs)
for key, value in attrs.items():
setattr(self, key, value)
[docs]
def copy(self) -> Element:
"""Return a shallow copy of the element"""
return copy(self)
[docs]
def deepcopy(self) -> Element:
"""Return a deep copy of the element"""
return deepcopy(self)
@property
def definition(self) -> tuple[str, tuple, dict]:
"""tuple (class_name, args, kwargs) defining the element"""
attrs = dict(self.items())
arguments = tuple(attrs.pop(
k, getattr(self, k)) for k in self._BUILD_ATTRIBUTES
)
defelem = self.__class__(*arguments)
keywords = dict(
(k, v)
for k, v in attrs.items()
if not numpy.array_equal(v, getattr(defelem, k, None))
)
return self.__class__.__name__, arguments, keywords
[docs]
def items(self) -> Generator[tuple[str, Any], None, None]:
"""Iterates through the data members"""
v = vars(self).copy()
for k in ["FamName", "Length", "PassMethod"]:
yield k, v.pop(k)
for k, v in sorted(v.items()):
yield k, v
[docs]
def is_compatible(self, other: Element) -> bool:
"""Checks if another :py:class:`Element` can be merged"""
return False
[docs]
def merge(self, other) -> None:
"""Merge another element"""
if not self.is_compatible(other):
badname = getattr(other, 'FamName', type(other))
raise TypeError("Cannot merge {0} and {1}".format(self.FamName, badname))
# noinspection PyMethodMayBeStatic
def _get_longt_motion(self):
return False
# noinspection PyMethodMayBeStatic
def _get_collective(self):
return False
@property
def longt_motion(self) -> bool:
""":py:obj:`True` if longitudinal motion is affected by the element"""
return self._get_longt_motion()
@property
def is_collective(self) -> bool:
""":py:obj:`True` if the element involves collective effects"""
return self._get_collective()
[docs]
class LongElement(Element):
"""Base class for long elements"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['Length']
def __init__(self, family_name: str, length: float, *args, **kwargs):
"""
Args:
family_name: Name of the element
length: Element length [m]
Other arguments and keywords are given to the base class
"""
kwargs.setdefault('Length', length)
# Ancestor may be either Element of ThinMultipole
# noinspection PyArgumentList
super(LongElement, self).__init__(family_name, *args, **kwargs)
def _part(self, fr, sumfr):
pp = self.copy()
pp.Length = fr * self.Length
if hasattr(self, 'KickAngle'):
pp.KickAngle = fr / sumfr * self.KickAngle
return pp
[docs]
def divide(self, frac) -> list[Element]:
def popattr(element, attr):
val = getattr(element, attr)
delattr(element, attr)
return attr, val
frac = numpy.asarray(frac, dtype=float)
el = self.copy()
# Remove entrance and exit attributes
fin = dict(popattr(el, key) for key in vars(self) if
key in self._entrance_fields)
fout = dict(popattr(el, key) for key in vars(self) if key in self._exit_fields)
# Split element
element_list = [el._part(f, numpy.sum(frac)) for f in frac]
# Restore entrance and exit attributes
for key, value in fin.items():
setattr(element_list[0], key, value)
for key, value in fout.items():
setattr(element_list[-1], key, value)
return element_list
[docs]
def is_compatible(self, other) -> bool:
def compatible_field(fieldname):
f1 = getattr(self, fieldname, None)
f2 = getattr(other, fieldname, None)
if f1 is None and f2 is None: # no such field
return True
elif f1 is None or f2 is None: # only one
return False
else: # both
return numpy.all(f1 == f2)
if not (type(other) is type(self) and self.PassMethod == other.PassMethod):
return False
for fname in ("RApertures", "EApertures"):
if not compatible_field(fname):
return False
return True
[docs]
def merge(self, other) -> None:
super().merge(other)
self.Length += other.Length
[docs]
class Marker(Element):
"""Marker element"""
[docs]
class Monitor(Element):
"""Monitor element"""
[docs]
class BeamMoments(Element):
"""Element to compute bunches mean and std"""
def __init__(self, family_name: str, **kwargs):
"""
Args:
family_name: Name of the element
Default PassMethod: ``BeamMomentsPass``
"""
kwargs.setdefault('PassMethod', 'BeamMomentsPass')
self._stds = numpy.zeros((6, 1, 1), order='F')
self._means = numpy.zeros((6, 1, 1), order='F')
super(BeamMoments, self).__init__(family_name, **kwargs)
[docs]
def set_buffers(self, nturns, nbunch):
self._stds = numpy.zeros((6, nbunch, nturns), order='F')
self._means = numpy.zeros((6, nbunch, nturns), order='F')
@property
def stds(self):
"""Beam 6d standard deviation"""
return self._stds
@property
def means(self):
"""Beam 6d center of mass"""
return self._means
[docs]
class SliceMoments(Element):
"""Element to compute slices mean and std"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['nslice']
_conversions = dict(Element._conversions, nslice=int)
def __init__(self, family_name: str, nslice: int, **kwargs):
"""
Args:
family_name: Name of the element
nslice: Number of slices
Keyword arguments:
startturn: Start turn of the acquisition (Default 0)
endturn: End turn of the acquisition (Default 1)
Default PassMethod: ``SliceMomentsPass``
"""
kwargs.setdefault('PassMethod', 'SliceMomentsPass')
self._startturn = kwargs.pop('startturn', 0)
self._endturn = kwargs.pop('endturn', 1)
super(SliceMoments, self).__init__(family_name, nslice=nslice, **kwargs)
self._nbunch = 1
self.startturn = self._startturn
self.endturn = self._endturn
self._dturns = self.endturn - self.startturn
self._stds = numpy.zeros((3, nslice, self._dturns), order='F')
self._means = numpy.zeros((3, nslice, self._dturns), order='F')
self._spos = numpy.zeros((nslice, self._dturns), order='F')
self._weights = numpy.zeros((nslice, self._dturns), order='F')
self.set_buffers(self._endturn, 1)
[docs]
def set_buffers(self, nturns, nbunch):
self.endturn = min(self.endturn, nturns)
self._dturns = self.endturn - self.startturn
self._nbunch = nbunch
self._stds = numpy.zeros((3, nbunch*self.nslice, self._dturns), order="F")
self._means = numpy.zeros((3, nbunch*self.nslice, self._dturns), order="F")
self._spos = numpy.zeros((nbunch*self.nslice, self._dturns), order="F")
self._weights = numpy.zeros((nbunch*self.nslice, self._dturns), order="F")
@property
def stds(self):
"""Slices x,y,dp standard deviation"""
return self._stds.reshape((3, self._nbunch, self.nslice, self._dturns))
@property
def means(self):
"""Slices x,y,dp center of mass"""
return self._means.reshape((3, self._nbunch, self.nslice, self._dturns))
@property
def spos(self):
"""Slices s position"""
return self._spos.reshape((self._nbunch, self.nslice, self._dturns))
@property
def weights(self):
"""Slices weights in mA if beam current >0,
otherwise fraction of total number of
particles in the bunch
"""
return self._weights.reshape((self._nbunch, self.nslice, self._dturns))
@property
def startturn(self):
"""Start turn of the acquisition"""
return self._startturn
@startturn.setter
def startturn(self, value):
if value < 0:
raise ValueError('startturn must be greater or equal to 0')
if value >= self._endturn:
raise ValueError('startturn must be smaller than endturn')
self._startturn = value
@property
def endturn(self):
"""End turn of the acquisition"""
return self._endturn
@endturn.setter
def endturn(self, value):
if value <= 0:
raise ValueError('endturn must be greater than 0')
if value <= self._startturn:
raise ValueError('endturn must be greater than startturn')
self._endturn = value
[docs]
class Aperture(Element):
"""Aperture element"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['Limits']
_conversions = dict(Element._conversions, Limits=lambda v: _array(v, (4,)))
def __init__(self, family_name, limits, **kwargs):
"""
Args:
family_name: Name of the element
limits: (4,) array of physical aperture:
[xmin, xmax, zmin, zmax] [m]
Default PassMethod: ``AperturePass``
"""
kwargs.setdefault('PassMethod', 'AperturePass')
super(Aperture, self).__init__(family_name, Limits=limits, **kwargs)
[docs]
class Drift(LongElement):
"""Drift space element"""
def __init__(self, family_name: str, length: float, **kwargs):
"""
Args:
family_name: Name of the element
length: Element length [m]
Default PassMethod: ``DriftPass``
"""
kwargs.setdefault('PassMethod', 'DriftPass')
super(Drift, self).__init__(family_name, length, **kwargs)
[docs]
def insert(self,
insert_list: Iterable[tuple[float, Optional[Element]]]) \
-> list[Element]:
"""insert elements inside a drift
Arguments:
insert_list: iterable, each item of insert_list is itself an
iterable with 2 objects:
1. the location where the center of the element
will be inserted, given as a fraction of the Drift length.
2. an element to be inserted at that location. If :py:obj:`None`,
the drift will be divided but no element will be inserted.
Returns:
elem_list: a list of elements.
Drifts with negative lengths may be generated if necessary.
Examples:
>>> Drift('dr', 2.0).insert(((0.25, None), (0.75, None)))
[Drift('dr', 0.5), Drift('dr', 1.0), Drift('dr', 0.5)]
>>> Drift('dr', 2.0).insert(((0.0, Marker('m1')),
... (0.5, Marker('m2'))))
[Marker('m1'), Drift('dr', 1.0), Marker('m2'), Drift('dr', 1.0)]
>>> Drift('dr', 2.0).insert(((0.5, Quadrupole('qp', 0.4, 0.0)),))
[Drift('dr', 0.8), Quadrupole('qp', 0.4), Drift('dr', 0.8)]
"""
frac, elements = zip(*insert_list)
lg = [0.0 if el is None else el.Length for el in elements]
fr = numpy.asarray(frac, dtype=float)
lg = 0.5 * numpy.asarray(lg, dtype=float) / self.Length
drfrac = numpy.hstack((fr - lg, 1.0)) - numpy.hstack((0.0, fr + lg))
long_elems = (drfrac != 0.0)
drifts = numpy.ndarray((len(drfrac),), dtype='O')
drifts[long_elems] = self.divide(drfrac[long_elems])
nline = len(drifts) + len(elements)
line = [None] * nline # type: list[Optional[Element]]
line[::2] = drifts
line[1::2] = elements
return [el for el in line if el is not None]
[docs]
class Collimator(Drift):
"""Collimator element"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['RApertures']
def __init__(self, family_name: str, length: float, limits, **kwargs):
"""
Args:
family_name: Name of the element
length: Element length [m]
limits: (4,) array of physical aperture:
[xmin, xmax, zmin, zmax] [m]
Default PassMethod: ``DriftPass``
"""
super(Collimator, self).__init__(family_name, length,
RApertures=limits, **kwargs)
[docs]
class ThinMultipole(Element):
"""Thin multipole element"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ["PolynomA", "PolynomB"]
def __init__(self, family_name: str, poly_a, poly_b, **kwargs):
"""
Args:
family_name: Name of the element
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
FieldScaling: Scaling factor applied to the magnetic field
(*PolynomA* and *PolynomB*)
Default PassMethod: ``ThinMPolePass``
"""
def getpol(poly):
nonzero = numpy.flatnonzero(poly != 0.0)
return poly, len(poly), nonzero[-1] if len(nonzero) > 0 else -1
def lengthen(poly, dl):
if dl > 0:
return numpy.concatenate((poly, numpy.zeros(dl)))
else:
return poly
# PolynomA and PolynomB and convert to ParamArray
prmpola = self._conversions["PolynomA"](kwargs.pop("PolynomA", poly_a))
prmpolb = self._conversions["PolynomB"](kwargs.pop("PolynomB", poly_b))
poly_a, len_a, ord_a = getpol(prmpola)
poly_b, len_b, ord_b = getpol(prmpolb)
deforder = max(getattr(self, 'DefaultOrder', 0), ord_a, ord_b)
# Remove MaxOrder
maxorder = kwargs.pop('MaxOrder', deforder)
kwargs.setdefault('PassMethod', 'ThinMPolePass')
super(ThinMultipole, self).__init__(family_name, **kwargs)
# Set MaxOrder while PolynomA and PolynomB are not set yet
super(ThinMultipole, self).__setattr__('MaxOrder', maxorder)
# Adjust polynom lengths and set them
len_ab = max(self.MaxOrder + 1, len_a, len_b)
self.PolynomA = lengthen(prmpola, len_ab - len_a)
self.PolynomB = lengthen(prmpolb, len_ab - len_b)
def __setattr__(self, key, value):
"""Check the compatibility of MaxOrder, PolynomA and PolynomB"""
polys = ('PolynomA', 'PolynomB')
if key in polys:
lmin = self.MaxOrder
if not len(value) > lmin:
raise ValueError(
'Length of {0} must be larger than {1}'.format(key, lmin))
elif key == 'MaxOrder':
intval = int(value)
lmax = min(len(getattr(self, k)) for k in polys)
if not intval < lmax:
raise ValueError("MaxOrder must be smaller than {0}".format(lmax))
super(ThinMultipole, self).__setattr__(key, value)
[docs]
class Multipole(_Radiative, LongElement, ThinMultipole):
"""Multipole element"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ["PolynomA", "PolynomB"]
_conversions = dict(ThinMultipole._conversions, K=float, H=float)
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(Multipole, self).__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
# noinspection PyPep8Naming
@property
def K(self) -> float:
"""Focusing strength [mˆ-2]"""
arr = self.PolynomB
return 0.0 if len(arr) < 2 else 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 arr[2]
# noinspection PyPep8Naming
@H.setter
def H(self, strength):
self.PolynomB[2] = strength
[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)
_entrance_fields = Multipole._entrance_fields + ['EntranceAngle',
'FringeInt1',
'FringeBendEntrance',
'FringeQuadEntrance']
_exit_fields = Multipole._exit_fields + ['ExitAngle',
'FringeInt2',
'FringeBendExit',
'FringeQuadExit']
DefaultOrder = 0
def __init__(self, family_name: str, length: float,
bending_angle: Optional[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: straight multipoles
PolynomA: skew multipoles
MaxOrder=0: Number of desired multipoles
NumIntSt=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(Dipole, self).__init__(family_name, length, [], [0.0, k], **kwargs)
[docs]
def items(self) -> Generator[tuple[str, Any], None, None]:
yield from super().items()
yield "K", vars(self)["PolynomB"][1]
def _part(self, fr, sumfr):
pp = super(Dipole, self)._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.e-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)
_entrance_fields = Multipole._entrance_fields + ['FringeQuadEntrance']
_exit_fields = Multipole._exit_fields + ['FringeQuadExit']
DefaultOrder = 1
def __init__(self, family_name: str, length: float,
k: Optional[float] = 0.0, **kwargs):
"""Quadrupole(FamName, Length, Strength=0, **keywords)
Args:
family_name: Name of the element
length: Element length [m]
k: Focusing strength [mˆ-2]
Keyword Arguments:
PolynomB: straight 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(Quadrupole, self).__init__(family_name, length, [], [0.0, k], **kwargs)
[docs]
def items(self) -> Generator[tuple[str, Any], None, None]:
yield from super().items()
yield "K", vars(self)["PolynomB"][1]
[docs]
class Sextupole(Multipole):
"""Sextupole element"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['H']
DefaultOrder = 2
def __init__(self, family_name: str, length: float,
h: Optional[float] = 0.0, **kwargs):
"""
Args:
family_name: Name of the element
length: Element length [m]
h: strength [mˆ-3]
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(Sextupole, self).__init__(family_name, length, [], [0.0, 0.0, h],
**kwargs)
[docs]
def items(self) -> Generator[tuple[str, Any], None, None]:
yield from super().items()
yield "H", vars(self)["PolynomB"][2]
[docs]
class Octupole(Multipole):
"""Octupole element, with no changes from multipole at present"""
_BUILD_ATTRIBUTES = Multipole._BUILD_ATTRIBUTES
DefaultOrder = 3
[docs]
class RFCavity(LongtMotion, LongElement):
"""RF cavity element"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['Voltage',
'Frequency',
'HarmNumber',
'Energy']
default_pass = {False: 'DriftPass', True: 'RFCavityPass'}
_conversions = dict(LongElement._conversions,
Voltage=float, Frequency=float,
HarmNumber=int, TimeLag=float)
def __init__(self, family_name: str, length: float, voltage: float,
frequency: float, harmonic_number: int, energy: float,
**kwargs):
"""
Args:
family_name: Name of the element
length: Element length [m]
voltage: RF voltage [V]
frequency: RF frequency [Hz]
harmonic_number:
energy: ring energy [eV]
Keyword Arguments:
TimeLag=0: Cavity time lag
Default PassMethod: ``RFCavityPass``
"""
kwargs.setdefault('TimeLag', 0.0)
kwargs.setdefault('PassMethod', self.default_pass[True])
super(RFCavity, self).__init__(family_name, length,
Voltage=voltage,
Frequency=frequency,
HarmNumber=harmonic_number,
Energy=energy, **kwargs)
def _part(self, fr, sumfr):
pp = super(RFCavity, self)._part(fr, sumfr)
pp.Voltage = fr * self.Voltage
return pp
[docs]
def is_compatible(self, other) -> bool:
return (super().is_compatible(other)
and self.Frequency == other.Frequency
and self.TimeLag == other.TimeLag)
[docs]
def merge(self, other) -> None:
super().merge(other)
self.Voltage += other.Voltage
def _get_longt_motion(self):
return self.PassMethod.endswith('CavityPass')
# noinspection PyShadowingNames
[docs]
def set_longt_motion(self, enable, new_pass=None, **kwargs):
if new_pass == 'auto':
new_pass = (self.default_pass[True] if enable else
('IdentityPass' if self.Length == 0 else 'DriftPass'))
return super().set_longt_motion(enable, new_pass=new_pass, **kwargs)
[docs]
class M66(Element):
"""Linear (6, 6) transfer matrix"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ["M66"]
_conversions = dict(Element._conversions, M66=_array66)
def __init__(self, family_name: str, m66=None, **kwargs):
"""
Args:
family_name: Name of the element
m66: Transfer matrix. Default: Identity matrix
Default PassMethod: ``Matrix66Pass``
"""
if m66 is None:
m66 = numpy.identity(6)
kwargs.setdefault('PassMethod', 'Matrix66Pass')
kwargs.setdefault("M66", m66)
super(M66, self).__init__(family_name, **kwargs)
[docs]
class SimpleQuantDiff(_DictLongtMotion, Element):
"""
Linear tracking element for a simplified quantum diffusion,
radiation damping and energy loss.
Note: The damping times are needed to compute the correct
kick for the emittance. Radiation damping is NOT applied.
"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES
default_pass = {False: 'IdentityPass', True: 'SimpleQuantDiffPass'}
def __init__(self, family_name: str, betax: float = 1.0,
betay: float = 1.0, emitx: float = 0.0,
emity: float = 0.0, espread: float = 0.0,
taux: float = 0.0, tauy: float = 0.0,
tauz: float = 0.0,
**kwargs):
"""
Args:
family_name: Name of the element
Optional Args:
betax: Horizontal beta function at element [m]
betay: Vertical beta function at element [m]
emitx: Horizontal equilibrium emittance [m.rad]
emity: Vertical equilibrium emittance [m.rad]
espread: Equilibrium energy spread
taux: Horizontal damping time [turns]
tauy: Vertical damping time [turns]
tauz: Longitudinal damping time [turns]
Default PassMethod: ``SimpleQuantDiffPass``
"""
kwargs.setdefault("PassMethod", self.default_pass[True])
assert taux >= 0.0, 'taux must be greater than or equal to 0'
self.taux = taux
assert tauy >= 0.0, 'tauy must be greater than or equal to 0'
self.tauy = tauy
assert tauz >= 0.0, 'tauz must be greater than or equal to 0'
self.tauz = tauz
assert emitx >= 0.0, 'emitx must be greater than or equal to 0'
self.emitx = emitx
if emitx > 0.0:
assert taux > 0.0, 'if emitx is given, taux must be non zero'
assert emity >= 0.0, 'emity must be greater than or equal to 0'
self.emity = emity
if emity > 0.0:
assert tauy > 0.0, 'if emity is given, tauy must be non zero'
assert espread >= 0.0, 'espread must be greater than or equal to 0'
self.espread = espread
if espread > 0.0:
assert tauz > 0.0, 'if espread is given, tauz must be non zero'
self.betax = betax
self.betay = betay
super(SimpleQuantDiff, self).__init__(family_name, **kwargs)
[docs]
class SimpleRadiation(_DictLongtMotion, Radiative, Element):
"""Simple radiation damping and energy loss"""
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES
_conversions = dict(Element._conversions, U0=float,
damp_mat_diag=lambda v: _array(v, shape=(6,)))
default_pass = {False: 'IdentityPass', True: 'SimpleRadiationPass'}
def __init__(self, family_name: str,
taux: float = 0.0, tauy: float = 0.0,
tauz: float = 0.0, U0: float = 0.0,
**kwargs):
"""
Args:
family_name: Name of the element
Optional Args:
taux: Horizontal damping time [turns]
tauy: Vertical damping time [turns]
tauz: Longitudinal damping time [turns]
U0: Energy loss per turn [eV]
Default PassMethod: ``SimpleRadiationPass``
"""
assert taux >= 0.0, 'taux must be greater than or equal to 0'
if taux == 0.0:
dampx = 1
else:
dampx = numpy.exp(-1/taux)
assert tauy >= 0.0, 'tauy must be greater than or equal to 0'
if tauy == 0.0:
dampy = 1
else:
dampy = numpy.exp(-1/tauy)
assert tauz >= 0.0, 'tauz must be greater than or equal to 0'
if tauz == 0.0:
dampz = 1
else:
dampz = numpy.exp(-1/tauz)
kwargs.setdefault('PassMethod', self.default_pass[True])
kwargs.setdefault("U0", U0)
kwargs.setdefault("damp_mat_diag",
numpy.array([dampx, dampx, dampy, dampy, dampz, dampz]))
super(SimpleRadiation, self).__init__(family_name, **kwargs)
[docs]
class Corrector(LongElement):
"""Corrector element"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['KickAngle']
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(Corrector, self).__init__(family_name, length,
KickAngle=kick_angle, **kwargs)
[docs]
class Wiggler(Radiative, LongElement):
"""Wiggler element
See atwiggler.m
"""
_BUILD_ATTRIBUTES = LongElement._BUILD_ATTRIBUTES + ['Lw', 'Bmax',
'Energy']
_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)
# noinspection PyPep8Naming
def __init__(self, family_name: str, length: float, wiggle_period: float,
b_max: float, energy: float, Nstep: Optional[int] = 5,
Nmeth: Optional[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: beam energy [eV]
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:
raise ValueError("Wiggler: length / wiggle_period is not an "
"integer. ({0}/{1}={2})".format(length,
wiggle_period,
n_wiggles))
super(Wiggler, self).__init__(family_name, length, Lw=wiggle_period,
Bmax=b_max, Nstep=Nstep, Nmeth=Nmeth,
By=By, Bx=Bx, Energy=energy, **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:
raise ValueError("Wiggler(H): kx^2 + kz^2 -ky^2 !=0, i = {0}".format(i))
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:
raise ValueError("Wiggler(V): ky^2 + kz^2 -kx^2 !=0, i = {0}".format(i))
self.NHharm = self.By.shape[1]
self.NVharm = self.Bx.shape[1]
[docs]
class QuantumDiffusion(_DictLongtMotion, Element):
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['Lmatp']
default_pass = {False: 'IdentityPass', True: 'QuantDiffPass'}
_conversions = dict(Element._conversions, Lmatp=_array66)
def __init__(self, family_name: str, lmatp: numpy.ndarray, **kwargs):
"""Quantum diffusion element
Args:
family_name: Name of the element
lmatp : Diffusion matrix for generation (see
:py:func:`.gen_quantdiff_elem`)
Default PassMethod: ``QuantDiffPass``
"""
kwargs.setdefault('PassMethod', self.default_pass[True])
super().__init__(family_name, Lmatp=lmatp, **kwargs)
[docs]
class EnergyLoss(_DictLongtMotion, Element):
_BUILD_ATTRIBUTES = Element._BUILD_ATTRIBUTES + ['EnergyLoss']
_conversions = dict(Element._conversions, EnergyLoss=float)
default_pass = {False: 'IdentityPass', True: 'EnergyLossRadPass'}
def __init__(self, family_name: str, energy_loss: float, **kwargs):
"""Energy loss element
Args:
family_name: Name of the element
energy_loss: Energy loss [eV]
"""
kwargs.setdefault('PassMethod', self.default_pass[False])
super().__init__(family_name, EnergyLoss=energy_loss, **kwargs)
Radiative.register(EnergyLoss)
[docs]
def build_class_map(): # Missing class aliases (Bend)
global CLASS_MAP
def subclasses_recursive(cl):
direct = cl.__subclasses__()
indirect = []
for subclass in direct:
indirect.extend(subclasses_recursive(subclass))
return frozenset([cl] + direct + indirect)
cls_list = subclasses_recursive(Element)
CLASS_MAP = {cls.__name__: cls for cls in cls_list}
[docs]
def get_class_map():
return CLASS_MAP
# build_class_map()
CLASS_MAP = dict((k, v) for k, v in locals().items()
if isinstance(v, type) and issubclass(v, Element))
