Merge pull request #71 from simonthor/particle_model_map

Enhance GenMultiDecay constructor

docs/GenMultiDecay_Tutorial.ipynb

@@ -13,8 +13,8 @@
     "# Tutorial for *GenMultiDecay* class\n",
     "This tutorial shows how ``phasespace.fromdecay.GenMultiDecay`` can be used.\n",
     "\n",
-    "In order to use this functionality, you need to install the extra `fromdecay`, for example through\n",
-    "``pip install phasespace[fromdecay]``.\n",
+    "In order to use this functionality, you need to install the extra dependencies, for example through\n",
+    "`pip install phasespace[fromdecay]`.\n",
     "\n",
     "This submodule makes it possible for `phasespace` and [`DecayLanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n",
     "More generally, `GenMultiDecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways."
@@ -35,7 +35,7 @@
     "\n",
     "import zfit\n",
     "from particle import Particle\n",
-    "from decaylanguage import DecFileParser, DecayChainViewer\n",
+    "from decaylanguage import DecFileParser, DecayChainViewer, DecayChain, DecayMode\n",
     "import tensorflow as tf\n",
     "\n",
     "from phasespace.fromdecay import GenMultiDecay"
@@ -98,6 +98,25 @@
     "DecayChainViewer(pi0_chain)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also create a decay using the `DecayChain` and `DecayMode` classes. However, a DecayChain can only contain one chain, i.e., a particle cannot decay in multiple ways."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dplus_decay = DecayMode(1, \"K- pi+ pi+ pi0\", model=\"PHSP\")\n",
+    "pi0_decay = DecayMode(1, \"gamma gamma\")\n",
+    "dplus_single = DecayChain(\"D+\", {\"D+\": dplus_decay, \"pi0\": pi0_decay})\n",
+    "DecayChainViewer(dplus_single.to_dict())"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -266,9 +285,45 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "If you want all $\\pi^0$ particles to decay with the same mass function, you do not need to specify the `zfit` parameter for each decay in the `dict`. Instead, one can pass the `particle_model_map` parameter to the constructor:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "GenMultiDecay.from_dict(dsplus_chain, particle_model_map={'pi0': 'gauss'})    # pi0 always decays with a gaussian mass distribution."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "When using `DecayChain`s, the syntax for specifying the mass function becomes cleaner:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dplus_decay = DecayMode(1, \"K- pi+ pi+ pi0\", model=\"PHSP\")  # The model parameter will be ignored by GenMultiDecay\n",
+    "pi0_decay = DecayMode(1, \"gamma gamma\", zfit=\"gauss\")   # Make pi0 have a gaussian mass distribution\n",
+    "dplus_single = DecayChain(\"D+\", {\"D+\": dplus_decay, \"pi0\": pi0_decay})\n",
+    "GenMultiDecay.from_dict(dplus_single.to_dict())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Custom mass functions\n",
     "The built-in supported mass function names are `gauss`, `bw`, and `relbw`, with `gauss` being the gaussian distribution, `bw` being the Breit-Wigner distribution, and `relbw` being the relativistic Breit-Wigner distribution.\n",
     "\n",
-    "If a non-supported value for the `zfit` parameter is used or if it is not specified, it will automatically use the relativistic Breit-Wigner distribution. This behavior can be changed by changing the value of `GenMultiDecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`.\n",
+    "If a non-supported value for the `zfit` parameter is not specified, it will automatically use the relativistic Breit-Wigner distribution. This behavior can be changed by changing the value of `GenMultiDecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`. If an invalid value for the `zfit` parameter is used, a `KeyError` is raised.\n",
     "\n",
     "It is also possible to add your own mass functions besides the built-in ones. You should then create a function that takes the mass and width of a particle and returns a mass function which with the [format](https://phasespace.readthedocs.io/en/stable/usage.html#resonances-with-variable-mass) that is used for all phasespace mass functions. Below is an example of a custom gaussian distribution (implemented in the same way as the built-in gaussian distribution), which uses `zfit` PDFs:"
    ]
@@ -362,4 +417,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}

phasespace/fromdecay/genmultidecay.py

@@ -2,7 +2,6 @@
 
 import itertools
 from collections.abc import Callable
-from typing import Union
 
 import tensorflow as tf
 import tensorflow.experimental.numpy as tnp
@@ -31,7 +30,8 @@ def from_dict(
         cls,
         dec_dict: dict,
         mass_converter: dict[str, Callable] = None,
-        tolerance: float = MASS_WIDTH_TOLERANCE,
+        tolerance: float = None,
+        particle_model_map: dict[str, str] = None,
     ):
         """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format.
 
@@ -45,9 +45,14 @@ def from_dict(
                 This dict will be combined with the predefined mass functions in this package.
                 See the Example below or the tutorial for how to use this parameter.
             tolerance: Minimum mass width of the particle to use a mass function instead of
-                assuming the mass to be constant. The default value is defined by the class variable
-                MASS_WIDTH_TOLERANCE and can be customized if desired.
-
+                assuming the mass to be constant. If None, the default value, defined by the class variable
+                MASS_WIDTH_TOLERANCE, will be used. This value can be customized if desired.
+            particle_model_map: A dict where the key is a particle name and the value is a mass function name.
+                If a particle is specified in the particle_model_map, then all appearances of this particle
+                in dec_dict will get the same mass function. This way, one does not have to manually add the
+                zfit parameter to every place where this particle appears in dec_dict. If the zfit parameter
+                is specified for a particle which is also included in particle_model_map, the zfit parameter
+                mass function name will be prioritized.
         Returns:
             The created GenMultiDecay object.
 
@@ -70,8 +75,13 @@ def from_dict(
                 {'bf': 0.016,
                  'fs': ['D+', 'gamma']}]}
 
-            If the D0 particle should have a mass distribution of a gaussian when it decays into K- and pi+
-            a `zfit` parameter can be added to its decay dict:
+            If the D0 particle should have a mass distribution of a gaussian when it decays, one can pass the
+            `particle_model_map` parameter to `from_dict`:
+            >>> dst_gen = GenMultiDecay.from_dict(dst_chain, particle_model_map={"D0": "gauss"})
+            This will then set the mass function of D0 to a gaussian for all its decays.
+
+            If more custom control is required, e.g., if D0 can decay in multiple ways and one of the decays
+            should have a specific mass function, a `zfit` parameter can be added to its decay dict:
             >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "gauss"
             >>> dst_chain
             {'D*+': [{'bf': 0.984,
@@ -81,11 +91,13 @@ def from_dict(
                     'pi+']},
                 {'bf': 0.016,
                 'fs': ['D+', 'gamma']}]}
-
             This dict can then be passed to `GenMultiDecay.from_dict`:
             >>> dst_gen = GenMultiDecay.from_dict(dst_chain)
+            This will now convert make the D0 particle have a gaussian mass function, only when it decays into
+            K- and pi+. In this case, there are no other ways that D0 can decay, so using `particle_model_map`
+            is a cleaner and easier option.
 
-            If the decay of the D0 particle instead should be modelled by a mass distribution that does not
+            If the decay of the D0 particle should be modelled by a mass distribution that does not
             come with the package, a custom distribution can be created:
             >>> def custom_gauss(mass, width):
             >>>     particle_mass = tf.cast(mass, tf.float64)
@@ -113,23 +125,29 @@ def from_dict(
             For a more in-depth tutorial, see the tutorial on GenMultiDecay in the
             [documentation](https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html).
         """
+        if tolerance is None:
+            tolerance = cls.MASS_WIDTH_TOLERANCE
+        if particle_model_map is None:
+            particle_model_map = {}
         if mass_converter is None:
             total_mass_converter = DEFAULT_CONVERTER
         else:
             # Combine the default mass functions specified with the mass functions from input.
             total_mass_converter = {**DEFAULT_CONVERTER, **mass_converter}
 
         gen_particles = _recursively_traverse(
-            dec_dict, total_mass_converter, tolerance=tolerance
+            dec_dict,
+            total_mass_converter,
+            particle_model_map=particle_model_map,
+            tolerance=tolerance,
         )
         return cls(gen_particles)
 
     def generate(
         self, n_events: int, normalize_weights: bool = True, **kwargs
-    ) -> (
-        tuple[list[tf.Tensor], list[tf.Tensor]]
-        | tuple[list[tf.Tensor], list[tf.Tensor], list[tf.Tensor]]
-    ):
+    ) -> tuple[list[tf.Tensor], list[tf.Tensor]] | tuple[
+        list[tf.Tensor], list[tf.Tensor], list[tf.Tensor]
+    ]:
         """Generate four-momentum vectors from the decay(s).
 
         Args:
@@ -197,13 +215,15 @@ def _get_particle_mass(
     name: str,
     mass_converter: dict[str, Callable],
     mass_func: str,
-    tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE,
+    tolerance: float,
 ) -> Callable | float:
     """Get mass or mass function of particle using the particle package.
 
     Args:
         name: Name of the particle. Name must be recognizable by the particle package.
-        tolerance : See _recursively_traverse
+        mass_converter: See _recursively_traverse
+        mass_func: See the name of the mass function, e.g., 'rel-bw'. Must be a valid key for mass_converter.
+        tolerance: See _recursively_traverse
 
     Returns:
         A function if the mass has a width smaller than tolerance. Otherwise, return a constant mass.
@@ -220,19 +240,25 @@ def _get_particle_mass(
 def _recursively_traverse(
     decaychain: dict,
     mass_converter: dict[str, Callable],
+    particle_model_map: dict[str, str],
+    tolerance: float,
     preexisting_particles: set[str] = None,
-    tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE,
 ) -> list[tuple[float, GenParticle]]:
     """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples.
 
     Args:
         decaychain: Decay chain with the format from DecayLanguage
+        mass_converter: Maps from mass function names to the actual callable mass function.
+        particle_model_map: See GenMultiDecay.from_dict.
         preexisting_particles: Names of all particles that have already been created.
         tolerance: Minimum mass width for a particle to set a non-constant mass to a particle.
+        If None, set to default value, given by GenMultiDecay.MASS_WIDTH_TOLERANCE
 
     Returns:
         The generated GenParticle instances, one for each possible way of the decay.
     """
+    if tolerance is None:
+        tolerance = GenMultiDecay.MASS_WIDTH_TOLERANCE
     # Get the only key inside the decaychain dict
     (original_mother_name,) = decaychain.keys()
 
@@ -264,8 +290,9 @@ def _recursively_traverse(
                 daughter = _recursively_traverse(
                     daughter_name,
                     mass_converter,
+                    particle_model_map,
+                    tolerance,
                     preexisting_particles,
-                    tolerance=tolerance,
                 )
             else:
                 raise TypeError(
@@ -279,10 +306,17 @@ def _recursively_traverse(
             if is_top_particle:
                 mother_mass = Particle.from_evtgen_name(original_mother_name).mass
             else:
+                if "zfit" in dm:
+                    mass_func = dm["zfit"]
+                elif original_mother_name in particle_model_map:
+                    mass_func = particle_model_map[original_mother_name]
+                else:
+                    mass_func = GenMultiDecay.DEFAULT_MASS_FUNC
+
                 mother_mass = _get_particle_mass(
                     original_mother_name,
                     mass_converter=mass_converter,
-                    mass_func=dm.get("zfit", GenMultiDecay.DEFAULT_MASS_FUNC),
+                    mass_func=mass_func,
                     tolerance=tolerance,
                 )
 

phasespace/fromdecay/mass_functions.py

@@ -2,14 +2,15 @@
 import zfit
 import zfit_physics as zphys
 
-
 # TODO refactor these mass functions using e.g. a decorator.
 #  Right now there is a lot of code repetition.
-def gauss(mass, width):
+
+
+def gauss_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def gauss(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs="")
@@ -18,14 +19,14 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator
         )
 
-    return mass_func
+    return gauss
 
 
-def breitwigner(mass, width):
+def breitwigner_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def bw(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zfit.pdf.Cauchy(m=particle_mass, gamma=particle_width, obs="")
@@ -34,14 +35,14 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator
         )
 
-    return mass_func
+    return bw
 
 
-def relativistic_breitwigner(mass, width):
+def relativistic_breitwigner_factory(mass, width):
     particle_mass = tf.cast(mass, tf.float64)
     particle_width = tf.cast(width, tf.float64)
 
-    def mass_func(min_mass, max_mass, n_events):
+    def relbw(min_mass, max_mass, n_events):
         min_mass = tf.cast(min_mass, tf.float64)
         max_mass = tf.cast(max_mass, tf.float64)
         pdf = zphys.pdf.RelativisticBreitWigner(
@@ -54,11 +55,11 @@ def mass_func(min_mass, max_mass, n_events):
             lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator
         )
 
-    return mass_func
+    return relbw
 
 
 DEFAULT_CONVERTER = {
-    "gauss": gauss,
-    "bw": breitwigner,
-    "relbw": relativistic_breitwigner,
+    "gauss": gauss_factory,
+    "bw": breitwigner_factory,
+    "relbw": relativistic_breitwigner_factory,
 }

tests/fromdecay/example_decay_chains.py

@@ -9,21 +9,14 @@
 
 # D+ particle with only one way of decaying
 dplus_decay = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="relbw")
-pi0_decay = DecayMode(1, "gamma gamma")
+pi0_decay = DecayMode(1, "gamma gamma", zfit="relbw")
 dplus_single = DecayChain("D+", {"D+": dplus_decay, "pi0": pi0_decay}).to_dict()
 
 # pi0 particle that can decay in 4 possible ways
 pi0_4branches = dfp.build_decay_chains("pi0")
 
 # D+ particle that decays into 4 particles, out of which one particle in turn decays in 4 different ways.
 dplus_4grandbranches = dfp.build_decay_chains("D+")
-# Specify different mass functions for the different decays of pi0
-mass_functions = ["relbw", "bw", "gauss"]
-
-for mass_function, decay_mode in zip(
-    mass_functions, dplus_4grandbranches["D+"][0]["fs"][-1]["pi0"]
-):
-    decay_mode["zfit"] = mass_function
 
 # D*+ particle that has multiple children, grandchild particles, many of which can decay in multiple ways.
 dstarplus_big_decay = dfp.build_decay_chains("D*+")