diff --git a/pyat/at/load/file_input.py b/pyat/at/load/file_input.py index 77a6b101e..a2be88847 100644 --- a/pyat/at/load/file_input.py +++ b/pyat/at/load/file_input.py @@ -116,7 +116,7 @@ def __repr__(self): @staticmethod def convert(name: str, *args, **params) -> list[Element]: - """Generate the AT element""" + """Generate the AT element, Most be overloaded for each specific element""" return [] # noinspection PyUnusedLocal @@ -145,7 +145,7 @@ class SequenceDescr(AnyDescr, list, abc.ABC): """Simple representation of a sequence of elements as a list""" @property - def length(self): + def length(self) -> float: return getattr(self, "l", 0.0) @@ -154,8 +154,9 @@ class BaseParser(dict): Analyses files with the following MAD-like format: - variable = value - label : command [,attribute=value] [,attribute=value]... + ``variable = value`` + + ``label : command [,attribute=value] [,attribute=value]...`` The parser builds a database of all the defined objects """ @@ -169,19 +170,19 @@ def __init__( *args, delimiter: Optional[str] = None, continuation: str = "\\", - linecomment: Union[None, str, Sequence[str]] = "#", + linecomment: Union[str, Sequence[str], None] = "#", blockcomment: Optional[tuple[str, str]] = None, endfile: Optional[str] = None, **kwargs, ): """ Args: - env: global namespace + env: global namespace used for evaluating commands delimiter: command delimiter continuation: command continuation character linecomment: Line comment character blockcomment: Block comment delimiter - endfile: End of input marker + endfile: "End of input" marker *args: dict initializer **kwargs: dict initializer """ @@ -201,12 +202,13 @@ def __init__( super().__init__(*args, **kwargs) def clear(self): + """Clean the database""" super().clear() self.update(self.kwargs) # noinspection PyUnusedLocal def evaluate(self, item, no_global: bool = False): - """Evaluate an expression using self as local namespace""" + """Evaluate an expression using *self* as local namespace""" return eval(_clean_expr(item), self.env, self) def _eval_cmd(self, cmdname: str, no_global: bool = False): @@ -255,7 +257,6 @@ def _raw_command( **kwargs, ): """Command execution""" - argparser = self._argument_parser.get(cmdname, _default_arg_parser) kwargs.update(argparser(self, arg) for arg in argnames) if label is None: @@ -274,7 +275,6 @@ def _format_statement(self, line: str) -> str: Overload this method for specific languages""" line, matches = protect(line, fence=('"', '"')) # protect the quoted parts line = "".join(line.split()).lower() # Remove all spaces, lower - # line = re.sub(r"(?<=[a-z_.])\.", "_", line) # Replace "." by "_" (line,) = restore(matches, line) return line diff --git a/pyat/at/load/madx.py b/pyat/at/load/madx.py index 84c1cc577..b889fe963 100644 --- a/pyat/at/load/madx.py +++ b/pyat/at/load/madx.py @@ -5,6 +5,7 @@ __all__ = ["MadxParser", "load_madx"] import functools +import warnings # functions known by MADX from math import pi, e, sqrt, exp, log, log10, sin, cos, tan # noqa: F401 @@ -42,7 +43,9 @@ def sinc(x: float) -> float: - return sin(x)/x + return sin(x) / x + + # ------------------- # Utility functions # ------------------- @@ -68,10 +71,10 @@ def wrapper(name, *args, tilt=None, **kwargs): return wrapper -def polyn(a): +def polyn(a: Sequence[float]) -> np.ndarray: """Convert polynomials from MADX to AT""" - def ref(n, t): + def ref(n: int, t: float): nonlocal f v = t / f f *= n + 1 @@ -449,22 +452,24 @@ def __init__( callfun, *args, particle="positron", - mass=emass, + mass=emass, # GeV charge=1.0, - energy=1.0, + energy=1.0, # GeV bcurrent=0.0, **kwargs, ): - atparticle = Particle(particle, rest_energy=1.0e09 * mass, charge=charge) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=UserWarning) + atparticle = Particle(particle, rest_energy=1.0e09 * mass, charge=charge) mass = 1.0e-09 * atparticle.rest_energy charge = atparticle.charge gamma = energy / mass beta = sqrt(1.0 - 1.0 / gamma / gamma) pc = beta * energy # GeV - kwargs.setdefault("gamma", gamma) + kwargs.setdefault("gamma", gamma) kwargs.setdefault("beta", beta) kwargs.setdefault("pc", pc) - kwargs.setdefault("brho", 1.0e09 * pc / abs(charge) / clight) + kwargs.setdefault("brho", 1.0e09 * pc / abs(charge) / clight) super().__init__( *args, source="beam", @@ -500,9 +505,39 @@ def expand(self, parser: MadxParser) -> dict: class MadxParser(UnorderedParser): - """MADX specific parser + # noinspection PyUnresolvedReferences + """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 1st file: + + >>> parser = at.MadxParser() + >>> parser.parse_file("file1") + + Parse another file: + + >>> parser.parse_file("file2") + + Get the variable "vkick" - After processing a file, the parser is a database containing all the MADX variables + >>> 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 """ _soft_eval = {"file", "refer"} @@ -557,16 +592,10 @@ def _beam_cmd( self[name] = beam def _format_statement(self, line: str) -> str: - - # def repl(match): - # return match.group().replace(".", "_") - line, matches = protect(line, fence=('"', '"')) line = "".join(line.split()).lower() # Remove all spaces, lower - # line = re.sub(r"[a-z][\w.]*", repl, line) # Replace "." by "_" line = line.replace("from", "frm") # Replace from by frm line = line.replace("{", "(").replace("}", ")") - # line = line.replace("->", ".") # for attribute access line = line.replace(":=", "=") # since we evaluate only once (line,) = restore(matches, line) return line @@ -619,7 +648,7 @@ def _get_beam(self, key: str): beam = self["beam"] return beam - def lattice(self, use="cell", **kwargs): + def lattice(self, use: str = "cell", **kwargs): """Create a lattice from the selected sequence Parameters: @@ -656,10 +685,10 @@ def gener(params, *args): def load_madx(*files: str, use: str = "cell", verbose=False, **kwargs) -> Lattice: - """Create a :py:class:`.Lattice` from a MADX file + """Create a :py:class:`.Lattice` from MAD-X files Parameters: - files: Names of one or several Mad files to process + files: Names of one or several MAD-X files use: Name of the MADX sequence or line containing the desired lattice. Default: ``cell`` verbose: Print details on processing diff --git a/pyat/test/test_file_parser.py b/pyat/test/test_file_parser.py index e026e1a56..03dd9704b 100644 --- a/pyat/test/test_file_parser.py +++ b/pyat/test/test_file_parser.py @@ -1,6 +1,6 @@ import pytest from math import sqrt # noqa: F401 -from at import UnorderedParser +from at import UnorderedParser, MadxParser __all__ = [] @@ -42,6 +42,40 @@ test_data = dict(data1=test_data1, data2=test_data2) +madx_data = """ +BEAM, PARTICLE='electron', RADIATE, ENERGY=6.0, SEQUENCE=RING; +BEAM, PARTICLE='positron', ENERGY=2.0; + +Q1: QUADRUPOLE, L:=QL ; ! check forward reference + +QF1: Q1, K1=0.5 ; +QD1: Q1, K1=-QF1->K1; ! check attribute access +QF1, TILT=0.001; ! check element update +SOL1: SOLENOID, L=0.5, K1S=3.0; +MULT1: MULTIPOLE, KNL={1.0, 1.0, 2.0, 6.0}; +HK1: HKICKER, L=0, KICK=0.001; +VK1: VKICKER, L=0, KICK=-0.001; +HVK2: KICKER, L=0, VKICK=VKICK; +BPM1: MONITOR, L=0.1; +BPM2: VMONITOR; +RFCAVITY, VOLT=5, FREQ=352.2, HARMON=31, AT=0.0; + +RING: SEQUENCE, L=9.0; +SOL1, AT=4.5; +MULT1, AT=4.8; +HK1, AT=4.8; +VK1, AT=4.8; +HVK2, AT=4.8; +BPM1, AT=4.9; +BPM2, AT=5; +ENDSEQUENCE; + +VALUE, EMASS; + +QL = sqrt(4)/2; ! check arithmetic +VKICK = 0.003; +""" + def command1(**kwargs): """Sample command testing that arguments are as expected""" @@ -58,7 +92,7 @@ def command1(**kwargs): @pytest.mark.parametrize( "delimiter, linecomment, data", [[";", ("!", "//"), "data1"], [None, "#", "data2"]] ) -def test_unordered_parser1(delimiter, linecomment, data): +def test_unordered_parser(delimiter, linecomment, data): parser = UnorderedParser( globals(), blockcomment=("/*", "*/"), @@ -67,3 +101,9 @@ def test_unordered_parser1(delimiter, linecomment, data): ) parser.parse_lines(test_data[data].splitlines()) assert parser["label"] == "done" + + +def test_madx_parser(): + parser = MadxParser() + parser.parse_lines(madx_data.splitlines()) + ring = parser.lattice(use="ring")