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
import warnings
from dipas.build import from_file
from dipas.elements import Quadrupole, HKicker, SBend
import dipas.test.sequences

warnings.simplefilter('ignore')

with resources.path(dipas.test.sequences, 'hades.seq') as path:
    lattice = from_file(path)

print(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:

print('HKicker', end='\n\n')
for kicker in lattice[HKicker]:
    print(kicker)

print('\nSBend', end='\n\n')
for sbend in lattice[SBend]:
    print(sbend)

HKicker

HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gte2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='ghadkx1')

SBend

Tilt(psi=tensor(0.3795),
     target=SBend(l=tensor(1.4726), angle=tensor(-0.1308), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), h1=tensor(0.), h2=tensor(0.), dk0=tensor(0.), label='ghadmu1')
    > Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None)
    > SBendBody(l=tensor(1.4726), angle=tensor(-0.1308), dk0=tensor(0.), label=None)
    > 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=SBend(l=tensor(1.4726), angle=tensor(-0.1311), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), h1=tensor(0.), h2=tensor(0.), dk0=tensor(0.), label='ghadmu2')
    > Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None)
    > SBendBody(l=tensor(1.4726), angle=tensor(-0.1311), dk0=tensor(0.), label=None)
    > Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None))

Note that the SBends 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:

print(lattice[Quadrupole, 5])
print(lattice[Quadrupole, 5] is lattice[Quadrupole][5])

Quadrupole(l=tensor(1.), k1=Parameter containing: tensor(0.0298, requires_grad=True), dk1=tensor(0.), label='gth1qd11')
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 dipas.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.), dkh=tensor(0.), dkv=tensor(0.), label='gte2kx1')
Tilt(psi=tensor(0.5000),
     target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1'))
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=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:

from dipas.elements import Offset

print('Elements with offset: ', lattice[Offset])

for element in lattice[Tilt]:
    print(element)

Elements with offset:  []
Tilt(psi=tensor(0.5000),
     target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1'))
Tilt(psi=tensor(0.3795),
     target=SBend(l=tensor(1.4726), angle=tensor(-0.1308), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), h1=tensor(0.), h2=tensor(0.), dk0=tensor(0.), label='ghadmu1')
    > Dipedge(l=tensor(0.), h=tensor(-0.0888), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None)
    > SBendBody(l=tensor(1.4726), angle=tensor(-0.1308), dk0=tensor(0.), label=None)
    > 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=SBend(l=tensor(1.4726), angle=tensor(-0.1311), e1=tensor(0.), e2=tensor(0.), fint=tensor(0.), fintx=tensor(0.), hgap=tensor(0.), h1=tensor(0.), h2=tensor(0.), dk0=tensor(0.), label='ghadmu2')
    > Dipedge(l=tensor(0.), h=tensor(-0.0891), e1=tensor(0.), fint=tensor(0.), hgap=tensor(0.), label=None)
    > SBendBody(l=tensor(1.4726), angle=tensor(-0.1311), dk0=tensor(0.), label=None)
    > 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 SBends from before and the HKicker that we tilted manually.

We can also select elements by their label:

print(lattice['gth1kx1'])
print(lattice['gth1kx1'] is lattice[HKicker, 1])

Tilt(psi=tensor(0.5000),
     target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1'))
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 dipas.elements import Drift

print(lattice[Drift, 0].label)
lattice[Drift, 1].label = 'pad_drift_0'
print(lattice['pad_drift_0'])
print(lattice['pad_drift_0', 1])

pad_drift_0
[Drift(l=tensor(3.9057), label='pad_drift_0'), Drift(l=tensor(0.8420), label='pad_drift_0')]
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.), dkh=tensor(0.), dkv=tensor(0.), label='gte2kx1')
Tilt(psi=tensor(0.5000),
     target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1'))
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth2kx1')
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=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:

print(lattice[4])  # Selecting the 5-th element.
print(lattice[19])  # Selecting the 20-th element.

Drift(l=tensor(0.3370), label='pad_drift_2')
Drift(l=tensor(0.5609), label='pad_drift_10')

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)).

print(lattice[:6], end='\n\n')
print(lattice[(Drift, 1):'gte1dg1'])

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

Segment(elements=[Drift(l=tensor(0.8420), label='pad_drift_0'),
 Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), dk1=tensor(0.), label='gte1qd11'),
 Drift(l=tensor(0.3370), label='pad_drift_2'),
 Monitor(l=tensor(0.), label='gte1dg1')])