Dcoumentation and tests

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
 

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

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")