"""Load a lattice from a Tracy file (.lat).
This is not complete but can parse the example files that I have.
This parser is quite similar to the Elegant parser in elegant.py.
"""
import logging as log
from os.path import abspath
import re
import numpy
from at.lattice.elements import (
Drift,
Dipole,
Quadrupole,
Sextupole,
Multipole,
Corrector,
Marker,
RFCavity,
)
from at.lattice import Lattice
from at.load import register_format, utils
__all__ = ["load_tracy"]
def create_drift(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
return Drift(name, length, **params)
# noinspection PyUnusedLocal
def create_marker(name, params, variables):
return Marker(name, **params)
def create_quad(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
params["NumIntSteps"] = parse_float(params.pop("n", 10), variables)
params["k"] = parse_float(params.pop("k"), variables)
params["PassMethod"] = "StrMPoleSymplectic4Pass"
return Quadrupole(name, length, **params)
def create_sext(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
params["NumIntSteps"] = parse_float(params.pop("n", 10), variables)
k2 = parse_float(params.pop("k", 0), variables)
return Sextupole(name, length, k2, **params)
def create_dipole(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
params["NumIntSteps"] = parse_float(params.pop("n", 10), variables)
params["PassMethod"] = "BndMPoleSymplectic4Pass"
params["BendingAngle"] = (parse_float(params.pop("t"), variables) / 180) * numpy.pi
params["EntranceAngle"] = (
parse_float(params.pop("t1"), variables) / 180
) * numpy.pi
params["ExitAngle"] = (parse_float(params.pop("t2"), variables) / 180) * numpy.pi
# Tracy is encoding gap plus fringe int in the 'gap' field.
# Since BndMPoleSymplectic4Pass only uses the product of FringeInt
# and gap we can substitute the following.
if "gap" in params:
params["FullGap"] = parse_float(params.pop("gap"), variables) / 2
params["FringeInt1"] = 1
params["FringeInt2"] = 1
k = parse_float(params.pop("k", 0), variables)
if "PolynomB" in params:
params["PolynomB"][1] = k
else:
params["PolynomB"] = [0, k, 0, 0]
return Dipole(name, length, **params)
def create_corrector(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
# Is there a property for this?
kick_angle = [0, 0]
return Corrector(name, length, kick_angle, **params)
def create_multipole(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
params["NumIntSteps"] = parse_float(params.pop("n", 10), variables)
poly_a = [0, 0, 0, 0]
poly_b = [0, 0, 0, 0]
return Multipole(name, length, poly_a, poly_b, **params)
def create_cavity(name, params, variables):
length = parse_float(params.pop("l", 0), variables)
voltage = parse_float(params.pop("voltage"), variables)
frequency = parse_float(params.pop("frequency"), variables)
params["Phi"] = parse_float(params.pop("phi", 0), variables)
harmonic_number = variables["harmonic_number"]
energy = variables["energy"]
return RFCavity(name, length, voltage, frequency, harmonic_number, energy, **params)
ELEMENT_MAP = {
"drift": create_drift,
"bending": create_dipole,
"quadrupole": create_quad,
"sextupole": create_sext,
"multipole": create_multipole,
"corrector": create_corrector,
"marker": create_marker,
"map": create_marker,
"beampositionmonitor": create_marker,
"cavity": create_cavity,
}
def tokenise_expression(expression):
tokens = []
current_token = ""
for char in expression:
if char.isspace(): # complete the current token
if current_token:
tokens.append(current_token)
current_token = ""
continue
if char in "+-" and current_token and current_token[-1] in "de":
# special case for 1e-4, 1.23e+3, 2.2d6:
# + or - are part of the current token, not their own token.
try:
# If d or e refers to a number, it must be possible
# to convert the preceding characters into a float.
float(current_token[:-1])
current_token += char
continue
except ValueError:
pass
if char in "+-*/()":
# standalone tokens: complete the current token and add this
if current_token:
tokens.append(current_token)
current_token = ""
tokens.append(char)
else:
current_token += char
if current_token:
tokens.append(current_token)
return tokens
def parse_float(expression, variables):
"""Evaluate the provided arithmetic expression substituting variables."""
log.debug("parse_float {}".format(expression))
try:
return float(expression)
except ValueError:
raw_tokens = tokenise_expression(expression)
tokens = []
for token in raw_tokens:
if token in variables:
tokens.append(variables[token])
else:
# Handle unusual format 2.2d6 which means 2.2e+6.
dformat_matches = re.match("(-?\\d+.?\\d*)d(-?\\d+)", token)
if dformat_matches:
mantissa, exponent = dformat_matches.groups()
tokens.append("{}e{}".format(mantissa, exponent))
else:
tokens.append(token)
log.debug(tokens)
def evaluate(tokens):
if len(tokens) == 1:
return float(tokens[0])
if len(tokens) == 2:
if tokens[0] == "+":
return float(tokens[1])
elif tokens[0] == "-":
return -float(tokens[1])
# Remove superfluous outer parentheses.
if tokens[0] == "(" and tokens[-1] == ")":
return evaluate(tokens[1:-1])
# First evaluate contents of parentheses.
try:
b1 = tokens.index("(")
b2 = len(tokens) - 1 - tokens[::-1].index(")")
return evaluate(
tokens[:b1] + [evaluate(tokens[b1 + 1 : b2])] + tokens[b2 + 1 :]
)
except ValueError:
# No open parentheses found.
pass
# Evaluate / and * from left to right.
for i, token in enumerate(tokens[:-1]):
if token == "/":
return evaluate(
tokens[: i - 1]
+ [float(tokens[i - 1]) / float(tokens[i + 1])]
+ tokens[i + 2 :]
)
if token == "*":
return evaluate(
tokens[: i - 1]
+ [float(tokens[i - 1]) * float(tokens[i + 1])]
+ tokens[i + 2 :]
)
# Evaluate + and - from left to right.
for i, token in enumerate(tokens[:-1]):
if token == "+":
return evaluate(
tokens[: i - 1]
+ [float(tokens[i - 1]) + float(tokens[i + 1])]
+ tokens[i + 2 :]
)
if token == "-":
return evaluate(
tokens[: i - 1]
+ [float(tokens[i - 1]) - float(tokens[i + 1])]
+ tokens[i + 2 :]
)
return evaluate(tokens)
def parse_lines(contents):
"""Return individual lines.
Remove comments and whitespace, convert to lowercase, and split on
semicolons.
"""
# Nested comments not handled.
in_comment = False
stripped_contents = ""
for char in contents.lower():
if char == "{":
in_comment = True
elif char == "}":
in_comment = False
elif not in_comment and not char.isspace():
stripped_contents += char
return stripped_contents.split(";")
def parse_chunk(value, elements, chunks):
log.debug("Parse_chunk %s", value)
chunk = []
parts = utils.split_ignoring_parentheses(value, ",")
log.debug(parts)
for part in parts:
if "symmetry" in part:
continue
if "inv" in part:
# Reverse a sequence of elements. When doing this, swap
# the entrance and exit angles for any dipoles.
chunk_to_invert = re.match("inv\\((.*)\\)", part).groups()[0]
inverted_chunk = []
for el in reversed(chunks[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, contents = part.split("*")
if contents.startswith("("):
assert contents[-1] == ")"
contents = contents[1:-1]
chunk.extend(int(num) * parse_chunk(contents, 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_tracy(contents, lattice_key, harmonic_number):
lines = parse_lines(contents)
assert lines[0] == "definelattice"
assert lines[-2] == "end"
assert lines[-1] == ""
variables = {"harmonic_number": harmonic_number}
elements = {}
chunks = {}
for line in lines[1:-2]:
if ":" not in line:
key, value = line.split("=")
if key == "energy":
value = parse_float(value, variables) * 1e9
try:
variables[key] = parse_float(value, variables)
except ValueError:
variables[key] = value
else:
key, value = line.split(":")
if value.split(",")[0].strip() in ELEMENT_MAP:
elements[key] = tracy_element_from_string(key, value, variables)
else:
chunk = parse_chunk(value, elements, chunks)
chunks[key] = chunk
return chunks[lattice_key], variables["energy"]
def parse_hom(hom_string, variables):
"""Parse 'hom' string from lattice file. Note that PolynomB
is before PolynomA.
hom(3, 1.2, 3.4) => [0, 0, 3.4], [0, 0, 1.2]
"""
hom_parts = [part.strip() for part in hom_string.split(",")]
assert len(hom_parts) % 3 == 0
max_order = max(int(p) for i, p in enumerate(hom_parts) if i % 3 == 0)
polynom_a = [0] * int(max_order)
polynom_b = [0] * int(max_order)
for i in range(len(hom_parts) // 3):
order = int(hom_parts[i * 3])
# PolynomB before PolynomA.
b = parse_float(hom_parts[i * 3 + 1], variables)
a = parse_float(hom_parts[i * 3 + 2], variables)
polynom_a[order - 1] = a
polynom_b[order - 1] = b
return polynom_a, polynom_b
def tracy_element_from_string(name, element_string, variables):
log.debug("Parsing tracy element {}".format(element_string))
parts = utils.split_ignoring_parentheses(element_string, ",")
params = {}
element_type = parts[0]
for part in parts[1:]:
try:
key, value = utils.split_ignoring_parentheses(part, "=")
except ValueError:
key, value = part, None
if value in variables:
value = variables[value]
if key == "hom":
assert value[0] == "("
assert value[-1] == ")"
polynom_a, polynom_b = parse_hom(value[1:-1], variables)
params["PolynomA"] = polynom_a
params["PolynomB"] = polynom_b
else:
params[key] = value
return ELEMENT_MAP[element_type](name, params, variables)
[docs]
def load_tracy(filename: str, **kwargs) -> Lattice:
"""Create a :py:class:`.Lattice` from a Tracy file
Parameters:
filename: Name of a Tracy file
Keyword Args:
use (str): Name of the variable containing the desired lattice.
Default: ``cell``
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:
def elem_iterator(params, tracy_file):
with open(params.setdefault("in_file", tracy_file)) as f:
contents = f.read()
element_lines, energy = expand_tracy(
contents, lattice_key, harmonic_number
)
params.setdefault("energy", energy)
for line in element_lines:
yield line
if "lattice_key" in kwargs: # process the deprecated 'lattice_key' keyword
kwargs.setdefault("use", kwargs.pop("lattice_key"))
harmonic_number = kwargs.pop("harmonic_number")
lattice_key = kwargs.pop("use", "cell")
return Lattice(abspath(filename), iterator=elem_iterator, **kwargs)
except Exception as e:
raise ValueError("Failed to load tracy " "lattice {}: {}".format(filename, e))
register_format(".lat", load_tracy, descr="Tracy format. See :py:func:`.load_tracy`.")
