madx input

pyat/at/lattice/elements.py

@@ -280,22 +280,17 @@ def __setattr__(self, key, value):
             raise
 
     def __str__(self):
-        first3 = ['FamName', 'Length', 'PassMethod']
-        attrs = dict(self.items())
-        keywords = ['\t{0} : {1!s}'.format(k, attrs.pop(k)) for k in first3]
-        keywords += ['\t{0} : {1!s}'.format(k, v) for k, v in attrs.items()]
-        return '\n'.join((type(self).__name__ + ':', '\n'.join(keywords)))
+        return "\n".join(
+            [self.__class__.__name__ + ":"]
+            + [f"{k:>14}: {v!s}" for k, v in self.items()]
+        )
 
     def __repr__(self):
-        attrs = dict(self.items())
-        arguments = [attrs.pop(k, getattr(self, k)) for k in
-                     self._BUILD_ATTRIBUTES]
-        defelem = self.__class__(*arguments)
-        keywords = ['{0!r}'.format(arg) for arg in arguments]
-        keywords += ['{0}={1!r}'.format(k, v) for k, v in sorted(attrs.items())
-                     if not numpy.array_equal(v, getattr(defelem, k, None))]
-        args = re.sub(r'\n\s*', ' ', ', '.join(keywords))
-        return '{0}({1})'.format(self.__class__.__name__, args)
+        clsname, args, kwargs = self.definition
+        keywords = [f"{arg!r}" for arg in args]
+        keywords += [f"{k}={v!r}" for k, v in kwargs.items()]
+        args = re.sub(r"\n\s*", " ", ", ".join(keywords))
+        return f"{clsname}({args})"
 
     def equals(self, other) -> bool:
         """Whether an element is equivalent to another.
@@ -371,9 +366,27 @@ def deepcopy(self) -> Element:
         """Return a deep copy of the element"""
         return deepcopy(self)
 
+    @property
+    def definition(self) -> tuple[str, tuple, dict]:
+        """tuple (class_name, args, kwargs) defining the element"""
+        attrs = dict(self.items())
+        arguments = tuple(attrs.pop(
+            k, getattr(self, k)) for k in self._BUILD_ATTRIBUTES
+        )
+        defelem = self.__class__(*arguments)
+        keywords = dict(
+            (k, v)
+            for k, v in attrs.items()
+            if not numpy.array_equal(v, getattr(defelem, k, None))
+        )
+        return self.__class__.__name__, arguments, keywords
+
     def items(self) -> Generator[tuple[str, Any], None, None]:
         """Iterates through the data members"""
-        for k, v in vars(self).items():
+        v = vars(self).copy()
+        for k in ["FamName", "Length", "PassMethod"]:
+            yield k, v.pop(k)
+        for k, v in sorted(v.items()):
             yield k, v
 
     def is_compatible(self, other: Element) -> bool:

pyat/at/load/__init__.py

@@ -8,3 +8,5 @@
 from .reprfile import *
 from .tracy import *
 from .elegant import *
+from .file_input import *
+from .madx import *

pyat/at/load/allfiles.py

@@ -13,40 +13,30 @@
 def load_lattice(filepath: str, **kwargs) -> Lattice:
     """Load a Lattice object from a file
 
-The file format is indicated by the filepath extension.
+The file format is indicated by the filepath extension. The file name is stored in
+the *in_file* Lattice attribute. The selected variable, if relevant, is stored
+in the *use* Lattice attribute.
 
 Parameters:
     filepath:           Name of the file
 
 Keyword Args:
+    use (str):          Name of the variable containing the desired lattice.
+      Default: if there is a single variable, use it, otherwise select ``"RING"``
     name (str):         Name of the lattice.
-      Default: taken from the file, or ``''``
+      Default: taken from the file, or ``""``
     energy (float):     Energy of the lattice
       (default: taken from the file)
-    periodicity (int]): Number of periods
+    periodicity (int):  Number of periods
       (default: taken from the file, or 1)
     *:                  All other keywords will be set as :py:class:`.Lattice`
       attributes
 
-Specific keywords for .mat files
-
-Keyword Args:
-    mat_key (str):      Name of the Matlab variable containing
-      the lattice. Default: Matlab variable name if there is only one,
-      otherwise ``'RING'``
-    check (bool):       Run coherence tests. Default: :py:obj:`True`
-    quiet (bool):       Suppress the warning for non-standard classes.
-      Default: :py:obj:`False`
-    keep_all (bool):    Keep Matlab RingParam elements as Markers.
-      Default: :py:obj:`False`
+Check the format-specific function for specific keyword arguments.
 
 Returns:
     lattice (Lattice):          New :py:class:`.Lattice` object
 
-See Also:
-    :py:func:`.load_mat`, :py:func:`.load_m`, :py:func:`.load_repr`,
-    :py:func:`.load_elegant`, :py:func:`.load_tracy`
-
 .. Admonition:: Known extensions are:
     """
     _, ext = os.path.splitext(filepath)

pyat/at/load/elegant.py

@@ -337,10 +337,8 @@ def load_elegant(filename: str, **kwargs) -> Lattice:
         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
+        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
@@ -354,7 +352,7 @@ def load_elegant(filename: str, **kwargs) -> Lattice:
         harmonic_number = kwargs.pop("harmonic_number")
 
         def elem_iterator(params, elegant_file):
-            with open(params.setdefault("elegant_file", elegant_file)) as f:
+            with open(params.setdefault("in_file", elegant_file)) as f:
                 contents = f.read()
                 element_lines = expand_elegant(
                     contents, lattice_key, energy, harmonic_number
@@ -370,4 +368,5 @@ def elem_iterator(params, elegant_file):
                          'lattice {}: {}'.format(filename, e))
 
 
-register_format(".lte", load_elegant, descr="Elegant format")
+register_format(
+    ".lte", load_elegant, descr="Elegant format. See :py:func:`.load_elegant`.")