r"""
Helper functions for working with AT lattices.
A lattice as understood by pyAT is a sequence of elements. These functions
are useful for working with these sequences.
.. _refpts:
**Selecting elements in a lattice:**
The *refpts* argument allows functions to select locations in the lattice. The
location is defined as the entrance of the selected element. *refpts* may be:
#. an integer in the range *[-len(lattice), len(lattice)-1]*
selecting an element according to the python indexing rules.
As a special case, *len(lattice)* is allowed and refers to the end
of the last element,
#. an ordered list of such integers without duplicates,
#. a numpy array of booleans of maximum length *len(lattice)+1*,
where selected elements are :py:obj:`True`,
#. :py:obj:`None`, meaning an empty selection,
#. :py:obj:`.All`, meaning "all possible reference points": the entrance of all
elements plus the end of the last element,
#. :py:obj:`.End`, selecting the end of the last element,
#. an element type, selecting all the elements of that type in
the lattice, e.g. :pycode:`at.Sextupole`,
#. a string, selecting all the elements whose `FamName` attribute matches it.
Unix shell-style wildcards are accepted, e.g. `"Q[FD]*"`,
#. a callable :pycode:`filtfunc` such that :pycode:`filtfunc(elem)`
is :py:obj:`True` for selected elements.
"""
import numpy
import functools
import re
from typing import Callable, Optional, Sequence, Iterator
from typing import Union, Tuple, List, Type
from enum import Enum
from itertools import compress
from fnmatch import fnmatch
from .elements import Element, Dipole
ElementFilter = Callable[[Element], bool]
BoolRefpts = numpy.ndarray
Uint32Refpts = numpy.ndarray
__all__ = ['All', 'End', 'AtError', 'AtWarning', 'axis_descr',
'check_radiation', 'check_6d',
'set_radiation', 'set_6d',
'make_copy', 'uint32_refpts', 'bool_refpts',
'get_uint32_index', 'get_bool_index',
'checkattr', 'checktype', 'checkname',
'get_elements', 'get_s_pos',
'refpts_count', 'refpts_iterator',
'set_shift', 'set_tilt', 'set_rotation',
'tilt_elem', 'shift_elem', 'rotate_elem',
'get_value_refpts', 'set_value_refpts', 'Refpts',
'get_geometry']
_axis_def = dict(
x=dict(index=0, label="x", unit=" [m]"),
xp=dict(index=1, label="x'", unit=" [rad]"),
y=dict(index=2, label="y", unit=" [m]"),
yp=dict(index=3, label="y'", unit=" [rad]"),
dp=dict(index=4, label=r"$\delta$", unit=""),
ct=dict(index=5, label=r"$\beta c \tau$", unit=" [m]"),
)
for vvv in [vv for vv in _axis_def.values()]:
_axis_def[vvv['index']] = vvv
[docs]class AtError(Exception):
pass
[docs]class AtWarning(UserWarning):
pass
_typ1 = "None, All, End, int, bool"
_typ2 = "None, All, End, int, bool, str, Type[Element], ElementFilter"
class RefptsCode(Enum):
All = 'All'
End = 'End'
RefIndex = Union[None, int, Sequence[int], bool, Sequence[bool], RefptsCode]
Refpts = Union[Type[Element], Element, ElementFilter, str, RefIndex]
#: :py:obj:`All` is a special value to be used as *refpts*. It means
#: "all possible reference points": the entrance of all elements plus the end
#: of the last element.
All = RefptsCode.All
#: :py:obj:`End` is a special value to be used as *refpts*. It refers to the
#: end of the last element.
End = RefptsCode.End
def _type_error(refpts, types):
if isinstance(refpts, numpy.ndarray):
tp = refpts.dtype.type
else:
tp = type(refpts)
return TypeError(
"Invalid refpts type {0}. Allowed types: {1}".format(tp, types))
# noinspection PyIncorrectDocstring
[docs]def axis_descr(*args, key=None) -> Tuple:
r"""axis_descr(axis [ ,axis], key=None)
Return a tuple containing for each input argument the requested information
Parameters:
axis (Union[int, str]): either an index in 0:6 or a string in
['x', 'xp', 'y', 'yp', 'dp', 'ct']
key: key in the coordinate description
dictionary, selecting the desired information. One of :
'index'
index in the standard AT coordinate vector
'label'
label for plot annotation
'unit'
coordinate unit
:py:obj:`None`
entire description dictionary
Returns:
descr (Tuple): requested information for each input argument.
Examples:
>>> axis_descr('x','dp', key='index')
(0, 4)
returns the indices in the standard coordinate vector
>>> dplabel, = axis_descr('dp', key='label')
>>> print(dplabel)
$\delta$
returns the coordinate label for plot annotation
>>> axis_descr('x','dp')
({'index': 0, 'label': 'x', 'unit': ' [m]'},
{'index': 4, 'label': '$\\delta$', 'unit': ''})
returns the entire description directories
"""
if key is None:
return tuple(_axis_def[k] for k in args)
else:
return tuple(_axis_def[k][key] for k in args)
def check_radiation(rad: bool) -> Callable:
r"""Deprecated decorator for optics functions (see :py:func:`check_6d`).
:meta private:
"""
return check_6d(rad)
def set_radiation(rad: bool) -> Callable:
r"""Deprecated decorator for optics functions (see :py:func:`set_6d`).
:meta private:
"""
return set_6d(rad)
[docs]def check_6d(is_6d: bool) -> Callable:
r"""Decorator for optics functions
Wraps a function like :pycode:`func(ring, *args, **kwargs)` where
:pycode:`ring` is a :py:class:`.Lattice` object. Raise :py:class:`.AtError`
if :pycode:`ring.is_6d` is not the desired *is_6d* state. No test is
performed if :pycode:`ring` is not a :py:class:`.Lattice`.
Parameters:
is_6d: Desired 6D state
Raises:
AtError: if :pycode:`ring.is_6d` is not *is_6d*
Example:
>>> @check_6d(False)
... def find_orbit4(ring, dct=None, guess=None, **kwargs):
...
Raises :py:class:`.AtError` if :pycode:`ring.is_6d` is :py:obj:`True`
See Also:
:py:func:`set_6d`
"""
def radiation_decorator(func):
@functools.wraps(func)
def wrapper(ring, *args, **kwargs):
ringrad = getattr(ring, 'is_6d', is_6d)
if ringrad != is_6d:
raise AtError('{0} needs "ring.is_6d" {1}'.format(
func.__name__, is_6d))
return func(ring, *args, **kwargs)
return wrapper
return radiation_decorator
[docs]def set_6d(is_6d: bool) -> Callable:
# noinspection PyUnresolvedReferences
r"""Decorator for optics functions
Wraps a function like :pycode:`func(ring, *args, **kwargs)` where
:pycode:`ring` is a :py:class:`.Lattice` object. If :pycode:`ring.is_6d`
is not the desired *is_6d* state, :pycode:`func()` will be called with a
modified copy of :pycode:`ring` satisfying the *is_6d* state.
Parameters:
is_6d: Desired 6D state
Example:
>>> @set_6d(True)
... def compute(ring)
...
is roughly equivalent to:
>>> if ring.is_6d:
... compute(ring)
... else:
... compute(ring.enable_6d(copy=True))
See Also:
:py:func:`check_6d`, :py:meth:`.Lattice.enable_6d`,
:py:meth:`.Lattice.disable_6d`
"""
if is_6d:
def setrad_decorator(func):
@functools.wraps(func)
def wrapper(ring, *args, **kwargs):
rg = ring if ring.is_6d else ring.enable_6d(copy=True)
return func(rg, *args, **kwargs)
return wrapper
else:
def setrad_decorator(func):
@functools.wraps(func)
def wrapper(ring, *args, **kwargs):
rg = ring.disable_6d(copy=True) if ring.is_6d else ring
return func(rg, *args, **kwargs)
return wrapper
return setrad_decorator
[docs]def make_copy(copy: bool) -> Callable:
r"""Decorator for optics functions
Wraps a function like :pycode:`func(ring, refpts, *args, **kwargs)` where
:pycode:`ring` is a :py:class:`.Lattice` object.
If *copy* is :py:obj:`False`, the function is not modified,
If *copy* is :py:obj:`True`, a shallow copy of :pycode:`ring` is done,
then the elements selected by *refpts* are deep-copied, then :pycode:`func`
is applied to the copy, and the new :pycode:`ring` is returned.
Parameters:
copy: If :py:obj:`True`, apply the decorated function to a copy of
:pycode:`ring`
"""
if copy:
def copy_decorator(func):
@functools.wraps(func)
def wrapper(ring, refpts, *args, **kwargs):
ring = ring.replace(refpts)
func(ring, refpts, *args, **kwargs)
return ring
return wrapper
else:
def copy_decorator(func):
return func
return copy_decorator
[docs]def uint32_refpts(refpts: RefIndex, n_elements: int,
endpoint: bool = True,
types: str = _typ1) -> Uint32Refpts:
r"""Return a :py:obj:`~numpy.uint32` array of element indices selecting
ring elements.
Parameters:
refpts: Element selector. *refpts* may be:
#. an integer or a sequence of integers
(0 indicating the first element),
#. a sequence of booleans marking the selected elements,
#. :py:obj:`None`, meaning empty selection,
#. :py:obj:`.All`, meaning "all possible reference points",
#. :py:obj:`.End`, selecting the end of the last element.
n_elements: Length of the sequence of elements
endpoint: if :py:obj:`True`, allow *n_elements* as a
special index, referring to the end of the last element.
types: Allowed types
Returns:
uint32_ref (Uint32Refpts): :py:obj:`~numpy.uint32` numpy array used
for indexing :py:class:`.Element`\ s in a lattice.
"""
refs = numpy.ravel(refpts)
if refpts is RefptsCode.All:
stop = n_elements+1 if endpoint else n_elements
return numpy.arange(stop, dtype=numpy.uint32)
elif refpts is RefptsCode.End:
if not endpoint:
raise IndexError('"End" index out of range')
return numpy.array([n_elements], dtype=numpy.uint32)
elif (refpts is None) or (refs.size == 0):
return numpy.array([], dtype=numpy.uint32)
elif numpy.issubdtype(refs.dtype, numpy.bool_):
return numpy.flatnonzero(refs).astype(numpy.uint32)
elif numpy.issubdtype(refs.dtype, numpy.integer):
# Handle negative indices
if endpoint:
refs = numpy.array([i if (i == n_elements) else i % n_elements
for i in refs], dtype=numpy.uint32)
else:
refs = numpy.array([i % n_elements
for i in refs], dtype=numpy.uint32)
# Check ascending
if refs.size > 1:
prev = refs[0]
for nxt in refs[1:]:
if nxt < prev:
raise IndexError('Index out of range or not in ascending'
' order')
elif nxt == prev:
raise IndexError('Duplicated index')
prev = nxt
return refs
else:
raise _type_error(refpts, types)
# noinspection PyIncorrectDocstring
[docs]def get_uint32_index(ring: Sequence[Element], refpts: Refpts,
endpoint: bool = True, regex=False) -> Uint32Refpts:
# noinspection PyUnresolvedReferences, PyShadowingNames
r"""Returns an integer array of element indices, selecting ring elements.
Parameters:
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
endpoint: if :py:obj:`True`, allow *len(ring)* as a
special index, referring to the end of the last element.
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
uint32_ref (Uint32Refpts): uint32 numpy array used for indexing
:py:class:`.Element`\ s in a lattice.
Examples:
>>> get_uint32_index(ring, at.Sextupole)
array([ 21, 27, 35, 89, 97, 103], dtype=uint32)
numpy array of indices of all :py:class:`.Sextupole`\ s
>>> get_uint32_index(ring, at.End)
array([121], dtype=uint32)
numpy array([:pycode:`len(ring)+1`])
>>> get_uint32_index(ring, at.checkattr('Frequency'))
array([0], dtype=uint32)
numpy array of indices of all elements having a 'Frequency'
attribute
"""
if isinstance(refpts, type):
checkfun = checktype(refpts)
elif callable(refpts):
checkfun = refpts
elif isinstance(refpts, Element):
checkfun = checktype(type(refpts))
elif isinstance(refpts, str):
checkfun = checkname(refpts, regex=regex)
else:
return uint32_refpts(refpts, len(ring), endpoint=endpoint, types=_typ2)
return numpy.fromiter((i for i, el in enumerate(ring) if checkfun(el)),
dtype=numpy.uint32)
[docs]def bool_refpts(refpts: RefIndex, n_elements: int,
endpoint: bool = True,
types: str = _typ1) -> BoolRefpts:
r"""Returns a :py:class:`bool` array of element indices, selecting ring
elements.
Parameters:
refpts: Element selector. *refpts* may be:
#. an integer or a sequence of integers
(0 indicating the first element),
#. a sequence of booleans marking the selected elements,
#. :py:obj:`None`, meaning empty selection,
#. :py:obj:`.All`, meaning "all possible reference points",
#. :py:obj:`.End`, selecting the end of the last element.
n_elements: Length of the lattice
endpoint: if :py:obj:`True`, allow *n_elements* as a
special index, referring to the end of the last element.
types: Allowed types
Returns:
bool_refs (BoolRefpts): A bool numpy array used for indexing
:py:class:`.Element`\ s in a lattice.
"""
refs = numpy.ravel(refpts)
stop = n_elements+1 if endpoint else n_elements
if refpts is RefptsCode.All:
return numpy.ones(stop, dtype=bool)
elif refpts is RefptsCode.End:
if not endpoint:
raise IndexError('"End" index out of range')
brefpts = numpy.zeros(stop, dtype=bool)
brefpts[n_elements] = True
return brefpts
elif (refpts is None) or (refs.size == 0):
return numpy.zeros(stop, dtype=bool)
elif numpy.issubdtype(refs.dtype, numpy.bool_):
diff = stop - refs.size
if diff <= 0:
return refs[:stop]
else:
return numpy.append(refs, numpy.zeros(diff, dtype=bool))
elif numpy.issubdtype(refs.dtype, numpy.integer):
brefpts = numpy.zeros(stop, dtype=bool)
brefpts[refs] = True
return brefpts
else:
raise _type_error(refpts, types)
# noinspection PyIncorrectDocstring
[docs]def get_bool_index(ring: Sequence[Element], refpts: Refpts,
endpoint: bool = True, regex=False) -> BoolRefpts:
# noinspection PyUnresolvedReferences, PyShadowingNames
r"""Returns a bool array of element indices, selecting ring elements.
Parameters:
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
endpoint: if :py:obj:`True`, allow *len(ring)* as a
special index, referring to the end of the last element.
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
bool_refs (BoolRefpts): A bool numpy array used for indexing
:py:class:`.Element`\ s in a lattice.
Examples:
>>> refpts = get_bool_index(ring, at.Quadrupole)
Returns a numpy array of booleans where all :py:class:`.Quadrupole`
are :py:obj:`True`
>>> refpts = get_bool_index(ring, "Q[FD]*")
Returns a numpy array of booleans where all elements whose *FamName*
matches "Q[FD]*" are :py:obj:`True`
>>> refpts = get_bool_index(ring, at.checkattr('K', 0.0))
Returns a numpy array of booleans where all elements whose *K*
attribute is 0.0 are :py:obj:`True`
>>> refpts = get_bool_index(ring, None)
Returns a numpy array of *len(ring)+1* :py:obj:`False` values
"""
if isinstance(refpts, type):
checkfun = checktype(refpts)
elif callable(refpts):
checkfun = refpts
elif isinstance(refpts, Element):
checkfun = checktype(type(refpts))
elif isinstance(refpts, str):
checkfun = checkname(refpts, regex=regex)
else:
return bool_refpts(refpts, len(ring), endpoint=endpoint, types=_typ2)
boolrefs = numpy.fromiter(map(checkfun, ring), dtype=bool, count=len(ring))
if endpoint:
boolrefs = numpy.append(boolrefs, False)
return boolrefs
[docs]def checkattr(attrname: str, attrvalue: Optional = None) \
-> ElementFilter:
# noinspection PyUnresolvedReferences
r"""Checks the presence or the value of an attribute
Returns a function to be used as an :py:class:`.Element` filter, which
checks the presence or the value of an attribute of the
provided :py:class:`.Element`.
This function can be used to extract from a ring all elements
having a given attribute.
Parameters:
attrname: Attribute name
attrvalue: Attribute value. If absent, the returned function checks
the presence of an *attrname* attribute. If present, the
returned function checks if :pycode:`attrname == attrvalue`.
Returns:
checkfun (ElementFilter): Element filter function
Examples:
>>> cavs = filter(checkattr('Frequency'), ring)
Returns an iterator over all elements in *ring* that have a
:pycode:`Frequency` attribute
>>> elts = filter(checkattr('K', 0.0), ring)
Returns an iterator over all elements in ring that have a
:pycode:`K` attribute equal to 0.0
"""
def testf(el):
try:
v = getattr(el, attrname)
return (attrvalue is None) or (v == attrvalue)
except AttributeError:
return False
return testf
[docs]def checktype(eltype: Union[type, Tuple[type, ...]]) -> ElementFilter:
# noinspection PyUnresolvedReferences
r"""Checks the type of an element
Returns a function to be used as an :py:class:`.Element` filter, which
checks the type of the provided :py:class:`.Element`.
This function can be used to extract from a ring all elements
having a given type.
Parameters:
eltype: Desired :py:class:`.Element` type
Returns:
checkfun (ElementFilter): Element filter function
Examples:
>>> qps = filter(checktype(at.Quadrupole), ring)
Returns an iterator over all quadrupoles in ring
"""
return lambda el: isinstance(el, eltype)
[docs]def checkname(pattern: str, regex=False) -> ElementFilter:
# noinspection PyUnresolvedReferences
r"""Checks the name of an element
Returns a function to be used as an :py:class:`.Element` filter,
which checks the name of the provided :py:class:`.Element`.
This function can be used to extract from a ring all elements
having a given name.
Parameters:
pattern: Desired :py:class:`.Element` name. Unix shell-style
wildcards are supported (see :py:func:`fnmatch.fnmatch`)
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
checkfun (ElementFilter): Element filter function
Examples:
>>> qps = filter(checkname('QF*'), ring)
Returns an iterator over all with name starting with ``QF``.
"""
if regex:
rgx = re.compile(pattern)
return lambda el: rgx.fullmatch(el.FamName)
else:
return lambda el: fnmatch(el.FamName, pattern)
[docs]def refpts_iterator(ring: Sequence[Element], refpts: Refpts, regex=False) \
-> Iterator[Element]:
r"""Return an iterator over selected elements in a lattice
Parameters:
ring: Lattice description
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
elem_iter (Iterator[Element]): Iterator over the elements in *ring*
selected by *refpts*.
"""
if isinstance(refpts, type):
checkfun = checktype(refpts)
elif callable(refpts):
checkfun = refpts
elif isinstance(refpts, Element):
checkfun = checktype(type(refpts))
elif isinstance(refpts, str):
checkfun = checkname(refpts, regex=regex)
else:
refs = numpy.ravel(refpts)
if refpts is RefptsCode.All:
return (el for el in ring)
elif refpts is RefptsCode.End:
raise IndexError('"End" is not allowed for endpoint=False')
elif (refpts is None) or (refs.size == 0):
return iter(())
elif numpy.issubdtype(refs.dtype, numpy.bool_):
return compress(ring, refs)
elif numpy.issubdtype(refs.dtype, numpy.integer):
return (ring[i] for i in refs)
else:
raise _type_error(refpts, _typ2)
return filter(checkfun, ring)
# noinspection PyUnusedLocal,PyIncorrectDocstring
[docs]def refpts_count(refpts: RefIndex, n_elements: int,
endpoint: bool = True,
types: str = _typ1) -> int:
r"""Returns the number of reference points
Parameters:
refpts: refpts may be:
#. an integer or a sequence of integers
(0 indicating the first element),
#. a sequence of booleans marking the selected elements,
#. :py:obj:`None`, meaning empty selection,
#. :py:obj:`.All`, meaning "all possible reference points",
#. :py:obj:`.End`, selecting the end of the last element.
n_elements: Lattice length
endpoint: if :py:obj:`True`, allow *n_elements* as a
special index, referring to the end of the last element.
Returns:
nrefs (int): The number of reference points
"""
refs = numpy.ravel(refpts)
if refpts is RefptsCode.All:
return n_elements+1 if endpoint else n_elements
elif refpts is RefptsCode.End:
if not endpoint:
raise IndexError('"End" index out of range')
return 1
elif (refpts is None) or (refs.size == 0):
return 0
elif numpy.issubdtype(refs.dtype, numpy.bool_):
return numpy.count_nonzero(refs)
elif numpy.issubdtype(refs.dtype, numpy.integer):
return len(refs)
else:
raise _type_error(refpts, types)
def _refcount(ring: Sequence[Element], refpts: Refpts,
endpoint: bool = True, regex=False) -> int:
# noinspection PyUnresolvedReferences, PyShadowingNames
r"""Returns the number of reference points
Parameters:
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
endpoint: if :py:obj:`True`, allow *len(ring)* as a
special index, referring to the end of the last element.
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
nrefs (int): The number of reference points
Examples:
>>> refpts = ring.refcount(at.Sextupole)
6
Returns the number of :py:class:`.Sextupole`\ s in the lattice
>>> refpts = ring.refcount(at.All)
122
Returns *len(ring)+1*
>>> refpts = ring.refcount(at.All, endpoint=False)
121
Returns *len(ring)*
"""
if isinstance(refpts, type):
checkfun = checktype(refpts)
elif callable(refpts):
checkfun = refpts
elif isinstance(refpts, Element):
checkfun = checktype(type(refpts))
elif isinstance(refpts, str):
checkfun = checkname(refpts, regex=regex)
else:
return refpts_count(refpts, len(ring), endpoint=endpoint, types=_typ2)
return len(list(filter(checkfun, ring)))
# noinspection PyUnusedLocal,PyIncorrectDocstring
[docs]def get_elements(ring: Sequence[Element], refpts: Refpts, regex=False) \
-> list:
r"""Returns a list of elements selected by *key*.
Deprecated: :pycode:`get_elements(ring, refpts)` is :pycode:`ring[refpts]`
Parameters:
ring: Lattice description
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
elem_list (list): list of :py:class:`.Element`\ s matching key
"""
return list(refpts_iterator(ring, refpts, regex=regex))
[docs]def get_value_refpts(ring: Sequence[Element], refpts: Refpts,
attrname: str, index: Optional[int] = None, regex=False):
r"""Extracts attribute values from selected
lattice :py:class:`.Element`\ s.
Parameters:
ring: Lattice description
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
attrname: Attribute name
index: index of the value to retrieve if *attrname* is
an array.
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
If :py:obj:`None` the full array is retrieved
Returns:
attrvalues: numpy Array of attribute values.
"""
if index is None:
def getf(elem):
return getattr(elem, attrname)
else:
def getf(elem):
return getattr(elem, attrname)[index]
return numpy.array([getf(elem) for elem in refpts_iterator(ring, refpts,
regex=regex)])
[docs]def set_value_refpts(ring: Sequence[Element], refpts: Refpts,
attrname: str, attrvalues, index: Optional[int] = None,
increment: Optional[bool] = False,
copy: Optional[bool] = False, regex=False):
r"""Set the values of an attribute of an array of elements based on
their refpts
Parameters:
ring: Lattice description
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
attrname: Attribute name
attrvalues: Attribute values
index: index of the value to set if *attrname* is
an array. if :py:obj:`None`, the full array is replaced by
*attrvalue*
increment: Add values to the initial values.
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
If :py:obj:`False` the initial value is replaced (Default)
copy: If :py:obj:`False`, the modification is done in-place,
If :py:obj:`True`, return a shallow copy of the lattice. Only the
modified elements are copied.
.. Caution::
a shallow copy means that all non-modified
elements are shared with the original lattice.
Any further modification will affect in both lattices.
"""
if index is None:
def setf(elem, value):
setattr(elem, attrname, value)
else:
def setf(elem, value):
getattr(elem, attrname)[index] = value
if increment:
attrvalues += get_value_refpts(ring, refpts,
attrname, index=index, regex=regex)
else:
attrvalues = numpy.broadcast_to(attrvalues,
(_refcount(ring, refpts, regex=regex),))
# noinspection PyShadowingNames
@make_copy(copy)
def apply(ring, refpts, values, regex):
for elm, val in zip(refpts_iterator(ring, refpts, regex=regex), values):
setf(elm, val)
return apply(ring, refpts, attrvalues, regex)
[docs]def get_s_pos(ring: Sequence[Element], refpts: Refpts = All, regex=False) \
-> Sequence[float]:
# noinspection PyUnresolvedReferences
r"""Returns the locations of selected elements
Parameters:
ring: Lattice description
refpts: Element selection key.
See ":ref:`Selecting elements in a lattice <refpts>`"
regex: Use regular expression for refpts string matching;
Default: False (Unix shell-style wildcards)
Returns:
s_pos: Array of locations of the elements selected by *refpts*
Example:
>>> get_s_pos(ring, at.End)
array([26.37428795])
Position at the end of the last element: length of the lattice
"""
# Positions at the end of each element.
s_pos = numpy.cumsum([getattr(el, 'Length', 0.0) for el in ring])
# Prepend position at the start of the first element.
s_pos = numpy.concatenate(([0.0], s_pos))
return s_pos[get_bool_index(ring, refpts, regex=regex)]
[docs]def rotate_elem(elem: Element, tilt: float = 0.0, pitch: float = 0.0,
yaw: float = 0.0, relative: bool = False) -> None:
r"""Set the tilt, pitch and yaw angle of an :py:class:`.Element`.
The tilt is a rotation around the *s*-axis, the pitch is a
rotation around the *x*-axis and the yaw is a rotation around
the *y*-axis.
A positive angle represent a clockwise rotation when looking in
the direction of the rotation axis.
The transformations are not all commmutative, the tilt is always the
last transformation applied.
If *relative* is :py:obj:`True`, the previous angle and shifts are
rebuilt form the *R* and *T* matrix and incremented by the input arguments.
The shift is always conserved regardless of the value of *relative*.
Parameters:
elem: Element to be tilted
tilt: Tilt angle [rad]
pitch: Pitch angle [rad]
yaw: Yaw angle [rad]
relative: If :py:obj:`True`, the rotation is added to the
previous one
"""
# noinspection PyShadowingNames
def _get_rm_tv(le, tilt, pitch, yaw):
tilt = numpy.around(tilt, decimals=15)
pitch = numpy.around(pitch, decimals=15)
yaw = numpy.around(yaw, decimals=15)
ct, st = numpy.cos(tilt), numpy.sin(tilt)
ap, ay = -0.5*le*numpy.sin(pitch), -0.5*le*numpy.sin(yaw)
rr1 = numpy.asfortranarray(numpy.diag([ct, ct, ct, ct, 1.0, 1.0]))
rr1[0, 2] = st
rr1[1, 3] = st
rr1[2, 0] = -st
rr1[3, 1] = -st
rr2 = rr1.T
t1 = numpy.array([ay, -yaw, ap, -pitch, 0, 0])
t2 = numpy.array([ay, yaw, ap, pitch, 0, 0])
rt1 = numpy.eye(6, order='F')
rt1[1, 4] = ct*t1[1]
rt1[3, 4] = ct*t1[3]
rt2 = numpy.eye(6, order='F')
rt2[1, 4] = ct*t2[1]
rt2[3, 4] = ct*t2[3]
return rr1 @ rt1, rt2 @ rr2, t1, t2
tilt0 = 0.0
pitch0 = 0.0
yaw0 = 0.0
t10 = numpy.zeros(6)
t20 = numpy.zeros(6)
if hasattr(elem, 'R1') and hasattr(elem, 'R2'):
rr10 = numpy.eye(6, order='F')
rr10[:4, :4] = elem.R1[:4, :4]
rt10 = rr10.T @ elem.R1
tilt0 = numpy.arctan2(rr10[0, 2], rr10[0, 0])
yaw0 = -rt10[1, 4]/rr10[0, 0]
pitch0 = -rt10[3, 4]/rr10[0, 0]
_, _, t10, t20 = _get_rm_tv(elem.Length, tilt0, pitch0, yaw0)
if hasattr(elem, 'T1') and hasattr(elem, 'T2'):
t10 = elem.T1-t10
t20 = elem.T2-t20
if relative:
tilt += tilt0
pitch += pitch0
yaw += yaw0
r1, r2, t1, t2 = _get_rm_tv(elem.Length, tilt, pitch, yaw)
elem.R1 = r1
elem.R2 = r2
elem.T1 = t1+t10
elem.T2 = t2+t20
[docs]def tilt_elem(elem: Element, rots: float, relative: bool = False) -> None:
r"""Set the tilt angle :math:`\theta` of an :py:class:`.Element`
The rotation matrices are stored in the :pycode:`R1` and :pycode:`R2`
attributes.
:math:`R_1=\begin{pmatrix} cos\theta & sin\theta \\
-sin\theta & cos\theta \end{pmatrix}`,
:math:`R_2=\begin{pmatrix} cos\theta & -sin\theta \\
sin\theta & cos\theta \end{pmatrix}`
Parameters:
elem: Element to be tilted
rots: Tilt angle :math:`\theta` [rd]. *rots* > 0 corresponds
to a corkscrew rotation of the element looking in the direction of
the beam
relative: If :py:obj:`True`, the rotation is added to the
existing one
"""
rotate_elem(elem, tilt=rots, relative=relative)
[docs]def shift_elem(elem: Element, deltax: float = 0.0, deltaz: float = 0.0,
relative: Optional[bool] = False) -> None:
r"""Sets the transverse displacement of an :py:class:`.Element`
The translation vectors are stored in the :pycode:`T1` and :pycode:`T2`
attributes.
Parameters:
elem: Element to be shifted
deltax: Horizontal displacement [m]
deltaz: Vertical displacement [m]
relative: If :py:obj:`True`, the translation is added to the
existing one
"""
tr = numpy.array([deltax, 0.0, deltaz, 0.0, 0.0, 0.0])
if relative and hasattr(elem, 'T1') and hasattr(elem, 'T2'):
elem.T1 -= tr
elem.T2 += tr
else:
elem.T1 = -tr
elem.T2 = tr
[docs]def set_rotation(ring: Sequence[Element], tilts=0.0,
pitches=0.0, yaws=0.0, relative=False) -> None:
r"""Sets the tilts of a list of elements.
Parameters:
ring: Lattice description
tilts: Sequence of tilt values as long as ring or
scalar tilt value applied to all elements, default=0
pitches: Sequence of pitch values as long as ring or
scalar tilt value applied to all elements, default=0
yaws: Sequence of yaw values as long as ring or
scalar tilt value applied to all elements, default=0
relative: If :py:obj:`True`, the rotations are added to the
existing ones
"""
tilts = numpy.broadcast_to(tilts, (len(ring),))
pitches = numpy.broadcast_to(pitches, (len(ring),))
yaws = numpy.broadcast_to(yaws, (len(ring),))
for el, tilt, pitch, yaw in zip(ring, tilts, pitches, yaws):
rotate_elem(el, tilt=tilt, pitch=pitch, yaw=yaw, relative=relative)
[docs]def set_tilt(ring: Sequence[Element], tilts, relative=False) -> None:
r"""Sets the tilts of a list of elements.
Parameters:
ring: Lattice description
tilts: Sequence of tilt values as long as ring or
scalar tilt value applied to all elements
relative: If :py:obj:`True`, the rotation is added to the
existing one
"""
tilts = numpy.broadcast_to(tilts, (len(ring),))
for el, tilt in zip(ring, tilts):
tilt_elem(el, tilt, relative=relative)
[docs]def set_shift(ring: Sequence[Element], dxs, dzs, relative=False) -> None:
r"""Sets the translations of a list of elements.
Parameters:
ring: Lattice description
dxs: Sequence of horizontal displacements values as long as
ring or scalar value applied to all elements [m]
dzs: Sequence of vertical displacements values as long as
ring or scalar value applied to all elements [m]
relative: If :py:obj:`True`, the displacement is added to the
existing one
"""
dxs = numpy.broadcast_to(dxs, (len(ring),))
dzs = numpy.broadcast_to(dzs, (len(ring),))
for el, dx, dy in zip(ring, dxs, dzs):
shift_elem(el, dx, dy, relative=relative)
[docs]def get_geometry(ring: List[Element],
start_coordinates: Tuple[float, float, float] = (0, 0, 0),
centered: bool = False):
# noinspection PyShadowingNames
r"""Compute the 2D ring geometry in cartesian coordinates
Parameters:
ring: Lattice description
start_coordinates: x ,y, angle at starting point
centered: if :py:obj:`True` the coordinates origin is the
center of the ring
Returns:
geomdata: recarray containing, x, y, angle
radius: machine radius
Example:
>>> geomdata, radius = get_geometry(ring)
"""
geom_dtype = [('x', numpy.float64),
('y', numpy.float64),
('angle', numpy.float64)]
geomdata = numpy.recarray((len(ring)+1, ), dtype=geom_dtype)
xx = numpy.zeros(len(ring)+1)
yy = numpy.zeros(len(ring)+1)
angle = numpy.zeros(len(ring)+1)
x, y, t = start_coordinates
x0, y0, t0 = start_coordinates
for ind, el in enumerate(ring+[ring[0]]):
ll = el.Length
if isinstance(el, Dipole) and el.BendingAngle != 0:
ang = 0.5 * el.BendingAngle
ll *= numpy.sin(ang)/ang
else:
ang = 0.0
t -= ang
x += ll * numpy.cos(t)
y += ll * numpy.sin(t)
t -= ang
xx[ind] = x
yy[ind] = y
angle[ind] = t
radius = get_s_pos(ring, len(ring))[0] / abs(t-t0) \
if t != 0.0 else 0.0
if centered:
xx += -abs(radius)*numpy.sin(t0) - x0
yy += abs(radius)*numpy.cos(t0) - y0
geomdata['x'] = xx
geomdata['y'] = yy
geomdata['angle'] = angle
return geomdata, radius