Add several methods to handle decay descriptors (#331)

* add DecayChain.to_string method plus tests

* style: pre-commit fixes

* Update src/decaylanguage/decay/decay.py

Co-authored-by: Eduardo Rodrigues <[email protected]>

* style: pre-commit fixes

* add example to docstring, fix typo in test

* Update src/decaylanguage/decay/decay.py

Co-authored-by: Eduardo Rodrigues <[email protected]>

* use DaughtersDict

* enforce list type

* functionality to expand decay chain dicts

* Update src/decaylanguage/dec/dec.py

Co-authored-by: Eduardo Rodrigues <[email protected]>

* Apply suggestions from code review

* add DescriptorFormat class for customising the format of descriptors

* Apply suggestions from code review

Co-authored-by: Eduardo Rodrigues <[email protected]>

* add FLATSQDALITZ to known EvtGen model names

* Apply suggestions from code review

* move DescriptorFormat test to tests/utils

* mention implicit use of aliases in docstring for DecFileParser.expand_decay_modes

* example of aliases in _expand_decay_modes

* add import statement to example to satisfy doctest(?)

* comply with RUF012

* use typing.Dict in typedef

* one more

* Update decay.py

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Eduardo Rodrigues <[email protected]>

src/decaylanguage/dec/dec.py

@@ -54,6 +54,7 @@
 from particle.converters import PDG2EvtGenNameMap
 
 from .. import data
+from ..decay.decay import _expand_decay_modes
 from ..utils import charge_conjugate_name
 from .enums import PhotosEnum
 
@@ -390,6 +391,17 @@ def dict_jetset_definitions(self) -> dict[str, dict[int, int | float | str]]:
         self._check_parsing()
         return get_jetset_definitions(self._parsed_dec_file)
 
+    def expand_decay_modes(self, particle: str) -> list[str]:
+        """
+        Return a list of expanded decay descriptors for the given (mother) particle.
+        The set of decay final states is effectively split and returned as a list.
+        NB: this implicitly reverts aliases back to the original (EvtGen) names.
+        """
+        self._check_parsing()
+        decay_chains = self.build_decay_chains(particle)
+        aliases = self.dict_aliases()
+        return _expand_decay_modes(decay_chains, aliases=aliases)
+
     def list_lineshape_definitions(self) -> list[tuple[list[str], int]]:
         """
         Return a list of all SetLineshapePW definitions in the input parsed file,

src/decaylanguage/decay/decay.py

@@ -8,13 +8,14 @@
 
 from collections import Counter
 from copy import deepcopy
+from itertools import product
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, TypeVar, Union
 
 from particle import PDGID, ParticleNotFound
 from particle.converters import EvtGenName2PDGIDBiMap
 from particle.exceptions import MatchingIDNotFound
 
-from ..utils import charge_conjugate_name
+from ..utils import DescriptorFormat, charge_conjugate_name
 
 Self_DaughtersDict = TypeVar("Self_DaughtersDict", bound="DaughtersDict")
 
@@ -446,7 +447,8 @@ def __str__(self) -> str:
 
 
 Self_DecayChain = TypeVar("Self_DecayChain", bound="DecayChain")
-DecayModeDict = Dict[str, List[Dict[str, Union[float, str, List[Any]]]]]
+DecayModeDict = Dict[str, Union[float, str, List[Any]]]
+DecayChainDict = Dict[str, List[DecayModeDict]]
 
 
 def _has_no_subdecay(ds: list[Any]) -> bool:
@@ -466,7 +468,7 @@ def _has_no_subdecay(ds: list[Any]) -> bool:
 
 
 def _build_decay_modes(
-    decay_modes: dict[str, DecayMode], dc_dict: DecayModeDict
+    decay_modes: dict[str, DecayMode], dc_dict: DecayChainDict
 ) -> None:
     """
     Internal recursive function that identifies and creates all `DecayMode` instances
@@ -505,6 +507,160 @@ def _build_decay_modes(
             decay_modes[mother] = DecayMode.from_dict(d)
 
 
+def _expand_decay_modes(
+    decay_chain: DecayChainDict,
+    *,
+    top: bool = True,
+    aliases: dict[str, str] | None = None,
+) -> list[str]:
+    """Given a dict with 1 key (the mother particle) whose value is a list of
+    decay modes, recursively replace all decay modes with decay descriptors.
+
+    Parameters
+    ----------
+    decay_chain: dict
+        A dict representing decay chains, such as returned by DecayChain.to_dict
+        or DecFileParser.build_decay_chains.
+    top: bool, optional, default=True
+        Whether the passed decay chain is the top-level or not (should usually
+        be True: only really set to False when recursing the function).
+    aliases: dict[str, str], optional, default={}
+        Mapping of names to replace. Useful when dealing with DecFiles that have
+        Alias statements.
+
+    Examples
+    --------
+    A simple example with no sub-decays:
+    {
+        "anti-D0": [
+            {
+              "bf": 1.0,
+              "fs": ["K+", "pi-"],
+              "model": "PHSP",
+              "model_params": ""
+            }
+        ]
+    }
+    becomes the dict
+    {
+        "anti-D0": [
+            "anti-D0 -> K+ pi-"
+        ]
+    }
+
+    A more complicated example with a sub-decay and more than one mode:
+    {
+        "anti-D*0": [
+            {
+                "bf": 0.619,
+                "fs": [
+                    {
+                        "anti-D0": [
+                            {
+                                "bf": 1.0,
+                                "fs": ["K+", "pi-"],
+                                "model": "PHSP",
+                                "model_params": ""
+                            }
+                        ]
+                    },
+                    "pi0"
+                ],
+                "model": "VSS",
+                "model_params": ""
+            },
+            {
+                "bf": 0.381,
+                "fs": [
+                    {
+                        "anti-D0": [
+                            {
+                                "bf": 1.0,
+                                "fs": ["K+", "pi-"],
+                                "model": "PHSP",
+                                "model_params": ""
+                            }
+                        ]
+                    },
+                    "gamma"
+                ],
+                "model": "VSP_PWAVE",
+                "model_params": ""
+            }
+        ]
+    }
+    becomes the dict
+    {
+        "anti-D*0": [
+            "anti-D*0 -> (anti-D0 -> K+ pi-) pi0",
+            "anti-D*0 -> (anti-D0 -> K+ pi-) gamma",
+        ]
+    }
+
+    and an example alias dict:
+    {"MyAntiD0": "anti-D0"}
+    can be used with
+    {
+        "MyAntiD0": [
+            {
+              "bf": 1.0,
+              "fs": ["K+", "pi-"],
+              "model": "PHSP",
+              "model_params": ""
+            }
+        ]
+    }
+    to result in
+    {
+        'MyAntiD0': [
+            'anti-D0 -> K+ pi-'
+        ]
+    }
+    """
+
+    def _get_modes(decay_chain: DecayChainDict) -> list[DecayModeDict]:
+        # The list of decay modes is the first (and only) value of the dict
+        assert len(decay_chain.values()) == 1
+        modes = list(decay_chain.values())
+        return modes[0]
+
+    def _get_fs(decay: DecayModeDict) -> list[Any]:
+        fs = decay["fs"]
+        if isinstance(fs, list):
+            return fs
+        raise TypeError(f"Expected list, not {type(fs)}")
+
+    # The mother particle is the first (and only) key of the dict
+    assert len(decay_chain.keys()) == 1
+    orig_mother = list(decay_chain.keys())[0]
+    mother = aliases.get(orig_mother, orig_mother) if aliases else orig_mother
+
+    for mode in _get_modes(decay_chain):
+        for fsp in _get_fs(mode):
+            if isinstance(fsp, dict):
+                _expand_decay_modes(fsp, top=False, aliases=aliases)
+
+    # Replace dicts with strings (decay descriptors)
+    expanded_modes = []
+    for mode in _get_modes(decay_chain):
+        fsp_options = []
+        for fsp in _get_fs(mode):
+            if isinstance(fsp, dict):
+                fsp_options += [_get_modes(fsp)]
+            elif isinstance(fsp, str):
+                fsp_options += [[fsp]]  # type: ignore[list-item]
+        for expanded_mode in product(*fsp_options):
+            # TODO: delegate descriptor-building to another function
+            #       allow for different conventions?
+            final_state = DaughtersDict(list(expanded_mode)).to_string()
+            descriptor = DescriptorFormat.format_descriptor(mother, final_state, top)
+            expanded_modes += [descriptor]
+
+    decay_chain[orig_mother] = expanded_modes  # type: ignore[assignment]
+
+    return expanded_modes
+
+
 class DecayChain:
     """
     Class holding a particle decay chain, which is typically a top-level decay
@@ -553,7 +709,7 @@ def __init__(self, mother: str, decays: dict[str, DecayMode]) -> None:
 
     @classmethod
     def from_dict(
-        cls: type[Self_DecayChain], decay_chain_dict: DecayModeDict
+        cls: type[Self_DecayChain], decay_chain_dict: DecayChainDict
     ) -> Self_DecayChain:
         """
         Constructor from a decay chain represented as a dictionary.
@@ -602,6 +758,26 @@ def ndecays(self) -> int:
         """
         return len(self.decays)
 
+    def to_string(self) -> str:
+        """
+        One-line string representation of the entire decay chain.
+        Sub-decays are enclosed in round parentheses.
+
+        Examples
+        --------
+        >>> dm1 = DecayMode(0.6770, "D0 pi+") # D*+
+        >>> dm2 = DecayMode(0.0124, "K_S0 pi0") # D0
+        >>> dm3 = DecayMode(0.692, "pi+ pi-") # K_S0
+        >>> dm4 = DecayMode(0.98823, "gamma gamma") # pi0
+        >>> dc = DecayChain("D*+", {"D*+":dm1, "D0":dm2, "K_S0":dm3, "pi0":dm4})
+        >>> print(dc.to_string())
+        D*+ -> (D0 -> (K_S0 -> pi+ pi-) (pi0 -> gamma gamma)) pi+
+        """
+        dc_dict = self.to_dict()
+        descriptors = _expand_decay_modes(dc_dict, top=True)
+        assert len(descriptors) == 1
+        return descriptors[0]
+
     def print_as_tree(self) -> None:  # pragma: no cover
         """
         Tree-structure like print of the entire decay chain.
@@ -654,7 +830,7 @@ def print_as_tree(self) -> None:  # pragma: no cover
 
         # TODO: simplify logic and perform further checks
         def _print(
-            decay_dict: dict[str, list[dict[str, float | str | list[Any]]]],
+            decay_dict: DecayChainDict,
             depth: int = 0,
             link: bool = False,
             last: bool = False,
@@ -686,7 +862,7 @@ def _print(
         dc_dict = self.to_dict()
         _print(dc_dict)
 
-    def to_dict(self) -> dict[str, list[dict[str, float | str | list[Any]]]]:
+    def to_dict(self) -> DecayChainDict:
         """
         Return the decay chain as a dictionary representation.
         The format is the same as `DecFileParser.build_decay_chains(...)`.
@@ -708,7 +884,7 @@ def to_dict(self) -> dict[str, list[dict[str, float | str | list[Any]]]]:
             'model_params': ''}]}
         """
 
-        # Ideally this would be a recursive type, DecayDict = dict[str, list[str | DecayDict]]
+        # Ideally this would be a recursive type, DecayDict = Dict[str, list[str | DecayDict]]
         DecayDict = Dict[str, List[Any]]
 
         def recursively_replace(mother: str) -> DecayDict:

src/decaylanguage/utils/__init__.py

@@ -7,9 +7,15 @@
 
 from .errors import LineFailure
 from .particleutils import charge_conjugate_name
-from .utilities import filter_lines, iter_flatten, split
+from .utilities import (
+    DescriptorFormat,
+    filter_lines,
+    iter_flatten,
+    split,
+)
 
 __all__ = (
+    "DescriptorFormat",
     "LineFailure",
     "iter_flatten",
     "split",