Add MADX parser (#39)

Author: felix-andreas

latticejson/__init__.py

@@ -1,6 +1,2 @@
 from .__about__ import __version__
-from .io import load, convert
-from .validate import validate
-from .convert import elegant_to_latticejson, latticejson_to_elegant
-from .migrate import migrate
-from .exceptions import UndefinedObjectError, UndefinedRPNVariableError
+from .io import load, save

latticejson/cli.py

@@ -5,14 +5,13 @@
 
 from . import __version__
 from .validate import validate_file
-from .io import convert as _convert
-from .parse import parse_elegant as _parse_elegant
-from .format import CompactJSONEncoder
+from . import io
+from .format import format_json
 from .migrate import migrate as _migrate
+from . import parse
 
 
-FORMATS = "json", "lte"
-dump_latticejson = lambda obj: json.dumps(obj, cls=CompactJSONEncoder, indent=4)
+FORMATS = "json", "lte", "madx"
 
 
 @click.group()
@@ -38,7 +37,7 @@ def cli():
 )
 def convert(file, from_, to):
     """Convert FILE (path or url) to another lattice file format."""
-    click.echo(_convert(file, from_, to))
+    click.echo(io.save_string(io.load(file, from_), to))
 
 
 @cli.command()
@@ -48,14 +47,6 @@ def validate(file):
     validate_file(file)
 
 
-@cli.command()
-@click.argument("file", type=click.Path(exists=True))
-def parse_elegant(file):
-    """Parse elegant file but do not convert to LatticeJSON."""
-    text = Path(file).read_text()
-    click.echo(dump_latticejson(_parse_elegant(text)))
-
-
 @cli.command()
 @click.argument("files", nargs=-1, type=click.Path(exists=True))
 @click.option(
@@ -70,7 +61,7 @@ def autoformat(files, dry_run):
         path.rglob("*.json") if path.is_dir() else (path,) for path in map(Path, files)
     ):
         latticejson = json.loads(path.read_text())
-        formatted = dump_latticejson(latticejson)
+        formatted = format_json(latticejson)
         click.secho(f"reformatted {path}", bold=True)
         if dry_run:
             click.echo(formatted)
@@ -88,4 +79,34 @@ def migrate(file, from_, to):
     initial_version = from_.split(".")
     final_version = to.split(".")
     latticejson = _migrate(json.loads(text), initial_version, final_version)
-    click.echo(dump_latticejson(latticejson))
+    click.echo(format_json(latticejson))
+
+
+@cli.group()
+def debug():
+    """Some useful commands for debugging/development."""
+    pass
+
+
+@debug.command()
+@click.argument("file", type=click.Path(exists=True))
+@click.option("--transform", "-t", is_flag=True, help="Print transformed tree.")
+def parse_elegant(file, transform):
+    """Print parse tree of elegant lattice file."""
+    text = Path(file).read_text()
+    if transform:
+        click.echo(format_json(parse.parse_elegant(text)))
+    else:
+        click.echo(parse.ELEGANT_PARSER.parse(text).pretty())
+
+
+@debug.command()
+@click.argument("file", type=click.Path(exists=True))
+@click.option("--transform", "-t", is_flag=True, help="Print transformed tree.")
+def parse_madx(file, transform):
+    """Print parse tree of madx lattice file."""
+    text = Path(file).read_text()
+    if transform:
+        click.echo(format_json(parse.parse_madx(text)))
+    else:
+        click.echo(parse.MADX_PARSER.parse(text).pretty())

latticejson/convert.py

@@ -1,56 +1,32 @@
-from typing import List, Tuple, Dict
+from typing import List, Dict
 from pathlib import Path
 import json
-import warnings
-
-from .validate import validate
-from .parse import parse_elegant
-
-
-LATTICEJSON_ELEGANT_MAP: Tuple[Tuple[str, Tuple[str, ...]], ...] = (
-    ("Drift", ("DRIF", "DRIFT")),
-    ("Dipole", ("CSBEND", "SBEND", "BEND")),
-    ("Quadrupole", ("KQUAD", "QUAD", "QUADRUPOLE")),
-    ("Sextupole", ("KSEXT", "SEXT", "SEXTUPOLE")),
-    ("Lattice", ("LINE",)),
-    ("length", ("L",)),
-    ("angle", ("ANGLE",)),
-    ("e1", ("E1",)),
-    ("e2", ("E2",)),
-    ("k1", ("K1",)),
-    ("k2", ("K2",)),
-)
-JSON_TO_ELE: Dict[str, str] = {x: y[0] for x, y in LATTICEJSON_ELEGANT_MAP}
-ELE_TO_JSON: Dict[str, str] = {y: x for x, tup in LATTICEJSON_ELEGANT_MAP for y in tup}
-
-
-def latticejson_to_elegant(lattice_dict) -> str:
-    """Convert LatticeJSON dict to elegant lattice file format.
-    :param dict: dict in LatticeJSON format
-    :return: string with in elegant lattice file format
-    """
-    elements = lattice_dict["elements"]
-    sub_lattices = lattice_dict["sub_lattices"]
+from warnings import warn
+from .parse import parse_elegant, parse_madx
+from .exceptions import UnknownElementWarning, UnknownAttributeWarning
 
-    strings = []
-    element_template = "{}: {}, {}".format
-    for name, (type_, attributes) in elements.items():
-        attrs = ", ".join(f"{JSON_TO_ELE[k]}={v}" for k, v in attributes.items())
-        elegant_type = JSON_TO_ELE[type_]
-        strings.append(element_template(name, elegant_type, attrs))
+NAME_MAP = json.loads((Path(__file__).parent / "map.json").read_text())["map"]
+JSON_TO_ELE = {x: y[0][0] for x, *y in NAME_MAP}
+ELE_TO_JSON = {y: x for x, *tup in NAME_MAP for y in tup[0]}
+JSON_TO_MADX = {x: y[1][0] for x, *y in NAME_MAP}
+MADX_TO_JSON = {y: x for x, *tup in NAME_MAP for y in tup[1]}
 
-    lattice_template = "{}: LINE=({})".format
-    for name in sort_lattices(sub_lattices):
-        strings.append(lattice_template(name, ", ".join(sub_lattices[name])))
 
-    name = lattice_dict["name"]
-    strings.append(lattice_template(name, ", ".join(lattice_dict["lattice"])))
-    strings.append("\n")
-    return "\n".join(strings)
+def from_elegant(string):
+    """Convert an elegant lattice file to a LatticeJSON dict.
 
+    :param str string: input lattice file as string
+    :param lattice_name: name of the lattice
+    :type str, optional
+    :param description: description of the lattice
+    :type str, optional
+    :return: dict in LatticeJSON format
+    """
+    return _map_names_from_elegant(parse_elegant(string))
 
-def elegant_to_latticejson(string):
-    """Convert an elegant lattice file to a LatticeJSON dict.
+
+def from_madx(string):
+    """Convert a MADX lattice file to a LatticeJSON dict.
 
     :param str string: input lattice file as string
     :param lattice_name: name of the lattice
@@ -59,30 +35,36 @@ def elegant_to_latticejson(string):
     :type str, optional
     :return: dict in LatticeJSON format
     """
-    elegant_dict = parse_elegant(string)
+    return _map_names_from_madx(parse_madx(string))
 
+
+def _map_names_from_madx(madx_dict: dict):
+    return _map_names(madx_dict, MADX_TO_JSON)
+
+
+def _map_names_from_elegant(elegant_dict: dict):
+    return _map_names(elegant_dict, ELE_TO_JSON)
+
+
+def _map_names(lattice_data: dict, name_map: dict):
     elements = {}
-    for name, (elegant_type, elegant_attributes) in elegant_dict["elements"].items():
-        latticejson_type = ELE_TO_JSON.get(elegant_type)
+    for name, (other_type, other_attributes) in lattice_data["elements"].items():
+        latticejson_type = name_map.get(other_type)
         if latticejson_type is None:
-            elements[name] = ["Drift", {"length": elegant_attributes.get("L", 0)}]
-            warnings.warn(
-                f"Replacing element {name} ({elegant_type}) with Drift.", stacklevel=2
-            )
+            elements[name] = ["Drift", {"length": other_attributes.get("L", 0)}]
+            warn(UnknownElementWarning(name, other_type), stacklevel=2)
             continue
 
         attributes = {}
         elements[name] = [latticejson_type, attributes]
-        for elegant_key, value in elegant_attributes.items():
-            latticejson_key = ELE_TO_JSON.get(elegant_key)
+        for other_key, value in other_attributes.items():
+            latticejson_key = name_map.get(other_key)
             if latticejson_key is not None:
                 attributes[latticejson_key] = value
             else:
-                warnings.warn(
-                    f"Ignoring attribute {elegant_key} of {name}.", stacklevel=2
-                )
+                warn(UnknownAttributeWarning(other_key, name), stacklevel=2)
 
-    lattices = elegant_dict["lattices"]
+    lattices = lattice_data["lattices"]
     lattice_name, main_lattice = lattices.popitem()  # use last lattice as main_lattice
     return dict(
         name=lattice_name,
@@ -92,24 +74,54 @@ def elegant_to_latticejson(string):
     )
 
 
-def sort_lattices_old(lattices: Dict[str, List[str]]) -> List[str]:
-    """Returns a sorted list of lattice names for a given dict of lattices."""
+def to_elegant(latticejson: dict) -> str:
+    """Convert a LatticeJSON dict to the elegant lattice file format.
 
-    lattices_copy = lattices.copy()
-    lattice_names = []
+    :param lattice_dict dict: dict in LatticeJSON format
+    :return: string with in elegant lattice file format
+    """
+    elements = latticejson["elements"]
+    sub_lattices = latticejson["sub_lattices"]
 
-    def _sort_lattices(name, arrangement: List[str]):
-        for child_name in arrangement:
-            if child_name in lattices_copy:
-                _sort_lattices(child_name, lattices_copy[child_name])
+    strings = [f"! TITLE: {latticejson['name']}"]
+    element_template = "{}: {}, {}".format
+    for name, (type_, attributes) in elements.items():
+        attrs = ", ".join(f"{JSON_TO_ELE[k]}={v}" for k, v in attributes.items())
+        elegant_type = JSON_TO_ELE[type_]
+        strings.append(element_template(name, elegant_type, attrs))
 
-        lattice_names.append(name)
-        lattices_copy.pop(name)
+    lattice_template = "{}: LINE=({})".format
+    for name in sort_lattices(sub_lattices):
+        strings.append(lattice_template(name, ", ".join(sub_lattices[name])))
+
+    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
+    strings.append("USE, __MAIN__\n")
+    return "\n".join(strings)
 
-    for name, arrangement in lattices.items():
-        _sort_lattices(name, arrangement)
 
-    return lattice_names
+def to_madx(latticejson: dict) -> str:
+    """Convert a LatticeJSON dict to the MADX lattice file format.
+
+    :param lattice_dict dict: dict in LatticeJSON format
+    :return: string with in elegant lattice file format
+    """
+    elements = latticejson["elements"]
+    sub_lattices = latticejson["sub_lattices"]
+
+    strings = [f"TITLE, \"{latticejson['name']}\""]
+    element_template = "{}: {}, {};".format
+    for name, (type_, attributes) in elements.items():
+        attrs = ", ".join(f"{JSON_TO_MADX[k]}={v}" for k, v in attributes.items())
+        elegant_type = JSON_TO_MADX[type_]
+        strings.append(element_template(name, elegant_type, attrs))
+
+    lattice_template = "{}: LINE=({});".format
+    for name in sort_lattices(sub_lattices):
+        strings.append(lattice_template(name, ", ".join(sub_lattices[name])))
+
+    strings.append(lattice_template("__MAIN__", ", ".join(latticejson["lattice"])))
+    strings.append("USE, __MAIN__;\n")
+    return "\n".join(strings)
 
 
 def sort_lattices(lattices: Dict[str, List[str]]) -> List[str]:

latticejson/elegant.lark

@@ -1,3 +1,4 @@
+// Elegant Grammar
 // This grammer is not 100% consistent with elegants parser:
 // - Elegants parser allows tokens to be split by the line continuation character "&".
 //   For example, it parses ANGLE=0.123&\n456 without an error.
@@ -9,26 +10,17 @@
 %ignore /!.*/            // ingore comments
 %ignore /[ \t\f]/+       // ingore whitespace
 %ignore /&[ \t\f]*\r?\n/ // line continuation
-%import common.ESCAPED_STRING
-
-INT         :  ["+" | "-"] DIGIT_PART
-FLOAT       :  ["+" | "-"] (EXP_FLOAT | POINT_FLOAT)
-EXP_FLOAT   :  (DIGIT_PART | POINT_FLOAT) EXPONENT
-POINT_FLOAT :  [DIGIT_PART] FRACTION | DIGIT_PART "."
-DIGIT_PART  :  /\d+/
-FRACTION    :  "." DIGIT_PART
-EXPONENT    :  ("e" | "E") ["+" | "-"] DIGIT_PART
-
-int         : INT
-u_int       : DIGIT_PART
-float       : FLOAT
+%import common (SIGNED_INT, SIGNED_FLOAT, NUMBER, ESCAPED_STRING, CNAME)
+
+int         : SIGNED_INT
+float       : SIGNED_FLOAT
 string      : ESCAPED_STRING
 word        : /\w+/
 name        : /\w+/ | "\"" /[\w:]+/ "\""
 
-file        : _NEWLINE* (statement _NEWLINE+)*
+start        : _NEWLINE* (statement _NEWLINE+)*
 _NEWLINE    : /[ \t\f]*\r?\n[ \t\f]*/
-?statement  : element | lattice | command | rpn_store
+?statement  : element | lattice | command | "%" assignment
 
 element     : name ":" [name] ("," attribute)* ","?
 attribute   : word "=" (int | float | string | word)
@@ -40,4 +32,19 @@ ref_name    : [int "*"] [/-/] ["\""] /[\w:]+/ ["\""]
 
 command     : name ["," word]
 
-rpn_store   : "%" /.+sto.+/
+// RPN Expression
+// 1. As there is no syntactic distinction between an escaped string and a variable,
+// it is possible that a collison can happen. In this case a variable is wrongly
+// identified as string.
+
+assignment  : expr "sto" CNAME
+?expr       : NUMBER        -> number
+            | CNAME         -> variable
+            | function
+            | binary
+!function   : expr ("exp" | "sin" | "cos" | "tan" | "asin" | "acos" | "atan")
+?binary     : expr expr "+" -> add
+            | expr expr "-" -> sub
+            | expr expr "*" -> mul
+            | expr expr "/" -> div
+?start_rpn   : assignment | expr // used to tested the rpn parser

latticejson/exceptions.py

@@ -12,8 +12,24 @@ def __init__(self, object_name, lattice_name):
         )
 
 
-class UndefinedRPNVariableError(Exception):
+class UnknownElementWarning(UserWarning):
+    """Raised if there is no equivalent LatticeJSON element."""
+
+    def __init__(self, name, type_, *args, **kwargs):
+        message = f"Replacing element {name} ({type_}) with Drift."
+        super().__init__(message, *args, **kwargs)
+
+
+class UnknownAttributeWarning(UserWarning):
+    """Raised if there is no equivalent LatticeJSON attribute."""
+
+    def __init__(self, attribute, element, *args, **kwargs):
+        message = f"Ignoring attribute {attribute} of {element}."
+        super().__init__(message, *args, **kwargs)
+
+
+class UndefinedVariableError(Exception):
     """Raised if a rpn variable is not defined."""
 
     def __init__(self, name, *args, **kwargs):
-        super().__init__(f"RPN variable {name} is not defined!", *args, **kwargs)
+        super().__init__(f"The variable '{name}' is not defined!", *args, **kwargs)

latticejson/format.py

@@ -1,6 +1,11 @@
 import json
 
 
+def format_json(obj):
+    """Compact version of json.dumps."""
+    return json.dumps(obj, cls=CompactJSONEncoder, indent=4)
+
+
 class CompactJSONEncoder(json.JSONEncoder):
     """A JSON Encoder which only indents the first two levels."""