Inspecting lattices

Once we have a lattice in form of a elements.Segment instance, such as returned from build.from_file we can inspect its elements in various ways:

>>> from importlib import resources
>>> from particleflow.build import from_file
>>> from particleflow.elements import Quadrupole, HKicker, SBend, Offset
>>> import particleflow.test.sequences
>>>
>>> with resources.path(particleflow.test.sequences, 'hades.seq') as path:
...     lattice = from_file(path)
...
>>> len(lattice[Quadrupole])
21

Here lattice[Quadrupole] returns a list containing all quadrupoles in the lattice. This can be done with any lattice element class:

>>> for kicker in lattice[HKicker]:
...     print(kicker)
...
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gte2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='ghadkx1')
>>>
>>> for sbend in lattice[SBend]:
...     print(sbend)
...
Tilt(psi=tensor(0.3795), target=Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None) | SBend(l=tensor(1.4726), angle=tensor(-0.1308), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), label='ghadmu1') | Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None))
Tilt(psi=tensor(-0.3795), target=Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None) | SBend(l=tensor(1.4726), angle=tensor(-0.1311), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), label='ghadmu2') | Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None))

Note that the SBend`s are tilted which is indicated by the wrapping `Tilt object. Also the two dipole edge elements are reported in terms of Dipedge elements.

We can select a specific element from a multi-element selection directly by providing a tuple:

>>> lattice[Quadrupole, 5]
Quadrupole(l=tensor(1.), k1=Parameter containing: tensor(0.0298, requires_grad=True), label='gth1qd11')
>>> lattice[Quadrupole, 5] is lattice[Quadrupole][5]
True

As shown, the same result can of course be obtained by indexing the resulting list of multiple elements. One case where tuples are the only way however is if we want to set a specific element in the sequence. For example if we want to tilt the second HKicker then we can do:

>>> from particleflow.elements import Tilt
>>>
>>> lattice[HKicker, 1] = Tilt(lattice[HKicker][1], psi=0.5)
>>> for kicker in lattice[HKicker]:
...     print(kicker)
...
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gte2kx1')
Tilt(psi=tensor(0.5000), target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1'))
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='ghadkx1')

Note that we specified HKicker, 1 because indices start at zero. Selections also work with modifiers such as Tilt or alignment errors such as Offset:

>>> lattice[Offset]
[]
>>> for element in lattice[Tilt]:
...     print(element)
...
Tilt(psi=tensor(0.5000), target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1'))
Tilt(psi=tensor(0.3795), target=Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None) | SBend(l=tensor(1.4726), angle=tensor(-0.1308), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), label='ghadmu1') | Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None))
Tilt(psi=tensor(-0.3795), target=Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None) | SBend(l=tensor(1.4726), angle=tensor(-0.1311), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), label='ghadmu2') | Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None))

There are no offset elements in the lattice but as we see there are three tilted elements: the two SBend`s from before and the `HKicker that we tilted manually.

We can also select elements by their label:

>>> lattice['gth1kx1']
Tilt(psi=tensor(0.5000), target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1'))
>>> lattice['gth1kx1'] is lattice[HKicker, 1]
True

If there are multiple elements that share a label, a list will be returned instead. Again we can use a tuple index to select a specific element:

>>> from particleflow.elements import Drift
>>>
>>> lattice[Drift, 0].label
'pad_drift_0'
>>> lattice[Drift, 1].label = 'pad_drift_0'
>>> lattice['pad_drift_0']
[Drift(l=tensor(3.9057), label='pad_drift_0'), Drift(l=tensor(0.8420), label='pad_drift_0')]
>>> lattice['pad_drift_0', 1]
Drift(l=tensor(0.8420), label='pad_drift_0')

By using regular expression patterns we can select all elements whose labels match the specified pattern:

>>> import re
>>>
>>> pattern = re.compile(r'[a-z0-9]+kx1')
>>> for element in lattice[pattern]:
...     print(element)
...
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gte2kx1')
Tilt(psi=tensor(0.5000), target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1'))
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='ghadkx1')

Here we need to use a compiled re object because strings will be interpreted as element labels, not patterns. An exception is if the string contains an asterisk * which will be interpreted as a shell-style wildcard pattern (internally it is converted to a regex while replacing * with .*?). Thus using lattice['*kx1'] selects all HKicker elements as before.

Last but not least we can select elements by their index position along the lattice:

>>> lattice[4]  # Selecting the 5-th element.
Drift(l=tensor(0.3370), label='pad_drift_2')
>>> lattice[19]  # Selecting the 20-th element.
Monitor(l=tensor(0.), label='gte2dg4')

Sub-segments can be selected by using slice syntax. Here the start and stop parameters must be unambiguous element identifiers, as described above (e.g. unique labels, or a multi-selector such as Quadrupole with an occurrence count, i.e. (Quadrupole, 5)).

>>> lattice[:6]
Segment(elements=[Drift(l=tensor(3.9057), label='pad_drift_0'),
 VKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gte1ky1'),
 Drift(l=tensor(0.8420), label='pad_drift_1'),
 Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), label='gte1qd11'),
 Drift(l=tensor(0.3370), label='pad_drift_2'),
 Monitor(l=tensor(0.), label='gte1dg1')])
>>>
>>> lattice[(Drift, 1):'gte1dg1']
Segment(elements=[Drift(l=tensor(0.8420), label='pad_drift_1'),
 Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), label='gte1qd11'),
 Drift(l=tensor(0.3370), label='pad_drift_2'),
 Monitor(l=tensor(0.), label='gte1dg1')])