at.lattice.utils#

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.

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:

  1. 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,

  2. an ordered list of such integers without duplicates,

  3. a numpy array of booleans of maximum length len(lattice)+1, where selected elements are True,

  4. None, meaning an empty selection,

  5. All, meaning “all possible reference points”: the entrance of all elements plus the end of the last element,

  6. End, selecting the end of the last element,

  7. an element type, selecting all the elements of that type in the lattice, e.g. at.Sextupole,

  8. a string, selecting all the elements whose FamName attribute matches it. Unix shell-style wildcards are accepted, e.g. “Q[FD]*”,

  9. a callable filtfunc such that filtfunc(elem) is True for selected elements.

Module Attributes

All

All is a special value to be used as refpts.

End

End is a special value to be used as refpts.

Functions

axis_descr(axis [ ,axis][, key])

Return a tuple containing for each input argument the requested information

check_radiation(rad)

Deprecated decorator for optics functions (see check_6d()).

check_6d(is_6d)

Decorator for optics functions

set_radiation(rad)

Deprecated decorator for optics functions (see set_6d()).

set_6d(is_6d)

Decorator for optics functions

make_copy(copy)

Decorator for optics functions

uint32_refpts(refpts, n_elements[, ...])

Return a uint32 array of element indices selecting ring elements.

bool_refpts(refpts, n_elements[, endpoint, ...])

Returns a bool array of element indices, selecting ring elements.

get_uint32_index(ring, refpts[, endpoint, regex])

Returns an integer array of element indices, selecting ring elements.

get_bool_index(ring, refpts[, endpoint, regex])

Returns a bool array of element indices, selecting ring elements.

checkattr(attrname[, attrvalue])

Checks the presence or the value of an attribute

checktype(eltype)

Checks the type of an element

checkname(pattern[, regex])

Checks the name of an element

get_elements(ring, refpts[, regex])

Returns a list of elements selected by key.

get_s_pos(ring[, refpts, regex])

Returns the locations of selected elements

refpts_count(refpts, n_elements[, endpoint, ...])

Returns the number of reference points

refpts_iterator(ring, refpts[, regex])

Return an iterator over selected elements in a lattice

set_shift(ring, dxs, dzs[, relative])

Sets the translations of a list of elements.

set_tilt(ring, tilts[, relative])

Sets the tilts of a list of elements.

set_rotation(ring[, tilts, pitches, yaws, ...])

Sets the tilts of a list of elements.

tilt_elem(elem, rots[, relative])

Set the tilt angle \(\theta\) of an Element

shift_elem(elem[, deltax, deltaz, relative])

Sets the transverse displacement of an Element

rotate_elem(elem[, tilt, pitch, yaw, relative])

Set the tilt, pitch and yaw angle of an Element.

get_value_refpts(ring, refpts, attrname[, ...])

Extracts attribute values from selected

set_value_refpts(ring, refpts, attrname, ...)

Set the values of an attribute of an array of elements based on their refpts

get_geometry(ring[, refpts, ...])

Compute the 2D ring geometry in cartesian coordinates

setval(attrname[, index])

Return a callable object which sets the value of item index of attribute attrname of its 1st argument to it 2nd orgument.

getval(attrname[, index])

Return a callable object which fetches item index of attribute attrname of its operand.

Exceptions

AtError

AtWarning

exception AtError[source]#

Bases: Exception

exception AtWarning[source]#

Bases: UserWarning

axis_descr(axis, [axis, ]key=None)[source]#

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

    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

bool_refpts(refpts, n_elements, endpoint=True, types='None, All, End, int, bool')[source]#

Returns a bool array of element indices, selecting ring elements.

Parameters:

  • refpts (None | int | Sequence[int] | bool | Sequence[bool] | RefptsCode) –

    Element selector. refpts may be:

    1. an integer or a sequence of integers (0 indicating the first element),

    2. a sequence of booleans marking the selected elements,

    3. None, meaning empty selection,

    4. All, meaning “all possible reference points”,

    5. End, selecting the end of the last element.

  • n_elements (int) – Length of the lattice

  • endpoint (bool) – if True, allow n_elements as a special index, referring to the end of the last element.

  • types (str) – Allowed types

Returns:

bool_refs (BoolRefpts) – A bool numpy array used for indexing Elements in a lattice.

check_6d(is_6d)[source]#

Decorator for optics functions

Wraps a function like func(ring, *args, **kwargs) where ring is a Lattice object. Raise AtError if ring.is_6d is not the desired is_6d state. No test is performed if ring is not a Lattice.

Parameters:

is_6d (bool) – Desired 6D state

Raises:

AtError – if ring.is_6d is not is_6d

Example

>>> @check_6d(False)
... def find_orbit4(ring, dct=None, guess=None, **kwargs):
        ...

Raises AtError if ring.is_6d is True

See also

set_6d()

checkattr(attrname, attrvalue=None)[source]#

Checks the presence or the value of an attribute

Returns a function to be used as an Element filter, which checks the presence or the value of an attribute of the provided Element. This function can be used to extract from a ring all elements having a given attribute.

Parameters:

  • attrname (str) – Attribute name

  • attrvalue (Optional) – Attribute value. If absent, the returned function checks the presence of an attrname attribute. If present, the returned function checks if 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 Frequency attribute

>>> elts = filter(checkattr('K', 0.0), ring)

Returns an iterator over all elements in ring that have a K attribute equal to 0.0

checkname(pattern, regex=False)[source]#

Checks the name of an element

Returns a function to be used as an Element filter, which checks the name of the provided Element. This function can be used to extract from a ring all elements having a given name.

Parameters:

  • pattern (str) – Desired Element name. Unix shell-style wildcards are supported (see fnmatch.fnmatch())

  • regex (bool) – Use regular expression for refpts string matching instead of 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.

checktype(eltype)[source]#

Checks the type of an element

Returns a function to be used as an Element filter, which checks the type of the provided Element. This function can be used to extract from a ring all elements having a given type.

Parameters:

eltype (type | Tuple[type, ...]) – Desired Element type

Returns:

checkfun (ElementFilter) – Element filter function

Examples

>>> qps = filter(checktype(at.Quadrupole), ring)

Returns an iterator over all quadrupoles in ring

get_bool_index(ring, refpts, endpoint=True, regex=False)[source]#

Returns a bool array of element indices, selecting ring elements.

Parameters:
Returns:

bool_refs (BoolRefpts) – A bool numpy array used for indexing Elements in a lattice.

Examples

>>> refpts = get_bool_index(ring, at.Quadrupole)

Returns a numpy array of booleans where all Quadrupole are True

>>> refpts = get_bool_index(ring, "Q[FD]*")

Returns a numpy array of booleans where all elements whose FamName matches “Q[FD]*” are 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 True

>>> refpts = get_bool_index(ring, None)

Returns a numpy array of len(ring)+1 False values

get_elements(ring, refpts, regex=False)[source]#

Returns a list of elements selected by key.

Deprecated: get_elements(ring, refpts) is ring[refpts]

Parameters:
Returns:

elem_list (list) – list of Elements matching key

get_geometry(ring, refpts=RefptsCode.All, start_coordinates=(0, 0, 0), centered=False, regex=False)[source]#

Compute the 2D ring geometry in cartesian coordinates

Parameters:
Returns:

  • geomdata – recarray containing, x, y, angle.

  • radius – machine radius at the beginning of the lattice.

    Attention

    This radius is different from the radius usually defined as \(C/2\pi\)

Example

>>> geomdata, radius = get_geometry(ring)
get_s_pos(ring, refpts=RefptsCode.All, regex=False)[source]#

Returns the locations of selected elements

Parameters:
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

get_uint32_index(ring, refpts, endpoint=True, regex=False)[source]#

Returns an integer array of element indices, selecting ring elements.

Parameters:
Returns:

uint32_ref (Uint32Refpts) – uint32 numpy array used for indexing Elements 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 Sextupoles

>>> get_uint32_index(ring, at.End)
array([121], dtype=uint32)

numpy array([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

get_value_refpts(ring, refpts, attrname, index=None, regex=False)[source]#
Extracts attribute values from selected

lattice Elements.

Parameters:
Returns:

attrvalues – numpy Array of attribute values.

getval(attrname, index=None)[source]#

Return a callable object which fetches item index of attribute attrname of its operand. Examples:

  • After f = getval('Length'), f(elem) returns elem.Length

  • After f = getval('PolynomB, index=1), f(elem) returns elem.PolynomB[1]

make_copy(copy)[source]#

Decorator for optics functions

Wraps a function like func(ring, refpts, *args, **kwargs) where ring is a Lattice object.

If copy is False, the function is not modified,

If copy is True, a shallow copy of ring is done, then the elements selected by refpts are deep-copied, then func is applied to the copy, and the new ring is returned.

Parameters:

copy (bool) – If True, apply the decorated function to a copy of ring

refpts_count(refpts, n_elements, endpoint=True, types='None, All, End, int, bool')[source]#

Returns the number of reference points

Parameters:

  • refpts (None | int | Sequence[int] | bool | Sequence[bool] | RefptsCode) –

    refpts may be:

    1. an integer or a sequence of integers (0 indicating the first element),

    2. a sequence of booleans marking the selected elements,

    3. None, meaning empty selection,

    4. All, meaning “all possible reference points”,

    5. End, selecting the end of the last element.

  • n_elements (int) – Lattice length

  • endpoint (bool) – if True, allow n_elements as a special index, referring to the end of the last element.

Returns:

nrefs (int) – The number of reference points

refpts_iterator(ring, refpts, regex=False)[source]#

Return an iterator over selected elements in a lattice

Parameters:
Returns:

elem_iter (Iterator[Element]) – Iterator over the elements in ring selected by refpts.

rotate_elem(elem, tilt=0.0, pitch=0.0, yaw=0.0, relative=False)[source]#

Set the tilt, pitch and yaw angle of an 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 represents a clockwise rotation when looking in the direction of the rotation axis.

The transformations are not all commmutative, the pitch and yaw are applied first and the tilt is always the last transformation applied. The element is rotated around its mid-point.

If relative is 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.

The transformations are applied by changing the particle coordinates at the entrance of the element and restoring them at the end. Following the small angles approximation the longitudinal shift of the particle coordinates is neglected and the element length is unchanged.

Parameters:

  • elem (Element) – Element to be tilted

  • tilt (float) – Tilt angle [rad]

  • pitch (float) – Pitch angle [rad]

  • yaw (float) – Yaw angle [rad]

  • relative (bool) – If True, the rotation is added to the previous one

set_6d(is_6d)[source]#

Decorator for optics functions

Wraps a function like func(ring, *args, **kwargs) where ring is a Lattice object. If ring.is_6d is not the desired is_6d state, func() will be called with a modified copy of ring satisfying the is_6d state.

Parameters:

is_6d (bool) – 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:

check_6d(), Lattice.enable_6d(), Lattice.disable_6d()

set_rotation(ring, tilts=0.0, pitches=0.0, yaws=0.0, relative=False)[source]#

Sets the tilts of a list of elements.

Parameters:

  • ring (Sequence[Element]) – 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 True, the rotations are added to the existing ones

set_shift(ring, dxs, dzs, relative=False)[source]#

Sets the translations of a list of elements.

Parameters:

  • ring (Sequence[Element]) – 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 True, the displacement is added to the existing one

set_tilt(ring, tilts, relative=False)[source]#

Sets the tilts of a list of elements.

Parameters:

  • ring (Sequence[Element]) – Lattice description

  • tilts – Sequence of tilt values as long as ring or scalar tilt value applied to all elements

  • relative – If True, the rotation is added to the existing one

set_value_refpts(ring, refpts, attrname, attrvalues, index=None, increment=False, copy=False, regex=False)[source]#

Set the values of an attribute of an array of elements based on their refpts

Parameters:

  • ring (Sequence[Element]) – Lattice description

  • refpts (Type[Element] | Element | Callable[[Element], bool] | str | None | int | Sequence[int] | bool | Sequence[bool] | RefptsCode) – Element selection key. See “Selecting elements in a lattice”

  • attrname (str) – Attribute name

  • attrvalues – Attribute values

  • index (int | None) – index of the value to set if attrname is an array. if None, the full array is replaced by attrvalue

  • increment (bool) – If True, attrvalues are added to the initial values.

  • regex (bool) – Use regular expression for refpts string matching instead of Unix shell-style wildcards.

  • copy (bool) –

    If False, the modification is done in-place,

    If 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 both lattices.

setval(attrname, index=None)[source]#

Return a callable object which sets the value of item index of attribute attrname of its 1st argument to it 2nd orgument.

  • After f = setval('Length'), f(elem, value) is equivalent to elem.Length = value

  • After f = setval('PolynomB, index=1), f(elem, value) is equivalent to elem.PolynomB[1] = value

shift_elem(elem, deltax=0.0, deltaz=0.0, relative=False)[source]#

Sets the transverse displacement of an Element

The translation vectors are stored in the T1 and T2 attributes.

Parameters:

  • elem (Element) – Element to be shifted

  • deltax (float) – Horizontal displacement [m]

  • deltaz (float) – Vertical displacement [m]

  • relative (bool) – If True, the translation is added to the existing one

tilt_elem(elem, rots, relative=False)[source]#

Set the tilt angle \(\theta\) of an Element

The rotation matrices are stored in the R1 and R2 attributes.

\(R_1=\begin{pmatrix} cos\theta & sin\theta \\ -sin\theta & cos\theta \end{pmatrix}\), \(R_2=\begin{pmatrix} cos\theta & -sin\theta \\ sin\theta & cos\theta \end{pmatrix}\)

Parameters:

  • elem (Element) – Element to be tilted

  • rots (float) – Tilt angle \(\theta\) [rd]. rots > 0 corresponds to a corkscrew rotation of the element looking in the direction of the beam

  • relative (bool) – If True, the rotation is added to the existing one

uint32_refpts(refpts, n_elements, endpoint=True, types='None, All, End, int, bool')[source]#

Return a uint32 array of element indices selecting ring elements.

Parameters:

  • refpts (None | int | Sequence[int] | bool | Sequence[bool] | RefptsCode) –

    Element selector. refpts may be:

    1. an integer or a sequence of integers (0 indicating the first element),

    2. a sequence of booleans marking the selected elements,

    3. None, meaning empty selection,

    4. All, meaning “all possible reference points”,

    5. End, selecting the end of the last element.

  • n_elements (int) – Length of the sequence of elements

  • endpoint (bool) – if True, allow n_elements as a special index, referring to the end of the last element.

  • types (str) – Allowed types

Returns:

uint32_ref (Uint32Refpts) – uint32 numpy array used for indexing Elements in a lattice.

All = RefptsCode.All#

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.

End = RefptsCode.End#

End is a special value to be used as refpts. It refers to the end of the last element.