Building lattices

Lattices can be built by parsing MADX scripts or programmatically using the API of the package.

Parsing MADX scripts

The main functions for parsing MADX scripts to lattices are build.from_file and build.from_script. The only difference is that the former expects the file name to the script and the latter the raw script as a string:

>>> from particleflow.build import from_file, from_script
>>>
>>> lattice = from_file('example.madx')
>>>
>>> with open('example.madx') as fh:  # alternatively
...     lattice = from_script(fh.read())
...

In case the MADX script contains an unknown element, a warning will be issued and the element is skipped. The supported elements can be found by inspecting the elements.elements dict; keys are MADX command names and values the corresponding PyTorch backend modules.

>>> from pprint import pprint
>>> from particleflow.elements import elements
>>>
>>> pprint(elements)
{'dipedge': <class 'particleflow.elements.Dipedge'>,
 'drift': <class 'particleflow.elements.Drift'>,
 'hkicker': <class 'particleflow.elements.HKicker'>,
 'hmonitor': <class 'particleflow.elements.HMonitor'>,
 'instrument': <class 'particleflow.elements.Instrument'>,
 'kicker': <class 'particleflow.elements.Kicker'>,
 'marker': <class 'particleflow.elements.Marker'>,
 'monitor': <class 'particleflow.elements.Monitor'>,
 'placeholder': <class 'particleflow.elements.Placeholder'>,
 'quadrupole': <class 'particleflow.elements.Quadrupole'>,
 'rbend': <class 'particleflow.elements.RBend'>,
 'sbend': <class 'particleflow.elements.SBend'>,
 'sextupole': <class 'particleflow.elements.Sextupole'>,
 'tkicker': <class 'particleflow.elements.TKicker'>,
 'vkicker': <class 'particleflow.elements.VKicker'>,
 'vmonitor': <class 'particleflow.elements.VMonitor'>}

Similarly we can check the supported alignment errors and aperture types:

>>> from particleflow.elements import alignment_errors, aperture_types
>>>
>>> pprint(alignment_errors)
{'dpsi': <class 'particleflow.elements.LongitudinalRoll'>,
 'dx': <class 'particleflow.elements.Offset'>,
 'dy': <class 'particleflow.elements.Offset'>,
 'mrex': <class 'particleflow.elements.BPMError'>,
 'mrey': <class 'particleflow.elements.BPMError'>,
 'mscalx': <class 'particleflow.elements.BPMError'>,
 'mscaly': <class 'particleflow.elements.BPMError'>,
 'tilt': <class 'particleflow.elements.Tilt'>}
>>>
>>> pprint(aperture_types)
{'circle': <class 'particleflow.elements.ApertureCircle'>,
 'ellipse': <class 'particleflow.elements.ApertureEllipse'>,
 'rectangle': <class 'particleflow.elements.ApertureRectangle'>,
 'rectellipse': <class 'particleflow.elements.ApertureRectEllipse'>}

As seen from the above script, a general MULTIPOLE is not yet supported and so attempting to load a script with such a definition will raise a warning:

>>> from importlib import resources
>>> from particleflow.build import from_file
>>> import particleflow.test.sequences
>>>
>>> with resources.path(particleflow.test.sequences, 'hades.seq') as path:
...     lattice = from_file(path)
...

This will issue a few warnings of the following form:

.../particleflow/build.py:174: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 8.6437999}, label='gts1mu1', base=None)

In order to not accidentally miss any such non-supported elements one can configure Python to raise an error whenever a warning is encountered (see the docs for more details):

>>> import warnings
>>>
>>> warnings.simplefilter('error')
>>>
>>> with resources.path(particleflow.test.sequences, 'hades.seq') as path:
...     lattice = from_file(path)
...

This will convert the previous warning into an error.

Using the build API

We can also build a lattice using the build.Lattice class:

>>> from particleflow.build import Lattice
>>>
>>> with Lattice(beam=dict(particle='proton', beta=0.6)) as lattice:
...     lattice.Drift(l=2)
...     lattice.Quadrupole(k1=0.25, l=1, label='q1')
...     lattice.Drift(l=3)
...     lattice.HKicker(kick=0.1, label='hk1')
...

When used as a context manager (i.e. inside with) we just need to invoke the various element functions in order to append them to the lattice.

We can get an overview of the lattice by printing it:

>>> print(lattice)
[    0.000000]  Drift(l=tensor(2.), label=None)
[    2.000000]  Quadrupole(l=tensor(1.), k1=tensor(0.2500), label='q1')
[    3.000000]  Drift(l=tensor(3.), label=None)
[    6.000000]  HKicker(l=tensor(0.), hkick=tensor(0.1000), vkick=tensor(0.), kick=tensor(0.1000), label='hk1')

The number in brackets [...] indicates the position along the lattice in meters, followed by a description of the element

Besides usage as a context manager other ways of adding elements exist:

>>> lattice = Lattice({'particle': 'proton', 'beta': 0.6})
>>> lattice += lattice.Drift(l=2)
>>> lattice.append(lattice.Quadrupole(k1=0.25, l=1, label='q1'))
>>> lattice += [lattice.Drift(l=3), lattice.HKicker(kick=0.1, label='hk1')]

This creates the same lattice as before. Note that because lattice is not used as a context manager, invoking the element functions, such as lattice.Quadrupole, will not automatically add the element to the lattice; we can do so via lattice += ..., lattice.append or lattice.extend.

We can also specify positions along the lattice directly, which will also take care of inserting implicit drift spaces:

>>> lattice = Lattice({'particle': 'proton', 'beta': 0.6})
>>> lattice[2.0] = lattice.Quadrupole(k1=0.25, l=1, label='q1')
>>> lattice['q1', 3.0] = lattice.HKicker(kick=0.1, label='hk1')

This again creates the same lattice as before. We can specify an absolute position along the lattice by just using a float or we can specify a position relative to another element by using a tuple and referring to the other element via its label.

Note

When using a relative position via tuple, the position is taken relative to the exit of the referred element.

After building the lattice in such a way there’s one step left to obtain the same result as via from_file or from_script. These methods return a elements.Segment instance which provides further functionality for tracking and conversion to thin elements for example. We can simply convert out lattice to a Segment as follows:

>>> from particleflow.elements import Segment
>>> lattice = Segment(lattice)  # `lattice` from before

Using the element types directly

Another option for building a lattice is to access the element classes directly. This can be done via elements.<cls_name> or by using the elements.elements dict which maps MADX command names to corresponding backend classes.

>>> from particleflow.build import Beam
>>> import particleflow.elements as elements
>>>
>>> beam = Beam(particle='proton', beta=0.6).to_dict()
>>> sequence = [
...     elements.Drift(l=2, beam=beam),
...     elements.Quadrupole(k1=0.25, l=1, beam=beam, label='q1'),
...     elements.elements['drift'](l=3, beam=beam),
...     elements.elements['hkicker'](kick=0.1, label='hk1')
... ]
>>> lattice = elements.Segment(sequence)

This creates the same lattice as in the previous section. Note that we had to use Beam(...).to_dict() and pass the result to the element classes. This is because the elements expect both beta and gamma in the beam dict and won’t compute it themselves. build.Beam however does the job for us:

>>> from pprint import pprint
>>> pprint(beam)
{'beta': 0.6,
 'brho': 2.34730408386391,
 'charge': 1,
 'energy': 1.1728401016249999,
 'gamma': 1.25,
 'mass': 0.9382720813,
 'particle': 'proton',
 'pc': 0.7037040609749998}

This work of taking care of the beam definition was done automatically by using the build.Lattice class as in the previous section.