r"""Using `MAD8`_ files with PyAT.
=================================
Using MAD8 files is similar to using MAD-X files: see
:py:mod:`the MAD-X documentation <at.load.madx>`
.. _mad8: https://mad8.web.cern.ch/user/mad.html
"""
from __future__ import annotations
__all__ = ["Mad8Parser", "load_mad8", "save_mad8"]
from pathlib import Path
from typing import ClassVar
# functions known by MAD-8
from math import pi, e, sqrt, exp, log, sin, cos, tan
from math import asin, atan2
import re
# constants known by MAD-8
from scipy.constants import c as clight, e as qelect
from scipy.constants import physical_constants as _cst
from ..lattice import Lattice, elements as elt
from .file_input import ignore_class
# noinspection PyProtectedMember
from .madx import _MadElement, _MadParser, _MadExporter
from .madx import mad_element, p_to_at, poly_to_mad
# Commands known by MAD8
# noinspection PyProtectedMember
from .madx import (
drift,
marker,
quadrupole,
sextupole,
octupole,
longmultipole,
sbend,
rbend,
kicker,
hkicker,
vkicker,
rfcavity,
monitor,
hmonitor,
vmonitor,
_Line,
_Sequence,
_Value,
)
twopi = 2 * pi
degrad = 180.0 / pi
raddeg = pi / 180.0
emass = 1.0e-03 * _cst["electron mass energy equivalent in MeV"][0] # [GeV]
pmass = 1.0e-03 * _cst["proton mass energy equivalent in MeV"][0] # [GeV]
_attr = re.compile(r"\[([a-zA-Z_][\w.:]*)]") # Identifier enclosed in square brackets
# noinspection PyPep8Naming
class multipole(_MadElement):
at2mad: ClassVar[dict[str, str]] = {}
klist = ("k0l", "k1l", "k2l", "k3l", "k4l", "k5l", "k6l", "k7l", "k8l", "k9l")
tlist = ("t0", "t1", "t2", "t3", "t4", "t5", "t6", "t7", "t8", "t9")
@mad_element
def to_at(self, **params):
polya = [0.0] * 10
polyb = [0.0] * 10
params.pop("l", None)
for k in list(params.keys()):
try:
order = self.klist.index(k)
except ValueError: # noqa: PERF203
pass
else:
strength = params.pop(k)
angle = (order + 1) * params.get(self.tlist[order], 0.0)
polyb[order] = strength * cos(-angle)
polya[order] = strength * sin(-angle)
return [
elt.ThinMultipole(
self.name, p_to_at(polya), p_to_at(polyb), **self.meval(params)
)
]
@classmethod
def from_at(cls, kwargs, factor=1.0):
el = super().from_at(kwargs)
maxorder = kwargs.pop("MaxOrder", -1) + 1
nlist = poly_to_mad(kwargs.pop("PolynomB", ())[: maxorder + 1], factor=factor)
slist = poly_to_mad(kwargs.pop("PolynomA", ())[: maxorder + 1], factor=factor)
for order, (va, vb) in enumerate(zip(slist, nlist, strict=True)):
if va != 0.0 or vb != 0.0:
el[multipole.klist[order]] = sqrt(va * va + vb * vb)
tilt = -atan2(va, vb) / (order + 1)
if tilt != 0.0:
el[multipole.tlist[order]] = tilt
return el
_mad8_env = {
# Constants
"_true_": True,
"_yes_": True,
"_t_": True,
"_on_": True,
"_false_": False,
"_no_": False,
"_f_": False,
"_off_": False,
"pi": pi,
"e": e,
"abs": abs,
"sqrt": sqrt,
"exp": exp,
"log": log,
"sin": sin,
"cos": cos,
"tan": tan,
"asin": asin,
"twopi": 2 * pi,
"degrad": 180.0 / pi,
"raddeg": pi / 180.0,
"emass": emass, # [GeV]
"pmass": pmass, # [GeV]
"clight": clight,
"qelect": qelect,
"centre": "centre",
"entry": "entry",
"exit": "exit",
# Elements
"drift": drift,
"marker": marker,
"quadrupole": quadrupole,
"sextupole": sextupole,
"octupole": octupole,
"multipole": multipole,
"longmultipole": longmultipole,
"sbend": sbend,
"rbend": rbend,
"kicker": kicker,
"hkicker": hkicker,
"vkicker": vkicker,
"rfcavity": rfcavity,
"monitor": monitor,
"hmonitor": hmonitor,
"vmonitor": vmonitor,
"sequence": _Sequence,
"line": _Line,
# Commands
"value": _Value(),
"__builtins__": {},
}
_ignore_names = [
"solenoid",
"rfmultipole",
"crabcavity",
"elseparator",
"collimator",
"tkicker",
]
_mad8_env.update(
(name, ignore_class(name, _MadElement, module=__name__)) for name in _ignore_names
)
[docs]
class Mad8Parser(_MadParser):
# noinspection PyUnresolvedReferences
r"""MAD-X specific parser.
The parser is a subclass of :py:class:`dict` and is database containing all the
MAD-X variables.
Example:
Parse a 1\ :sup:`st` file:
>>> parser = at.Mad8Parser()
>>> parser.parse_file("file1")
Parse another file:
>>> parser.parse_file("file2")
Get the variable "vkick"
>>> parser["vkick"]
0.003
Define a new variable:
>>> parser["hkick"] = -0.0024
Get the "qf1" element
>>> parser["qf1"]
quadrupole(name=qf1, l=1.0, k1=0.5, tilt=0.001)
Generate an AT :py:class:`.Lattice` from the "ring" sequence
>>> ring = parser.lattice(use="ring") # generate an AT Lattice
"""
_continuation = "&"
_blockcomment = ("comment", "endcomment")
def __init__(self, *filenames: str, **kwargs):
"""
Args:
*filenames: files to be read at initialisation
strict: If :py:obj:`False`, assign 0 to undefined variables
verbose: If :py:obj:`True`, print details on the processing
**kwargs: Initial variable definitions.
"""
super().__init__(_mad8_env, *filenames, **kwargs)
def _format_command(self, expr: str) -> str:
"""Evaluate an expression using *self* as local namespace."""
expr = super()._format_command(expr)
expr = _attr.sub(r".\1", expr) # Attribute access: VAR[ATTR]
expr = expr.replace("^", "**") # Exponentiation
return expr
[docs]
def load_mad8(
*files: str | Path,
use: str = "ring",
strict: bool = True,
verbose: bool = False,
parameterised: bool = False,
**kwargs,
) -> Lattice:
"""Create a :py:class:`.Lattice` from MAD8 files.
- The *energy* and *particle* of the generated lattice are taken from the MAD8
``BEAM`` object, using the MAD8 default parameters: positrons at 1 Gev.
These parameters are overloaded by the value given in the *energy* and
*particle* keyword arguments.
- The radiation state is given by the ``RADIATE`` flag of the ``BEAM`` object,
using the AT defaults: RF cavities active, synchrotron radiation in dipoles and
quadrupoles.
- Long elements are split according to the default AT value for *NumIntSteps* (10).
Parameters:
files: Names of one or several MAD8 files
strict: If :py:obj:`False`, assign 0 to undefined variables
use: Name of the MAD8 sequence or line containing the desired
lattice. Default: ``ring``
verbose: If :py:obj:`True`, print details on the processing
Keyword Args:
name (str): Name of the lattice. Default: MAD8 sequence name.
particle(Particle): Circulating particle. Default: from MAD8
energy (float): Energy of the lattice [eV]. Default: from MAD8
periodicity(int): Number of periods. Default: 1
*: Other keywords will be used as initial variable definitions
Returns:
lattice (Lattice): New :py:class:`.Lattice` object
"""
parser = Mad8Parser(strict=strict, verbose=verbose)
params = {
key: kwargs.pop(key)
for key in ("name", "particle", "energy", "periodicity")
if key in kwargs
}
parser.parse_files(*files, **kwargs)
return parser.lattice(use=use, parameterised=parameterised, **params)
class _Mad8Exporter(_MadExporter):
delimiter: ClassVar[str] = ""
continuation: ClassVar[str] = "&"
bool_fmt: ClassVar[dict[bool, str]] = {False: ".FALSE.", True: ".TRUE."}
AT2MAD: ClassVar[dict] = {
elt.Quadrupole: quadrupole,
elt.Sextupole: sextupole,
elt.Octupole: octupole,
elt.ThinMultipole: multipole,
elt.Multipole: longmultipole,
elt.RFCavity: rfcavity,
elt.Drift: drift,
elt.Bend: sbend,
elt.Marker: marker,
elt.Monitor: monitor,
elt.Corrector: kicker,
}
[docs]
def save_mad8(ring: Lattice, filename: str | Path | None = None, **kwargs):
"""Save a :py:class:`.Lattice` as a MAD8 file.
Args:
ring: lattice
filename: file to be created. If None, write to sys.stdout
Keyword Args:
use (str | None): name of the created SEQUENCE of LINE.
Default: name of the PyAT lattice
use_line (bool): If True, use a MAD "LINE" format. Otherwise, use
a MAD "SEQUENCE"
"""
exporter = _Mad8Exporter(ring, **kwargs)
exporter.export(filename)
