From c80b2df2869d488defab640d9f5bbbce4650589e Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 22 Dec 2021 21:54:32 +0100 Subject: [PATCH 01/28] Remove unnecessary typing import --- phasespace/fromdecay/genmultidecay.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 474dff91..47b42135 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/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 @@ -126,10 +125,9 @@ def from_dict( 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: From 1635ecdd7d973b07a8fb938dcdf26e2d6f3898c8 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 22 Dec 2021 22:10:12 +0100 Subject: [PATCH 02/28] Fix bug with class variable changes not having an effect. --- phasespace/fromdecay/genmultidecay.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 47b42135..c527d7f8 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -44,8 +44,8 @@ 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. Returns: The created GenMultiDecay object. @@ -112,6 +112,9 @@ 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 mass_converter is None: total_mass_converter = DEFAULT_CONVERTER else: @@ -219,7 +222,7 @@ def _recursively_traverse( decaychain: dict, mass_converter: dict[str, Callable], preexisting_particles: set[str] = None, - tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE, + tolerance: float = None, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -227,10 +230,13 @@ def _recursively_traverse( decaychain: Decay chain with the format from DecayLanguage 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() From 69c2d77940ca21e66f99c80b9324d77f82d536b8 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 19 Jan 2022 23:08:47 +0100 Subject: [PATCH 03/28] Add preliminary code for particle_model_map However, some tests still fail. --- phasespace/fromdecay/genmultidecay.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index c527d7f8..46bdb713 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -31,6 +31,7 @@ def from_dict( dec_dict: dict, mass_converter: dict[str, Callable] = None, tolerance: float = MASS_WIDTH_TOLERANCE, + particle_model_map: dict[str, str] = None, ): """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format. @@ -46,7 +47,7 @@ def from_dict( tolerance: Minimum mass width of the particle to use a mass function instead of 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: TODO Returns: The created GenMultiDecay object. @@ -114,7 +115,8 @@ def from_dict( """ 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: @@ -122,7 +124,10 @@ def from_dict( total_mass_converter = {**DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( - dec_dict, total_mass_converter, tolerance=tolerance + dec_dict, + total_mass_converter, + tolerance=tolerance, + particle_model_map=particle_model_map, ) return cls(gen_particles) @@ -221,6 +226,7 @@ def _get_particle_mass( def _recursively_traverse( decaychain: dict, mass_converter: dict[str, Callable], + particle_model_map: dict[str, str], preexisting_particles: set[str] = None, tolerance: float = None, ) -> list[tuple[float, GenParticle]]: @@ -283,10 +289,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, ) From 0c45899b5ea60c2a260097cdf2ace4b8b2f5fb19 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 19 Jan 2022 23:23:25 +0100 Subject: [PATCH 04/28] Add test for ValueError that can be raised --- tests/fromdecay/test_fromdecay.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 61a2487d..d46001e3 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest +from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal from phasespace.fromdecay import GenMultiDecay @@ -34,6 +36,16 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: return all_return_args +def test_invalid_chain(): + """Test that a ValueError is raised when a value in the fs key is not a str or dict.""" + dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="rel-BW") + dm2 = DecayMode(1, "gamma gamma") + dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() + dc["D+"][0]["fs"][0] = 1 + with pytest.raises(TypeError): + GenMultiDecay.from_dict(dc) + + def test_single_chain(): """Test converting a DecayLanguage dict with only one possible decay.""" container = GenMultiDecay.from_dict( From 9bdeaa5b4f10796f9280740ebfb0edab1071ee0d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 25 Jan 2022 23:12:05 +0100 Subject: [PATCH 05/28] Fix bug caused by wrong order of arguments --- phasespace/fromdecay/genmultidecay.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 46bdb713..9d3b451b 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -126,8 +126,8 @@ def from_dict( gen_particles = _recursively_traverse( dec_dict, total_mass_converter, - tolerance=tolerance, particle_model_map=particle_model_map, + tolerance=tolerance, ) return cls(gen_particles) @@ -203,7 +203,7 @@ 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. @@ -227,8 +227,8 @@ 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 = None, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -274,8 +274,9 @@ def _recursively_traverse( daughter = _recursively_traverse( daughter_name, mass_converter, + particle_model_map, + tolerance, preexisting_particles, - tolerance=tolerance, ) else: raise TypeError( From e4aa2c4b640552e49477f8acdd6a0435c499759d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 26 Jan 2022 22:51:34 +0100 Subject: [PATCH 06/28] Add more parameters in docstrings. Especially the particle_model_map parameter. --- phasespace/fromdecay/genmultidecay.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 9d3b451b..55038cc7 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -47,7 +47,12 @@ def from_dict( tolerance: Minimum mass width of the particle to use a mass function instead of 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: TODO + 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. @@ -108,6 +113,7 @@ def from_dict( One can then pass the `custom_gauss` function and its name (in this case "custom_gauss") as a `dict`to `from_dict` as the mass_converter parameter: >>> dst_gen = GenMultiDecay.from_dict(dst_chain, mass_converter={"custom_gauss": custom_gauss}) + TODO add example with particle_model_map Notes: For a more in-depth tutorial, see the tutorial on GenMultiDecay in the @@ -209,7 +215,9 @@ def _get_particle_mass( 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. @@ -234,6 +242,8 @@ def _recursively_traverse( 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 From f3d80feaa318dae101128e6d84d1fa486ea49c5f Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 29 Jan 2022 20:36:17 +0100 Subject: [PATCH 07/28] Set function names to make the code testable and add test for dicts created from a DecayChain. --- phasespace/fromdecay/mass_functions.py | 9 ++++++++- tests/fromdecay/test_fromdecay.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index f4ea260b..b841d7a9 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -2,9 +2,10 @@ 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): particle_mass = tf.cast(mass, tf.float64) particle_width = tf.cast(width, tf.float64) @@ -18,6 +19,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) + mass_func.__name__ = "gauss" + return mass_func @@ -34,6 +37,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) + mass_func.__name__ = "breitwigner" + return mass_func @@ -54,6 +59,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) + mass_func.__name__ = "relativistic_breitwigner" + return mass_func diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index d46001e3..f496bcda 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -3,6 +3,7 @@ import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal +from particle import Particle from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER @@ -46,6 +47,27 @@ def test_invalid_chain(): GenMultiDecay.from_dict(dc) +def test_decay_chain_dict(): + """Create a DecayChain object and convert it to a dict.""" + dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP") + dm2 = DecayMode(1, "gamma gamma", zfit="relbw") + dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() + container = GenMultiDecay.from_dict( + dc, tolerance=Particle.from_evtgen_name("pi0").width / 1.1 + ) + + assert len(container.gen_particles) == 1 + gen_particle = container.gen_particles[0][1] + assert gen_particle.name == "D+" + assert gen_particle.has_fixed_mass + + children_name = [p.name for p in gen_particle.children] + assert set(children_name) == {"K-", "pi+", "pi+ [0]", "pi0"} + pi0 = gen_particle.children[children_name.index("pi0")] + assert {p.name for p in pi0.children} == {"gamma", "gamma [0]"} + assert pi0._mass.__name__ == "relativistic_breitwigner" + + def test_single_chain(): """Test converting a DecayLanguage dict with only one possible decay.""" container = GenMultiDecay.from_dict( From 4df32b0e654d68e252fa3a85c19a0d2d4971c227 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 29 Jan 2022 20:45:41 +0100 Subject: [PATCH 08/28] Remove a test that was almost identical --- tests/fromdecay/example_decay_chains.py | 2 +- tests/fromdecay/test_fromdecay.py | 29 +++++-------------------- 2 files changed, 7 insertions(+), 24 deletions(-) diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py index c04d867d..3f900e76 100644 --- a/tests/fromdecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -9,7 +9,7 @@ # 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 diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index f496bcda..cc81767f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -3,7 +3,6 @@ import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal -from particle import Particle from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER @@ -47,29 +46,12 @@ def test_invalid_chain(): GenMultiDecay.from_dict(dc) -def test_decay_chain_dict(): - """Create a DecayChain object and convert it to a dict.""" - dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP") - dm2 = DecayMode(1, "gamma gamma", zfit="relbw") - dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() - container = GenMultiDecay.from_dict( - dc, tolerance=Particle.from_evtgen_name("pi0").width / 1.1 - ) - - assert len(container.gen_particles) == 1 - gen_particle = container.gen_particles[0][1] - assert gen_particle.name == "D+" - assert gen_particle.has_fixed_mass - - children_name = [p.name for p in gen_particle.children] - assert set(children_name) == {"K-", "pi+", "pi+ [0]", "pi0"} - pi0 = gen_particle.children[children_name.index("pi0")] - assert {p.name for p in pi0.children} == {"gamma", "gamma [0]"} - assert pi0._mass.__name__ == "relativistic_breitwigner" - - def test_single_chain(): - """Test converting a DecayLanguage dict with only one possible decay.""" + """Test converting a DecayLanguage dict with only one possible decay. + + Since dplus_single is constructed using DecayChain.to_dict, this also tests that the code works dicts + created from DecayChains, not just .dec files. + """ container = GenMultiDecay.from_dict( example_decay_chains.dplus_single, tolerance=1e-10 ) @@ -82,6 +64,7 @@ def test_single_chain(): for p in gen.children: if "pi0" == p.name[:3]: assert not p.has_fixed_mass + assert p._mass.__name__ == "relativistic_breitwigner" else: assert p.has_fixed_mass From 5c70b914ca5602c0595e68a883432e5012152dcf Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 31 Jan 2022 21:17:07 +0100 Subject: [PATCH 09/28] Add test for particle_model_map Additionally, the function __name__ parameters were renamed to their corresponding keys in DEFAULT_CONVERTER. --- phasespace/fromdecay/mass_functions.py | 4 +-- tests/fromdecay/example_decay_chains.py | 7 ---- tests/fromdecay/test_fromdecay.py | 47 ++++++++++++++++++++++--- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index b841d7a9..d4903817 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -37,7 +37,7 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) - mass_func.__name__ = "breitwigner" + mass_func.__name__ = "bw" return mass_func @@ -59,7 +59,7 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) - mass_func.__name__ = "relativistic_breitwigner" + mass_func.__name__ = "relbw" return mass_func diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py index 3f900e76..934cfd65 100644 --- a/tests/fromdecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -17,13 +17,6 @@ # 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*+") diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index cc81767f..29f7b26f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,5 +1,7 @@ from __future__ import annotations +from copy import deepcopy + import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal @@ -11,7 +13,7 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: - """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate. + """Helper function that tests whether the normalize_weights argument works for GenMultiDecay.generate. Args: full_decay: full_decay.generate will be called. @@ -64,7 +66,7 @@ def test_single_chain(): for p in gen.children: if "pi0" == p.name[:3]: assert not p.has_fixed_mass - assert p._mass.__name__ == "relativistic_breitwigner" + assert p._mass.__name__ == "relbw" else: assert p.has_fixed_mass @@ -90,13 +92,50 @@ def test_branching_children(): def test_branching_grandchilden(): """Test converting a DecayLanguage dict where children to the mother particle can decay in many ways.""" - container = GenMultiDecay.from_dict(example_decay_chains.dplus_4grandbranches) + # Specify different mass functions for the different decays of pi0 + decay_dict = deepcopy(example_decay_chains.dplus_4grandbranches) + + # Add different zfit parameters to all pi0 decays. The fourth decay has no zfit parameter + for mass_function, decay_mode in zip( + ("relbw", "bw", "gauss"), decay_dict["D+"][0]["fs"][-1]["pi0"] + ): + decay_mode["zfit"] = mass_function + + container = GenMultiDecay.from_dict(decay_dict, tolerance=1e-10) + output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) + + for p, mass_func in zip( + output_decays, ("relbw", "bw", "gauss", GenMultiDecay.DEFAULT_MASS_FUNC) + ): + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name == "pi0" + # Check that the zfit parameter assigns the correct mass function + assert gen_particle.children[-1]._mass.__name__ == mass_func + + check_norm(container, n_events=1) + check_norm(container, n_events=100) + + +def test_particle_model_map(): + """Test that the particle_model_map parameter works as intended.""" + container = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches, + particle_model_map={"pi0": "bw"}, + tolerance=1e-10, + ) + output_decays = container.gen_particles + assert len(output_decays) == 4 + assert_almost_equal(sum(d[0] for d in output_decays), 1) + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert gen_particle.children[-1]._mass.__name__ == "bw" check_norm(container, n_events=1) check_norm(container, n_events=100) - # TODO add more asserts here def test_mass_converter(): From 248878340214fe9adfcd24f791e257e799d7cf2e Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 5 Feb 2022 19:05:57 +0100 Subject: [PATCH 10/28] Add tutorial for using DecayChain and particle_model_map --- docs/GenMultiDecay_Tutorial.ipynb | 63 +++++++++++++++++++++++++++++-- 1 file changed, 59 insertions(+), 4 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index 8bb26c39..b42d4071 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/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,6 +285,42 @@ "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", @@ -362,4 +417,4 @@ }, "nbformat": 4, "nbformat_minor": 1 -} \ No newline at end of file +} From e785463dfff5cbb4b3f7d82854f759b54f96f3bf Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 16 Feb 2022 23:59:08 +0100 Subject: [PATCH 11/28] Add tests for class variables and fix bug related to these. --- phasespace/fromdecay/genmultidecay.py | 2 +- tests/fromdecay/test_fromdecay.py | 32 +++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 55038cc7..b4e7d50d 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -30,7 +30,7 @@ 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. diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 29f7b26f..44c5ae02 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -161,9 +161,41 @@ def test_mass_converter(): def test_big_decay(): + """Create a GenMultiDecay object from a large dict with many branches and subbranches.""" container = GenMultiDecay.from_dict(example_decay_chains.dstarplus_big_decay) output_decays = container.gen_particles assert_almost_equal(sum(d[0] for d in output_decays), 1) check_norm(container, n_events=1) check_norm(container, n_events=100) # TODO add more asserts here + + +def test_mass_width_tolerance(): + """Test changing the MASS_WIDTH_TOLERANCE class variable.""" + GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10 + output_decays = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches + ).gen_particles + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert not gen_particle.children[-1].has_fixed_mass + # Restore class variable to not affect other tests + GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10 + + +def test_default_mass_func(): + """Test changing the DEFAULT_MASS_FUNC class variable.""" + GenMultiDecay.DEFAULT_MASS_FUNC = "bw" + output_decays = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches, tolerance=1e-10 + ).gen_particles + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert gen_particle.children[-1]._mass.__name__ == "bw" + + # Restore class variable to not affect other tests + GenMultiDecay.DEFAULT_MASS_FUNC = "bw" From 2fc3c0ab63cff8488507900272c7b3da2bae609d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 26 Feb 2022 12:54:04 +0100 Subject: [PATCH 12/28] Add example of particle_model_map usage in the from_dict docstring. --- phasespace/fromdecay/genmultidecay.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index b4e7d50d..081edbf5 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -75,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, @@ -86,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,7 +120,6 @@ def from_dict( One can then pass the `custom_gauss` function and its name (in this case "custom_gauss") as a `dict`to `from_dict` as the mass_converter parameter: >>> dst_gen = GenMultiDecay.from_dict(dst_chain, mass_converter={"custom_gauss": custom_gauss}) - TODO add example with particle_model_map Notes: For a more in-depth tutorial, see the tutorial on GenMultiDecay in the From 4a4d1b291d18905ace955fad66595fd8ed740cb5 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 20 Mar 2022 19:12:19 +0100 Subject: [PATCH 13/28] Fix failing tets with deepcopy --- tests/fromdecay/test_fromdecay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 44c5ae02..f1349a04 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -140,7 +140,7 @@ def test_particle_model_map(): def test_mass_converter(): """Test that the mass_converter parameter works as intended.""" - dplus_4grandbranches_massfunc = example_decay_chains.dplus_4grandbranches.copy() + dplus_4grandbranches_massfunc = deepcopy(example_decay_chains.dplus_4grandbranches) dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" container = GenMultiDecay.from_dict( dplus_4grandbranches_massfunc, From 719659efeb74ca661927524b5fc3060d11969080 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 22 Dec 2021 21:54:32 +0100 Subject: [PATCH 14/28] Remove unnecessary typing import --- phasespace/fromdecay/genmultidecay.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 474dff91..47b42135 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/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 @@ -126,10 +125,9 @@ def from_dict( 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: From 26d7cd1f83354183528d8a320d7c2aae1b4aefab Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 22 Dec 2021 22:10:12 +0100 Subject: [PATCH 15/28] Fix bug with class variable changes not having an effect. --- phasespace/fromdecay/genmultidecay.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 47b42135..c527d7f8 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -44,8 +44,8 @@ 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. Returns: The created GenMultiDecay object. @@ -112,6 +112,9 @@ 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 mass_converter is None: total_mass_converter = DEFAULT_CONVERTER else: @@ -219,7 +222,7 @@ def _recursively_traverse( decaychain: dict, mass_converter: dict[str, Callable], preexisting_particles: set[str] = None, - tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE, + tolerance: float = None, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -227,10 +230,13 @@ def _recursively_traverse( decaychain: Decay chain with the format from DecayLanguage 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() From fabaabf3696fc52b9d76d55a1aa5a6e24d090adf Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 19 Jan 2022 23:08:47 +0100 Subject: [PATCH 16/28] Add preliminary code for particle_model_map However, some tests still fail. --- phasespace/fromdecay/genmultidecay.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index c527d7f8..46bdb713 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -31,6 +31,7 @@ def from_dict( dec_dict: dict, mass_converter: dict[str, Callable] = None, tolerance: float = MASS_WIDTH_TOLERANCE, + particle_model_map: dict[str, str] = None, ): """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format. @@ -46,7 +47,7 @@ def from_dict( tolerance: Minimum mass width of the particle to use a mass function instead of 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: TODO Returns: The created GenMultiDecay object. @@ -114,7 +115,8 @@ def from_dict( """ 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: @@ -122,7 +124,10 @@ def from_dict( total_mass_converter = {**DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( - dec_dict, total_mass_converter, tolerance=tolerance + dec_dict, + total_mass_converter, + tolerance=tolerance, + particle_model_map=particle_model_map, ) return cls(gen_particles) @@ -221,6 +226,7 @@ def _get_particle_mass( def _recursively_traverse( decaychain: dict, mass_converter: dict[str, Callable], + particle_model_map: dict[str, str], preexisting_particles: set[str] = None, tolerance: float = None, ) -> list[tuple[float, GenParticle]]: @@ -283,10 +289,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, ) From e8994e55a0d08cbc0fb2f09e74ed3a5765f0a866 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 19 Jan 2022 23:23:25 +0100 Subject: [PATCH 17/28] Add test for ValueError that can be raised --- tests/fromdecay/test_fromdecay.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 61a2487d..d46001e3 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest +from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal from phasespace.fromdecay import GenMultiDecay @@ -34,6 +36,16 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: return all_return_args +def test_invalid_chain(): + """Test that a ValueError is raised when a value in the fs key is not a str or dict.""" + dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="rel-BW") + dm2 = DecayMode(1, "gamma gamma") + dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() + dc["D+"][0]["fs"][0] = 1 + with pytest.raises(TypeError): + GenMultiDecay.from_dict(dc) + + def test_single_chain(): """Test converting a DecayLanguage dict with only one possible decay.""" container = GenMultiDecay.from_dict( From e1b2b795da77319866823be0a57bae98dcbbd144 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 25 Jan 2022 23:12:05 +0100 Subject: [PATCH 18/28] Fix bug caused by wrong order of arguments --- phasespace/fromdecay/genmultidecay.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 46bdb713..9d3b451b 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -126,8 +126,8 @@ def from_dict( gen_particles = _recursively_traverse( dec_dict, total_mass_converter, - tolerance=tolerance, particle_model_map=particle_model_map, + tolerance=tolerance, ) return cls(gen_particles) @@ -203,7 +203,7 @@ 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. @@ -227,8 +227,8 @@ 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 = None, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -274,8 +274,9 @@ def _recursively_traverse( daughter = _recursively_traverse( daughter_name, mass_converter, + particle_model_map, + tolerance, preexisting_particles, - tolerance=tolerance, ) else: raise TypeError( From 7b27a9f685dc8640c9d6e4d57f7da95d3ace5952 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 26 Jan 2022 22:51:34 +0100 Subject: [PATCH 19/28] Add more parameters in docstrings. Especially the particle_model_map parameter. --- phasespace/fromdecay/genmultidecay.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 9d3b451b..55038cc7 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -47,7 +47,12 @@ def from_dict( tolerance: Minimum mass width of the particle to use a mass function instead of 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: TODO + 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. @@ -108,6 +113,7 @@ def from_dict( One can then pass the `custom_gauss` function and its name (in this case "custom_gauss") as a `dict`to `from_dict` as the mass_converter parameter: >>> dst_gen = GenMultiDecay.from_dict(dst_chain, mass_converter={"custom_gauss": custom_gauss}) + TODO add example with particle_model_map Notes: For a more in-depth tutorial, see the tutorial on GenMultiDecay in the @@ -209,7 +215,9 @@ def _get_particle_mass( 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. @@ -234,6 +242,8 @@ def _recursively_traverse( 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 From d64ddda69f727a8e3dd943ca939cd2cdf1f1907d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 29 Jan 2022 20:36:17 +0100 Subject: [PATCH 20/28] Set function names to make the code testable and add test for dicts created from a DecayChain. --- phasespace/fromdecay/mass_functions.py | 9 ++++++++- tests/fromdecay/test_fromdecay.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index f4ea260b..b841d7a9 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -2,9 +2,10 @@ 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): particle_mass = tf.cast(mass, tf.float64) particle_width = tf.cast(width, tf.float64) @@ -18,6 +19,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) + mass_func.__name__ = "gauss" + return mass_func @@ -34,6 +37,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) + mass_func.__name__ = "breitwigner" + return mass_func @@ -54,6 +59,8 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) + mass_func.__name__ = "relativistic_breitwigner" + return mass_func diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index d46001e3..f496bcda 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -3,6 +3,7 @@ import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal +from particle import Particle from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER @@ -46,6 +47,27 @@ def test_invalid_chain(): GenMultiDecay.from_dict(dc) +def test_decay_chain_dict(): + """Create a DecayChain object and convert it to a dict.""" + dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP") + dm2 = DecayMode(1, "gamma gamma", zfit="relbw") + dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() + container = GenMultiDecay.from_dict( + dc, tolerance=Particle.from_evtgen_name("pi0").width / 1.1 + ) + + assert len(container.gen_particles) == 1 + gen_particle = container.gen_particles[0][1] + assert gen_particle.name == "D+" + assert gen_particle.has_fixed_mass + + children_name = [p.name for p in gen_particle.children] + assert set(children_name) == {"K-", "pi+", "pi+ [0]", "pi0"} + pi0 = gen_particle.children[children_name.index("pi0")] + assert {p.name for p in pi0.children} == {"gamma", "gamma [0]"} + assert pi0._mass.__name__ == "relativistic_breitwigner" + + def test_single_chain(): """Test converting a DecayLanguage dict with only one possible decay.""" container = GenMultiDecay.from_dict( From 880a21f1501b52f8595540d629b23a205b225e11 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 29 Jan 2022 20:45:41 +0100 Subject: [PATCH 21/28] Remove a test that was almost identical --- tests/fromdecay/example_decay_chains.py | 2 +- tests/fromdecay/test_fromdecay.py | 29 +++++-------------------- 2 files changed, 7 insertions(+), 24 deletions(-) diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py index c04d867d..3f900e76 100644 --- a/tests/fromdecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -9,7 +9,7 @@ # 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 diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index f496bcda..cc81767f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -3,7 +3,6 @@ import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal -from particle import Particle from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER @@ -47,29 +46,12 @@ def test_invalid_chain(): GenMultiDecay.from_dict(dc) -def test_decay_chain_dict(): - """Create a DecayChain object and convert it to a dict.""" - dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP") - dm2 = DecayMode(1, "gamma gamma", zfit="relbw") - dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() - container = GenMultiDecay.from_dict( - dc, tolerance=Particle.from_evtgen_name("pi0").width / 1.1 - ) - - assert len(container.gen_particles) == 1 - gen_particle = container.gen_particles[0][1] - assert gen_particle.name == "D+" - assert gen_particle.has_fixed_mass - - children_name = [p.name for p in gen_particle.children] - assert set(children_name) == {"K-", "pi+", "pi+ [0]", "pi0"} - pi0 = gen_particle.children[children_name.index("pi0")] - assert {p.name for p in pi0.children} == {"gamma", "gamma [0]"} - assert pi0._mass.__name__ == "relativistic_breitwigner" - - def test_single_chain(): - """Test converting a DecayLanguage dict with only one possible decay.""" + """Test converting a DecayLanguage dict with only one possible decay. + + Since dplus_single is constructed using DecayChain.to_dict, this also tests that the code works dicts + created from DecayChains, not just .dec files. + """ container = GenMultiDecay.from_dict( example_decay_chains.dplus_single, tolerance=1e-10 ) @@ -82,6 +64,7 @@ def test_single_chain(): for p in gen.children: if "pi0" == p.name[:3]: assert not p.has_fixed_mass + assert p._mass.__name__ == "relativistic_breitwigner" else: assert p.has_fixed_mass From 2909110158934ff9ccde35a268fe1941aa051dd5 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 31 Jan 2022 21:17:07 +0100 Subject: [PATCH 22/28] Add test for particle_model_map Additionally, the function __name__ parameters were renamed to their corresponding keys in DEFAULT_CONVERTER. --- phasespace/fromdecay/mass_functions.py | 4 +-- tests/fromdecay/example_decay_chains.py | 7 ---- tests/fromdecay/test_fromdecay.py | 47 ++++++++++++++++++++++--- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index b841d7a9..d4903817 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -37,7 +37,7 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) - mass_func.__name__ = "breitwigner" + mass_func.__name__ = "bw" return mass_func @@ -59,7 +59,7 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) - mass_func.__name__ = "relativistic_breitwigner" + mass_func.__name__ = "relbw" return mass_func diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py index 3f900e76..934cfd65 100644 --- a/tests/fromdecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -17,13 +17,6 @@ # 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*+") diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index cc81767f..29f7b26f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,5 +1,7 @@ from __future__ import annotations +from copy import deepcopy + import pytest from decaylanguage import DecayChain, DecayMode from numpy.testing import assert_almost_equal @@ -11,7 +13,7 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: - """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate. + """Helper function that tests whether the normalize_weights argument works for GenMultiDecay.generate. Args: full_decay: full_decay.generate will be called. @@ -64,7 +66,7 @@ def test_single_chain(): for p in gen.children: if "pi0" == p.name[:3]: assert not p.has_fixed_mass - assert p._mass.__name__ == "relativistic_breitwigner" + assert p._mass.__name__ == "relbw" else: assert p.has_fixed_mass @@ -90,13 +92,50 @@ def test_branching_children(): def test_branching_grandchilden(): """Test converting a DecayLanguage dict where children to the mother particle can decay in many ways.""" - container = GenMultiDecay.from_dict(example_decay_chains.dplus_4grandbranches) + # Specify different mass functions for the different decays of pi0 + decay_dict = deepcopy(example_decay_chains.dplus_4grandbranches) + + # Add different zfit parameters to all pi0 decays. The fourth decay has no zfit parameter + for mass_function, decay_mode in zip( + ("relbw", "bw", "gauss"), decay_dict["D+"][0]["fs"][-1]["pi0"] + ): + decay_mode["zfit"] = mass_function + + container = GenMultiDecay.from_dict(decay_dict, tolerance=1e-10) + output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) + + for p, mass_func in zip( + output_decays, ("relbw", "bw", "gauss", GenMultiDecay.DEFAULT_MASS_FUNC) + ): + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name == "pi0" + # Check that the zfit parameter assigns the correct mass function + assert gen_particle.children[-1]._mass.__name__ == mass_func + + check_norm(container, n_events=1) + check_norm(container, n_events=100) + + +def test_particle_model_map(): + """Test that the particle_model_map parameter works as intended.""" + container = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches, + particle_model_map={"pi0": "bw"}, + tolerance=1e-10, + ) + output_decays = container.gen_particles + assert len(output_decays) == 4 + assert_almost_equal(sum(d[0] for d in output_decays), 1) + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert gen_particle.children[-1]._mass.__name__ == "bw" check_norm(container, n_events=1) check_norm(container, n_events=100) - # TODO add more asserts here def test_mass_converter(): From bd81dd3beee88313d4a5116eaf319d739b291541 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 5 Feb 2022 19:05:57 +0100 Subject: [PATCH 23/28] Add tutorial for using DecayChain and particle_model_map --- docs/GenMultiDecay_Tutorial.ipynb | 63 +++++++++++++++++++++++++++++-- 1 file changed, 59 insertions(+), 4 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index 8bb26c39..b42d4071 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/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,6 +285,42 @@ "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", @@ -362,4 +417,4 @@ }, "nbformat": 4, "nbformat_minor": 1 -} \ No newline at end of file +} From a86973d7a578eb5fde17eec0f2cbf519c6c61d08 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 16 Feb 2022 23:59:08 +0100 Subject: [PATCH 24/28] Add tests for class variables and fix bug related to these. --- phasespace/fromdecay/genmultidecay.py | 2 +- tests/fromdecay/test_fromdecay.py | 32 +++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 55038cc7..b4e7d50d 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -30,7 +30,7 @@ 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. diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 29f7b26f..44c5ae02 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -161,9 +161,41 @@ def test_mass_converter(): def test_big_decay(): + """Create a GenMultiDecay object from a large dict with many branches and subbranches.""" container = GenMultiDecay.from_dict(example_decay_chains.dstarplus_big_decay) output_decays = container.gen_particles assert_almost_equal(sum(d[0] for d in output_decays), 1) check_norm(container, n_events=1) check_norm(container, n_events=100) # TODO add more asserts here + + +def test_mass_width_tolerance(): + """Test changing the MASS_WIDTH_TOLERANCE class variable.""" + GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10 + output_decays = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches + ).gen_particles + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert not gen_particle.children[-1].has_fixed_mass + # Restore class variable to not affect other tests + GenMultiDecay.MASS_WIDTH_TOLERANCE = 1e-10 + + +def test_default_mass_func(): + """Test changing the DEFAULT_MASS_FUNC class variable.""" + GenMultiDecay.DEFAULT_MASS_FUNC = "bw" + output_decays = GenMultiDecay.from_dict( + example_decay_chains.dplus_4grandbranches, tolerance=1e-10 + ).gen_particles + for p in output_decays: + gen_particle = p[1] # Ignore probability + assert gen_particle.children[-1].name[:3] == "pi0" + # Check that particle_model_map has assigned the bw mass function to all pi0 decays. + assert gen_particle.children[-1]._mass.__name__ == "bw" + + # Restore class variable to not affect other tests + GenMultiDecay.DEFAULT_MASS_FUNC = "bw" From 7d6a276c0b13506bd126e6cabbd94d3ab9c5cbea Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 26 Feb 2022 12:54:04 +0100 Subject: [PATCH 25/28] Add example of particle_model_map usage in the from_dict docstring. --- phasespace/fromdecay/genmultidecay.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index b4e7d50d..081edbf5 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -75,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, @@ -86,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,7 +120,6 @@ def from_dict( One can then pass the `custom_gauss` function and its name (in this case "custom_gauss") as a `dict`to `from_dict` as the mass_converter parameter: >>> dst_gen = GenMultiDecay.from_dict(dst_chain, mass_converter={"custom_gauss": custom_gauss}) - TODO add example with particle_model_map Notes: For a more in-depth tutorial, see the tutorial on GenMultiDecay in the From d07970f30d88cc39f6abe67c6b08b21c6efe961e Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 20 Mar 2022 19:12:19 +0100 Subject: [PATCH 26/28] Fix failing tets with deepcopy --- tests/fromdecay/test_fromdecay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 44c5ae02..f1349a04 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -140,7 +140,7 @@ def test_particle_model_map(): def test_mass_converter(): """Test that the mass_converter parameter works as intended.""" - dplus_4grandbranches_massfunc = example_decay_chains.dplus_4grandbranches.copy() + dplus_4grandbranches_massfunc = deepcopy(example_decay_chains.dplus_4grandbranches) dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" container = GenMultiDecay.from_dict( dplus_4grandbranches_massfunc, From 65e8ee23eb631494cf73bd12736192ab2c7a64de Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Apr 2022 11:43:15 +0200 Subject: [PATCH 27/28] Remove manual change of __name__ attribute for mass functions --- phasespace/fromdecay/mass_functions.py | 30 +++++++++++--------------- tests/fromdecay/test_mass_functions.py | 3 ++- 2 files changed, 14 insertions(+), 19 deletions(-) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index d4903817..6767414b 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -6,11 +6,11 @@ # 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="") @@ -19,16 +19,14 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) - mass_func.__name__ = "gauss" + return gauss - return mass_func - -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="") @@ -37,16 +35,14 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator ) - mass_func.__name__ = "bw" - - 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( @@ -59,13 +55,11 @@ def mass_func(min_mass, max_mass, n_events): lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) - mass_func.__name__ = "relbw" - - return mass_func + return relbw DEFAULT_CONVERTER = { - "gauss": gauss, - "bw": breitwigner, - "relbw": relativistic_breitwigner, + "gauss": gauss_factory, + "bw": breitwigner_factory, + "relbw": relativistic_breitwigner_factory, } diff --git a/tests/fromdecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py index 339cbe44..e93e4c2f 100644 --- a/tests/fromdecay/test_mass_functions.py +++ b/tests/fromdecay/test_mass_functions.py @@ -40,7 +40,8 @@ def ref_mass_func(min_mass, max_mass, n_events): @pytest.mark.parametrize( - "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) + "function", + (mf.gauss_factory, mf.breitwigner_factory, mf.relativistic_breitwigner_factory), ) @pytest.mark.parametrize("size", (1, 10)) def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): From a32d54a056437d33d3a0baa50d0db45cf8d287a6 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Apr 2022 11:52:53 +0200 Subject: [PATCH 28/28] Add test for when the zfit parameter is invalid and update the tutorial for a correct description of this behaviour. --- docs/GenMultiDecay_Tutorial.ipynb | 2 +- tests/fromdecay/test_fromdecay.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index b42d4071..8d88b6ed 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -323,7 +323,7 @@ "#### 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:" ] diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index f1349a04..63574308 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -48,6 +48,15 @@ def test_invalid_chain(): GenMultiDecay.from_dict(dc) +def test_invalid_mass_func_name(): + """Test if an invalid mass function name as the zfit parameter raises a KeyError.""" + dm1 = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP") + dm2 = DecayMode(1, "gamma gamma", zfit="invalid name") + dc = DecayChain("D+", {"D+": dm1, "pi0": dm2}).to_dict() + with pytest.raises(KeyError): + GenMultiDecay.from_dict(dc, tolerance=1e-10) + + def test_single_chain(): """Test converting a DecayLanguage dict with only one possible decay.