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())
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 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:560: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 8.6437999}, label='gts1mu1', base=None)
warnings.warn(f'Skipping element (no equivalent implementation found): {command}')
/home/dominik/Projects/DiPAS/dipas/build.py:560: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 28.6437973}, label='gte3mu1', base=None)
warnings.warn(f'Skipping element (no equivalent implementation found): {command}')
/home/dominik/Projects/DiPAS/dipas/build.py:560: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 52.4014301}, label='ghhtmu1', base=None)
warnings.warn(f'Skipping element (no equivalent implementation found): {command}')
/home/dominik/Projects/DiPAS/dipas/build.py:560: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 100.2795473}, label='gth3mu1', base=None)
warnings.warn(f'Skipping element (no equivalent implementation found): {command}')
/home/dominik/Projects/DiPAS/dipas/build.py:560: UserWarning: Skipping element (no equivalent implementation found): Command(keyword='multipole', local_attributes={'knl': array([0.]), 'at': 125.7672306}, label='gtp1mu1', base=None)
warnings.warn(f'Skipping element (no equivalent implementation found): {command}')
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=None)
[ 2.000000] Quadrupole(l=tensor(1.), k1=tensor(0.2500), dk1=tensor(0.), 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:
[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.34730408386391,
'charge': 1,
'energy': 1.1728401016249999,
'gamma': 1.25,
'mass': 0.9382720813,
'particle': 'proton',
'pc': 0.7037040609749998}
Taking care of the beam definition was done automatically by using the build.Lattice
class as in the previous section.