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 dipas.build import from_file, from_script

lattice = from_file('example.madx')

with open('example.madx') as fh:  # alternatively build from script string
    lattice = from_script(fh.read())

The documentation of the dipas.madx.parser module contains detailed information on how to customize the parsing behavior.

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 are the corresponding PyTorch backend modules.

[1]:
from pprint import pprint
from dipas.elements import elements

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

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

[2]:
from dipas.elements import alignment_errors, aperture_types

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

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

[3]:
from importlib import resources
from dipas.build import from_file
import dipas.test.sequences

with resources.path(dipas.test.sequences, 'hades.seq') as path:
    lattice = from_file(path)
/home/dominik/Projects/DiPAS/dipas/build.py:582: UnknownElementTypeWarning: Unknown element type got replaced with Drift: Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 8.6437999}, label='gts1mu1', base=None, line_number=50)
  warnings.warn(f'Unknown element type got replaced with Drift: {command}', category=UnknownElementTypeWarning)
/home/dominik/Projects/DiPAS/dipas/elements.py:468: UnknownParametersWarning: Unknown parameters for element of type <class 'dipas.elements.Drift'>: {'knl': tensor([0.])}
  warnings.warn(f'Unknown parameters for element of type {type(self)}: {kwargs}', category=UnknownParametersWarning)
/home/dominik/Projects/DiPAS/dipas/build.py:582: UnknownElementTypeWarning: Unknown element type got replaced with Drift: Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 28.6437973}, label='gte3mu1', base=None, line_number=57)
  warnings.warn(f'Unknown element type got replaced with Drift: {command}', category=UnknownElementTypeWarning)
/home/dominik/Projects/DiPAS/dipas/build.py:582: UnknownElementTypeWarning: Unknown element type got replaced with Drift: Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 52.4014301}, label='ghhtmu1', base=None, line_number=64)
  warnings.warn(f'Unknown element type got replaced with Drift: {command}', category=UnknownElementTypeWarning)
/home/dominik/Projects/DiPAS/dipas/build.py:582: UnknownElementTypeWarning: Unknown element type got replaced with Drift: Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 100.2795473}, label='gth3mu1', base=None, line_number=75)
  warnings.warn(f'Unknown element type got replaced with Drift: {command}', category=UnknownElementTypeWarning)
/home/dominik/Projects/DiPAS/dipas/build.py:582: UnknownElementTypeWarning: Unknown element type got replaced with Drift: Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 125.7672306}, label='gtp1mu1', base=None, line_number=86)
  warnings.warn(f'Unknown element type got replaced with Drift: {command}', category=UnknownElementTypeWarning)

This issues a few warnings of the following form:

.../dipas/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(dipas.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:

[4]:
from dipas.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:

[5]:
print(lattice)
[    0.000000]  Drift(l=tensor(2.), label='e1')
[    2.000000]  Quadrupole(l=tensor(1.), k1=tensor(0.2500), dk1=tensor(0.), label='q1')
[    3.000000]  Drift(l=tensor(3.), label='e3')
[    6.000000]  HKicker(l=tensor(0.), hkick=tensor(0.1000), vkick=tensor(0.), kick=tensor(0.1000), dkh=tensor(0.), dkv=tensor(0.), 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:

[6]:
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:

[7]:
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 build.from_file or build.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 our lattice to a Segment as follows:

[8]:
from dipas.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:

[9]:
from dipas.build import Beam
import dipas.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:

[10]:
from pprint import pprint

pprint(beam)
{'beta': 0.6,
 'brho': 2.3473041010257827,
 'charge': 1,
 'energy': 1.1728401102,
 'gamma': 1.25,
 'mass': 0.93827208816,
 'particle': 'proton',
 'pc': 0.7037040661199998}

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