Save and load lattices in JSON format (#766)

* Handle JSON files in python and Matlab
* Added tests for save and load

atmat/lattice/atdivelem.m

@@ -37,8 +37,8 @@
     line=atsetfieldvalues(line,'ExitAngle',0.0);
 end
 if isfield(elem,'KickAngle')
-    line=atsetfieldvalues(line,'KickAngle',{1,1},el.KickAngle(1,1)*frac(:)/sum(frac));
-    line=atsetfieldvalues(line,'KickAngle',{1,2},el.KickAngle(1,2)*frac(:)/sum(frac));
+    line=atsetfieldvalues(line,'KickAngle',{1},el.KickAngle(1)*frac(:)/sum(frac));
+    line=atsetfieldvalues(line,'KickAngle',{2},el.KickAngle(2)*frac(:)/sum(frac));
 end
 
 line{1}=mvfield(line{1},entrancef);     % Set back entrance fields

atmat/lattice/atloadlattice.m

@@ -21,12 +21,16 @@
 %           the variable name must be specified using the 'matkey' keyword.
 %
 %   .m      Matlab function. The function must output a valid AT structure.
+%   .json   JSON file
+%
+%see also atwritem, atwritejson
 
 persistent link_table
 
 if isempty(link_table)
     link_table.mat=@load_mat;
     link_table.m=@load_m;
+    link_table.json=@load_json;
 end
 
 [~,~,fext]=fileparts(fspec);
@@ -57,7 +61,7 @@
         dt=load(fpath);
         vnames=fieldnames(dt);
         key='RING';
-        if length(vnames) == 1
+        if isscalar(vnames)
             key=vnames{1};
         else
             for v={'ring','lattice'}
@@ -75,7 +79,31 @@
             error('AT:load','Cannot find variable %s\nmatkey must be in: %s',...
                 key, strjoin(vnames,', '));
         end
+    end
 
+    function [lattice, opts]=load_json(fpath, opts)
+        data=jsondecode(fileread(fpath));
+        % File signature for later use
+        try
+            atjson=data.atjson;
+        catch
+            atjson=1;
+        end
+        props=data.properties;
+        name=props.name;
+        energy=props.energy;
+        periodicity=props.periodicity;
+        particle=atparticle.loadobj(props.particle);
+        harmnumber=props.harmonic_number;
+        props=rmfield(props,{'name','energy','periodicity','particle','harmonic_number'});
+        args=[fieldnames(props) struct2cell(props)]';
+        lattice=atSetRingProperties(data.elements,...
+            'FamName', name,...
+            'Energy', energy,...
+            'Periodicity', periodicity,...
+            'Particle', particle,...
+            'HarmNumber', harmnumber, ...
+            args{:});
     end
 
 end

atmat/lattice/atwritejson.m

@@ -0,0 +1,68 @@
+function varargout=atwritejson(ring, varargin)
+%ATWRITEJSON Create a JSON file to store an AT lattice
+%
+%JS=ATWRITEJSON(RING)
+%   Return the JSON representation of RING as a character array
+%
+%ATWRITEJSON(RING, FILENAME)
+%   Write the JSON representation of RING to the file FILENAME
+%
+%ATWRITEJSON(RING, ..., 'compact', true)
+%   If compact is true, write a compact JSON file (no linefeeds)
+%
+%see also atloadlattice
+
+[compact, varargs]=getoption(varargin, 'compact', false);
+[filename, ~]=getargs(varargs,[]);
+
+if ~isempty(filename)
+    %get filename information
+    [pname,fname,ext]=fileparts(filename);
+
+    %Check file extension
+    if isempty(ext), ext='.json'; end
+
+    % Open file to be written
+    [fid,mess]=fopen(fullfile(pname,[fname ext]),'wt');
+
+    if fid==-1
+        error('AT:FileErr','Cannot Create file %s\n%s',fn,mess);
+    else
+        fprintf(fid, sjson(ring));
+        fclose(fid);
+    end
+    varargout={};
+else
+    varargout={sjson(ring)};
+end
+
+    function jsondata=sjson(ring)
+        ok=~atgetcells(ring, 'Class', 'RingParam');
+        data.atjson= 1;
+        data.elements=ring(ok);
+        data.properties=get_params(ring);
+        jsondata=jsonencode(data, 'PrettyPrint', ~compact);
+    end
+
+    function prms=get_params(ring)
+        % Get "standard" properties
+        [name, energy, part, periodicity, harmonic_number]=...
+            atGetRingProperties(ring,'FamName', 'Energy', 'Particle',...
+            'Periodicity', 'HarmNumber');
+        prms=struct('name', name, 'energy', energy, 'periodicity', periodicity,...
+            'particle', saveobj(part), 'harmonic_number', harmonic_number);
+        % Add user-defined properties
+        idx=atlocateparam(ring);
+        if ~isempty(idx)
+            flist={'FamName','PassMethod','Length','Class',...
+                'Energy', 'Particle','Periodicity','cell_harmnumber'};
+            present=isfield(ring{idx}, flist);
+            p2=rmfield(ring{idx},flist(present));
+            for nm=fieldnames(p2)'
+                na=nm{1};
+                prms.(na)=p2.(na);
+            end
+        end
+    end
+
+end

pyat/at/lattice/elements.py

@@ -302,26 +302,17 @@ def __setattr__(self, key, value):
             super(Element, self).__setattr__(key, value)
 
     def __str__(self):
-        first3 = ["FamName", "Length", "PassMethod"]
-        # Get values and parameter objects
-        attrs = dict(self.items())
-        keywords = [f"\t{k} : {attrs.pop(k)!s}" for k in first3]
-        keywords += [f"\t{k} : {v!s}" 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):
-        # Get values only, even for parameters
-        attrs = dict((k, getattr(self, k)) for k, v in self.items())
-        arguments = [attrs.pop(k) for k in self._BUILD_ATTRIBUTES]
-        defelem = self.__class__(*arguments)
-        keywords = [f"{v!r}" for v in arguments]
-        keywords += [
-            f"{k}={v!r}"
-            for k, v in sorted(attrs.items())
-            if not numpy.array_equal(v, getattr(defelem, k, None))
-        ]
+        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 "{0}({1})".format(self.__class__.__name__, args)
+        return f"{clsname}({args})"
 
     def equals(self, other) -> bool:
         """Whether an element is equivalent to another.
@@ -399,10 +390,28 @@ 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"""
-        # Properties may be added by overloading this method
-        yield from 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:
         """Checks if another :py:class:`Element` can be merged"""

pyat/at/load/__init__.py

@@ -8,3 +8,4 @@
 from .reprfile import *
 from .tracy import *
 from .elegant import *
+from .json import *

pyat/at/load/allfiles.py

@@ -1,10 +1,16 @@
 """Generic function to save and load python AT lattices. The format is
 determined by the file extension
 """
+
+from __future__ import annotations
+
+__all__ = ["load_lattice", "save_lattice", "register_format"]
+
 import os.path
-from at.lattice import Lattice
+from collections.abc import Callable
+from typing import Optional
 
-__all__ = ['load_lattice', 'save_lattice', 'register_format']
+from at.lattice import Lattice
 
 _load_extension = {}
 _save_extension = {}
@@ -13,41 +19,31 @@
 def load_lattice(filepath: str, **kwargs) -> Lattice:
     """Load a Lattice object from a file
 
-The file format is indicated by the filepath extension.
-
-Parameters:
-    filepath:           Name of the file
-
-Keyword Args:
-    name (str):         Name of the lattice.
-      Default: taken from the file, or ``''``
-    energy (float):     Energy of the lattice
-      (default: taken from the file)
-    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`
-
-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:
+    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 ``""``
+        energy (float):     Energy of the lattice
+          (default: taken from the file)
+        periodicity (int):  Number of periods
+          (default: taken from the file, or 1)
+        *:                  All other keywords will be set as :py:class:`.Lattice`
+          attributes
+
+    Returns:
+        lattice (Lattice):          New :py:class:`.Lattice` object
+
+    Check the format-specific function for specific keyword arguments:
+
+    .. Admonition:: Known extensions are:
     """
     _, ext = os.path.splitext(filepath)
     try:
@@ -58,25 +54,18 @@ def load_lattice(filepath: str, **kwargs) -> Lattice:
         return load_func(filepath, **kwargs)
 
 
-def save_lattice(ring: Lattice, filepath: str, **kwargs):
+def save_lattice(ring: Lattice, filepath: str, **kwargs) -> None:
     """Save a Lattice object
 
-The file format is indicated by the filepath extension.
-
-Parameters:
-    ring:               Lattice description
-    filepath:           Name of the file
+    The file format is indicated by the filepath extension.
 
-Specific keywords for .mat files
-
-Keyword Args:
-    mat_key (str):      Name of the Matlab variable containing the lattice.
-      Default: ``'RING'``
+    Parameters:
+        ring:               Lattice description
+        filepath:           Name of the file
 
-See Also:
-    :py:func:`.save_mat`, :py:func:`.save_m`, :py:func:`.save_repr`
+    Check the format-specific function for specific keyword arguments:
 
-.. Admonition:: Known extensions are:
+    .. Admonition:: Known extensions are:
     """
     _, ext = os.path.splitext(filepath)
     try:
@@ -87,24 +76,26 @@ def save_lattice(ring: Lattice, filepath: str, **kwargs):
         return save_func(ring, filepath, **kwargs)
 
 
-def register_format(extension: str, load_func=None, save_func=None,
-                    descr: str = ''):
+def register_format(
+    extension: str,
+    load_func: Optional[Callable[..., Lattice]] = None,
+    save_func: Optional[Callable[..., None]] = None,
+    descr: str = "",
+):
     """Register format-specific processing functions
 
     Parameters:
         extension:      File extension string.
-        load_func:      load function. Default: :py:obj:`None`
-        save_func:      save_lattice function Default: :py:obj:`None`
-        descr:          File type description
+        load_func:      load function.
+        save_func:      save function.
+        descr:          File type description.
     """
     if load_func is not None:
         _load_extension[extension] = load_func
-        load_lattice.__doc__ += '\n    {0:<10}'\
-                                '\n        {1}\n'.format(extension, descr)
+        load_lattice.__doc__ += f"\n        {extension:<10}\n            {descr}\n"
     if save_func is not None:
         _save_extension[extension] = save_func
-        save_lattice.__doc__ += '\n    {0:<10}'\
-                                '\n        {1}\n'.format(extension, descr)
+        save_lattice.__doc__ += f"\n        {extension:<10}\n            {descr}\n"
 
 
 Lattice.load = staticmethod(load_lattice)

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