Simplify the decay file parsing grammar and implement more get methods for EvtGen keywords + unit tests

src/decaylanguage/data/decfile.lark

@@ -6,31 +6,37 @@
 start : _NEWLINE* (line _NEWLINE+)* ("End" _NEWLINE+)?
 ?line : define | particle_def | pythia_def | jetset_def | ls_def | model_alias | alias | chargeconj | commands | decay | cdecay | copydecay | setlspw | setlsbw | changemasslimit | inc_factor
 
-pythia_def : ("PythiaBothParam" | "PythiaAliasParam") LABEL ":" LABEL "=" (LABEL | SIGNED_NUMBER)
+pythia_def : LABEL_PYTHIA_PARAMETERS LABEL ":" LABEL "=" (LABEL | SIGNED_NUMBER)
+
+LABEL_PYTHIA_PARAMETERS : "PythiaBothParam" | "PythiaAliasParam"
 
 jetset_def : "JetSetPar" LABEL "=" SIGNED_NUMBER
 
-ls_def : ("LSFLAT" | "LSNONRELBW") label // Choose a lineshape for a particle
+ls_def : LABEL_LINESHAPE LABEL // Choose a lineshape for a particle
+
+LABEL_LINESHAPE : "LSFLAT" | "LSNONRELBW" | "LSMANYDELTAFUNC"
 
 inc_factor: ("IncludeBirthFactor" | "IncludeDecayFactor") label ("yes" | "no") // Presence of the birth/decay factor and form-factor
 
-setlsbw : "BlattWeisskopf" label value // Set Blatt-Weisskopf barrier factor for a lineshape
+setlsbw : "BlattWeisskopf" LABEL SIGNED_NUMBER // Set Blatt-Weisskopf barrier factor for a lineshape
 
 setlspw : "SetLineshapePW" LABEL LABEL LABEL INT // Redefine Partial Wave for label -> label label
 
-cdecay : "CDecay" label
+cdecay : "CDecay" LABEL
 
-define : "Define" label value
+define : "Define" LABEL SIGNED_NUMBER
 
 particle_def: "Particle" label value value // Set the mass and width of a particle (in GeV)
 
 model_alias : "ModelAlias" model_label model
 
-alias : "Alias" label label
+alias : "Alias" LABEL LABEL
+
+chargeconj : "ChargeConj" LABEL LABEL
 
-chargeconj : "ChargeConj" label label
+changemasslimit : LABEL_CHANGE_MASS LABEL SIGNED_NUMBER // Set upper/lower mass cuts on a lineshape
 
-changemasslimit : ("ChangeMassMin" | "ChangeMassMax") label value // Set upper/lower mass cuts on a lineshape
+LABEL_CHANGE_MASS : "ChangeMassMin" | "ChangeMassMax"
 
 ?commands : global_photos
 

src/decaylanguage/dec/dec.py

@@ -382,15 +382,16 @@ def get_particle_property_definitions(self) -> dict[str, dict[str, float]]:
         self._check_parsing()
         return get_particle_property_definitions(self._parsed_dec_file)
 
-    def dict_pythia_definitions(self) -> dict[str, str | float]:
+    def dict_pythia_definitions(self) -> dict[str, dict[str, str | float]]:
         """
-        Return a dictionary of all Pythia definitions, of type
-        <PYTHIA_DEF> = "PythiaBothParam" or "PythiaAliasParam",
-        in the input parsed file, of the form
+        Return a dictionary of all Pythia definitions, with keys corresponding to the types
+        <PYTHIA_DEF> = "PythiaBothParam" and/or "PythiaAliasParam",
+        set in the input parsed file with statements of the form
         "<PYTHIA_DEF> <NAME>=<LABEL>"
         or
         "<PYTHIA_DEF> <NAME>=<NUMBER>",
-        as {'NAME1': 'LABEL1', 'NAME2': VALUE2, ...}.
+        as {"PythiaAliasParam": {'NAME1': 'LABEL1', 'NAME2': VALUE2, ...},
+            "PythiaBothParam": {'NAME3': 'LABEL3', 'NAME4': VALUE4, ...}}.
         """
         self._check_parsing()
         return get_pythia_definitions(self._parsed_dec_file)
@@ -418,7 +419,26 @@ def expand_decay_modes(self, particle: str) -> list[str]:
         aliases = self.dict_aliases()
         return _expand_decay_modes(decay_chains, aliases=aliases)
 
-    def list_lineshape_definitions(self) -> list[tuple[list[str], int]]:
+    def dict_lineshape_settings(self) -> dict[str, dict[str, str | float]]:
+        """
+        Return a dictionary of all lineshape settings,
+        with keys corresponding to particle names or aliases, as
+        {PARTICLE1: {'lineshape': 'NAME1',  # E.g. "LSFLAT" or "LSNONRELBW"
+          'BlattWeisskopf': VALUE1,
+          'ChangeMassMin': VALUE12,
+          'ChangeMassMax': VALUE13},
+         PARTICLE2: {'lineshape': 'NAME2',
+          'BlattWeisskopf': VALUE2,
+          'ChangeMassMin': VALUE22,
+          'ChangeMassMax': VALUE23},
+         ...
+        }
+        where not all "sub-dictionaries" may contain all and/or the same keys.
+        """
+        self._check_parsing()
+        return get_lineshape_settings(self._parsed_dec_file)
+
+    def list_lineshapePW_definitions(self) -> list[tuple[list[str], int]]:
         """
         Return a list of all SetLineshapePW definitions in the input parsed file,
         of the form
@@ -429,7 +449,7 @@ def list_lineshape_definitions(self) -> list[tuple[list[str], int]]:
         ...]
         """
         self._check_parsing()
-        return get_lineshape_definitions(self._parsed_dec_file)
+        return get_lineshapePW_definitions(self._parsed_dec_file)
 
     def global_photos_flag(self) -> int:
         """
@@ -1390,8 +1410,7 @@ def get_charge_conjugate_decays(parsed_file: Tree) -> list[str]:
 
     try:
         return sorted(
-            tree.children[0].children[0].value
-            for tree in parsed_file.find_data("cdecay")
+            tree.children[0].value for tree in parsed_file.find_data("cdecay")
         )
     except Exception as err:
         raise RuntimeError(
@@ -1434,9 +1453,7 @@ def get_definitions(parsed_file: Tree) -> dict[str, float]:
     """
     try:
         return {
-            tree.children[0]
-            .children[0]
-            .value: float(tree.children[1].children[0].value)
+            tree.children[0].value: float(tree.children[1].value)
             for tree in parsed_file.find_data("define")
         }
     except Exception as err:
@@ -1485,7 +1502,7 @@ def get_aliases(parsed_file: Tree) -> dict[str, str]:
     """
     try:
         return {
-            tree.children[0].children[0].value: tree.children[1].children[0].value
+            tree.children[0].value: tree.children[1].value
             for tree in parsed_file.find_data("alias")
         }
     except Exception as err:
@@ -1507,7 +1524,7 @@ def get_charge_conjugate_defs(parsed_file: Tree) -> dict[str, str]:
     """
     try:
         return {
-            tree.children[0].children[0].value: tree.children[1].children[0].value
+            tree.children[0].value: tree.children[1].value
             for tree in parsed_file.find_data("chargeconj")
         }
     except Exception as err:
@@ -1544,28 +1561,41 @@ def get_particle_property_definitions(parsed_file: Tree) -> dict[str, dict[str,
         ) from err
 
 
-def get_pythia_definitions(parsed_file: Tree) -> dict[str, str | float]:
+def get_pythia_definitions(parsed_file: Tree) -> dict[str, dict[str, str | float]]:
     """
-    Return a dictionary of all Pythia definitions, of type
-    <PYTHIA_DEF> = "PythiaBothParam" or "PythiaAliasParam",
-    in the input parsed file, of the form
+    Return a dictionary of all Pythia definitions, with keys corresponding to the types
+    <PYTHIA_DEF> = "PythiaBothParam" and/or "PythiaAliasParam",
+    set in the input parsed file with statements of the form
     "<PYTHIA_DEF> <NAME>=<LABEL>"
     or
     "<PYTHIA_DEF> <NAME>=<NUMBER>",
-    as {'NAME1': 'LABEL1', 'NAME2': VALUE2, ...}.
+    as {"PythiaAliasParam": {'NAME1': 'LABEL1', 'NAME2': VALUE2, ...},
+        "PythiaBothParam": {'NAME3': 'LABEL3', 'NAME4': VALUE4, ...}}.
 
     Parameters
     ----------
     parsed_file: Lark Tree instance
         Input parsed file.
     """
+    d: dict[str, dict[str, str | float]] = {}
+
     try:
-        return {
-            f"{tree.children[0].value}:{tree.children[1].value}": _str_or_float(
-                tree.children[2].value
-            )
-            for tree in parsed_file.find_data("pythia_def")
-        }
+        for tree in parsed_file.find_data("pythia_def"):
+            if tree.children[0].value in d:
+                d[tree.children[0].value].update(
+                    {
+                        f"{tree.children[1].value}:{tree.children[2].value}": _str_or_float(
+                            tree.children[3].value
+                        )
+                    }
+                )
+            else:
+                d[tree.children[0].value] = {
+                    f"{tree.children[1].value}:{tree.children[2].value}": _str_or_float(
+                        tree.children[3].value
+                    )
+                }
+        return d
     except Exception as err:
         raise RuntimeError(
             "Input parsed file does not seem to have the expected structure."
@@ -1631,7 +1661,77 @@ def to_int_or_float(n: str) -> int | float | str:
         ) from err
 
 
-def get_lineshape_definitions(
+def get_lineshape_settings(
+    parsed_file: Tree,
+) -> dict[str, dict[str, str | float]]:
+    """
+     Return a dictionary of all lineshape settings,
+     with keys corresponding to particle names or aliases, as
+     {PARTICLE1: {'lineshape': 'NAME1',  # E.g. "LSFLAT" or "LSNONRELBW"
+       'BlattWeisskopf': VALUE11,
+       'ChangeMassMin': VALUE12,
+       'ChangeMassMax': VALUE13},
+      PARTICLE2: {'lineshape': 'NAME2',
+       'BlattWeisskopf': VALUE21,
+       'ChangeMassMin': VALUE22,
+       'ChangeMassMax': VALUE23},
+       ...
+     }
+     where not all "sub-dictionaries" may contain all and/or the same keys.
+
+    Parameters
+     ----------
+     parsed_file: Lark Tree instance
+         Input parsed file.
+    """
+    d: dict[str, dict[str, str | float]] = {}
+    try:
+        # Lineshape definitions
+        for tree in parsed_file.find_data("ls_def"):
+            if tree.children[1].value not in d:
+                d[tree.children[1].value] = {"lineshape": tree.children[0].value}
+            else:
+                raise RuntimeError(
+                    "Input parsed file does not seem to have the expected structure for the lineshape definitions."
+                ) from None
+
+        # Blatt-Weisskopf barrier factor for a lineshape
+        for tree in parsed_file.find_data("setlsbw"):
+            particle_or_alias = tree.children[0].value
+            if particle_or_alias in d:
+                if "BlattWeisskopf" in d[particle_or_alias]:
+                    raise RuntimeError(
+                        f"The Blatt-Weisskopf barrier factor for particle/alias {particle_or_alias} seems to be redefined."
+                    ) from None
+                d[particle_or_alias]["BlattWeisskopf"] = float(tree.children[1].value)
+            else:
+                d[particle_or_alias] = {"BlattWeisskopf": float(tree.children[1].value)}
+
+        # Upper and lower mass cuts on a lineshape
+        for tree in parsed_file.find_data("changemasslimit"):
+            particle_or_alias = tree.children[1].value
+            if particle_or_alias in d:
+                if tree.children[0].value in d[particle_or_alias]:
+                    raise RuntimeError(
+                        f"The upper/lower mass cut on the lineshape for particle/alias {particle_or_alias} seems to be redefined."
+                    ) from None
+                d[particle_or_alias][f"{tree.children[0].value}"] = float(
+                    tree.children[2].value
+                )
+            else:
+                d[particle_or_alias] = {
+                    f"{tree.children[0].value}": float(tree.children[2].value)
+                }
+
+        return d
+
+    except Exception as err:
+        raise RuntimeError(
+            "Input parsed file does not seem to have the expected structure."
+        ) from err
+
+
+def get_lineshapePW_definitions(
     parsed_file: Tree,
 ) -> list[tuple[list[str], int]]:
     """

tests/dec/test_dec.py

@@ -202,12 +202,16 @@ def test_pythia_definitions_parsing():
     p.parse()
 
     assert p.dict_pythia_definitions() == {
-        "ParticleDecays:mixB": "off",
-        "Init:showChangedSettings": "off",
-        "Init:showChangedParticleData": "off",
-        "Next:numberShowEvent": 0.0,
-        "ParticleDecays:sophisticatedTau": 3,
-        "ParticleDecays:tauPolarization": -1.0,
+        "PythiaAliasParam": {
+            "ParticleDecays:sophisticatedTau": 3.0,
+            "ParticleDecays:tauPolarization": -1.0,
+        },
+        "PythiaBothParam": {
+            "Init:showChangedParticleData": "off",
+            "Init:showChangedSettings": "off",
+            "Next:numberShowEvent": 0.0,
+            "ParticleDecays:mixB": "off",
+        },
     }
 
 
@@ -223,11 +227,36 @@ def test_jetset_definitions_parsing():
     }
 
 
-def test_list_lineshape_definitions():
+def test_dict_lineshape_settings():
     p = DecFileParser(DIR / "../data/defs-aliases-chargeconj.dec")
     p.parse()
 
-    assert p.list_lineshape_definitions() == [
+    assert p.dict_lineshape_settings() == {
+        "MyK*0": {
+            "lineshape": "LSNONRELBW",
+            "BlattWeisskopf": 0.0,
+            "ChangeMassMin": 0.5,
+            "ChangeMassMax": 3.5,
+        },
+        "MyPhi": {
+            "lineshape": "LSNONRELBW",
+            "BlattWeisskopf": 0.0,
+            "ChangeMassMin": 1.0,
+            "ChangeMassMax": 1.04,
+        },
+        "MyKS0pipi": {
+            "lineshape": "LSFLAT",
+            "ChangeMassMin": 1.1,
+            "ChangeMassMax": 2.4,
+        },
+    }
+
+
+def test_list_lineshapePW_definitions():
+    p = DecFileParser(DIR / "../data/defs-aliases-chargeconj.dec")
+    p.parse()
+
+    assert p.list_lineshapePW_definitions() == [
         (["D_1+", "D*+", "pi0"], 2),
         (["D_1+", "D*0", "pi+"], 2),
         (["D_1-", "D*-", "pi0"], 2),