diff --git a/src/decaylanguage/dec/dec.py b/src/decaylanguage/dec/dec.py index e6cf7a0a..10f45448 100644 --- a/src/decaylanguage/dec/dec.py +++ b/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, diff --git a/src/decaylanguage/decay/decay.py b/src/decaylanguage/decay/decay.py index d756ccb5..99388b7c 100644 --- a/src/decaylanguage/decay/decay.py +++ b/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: diff --git a/src/decaylanguage/utils/__init__.py b/src/decaylanguage/utils/__init__.py index f48f6bcd..e15d7ae2 100644 --- a/src/decaylanguage/utils/__init__.py +++ b/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", diff --git a/src/decaylanguage/utils/utilities.py b/src/decaylanguage/utils/utilities.py index 24f6b861..7dc6e9a0 100644 --- a/src/decaylanguage/utils/utilities.py +++ b/src/decaylanguage/utils/utilities.py @@ -5,7 +5,9 @@ from __future__ import annotations -from typing import Any, Iterator, Pattern +import string +from copy import copy +from typing import Any, ClassVar, Iterator, Pattern def iter_flatten(iterable: list[str] | tuple[str, ...]) -> Iterator[str]: @@ -57,3 +59,102 @@ def filter_lines( output = [match.groupdict() for match in matches if match is not None] new_inp = [ln for ln in inp if matcher.match(ln) is None] return output, new_inp + + +class DescriptorFormat: + """ + Class to help with setting the decay descriptor format. The format is stored as a + class-level variable: `DescriptorFormat.config`. + + Examples + -------- + >>> from decaylanguage import DecayMode, DecayChain + >>> 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}) + >>> with DescriptorFormat("{mother} --> {daughters}", "[{mother} --> {daughters}]"): dc.to_string() + ... + 'D*+ --> [D0 --> [K_S0 --> pi+ pi-] [pi0 --> gamma gamma]] pi+' + >>> with DescriptorFormat("{mother} => {daughters}", "{mother} (=> {daughters})"): dc.to_string(); + ... + 'D*+ => D0 (=> K_S0 (=> pi+ pi-) pi0 (=> gamma gamma)) pi+' + >>> dc.to_string() + 'D*+ -> (D0 -> (K_S0 -> pi+ pi-) (pi0 -> gamma gamma)) pi+' + """ + + config: ClassVar[dict[str, str]] = { + "decay_pattern": "{mother} -> {daughters}", + "sub_decay_pattern": "({mother} -> {daughters})", + } + + def __init__(self, decay_pattern: str, sub_decay_pattern: str) -> None: + self.new_config = { + "decay_pattern": decay_pattern, + "sub_decay_pattern": sub_decay_pattern, + } + self.old_config = copy(DescriptorFormat.config) + + def __enter__(self) -> None: + self.set_config(**self.new_config) + + def __exit__(self, *args: list[Any]) -> None: + self.set_config(**self.old_config) + + @staticmethod + def set_config(decay_pattern: str, sub_decay_pattern: str) -> None: + """ + Configure the descriptor patterns after checking that each pattern + has named-wildcards "mother" and "daughters". + + Parameters + ---------- + + decay_pattern: str + Format-string expression for a top-level decay, + e.g. "{mother} -> {daughters}" + sub_decay_pattern: str + Format-string expression for a sub-decay, + e.g. "({mother} -> {daughters}" + """ + new_config = { + "decay_pattern": decay_pattern, + "sub_decay_pattern": sub_decay_pattern, + } + expected_wildcards = {"mother", "daughters"} + for pattern in new_config.values(): + wildcards = { + t[1] for t in string.Formatter().parse(pattern) if isinstance(t[1], str) + } + if wildcards != expected_wildcards: + error_msg = ( + "The pattern should only have the wildcards " + f"{expected_wildcards}, while '{pattern}' has the wildcards " + f"{wildcards}." + ) + raise ValueError(error_msg) + DescriptorFormat.config = new_config + + @staticmethod + def format_descriptor(mother: str, daughters: str, top: bool = True) -> str: + """ + Apply the format to one "layer" of the decay. Does not handle nested + decays itself. It is assumed that the `daughters` string already contains + any sub-decays. + + Parameters + ---------- + + mother: str + The decaying particle. + daughters: str + The final-state particles. + """ + args = { + "mother": mother, + "daughters": daughters, + } + if top: + return DescriptorFormat.config["decay_pattern"].format(**args) + return DescriptorFormat.config["sub_decay_pattern"].format(**args) diff --git a/tests/dec/test_dec.py b/tests/dec/test_dec.py index e7e280c0..f3da7275 100644 --- a/tests/dec/test_dec.py +++ b/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", diff --git a/tests/decay/test_descriptor.py b/tests/decay/test_descriptor.py new file mode 100644 index 00000000..51471a92 --- /dev/null +++ b/tests/decay/test_descriptor.py @@ -0,0 +1,47 @@ +# Copyright (c) 2018-2023, Eduardo Rodrigues and Henry Schreiner. +# +# Distributed under the 3-clause BSD license, see accompanying file LICENSE +# or https://github.com/scikit-hep/decaylanguage for details. + +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 diff --git a/tests/utils/test_utilities.py b/tests/utils/test_utilities.py new file mode 100644 index 00000000..8b652ae0 --- /dev/null +++ b/tests/utils/test_utilities.py @@ -0,0 +1,45 @@ +# Copyright (c) 2018-2023, Eduardo Rodrigues and Henry Schreiner. +# +# Distributed under the 3-clause BSD license, see accompanying file LICENSE +# or https://github.com/scikit-hep/decaylanguage for details. + +from __future__ import annotations + +import pytest + +from decaylanguage import DecayChain, DecayMode +from decaylanguage.utils import DescriptorFormat + +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 + + +@pytest.mark.parametrize( + ("decay_pattern", "sub_decay_pattern", "expected"), + [ + ( + "{mother} -> {daughters}", + "({mother} -> {daughters})", + "D*+ -> (D0 -> (K_S0 -> pi+ pi-) (pi0 -> gamma gamma)) pi+", + ), + ( + "{mother} --> {daughters}", + "[{mother} --> {daughters}]", + "D*+ --> [D0 --> [K_S0 --> pi+ pi-] [pi0 --> gamma gamma]] pi+", + ), + ( + "{mother} => {daughters}", + "{mother} (=> {daughters})", + "D*+ => D0 (=> K_S0 (=> pi+ pi-) pi0 (=> gamma gamma)) pi+", + ), + ], +) +def test_set_descriptor_pattern( + decay_pattern: str, sub_decay_pattern: str, expected: str +): + dc = DecayChain("D*+", {"D*+": dm1, "D0": dm2, "K_S0": dm3, "pi0": dm4}) + with DescriptorFormat(decay_pattern, sub_decay_pattern): + descriptor = dc.to_string() + assert descriptor == expected