particleflow.elements module¶

Accelerator lattice elements represented by PyTorch Modules.

The various element classes are made available together with the corresponding MADX command name in the following dicts.

particleflow.elements.elements¶

Maps MADX command names to corresponding Element backend classes.

Type: dict

particleflow.elements.alignment_errors¶

Maps MADX EALIGN command attributes to corresponding AlignmentError backend classes.

Type: dict

particleflow.elements.aperture_types¶

Maps MADX apertypes definitions to corresponding Aperture backend classes.

Type: dict

class particleflow.elements.ApertureCircle(aperture: float = <MagicMock name='mock.inf' id='140244239446536'>, offset: Tuple[float, float] = (0.0, 0.0), padding: Union[float, Tuple] = 0.0)¶

Bases: particleflow.elements.ApertureEllipse

Circular aperture.

Parameters: aperture (float) – Radius of the circle.

class particleflow.elements.ApertureEllipse(aperture: Tuple[float, float] = (<MagicMock name='mock.inf' id='140244239446536'>, <MagicMock name='mock.inf' id='140244239446536'>), offset: Tuple[float, float] = (0.0, 0.0), padding: Union[float, Tuple] = 0.0)¶

Bases: particleflow.elements.Aperture

Elliptical aperture.

Parameters: aperture (tuple) – Horizontal and vertical semi-axes of the ellipse.

loss(xy: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶

Compute loss values for the given xy-positions.

Parameters: xy (torch.Tensor) – Tensor of shape (2, N) where N is the number of particles and rows are x- and y-positions respectively.
Returns: loss_val – Tensor of shape (N,), zero where positions are inside the aperture (including exactly at the aperture) and greater than zero where positions are outside the aperture.
Return type: torch.Tensor

class particleflow.elements.ApertureRectangle(aperture: Tuple[float, float] = (<MagicMock name='mock.inf' id='140244239446536'>, <MagicMock name='mock.inf' id='140244239446536'>), offset: Tuple[float, float] = (0.0, 0.0), padding: Union[float, Tuple] = 0.0)¶

Bases: particleflow.elements.Aperture

Rectangular aperture.

Parameters: aperture (tuple) – Half width (horizontal) and half height (vertical) of the rectangle.

loss(xy: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶

Compute loss values for the given xy-positions.

Parameters: xy (torch.Tensor) – Tensor of shape (2, N) where N is the number of particles and rows are x- and y-positions respectively.
Returns: loss_val – Tensor of shape (N,), zero where positions are inside the aperture (including exactly at the aperture) and greater than zero where positions are outside the aperture.
Return type: torch.Tensor

class particleflow.elements.ApertureRectEllipse(aperture: Tuple[float, float, float, float] = (<MagicMock name='mock.inf' id='140244239446536'>, <MagicMock name='mock.inf' id='140244239446536'>, <MagicMock name='mock.inf' id='140244239446536'>, <MagicMock name='mock.inf' id='140244239446536'>), offset: Tuple[float, float] = (0.0, 0.0), padding: Union[float, Tuple] = 0.0)¶

Bases: particleflow.elements.Aperture

Rectangular and elliptical aperture overlaid.

Parameters: aperture (tuple) – Half width (horizontal) and half height (vertical) of the rectangle followed by horizontal and vertical semi-axes of the ellipse.

loss(xy: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶

Compute loss values for the given xy-positions.

Parameters: xy (torch.Tensor) – Tensor of shape (2, N) where N is the number of particles and rows are x- and y-positions respectively.
Returns: loss_val – Tensor of shape (N,), zero where positions are inside the aperture (including exactly at the aperture) and greater than zero where positions are outside the aperture.
Return type: torch.Tensor

class particleflow.elements.Marker(**kwargs)¶

Bases: particleflow.elements.Element

Marker element.

exact(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Exact analytic solution for tracking through the element.

linear(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Linear tracking through the element.

class particleflow.elements.Drift(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Element

Drift space.

Note

Unlike in MADX, drift spaces feature aperture checks here.

Parameters: l (float) – Length of the drift [m].

exact(x)¶: Exact analytic solution for tracking through the element.

makethin(n: int, *, style: Optional[str] = None) → particleflow.elements.Drift¶: Drift spaces are not affected by makethin.

class particleflow.elements.Instrument(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Drift

A place holder for any type of beam instrumentation.

class particleflow.elements.Placeholder(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Drift

A place holder for any type of element.

class particleflow.elements.Monitor(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Drift

Beam position monitor.

class particleflow.elements.HMonitor(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Monitor

Beam position monitor for measuring horizontal beam position.

class particleflow.elements.VMonitor(l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Monitor

Beam position monitor for measuring horizontal beam position.

class particleflow.elements.Kicker(hkick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, vkick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, **kwargs)¶

Bases: particleflow.elements.Element

Combined horizontal and vertical kicker magnet.

Parameters

hkick (torch.Tensor or torch.nn.Parameter) – Horizontal kick [rad].
vkick (torch.Tensor or torch.nn.Parameter) – Vertical kick [rad].

property d¶

linear(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Linear tracking through the element.

class particleflow.elements.HKicker(kick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], **kwargs)¶

Bases: particleflow.elements.Kicker

Horizontal kicker magnet.

property kick¶

class particleflow.elements.VKicker(kick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], **kwargs)¶

Bases: particleflow.elements.Kicker

Vertical kicker magnet.

property kick¶

class particleflow.elements.TKicker(hkick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, vkick: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, **kwargs)¶

Bases: particleflow.elements.Kicker

Similar to Kicker (see Chapter 10.12, MADX User’s Guide).

class particleflow.elements.Quadrupole(k1: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Element

Quadrupole magnet.

Whether this is a (horizontally) focusing or defocusing magnet will be determined from the value of k1 (k1 > 0 indicates a horizontally focusing quadrupole). Hence, in case k1 is a Parameter, it always must be non-zero. For that reason it is convenient to use boundaries e.g. [eps, k1_max] with a small number eps (e.g. 1e-16) and clip the value of k1 accordingly. If k1 is not a Parameter it can be zero and a corresponding Drift transformation will be used.

k1¶

Normalized quadrupole gradient [1/m^2].

Type: torch.Tensor or Parameter

property R¶

makethin(n: int, *, style: Optional[str] = None) → Union[particleflow.elements.Element, particleflow.elements.ThinSegment]¶

Transform element to a sequence of n thin elements using the requested slicing style.

Important

If optimization of parameters is required then it is important to call makethin before each forward pass in order to always use the up-to-date parameter values from the original elements. Calling makethin only once at the beginning and then reusing the resulting sequence on every iteration would reuse the initial parameter values and hence make the optimization ineffective.

Parameters

n (int) – The number of slices (thin elements).
style (str, optional) – The slicing style to be used. For available styles see create_thin_sequence().

Returns

A segment containing the slices (thin elements) separated by drifts.

Return type

Segment

See also

Element.create_thin_sequence()

class particleflow.elements.Sextupole(k2: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], l: float, beam: dict, **kwargs)¶

Bases: particleflow.elements.Element

Sextupole magnet.

k2¶

Normalized sextupole coefficient [1/m^3].

Type: torch.Tensor or Parameter

class particleflow.elements.SBend(angle: float, l: float, beam: dict, e1: float = 0.0, e2: float = 0.0, fint: Union[float, bool] = 0.0, fintx: Optional[float] = None, hgap: float = 0.0, **kwargs)¶

Bases: particleflow.elements.Element

Sector bending magnet.

Parameters

angle (float) – Bending angle of the dipole [rad].
e1 (float) – Rotation angle for the entrance pole face [rad]. e1 = e2 = angle/2 turns an SBEND into a RBEND.
e2 (float) – Rotation angle for the exit pole face [rad].
fint (float) – Fringing field integral at entrance. If fintx is not specified then fint is also used at the exit.
fintx (float) – Fringing field integral at exit.
hgap (float) – Half gap of the magnet [m].

makethin(n: int, *, style: Optional[str] = None) → particleflow.elements.Segment¶

Transform element to a sequence of n thin elements using the requested slicing style.

Important

If optimization of parameters is required then it is important to call makethin before each forward pass in order to always use the up-to-date parameter values from the original elements. Calling makethin only once at the beginning and then reusing the resulting sequence on every iteration would reuse the initial parameter values and hence make the optimization ineffective.

Parameters

n (int) – The number of slices (thin elements).
style (str, optional) – The slicing style to be used. For available styles see create_thin_sequence().

Returns

A segment containing the slices (thin elements) separated by drifts.

Return type

Segment

See also

Element.create_thin_sequence()

class particleflow.elements.RBend(angle: float, l: float, beam: dict, e1: float = 0.0, e2: float = 0.0, fint: float = 0.0, fintx: Optional[float] = None, hgap: float = 0.0, **kwargs)¶: Bases: particleflow.elements.SBend

class particleflow.elements.Dipedge(h: float, e1: float, fint: float, hgap: float, **kwargs)¶

Bases: particleflow.elements.Element

Fringing fields at the entrance and exit of dipole magnets.

Parameters

h (float) – Curvature of the associated dipole magnet body.
e1 (float) – The rotation angle of the pole face.
fint (float) – The fringing field integral.
hgap (float) – The half gap height of the associated dipole magnet.

second_order(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Second order tracking through the element.

class particleflow.elements.ThinQuadrupole(k1l: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], **kwargs)¶

Bases: particleflow.elements.Element

Thin lens representation of a quadrupole magnet.

Parameters: k1l (torch.Tensor or torch.nn.Parameter) –

property R¶

makethin(n: int, *, style: Optional[str] = None) → List[particleflow.elements.Element]¶

Transform element to a sequence of n thin elements using the requested slicing style.

Important

If optimization of parameters is required then it is important to call makethin before each forward pass in order to always use the up-to-date parameter values from the original elements. Calling makethin only once at the beginning and then reusing the resulting sequence on every iteration would reuse the initial parameter values and hence make the optimization ineffective.

Parameters

n (int) – The number of slices (thin elements).
style (str, optional) – The slicing style to be used. For available styles see create_thin_sequence().

Returns

A segment containing the slices (thin elements) separated by drifts.

Return type

Segment

See also

Element.create_thin_sequence()

class particleflow.elements.ThinSextupole(k2l: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>], **kwargs)¶

Bases: particleflow.elements.Element

Thin lens representation of a sextupole magnet.

Parameters: k2l (torch.Tensor or torch.nn.Parameter) –

makethin(n: int, *, style: Optional[str] = None) → List[particleflow.elements.Element]¶

Transform element to a sequence of n thin elements using the requested slicing style.

Important

If optimization of parameters is required then it is important to call makethin before each forward pass in order to always use the up-to-date parameter values from the original elements. Calling makethin only once at the beginning and then reusing the resulting sequence on every iteration would reuse the initial parameter values and hence make the optimization ineffective.

Parameters

n (int) – The number of slices (thin elements).
style (str, optional) – The slicing style to be used. For available styles see create_thin_sequence().

Returns

A segment containing the slices (thin elements) separated by drifts.

Return type

Segment

See also

Element.create_thin_sequence()

class particleflow.elements.Tilt(target: Union[particleflow.elements.Element, particleflow.elements.AlignmentError], psi: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0.0)¶

Bases: particleflow.elements.LongitudinalRoll

The tilt of an element represents the roll about the longitudinal axis.

Note

MADX uses a right-handed coordinate system (see Fig. 1.1, MADX User’s Guide), therefore x = x*cos(psi) - y*sin(psi) describes a clockwise rotation of the trajectory.

psi¶

Rotation angle about s-axis.

Type: torch.Tensor or Parameter

Notes

Tilt is only a subclass of AlignmentError for technical reasons and has no meaning beyond that. A Tilt is not considered an alignment error from the simulation point of view.

triggers = ('tilt',)¶

class particleflow.elements.Offset(target: Union[particleflow.elements.Element, particleflow.elements.AlignmentError], dx: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0.0, dy: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0.0)¶

Bases: particleflow.elements.AlignmentError

AlignmentError representing the xy-offset of an element.

dx¶

Horizontal offset.

Type: torch.Tensor or Parameter

dy¶

Vertical offset.

Type: torch.Tensor or Parameter

property d¶

property d_inv¶

triggers = ('dx', 'dy')¶

class particleflow.elements.LongitudinalRoll(target: Union[particleflow.elements.Element, particleflow.elements.AlignmentError], psi: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0.0)¶

Bases: particleflow.elements.AlignmentError

AlignmentError representing the roll about the longitudinal axis of an element.

Note

MADX uses a right-handed coordinate system (see Fig. 1.1, MADX User’s Guide), therefore x = x*cos(psi) - y*sin(psi) describes a clockwise rotation of the trajectory.

psi¶

Rotation angle about s-axis.

Type: torch.Tensor or Parameter

property R¶

property R_inv¶

triggers = ('dpsi',)¶

class particleflow.elements.BPMError(target: Union[particleflow.elements.Element, particleflow.elements.AlignmentError], ax: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, ay: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, rx: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0, ry: Union[float, <MagicMock name='mock.nn.Parameter' id='140244239883344'>] = 0)¶

Bases: particleflow.elements.AlignmentError

BPM readout errors.

The actual BPM reading is computed as (for both horizontal and vertical plane separately): rx * (x + ax).

ax¶

Horizontal absolute read error.

Type: torch.Tensor or Parameter

ay¶

Vertical absolute read error.

Type: torch.Tensor or Parameter

rx¶

Horizontal relative read error.

Type: torch.Tensor or Parameter

ry¶

Vertical relative read error.

Type: torch.Tensor or Parameter

enter(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Applies linear coordinate transformation at the entrance of the wrapped element.

exit(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶: Applies linear coordinate transformation at the exit of the wrapped element.

readout(x: <MagicMock name='mock.Tensor' id='140244238396048'>) → <MagicMock name='mock.Tensor' id='140244238396048'>¶

Return BPM readings for the given coordinates.

Parameters: x (torch.Tensor) – 6D phase-space coordinates of shape (6, N).
Returns: xy – BPM readings in x- and y-dimension of shape (2, N).
Return type: torch.Tensor

triggers = ('mrex', 'mrey', 'mscalx', 'mscaly')¶

class particleflow.elements.Segment(elements: Sequence[Union[particleflow.elements.Element, particleflow.elements.AlignmentError, particleflow.elements.Segment]])¶

Bases: particleflow.elements.Module

Wrapper class representing a sequence of elements (possibly a segment of the lattice).

Elements or sub-segments can be selected via __getitem__, i.e. segment[item] notation. Here item can be one of the following:

int - indicating the index position in the segment.
str - will be compared for equality against element labels; if a single element with that label is found it is returned otherwise all elements with that label are returned as a list. An exception are strings containing an asterisk which will be interpreted as a shell-style wildcard and converted to a corresponding regex Pattern.
Pattern - will be matched against element labels; a list of all matching elements is returned.
instance of Element or AlignmentError - the element itself (possibly wrapped by other AlignmentError instances is returned.
subclass of Element or AlignmentError - a list of elements of that type (possibly wrapped by AlignmentError instances) is returned.
tuple - must contain two elements, the first being one of the above types and the second an integer; the first element is used to select a list of matching elements and the second integer element is used to select the corresponding element from the resulting list.
slice - start and stop indices can be any of the above types that selects exactly a single element or None; a corresponding sub-Segment is returned. The step parameter of the slice is ignored. In case the stop marker is not None, the corresponding element is included in the selection.

An element of a segment can be updated by using __setitem__, i.e. segment[item] = ... notation, where item can be any of the above types that selects exactly a single element.

elements¶

Type: list of Element or AlignmentError or Segment

property R¶

property d¶

flat() → particleflow.elements.Segment¶: Convenience function wrapping Segment.flatten.

flatten() → Iterator[Union[particleflow.elements.Element, particleflow.elements.AlignmentError]]¶: Retrieve a flat representation of the segment (with sub-segments flattened as well).

forward(x: <MagicMock name='mock.Tensor' id='140244238396048'>, method='forward', *, aperture: bool = False, exact_drift: bool = True, observe: Union[str, Pattern, Type[Union[Element, AlignmentError]], Sequence[Union[str, Pattern, Type[Union[Element, AlignmentError]]]], None] = None, recloss: Union[bool, str, Pattern, Type[Union[Element, AlignmentError]], Sequence[Union[str, Pattern, Type[Union[Element, AlignmentError]]]], None] = None, loss_func: Optional[Callable[[<MagicMock name='mock.Tensor' id='140244238396048'>], <MagicMock name='mock.Tensor' id='140244238396048'>]] = None) → Union[<MagicMock name='mock.Tensor' id='140244238396048'>, Tuple]¶

Track the given particles through the segment.

Parameters

x (torch.Tensor) – Shape (6, N) where N is the number of particles.
method (str) – Method name which will be used for the lattice elements to perform tracking.
aperture (bool) – Determines whether aperture checks are performed (and thus particles marked lost / excluded from tracking).
exact_drift (bool) – If true (default) then Drift`s will always be tracked through via `exact, no matter what method is.
observe (sequence of {str or re.Pattern or subclass of Element}) – Indicates relevant observation points; the (intermediary) positions at these places will be returned. Items are matched against element labels (see function match_element).
recloss (bool or “sum” or {str or re.Pattern or subclass of Element} or a sequence thereof) – If “sum” then the particle loss at each element will be recorded and summed into a single variable which will be returned. If a sequence is given it must contain element labels or regex patterns and the loss will be recorded only at the corresponding elements (similar to observe for positions). If True then the loss will be recorded at all elements; if False the loss is not recorded at all. A true value for this parameter will automatically set aperture to True.
loss_func (callable) – This parameter can be used to supply a function for transforming the returned loss at each element. This can be useful if some variation of the loss is to be accumulated into a single tensor. The function receives the loss tensor as returned by Aperture.loss as an input and should return a tensor of any shape. If not supplied this function defaults to torch.sum if recloss == “accumulate” and the identity if recloss == "history". Note that if a function is supplied it needs to cover all the steps, also the summation in case recloss == "accumulate" (the default only applies if no function is given).

Returns

x (torch.Tensor) – Shape (6, M) where M is the number of particles that reached the end of the segment (M can be different from N in case aperture checks are performed).
history (dict) – The (intermediary) positions at the specified observation points. Keys are element labels and values are positions of shape (6, M_i) where M_i is the number of particles that reached that element. This is only returned if observe is true.
loss (torch.Tensor or dict) – If recloss == "accumulate" the loss value accumulated for each element (i.e. the sum of all individual loss values) is returned. Otherwise, if recloss is true, a dict mapping element labels to recorded loss values is returned.

get_element_index(marker: Union[int, str, Element, AlignmentError, Tuple[Union[str, Pattern, Type[Union[Element, AlignmentError]]], int]]) → int¶

makethin(n: Union[int, Dict[Union[str, Pattern[AnyStr], Type[Element], None], int]], *, style: Union[str, Dict[Union[str, Pattern[AnyStr], Type[Element], None], str], None] = None) → particleflow.elements.Segment¶

Retrieve a thin representation for each element of the segment except for those indicated by exclude.

Parameters

n (int or dict) –
Number of slices (thin elements). If int this applies to all elements. If dict then keys should be either element classes (subclasses of Element) or strings (indicating element names) or regex patterns (indicating element name patterns). Values should be the corresponding number of slices that are applied to a matching element. The first matching criterion is considered. If a default value is desired it can be provided in various ways:
1. Using {None: default_value}.
2. Using a regex that matches any label (given that all elements have a label different from None).
3. Using the class Element.
Note that for options 2. and 3. these should come at the very end of the dict, otherwise they will override any subsequent definitions.
style (str or dict) – Slicing style per element. Works similar to n. See Element.makethin() for more information about available styles.

Returns

Containing thin elements (or sub-segments).

Return type

Segment

See also

find_matching_criterion(): For details about how keys in n can be used to fine-tune the slice number per element.
Element.makethin(): For available slicing styles.

unstack() → Iterator[particleflow.elements.Segment]¶: Yields the segment itself (this method exists for consistency with AlignmentError).