Add several methods to handle decay descriptors

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,16 @@ 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.
+        """
+        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,7 +8,8 @@
 
 from collections import Counter
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, TypeVar, Union
+from itertools import product
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, TypeVar, Union
 
 from particle import PDGID, ParticleNotFound
 from particle.converters import EvtGenName2PDGIDBiMap
@@ -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,142 @@ 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
+
+    top: bool, optional, default=True
+
+    aliases: dict[str, str], optional, default={}
+
+
+    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",
+        ]
+    }
+    """
+
+    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 = f"{mother} -> {final_state}"
+            if not top:
+                descriptor = f"({descriptor})"
+            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 +691,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 +740,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 +812,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 +844,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(...)`.
@@ -709,7 +867,7 @@ def to_dict(self) -> dict[str, list[dict[str, float | str | list[Any]]]]:
         """
 
         # Ideally this would be a recursive type, DecayDict = dict[str, list[str | DecayDict]]
-        DecayDict = Dict[str, List[Any]]
+        DecayDict = dict[str, list[Any]]
 
         def recursively_replace(mother: str) -> DecayDict:
             dm = self.decays[mother].to_dict()

tests/dec/test_dec.py

@@ -537,6 +537,48 @@ def test_build_decay_chains():
     assert p.build_decay_chains("D+", stable_particles=["pi0"]) == output
 
 
+def test_expand_decay_chains():
+    p = DecFileParser(DIR / "../data/test_example_Dst.dec")
+    p.parse()
+
+    output = {
+        "D*+": [
+            "D*+ -> (D0 -> K- pi+) pi+",
+            "D*+ -> (D+ -> (pi0 -> gamma gamma) K- pi+ pi+) (pi0 -> gamma gamma)",
+            "D*+ -> (D+ -> (pi0 -> gamma gamma) K- pi+ pi+) (pi0 -> e+ e- gamma)",
+            "D*+ -> (D+ -> (pi0 -> gamma gamma) K- pi+ pi+) (pi0 -> e+ e+ e- e-)",
+            "D*+ -> (D+ -> (pi0 -> gamma gamma) K- pi+ pi+) (pi0 -> e+ e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e- gamma) K- pi+ pi+) (pi0 -> gamma gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e- gamma) K- pi+ pi+) (pi0 -> e+ e- gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e- gamma) K- pi+ pi+) (pi0 -> e+ e+ e- e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e- gamma) K- pi+ pi+) (pi0 -> e+ e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e+ e- e-) K- pi+ pi+) (pi0 -> gamma gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e+ e- e-) K- pi+ pi+) (pi0 -> e+ e- gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e+ e- e-) K- pi+ pi+) (pi0 -> e+ e+ e- e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e+ e- e-) K- pi+ pi+) (pi0 -> e+ e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e-) K- pi+ pi+) (pi0 -> gamma gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e-) K- pi+ pi+) (pi0 -> e+ e- gamma)",
+            "D*+ -> (D+ -> (pi0 -> e+ e-) K- pi+ pi+) (pi0 -> e+ e+ e- e-)",
+            "D*+ -> (D+ -> (pi0 -> e+ e-) K- pi+ pi+) (pi0 -> e+ e-)",
+            "D*+ -> (D+ -> (pi0 -> gamma gamma) K- pi+ pi+) gamma",
+            "D*+ -> (D+ -> (pi0 -> e+ e- gamma) K- pi+ pi+) gamma",
+            "D*+ -> (D+ -> (pi0 -> e+ e+ e- e-) K- pi+ pi+) gamma",
+            "D*+ -> (D+ -> (pi0 -> e+ e-) K- pi+ pi+) gamma",
+        ],
+        "D*-": [  # NB: decays of anti-D0 and D- are not defined in this DecFile
+            "D*- -> anti-D0 pi-",
+            "D*- -> (pi0 -> gamma gamma) D-",
+            "D*- -> (pi0 -> e+ e- gamma) D-",
+            "D*- -> (pi0 -> e+ e+ e- e-) D-",
+            "D*- -> (pi0 -> e+ e-) D-",
+            "D*- -> D- gamma",
+        ],
+    }
+    for particle in output:
+        all_decays = p.expand_decay_modes(particle)
+        assert set(all_decays) == set(output[particle])
+
+
 def test_Lark_ModelNameCleanup_Transformer_no_params():
     t = Tree(
         "decay",

tests/decay/test_descriptor.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import pytest
+
+from decaylanguage import DecayChain, DecayMode
+
+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
+dm5 = DecayMode(0.0105, "D- tau+ nu_tau")  # B0
+dm6 = DecayMode(0.0938, "K+ pi- pi-")  # D-
+dm7 = DecayMode(0.0931, "pi+ pi+ pi- anti-nu_tau")  # tau+
+dm8 = DecayMode(1.85e-6, "phi phi'")  #  B_s0
+dm9a = DecayMode(0.491, "K+ K-")  #  phi
+dm9b = DecayMode(0.154, "pi+ pi- pi0")  #  phi
+
+
+@pytest.mark.parametrize(
+    ("dc", "expected"),
+    [
+        (
+            DecayChain("D0", {"D0": dm2, "K_S0": dm3, "pi0": dm4}),
+            "D0 -> (K_S0 -> pi+ pi-) (pi0 -> gamma gamma)",
+        ),
+        (
+            DecayChain("D*+", {"D*+": dm1, "D0": dm2, "K_S0": dm3, "pi0": dm4}),
+            "D*+ -> (D0 -> (K_S0 -> pi+ pi-) (pi0 -> gamma gamma)) pi+",
+        ),
+        (
+            DecayChain("B0", {"B0": dm5, "D-": dm6, "tau+": dm7}),
+            "B0 -> (D- -> K+ pi- pi-) (tau+ -> anti-nu_tau pi+ pi+ pi-) nu_tau",
+        ),
+        (
+            DecayChain("B_s0", {"B_s0": dm8, "phi": dm9a, "phi'": dm9b}),
+            "B_s0 -> (phi -> K+ K-) (phi' -> pi+ pi- pi0)",
+        ),
+    ],
+)
+def test_descriptor(dc: DecayChain, expected: str):
+    descriptor = dc.to_string()
+    assert descriptor == expected