"""Load a lattice from an Elegant file (.lte).
This is not complete but can parse the example files that I have.
This parser is quite similar to the Tracy parser in tracy.py.
The Elegant file format is described briefly `here
<https://ops.aps.anl.gov/manuals/elegant_latest/elegantse9.html#x113-1120009>`_.
It is similar to the `MAD-X format <http://madx.web.cern.ch/madx/>`_.
Note that Elegant scales magnet polynomials in a different way
to AT, so the parsed coefficients need to be divided by n! for
the coefficient of order n.
"""
import logging as log
from os.path import abspath
import re
from at.lattice.elements import (
Aperture,
Corrector,
Drift,
Dipole,
Marker,
Multipole,
Octupole,
Quadrupole,
RFCavity,
Sextupole,
)
from at.lattice import Lattice
from at.load import register_format, utils
__all__ = ['load_elegant']
# noinspection PyUnusedLocal
def create_drift(name, params, energy, harmonic_number):
length = params.pop("l", 0)
return Drift(name, length, **params)
# noinspection PyUnusedLocal
def create_marker(name, params, energy, harmonic_number):
return Marker(name, **params)
# noinspection PyUnusedLocal
def create_aperture(name, params, energy, harmonic_number):
x_lim = float(params.get('x_max'))
y_lim = float(params.get('y_max'))
limits = [-x_lim, x_lim, -y_lim, y_lim]
return Aperture(name, limits)
# noinspection PyUnusedLocal
def create_quad(name, params, energy, harmonic_number):
length = params.pop("l", 0)
params["NumIntSteps"] = params.pop("n_kicks", 10)
params["k"] = float(params.pop("k1"))
params["PassMethod"] = "StrMPoleSymplectic4Pass"
return Quadrupole(name, length, **params)
# noinspection PyUnusedLocal
def create_sext(name, params, energy, harmonic_number):
length = params.pop("l", 0)
params["NumIntSteps"] = params.pop("n_kicks", 10)
k2 = float(params.pop("k2", 0))
return Sextupole(name, length, k2 / 2, **params)
# noinspection PyUnusedLocal
def create_oct(name, params, energy, harmonic_number):
length = params.pop("l", 0)
params["NumIntSteps"] = params.pop("n_kicks", 10)
k3 = float(params.pop("k3", 0))
PolynomA = [0, 0, 0, 0]
PolynomB = [0, 0, 0, k3 / 6]
return Octupole(name, length, PolynomA, PolynomB, **params)
# noinspection PyUnusedLocal
def create_multipole(name, params, energy, harmonic_number):
def factorial(x, acc=1):
if x == 0:
return acc
else:
return factorial(x - 1, acc * x)
length = params.pop("l", 0)
PolynomA = [0, 0, 0, 0]
PolynomB = [0, 0, 0, 0]
if "knl" in params:
order = int(float(params.pop("order")))
PolynomB[order] = float(params.pop("knl")) / factorial(order)
if "ksl" in params:
order = int(float(params.pop("order")))
PolynomA[order] = float(params.pop("ksl")) / factorial(order)
return Multipole(name, length, PolynomA, PolynomB, **params)
# noinspection PyUnusedLocal
def create_dipole(name, params, energy, harmonic_number):
length = params.pop("l", 0)
params["NumIntSteps"] = params.pop("n_kicks", 10)
params["PassMethod"] = "BndMPoleSymplectic4Pass"
params["BendingAngle"] = float(params.pop("angle"))
params["EntranceAngle"] = float(params.pop("e1"))
params["ExitAngle"] = float(params.pop("e2"))
if "hgap" in params:
params["FullGap"] = float(params.pop("hgap")) * 2
# What should we do if no fint property?
fint = params.pop("fint", 1)
params["FringeInt1"] = fint
params["FringeInt2"] = fint
k1 = float(params.pop("k1", 0))
k2 = float(params.pop("k2", 0))
k3 = float(params.pop("k3", 0))
k4 = float(params.pop("k4", 0))
params["PolynomB"] = [0, k1, k2 / 2, k3 / 6, k4 / 24]
return Dipole(name, length, **params)
# noinspection PyUnusedLocal
def create_corrector(name, params, energy, harmonic_number):
length = params.pop("l", 0)
hkick = params.pop("hkick", 0)
vkick = params.pop("vkick", 0)
kick_angle = [hkick, vkick]
return Corrector(name, length, kick_angle, **params)
def create_cavity(name, params, energy, harmonic_number):
length = params.pop("l", 0)
voltage = params.pop("volt")
frequency = params.pop("freq")
params["Phi"] = params.pop("phase")
return RFCavity(name, length, voltage,
frequency, harmonic_number, energy, **params)
ELEMENT_MAP = {
"drift": create_drift,
"drif": create_drift,
"edrift": create_drift,
# This should be a wiggler.
"cwiggler": create_drift,
"csben": create_dipole,
"csbend": create_dipole,
"csrcsben": create_dipole,
"quadrupole": create_quad,
"kquad": create_quad,
"ksext": create_sext,
"koct": create_oct,
"mult": create_multipole,
"multipole": create_multipole,
"kicker": create_corrector,
"hkick": create_corrector,
"vkick": create_corrector,
"mark": create_marker,
"marker": create_marker,
"malign": create_marker,
"recirc": create_marker,
"sreffects": create_marker,
"rcol": create_marker,
"watch": create_marker,
"charge": create_marker,
"monitor": create_marker,
"moni": create_marker,
"maxamp": create_aperture,
"rfca": create_cavity,
}
def parse_lines(contents):
"""Return individual lines.
Remove comments and whitespace and convert to lowercase.
"""
lines = []
current_line = ""
for line in contents.splitlines():
line = line.lower()
line = line.split("!")[0]
line = line.strip()
if line:
if line.endswith("&"):
current_line += line[:-1]
else:
lines.append(current_line + line)
current_line = ""
return lines
def parse_chunk(value, elements, chunks):
"""Parse a non-element part of a lattice file.
That part can reference already-parsed elements and chunks.
if
chunks = {"x": ["a", "b"]}
and
elements = {"a": Marker(...), "b": Quadrupole(...)}
line(a,b) => [a, b]
-x => [b, a]
2*x => [a, b, a, b]
"""
chunk = []
parts = utils.split_ignoring_parentheses(value, ",")
for part in parts:
# What's this?
if "symmetry" in part:
continue
if "line" in part:
line_parts = re.match("line\\s*=\\s*\\((.*)\\)", part).groups()[0]
for p in line_parts.split(","):
p = p.strip()
chunk.extend(parse_chunk(p, elements, chunks))
elif part.startswith("-"):
# Reverse a sequence of elements. When doing this, swap
# the entrance and exit angles for any dipoles.
chunk_name = part[1:]
try:
chunk_to_invert = chunks[chunk_name]
except KeyError:
# You can reverse just a single element: probably a
# bend, as there's no reason to invert other elements.
chunk_to_invert = [elements[chunk_name]]
inverted_chunk = []
for el in reversed(chunk_to_invert):
if el.__class__ == Dipole:
inverted_dipole = el.copy()
inverted_dipole.EntranceAngle = el.ExitAngle
inverted_dipole.ExitAngle = el.EntranceAngle
inverted_chunk.append(inverted_dipole)
else:
inverted_chunk.append(el)
chunk.extend(inverted_chunk)
elif "*" in part:
num, chunk_name = part.split("*")
if chunk_name[0] == "(":
assert chunk_name[-1] == ")"
chunk_name = chunk_name[1:-1]
chunk.extend(int(num) * parse_chunk(chunk_name, elements, chunks))
elif part in elements:
chunk.append(elements[part])
elif part in chunks:
chunk.extend(chunks[part])
else:
raise Exception("part {} not understood".format(part))
return chunk
def expand_elegant(contents, lattice_key, energy, harmonic_number):
lines = parse_lines(contents)
variables = {"energy": energy, "harmonic_number": harmonic_number}
elements = {}
chunks = {}
for line in lines:
if ":" not in line:
key, value = line.split("=")
variables[key.strip()] = value.strip()
else:
key, value = line.split(":")
key = key.strip()
value = value.strip()
if value.split(",")[0] in ELEMENT_MAP:
elements[key] = elegant_element_from_string(key,
value,
variables)
else:
chunk = parse_chunk(value, elements, chunks)
chunks[key] = chunk
return chunks[lattice_key.lower()]
def handle_value(value):
value = value.strip()
if value.startswith('"'):
assert value.endswith('"')
value = value[1:-1]
value = value.split()
if len(value) > 1:
# Handle basic arithmetic e.g. "0.04 2 /" -> 0.02
assert len(value) == 3
if value[2] == "/":
value = float(value[0]) / float(value[1])
elif value[2] == "*":
value = float(value[0]) * float(value[1])
elif value[2] == "-":
value = float(value[0]) - float(value[1])
elif value[2] == "+":
value = float(value[0]) + float(value[1])
else:
value = value[0]
return value
def elegant_element_from_string(name, element_string, variables):
"""Create element from Elegant's string representation.
e.g. drift,l=0.045 => Drift(name, 0.045)
"""
log.debug("Parsing elegant element {}".format(element_string))
parts = utils.split_ignoring_parentheses(element_string, ",")
params = {}
element_type = parts[0]
for part in parts[1:]:
key, value = utils.split_ignoring_parentheses(part, "=")
key = key.strip()
value = handle_value(value)
if value in variables:
value = variables[value]
params[key] = value
energy = variables["energy"]
harmonic_number = variables["harmonic_number"]
return ELEMENT_MAP[element_type](name, params, energy, harmonic_number)
[docs]
def load_elegant(filename: str, **kwargs) -> Lattice:
"""Create a :py:class:`.Lattice` from an Elegant file
Parameters:
filename: Name of an Elegant file
Keyword Args:
name (str): Name of the lattice. Default: taken from
the file.
energy (float): Energy of the lattice [eV]
periodicity(int): Number of periods. Default: taken from the elements, or 1
*: All other keywords will be set as Lattice attributes
Returns:
lattice (Lattice): New :py:class:`.Lattice` object
See Also:
:py:func:`.load_lattice` for a generic lattice-loading function.
"""
try:
energy = kwargs.pop("energy")
lattice_key = kwargs.pop("lattice_key")
harmonic_number = kwargs.pop("harmonic_number")
def elem_iterator(params, elegant_file):
with open(params.setdefault("in_file", elegant_file)) as f:
contents = f.read()
element_lines = expand_elegant(
contents, lattice_key, energy, harmonic_number
)
params.setdefault("energy", energy)
for line in element_lines:
yield line
lat = Lattice(abspath(filename), iterator=elem_iterator, **kwargs)
return lat
except Exception as e:
raise ValueError('Failed to load elegant '
'lattice {}: {}'.format(filename, e))
register_format(
".lte", load_elegant, descr="Elegant format. See :py:func:`.load_elegant`.")
