From 9254e8f4eeff602fce744d5638a39c4944752a50 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 21 Sep 2021 16:21:48 +0200 Subject: [PATCH 01/71] Add files for compatibility between phasespace and decaylanguage --- phasespace/fulldecay/__init__.py | 14 ++ phasespace/fulldecay/fulldecay.py | 225 ++++++++++++++++++++++++ phasespace/fulldecay/mass_functions.py | 50 ++++++ tests/fulldecay/example_decay_chains.py | 51 ++++++ tests/fulldecay/test_fulldecay.py | 88 +++++++++ tests/fulldecay/test_mass_functions.py | 52 ++++++ 6 files changed, 480 insertions(+) create mode 100644 phasespace/fulldecay/__init__.py create mode 100644 phasespace/fulldecay/fulldecay.py create mode 100644 phasespace/fulldecay/mass_functions.py create mode 100644 tests/fulldecay/example_decay_chains.py create mode 100644 tests/fulldecay/test_fulldecay.py create mode 100644 tests/fulldecay/test_mass_functions.py diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py new file mode 100644 index 00000000..c48943b8 --- /dev/null +++ b/phasespace/fulldecay/__init__.py @@ -0,0 +1,14 @@ +from .fulldecay import FullDecay + +import sys +try: + from particle import Particle + import zfit + import zfit_physics as zphys +except ModuleNotFoundError: + print( + "the fulldecay functionality in phasespace requires particle and zfit-physics. " + "Either install phasespace[fulldecay] or particle and zfit-physics.", + file=sys.stderr, + ) + raise diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py new file mode 100644 index 00000000..aa1e04f4 --- /dev/null +++ b/phasespace/fulldecay/fulldecay.py @@ -0,0 +1,225 @@ +from phasespace import GenParticle +import tensorflow as tf +import tensorflow.experimental.numpy as tnp +from particle import Particle + +from .mass_functions import _DEFAULT_CONVERTER + +from typing import Callable, Union +import itertools + +_MASS_WIDTH_TOLERANCE = 0.01 +_DEFAULT_MASS_FUNC = 'rel-BW' + + +class FullDecay: + """ + A container that works like GenParticle that can handle multiple decays + """ + def __init__(self, gen_particles: list[tuple[float, GenParticle]]): + """ + Create an instance of FullDecay + + Parameters + ---------- + gen_particles : list[tuple[float, GenParticle]] + All the GenParticles and their corresponding probabilities. + The list must be of the format [[probability, GenParticle instance], [probability, ... + Notes + ----- + Input format might change + """ + self.gen_particles = gen_particles + + @classmethod + def from_dict(cls, dec_dict: dict, mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE): + """ + Create a FullDecay instance from a dict in the decaylanguage format. + + Parameters + ---------- + dec_dict : dict + The input dict from which the FullDecay object will be created from. + mass_converter : dict[str, Callable] + A dict with mass function names and their corresponding mass functions. + These functions should take the average particle mass and the mass width as inputs + and return a mass function that phasespace can understand. + This dict will be combined with the predefined mass functions in this package. + tolerance : float + Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. + + Returns + ------- + FullDecay + The created FullDecay object. + """ + if mass_converter is None: + total_mass_converter = _DEFAULT_CONVERTER + else: + # Combine the mass functions specified by the package to the mass functions specified from the input. + total_mass_converter = {**_DEFAULT_CONVERTER, **mass_converter} + + gen_particles = _recursively_traverse(dec_dict, total_mass_converter, tolerance=tolerance) + return cls(gen_particles) + + def generate(self, n_events: int, normalize_weights: bool = False, + **kwargs) -> Union[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). + + Parameters + ---------- + n_events : int + Total number of events combined, for all the decays. + normalize_weights : bool + Normalize weights according to all events generated. This also changes the return values. + See the phasespace documentation for more details. + kwargs + Additional parameters passed to all calls of GenParticle.generate + + Returns + ------- + The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. + However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding + to the return arguments from the corresponding GenParticle instances in self.gen_particles. + Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. + """ + # Input to tf.random.categorical must be 2D + rand_i = tf.random.categorical(tnp.log([[dm[0] for dm in self.gen_particles]]), n_events) + # Input to tf.unique_with_counts must be 1D + dec_indices, _, counts = tf.unique_with_counts(rand_i[0]) + counts = tf.cast(counts, tf.int64) + weights, max_weights, events = [], [], [] + for i, n in zip(dec_indices, counts): + weight, max_weight, four_vectors = self.gen_particles[i][1].generate(n, normalize_weights=False, **kwargs) + weights.append(weight) + max_weights.append(max_weight) + events.append(four_vectors) + + if normalize_weights: + total_max = tnp.max([tnp.max(mw) for mw in max_weights]) + normed_weights = [w / total_max for w in weights] + return normed_weights, events + + return weights, max_weights, events + + +def _unique_name(name: str, preexisting_particles: set[str]) -> str: + """ + Create a string that does not exist in preexisting_particles based on name. + + Parameters + ---------- + name : str + Name that should be + preexisting_particles : set[str] + Preexisting names + + Returns + ------- + name : str + Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i will begin at 0 + and increase until the name is not preexisting_particles. + """ + if name not in preexisting_particles: + preexisting_particles.add(name) + return name + + name += ' [0]' + i = 1 + while name in preexisting_particles: + name = name[:name.rfind('[')] + f'[{str(i)}]' + i += 1 + preexisting_particles.add(name) + return name + + +def _get_particle_mass(name: str, mass_converter: dict[str, Callable], mass_func: str, + tolerance: float = _MASS_WIDTH_TOLERANCE) -> Union[Callable, float]: + """ + Get mass or mass function of particle using the particle package. + Parameters + ---------- + name : str + Name of the particle. Name must be recognizable by the particle package. + tolerance : float + See _recursively_traverse + + Returns + ------- + Callable, float + Returns a function if the mass has a width smaller than tolerance. + Otherwise, return a constant mass. + TODO try to cache results for this function in the future for speedup. + """ + particle = Particle.from_evtgen_name(name) + + if particle.width <= tolerance: + return tf.cast(particle.mass, tf.float64) + # If name does not exist in the predefined mass distributions, use Breit-Wigner + return mass_converter[mass_func](mass=particle.mass, width=particle.width) + + +def _recursively_traverse(decaychain: dict, mass_converter: dict[str, Callable], + preexisting_particles: set[str] = None, tolerance: float = _MASS_WIDTH_TOLERANCE) -> list[tuple[float, GenParticle]]: + """ + Create all possible GenParticles by recursively traversing a dict from decaylanguage. + + Parameters + ---------- + decaychain: dict + Decay chain with the format from decaylanguage + preexisting_particles : set + names of all particles that have already been created. + tolerance : float + Minimum mass width for a particle to set a non-constant mass to a particle. + + Returns + ------- + list[tuple[float, GenParticle]] + The generated particle + """ + original_mother_name = list(decaychain.keys())[0] # Get the only key inside the dict + + if preexisting_particles is None: + preexisting_particles = set() + is_top_particle = True + else: + is_top_particle = False + + # This is in the form of dicts + decay_modes = decaychain[original_mother_name] + mother_name = _unique_name(original_mother_name, preexisting_particles) + # This will contain GenParticle instances and their probabilities + all_decays = [] + for dm in decay_modes: + dm_probability = dm['bf'] + daughter_particles = dm['fs'] + daughter_gens = [] + + for daughter_name in daughter_particles: + if isinstance(daughter_name, str): + # Always use constant mass for stable particles + daughter = GenParticle(_unique_name(daughter_name, preexisting_particles), + Particle.from_evtgen_name(daughter_name).mass) + daughter = [(1., daughter)] + elif isinstance(daughter_name, dict): + daughter = _recursively_traverse(daughter_name, mass_converter, preexisting_particles, tolerance=tolerance) + else: + raise TypeError(f'Expected elements in decaychain["fs"] to only be str or dict ' + f'but found an instance of type {type(daughter_name)}') + daughter_gens.append(daughter) + + for daughter_combination in itertools.product(*daughter_gens): + p = tnp.prod([decay[0] for decay in daughter_combination]) * dm_probability + if is_top_particle: + mother_mass = Particle.from_evtgen_name(original_mother_name).mass + else: + mother_mass = _get_particle_mass(original_mother_name, mass_converter=mass_converter, + mass_func=dm.get('zfit', _DEFAULT_MASS_FUNC), tolerance=tolerance) + + one_decay = GenParticle(mother_name, mother_mass).set_children( + *[decay[1] for decay in daughter_combination]) + all_decays.append((p, one_decay)) + + return all_decays diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py new file mode 100644 index 00000000..c6130fdd --- /dev/null +++ b/phasespace/fulldecay/mass_functions.py @@ -0,0 +1,50 @@ +import tensorflow as tf + +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) + + def mass_func(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='') + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator) + return mass_func + + +def breitwigner(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): + 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='') + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator) + return mass_func + + +def relativistic_breitwigner(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): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + pdf = zphys.pdf.RelativisticBreitWigner(m=particle_mass, gamma=particle_width, obs='') + iterator = tf.stack([min_mass, max_mass], axis=-1) + # TODO this works with map_fn but not with vectorized_map for some reason. + # Does not work for e.g., zfit.pdf.CrystalBall either + return tf.map_fn(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator) + return mass_func + + +_DEFAULT_CONVERTER = {'gauss': gauss, 'BW': breitwigner, 'rel-BW': relativistic_breitwigner} diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py new file mode 100644 index 00000000..34ec2b3a --- /dev/null +++ b/tests/fulldecay/example_decay_chains.py @@ -0,0 +1,51 @@ +# A D+ particle with only one way of decaying +dplus_single = {'D+': [{'bf': 1, + 'fs': ['K-', 'pi+', 'pi+', + {'pi0': [{'bf': 1, 'fs': ['gamma', 'gamma']}]}, + ], + 'model': 'PHSP', 'model_params': '', 'zfit': 'rel-BW'}] + } + +pi0_4branches = {'pi0': [{'bf': 0.988228297, 'fs': ['gamma', 'gamma'], 'zfit': 'BW'}, + {'bf': 0.011738247, 'fs': ['e+', 'e-', 'gamma'], 'zfit': 'gauss'}, + {'bf': 3.3392e-5, 'fs': ['e+', 'e+', 'e-', 'e-'], 'zfit': 'rel-BW'}, + {'bf': 6.5e-8, 'fs': ['e+', 'e-']}]} + +dplus_4grandbranches = {'D+': [{'bf': 1.0, + 'fs': ['K-', + 'pi+', + 'pi+', + pi0_4branches], + 'model': 'PHSP', + 'model_params': ''}]} + +dstarplus_big_decay = {'D*+': [{'bf': 0.677, + 'fs': [{'D0': [{'bf': 1.0, + 'fs': ['K-', 'pi+'], + 'model': 'PHSP', + 'model_params': ''}]}, + 'pi+'], + 'model': 'VSS', + 'model_params': ''}, + {'bf': 0.307, + 'fs': [{'D+': [{'bf': 1.0, + 'fs': ['K-', + 'pi+', + 'pi+', + pi0_4branches], + 'model': 'PHSP', + 'model_params': ''}]}, + pi0_4branches], + 'model': 'VSS', + 'model_params': ''}, + {'bf': 0.016, + 'fs': [{'D+': [{'bf': 1.0, + 'fs': ['K-', + 'pi+', + 'pi+', + pi0_4branches], + 'model': 'PHSP', + 'model_params': ''}]}, + 'gamma'], + 'model': 'VSP_PWAVE', + 'model_params': ''}]} \ No newline at end of file diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py new file mode 100644 index 00000000..35c4451e --- /dev/null +++ b/tests/fulldecay/test_fulldecay.py @@ -0,0 +1,88 @@ +from phasespace.fulldecay import FullDecay +from example_decay_chains import * # TODO remove * since it is bad practice? + +from numpy.testing import assert_almost_equal + + +def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: + """ + Checks whether the normalize_weights argument works for FullDecay.generate + + Parameters + ---------- + full_decay : FullDecay + full_decay.generate will be called. + kwargs + Additional parameters passed to generate. + + Returns + ------- + list[tuple] + All the values returned by generate, both times. + The return arguments from normalize_weights=True is the first element in the returned list. + + Notes + ----- + The function is called check_norm instead of test_norm since it is used by other functions and is not a stand-alone test. + """ + all_return_args = [] + for norm in (True, False): + return_args = full_decay.generate(normalize_weights=norm, **kwargs) + assert len(return_args) == 2 if norm else 3 + assert sum(len(w) for w in return_args[0]) == kwargs['n_events'] + if not norm: + assert all(len(w) == len(mw) for w, mw in zip(return_args[0], return_args[1])) + + all_return_args.append(return_args) + + return all_return_args + + +def test_single_chain(): + """Test converting a decaylanguage dict with only one possible decay.""" + container = FullDecay.from_dict(dplus_single, tolerance=1e-10) + output_decay = container.gen_particles + assert len(output_decay) == 1 + prob, gen = output_decay[0] + assert_almost_equal(prob, 1) + assert gen.name == 'D+' + assert {p.name for p in gen.children} == {'K-', 'pi+', 'pi+ [0]', 'pi0'} + for p in gen.children: + if 'pi0' == p.name[:3]: + assert not p.has_fixed_mass + else: + assert p.has_fixed_mass + + check_norm(container, n_events=1) + (normed_weights, decay_list), _ = check_norm(container, n_events=100) + assert len(decay_list) == 1 + events = decay_list[0] + assert set(events.keys()) == {'K-', 'pi+', 'pi+ [0]', 'pi0', 'gamma', 'gamma [0]'} + assert all(len(p) == 100 for p in events.values()) + + +def test_branching_children(): + container = FullDecay.from_dict(pi0_4branches, 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) + check_norm(container, n_events=1) + (normed_weights, events), _ = check_norm(container, n_events=100) + + +def test_branching_grandchilden(): + container = FullDecay.from_dict(dplus_4grandbranches) + output_decays = container.gen_particles + assert_almost_equal(sum(d[0] for d in output_decays), 1) + check_norm(container, n_events=1) + (normed_weights, events), _ = check_norm(container, n_events=100) + # TODO add more asserts here + + +def test_big_decay(): + container = FullDecay.from_dict(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) + (normed_weights, events), _ = check_norm(container, n_events=100) + # TODO add more asserts here diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py new file mode 100644 index 00000000..ba73b79a --- /dev/null +++ b/tests/fulldecay/test_mass_functions.py @@ -0,0 +1,52 @@ +import pytest +import tensorflow as tf +import tensorflow_probability as tfp +import phasespace.fulldecay.mass_functions as mf + +from typing import Callable + + +KSTARZ_MASS = 895.81 +KSTARZ_WIDTH = 47.4 +def ref_mass_func(min_mass, max_mass, n_events): + """ + Reference mass function used to compare the behavior of the actual mass functions. + + Parameters + ---------- + min_mass + max_mass + n_events + + Returns + ------- + kstar_mass + Mass generated + + Notes + ----- + Code taken from phasespace documentation. + """ + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, + scale=kstar_width_cast, + low=min_mass, + high=max_mass).sample() + return kstar_mass + + +@pytest.mark.parametrize('function', (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner)) +@pytest.mark.parametrize('size', (1, 10)) +def test_shape(function: Callable, size: int, params: tuple = (1., 1.)): + rng = tf.random.Generator.from_seed(1234) + min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) + min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) + assert tf.reduce_all(min_mass <= max_mass) + ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) + sample = function(*params)(min_mass, max_mass, len(min_mass)) + assert sample.shape[0] == ref_sample.shape[0] + assert all(i <= 1 for i in sample.shape[1:]) # Sample.shape have extra dimensions with just 1 or 0, but code still seems to work From 79d1ccf02ee0457ce372a553ed43524f8c66a3e4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Sep 2021 14:33:07 +0000 Subject: [PATCH 02/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- phasespace/fulldecay/__init__.py | 5 +- phasespace/fulldecay/fulldecay.py | 118 +++++++++++++------- phasespace/fulldecay/mass_functions.py | 30 +++-- tests/fulldecay/example_decay_chains.py | 139 ++++++++++++++++-------- tests/fulldecay/test_fulldecay.py | 21 ++-- tests/fulldecay/test_mass_functions.py | 30 ++--- 6 files changed, 221 insertions(+), 122 deletions(-) diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py index c48943b8..0ad22ad9 100644 --- a/phasespace/fulldecay/__init__.py +++ b/phasespace/fulldecay/__init__.py @@ -1,10 +1,11 @@ +import sys + from .fulldecay import FullDecay -import sys try: - from particle import Particle import zfit import zfit_physics as zphys + from particle import Particle except ModuleNotFoundError: print( "the fulldecay functionality in phasespace requires particle and zfit-physics. " diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index aa1e04f4..1b4705cc 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -1,24 +1,23 @@ -from phasespace import GenParticle +import itertools +from typing import Callable, Union + import tensorflow as tf import tensorflow.experimental.numpy as tnp from particle import Particle -from .mass_functions import _DEFAULT_CONVERTER +from phasespace import GenParticle -from typing import Callable, Union -import itertools +from .mass_functions import _DEFAULT_CONVERTER _MASS_WIDTH_TOLERANCE = 0.01 -_DEFAULT_MASS_FUNC = 'rel-BW' +_DEFAULT_MASS_FUNC = "rel-BW" class FullDecay: - """ - A container that works like GenParticle that can handle multiple decays - """ + """A container that works like GenParticle that can handle multiple decays.""" + def __init__(self, gen_particles: list[tuple[float, GenParticle]]): - """ - Create an instance of FullDecay + """Create an instance of FullDecay. Parameters ---------- @@ -32,9 +31,13 @@ def __init__(self, gen_particles: list[tuple[float, GenParticle]]): self.gen_particles = gen_particles @classmethod - def from_dict(cls, dec_dict: dict, mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE): - """ - Create a FullDecay instance from a dict in the decaylanguage format. + def from_dict( + cls, + dec_dict: dict, + mass_converter: dict[str, Callable] = None, + tolerance: float = _MASS_WIDTH_TOLERANCE, + ): + """Create a FullDecay instance from a dict in the decaylanguage format. Parameters ---------- @@ -59,13 +62,18 @@ def from_dict(cls, dec_dict: dict, mass_converter: dict[str, Callable] = None, t # Combine the mass functions specified by the package to the mass functions specified from the input. total_mass_converter = {**_DEFAULT_CONVERTER, **mass_converter} - gen_particles = _recursively_traverse(dec_dict, total_mass_converter, tolerance=tolerance) + gen_particles = _recursively_traverse( + dec_dict, total_mass_converter, tolerance=tolerance + ) return cls(gen_particles) - def generate(self, n_events: int, normalize_weights: bool = False, - **kwargs) -> Union[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). + def generate( + self, n_events: int, normalize_weights: bool = False, **kwargs + ) -> Union[ + 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). Parameters ---------- @@ -85,13 +93,17 @@ def generate(self, n_events: int, normalize_weights: bool = False, Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. """ # Input to tf.random.categorical must be 2D - rand_i = tf.random.categorical(tnp.log([[dm[0] for dm in self.gen_particles]]), n_events) + rand_i = tf.random.categorical( + tnp.log([[dm[0] for dm in self.gen_particles]]), n_events + ) # Input to tf.unique_with_counts must be 1D dec_indices, _, counts = tf.unique_with_counts(rand_i[0]) counts = tf.cast(counts, tf.int64) weights, max_weights, events = [], [], [] for i, n in zip(dec_indices, counts): - weight, max_weight, four_vectors = self.gen_particles[i][1].generate(n, normalize_weights=False, **kwargs) + weight, max_weight, four_vectors = self.gen_particles[i][1].generate( + n, normalize_weights=False, **kwargs + ) weights.append(weight) max_weights.append(max_weight) events.append(four_vectors) @@ -105,8 +117,7 @@ def generate(self, n_events: int, normalize_weights: bool = False, def _unique_name(name: str, preexisting_particles: set[str]) -> str: - """ - Create a string that does not exist in preexisting_particles based on name. + """Create a string that does not exist in preexisting_particles based on name. Parameters ---------- @@ -125,17 +136,21 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: preexisting_particles.add(name) return name - name += ' [0]' + name += " [0]" i = 1 while name in preexisting_particles: - name = name[:name.rfind('[')] + f'[{str(i)}]' + name = name[: name.rfind("[")] + f"[{str(i)}]" i += 1 preexisting_particles.add(name) return name -def _get_particle_mass(name: str, mass_converter: dict[str, Callable], mass_func: str, - tolerance: float = _MASS_WIDTH_TOLERANCE) -> Union[Callable, float]: +def _get_particle_mass( + name: str, + mass_converter: dict[str, Callable], + mass_func: str, + tolerance: float = _MASS_WIDTH_TOLERANCE, +) -> Union[Callable, float]: """ Get mass or mass function of particle using the particle package. Parameters @@ -160,10 +175,13 @@ def _get_particle_mass(name: str, mass_converter: dict[str, Callable], mass_func return mass_converter[mass_func](mass=particle.mass, width=particle.width) -def _recursively_traverse(decaychain: dict, mass_converter: dict[str, Callable], - preexisting_particles: set[str] = None, tolerance: float = _MASS_WIDTH_TOLERANCE) -> list[tuple[float, GenParticle]]: - """ - Create all possible GenParticles by recursively traversing a dict from decaylanguage. +def _recursively_traverse( + decaychain: dict, + mass_converter: dict[str, Callable], + preexisting_particles: set[str] = None, + tolerance: float = _MASS_WIDTH_TOLERANCE, +) -> list[tuple[float, GenParticle]]: + """Create all possible GenParticles by recursively traversing a dict from decaylanguage. Parameters ---------- @@ -179,7 +197,9 @@ def _recursively_traverse(decaychain: dict, mass_converter: dict[str, Callable], list[tuple[float, GenParticle]] The generated particle """ - original_mother_name = list(decaychain.keys())[0] # Get the only key inside the dict + original_mother_name = list(decaychain.keys())[ + 0 + ] # Get the only key inside the dict if preexisting_particles is None: preexisting_particles = set() @@ -193,21 +213,30 @@ def _recursively_traverse(decaychain: dict, mass_converter: dict[str, Callable], # This will contain GenParticle instances and their probabilities all_decays = [] for dm in decay_modes: - dm_probability = dm['bf'] - daughter_particles = dm['fs'] + dm_probability = dm["bf"] + daughter_particles = dm["fs"] daughter_gens = [] for daughter_name in daughter_particles: if isinstance(daughter_name, str): # Always use constant mass for stable particles - daughter = GenParticle(_unique_name(daughter_name, preexisting_particles), - Particle.from_evtgen_name(daughter_name).mass) - daughter = [(1., daughter)] + daughter = GenParticle( + _unique_name(daughter_name, preexisting_particles), + Particle.from_evtgen_name(daughter_name).mass, + ) + daughter = [(1.0, daughter)] elif isinstance(daughter_name, dict): - daughter = _recursively_traverse(daughter_name, mass_converter, preexisting_particles, tolerance=tolerance) + daughter = _recursively_traverse( + daughter_name, + mass_converter, + preexisting_particles, + tolerance=tolerance, + ) else: - raise TypeError(f'Expected elements in decaychain["fs"] to only be str or dict ' - f'but found an instance of type {type(daughter_name)}') + raise TypeError( + f'Expected elements in decaychain["fs"] to only be str or dict ' + f"but found an instance of type {type(daughter_name)}" + ) daughter_gens.append(daughter) for daughter_combination in itertools.product(*daughter_gens): @@ -215,11 +244,16 @@ def _recursively_traverse(decaychain: dict, mass_converter: dict[str, Callable], if is_top_particle: mother_mass = Particle.from_evtgen_name(original_mother_name).mass else: - mother_mass = _get_particle_mass(original_mother_name, mass_converter=mass_converter, - mass_func=dm.get('zfit', _DEFAULT_MASS_FUNC), tolerance=tolerance) + mother_mass = _get_particle_mass( + original_mother_name, + mass_converter=mass_converter, + mass_func=dm.get("zfit", _DEFAULT_MASS_FUNC), + tolerance=tolerance, + ) one_decay = GenParticle(mother_name, mother_mass).set_children( - *[decay[1] for decay in daughter_combination]) + *(decay[1] for decay in daughter_combination) + ) all_decays.append((p, one_decay)) return all_decays diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py index c6130fdd..00fcde09 100644 --- a/phasespace/fulldecay/mass_functions.py +++ b/phasespace/fulldecay/mass_functions.py @@ -1,5 +1,4 @@ import tensorflow as tf - import zfit import zfit_physics as zphys @@ -13,9 +12,12 @@ def gauss(mass, width): def mass_func(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='') + pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs="") iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + return mass_func @@ -26,9 +28,12 @@ def breitwigner(mass, width): def mass_func(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='') + pdf = zfit.pdf.Cauchy(m=particle_mass, gamma=particle_width, obs="") iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + return mass_func @@ -39,12 +44,21 @@ def relativistic_breitwigner(mass, width): def mass_func(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(m=particle_mass, gamma=particle_width, obs='') + pdf = zphys.pdf.RelativisticBreitWigner( + m=particle_mass, gamma=particle_width, obs="" + ) iterator = tf.stack([min_mass, max_mass], axis=-1) # TODO this works with map_fn but not with vectorized_map for some reason. # Does not work for e.g., zfit.pdf.CrystalBall either - return tf.map_fn(lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator) + return tf.map_fn( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator + ) + return mass_func -_DEFAULT_CONVERTER = {'gauss': gauss, 'BW': breitwigner, 'rel-BW': relativistic_breitwigner} +_DEFAULT_CONVERTER = { + "gauss": gauss, + "BW": breitwigner, + "rel-BW": relativistic_breitwigner, +} diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index 34ec2b3a..aee3c638 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -1,51 +1,96 @@ # A D+ particle with only one way of decaying -dplus_single = {'D+': [{'bf': 1, - 'fs': ['K-', 'pi+', 'pi+', - {'pi0': [{'bf': 1, 'fs': ['gamma', 'gamma']}]}, - ], - 'model': 'PHSP', 'model_params': '', 'zfit': 'rel-BW'}] - } +dplus_single = { + "D+": [ + { + "bf": 1, + "fs": [ + "K-", + "pi+", + "pi+", + {"pi0": [{"bf": 1, "fs": ["gamma", "gamma"]}]}, + ], + "model": "PHSP", + "model_params": "", + "zfit": "rel-BW", + } + ] +} -pi0_4branches = {'pi0': [{'bf': 0.988228297, 'fs': ['gamma', 'gamma'], 'zfit': 'BW'}, - {'bf': 0.011738247, 'fs': ['e+', 'e-', 'gamma'], 'zfit': 'gauss'}, - {'bf': 3.3392e-5, 'fs': ['e+', 'e+', 'e-', 'e-'], 'zfit': 'rel-BW'}, - {'bf': 6.5e-8, 'fs': ['e+', 'e-']}]} +pi0_4branches = { + "pi0": [ + {"bf": 0.988228297, "fs": ["gamma", "gamma"], "zfit": "BW"}, + {"bf": 0.011738247, "fs": ["e+", "e-", "gamma"], "zfit": "gauss"}, + {"bf": 3.3392e-5, "fs": ["e+", "e+", "e-", "e-"], "zfit": "rel-BW"}, + {"bf": 6.5e-8, "fs": ["e+", "e-"]}, + ] +} -dplus_4grandbranches = {'D+': [{'bf': 1.0, - 'fs': ['K-', - 'pi+', - 'pi+', - pi0_4branches], - 'model': 'PHSP', - 'model_params': ''}]} +dplus_4grandbranches = { + "D+": [ + { + "bf": 1.0, + "fs": ["K-", "pi+", "pi+", pi0_4branches], + "model": "PHSP", + "model_params": "", + } + ] +} -dstarplus_big_decay = {'D*+': [{'bf': 0.677, - 'fs': [{'D0': [{'bf': 1.0, - 'fs': ['K-', 'pi+'], - 'model': 'PHSP', - 'model_params': ''}]}, - 'pi+'], - 'model': 'VSS', - 'model_params': ''}, - {'bf': 0.307, - 'fs': [{'D+': [{'bf': 1.0, - 'fs': ['K-', - 'pi+', - 'pi+', - pi0_4branches], - 'model': 'PHSP', - 'model_params': ''}]}, - pi0_4branches], - 'model': 'VSS', - 'model_params': ''}, - {'bf': 0.016, - 'fs': [{'D+': [{'bf': 1.0, - 'fs': ['K-', - 'pi+', - 'pi+', - pi0_4branches], - 'model': 'PHSP', - 'model_params': ''}]}, - 'gamma'], - 'model': 'VSP_PWAVE', - 'model_params': ''}]} \ No newline at end of file +dstarplus_big_decay = { + "D*+": [ + { + "bf": 0.677, + "fs": [ + { + "D0": [ + { + "bf": 1.0, + "fs": ["K-", "pi+"], + "model": "PHSP", + "model_params": "", + } + ] + }, + "pi+", + ], + "model": "VSS", + "model_params": "", + }, + { + "bf": 0.307, + "fs": [ + { + "D+": [ + { + "bf": 1.0, + "fs": ["K-", "pi+", "pi+", pi0_4branches], + "model": "PHSP", + "model_params": "", + } + ] + }, + pi0_4branches, + ], + "model": "VSS", + "model_params": "", + }, + { + "bf": 0.016, + "fs": [ + { + "D+": [ + { + "bf": 1.0, + "fs": ["K-", "pi+", "pi+", pi0_4branches], + "model": "PHSP", + "model_params": "", + } + ] + }, + "gamma", + ], + "model": "VSP_PWAVE", + "model_params": "", + }, + ] +} diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index 35c4451e..3874ac80 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -1,12 +1,11 @@ -from phasespace.fulldecay import FullDecay from example_decay_chains import * # TODO remove * since it is bad practice? - from numpy.testing import assert_almost_equal +from phasespace.fulldecay import FullDecay + def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: - """ - Checks whether the normalize_weights argument works for FullDecay.generate + """Checks whether the normalize_weights argument works for FullDecay.generate. Parameters ---------- @@ -29,9 +28,11 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: for norm in (True, False): return_args = full_decay.generate(normalize_weights=norm, **kwargs) assert len(return_args) == 2 if norm else 3 - assert sum(len(w) for w in return_args[0]) == kwargs['n_events'] + assert sum(len(w) for w in return_args[0]) == kwargs["n_events"] if not norm: - assert all(len(w) == len(mw) for w, mw in zip(return_args[0], return_args[1])) + assert all( + len(w) == len(mw) for w, mw in zip(return_args[0], return_args[1]) + ) all_return_args.append(return_args) @@ -45,10 +46,10 @@ def test_single_chain(): assert len(output_decay) == 1 prob, gen = output_decay[0] assert_almost_equal(prob, 1) - assert gen.name == 'D+' - assert {p.name for p in gen.children} == {'K-', 'pi+', 'pi+ [0]', 'pi0'} + assert gen.name == "D+" + assert {p.name for p in gen.children} == {"K-", "pi+", "pi+ [0]", "pi0"} for p in gen.children: - if 'pi0' == p.name[:3]: + if "pi0" == p.name[:3]: assert not p.has_fixed_mass else: assert p.has_fixed_mass @@ -57,7 +58,7 @@ def test_single_chain(): (normed_weights, decay_list), _ = check_norm(container, n_events=100) assert len(decay_list) == 1 events = decay_list[0] - assert set(events.keys()) == {'K-', 'pi+', 'pi+ [0]', 'pi0', 'gamma', 'gamma [0]'} + assert set(events.keys()) == {"K-", "pi+", "pi+ [0]", "pi0", "gamma", "gamma [0]"} assert all(len(p) == 100 for p in events.values()) diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py index ba73b79a..2ef1e23b 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fulldecay/test_mass_functions.py @@ -1,16 +1,17 @@ +from typing import Callable + import pytest import tensorflow as tf import tensorflow_probability as tfp -import phasespace.fulldecay.mass_functions as mf - -from typing import Callable +import phasespace.fulldecay.mass_functions as mf KSTARZ_MASS = 895.81 KSTARZ_WIDTH = 47.4 + + def ref_mass_func(min_mass, max_mass, n_events): - """ - Reference mass function used to compare the behavior of the actual mass functions. + """Reference mass function used to compare the behavior of the actual mass functions. Parameters ---------- @@ -32,16 +33,17 @@ def ref_mass_func(min_mass, max_mass, n_events): kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, - scale=kstar_width_cast, - low=min_mass, - high=max_mass).sample() + kstar_mass = tfp.distributions.TruncatedNormal( + loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass + ).sample() return kstar_mass -@pytest.mark.parametrize('function', (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner)) -@pytest.mark.parametrize('size', (1, 10)) -def test_shape(function: Callable, size: int, params: tuple = (1., 1.)): +@pytest.mark.parametrize( + "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) +) +@pytest.mark.parametrize("size", (1, 10)) +def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): rng = tf.random.Generator.from_seed(1234) min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) @@ -49,4 +51,6 @@ def test_shape(function: Callable, size: int, params: tuple = (1., 1.)): ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) sample = function(*params)(min_mass, max_mass, len(min_mass)) assert sample.shape[0] == ref_sample.shape[0] - assert all(i <= 1 for i in sample.shape[1:]) # Sample.shape have extra dimensions with just 1 or 0, but code still seems to work + assert all( + i <= 1 for i in sample.shape[1:] + ) # Sample.shape have extra dimensions with just 1 or 0, but code still seems to work From 6d5e448b990e4d7f31cccdd2be435044422c9c06 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 26 Sep 2021 18:26:58 +0200 Subject: [PATCH 03/71] Syntactic update to code for extracting mother name --- phasespace/fulldecay/fulldecay.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 1b4705cc..489cb03b 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -197,9 +197,8 @@ def _recursively_traverse( list[tuple[float, GenParticle]] The generated particle """ - original_mother_name = list(decaychain.keys())[ - 0 - ] # Get the only key inside the dict + # Get the only key inside the decaychain dict + original_mother_name, = decaychain.keys() if preexisting_particles is None: preexisting_particles = set() From ca5c115a67f17af73cdb4d11f59d248dbf102da0 Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Mon, 27 Sep 2021 14:23:29 +0200 Subject: [PATCH 04/71] Better error handling for module imports Co-authored-by: Jonas Eschle --- phasespace/fulldecay/__init__.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py index 0ad22ad9..c5c74cd5 100644 --- a/phasespace/fulldecay/__init__.py +++ b/phasespace/fulldecay/__init__.py @@ -6,10 +6,10 @@ import zfit import zfit_physics as zphys from particle import Particle -except ModuleNotFoundError: - print( - "the fulldecay functionality in phasespace requires particle and zfit-physics. " +except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "The fulldecay functionality in phasespace requires particle and zfit-physics. " "Either install phasespace[fulldecay] or particle and zfit-physics.", file=sys.stderr, - ) - raise + ) from error + From 54a789c2906a7a0f2ca67b8b6ce32891a651ccbf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Sep 2021 12:23:46 +0000 Subject: [PATCH 05/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- phasespace/fulldecay/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py index c5c74cd5..42f877ee 100644 --- a/phasespace/fulldecay/__init__.py +++ b/phasespace/fulldecay/__init__.py @@ -12,4 +12,3 @@ "Either install phasespace[fulldecay] or particle and zfit-physics.", file=sys.stderr, ) from error - From 041950197f842e3386d85c68ff5c2852a7f4b37f Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 14:29:44 +0200 Subject: [PATCH 06/71] Add pytest assert introspection to check_norm helper function --- tests/fulldecay/__init__.py | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 tests/fulldecay/__init__.py diff --git a/tests/fulldecay/__init__.py b/tests/fulldecay/__init__.py new file mode 100644 index 00000000..8b6e34cb --- /dev/null +++ b/tests/fulldecay/__init__.py @@ -0,0 +1,3 @@ +import pytest + +pytest.register_assert_rewrite('fulldecay.test_fulldecay.check_norm') From 968cf0b7a61259e1042ba6886c8941403dbf7e88 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 21:49:12 +0200 Subject: [PATCH 07/71] Better assertion errors for check_norm --- tests/fulldecay/__init__.py | 3 ++- tests/fulldecay/test_fulldecay.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/fulldecay/__init__.py b/tests/fulldecay/__init__.py index 8b6e34cb..cecc0a43 100644 --- a/tests/fulldecay/__init__.py +++ b/tests/fulldecay/__init__.py @@ -1,3 +1,4 @@ import pytest -pytest.register_assert_rewrite('fulldecay.test_fulldecay.check_norm') +# This makes it so that assert errors are more helpful for e.g., the check_norm helper function +pytest.register_assert_rewrite('fulldecay.test_fulldecay') diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index 3874ac80..46b853c3 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -1,4 +1,4 @@ -from example_decay_chains import * # TODO remove * since it is bad practice? +from .example_decay_chains import * # TODO remove * since it is bad practice? from numpy.testing import assert_almost_equal from phasespace.fulldecay import FullDecay @@ -24,6 +24,8 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: ----- The function is called check_norm instead of test_norm since it is used by other functions and is not a stand-alone test. """ + a = 2 + assert a == 1 all_return_args = [] for norm in (True, False): return_args = full_decay.generate(normalize_weights=norm, **kwargs) From 03a8c805e10d944b123f6358056c47201c5ff4c6 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 21:55:16 +0200 Subject: [PATCH 08/71] Remove forgotten dummy assert --- tests/fulldecay/test_fulldecay.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index 46b853c3..bc3d7651 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -24,8 +24,6 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: ----- The function is called check_norm instead of test_norm since it is used by other functions and is not a stand-alone test. """ - a = 2 - assert a == 1 all_return_args = [] for norm in (True, False): return_args = full_decay.generate(normalize_weights=norm, **kwargs) From 60ab98243c8cf265604594db78b28d46bfd824fd Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 21:56:39 +0200 Subject: [PATCH 09/71] Remove hardcoded K*0 mass and width with call to particle --- tests/fulldecay/test_mass_functions.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py index 2ef1e23b..342b1bf5 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fulldecay/test_mass_functions.py @@ -3,11 +3,13 @@ import pytest import tensorflow as tf import tensorflow_probability as tfp +from particle import Particle import phasespace.fulldecay.mass_functions as mf -KSTARZ_MASS = 895.81 -KSTARZ_WIDTH = 47.4 +_kstarz = Particle.from_evtgen_name('K*0') +KSTARZ_MASS = _kstarz.mass +KSTARZ_WIDTH = _kstarz.width def ref_mass_func(min_mass, max_mass, n_events): @@ -53,4 +55,4 @@ def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): assert sample.shape[0] == ref_sample.shape[0] assert all( i <= 1 for i in sample.shape[1:] - ) # Sample.shape have extra dimensions with just 1 or 0, but code still seems to work + ) # Sample.shape have extra dimensions with just 1 or 0, which can be ignored From 8627d134c5ec6935ce281b8c54f0e0c0cedeb54b Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 22:01:06 +0200 Subject: [PATCH 10/71] Move docstring for FullDecay to __init__ --- phasespace/fulldecay/fulldecay.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 489cb03b..98ce8db2 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -14,10 +14,9 @@ class FullDecay: - """A container that works like GenParticle that can handle multiple decays.""" - def __init__(self, gen_particles: list[tuple[float, GenParticle]]): - """Create an instance of FullDecay. + """ + A container that works like GenParticle that can handle multiple decays. Can be created from Parameters ---------- From 436a6d4561a38b0fecddcda683e882539a746c4c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Sep 2021 20:01:38 +0000 Subject: [PATCH 11/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- phasespace/fulldecay/fulldecay.py | 5 ++--- tests/fulldecay/__init__.py | 2 +- tests/fulldecay/test_fulldecay.py | 3 ++- tests/fulldecay/test_mass_functions.py | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 98ce8db2..813752e3 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -15,8 +15,7 @@ class FullDecay: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): - """ - A container that works like GenParticle that can handle multiple decays. Can be created from + """A container that works like GenParticle that can handle multiple decays. Can be created from. Parameters ---------- @@ -197,7 +196,7 @@ def _recursively_traverse( The generated particle """ # Get the only key inside the decaychain dict - original_mother_name, = decaychain.keys() + (original_mother_name,) = decaychain.keys() if preexisting_particles is None: preexisting_particles = set() diff --git a/tests/fulldecay/__init__.py b/tests/fulldecay/__init__.py index cecc0a43..b98e9386 100644 --- a/tests/fulldecay/__init__.py +++ b/tests/fulldecay/__init__.py @@ -1,4 +1,4 @@ import pytest # This makes it so that assert errors are more helpful for e.g., the check_norm helper function -pytest.register_assert_rewrite('fulldecay.test_fulldecay') +pytest.register_assert_rewrite("fulldecay.test_fulldecay") diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index bc3d7651..67bec9c1 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -1,8 +1,9 @@ -from .example_decay_chains import * # TODO remove * since it is bad practice? from numpy.testing import assert_almost_equal from phasespace.fulldecay import FullDecay +from .example_decay_chains import * # TODO remove * since it is bad practice? + def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: """Checks whether the normalize_weights argument works for FullDecay.generate. diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py index 342b1bf5..3ec2492b 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fulldecay/test_mass_functions.py @@ -7,7 +7,7 @@ import phasespace.fulldecay.mass_functions as mf -_kstarz = Particle.from_evtgen_name('K*0') +_kstarz = Particle.from_evtgen_name("K*0") KSTARZ_MASS = _kstarz.mass KSTARZ_WIDTH = _kstarz.width From 06f0713f1a22ca0bda102c3f4cccc69d0138b3d6 Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Mon, 27 Sep 2021 22:08:25 +0200 Subject: [PATCH 12/71] Make comment about map_fn more descriptive Co-authored-by: Jonas Eschle --- phasespace/fulldecay/mass_functions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py index 00fcde09..0567696e 100644 --- a/phasespace/fulldecay/mass_functions.py +++ b/phasespace/fulldecay/mass_functions.py @@ -48,7 +48,8 @@ def mass_func(min_mass, max_mass, n_events): m=particle_mass, gamma=particle_width, obs="" ) iterator = tf.stack([min_mass, max_mass], axis=-1) - # TODO this works with map_fn but not with vectorized_map for some reason. + # TODO this works with map_fn but not with vectorized_map as no analytic sampling is available. + # Does not work for e.g., zfit.pdf.CrystalBall either return tf.map_fn( lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator From f5f0ca93f89b3b05eaf915169e9a38c6ffb3185a Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Mon, 27 Sep 2021 22:11:03 +0200 Subject: [PATCH 13/71] Remove average from `from_dict` docstring Co-authored-by: Jonas Eschle --- phasespace/fulldecay/fulldecay.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 813752e3..8ef1149e 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -43,7 +43,8 @@ def from_dict( The input dict from which the FullDecay object will be created from. mass_converter : dict[str, Callable] A dict with mass function names and their corresponding mass functions. - These functions should take the average particle mass and the mass width as inputs + These functions should take the particle mass and the mass width as inputs + and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. tolerance : float From 6d4268d80d6d93b9181495235826b532701f1382 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 22:28:00 +0200 Subject: [PATCH 14/71] Remove part of comment --- phasespace/fulldecay/mass_functions.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py index 0567696e..cca0ccb9 100644 --- a/phasespace/fulldecay/mass_functions.py +++ b/phasespace/fulldecay/mass_functions.py @@ -48,9 +48,8 @@ def mass_func(min_mass, max_mass, n_events): m=particle_mass, gamma=particle_width, obs="" ) iterator = tf.stack([min_mass, max_mass], axis=-1) - # TODO this works with map_fn but not with vectorized_map as no analytic sampling is available. - # Does not work for e.g., zfit.pdf.CrystalBall either + # TODO this works with map_fn but not with vectorized_map as no analytic sampling is available. return tf.map_fn( lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) From 9e58f4a649df28189e9fcb4803b8d6b1219ad42f Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Mon, 27 Sep 2021 22:28:53 +0200 Subject: [PATCH 15/71] Update docstrings to Google style instead of NumPy style. Also made some minor clarifications to the docstrings. --- phasespace/fulldecay/fulldecay.py | 121 +++++++++---------------- tests/fulldecay/test_fulldecay.py | 24 ++--- tests/fulldecay/test_mass_functions.py | 24 ++--- 3 files changed, 63 insertions(+), 106 deletions(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 8ef1149e..6ee9963b 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -17,14 +17,12 @@ class FullDecay: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): """A container that works like GenParticle that can handle multiple decays. Can be created from. - Parameters - ---------- - gen_particles : list[tuple[float, GenParticle]] - All the GenParticles and their corresponding probabilities. + Args: + gen_particles: All the GenParticles and their corresponding probabilities. The list must be of the format [[probability, GenParticle instance], [probability, ... - Notes - ----- - Input format might change + + Notes: + Input format might change """ self.gen_particles = gen_particles @@ -37,22 +35,15 @@ def from_dict( ): """Create a FullDecay instance from a dict in the decaylanguage format. - Parameters - ---------- - dec_dict : dict - The input dict from which the FullDecay object will be created from. - mass_converter : dict[str, Callable] - A dict with mass function names and their corresponding mass functions. - These functions should take the particle mass and the mass width as inputs - - and return a mass function that phasespace can understand. - This dict will be combined with the predefined mass functions in this package. - tolerance : float - Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. - - Returns - ------- - FullDecay + Args: + dec_dict: The input dict from which the FullDecay object will be created from. + mass_converter: A dict with mass function names and their corresponding mass functions. + These functions should take the particle mass and the mass width as inputs + and return a mass function that phasespace can understand. + This dict will be combined with the predefined mass functions in this package. + tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. + + Returns: The created FullDecay object. """ if mass_converter is None: @@ -74,22 +65,17 @@ def generate( ]: """Generate four-momentum vectors from the decay(s). - Parameters - ---------- - n_events : int - Total number of events combined, for all the decays. - normalize_weights : bool - Normalize weights according to all events generated. This also changes the return values. - See the phasespace documentation for more details. - kwargs - Additional parameters passed to all calls of GenParticle.generate - - Returns - ------- - The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. - However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding - to the return arguments from the corresponding GenParticle instances in self.gen_particles. - Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. + Args: + n_events: Total number of events combined, for all the decays. + normalize_weights: Normalize weights according to all events generated. This also changes the return values. + See the phasespace documentation for more details. + kwargs: Additional parameters passed to all calls of GenParticle.generate + + Returns: + The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. + However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding + to the return arguments from the corresponding GenParticle instances in self.gen_particles. + Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. """ # Input to tf.random.categorical must be 2D rand_i = tf.random.categorical( @@ -118,18 +104,13 @@ def generate( def _unique_name(name: str, preexisting_particles: set[str]) -> str: """Create a string that does not exist in preexisting_particles based on name. - Parameters - ---------- - name : str - Name that should be - preexisting_particles : set[str] - Preexisting names - - Returns - ------- - name : str - Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i will begin at 0 - and increase until the name is not preexisting_particles. + Args: + name: Original name + preexisting_particles: Names that the particle cannot have as name. + + Returns: + name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i begins at 0 + and increases until the name is not preexisting_particles. """ if name not in preexisting_particles: preexisting_particles.add(name) @@ -152,18 +133,12 @@ def _get_particle_mass( ) -> Union[Callable, float]: """ Get mass or mass function of particle using the particle package. - Parameters - ---------- - name : str - Name of the particle. Name must be recognizable by the particle package. - tolerance : float - See _recursively_traverse - - Returns - ------- - Callable, float - Returns a function if the mass has a width smaller than tolerance. - Otherwise, return a constant mass. + Args: + name: Name of the particle. Name must be recognizable by the particle package. + tolerance : See _recursively_traverse + + Returns: + A function if the mass has a width smaller than tolerance. Otherwise, return a constant mass. TODO try to cache results for this function in the future for speedup. """ particle = Particle.from_evtgen_name(name) @@ -182,19 +157,13 @@ def _recursively_traverse( ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from decaylanguage. - Parameters - ---------- - decaychain: dict - Decay chain with the format from decaylanguage - preexisting_particles : set - names of all particles that have already been created. - tolerance : float - Minimum mass width for a particle to set a non-constant mass to a particle. - - Returns - ------- - list[tuple[float, GenParticle]] - The generated particle + Args: + 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. + + Returns: + The generated GenParticle instances, one for each possible way of the decay. """ # Get the only key inside the decaychain dict (original_mother_name,) = decaychain.keys() diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index 67bec9c1..a26142a9 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -6,24 +6,14 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: - """Checks whether the normalize_weights argument works for FullDecay.generate. - - Parameters - ---------- - full_decay : FullDecay - full_decay.generate will be called. - kwargs - Additional parameters passed to generate. - - Returns - ------- - list[tuple] + """Helper function that checks whether the normalize_weights argument works for FullDecay.generate. + Args: + full_decay: full_decay.generate will be called. + kwargs: Additional parameters passed to generate. + + Returns: All the values returned by generate, both times. The return arguments from normalize_weights=True is the first element in the returned list. - - Notes - ----- - The function is called check_norm instead of test_norm since it is used by other functions and is not a stand-alone test. """ all_return_args = [] for norm in (True, False): @@ -64,6 +54,7 @@ def test_single_chain(): def test_branching_children(): + """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" container = FullDecay.from_dict(pi0_4branches, tolerance=1e-10) output_decays = container.gen_particles assert len(output_decays) == 4 @@ -73,6 +64,7 @@ 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 = FullDecay.from_dict(dplus_4grandbranches) output_decays = container.gen_particles assert_almost_equal(sum(d[0] for d in output_decays), 1) diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py index 3ec2492b..a7c12d3e 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fulldecay/test_mass_functions.py @@ -15,20 +15,16 @@ def ref_mass_func(min_mass, max_mass, n_events): """Reference mass function used to compare the behavior of the actual mass functions. - Parameters - ---------- - min_mass - max_mass - n_events - - Returns - ------- - kstar_mass - Mass generated - - Notes - ----- - Code taken from phasespace documentation. + Args: + min_mass: lower limit of mass. + max_mass: upper limit of mass. + n_events: number of mass values that should be generated. + + Returns: + kstar_mass: Generated mass. + + Notes: + Code taken from phasespace documentation. """ min_mass = tf.cast(min_mass, tf.float64) max_mass = tf.cast(max_mass, tf.float64) From 33ef57b88ac8afdf3363447ec4f377527802fdd5 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 16:05:48 +0200 Subject: [PATCH 16/71] Rename mass functions to use lowercase. --- phasespace/fulldecay/fulldecay.py | 2 +- phasespace/fulldecay/mass_functions.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 6ee9963b..4f26d4e0 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -10,7 +10,7 @@ from .mass_functions import _DEFAULT_CONVERTER _MASS_WIDTH_TOLERANCE = 0.01 -_DEFAULT_MASS_FUNC = "rel-BW" +_DEFAULT_MASS_FUNC = "relbw" class FullDecay: diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py index cca0ccb9..80738fb9 100644 --- a/phasespace/fulldecay/mass_functions.py +++ b/phasespace/fulldecay/mass_functions.py @@ -49,7 +49,7 @@ def mass_func(min_mass, max_mass, n_events): ) iterator = tf.stack([min_mass, max_mass], axis=-1) - # TODO this works with map_fn but not with vectorized_map as no analytic sampling is available. + # this works with map_fn but not with vectorized_map as no analytic sampling is available. return tf.map_fn( lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator ) @@ -59,6 +59,6 @@ def mass_func(min_mass, max_mass, n_events): _DEFAULT_CONVERTER = { "gauss": gauss, - "BW": breitwigner, - "rel-BW": relativistic_breitwigner, + "bw": breitwigner, + "relbw": relativistic_breitwigner, } From 96aea999fa3bc10a04b8e9e2a24030d294d2ac24 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 16:43:19 +0200 Subject: [PATCH 17/71] Move most decays into a decfile --- tests/fulldecay/example_decay_chains.py | 110 ++++-------------------- tests/fulldecay/example_decays.dec | 31 +++++++ 2 files changed, 48 insertions(+), 93 deletions(-) create mode 100644 tests/fulldecay/example_decays.dec diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index aee3c638..2c6167c0 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -1,96 +1,20 @@ -# A D+ particle with only one way of decaying -dplus_single = { - "D+": [ - { - "bf": 1, - "fs": [ - "K-", - "pi+", - "pi+", - {"pi0": [{"bf": 1, "fs": ["gamma", "gamma"]}]}, - ], - "model": "PHSP", - "model_params": "", - "zfit": "rel-BW", - } - ] -} +from decaylanguage import DecayMode, DecayChain, DecFileParser -pi0_4branches = { - "pi0": [ - {"bf": 0.988228297, "fs": ["gamma", "gamma"], "zfit": "BW"}, - {"bf": 0.011738247, "fs": ["e+", "e-", "gamma"], "zfit": "gauss"}, - {"bf": 3.3392e-5, "fs": ["e+", "e+", "e-", "e-"], "zfit": "rel-BW"}, - {"bf": 6.5e-8, "fs": ["e+", "e-"]}, - ] -} +dfp = DecFileParser("example_decays.dec") +dfp.parse() -dplus_4grandbranches = { - "D+": [ - { - "bf": 1.0, - "fs": ["K-", "pi+", "pi+", pi0_4branches], - "model": "PHSP", - "model_params": "", - } - ] -} +# 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') +dplus_single = DecayChain('D+', {'D+': dplus_decay, 'pi0': pi0_decay}).to_dict() -dstarplus_big_decay = { - "D*+": [ - { - "bf": 0.677, - "fs": [ - { - "D0": [ - { - "bf": 1.0, - "fs": ["K-", "pi+"], - "model": "PHSP", - "model_params": "", - } - ] - }, - "pi+", - ], - "model": "VSS", - "model_params": "", - }, - { - "bf": 0.307, - "fs": [ - { - "D+": [ - { - "bf": 1.0, - "fs": ["K-", "pi+", "pi+", pi0_4branches], - "model": "PHSP", - "model_params": "", - } - ] - }, - pi0_4branches, - ], - "model": "VSS", - "model_params": "", - }, - { - "bf": 0.016, - "fs": [ - { - "D+": [ - { - "bf": 1.0, - "fs": ["K-", "pi+", "pi+", pi0_4branches], - "model": "PHSP", - "model_params": "", - } - ] - }, - "gamma", - ], - "model": "VSP_PWAVE", - "model_params": "", - }, - ] -} + +# pi0 particle that can decay in 4 possible ways +pi0_4branches = dfp.build_decay_chains('pi0') +# TODO add zfit mass functions + +# 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+") + +# D*+ particle that has multiple child particles, grandchild particles, many of which can decay in multiple ways. +dstarplus_big_decay = dfp.build_decay_chains("D*+") diff --git a/tests/fulldecay/example_decays.dec b/tests/fulldecay/example_decays.dec new file mode 100644 index 00000000..c1eae76e --- /dev/null +++ b/tests/fulldecay/example_decays.dec @@ -0,0 +1,31 @@ +# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec +# Example decay chain for testing purposes +# Considered by itself, this file in in fact incomplete, +# as there are no instructions on how to decay the anti-D0 and the D-! + +Decay D*+ +0.6770 D0 pi+ VSS; +0.3070 D+ pi0 VSS; +0.0160 D+ gamma VSP_PWAVE; +Enddecay + +Decay D*- +0.6770 anti-D0 pi- VSS; +0.3070 D- pi0 VSS; +0.0160 D- gamma VSP_PWAVE; +Enddecay + +Decay D0 +1.0 K- pi+ PHSP; +Enddecay + +Decay D+ +1.0 K- pi+ pi+ pi0 PHSP; +Enddecay + +Decay pi0 +0.988228297 gamma gamma PHSP; +0.011738247 e+ e- gamma PI0_DALITZ; +0.000033392 e+ e+ e- e- PHSP; +0.000000065 e+ e- PHSP; +Enddecay From e610d11aaac3ae790efca0f236f1a990c0242d04 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 16:54:41 +0200 Subject: [PATCH 18/71] Add mass functions and make tests pass --- tests/fulldecay/example_decay_chains.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index 2c6167c0..333e7c50 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -1,6 +1,9 @@ from decaylanguage import DecayMode, DecayChain, DecFileParser +import os.path -dfp = DecFileParser("example_decays.dec") +script_dir = os.path.dirname(os.path.abspath(__file__)) + +dfp = DecFileParser(f"{script_dir}/example_decays.dec") dfp.parse() # D+ particle with only one way of decaying @@ -11,7 +14,10 @@ # pi0 particle that can decay in 4 possible ways pi0_4branches = dfp.build_decay_chains('pi0') -# TODO add zfit mass functions +# Specify different mass functions for the different decays of pi0 +mass_functions = ["relbw", "bw", "gauss"] +for mass_function, decay_mode in zip(mass_functions, pi0_4branches["pi0"]): + decay_mode["zfit"] = mass_function # 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+") From 77830d13de3d09e56448086a083ef7b4028d2140 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 16:55:14 +0200 Subject: [PATCH 19/71] Remove newline --- tests/fulldecay/example_decay_chains.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index 333e7c50..5dd8b2b6 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -11,7 +11,6 @@ pi0_decay = DecayMode(1, 'gamma gamma') 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') # Specify different mass functions for the different decays of pi0 From f5a23ec547317c884eecf263fdb6db32338ae03f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 3 Oct 2021 14:55:45 +0000 Subject: [PATCH 20/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/fulldecay/example_decay_chains.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index 5dd8b2b6..efe1a69a 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -1,18 +1,19 @@ -from decaylanguage import DecayMode, DecayChain, DecFileParser import os.path +from decaylanguage import DecayChain, DecayMode, DecFileParser + script_dir = os.path.dirname(os.path.abspath(__file__)) dfp = DecFileParser(f"{script_dir}/example_decays.dec") dfp.parse() # 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') -dplus_single = DecayChain('D+', {'D+': dplus_decay, 'pi0': pi0_decay}).to_dict() +dplus_decay = DecayMode(1, "K- pi+ pi+ pi0", model="PHSP", zfit="relbw") +pi0_decay = DecayMode(1, "gamma gamma") +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') +pi0_4branches = dfp.build_decay_chains("pi0") # Specify different mass functions for the different decays of pi0 mass_functions = ["relbw", "bw", "gauss"] for mass_function, decay_mode in zip(mass_functions, pi0_4branches["pi0"]): From 8e31148f2ed8599b3399cf17f50cdc85d12ba184 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 23:08:59 +0200 Subject: [PATCH 21/71] Add test for mass_converter --- tests/fulldecay/test_fulldecay.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index a26142a9..d0f11c1c 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -1,6 +1,7 @@ from numpy.testing import assert_almost_equal from phasespace.fulldecay import FullDecay +from phasespace.fulldecay.mass_functions import _DEFAULT_CONVERTER from .example_decay_chains import * # TODO remove * since it is bad practice? @@ -63,6 +64,22 @@ def test_branching_children(): (normed_weights, events), _ = check_norm(container, n_events=100) +def test_mass_converter(): + """Test that the mass_converter parameter works as intended""" + pi0_4branches_copy = pi0_4branches.copy() + pi0_4branches_copy[-1]["zfit"] = "rel-BW" + container = FullDecay.from_dict(pi0_4branches, tolerance=1e-10, + mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}) + + output_decays = container.gen_particles + assert len(output_decays) == 4 + assert_almost_equal(sum(d[0] for d in output_decays), 1) + assert all(not decay.has_fixed_mass() for decay in output_decays) + + check_norm(container, n_events=1) + (normed_weights, events), _ = check_norm(container, n_events=100) + + def test_branching_grandchilden(): """Test converting a decaylanguage dict where children to the mother particle can decay in many ways.""" container = FullDecay.from_dict(dplus_4grandbranches) From 090e65b6d0a5f9375e57bbac38c80388cf8e7ec2 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 3 Oct 2021 23:39:45 +0200 Subject: [PATCH 22/71] Fix test with mass_converter --- tests/fulldecay/example_decay_chains.py | 9 ++++---- tests/fulldecay/test_fulldecay.py | 30 ++++++++++++++----------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index efe1a69a..2ed94c63 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -14,13 +14,14 @@ # pi0 particle that can decay in 4 possible ways pi0_4branches = dfp.build_decay_chains("pi0") -# Specify different mass functions for the different decays of pi0 -mass_functions = ["relbw", "bw", "gauss"] -for mass_function, decay_mode in zip(mass_functions, pi0_4branches["pi0"]): - decay_mode["zfit"] = mass_function # 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 child particles, grandchild particles, many of which can decay in multiple ways. dstarplus_big_decay = dfp.build_decay_chains("D*+") diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index d0f11c1c..e309965f 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -64,30 +64,34 @@ def test_branching_children(): (normed_weights, events), _ = check_norm(container, n_events=100) -def test_mass_converter(): - """Test that the mass_converter parameter works as intended""" - pi0_4branches_copy = pi0_4branches.copy() - pi0_4branches_copy[-1]["zfit"] = "rel-BW" - container = FullDecay.from_dict(pi0_4branches, tolerance=1e-10, - mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}) - +def test_branching_grandchilden(): + """Test converting a decaylanguage dict where children to the mother particle can decay in many ways.""" + container = FullDecay.from_dict(dplus_4grandbranches) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) - assert all(not decay.has_fixed_mass() for decay in output_decays) - check_norm(container, n_events=1) (normed_weights, events), _ = check_norm(container, n_events=100) + # TODO add more asserts here -def test_branching_grandchilden(): - """Test converting a decaylanguage dict where children to the mother particle can decay in many ways.""" - container = FullDecay.from_dict(dplus_4grandbranches) +def test_mass_converter(): + """Test that the mass_converter parameter works as intended""" + dplus_4grandbranches_massfunc = dplus_4grandbranches.copy() + dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" + container = FullDecay.from_dict(dplus_4grandbranches_massfunc, tolerance=1e-10, + mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}) output_decays = container.gen_particles + assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) + + for decay in output_decays: + for child in decay[1].children: + if "pi0" in child.name: + assert not child.has_fixed_mass + check_norm(container, n_events=1) (normed_weights, events), _ = check_norm(container, n_events=100) - # TODO add more asserts here def test_big_decay(): From 3879e0415cddaaf098a2dcea1e80676fdba8ad5c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 3 Oct 2021 21:40:12 +0000 Subject: [PATCH 23/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/fulldecay/example_decay_chains.py | 4 +++- tests/fulldecay/test_fulldecay.py | 9 ++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py index 2ed94c63..ef71fda6 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -20,7 +20,9 @@ # 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"]): +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 child particles, grandchild particles, many of which can decay in multiple ways. diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py index e309965f..4baa7568 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -76,11 +76,14 @@ def test_branching_grandchilden(): def test_mass_converter(): - """Test that the mass_converter parameter works as intended""" + """Test that the mass_converter parameter works as intended.""" dplus_4grandbranches_massfunc = dplus_4grandbranches.copy() dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" - container = FullDecay.from_dict(dplus_4grandbranches_massfunc, tolerance=1e-10, - mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}) + container = FullDecay.from_dict( + dplus_4grandbranches_massfunc, + tolerance=1e-10, + mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}, + ) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) From 51e5ecb373acf1743876c834b8ba82a5648ac570 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Thu, 7 Oct 2021 11:38:40 +0200 Subject: [PATCH 24/71] Update default value --- phasespace/fulldecay/fulldecay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py index 4f26d4e0..76e755af 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -58,7 +58,7 @@ def from_dict( return cls(gen_particles) def generate( - self, n_events: int, normalize_weights: bool = False, **kwargs + self, n_events: int, normalize_weights: bool = True, **kwargs ) -> Union[ tuple[list[tf.Tensor], list[tf.Tensor]], tuple[list[tf.Tensor], list[tf.Tensor], list[tf.Tensor]], From bf6f19a90facc76761899dd708b3387b405482bb Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Thu, 7 Oct 2021 11:55:36 +0200 Subject: [PATCH 25/71] Begin working on documentation notebook --- docs/fromdecay.ipynb | 323 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 323 insertions(+) create mode 100644 docs/fromdecay.ipynb diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb new file mode 100644 index 00000000..b2788912 --- /dev/null +++ b/docs/fromdecay.ipynb @@ -0,0 +1,323 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": true, + "pycharm": { + "is_executing": true, + "name": "#%% md\n" + } + }, + "source": [ + "# Tutorial for `fromdecay` functionality\n", + "This tutorial shows how `phasespace.fromdecay` can be used.\n", + "\n", + "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", + "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Import libraries\n", + "from pprint import pprint\n", + "from decaylanguage import DecFileParser, DecayChainViewer\n", + "# TODO rename\n", + "from phasespace.fulldecay import FullDecay" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", + "\n", + "We will begin by parsing a .dec file using DecayLanguage:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "parser = DecFileParser('../tests/fulldecay/example_decays.dec')\n", + "parser.parse()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'pi0': [{'bf': 0.988228297,\n", + " 'fs': ['gamma', 'gamma'],\n", + " 'model': 'PHSP',\n", + " 'model_params': ''},\n", + " {'bf': 0.011738247,\n", + " 'fs': ['e+', 'e-', 'gamma'],\n", + " 'model': 'PI0_DALITZ',\n", + " 'model_params': ''},\n", + " {'bf': 3.3392e-05,\n", + " 'fs': ['e+', 'e+', 'e-', 'e-'],\n", + " 'model': 'PHSP',\n", + " 'model_params': ''},\n", + " {'bf': 6.5e-08,\n", + " 'fs': ['e+', 'e-'],\n", + " 'model': 'PHSP',\n", + " 'model_params': ''}]}\n" + ] + } + ], + "source": [ + "pi0_chain = parser.build_decay_chains(\"pi0\")\n", + "pprint(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "DecayChainGraph\n", + "\n", + "\n", + "\n", + "mother\n", + "\n", + "\n", + "π\n", + "0\n", + "\n", + "\n", + "\n", + "dec0\n", + "\n", + "\n", + "γ\n", + "γ\n", + "\n", + "\n", + "\n", + "mother->dec0\n", + "\n", + "\n", + "0.988228297\n", + "\n", + "\n", + "\n", + "dec1\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "-\n", + "γ\n", + "\n", + "\n", + "\n", + "mother->dec1\n", + "\n", + "\n", + "0.011738247\n", + "\n", + "\n", + "\n", + "dec2\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "+\n", + "e\n", + "-\n", + "e\n", + "-\n", + "\n", + "\n", + "\n", + "mother->dec2\n", + "\n", + "\n", + "3.3392e-05\n", + "\n", + "\n", + "\n", + "dec3\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "-\n", + "\n", + "\n", + "\n", + "mother->dec3\n", + "\n", + "\n", + "6.5e-08\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "DecayChainViewer(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_decay = FullDecay.from_dict(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One can then simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [], + "source": [ + "weights, events = pi0_decay.generate(n_events=10_000, normalize_weights=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "There is a probability of 0.988228297 that pi0 decays into gamma, gamma [0]\n", + "There is a probability of 0.011738247 that pi0 decays into e+, e-, gamma [1]\n", + "There is a probability of 3.3392e-05 that pi0 decays into e+ [0], e+ [1], e- [0], e- [1]\n", + "There is a probability of 6.5e-08 that pi0 decays into e+ [2], e- [2]\n" + ] + } + ], + "source": [ + "for probability, particle in gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When calling the `FullDecay.generate` method, it internally calls the generate method on all of these GenParticle instances. These are then placed in a list, which is returned. This is why the returned weights and events from `FullDecay.generate` are lists." + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [], + "source": [ + "# TODO add more about mass functions, the zfit parameter, and tolerance" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} From 4d9b009dec0c649e99ac6bb81e1bcda7c1c91130 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Thu, 7 Oct 2021 22:42:14 +0200 Subject: [PATCH 26/71] Add docs about tolerance and improve former parts --- docs/fromdecay.ipynb | 311 ++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 290 insertions(+), 21 deletions(-) diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index b2788912..335dee2e 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -19,7 +19,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 18, "metadata": { "pycharm": { "name": "#%%\n" @@ -29,6 +29,7 @@ "source": [ "# Import libraries\n", "from pprint import pprint\n", + "from particle import Particle\n", "from decaylanguage import DecFileParser, DecayChainViewer\n", "# TODO rename\n", "from phasespace.fulldecay import FullDecay" @@ -38,6 +39,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ + "## Quick Intro to DecayLanguage\n", "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", "\n", "We will begin by parsing a .dec file using DecayLanguage:" @@ -66,7 +68,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 4, "metadata": {}, "outputs": [ { @@ -207,7 +209,7 @@ "\n" ], "text/plain": [ - "" + "" ] }, "execution_count": 5, @@ -223,12 +225,13 @@ "cell_type": "markdown", "metadata": {}, "source": [ + "## Creating a FullDecay object\n", "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" ] }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 7, "metadata": {}, "outputs": [], "source": [ @@ -239,63 +242,329 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "One can then simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method." + "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." ] }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 11, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "There is a probability of 0.988228297 that pi0 decays into gamma, gamma [0]\n", + "There is a probability of 0.011738247 that pi0 decays into e+, e-, gamma [1]\n", + "There is a probability of 3.3392e-05 that pi0 decays into e+ [0], e+ [1], e- [0], e- [1]\n", + "There is a probability of 6.5e-08 that pi0 decays into e+ [2], e- [2]\n" + ] + } + ], "source": [ - "weights, events = pi0_decay.generate(n_events=10_000, normalize_weights=True)" + "for probability, particle in pi0_decay.gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", "\n", - "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." ] }, { "cell_type": "code", - "execution_count": 30, + "execution_count": 16, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "There is a probability of 0.988228297 that pi0 decays into gamma, gamma [0]\n", - "There is a probability of 0.011738247 that pi0 decays into e+, e-, gamma [1]\n", - "There is a probability of 3.3392e-05 that pi0 decays into e+ [0], e+ [1], e- [0], e- [1]\n", - "There is a probability of 6.5e-08 that pi0 decays into e+ [2], e- [2]\n" + "Number of events for each decay mode: 9888, 112\n" ] } ], "source": [ - "for probability, particle in gen_particles:\n", - " print(f\"There is a probability of {probability} \"\n", - " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + "weights, events = pi0_decay.generate(n_events=10_000)\n", + "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can confirm that the counts above are close to the expected counts based on the probabilities. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Changing mass settings FullDecay\n", + "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", + "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", + "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", + "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", + "\n", + "### tolerance\n", + "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", + "This will be demonsttrated with the decay below:" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "DecayChainGraph\n", + "\n", + "\n", + "\n", + "mother\n", + "\n", + "\n", + "D\n", + "*\n", + "(2010)\n", + "+\n", + "\n", + "\n", + "\n", + "dec45\n", + "\n", + "\n", + "D\n", + "0\n", + "\n", + "π\n", + "+\n", + "\n", + "\n", + "\n", + "mother->dec45\n", + "\n", + "\n", + "0.677\n", + "\n", + "\n", + "\n", + "dec47\n", + "\n", + "\n", + "D\n", + "+\n", + "\n", + "π\n", + "0\n", + "\n", + "\n", + "\n", + "mother->dec47\n", + "\n", + "\n", + "0.307\n", + "\n", + "\n", + "\n", + "dec52\n", + "\n", + "\n", + "D\n", + "+\n", + "γ\n", + "\n", + "\n", + "\n", + "mother->dec52\n", + "\n", + "\n", + "0.016\n", + "\n", + "\n", + "\n", + "dec46\n", + "\n", + "\n", + "K\n", + "-\n", + "π\n", + "+\n", + "\n", + "\n", + "\n", + "dec45:p0->dec46\n", + "\n", + "\n", + "1.0\n", + "\n", + "\n", + "\n", + "dec48\n", + "\n", + "\n", + "γ\n", + "γ\n", + "\n", + "\n", + "\n", + "dec47:p1->dec48\n", + "\n", + "\n", + "0.988228297\n", + "\n", + "\n", + "\n", + "dec49\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "-\n", + "γ\n", + "\n", + "\n", + "\n", + "dec47:p1->dec49\n", + "\n", + "\n", + "0.011738247\n", + "\n", + "\n", + "\n", + "dec50\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "+\n", + "e\n", + "-\n", + "e\n", + "-\n", + "\n", + "\n", + "\n", + "dec47:p1->dec50\n", + "\n", + "\n", + "3.3392e-05\n", + "\n", + "\n", + "\n", + "dec51\n", + "\n", + "\n", + "e\n", + "+\n", + "e\n", + "-\n", + "\n", + "\n", + "\n", + "dec47:p1->dec51\n", + "\n", + "\n", + "6.5e-08\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 26, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", + "DecayChainViewer(dsplus_chain)" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pi0 width = 7.81e-06\n", + "D0 width = 1.605e-09\n" + ] + } + ], + "source": [ + "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", + " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$\\pi^0$ has a greater width than $D^0$. \n", + "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": {}, + "outputs": [], + "source": [ + "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", + "# Loop over D0 and pi+ particles, see graph above\n", + "for particle in dstar_decay.gen_particles[0][1].children:\n", + " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", + " assert particle.has_fixed_mass\n", + " \n", + "# Loop over D+ and pi0. See above.\n", + "for particle in dstar_decay.gen_particles[1][1].children:\n", + " if particle.name == \"pi0\":\n", + " assert not particle.has_fixed_mass" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "When calling the `FullDecay.generate` method, it internally calls the generate method on all of these GenParticle instances. These are then placed in a list, which is returned. This is why the returned weights and events from `FullDecay.generate` are lists." + "### mass_converter\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed." ] }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 45, "metadata": {}, "outputs": [], "source": [ - "# TODO add more about mass functions, the zfit parameter, and tolerance" + "# TODO add more about mass functions, the zfit parameter" ] } ], From be38a78cc4ba6f235b0c10c524e0cbd34a3c8aee Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Fri, 8 Oct 2021 11:25:47 +0200 Subject: [PATCH 27/71] chore: add pre-commit hook to clean jupyter notebook --- .pre-commit-config.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 672d082a..63566003 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -84,3 +84,8 @@ repos: rev: v2.3.2 hooks: - id: prettier + + - repo: https://github.com/roy-ht/pre-commit-jupyter + rev: v1.2.1 + hooks: + - id: jupyter-notebook-cleanup From 79142cc3736e18361e94f4beaba885a976fa9955 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 Oct 2021 09:26:03 +0000 Subject: [PATCH 28/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .pre-commit-config.yaml | 2 +- docs/fromdecay.ipynb | 368 ++-------------------------------------- 2 files changed, 18 insertions(+), 352 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 63566003..cdd065df 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -84,7 +84,7 @@ repos: rev: v2.3.2 hooks: - id: prettier - + - repo: https://github.com/roy-ht/pre-commit-jupyter rev: v1.2.1 hooks: diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index 335dee2e..c51f94b8 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -19,7 +19,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": { "pycharm": { "name": "#%%\n" @@ -47,7 +47,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": { "pycharm": { "name": "#%%\n" @@ -68,32 +68,9 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'pi0': [{'bf': 0.988228297,\n", - " 'fs': ['gamma', 'gamma'],\n", - " 'model': 'PHSP',\n", - " 'model_params': ''},\n", - " {'bf': 0.011738247,\n", - " 'fs': ['e+', 'e-', 'gamma'],\n", - " 'model': 'PI0_DALITZ',\n", - " 'model_params': ''},\n", - " {'bf': 3.3392e-05,\n", - " 'fs': ['e+', 'e+', 'e-', 'e-'],\n", - " 'model': 'PHSP',\n", - " 'model_params': ''},\n", - " {'bf': 6.5e-08,\n", - " 'fs': ['e+', 'e-'],\n", - " 'model': 'PHSP',\n", - " 'model_params': ''}]}\n" - ] - } - ], + "outputs": [], "source": [ "pi0_chain = parser.build_decay_chains(\"pi0\")\n", "pprint(pi0_chain)" @@ -108,115 +85,9 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "DecayChainGraph\n", - "\n", - "\n", - "\n", - "mother\n", - "\n", - "\n", - "π\n", - "0\n", - "\n", - "\n", - "\n", - "dec0\n", - "\n", - "\n", - "γ\n", - "γ\n", - "\n", - "\n", - "\n", - "mother->dec0\n", - "\n", - "\n", - "0.988228297\n", - "\n", - "\n", - "\n", - "dec1\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "-\n", - "γ\n", - "\n", - "\n", - "\n", - "mother->dec1\n", - "\n", - "\n", - "0.011738247\n", - "\n", - "\n", - "\n", - "dec2\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "+\n", - "e\n", - "-\n", - "e\n", - "-\n", - "\n", - "\n", - "\n", - "mother->dec2\n", - "\n", - "\n", - "3.3392e-05\n", - "\n", - "\n", - "\n", - "dec3\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "-\n", - "\n", - "\n", - "\n", - "mother->dec3\n", - "\n", - "\n", - "6.5e-08\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "DecayChainViewer(pi0_chain)" ] @@ -231,7 +102,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -249,20 +120,9 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "There is a probability of 0.988228297 that pi0 decays into gamma, gamma [0]\n", - "There is a probability of 0.011738247 that pi0 decays into e+, e-, gamma [1]\n", - "There is a probability of 3.3392e-05 that pi0 decays into e+ [0], e+ [1], e- [0], e- [1]\n", - "There is a probability of 6.5e-08 that pi0 decays into e+ [2], e- [2]\n" - ] - } - ], + "outputs": [], "source": [ "for probability, particle in pi0_decay.gen_particles:\n", " print(f\"There is a probability of {probability} \"\n", @@ -280,17 +140,9 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Number of events for each decay mode: 9888, 112\n" - ] - } - ], + "outputs": [], "source": [ "weights, events = pi0_decay.generate(n_events=10_000)\n", "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" @@ -320,186 +172,9 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "DecayChainGraph\n", - "\n", - "\n", - "\n", - "mother\n", - "\n", - "\n", - "D\n", - "*\n", - "(2010)\n", - "+\n", - "\n", - "\n", - "\n", - "dec45\n", - "\n", - "\n", - "D\n", - "0\n", - "\n", - "π\n", - "+\n", - "\n", - "\n", - "\n", - "mother->dec45\n", - "\n", - "\n", - "0.677\n", - "\n", - "\n", - "\n", - "dec47\n", - "\n", - "\n", - "D\n", - "+\n", - "\n", - "π\n", - "0\n", - "\n", - "\n", - "\n", - "mother->dec47\n", - "\n", - "\n", - "0.307\n", - "\n", - "\n", - "\n", - "dec52\n", - "\n", - "\n", - "D\n", - "+\n", - "γ\n", - "\n", - "\n", - "\n", - "mother->dec52\n", - "\n", - "\n", - "0.016\n", - "\n", - "\n", - "\n", - "dec46\n", - "\n", - "\n", - "K\n", - "-\n", - "π\n", - "+\n", - "\n", - "\n", - "\n", - "dec45:p0->dec46\n", - "\n", - "\n", - "1.0\n", - "\n", - "\n", - "\n", - "dec48\n", - "\n", - "\n", - "γ\n", - "γ\n", - "\n", - "\n", - "\n", - "dec47:p1->dec48\n", - "\n", - "\n", - "0.988228297\n", - "\n", - "\n", - "\n", - "dec49\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "-\n", - "γ\n", - "\n", - "\n", - "\n", - "dec47:p1->dec49\n", - "\n", - "\n", - "0.011738247\n", - "\n", - "\n", - "\n", - "dec50\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "+\n", - "e\n", - "-\n", - "e\n", - "-\n", - "\n", - "\n", - "\n", - "dec47:p1->dec50\n", - "\n", - "\n", - "3.3392e-05\n", - "\n", - "\n", - "\n", - "dec51\n", - "\n", - "\n", - "e\n", - "+\n", - "e\n", - "-\n", - "\n", - "\n", - "\n", - "dec47:p1->dec51\n", - "\n", - "\n", - "6.5e-08\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 26, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", "DecayChainViewer(dsplus_chain)" @@ -507,18 +182,9 @@ }, { "cell_type": "code", - "execution_count": 39, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "pi0 width = 7.81e-06\n", - "D0 width = 1.605e-09\n" - ] - } - ], + "outputs": [], "source": [ "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" @@ -534,7 +200,7 @@ }, { "cell_type": "code", - "execution_count": 44, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -560,7 +226,7 @@ }, { "cell_type": "code", - "execution_count": 45, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ From d76a96704c778022760f1f3d0cfd3b11a9b1b986 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 10 Oct 2021 16:59:40 +0200 Subject: [PATCH 29/71] Write about all features in fulldecay --- docs/fromdecay.ipynb | 108 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 104 insertions(+), 4 deletions(-) diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index c51f94b8..ee248d7d 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -29,6 +29,8 @@ "source": [ "# Import libraries\n", "from pprint import pprint\n", + "\n", + "import zfit\n", "from particle import Particle\n", "from decaylanguage import DecFileParser, DecayChainViewer\n", "# TODO rename\n", @@ -165,7 +167,7 @@ "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", "\n", - "### tolerance\n", + "### Constant vs variable mass\n", "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", "This will be demonsttrated with the decay below:" ] @@ -220,8 +222,82 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### mass_converter\n", - "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed." + "### Configuring mass fucntions\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_custom_mass_func = dsplus_chain.copy()\n", + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def custom_gauss(mass, width):\n", + " particle_mass = tf.cast(mass, tf.float64)\n", + " particle_width = tf.cast(width, tf.float64)\n", + " \n", + " # This is the actual mass function that will be returned\n", + " def mass_func(min_mass, max_mass, n_events):\n", + " min_mass = tf.cast(min_mass, tf.float64)\n", + " max_mass = tf.cast(max_mass, tf.float64)\n", + " # Use a zfit PDF\n", + " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", + " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", + " return tf.vectorized_map(\n", + " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", + " )\n", + "\n", + " return mass_func" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." ] }, { @@ -230,8 +306,32 @@ "metadata": {}, "outputs": [], "source": [ - "# TODO add more about mass functions, the zfit parameter" + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "\n", + "# Set the mass function of pi0 to the custom gaussian distribution \n", + "# when it decays into an electron-positron pair and a photon (gamma)\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { From 02549c5f31750d16c43768ffcb8b2f4a0a8e0ee8 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 10 Oct 2021 17:35:52 +0200 Subject: [PATCH 30/71] Rename fuldecay to fromdecay. This renaming convention does however not include the FullDecay class, since FromDecay seems like a strange class name. This might change. --- .git_archival.txt | 2 +- .gitattributes | 10 +- .gitignore | 224 +-- AUTHORS.rst | 36 +- CHANGELOG.rst | 248 +-- CONTRIBUTING.rst | 250 +-- LICENSE | 48 +- MANIFEST.in | 28 +- README.rst | 466 ++--- benchmark/bench_phasespace.py | 280 +-- benchmark/bench_tgenphasespace.cxx | 42 +- benchmark/monitoring.py | 152 +- data/download_test_files.py | 98 +- docs/Makefile | 40 +- docs/authors.rst | 2 +- docs/conf.py | 510 +++--- docs/contributing.rst | 2 +- docs/history.rst | 2 +- docs/index.rst | 34 +- docs/make.bat | 72 +- docs/phasespace.rst | 36 +- docs/usage.rst | 406 ++--- paper/paper.bib | 230 +-- phasespace/__init__.py | 52 +- phasespace/backend.py | 42 +- phasespace/fromdecay/__init__.py | 14 + .../{fulldecay => fromdecay}/fulldecay.py | 23 +- .../mass_functions.py | 0 phasespace/fulldecay/__init__.py | 14 - phasespace/kinematics.py | 304 ++-- phasespace/phasespace.py | 1552 ++++++++--------- phasespace/random.py | 82 +- pyproject.toml | 26 +- scripts/prepare_test_samples.cxx | 246 +-- setup.cfg | 190 +- setup.py | 40 +- tests/conftest.py | 8 +- tests/{fulldecay => fromdecay}/__init__.py | 8 +- .../example_decay_chains.py | 2 +- .../example_decays.dec | 62 +- .../test_fulldecay.py | 16 +- .../test_mass_functions.py | 2 +- tests/helpers/decays.py | 160 +- tests/helpers/plotting.py | 56 +- tests/helpers/rapidsim.py | 222 +-- tests/test_chain.py | 274 +-- tests/test_generate.py | 172 +- tests/test_nbody_decay.py | 128 +- tests/test_physics.py | 822 ++++----- tests/test_random.py | 54 +- 50 files changed, 3897 insertions(+), 3892 deletions(-) create mode 100644 phasespace/fromdecay/__init__.py rename phasespace/{fulldecay => fromdecay}/fulldecay.py (90%) rename phasespace/{fulldecay => fromdecay}/mass_functions.py (100%) delete mode 100644 phasespace/fulldecay/__init__.py rename tests/{fulldecay => fromdecay}/__init__.py (63%) rename tests/{fulldecay => fromdecay}/example_decay_chains.py (87%) rename tests/{fulldecay => fromdecay}/example_decays.dec (96%) rename tests/{fulldecay => fromdecay}/test_fulldecay.py (84%) rename tests/{fulldecay => fromdecay}/test_mass_functions.py (94%) diff --git a/.git_archival.txt b/.git_archival.txt index 95cb3eea..2e6e85ba 100644 --- a/.git_archival.txt +++ b/.git_archival.txt @@ -1 +1 @@ -ref-names: $Format:%D$ +ref-names: $Format:%D$ diff --git a/.gitattributes b/.gitattributes index 7f892985..4663edcf 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,5 +1,5 @@ -data/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2K1Gamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2KstGamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text -.git_archival.txt export-subst +data/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2K1Gamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2KstGamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text +.git_archival.txt export-subst diff --git a/.gitignore b/.gitignore index 363a46df..47e05415 100644 --- a/.gitignore +++ b/.gitignore @@ -1,112 +1,112 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -env/ -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# pyenv -.python-version - -# celery beat schedule file -celerybeat-schedule - -# SageMath parsed files -*.sage.py - -# dotenv -.env - -# virtualenv -.venv -venv/ -ENV/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy and pyright -.mypy_cache/ -/typings/ -pyrightconfig.json - -# specific -*/*_cxx* -tests/plots -/data/backup/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root -/data/backup/B2K1Gamma_RapidSim_7TeV_Tree.root -/data/backup/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root -/data/backup/B2KstGamma_RapidSim_7TeV_Tree.root +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# dotenv +.env + +# virtualenv +.venv +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy and pyright +.mypy_cache/ +/typings/ +pyrightconfig.json + +# specific +*/*_cxx* +tests/plots +/data/backup/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root +/data/backup/B2K1Gamma_RapidSim_7TeV_Tree.root +/data/backup/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root +/data/backup/B2KstGamma_RapidSim_7TeV_Tree.root diff --git a/AUTHORS.rst b/AUTHORS.rst index 11a8d85e..e98f898f 100644 --- a/AUTHORS.rst +++ b/AUTHORS.rst @@ -1,18 +1,18 @@ -========== -Credits -========== - -Development Lead ----------------- - -* Albert Puig Navarro - -Core Developers ---------------- - -* Jonas Eschle - -Contributors ------------- - -None yet. Why not be the first? +========== +Credits +========== + +Development Lead +---------------- + +* Albert Puig Navarro + +Core Developers +--------------- + +* Jonas Eschle + +Contributors +------------ + +None yet. Why not be the first? diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3dbc8e5c..f0339867 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,124 +1,124 @@ -********* -Changelog -********* - -Develop -========== - - -Major Features and Improvements -------------------------------- - -Behavioral changes ------------------- - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - - -Thanks ------- - -1.4.1 (27.08.2021) -================== - -Requirement changes -------------------- -- Losen restriction on TensorFlow, allow version 2.6 (and 2.5) - -1.4.0 (11.06.2021) -================== - -Requirement changes -------------------- -- require TensorFlow 2.5 as 2.4 breaks some functionality - -1.3.0 (28.05.2021) -=================== - - -Major Features and Improvements -------------------------------- - -- Support Python 3.9 -- Support TensorFlow 2.5 -- improved compilation in tf.functions, use of XLA where applicable -- developer: modernization of setup, CI and more - -Thanks ------- - -- Remco de Boer for many commits and cleanups - -1.2.0 (17.12.20) -================ - - -Major Features and Improvements -------------------------------- - -- Python 3.8 support -- Allow eager execution by setting with `tf.config.run_functions_eagerly(True)` - or the environment variable "PHASESPACE_EAGER" -- Deterministic random number generation via seed - or `tf.random.Generator` instance - -Behavioral changes ------------------- - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - -- tighten TensorFlow to 2.3/2.4 -- tighten TensorFlow Probability to 0.11/0.12 - -Thanks ------- -- Remco de Boer and Stefan Pflüger for discussions on random number genration - -1.1.0 (27.1.2020) -================= - -This release switched to TensorFlow 2.0 eager mode. Please upgrade your TensorFlow installation if possible and change -your code (minimal changes) as described under "Behavioral changes". -In case this is currently impossible to do, please downgrade to < 1.1.0. - -Major Features and Improvements -------------------------------- - - full TF2 compatibility - -Behavioral changes ------------------- - - `generate` now returns an eager Tensor. This is basically a numpy array wrapped by TensorFlow. - To explicitly convert it to a numpy array, use the `numpy()` method of the eager Tensor. - - `generate_tensor` is now depreceated, `generate` can directly be used instead. - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - - requires now TensorFlow >= 2.0.0 - - -Thanks ------- - - -1.0.4 (13-10-2019) -========================== - - -Major Features and Improvements -------------------------------- - -Release to conda-forge, thanks to Chris Burr +********* +Changelog +********* + +Develop +========== + + +Major Features and Improvements +------------------------------- + +Behavioral changes +------------------ + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + + +Thanks +------ + +1.4.1 (27.08.2021) +================== + +Requirement changes +------------------- +- Losen restriction on TensorFlow, allow version 2.6 (and 2.5) + +1.4.0 (11.06.2021) +================== + +Requirement changes +------------------- +- require TensorFlow 2.5 as 2.4 breaks some functionality + +1.3.0 (28.05.2021) +=================== + + +Major Features and Improvements +------------------------------- + +- Support Python 3.9 +- Support TensorFlow 2.5 +- improved compilation in tf.functions, use of XLA where applicable +- developer: modernization of setup, CI and more + +Thanks +------ + +- Remco de Boer for many commits and cleanups + +1.2.0 (17.12.20) +================ + + +Major Features and Improvements +------------------------------- + +- Python 3.8 support +- Allow eager execution by setting with `tf.config.run_functions_eagerly(True)` + or the environment variable "PHASESPACE_EAGER" +- Deterministic random number generation via seed + or `tf.random.Generator` instance + +Behavioral changes +------------------ + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + +- tighten TensorFlow to 2.3/2.4 +- tighten TensorFlow Probability to 0.11/0.12 + +Thanks +------ +- Remco de Boer and Stefan Pflüger for discussions on random number genration + +1.1.0 (27.1.2020) +================= + +This release switched to TensorFlow 2.0 eager mode. Please upgrade your TensorFlow installation if possible and change +your code (minimal changes) as described under "Behavioral changes". +In case this is currently impossible to do, please downgrade to < 1.1.0. + +Major Features and Improvements +------------------------------- + - full TF2 compatibility + +Behavioral changes +------------------ + - `generate` now returns an eager Tensor. This is basically a numpy array wrapped by TensorFlow. + To explicitly convert it to a numpy array, use the `numpy()` method of the eager Tensor. + - `generate_tensor` is now depreceated, `generate` can directly be used instead. + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + - requires now TensorFlow >= 2.0.0 + + +Thanks +------ + + +1.0.4 (13-10-2019) +========================== + + +Major Features and Improvements +------------------------------- + +Release to conda-forge, thanks to Chris Burr diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 88d4d412..88e9e2ca 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -1,125 +1,125 @@ -.. highlight:: shell - -============ -Contributing -============ - -Contributions are welcome, and they are greatly appreciated! Every little bit -helps, and credit will always be given. - -You can contribute in many ways: - -Types of Contributions ----------------------- - -Report Bugs -~~~~~~~~~~~ - -Report bugs at https://github.com/zfit/phasespace/issues. - -If you are reporting a bug, please include: - -* Your operating system name and version. -* Any details about your local setup that might be helpful in troubleshooting. -* Detailed steps to reproduce the bug. - -Fix Bugs -~~~~~~~~ - -Look through the GitHub issues for bugs. Anything tagged with "bug" and "help -wanted" is open to whoever wants to implement it. - -Implement Features -~~~~~~~~~~~~~~~~~~ - -Look through the GitHub issues for features. Anything tagged with "enhancement" -and "help wanted" is open to whoever wants to implement it. - -Write Documentation -~~~~~~~~~~~~~~~~~~~ - -TensorFlow PhaseSpace could always use more documentation, whether as part of the -official TensorFlow PhaseSpace docs, in docstrings, or even on the web in blog posts, -articles, and such. - -Submit Feedback -~~~~~~~~~~~~~~~ - -The best way to send feedback is to file an issue at https://github.com/zfit/phasespace/issues. - -If you are proposing a feature: - -* Explain in detail how it would work. -* Keep the scope as narrow as possible, to make it easier to implement. -* Remember that this is a volunteer-driven project, and that contributions - are welcome :) - -Get Started! ------------- - -Ready to contribute? Here's how to set up `phasespace` for local development. - -1. Fork the `phasespace` repo on GitHub. -2. Clone your fork locally:: - - $ git clone git@github.com:your_name_here/phasespace.git - -3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: - - $ mkvirtualenv phasespace - $ cd phasespace/ - $ pip install -e .[dev] - -4. Create a branch for local development:: - - $ git checkout -b name-of-your-bugfix-or-feature - - Now you can make your changes locally. - -5. When you're done making changes, check that your changes pass pre-commit and the - tests: - - $ pytest - $ pre-commit run -a - -6. Commit your changes and push your branch to GitHub:: - - $ git add . - $ git commit -m "Your detailed description of your changes." - $ git push origin name-of-your-bugfix-or-feature - -7. Submit a pull request through the GitHub website. - -Pull Request Guidelines ------------------------ - -Before you submit a pull request, check that it meets these guidelines: - -1. The pull request should include tests. -2. If the pull request adds functionality, the docs should be updated. Put - your new functionality into a function with a docstring, and add the - feature to the list in README.rst. -3. The pull request should work for Python 2.7, 3.4, 3.5 and 3.6, and for PyPy. Check - https://github.com/zfit/phasespace/actions/workflows/ci.yml - and make sure that the tests pass for all supported Python versions. - -Tips ----- - -To run a subset of tests (for example those in `tests/test_generate.py`):: - - - $ pytest -k test_generate - -Deploying ---------- - -A reminder for the maintainers on how to deploy. -Make sure all your changes are committed (including an entry in HISTORY.rst). -Then run:: - -$ bumpversion patch # possible: major / minor / patch -$ git push -$ git push --tags - -GitHub Actions will then deploy to PyPI if tests pass. +.. highlight:: shell + +============ +Contributing +============ + +Contributions are welcome, and they are greatly appreciated! Every little bit +helps, and credit will always be given. + +You can contribute in many ways: + +Types of Contributions +---------------------- + +Report Bugs +~~~~~~~~~~~ + +Report bugs at https://github.com/zfit/phasespace/issues. + +If you are reporting a bug, please include: + +* Your operating system name and version. +* Any details about your local setup that might be helpful in troubleshooting. +* Detailed steps to reproduce the bug. + +Fix Bugs +~~~~~~~~ + +Look through the GitHub issues for bugs. Anything tagged with "bug" and "help +wanted" is open to whoever wants to implement it. + +Implement Features +~~~~~~~~~~~~~~~~~~ + +Look through the GitHub issues for features. Anything tagged with "enhancement" +and "help wanted" is open to whoever wants to implement it. + +Write Documentation +~~~~~~~~~~~~~~~~~~~ + +TensorFlow PhaseSpace could always use more documentation, whether as part of the +official TensorFlow PhaseSpace docs, in docstrings, or even on the web in blog posts, +articles, and such. + +Submit Feedback +~~~~~~~~~~~~~~~ + +The best way to send feedback is to file an issue at https://github.com/zfit/phasespace/issues. + +If you are proposing a feature: + +* Explain in detail how it would work. +* Keep the scope as narrow as possible, to make it easier to implement. +* Remember that this is a volunteer-driven project, and that contributions + are welcome :) + +Get Started! +------------ + +Ready to contribute? Here's how to set up `phasespace` for local development. + +1. Fork the `phasespace` repo on GitHub. +2. Clone your fork locally:: + + $ git clone git@github.com:your_name_here/phasespace.git + +3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: + + $ mkvirtualenv phasespace + $ cd phasespace/ + $ pip install -e .[dev] + +4. Create a branch for local development:: + + $ git checkout -b name-of-your-bugfix-or-feature + + Now you can make your changes locally. + +5. When you're done making changes, check that your changes pass pre-commit and the + tests: + + $ pytest + $ pre-commit run -a + +6. Commit your changes and push your branch to GitHub:: + + $ git add . + $ git commit -m "Your detailed description of your changes." + $ git push origin name-of-your-bugfix-or-feature + +7. Submit a pull request through the GitHub website. + +Pull Request Guidelines +----------------------- + +Before you submit a pull request, check that it meets these guidelines: + +1. The pull request should include tests. +2. If the pull request adds functionality, the docs should be updated. Put + your new functionality into a function with a docstring, and add the + feature to the list in README.rst. +3. The pull request should work for Python 2.7, 3.4, 3.5 and 3.6, and for PyPy. Check + https://github.com/zfit/phasespace/actions/workflows/ci.yml + and make sure that the tests pass for all supported Python versions. + +Tips +---- + +To run a subset of tests (for example those in `tests/test_generate.py`):: + + + $ pytest -k test_generate + +Deploying +--------- + +A reminder for the maintainers on how to deploy. +Make sure all your changes are committed (including an entry in HISTORY.rst). +Then run:: + +$ bumpversion patch # possible: major / minor / patch +$ git push +$ git push --tags + +GitHub Actions will then deploy to PyPI if tests pass. diff --git a/LICENSE b/LICENSE index 5b001265..d3037b76 100644 --- a/LICENSE +++ b/LICENSE @@ -1,24 +1,24 @@ -Copyright (c) 2019, zfit -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - * Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - * Neither the name of the nor the - names of its contributors may be used to endorse or promote products - derived from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL BE LIABLE FOR ANY -DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND -ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +Copyright (c) 2019, zfit +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/MANIFEST.in b/MANIFEST.in index 534b0b22..a383b4ba 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,14 +1,14 @@ -include AUTHORS.rst -include CONTRIBUTING.rst -include CHANGELOG.rst -include LICENSE -include README.rst -include pyproject.toml -include *.rst -include *.txt - -recursive-include tests * -recursive-exclude * __pycache__ -recursive-exclude * *.py[co] - -recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif +include AUTHORS.rst +include CONTRIBUTING.rst +include CHANGELOG.rst +include LICENSE +include README.rst +include pyproject.toml +include *.rst +include *.txt + +recursive-include tests * +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif diff --git a/README.rst b/README.rst index 1079a8e8..b7b4fe72 100644 --- a/README.rst +++ b/README.rst @@ -1,233 +1,233 @@ -******************************* -PhaseSpace -******************************* - -.. image:: https://joss.theoj.org/papers/10.21105/joss.01570/status.svg - :target: https://doi.org/10.21105/joss.01570 -.. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.2591993.svg - :target: https://doi.org/10.5281/zenodo.2591993 -.. image:: https://img.shields.io/pypi/status/phasespace.svg - :target: https://pypi.org/project/phasespace/ -.. image:: https://img.shields.io/pypi/pyversions/phasespace.svg - :target: https://pypi.org/project/phasespace/ -.. image:: https://github.com/zfit/phasespace/workflows/tests/badge.svg - :target: https://github.com/zfit/phasespace/actions/workflows/ci.yml?query=branch%3Amaster -.. image:: https://codecov.io/gh/zfit/phasespace/branch/master/graph/badge.svg - :target: https://codecov.io/gh/zfit/phasespace -.. image:: https://readthedocs.org/projects/phasespace/badge/?version=stable - :target: https://phasespace.readthedocs.io/en/latest/?badge=stable - :alt: Documentation Status -.. image:: https://badges.gitter.im/zfit/phasespace.svg - :target: https://gitter.im/zfit/phasespace?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge - :alt: Gitter chat - -Python implementation of the Raubold and Lynch method for `n`-body events using -TensorFlow as a backend. - -The code is based on the GENBOD function (W515 from CERNLIB), documented in [1] -and tries to follow it as closely as possible. - -Detailed documentation, including the API, can be found in https://phasespace.readthedocs.io. -Don't hesitate to join our `gitter`_ channel for questions and comments. - -If you use phasespace in a scientific publication we would appreciate citations to the `JOSS`_ publication: - -.. code-block:: bibtex - - @article{puig_eschle_phasespace-2019, - title = {phasespace: n-body phase space generation in Python}, - doi = {10.21105/joss.01570}, - url = {https://doi.org/10.21105/joss.01570}, - year = {2019}, - month = {oct}, - publisher = {The Open Journal}, - author = {Albert Puig and Jonas Eschle}, - journal = {Journal of Open Source Software} - } - -Free software: BSD-3-Clause. - -[1] F. James, Monte Carlo Phase Space, CERN 68-15 (1968) - -.. _JOSS: https://joss.theoj.org/papers/10.21105/joss.01570 -.. _Gitter: https://gitter.im/zfit/phasespace - - -Why? -==== -Lately, data analysis in High Energy Physics (HEP), traditionally performed within the `ROOT`_ ecosystem, -has been moving more and more towards Python. -The possibility of carrying out purely Python-based analyses has become real thanks to the -development of many open source Python packages, -which have allowed to replace most ROOT functionality with Python-based packages. - -One of the aspects where this is still not possible is in the random generation of `n`-body phase space events, -which are widely used in the field, for example to study kinematics -of the particle decays of interest, or to perform importance sampling in the case of complex amplitude models. -This has been traditionally done with the `TGenPhaseSpace`_ class, which is based of the GENBOD function of the -CERNLIB FORTRAN libraries and which requires a full working ROOT installation. - -This package aims to address this issue by providing a TensorFlow-based implementation of such a function -to generate `n`-body decays without requiring a ROOT installation. -Additionally, an oft-needed functionality to generate complex decay chains, not included in ``TGenPhaseSpace``, -is also offered, leaving room for decaying resonances (which don't have a fixed mass, but can be seen as a -broad peak). - -.. _ROOT: https://root.cern.ch -.. _TGenPhaseSpace: https://root.cern.ch/doc/master/classTGenPhaseSpace.html - -Installing -========== - -``phasespace`` is available on conda-forge and pip. - -To install ``phasespace`` with conda, run: - - -.. code-block:: console - - $ conda install phasespace -c conda-forge - -To install with pip: - -.. code-block:: console - - $ pip install phasespace - -This is the preferred method to install ``phasespace``, as it will always install the most recent stable release. - -For the newest development version, which may be unstable, you can install the version from git with - -.. code-block:: console - - $ pip install git+https://github.com/zfit/phasespace - - -How to use -========== - -The generation of simple `n`-body decays can be done using the ``nbody_decay`` shortcut to create a decay chain -with a very simple interface: one needs to pass the mass of the top particle and the masses of the children particle as -a list, optionally giving the names of the particles. Then, the `generate` method can be used to produce the desired sample. -For example, to generate :math:`B^0\to K\pi`, we would do: - -.. code-block:: python - - import phasespace - - B0_MASS = 5279.65 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - weights, particles = phasespace.nbody_decay(B0_MASS, - [PION_MASS, KAON_MASS]).generate(n_events=1000) - -Behind the scenes, this function runs the TensorFlow graph. It returns `tf.Tensor`, which, as TensorFlow 2.x is in eager mode, -is basically a numpy array. Any `tf.Tensor` can be explicitly converted to a numpy array by calling `tf.Tensor.numpy()` on it. -The `generate` function returns a `tf.Tensor` of 1000 elements in the case of ``weights`` and a list of -``n particles`` (2) arrays of (1000, 4) shape, -where each of the 4-dimensions corresponds to one of the components of the generated Lorentz 4-vector. -All particles are generated in the rest frame of the top particle; boosting to a certain momentum (or list of momenta) can be -achieved by passing the momenta to the ``boost_to`` argument. - -Sequential decays can be handled with the ``GenParticle`` class (used internally by ``generate``) and its ``set_children`` method. -As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which :math:`K^*\to K\pi`, we would write: - -.. code-block:: python - - from phasespace import GenParticle - - B0_MASS = 5279.65 - KSTARZ_MASS = 895.55 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - kaon = GenParticle('K+', KAON_MASS) - pion = GenParticle('pi-', PION_MASS) - kstar = GenParticle('K*', KSTARZ_MASS).set_children(kaon, pion) - gamma = GenParticle('gamma', 0) - bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) - - weights, particles = bz.generate(n_events=1000) - -Where we have used the fact that ``set_children`` returns the parent particle. -In this case, ``particles`` is a ``dict`` with the particle names as keys: - -.. code-block:: pycon - - >>> particles - {'K*': array([[ 1732.79325872, -1632.88873127, 950.85807735, 2715.78804872], - [-1633.95329448, 239.88921123, -1961.0402768 , 2715.78804872], - [ 407.15613764, -2236.6569286 , -1185.16616251, 2715.78804872], - ..., - [ 1091.64603395, -1301.78721269, 1920.07503991, 2715.78804872], - [ -517.3125083 , 1901.39296899, 1640.15905194, 2715.78804872], - [ 656.56413668, -804.76922982, 2343.99214816, 2715.78804872]]), - 'K+': array([[ 750.08077976, -547.22569019, 224.6920906 , 1075.30490935], - [-1499.90049089, 289.19714633, -1935.27960292, 2514.43047106], - [ 97.64746732, -1236.68112923, -381.09526192, 1388.47607911], - ..., - [ 508.66157459, -917.93523639, 1474.7064148 , 1876.11771642], - [ -212.28646168, 540.26381432, 610.86656669, 976.63988936], - [ 177.16656666, -535.98777569, 946.12636904, 1207.28744488]]), - 'gamma': array([[-1732.79325872, 1632.88873127, -950.85807735, 2563.79195128], - [ 1633.95329448, -239.88921123, 1961.0402768 , 2563.79195128], - [ -407.15613764, 2236.6569286 , 1185.16616251, 2563.79195128], - ..., - [-1091.64603395, 1301.78721269, -1920.07503991, 2563.79195128], - [ 517.3125083 , -1901.39296899, -1640.15905194, 2563.79195128], - [ -656.56413668, 804.76922982, -2343.99214816, 2563.79195128]]), - 'pi-': array([[ 982.71247896, -1085.66304109, 726.16598675, 1640.48313937], - [ -134.0528036 , -49.3079351 , -25.76067389, 201.35757766], - [ 309.50867032, -999.97579937, -804.0709006 , 1327.31196961], - ..., - [ 582.98445936, -383.85197629, 445.36862511, 839.6703323 ], - [ -305.02604662, 1361.12915468, 1029.29248526, 1739.14815935], - [ 479.39757002, -268.78145413, 1397.86577911, 1508.50060384]])} - -The `GenParticle` class is able to cache the graphs so it is possible to generate in a loop -without overhead: - -.. code-block:: pycon - - for i in range(10): - weights, particles = bz.generate(n_events=1000) - ... - (do something with weights and particles) - ... - -This way of generating is recommended in the case of large samples, as it allows to benefit from -parallelisation while at the same time keep the memory usage low. - -If we want to operate with the TensorFlow graph instead, we can use the `generate_tensor` method -of `GenParticle`, which has the same signature as `generate`. - -More examples can be found in the ``tests`` folder and in the `documentation`_. - -.. _documentation: https://phasespace.readthedocs.io/en/latest/usage.html - - -Physics validation -================== - -Physics validation is performed continuously in the included tests (``tests/test_physics.py``), run through GitHub Actions. -This validation is performed at two levels: - -- In simple `n`-body decays, the results of ``phasespace`` are checked against ``TGenPhaseSpace``. -- For sequential decays, the results of ``phasespace`` are checked against `RapidSim`_, a "fast Monte Carlo generator - for simulation of heavy-quark hadron decays". - In the case of resonances, differences are expected because our tests don't include proper modelling of their - mass shape, as it would require the introduction of - further dependencies. However, the results of the comparison can be expected visually. - -The results of all physics validation performed by the ``tests_physics.py`` test are written in ``tests/plots``. - -.. _RapidSim: https://github.com/gcowan/RapidSim/ - - -Contributing -============ - -Contributions are always welcome, please have a look at the `Contributing guide`_. - -.. _Contributing guide: CONTRIBUTING.rst +******************************* +PhaseSpace +******************************* + +.. image:: https://joss.theoj.org/papers/10.21105/joss.01570/status.svg + :target: https://doi.org/10.21105/joss.01570 +.. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.2591993.svg + :target: https://doi.org/10.5281/zenodo.2591993 +.. image:: https://img.shields.io/pypi/status/phasespace.svg + :target: https://pypi.org/project/phasespace/ +.. image:: https://img.shields.io/pypi/pyversions/phasespace.svg + :target: https://pypi.org/project/phasespace/ +.. image:: https://github.com/zfit/phasespace/workflows/tests/badge.svg + :target: https://github.com/zfit/phasespace/actions/workflows/ci.yml?query=branch%3Amaster +.. image:: https://codecov.io/gh/zfit/phasespace/branch/master/graph/badge.svg + :target: https://codecov.io/gh/zfit/phasespace +.. image:: https://readthedocs.org/projects/phasespace/badge/?version=stable + :target: https://phasespace.readthedocs.io/en/latest/?badge=stable + :alt: Documentation Status +.. image:: https://badges.gitter.im/zfit/phasespace.svg + :target: https://gitter.im/zfit/phasespace?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge + :alt: Gitter chat + +Python implementation of the Raubold and Lynch method for `n`-body events using +TensorFlow as a backend. + +The code is based on the GENBOD function (W515 from CERNLIB), documented in [1] +and tries to follow it as closely as possible. + +Detailed documentation, including the API, can be found in https://phasespace.readthedocs.io. +Don't hesitate to join our `gitter`_ channel for questions and comments. + +If you use phasespace in a scientific publication we would appreciate citations to the `JOSS`_ publication: + +.. code-block:: bibtex + + @article{puig_eschle_phasespace-2019, + title = {phasespace: n-body phase space generation in Python}, + doi = {10.21105/joss.01570}, + url = {https://doi.org/10.21105/joss.01570}, + year = {2019}, + month = {oct}, + publisher = {The Open Journal}, + author = {Albert Puig and Jonas Eschle}, + journal = {Journal of Open Source Software} + } + +Free software: BSD-3-Clause. + +[1] F. James, Monte Carlo Phase Space, CERN 68-15 (1968) + +.. _JOSS: https://joss.theoj.org/papers/10.21105/joss.01570 +.. _Gitter: https://gitter.im/zfit/phasespace + + +Why? +==== +Lately, data analysis in High Energy Physics (HEP), traditionally performed within the `ROOT`_ ecosystem, +has been moving more and more towards Python. +The possibility of carrying out purely Python-based analyses has become real thanks to the +development of many open source Python packages, +which have allowed to replace most ROOT functionality with Python-based packages. + +One of the aspects where this is still not possible is in the random generation of `n`-body phase space events, +which are widely used in the field, for example to study kinematics +of the particle decays of interest, or to perform importance sampling in the case of complex amplitude models. +This has been traditionally done with the `TGenPhaseSpace`_ class, which is based of the GENBOD function of the +CERNLIB FORTRAN libraries and which requires a full working ROOT installation. + +This package aims to address this issue by providing a TensorFlow-based implementation of such a function +to generate `n`-body decays without requiring a ROOT installation. +Additionally, an oft-needed functionality to generate complex decay chains, not included in ``TGenPhaseSpace``, +is also offered, leaving room for decaying resonances (which don't have a fixed mass, but can be seen as a +broad peak). + +.. _ROOT: https://root.cern.ch +.. _TGenPhaseSpace: https://root.cern.ch/doc/master/classTGenPhaseSpace.html + +Installing +========== + +``phasespace`` is available on conda-forge and pip. + +To install ``phasespace`` with conda, run: + + +.. code-block:: console + + $ conda install phasespace -c conda-forge + +To install with pip: + +.. code-block:: console + + $ pip install phasespace + +This is the preferred method to install ``phasespace``, as it will always install the most recent stable release. + +For the newest development version, which may be unstable, you can install the version from git with + +.. code-block:: console + + $ pip install git+https://github.com/zfit/phasespace + + +How to use +========== + +The generation of simple `n`-body decays can be done using the ``nbody_decay`` shortcut to create a decay chain +with a very simple interface: one needs to pass the mass of the top particle and the masses of the children particle as +a list, optionally giving the names of the particles. Then, the `generate` method can be used to produce the desired sample. +For example, to generate :math:`B^0\to K\pi`, we would do: + +.. code-block:: python + + import phasespace + + B0_MASS = 5279.65 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + weights, particles = phasespace.nbody_decay(B0_MASS, + [PION_MASS, KAON_MASS]).generate(n_events=1000) + +Behind the scenes, this function runs the TensorFlow graph. It returns `tf.Tensor`, which, as TensorFlow 2.x is in eager mode, +is basically a numpy array. Any `tf.Tensor` can be explicitly converted to a numpy array by calling `tf.Tensor.numpy()` on it. +The `generate` function returns a `tf.Tensor` of 1000 elements in the case of ``weights`` and a list of +``n particles`` (2) arrays of (1000, 4) shape, +where each of the 4-dimensions corresponds to one of the components of the generated Lorentz 4-vector. +All particles are generated in the rest frame of the top particle; boosting to a certain momentum (or list of momenta) can be +achieved by passing the momenta to the ``boost_to`` argument. + +Sequential decays can be handled with the ``GenParticle`` class (used internally by ``generate``) and its ``set_children`` method. +As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which :math:`K^*\to K\pi`, we would write: + +.. code-block:: python + + from phasespace import GenParticle + + B0_MASS = 5279.65 + KSTARZ_MASS = 895.55 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + kaon = GenParticle('K+', KAON_MASS) + pion = GenParticle('pi-', PION_MASS) + kstar = GenParticle('K*', KSTARZ_MASS).set_children(kaon, pion) + gamma = GenParticle('gamma', 0) + bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) + + weights, particles = bz.generate(n_events=1000) + +Where we have used the fact that ``set_children`` returns the parent particle. +In this case, ``particles`` is a ``dict`` with the particle names as keys: + +.. code-block:: pycon + + >>> particles + {'K*': array([[ 1732.79325872, -1632.88873127, 950.85807735, 2715.78804872], + [-1633.95329448, 239.88921123, -1961.0402768 , 2715.78804872], + [ 407.15613764, -2236.6569286 , -1185.16616251, 2715.78804872], + ..., + [ 1091.64603395, -1301.78721269, 1920.07503991, 2715.78804872], + [ -517.3125083 , 1901.39296899, 1640.15905194, 2715.78804872], + [ 656.56413668, -804.76922982, 2343.99214816, 2715.78804872]]), + 'K+': array([[ 750.08077976, -547.22569019, 224.6920906 , 1075.30490935], + [-1499.90049089, 289.19714633, -1935.27960292, 2514.43047106], + [ 97.64746732, -1236.68112923, -381.09526192, 1388.47607911], + ..., + [ 508.66157459, -917.93523639, 1474.7064148 , 1876.11771642], + [ -212.28646168, 540.26381432, 610.86656669, 976.63988936], + [ 177.16656666, -535.98777569, 946.12636904, 1207.28744488]]), + 'gamma': array([[-1732.79325872, 1632.88873127, -950.85807735, 2563.79195128], + [ 1633.95329448, -239.88921123, 1961.0402768 , 2563.79195128], + [ -407.15613764, 2236.6569286 , 1185.16616251, 2563.79195128], + ..., + [-1091.64603395, 1301.78721269, -1920.07503991, 2563.79195128], + [ 517.3125083 , -1901.39296899, -1640.15905194, 2563.79195128], + [ -656.56413668, 804.76922982, -2343.99214816, 2563.79195128]]), + 'pi-': array([[ 982.71247896, -1085.66304109, 726.16598675, 1640.48313937], + [ -134.0528036 , -49.3079351 , -25.76067389, 201.35757766], + [ 309.50867032, -999.97579937, -804.0709006 , 1327.31196961], + ..., + [ 582.98445936, -383.85197629, 445.36862511, 839.6703323 ], + [ -305.02604662, 1361.12915468, 1029.29248526, 1739.14815935], + [ 479.39757002, -268.78145413, 1397.86577911, 1508.50060384]])} + +The `GenParticle` class is able to cache the graphs so it is possible to generate in a loop +without overhead: + +.. code-block:: pycon + + for i in range(10): + weights, particles = bz.generate(n_events=1000) + ... + (do something with weights and particles) + ... + +This way of generating is recommended in the case of large samples, as it allows to benefit from +parallelisation while at the same time keep the memory usage low. + +If we want to operate with the TensorFlow graph instead, we can use the `generate_tensor` method +of `GenParticle`, which has the same signature as `generate`. + +More examples can be found in the ``tests`` folder and in the `documentation`_. + +.. _documentation: https://phasespace.readthedocs.io/en/latest/usage.html + + +Physics validation +================== + +Physics validation is performed continuously in the included tests (``tests/test_physics.py``), run through GitHub Actions. +This validation is performed at two levels: + +- In simple `n`-body decays, the results of ``phasespace`` are checked against ``TGenPhaseSpace``. +- For sequential decays, the results of ``phasespace`` are checked against `RapidSim`_, a "fast Monte Carlo generator + for simulation of heavy-quark hadron decays". + In the case of resonances, differences are expected because our tests don't include proper modelling of their + mass shape, as it would require the introduction of + further dependencies. However, the results of the comparison can be expected visually. + +The results of all physics validation performed by the ``tests_physics.py`` test are written in ``tests/plots``. + +.. _RapidSim: https://github.com/gcowan/RapidSim/ + + +Contributing +============ + +Contributions are always welcome, please have a look at the `Contributing guide`_. + +.. _Contributing guide: CONTRIBUTING.rst diff --git a/benchmark/bench_phasespace.py b/benchmark/bench_phasespace.py index d9cae3b0..52bcc4f2 100644 --- a/benchmark/bench_phasespace.py +++ b/benchmark/bench_phasespace.py @@ -1,140 +1,140 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file bench_phasespace.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Benchmark phasespace.""" - -import os -import sys -from timeit import default_timer - -import tensorflow as tf - -from phasespace import phasespace - -sys.path.append(os.path.dirname(__file__)) - - -def memory_usage(): - """Get memory usage of current process in MiB. - - Tries to use :mod:`psutil`, if possible, otherwise fallback to calling - ``ps`` directly. - - Return: - float: Memory usage of the current process. - """ - pid = os.getpid() - try: - import psutil - - process = psutil.Process(pid) - mem = process.memory_info()[0] / float(2 ** 20) - except ImportError: - import subprocess - - out = ( - subprocess.Popen(["ps", "v", "-p", str(pid)], stdout=subprocess.PIPE) - .communicate()[0] - .split(b"\n") - ) - vsz_index = out[0].split().index(b"RSS") - mem = float(out[1].split()[vsz_index]) / 1024 - return mem - - -# pylint: disable=too-few-public-methods -class Timer: - """Time the code placed inside its context. - - Taken from http://coreygoldberg.blogspot.ch/2012/06/python-timer-class-context-manager-for.html - - Attributes: - verbose (bool): Print the elapsed time at context exit? - start (float): Start time in seconds since Epoch Time. Value set - to 0 if not run. - elapsed (float): Elapsed seconds in the timer. Value set to - 0 if not run. - - Arguments: - verbose (bool, optional): Print the elapsed time at - context exit? Defaults to False. - """ - - def __init__(self, verbose=False, n=1): - """Initialize the timer.""" - self.verbose = verbose - self.n = n - self._timer = default_timer - self.start = 0 - self.elapsed = 0 - - def __enter__(self): - self.start = self._timer() - return self - - def __exit__(self, *args): - self.elapsed = self._timer() - self.start - if self.verbose: - print(f"Elapsed time: {self.elapsed * 1000.0 / self.n} ms") - - -# EOF - - -# to play around with optimization, no big effect though -NUM_PARALLEL_EXEC_UNITS = 1 -# config = tf.ConfigProto( -# intra_op_parallelism_threads=NUM_PARALLEL_EXEC_UNITS, -# inter_op_parallelism_threads=1, -# allow_soft_placement=True, -# device_count={"CPU": NUM_PARALLEL_EXEC_UNITS}, -# ) - -B_MASS = 5279.0 -B_AT_REST = tf.stack((0.0, 0.0, 0.0, B_MASS), axis=-1) -PION_MASS = 139.6 - -N_EVENTS = 1000000 -CHUNK_SIZE = int(N_EVENTS) - -n_runs = 10 - - -# N_EVENTS_VAR = tf.Variable(initial_value=N_EVENTS) -# CHUNK_SIZE_VAR = tf.Variable(initial_value=CHUNK_SIZE) - - -def test_three_body(): - """Test B -> pi pi pi decay.""" - with Timer(verbose=True): - print("Initial run (may takes more time than consequent runs)") - do_run() # to get rid of initial overhead - print("starting benchmark") - with Timer(verbose=True, n=n_runs): - for _ in range(n_runs): - # CHUNK_SIZE_VAR.assign(CHUNK_SIZE + 1) # +1 to make sure we're not using any trivial caching - samples = do_run() - - print(f"nevents produced {samples[0][0].shape}") - print("Shape of one particle momentum", samples[0][1]["p_0"].shape) - - -decay = phasespace.nbody_decay( - B_MASS, - [PION_MASS, PION_MASS, PION_MASS], -) - - -# tf.config.run_functions_eagerly(True) -@tf.function(autograph=False) -def do_run(): - return [decay.generate(N_EVENTS) for _ in range(0, N_EVENTS, CHUNK_SIZE)] - - -if __name__ == "__main__": - test_three_body() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file bench_phasespace.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Benchmark phasespace.""" + +import os +import sys +from timeit import default_timer + +import tensorflow as tf + +from phasespace import phasespace + +sys.path.append(os.path.dirname(__file__)) + + +def memory_usage(): + """Get memory usage of current process in MiB. + + Tries to use :mod:`psutil`, if possible, otherwise fallback to calling + ``ps`` directly. + + Return: + float: Memory usage of the current process. + """ + pid = os.getpid() + try: + import psutil + + process = psutil.Process(pid) + mem = process.memory_info()[0] / float(2 ** 20) + except ImportError: + import subprocess + + out = ( + subprocess.Popen(["ps", "v", "-p", str(pid)], stdout=subprocess.PIPE) + .communicate()[0] + .split(b"\n") + ) + vsz_index = out[0].split().index(b"RSS") + mem = float(out[1].split()[vsz_index]) / 1024 + return mem + + +# pylint: disable=too-few-public-methods +class Timer: + """Time the code placed inside its context. + + Taken from http://coreygoldberg.blogspot.ch/2012/06/python-timer-class-context-manager-for.html + + Attributes: + verbose (bool): Print the elapsed time at context exit? + start (float): Start time in seconds since Epoch Time. Value set + to 0 if not run. + elapsed (float): Elapsed seconds in the timer. Value set to + 0 if not run. + + Arguments: + verbose (bool, optional): Print the elapsed time at + context exit? Defaults to False. + """ + + def __init__(self, verbose=False, n=1): + """Initialize the timer.""" + self.verbose = verbose + self.n = n + self._timer = default_timer + self.start = 0 + self.elapsed = 0 + + def __enter__(self): + self.start = self._timer() + return self + + def __exit__(self, *args): + self.elapsed = self._timer() - self.start + if self.verbose: + print(f"Elapsed time: {self.elapsed * 1000.0 / self.n} ms") + + +# EOF + + +# to play around with optimization, no big effect though +NUM_PARALLEL_EXEC_UNITS = 1 +# config = tf.ConfigProto( +# intra_op_parallelism_threads=NUM_PARALLEL_EXEC_UNITS, +# inter_op_parallelism_threads=1, +# allow_soft_placement=True, +# device_count={"CPU": NUM_PARALLEL_EXEC_UNITS}, +# ) + +B_MASS = 5279.0 +B_AT_REST = tf.stack((0.0, 0.0, 0.0, B_MASS), axis=-1) +PION_MASS = 139.6 + +N_EVENTS = 1000000 +CHUNK_SIZE = int(N_EVENTS) + +n_runs = 10 + + +# N_EVENTS_VAR = tf.Variable(initial_value=N_EVENTS) +# CHUNK_SIZE_VAR = tf.Variable(initial_value=CHUNK_SIZE) + + +def test_three_body(): + """Test B -> pi pi pi decay.""" + with Timer(verbose=True): + print("Initial run (may takes more time than consequent runs)") + do_run() # to get rid of initial overhead + print("starting benchmark") + with Timer(verbose=True, n=n_runs): + for _ in range(n_runs): + # CHUNK_SIZE_VAR.assign(CHUNK_SIZE + 1) # +1 to make sure we're not using any trivial caching + samples = do_run() + + print(f"nevents produced {samples[0][0].shape}") + print("Shape of one particle momentum", samples[0][1]["p_0"].shape) + + +decay = phasespace.nbody_decay( + B_MASS, + [PION_MASS, PION_MASS, PION_MASS], +) + + +# tf.config.run_functions_eagerly(True) +@tf.function(autograph=False) +def do_run(): + return [decay.generate(N_EVENTS) for _ in range(0, N_EVENTS, CHUNK_SIZE)] + + +if __name__ == "__main__": + test_three_body() + +# EOF diff --git a/benchmark/bench_tgenphasespace.cxx b/benchmark/bench_tgenphasespace.cxx index 5f6211da..4001a587 100644 --- a/benchmark/bench_tgenphasespace.cxx +++ b/benchmark/bench_tgenphasespace.cxx @@ -1,21 +1,21 @@ -#include "TROOT.h" -#include "TSystem.h" -#include "TFile.h" -#include "TTree.h" -#include "TLorentzVector.h" -#include "TGenPhaseSpace.h" - -Int_t N_EVENTS = 1000000; - -int bench_tgenphasespace() -{ - TLorentzVector B(0.0, 0.0, 0.0, 5279.0); - Double_t masses[3] = {139.6, 139.6, 139.6}; - - if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); - TGenPhaseSpace event; - event.SetDecay(B, 3, masses); - - for (Int_t n=0; nGetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); + TGenPhaseSpace event; + event.SetDecay(B, 3, masses); + + for (Int_t n=0; n google style - -# Napoleon settings (convert numpy/google docstrings to proper ReST -napoleon_google_docstring = not using_numpy_style -napoleon_numpy_docstring = using_numpy_style -napoleon_include_init_with_doc = False -napoleon_include_private_with_doc = False -napoleon_include_special_with_doc = True -napoleon_use_admonition_for_examples = False -napoleon_use_admonition_for_notes = False -napoleon_use_admonition_for_references = False -napoleon_use_ivar = False -napoleon_use_param = True -napoleon_use_rtype = True - -# Add any paths that contain templates here, relative to this directory. -templates_path = ["_templates"] - -# The suffix(es) of source filenames. -# You can specify multiple suffix as a list of string: -# -# source_suffix = ['.rst', '.md'] -source_suffix = ".rst" - -# The master toctree document. -master_doc = "index" - -# General information about the project. -project = "TensorFlow PhaseSpace" -copyright = "2019, Albert Puig Navarro" -author = "Albert Puig Navarro" - -# The version info for the project you're documenting, acts as replacement -# for |version| and |release|, also used in various other places throughout -# the built documents. -# -# The short X.Y version. -version = phasespace.__version__ -# The full version, including alpha/beta/rc tags. -release = phasespace.__version__ - -# The language for content autogenerated by Sphinx. Refer to documentation -# for a list of supported languages. -# -# This is also used if you do content translation via gettext catalogs. -# Usually you set "language" from the command line for these cases. -language = "en" - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - -# The name of the Pygments (syntax highlighting) style to use. -pygments_style = "sphinx" - -# If true, `todo` and `todoList` produce output, else they produce nothing. -todo_include_todos = False - -# makes the jupyter extension executable -jupyter_sphinx_thebelab_config = { - "requestKernel": True, - "binderOptions": { - "repo": "zfit/phasespace", - "binderUrl": "https://mybinder.org", - "repoProvider": "github", - }, -} - -# -- Options for HTML output ------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "bootstrap" -html_theme_path = sphinx_bootstrap_theme.get_html_theme_path() -html_theme_options = { - # Navigation bar title. (Default: ``project`` value) - "navbar_title": "phasespace", - # Tab name for entire site. (Default: "Site") - # 'navbar_site_name': "Docs", - # 'navbar_site_name': "Overview", - # A list of tuples containing pages or urls to link to. - # Valid tuples should be in the following forms: - # (name, page) # a link to a page - # (name, "/aa/bb", 1) # a link to an arbitrary relative url - # (name, "http://example.com", True) # arbitrary absolute url - # Note the "1" or "True" value above as the third argument to indicate - # an arbitrary url. - "navbar_links": [ - ("Phasespace", "index"), - ("Usage", "usage"), - ("API", "phasespace"), - ("Contributing", "contributing"), - # ("Link", "http://example.com", True), - ], - # Render the next and previous page links in navbar. (Default: true) - "navbar_sidebarrel": False, - # Render the current pages TOC in the navbar. (Default: true) - "navbar_pagenav": False, - # Tab name for the current pages TOC. (Default: "Page") - # 'navbar_pagenav_name': "Page", - # Global TOC depth for "site" navbar tab. (Default: 1) - # Switching to -1 shows all levels. - "globaltoc_depth": 1, - # Include hidden TOCs in Site navbar? - # - # Note: If this is "false", you cannot have mixed ``:hidden:`` and - # non-hidden ``toctree`` directives in the same page, or else the build - # will break. - # - # Values: "true" (default) or "false" - "globaltoc_includehidden": "true", - # HTML navbar class (Default: "navbar") to attach to
element. - # For black navbar, do "navbar navbar-inverse" - # 'navbar_class': "navbar navbar-inverse", - "navbar_class": "navbar", - # Fix navigation bar to top of page? - # Values: "true" (default) or "false" - "navbar_fixed_top": "true", - # Location of link to source. - # Options are "nav" (default), "footer" or anything else to exclude. - # 'source_link_position': "nav", - "source_link_position": False, - # Bootswatch (http://bootswatch.com/) theme. - # - # Options are nothing (default) or the name of a valid theme - # such as "cosmo" or "sandstone". - # - # The set of valid themes depend on the version of Bootstrap - # that's used (the next config option). - # - # Currently, the supported themes are: - # - Bootstrap 2: https://bootswatch.com/2 - # - Bootstrap 3: https://bootswatch.com/3 - "bootswatch_theme": "flatly", - # Choose Bootstrap version. - # Values: "3" (default) or "2" (in quotes) - "bootstrap_version": "4", -} -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = [] - -# -- Options for HTMLHelp output --------------------------------------- - -# Output file base name for HTML help builder. -htmlhelp_basename = "phasespacedoc" - -# -- Options for LaTeX output ------------------------------------------ - -latex_elements = { - # The paper size ('letterpaper' or 'a4paper'). - # - # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). - # - # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. - # - # 'preamble': '', - # Latex figure (float) alignment - # - # 'figure_align': 'htbp', -} - -# Grouping the document tree into LaTeX files. List of tuples -# (source start file, target name, title, author, documentclass -# [howto, manual, or own class]). -latex_documents = [ - ( - master_doc, - "phasespace.tex", - "TensorFlow PhaseSpace Documentation", - "Albert Puig Navarro", - "manual", - ), -] - -# -- Options for manual page output ------------------------------------ - -# One entry per manual page. List of tuples -# (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, "phasespace", "TensorFlow PhaseSpace Documentation", [author], 1) -] - -# -- Options for Texinfo output ---------------------------------------- - -# Grouping the document tree into Texinfo files. List of tuples -# (source start file, target name, title, author, -# dir menu entry, description, category) -texinfo_documents = [ - ( - master_doc, - "phasespace", - "TensorFlow PhaseSpace Documentation", - author, - "phasespace", - "One line description of project.", - "Miscellaneous", - ), -] +#!/usr/bin/env python +# +# phasespace documentation build configuration file, created by +# sphinx-quickstart on Fri Jun 9 13:47:02 2017. +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +# If extensions (or modules to document with autodoc) are in another +# directory, add these directories to sys.path here. If the directory is +# relative to the documentation root, use os.path.abspath to make it +# absolute, like shown here. +# + +import sphinx_bootstrap_theme + +import phasespace + +# -- General configuration --------------------------------------------- + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# use for classes the class and the __init__ docs combined +autoclass_content = "both" + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones. +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.viewcode", + "sphinx.ext.napoleon", + "sphinx.ext.mathjax", + "sphinx_math_dollar", + "jupyter_sphinx", +] + + +mathjax_config = { + "tex2jax": { + "inlineMath": [["\\(", "\\)"]], + "displayMath": [["\\[", "\\]"]], + }, +} +using_numpy_style = False # False -> google style + +# Napoleon settings (convert numpy/google docstrings to proper ReST +napoleon_google_docstring = not using_numpy_style +napoleon_numpy_docstring = using_numpy_style +napoleon_include_init_with_doc = False +napoleon_include_private_with_doc = False +napoleon_include_special_with_doc = True +napoleon_use_admonition_for_examples = False +napoleon_use_admonition_for_notes = False +napoleon_use_admonition_for_references = False +napoleon_use_ivar = False +napoleon_use_param = True +napoleon_use_rtype = True + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# +# source_suffix = ['.rst', '.md'] +source_suffix = ".rst" + +# The master toctree document. +master_doc = "index" + +# General information about the project. +project = "TensorFlow PhaseSpace" +copyright = "2019, Albert Puig Navarro" +author = "Albert Puig Navarro" + +# The version info for the project you're documenting, acts as replacement +# for |version| and |release|, also used in various other places throughout +# the built documents. +# +# The short X.Y version. +version = phasespace.__version__ +# The full version, including alpha/beta/rc tags. +release = phasespace.__version__ + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This patterns also effect to html_static_path and html_extra_path +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = False + +# makes the jupyter extension executable +jupyter_sphinx_thebelab_config = { + "requestKernel": True, + "binderOptions": { + "repo": "zfit/phasespace", + "binderUrl": "https://mybinder.org", + "repoProvider": "github", + }, +} + +# -- Options for HTML output ------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "bootstrap" +html_theme_path = sphinx_bootstrap_theme.get_html_theme_path() +html_theme_options = { + # Navigation bar title. (Default: ``project`` value) + "navbar_title": "phasespace", + # Tab name for entire site. (Default: "Site") + # 'navbar_site_name': "Docs", + # 'navbar_site_name': "Overview", + # A list of tuples containing pages or urls to link to. + # Valid tuples should be in the following forms: + # (name, page) # a link to a page + # (name, "/aa/bb", 1) # a link to an arbitrary relative url + # (name, "http://example.com", True) # arbitrary absolute url + # Note the "1" or "True" value above as the third argument to indicate + # an arbitrary url. + "navbar_links": [ + ("Phasespace", "index"), + ("Usage", "usage"), + ("API", "phasespace"), + ("Contributing", "contributing"), + # ("Link", "http://example.com", True), + ], + # Render the next and previous page links in navbar. (Default: true) + "navbar_sidebarrel": False, + # Render the current pages TOC in the navbar. (Default: true) + "navbar_pagenav": False, + # Tab name for the current pages TOC. (Default: "Page") + # 'navbar_pagenav_name': "Page", + # Global TOC depth for "site" navbar tab. (Default: 1) + # Switching to -1 shows all levels. + "globaltoc_depth": 1, + # Include hidden TOCs in Site navbar? + # + # Note: If this is "false", you cannot have mixed ``:hidden:`` and + # non-hidden ``toctree`` directives in the same page, or else the build + # will break. + # + # Values: "true" (default) or "false" + "globaltoc_includehidden": "true", + # HTML navbar class (Default: "navbar") to attach to
element. + # For black navbar, do "navbar navbar-inverse" + # 'navbar_class': "navbar navbar-inverse", + "navbar_class": "navbar", + # Fix navigation bar to top of page? + # Values: "true" (default) or "false" + "navbar_fixed_top": "true", + # Location of link to source. + # Options are "nav" (default), "footer" or anything else to exclude. + # 'source_link_position': "nav", + "source_link_position": False, + # Bootswatch (http://bootswatch.com/) theme. + # + # Options are nothing (default) or the name of a valid theme + # such as "cosmo" or "sandstone". + # + # The set of valid themes depend on the version of Bootstrap + # that's used (the next config option). + # + # Currently, the supported themes are: + # - Bootstrap 2: https://bootswatch.com/2 + # - Bootstrap 3: https://bootswatch.com/3 + "bootswatch_theme": "flatly", + # Choose Bootstrap version. + # Values: "3" (default) or "2" (in quotes) + "bootstrap_version": "4", +} +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = [] + +# -- Options for HTMLHelp output --------------------------------------- + +# Output file base name for HTML help builder. +htmlhelp_basename = "phasespacedoc" + +# -- Options for LaTeX output ------------------------------------------ + +latex_elements = { + # The paper size ('letterpaper' or 'a4paper'). + # + # 'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + # + # 'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + # + # 'preamble': '', + # Latex figure (float) alignment + # + # 'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, author, documentclass +# [howto, manual, or own class]). +latex_documents = [ + ( + master_doc, + "phasespace.tex", + "TensorFlow PhaseSpace Documentation", + "Albert Puig Navarro", + "manual", + ), +] + +# -- Options for manual page output ------------------------------------ + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + (master_doc, "phasespace", "TensorFlow PhaseSpace Documentation", [author], 1) +] + +# -- Options for Texinfo output ---------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ( + master_doc, + "phasespace", + "TensorFlow PhaseSpace Documentation", + author, + "phasespace", + "One line description of project.", + "Miscellaneous", + ), +] diff --git a/docs/contributing.rst b/docs/contributing.rst index e582053e..819f45e0 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -1 +1 @@ -.. include:: ../CONTRIBUTING.rst +.. include:: ../CONTRIBUTING.rst diff --git a/docs/history.rst b/docs/history.rst index 565b0521..9ad7a341 100644 --- a/docs/history.rst +++ b/docs/history.rst @@ -1 +1 @@ -.. include:: ../CHANGELOG.rst +.. include:: ../CHANGELOG.rst diff --git a/docs/index.rst b/docs/index.rst index 24a715e8..232f91e1 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,17 +1,17 @@ - -.. include:: ../README.rst - -.. include:: authors.rst - -================= -Table of Contents -================= - -.. toctree:: - :maxdepth: 1 - - usage - API - contributing - authors - history + +.. include:: ../README.rst + +.. include:: authors.rst + +================= +Table of Contents +================= + +.. toctree:: + :maxdepth: 1 + + usage + API + contributing + authors + history diff --git a/docs/make.bat b/docs/make.bat index 31e560f0..25284887 100644 --- a/docs/make.bat +++ b/docs/make.bat @@ -1,36 +1,36 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=python -msphinx -) -set SOURCEDIR=. -set BUILDDIR=_build -set SPHINXPROJ=phasespace - -if "%1" == "" goto help - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The Sphinx module was not found. Make sure you have Sphinx installed, - echo.then set the SPHINXBUILD environment variable to point to the full - echo.path of the 'sphinx-build' executable. Alternatively you may add the - echo.Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.http://sphinx-doc.org/ - exit /b 1 -) - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% - -:end -popd +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=python -msphinx +) +set SOURCEDIR=. +set BUILDDIR=_build +set SPHINXPROJ=phasespace + +if "%1" == "" goto help + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The Sphinx module was not found. Make sure you have Sphinx installed, + echo.then set the SPHINXBUILD environment variable to point to the full + echo.path of the 'sphinx-build' executable. Alternatively you may add the + echo.Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 +) + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% + +:end +popd diff --git a/docs/phasespace.rst b/docs/phasespace.rst index ab905f15..02104012 100644 --- a/docs/phasespace.rst +++ b/docs/phasespace.rst @@ -1,18 +1,18 @@ -phasespace package -==================== - -phasespace.phasespace module --------------------------------- - -.. automodule:: phasespace.phasespace - :members: - :undoc-members: - :show-inheritance: - -phasespace.kinematics module ------------------------------- - -.. automodule:: phasespace.kinematics - :members: - :undoc-members: - :show-inheritance: +phasespace package +==================== + +phasespace.phasespace module +-------------------------------- + +.. automodule:: phasespace.phasespace + :members: + :undoc-members: + :show-inheritance: + +phasespace.kinematics module +------------------------------ + +.. automodule:: phasespace.kinematics + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/usage.rst b/docs/usage.rst index ee6f0f38..11232063 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -1,203 +1,203 @@ -===== -Usage -===== - -The base of ``phasespace`` is the ``GenParticle`` object. -This object, which represents a particle, either stable or decaying, has only one mandatory argument, its name. - -In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either -a number or ``tf.constant``, or a function. -Functions are used to specify the mass of particles such as resonances, which are not fixed but vary according to -a broad distribution. -These mass functions get three arguments, and must return a ``TensorFlow`` Tensor: - -- The minimum mass allowed by the decay chain, which will be of shape `(n_events,)`. -- The maximum mass available, which will be of shape `(n_events,)`. -- The number of events to generate. - -This function signature allows to handle threshold effects cleanly, giving enough information to produce kinematically -allowed decays (NB: ``phasespace`` will throw an error if a kinematically forbidden decay is requested). - -A simple example --------------------------- - - -With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the ``GenParticle`` -class (which returns the class itself). As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which -:math:`K^*\to K\pi` with a fixed mass, we would write: - -.. jupyter-execute:: - - from phasespace import GenParticle - - B0_MASS = 5279.58 - KSTARZ_MASS = 895.81 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - pion = GenParticle('pi-', PION_MASS) - kaon = GenParticle('K+', KAON_MASS) - kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon) - gamma = GenParticle('gamma', 0) - bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) - -.. thebe-button:: Run this interactively - - -Phasespace events can be generated using the ``generate`` method, which gets the number of events to generate as input. -The method returns: - -- The normalized weights of each event, as an array of dimension (n_events,). -- The 4-momenta of the generated particles as values of a dictionary with the particle name as key. These momenta - are expressed as arrays of dimension (n_events, 4). - -.. jupyter-execute:: - - N_EVENTS = 1000 - - weights, particles = bz.generate(n_events=N_EVENTS) - -The ``generate`` method return an eager ``Tensor``: this is basically a wrapped numpy array. With ``Tensor.numpy()``, -it can always directly be converted to a numpy array (if really needed). - -Boosting the particles --------------------------- - - -The particles are generated in the rest frame of the top particle. -To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in both -``generate`` and ~`tf.Tensor`. This latter approach can be useful if the momentum of the top particle -is generated according to some distribution, for example the kinematics of the LHC (see ``test_kstargamma_kstarnonresonant_lhc`` -and ``test_k1gamma_kstarnonresonant_lhc`` in ``tests/test_physics.py`` to see how this could be done). - -Weights --------------------------- - - -Additionally, it is possible to obtain the unnormalized weights by using the ``generate_unnormalized`` flag in -``generate``. In this case, the method returns the unnormalized weights, the per-event maximum weight -and the particle dictionary. - -.. jupyter-execute:: - - print(particles) - - -It is worth noting that the graph generation is cached even when using ``generate``, so iterative generation -can be performed using normal python loops without loss in performance: - -.. jupyter-execute:: - - for i in range(5): - weights, particles = bz.generate(n_events=100) - # ... - # (do something with weights and particles) - # ... - - - -Resonances with variable mass ------------------------------- - - -To generate the mass of a resonance, we need to give a function as its mass instead of a floating number. -This function should take as input the per-event lower mass allowed, per-event upper mass allowed and the number of -events, and should return a ~`tf.Tensor` with the generated masses and shape (nevents,). Well suited for this task -are the `TensorFlow Probability distributions `_ -or, for more customized mass shapes, the -`zfit pdfs `_ (currently an -*experimental feature* is needed, contact the `zfit developers `_ to learn more). - -Following with the same example as above, and approximating the resonance shape by a gaussian, we could -write the :math:`B^{0}\to K^{*}\gamma` decay chain as (more details can be found in ``tests/helpers/decays.py``): - -.. jupyter-execute:: - :hide-output: - - import tensorflow as tf - import tensorflow_probability as tfp - from phasespace import GenParticle - - KSTARZ_MASS = 895.81 - KSTARZ_WIDTH = 47.4 - - def kstar_mass(min_mass, max_mass, n_events): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - if KSTARZ_WIDTH > 0: - kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, - scale=kstar_width_cast, - low=min_mass, - high=max_mass).sample() - return kstar_mass - - bz = GenParticle('B0', B0_MASS).set_children(GenParticle('K*0', mass=kstar_mass) - .set_children(GenParticle('K+', mass=KAON_MASS), - GenParticle('pi-', mass=PION_MASS)), - GenParticle('gamma', mass=0.0)) - - bz.generate(n_events=500) - - -Shortcut for simple decays --------------------------- - -The generation of simple `n`-body decay chains can be done using the ``nbody_decay`` function of ``phasespace``, which takes - -- The mass of the top particle. -- The mass of children particles as a list. -- The name of the top particle (optional). -- The names of the children particles (optional). - -If the names are not given, `top` and `p_{i}` are assigned. For example, to generate :math:`B^0\to K\pi`, one would do: - -.. jupyter-execute:: - - import phasespace - - N_EVENTS = 1000 - - B0_MASS = 5279.58 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS], - top_name="B0", names=["pi", "K"]) - weights, particles = decay.generate(n_events=N_EVENTS) - -In this example, ``decay`` is simply a ``GenParticle`` with the corresponding children. - - -Eager execution ---------------- - -By default, `phasespace` uses TensorFlow to build a graph of the computations. This is usually more -performant, especially if used multiple times. However, this has the disadvantage that _inside_ -`phasespac`, the actual values are not computed on Python runtime, e.g. if a breakpoint is set -the values of a `tf.Tensor` won't be available. - -TensorFlow (since version 2.0) however can easily switch to so called "eager execution": in this -mode, it behaves the same as Numpy; values are computed instantly and the Python code is not only -executed once but every time. - -To switch this on or off, the global flag in TensorFlow `tf.config.run_functions_eagerly(True)` or -the enviroment variable "PHASESPACE_EAGER" (which switches this flag) can be used. - -Random numbers --------------- - -The random number generation inside `phasespace` is transparent in order to allow for deterministic -behavior if desired. A function that uses random number generation inside always takes a `seed` (or `rng`) -argument. The behavior is as follows - -- if no seed is given, the global random number generator of TensorFlow will be used. Setting this - instance explicitly or by setting the seed via `tf.random.set_seed` allows for a deterministic - execution of a whole _script_. -- if the seed is a number it will be used to create a random number generator from this. Using the - same seed again will result in the same output. -- if the seed is an instance of :py:class:`tf.random.Generator`, this instance will directly be used - and advances an undefined number of steps. +===== +Usage +===== + +The base of ``phasespace`` is the ``GenParticle`` object. +This object, which represents a particle, either stable or decaying, has only one mandatory argument, its name. + +In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either +a number or ``tf.constant``, or a function. +Functions are used to specify the mass of particles such as resonances, which are not fixed but vary according to +a broad distribution. +These mass functions get three arguments, and must return a ``TensorFlow`` Tensor: + +- The minimum mass allowed by the decay chain, which will be of shape `(n_events,)`. +- The maximum mass available, which will be of shape `(n_events,)`. +- The number of events to generate. + +This function signature allows to handle threshold effects cleanly, giving enough information to produce kinematically +allowed decays (NB: ``phasespace`` will throw an error if a kinematically forbidden decay is requested). + +A simple example +-------------------------- + + +With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the ``GenParticle`` +class (which returns the class itself). As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which +:math:`K^*\to K\pi` with a fixed mass, we would write: + +.. jupyter-execute:: + + from phasespace import GenParticle + + B0_MASS = 5279.58 + KSTARZ_MASS = 895.81 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + pion = GenParticle('pi-', PION_MASS) + kaon = GenParticle('K+', KAON_MASS) + kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon) + gamma = GenParticle('gamma', 0) + bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) + +.. thebe-button:: Run this interactively + + +Phasespace events can be generated using the ``generate`` method, which gets the number of events to generate as input. +The method returns: + +- The normalized weights of each event, as an array of dimension (n_events,). +- The 4-momenta of the generated particles as values of a dictionary with the particle name as key. These momenta + are expressed as arrays of dimension (n_events, 4). + +.. jupyter-execute:: + + N_EVENTS = 1000 + + weights, particles = bz.generate(n_events=N_EVENTS) + +The ``generate`` method return an eager ``Tensor``: this is basically a wrapped numpy array. With ``Tensor.numpy()``, +it can always directly be converted to a numpy array (if really needed). + +Boosting the particles +-------------------------- + + +The particles are generated in the rest frame of the top particle. +To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in both +``generate`` and ~`tf.Tensor`. This latter approach can be useful if the momentum of the top particle +is generated according to some distribution, for example the kinematics of the LHC (see ``test_kstargamma_kstarnonresonant_lhc`` +and ``test_k1gamma_kstarnonresonant_lhc`` in ``tests/test_physics.py`` to see how this could be done). + +Weights +-------------------------- + + +Additionally, it is possible to obtain the unnormalized weights by using the ``generate_unnormalized`` flag in +``generate``. In this case, the method returns the unnormalized weights, the per-event maximum weight +and the particle dictionary. + +.. jupyter-execute:: + + print(particles) + + +It is worth noting that the graph generation is cached even when using ``generate``, so iterative generation +can be performed using normal python loops without loss in performance: + +.. jupyter-execute:: + + for i in range(5): + weights, particles = bz.generate(n_events=100) + # ... + # (do something with weights and particles) + # ... + + + +Resonances with variable mass +------------------------------ + + +To generate the mass of a resonance, we need to give a function as its mass instead of a floating number. +This function should take as input the per-event lower mass allowed, per-event upper mass allowed and the number of +events, and should return a ~`tf.Tensor` with the generated masses and shape (nevents,). Well suited for this task +are the `TensorFlow Probability distributions `_ +or, for more customized mass shapes, the +`zfit pdfs `_ (currently an +*experimental feature* is needed, contact the `zfit developers `_ to learn more). + +Following with the same example as above, and approximating the resonance shape by a gaussian, we could +write the :math:`B^{0}\to K^{*}\gamma` decay chain as (more details can be found in ``tests/helpers/decays.py``): + +.. jupyter-execute:: + :hide-output: + + import tensorflow as tf + import tensorflow_probability as tfp + from phasespace import GenParticle + + KSTARZ_MASS = 895.81 + KSTARZ_WIDTH = 47.4 + + def kstar_mass(min_mass, max_mass, n_events): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + if KSTARZ_WIDTH > 0: + kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, + scale=kstar_width_cast, + low=min_mass, + high=max_mass).sample() + return kstar_mass + + bz = GenParticle('B0', B0_MASS).set_children(GenParticle('K*0', mass=kstar_mass) + .set_children(GenParticle('K+', mass=KAON_MASS), + GenParticle('pi-', mass=PION_MASS)), + GenParticle('gamma', mass=0.0)) + + bz.generate(n_events=500) + + +Shortcut for simple decays +-------------------------- + +The generation of simple `n`-body decay chains can be done using the ``nbody_decay`` function of ``phasespace``, which takes + +- The mass of the top particle. +- The mass of children particles as a list. +- The name of the top particle (optional). +- The names of the children particles (optional). + +If the names are not given, `top` and `p_{i}` are assigned. For example, to generate :math:`B^0\to K\pi`, one would do: + +.. jupyter-execute:: + + import phasespace + + N_EVENTS = 1000 + + B0_MASS = 5279.58 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS], + top_name="B0", names=["pi", "K"]) + weights, particles = decay.generate(n_events=N_EVENTS) + +In this example, ``decay`` is simply a ``GenParticle`` with the corresponding children. + + +Eager execution +--------------- + +By default, `phasespace` uses TensorFlow to build a graph of the computations. This is usually more +performant, especially if used multiple times. However, this has the disadvantage that _inside_ +`phasespac`, the actual values are not computed on Python runtime, e.g. if a breakpoint is set +the values of a `tf.Tensor` won't be available. + +TensorFlow (since version 2.0) however can easily switch to so called "eager execution": in this +mode, it behaves the same as Numpy; values are computed instantly and the Python code is not only +executed once but every time. + +To switch this on or off, the global flag in TensorFlow `tf.config.run_functions_eagerly(True)` or +the enviroment variable "PHASESPACE_EAGER" (which switches this flag) can be used. + +Random numbers +-------------- + +The random number generation inside `phasespace` is transparent in order to allow for deterministic +behavior if desired. A function that uses random number generation inside always takes a `seed` (or `rng`) +argument. The behavior is as follows + +- if no seed is given, the global random number generator of TensorFlow will be used. Setting this + instance explicitly or by setting the seed via `tf.random.set_seed` allows for a deterministic + execution of a whole _script_. +- if the seed is a number it will be used to create a random number generator from this. Using the + same seed again will result in the same output. +- if the seed is an instance of :py:class:`tf.random.Generator`, this instance will directly be used + and advances an undefined number of steps. diff --git a/paper/paper.bib b/paper/paper.bib index 7dd5b40c..88950958 100644 --- a/paper/paper.bib +++ b/paper/paper.bib @@ -1,115 +1,115 @@ -@article{Brun:1997pa, - author = "Brun, R. and Rademakers, F.", - title = "{ROOT: An object oriented data analysis framework}", - booktitle = "{New computing techniques in physics research V. - Proceedings, 5th International Workshop, AIHENP '96, - Lausanne, Switzerland, September 2-6, 1996}", - journal = "Nucl. Instrum. Meth.", - volume = "A389", - year = "1997", - pages = "81-86", - doi = "10.1016/S0168-9002(97)00048-X", - notes = "{See also \url{http://root.cern.ch/}}", - SLACcitation = "%%CITATION = NUIMA,A389,81;%%" -} - - -@article{Cowan:2016tnm, - author = "Cowan, G. A. and Craik, D. C. and Needham, M. D.", - title = "{RapidSim: an application for the fast simulation of - heavy-quark hadron decays}", - journal = "Comput. Phys. Commun.", - volume = "214", - year = "2017", - pages = "239-246", - doi = "10.1016/j.cpc.2017.01.029", - eprint = "1612.07489", - archivePrefix = "arXiv", - primaryClass = "hep-ex", - SLACcitation = "%%CITATION = ARXIV:1612.07489;%%" -} - -@article{James:1968gu, - author = "James, F.", - title = "{Monte-Carlo phase space}", - year = "1968", - reportNumber = "CERN-68-15", - SLACcitation = "%%CITATION = CERN-68-15;%%" -} - -@article{Poluektov:2266468, - author = "Poluektov, Anton", - title = "{Performing amplitude fits with TensorFlow: LHCb - experience. HEP analysis ecosystem workshop}", - month = "May", - year = "2017", - reportNumber = "LHCb-TALK-2017-134", - url = "https://cds.cern.ch/record/2266468", -} - -@misc{tensorflow2015-whitepaper, -title={ {TensorFlow}: Large-Scale Machine Learning on Heterogeneous Systems}, -url={https://www.tensorflow.org/}, -note={Software available from tensorflow.org}, -author={ - Mart\'{\i}n~Abadi and - Ashish~Agarwal and - Paul~Barham and - Eugene~Brevdo and - Zhifeng~Chen and - Craig~Citro and - Greg~S.~Corrado and - Andy~Davis and - Jeffrey~Dean and - Matthieu~Devin and - Sanjay~Ghemawat and - Ian~Goodfellow and - Andrew~Harp and - Geoffrey~Irving and - Michael~Isard and - Yangqing Jia and - Rafal~Jozefowicz and - Lukasz~Kaiser and - Manjunath~Kudlur and - Josh~Levenberg and - Dandelion~Man\'{e} and - Rajat~Monga and - Sherry~Moore and - Derek~Murray and - Chris~Olah and - Mike~Schuster and - Jonathon~Shlens and - Benoit~Steiner and - Ilya~Sutskever and - Kunal~Talwar and - Paul~Tucker and - Vincent~Vanhoucke and - Vijay~Vasudevan and - Fernanda~Vi\'{e}gas and - Oriol~Vinyals and - Pete~Warden and - Martin~Wattenberg and - Martin~Wicke and - Yuan~Yu and - Xiaoqiang~Zheng}, - year={2015}, -} - - - -@misc{zfit, - Author = {Jonas Eschle and Albert {Puig Navarro} and Rafael {Silva Coutinho}}, - Date-Added = {2019-03-12 14:59:41 +0100}, - Date-Modified = {2019-03-21 14:39:21 -0400}, - Doi = {10.5281/zenodo.2602043}, - Title = {{zfit: scalable pythonic fitting}}, - Year = {2019}} - - -@article{zenodo, - title={phasespace: n-body phase space generation in Python}, - DOI={10.5281/zenodo.2591993}, - publisher={Zenodo}, - author={Albert {Puig Navarro} and Jonas Eschle}, - year={2019}, - month={Mar}} +@article{Brun:1997pa, + author = "Brun, R. and Rademakers, F.", + title = "{ROOT: An object oriented data analysis framework}", + booktitle = "{New computing techniques in physics research V. + Proceedings, 5th International Workshop, AIHENP '96, + Lausanne, Switzerland, September 2-6, 1996}", + journal = "Nucl. Instrum. Meth.", + volume = "A389", + year = "1997", + pages = "81-86", + doi = "10.1016/S0168-9002(97)00048-X", + notes = "{See also \url{http://root.cern.ch/}}", + SLACcitation = "%%CITATION = NUIMA,A389,81;%%" +} + + +@article{Cowan:2016tnm, + author = "Cowan, G. A. and Craik, D. C. and Needham, M. D.", + title = "{RapidSim: an application for the fast simulation of + heavy-quark hadron decays}", + journal = "Comput. Phys. Commun.", + volume = "214", + year = "2017", + pages = "239-246", + doi = "10.1016/j.cpc.2017.01.029", + eprint = "1612.07489", + archivePrefix = "arXiv", + primaryClass = "hep-ex", + SLACcitation = "%%CITATION = ARXIV:1612.07489;%%" +} + +@article{James:1968gu, + author = "James, F.", + title = "{Monte-Carlo phase space}", + year = "1968", + reportNumber = "CERN-68-15", + SLACcitation = "%%CITATION = CERN-68-15;%%" +} + +@article{Poluektov:2266468, + author = "Poluektov, Anton", + title = "{Performing amplitude fits with TensorFlow: LHCb + experience. HEP analysis ecosystem workshop}", + month = "May", + year = "2017", + reportNumber = "LHCb-TALK-2017-134", + url = "https://cds.cern.ch/record/2266468", +} + +@misc{tensorflow2015-whitepaper, +title={ {TensorFlow}: Large-Scale Machine Learning on Heterogeneous Systems}, +url={https://www.tensorflow.org/}, +note={Software available from tensorflow.org}, +author={ + Mart\'{\i}n~Abadi and + Ashish~Agarwal and + Paul~Barham and + Eugene~Brevdo and + Zhifeng~Chen and + Craig~Citro and + Greg~S.~Corrado and + Andy~Davis and + Jeffrey~Dean and + Matthieu~Devin and + Sanjay~Ghemawat and + Ian~Goodfellow and + Andrew~Harp and + Geoffrey~Irving and + Michael~Isard and + Yangqing Jia and + Rafal~Jozefowicz and + Lukasz~Kaiser and + Manjunath~Kudlur and + Josh~Levenberg and + Dandelion~Man\'{e} and + Rajat~Monga and + Sherry~Moore and + Derek~Murray and + Chris~Olah and + Mike~Schuster and + Jonathon~Shlens and + Benoit~Steiner and + Ilya~Sutskever and + Kunal~Talwar and + Paul~Tucker and + Vincent~Vanhoucke and + Vijay~Vasudevan and + Fernanda~Vi\'{e}gas and + Oriol~Vinyals and + Pete~Warden and + Martin~Wattenberg and + Martin~Wicke and + Yuan~Yu and + Xiaoqiang~Zheng}, + year={2015}, +} + + + +@misc{zfit, + Author = {Jonas Eschle and Albert {Puig Navarro} and Rafael {Silva Coutinho}}, + Date-Added = {2019-03-12 14:59:41 +0100}, + Date-Modified = {2019-03-21 14:39:21 -0400}, + Doi = {10.5281/zenodo.2602043}, + Title = {{zfit: scalable pythonic fitting}}, + Year = {2019}} + + +@article{zenodo, + title={phasespace: n-body phase space generation in Python}, + DOI={10.5281/zenodo.2591993}, + publisher={Zenodo}, + author={Albert {Puig Navarro} and Jonas Eschle}, + year={2019}, + month={Mar}} diff --git a/phasespace/__init__.py b/phasespace/__init__.py index 5132c914..5c190557 100644 --- a/phasespace/__init__.py +++ b/phasespace/__init__.py @@ -1,26 +1,26 @@ -"""Top-level package for TensorFlow PhaseSpace.""" -from pkg_resources import get_distribution - -__author__ = """Albert Puig Navarro""" -__email__ = "apuignav@gmail.com" -__version__ = get_distribution(__name__).version -__maintainer__ = "zfit" - -__credits__ = ["Jonas Eschle "] - -__all__ = ["nbody_decay", "GenParticle", "random"] - -import tensorflow as tf - -from . import random -from .phasespace import GenParticle, nbody_decay - - -def _set_eager_mode(): - import os - - is_eager = bool(os.environ.get("PHASESPACE_EAGER")) - tf.config.run_functions_eagerly(is_eager) - - -_set_eager_mode() +"""Top-level package for TensorFlow PhaseSpace.""" +from pkg_resources import get_distribution + +__author__ = """Albert Puig Navarro""" +__email__ = "apuignav@gmail.com" +__version__ = get_distribution(__name__).version +__maintainer__ = "zfit" + +__credits__ = ["Jonas Eschle "] + +__all__ = ["nbody_decay", "GenParticle", "random"] + +import tensorflow as tf + +from . import random +from .phasespace import GenParticle, nbody_decay + + +def _set_eager_mode(): + import os + + is_eager = bool(os.environ.get("PHASESPACE_EAGER")) + tf.config.run_functions_eagerly(is_eager) + + +_set_eager_mode() diff --git a/phasespace/backend.py b/phasespace/backend.py index 25e9498b..afb55167 100644 --- a/phasespace/backend.py +++ b/phasespace/backend.py @@ -1,21 +1,21 @@ -import tensorflow as tf - -RELAX_SHAPES = True -if int(tf.__version__.split(".")[1]) < 5: # smaller than 2.5 - jit_compile_argname = "experimental_compile" -else: - jit_compile_argname = "jit_compile" -function = tf.function( - autograph=False, - experimental_relax_shapes=RELAX_SHAPES, - **{jit_compile_argname: False} -) -function_jit = tf.function( - autograph=False, - experimental_relax_shapes=RELAX_SHAPES, - **{jit_compile_argname: True} -) - -function_jit_fixedshape = tf.function( - autograph=False, experimental_relax_shapes=False, **{jit_compile_argname: True} -) +import tensorflow as tf + +RELAX_SHAPES = True +if int(tf.__version__.split(".")[1]) < 5: # smaller than 2.5 + jit_compile_argname = "experimental_compile" +else: + jit_compile_argname = "jit_compile" +function = tf.function( + autograph=False, + experimental_relax_shapes=RELAX_SHAPES, + **{jit_compile_argname: False} +) +function_jit = tf.function( + autograph=False, + experimental_relax_shapes=RELAX_SHAPES, + **{jit_compile_argname: True} +) + +function_jit_fixedshape = tf.function( + autograph=False, experimental_relax_shapes=False, **{jit_compile_argname: True} +) diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py new file mode 100644 index 00000000..38788a14 --- /dev/null +++ b/phasespace/fromdecay/__init__.py @@ -0,0 +1,14 @@ +import sys + +from .fulldecay import FullDecay # noqa: F401 + +try: + import zfit # noqa: F401 + import zfit_physics as zphys # noqa: F401 + from particle import Particle # noqa: F401 +except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "The fromdecay functionality in phasespace requires particle and zfit-physics. " + "Either install phasespace[fromdecay] or particle and zfit-physics.", + file=sys.stderr, + ) from error diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fromdecay/fulldecay.py similarity index 90% rename from phasespace/fulldecay/fulldecay.py rename to phasespace/fromdecay/fulldecay.py index 76e755af..2334d8be 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fromdecay/fulldecay.py @@ -41,7 +41,8 @@ def from_dict( These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. - tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. + tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass + to be constant. Returns: The created FullDecay object. @@ -49,7 +50,7 @@ def from_dict( if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER else: - # Combine the mass functions specified by the package to the mass functions specified from the input. + # Combine the default mass functions with the mass functions specified from the input. total_mass_converter = {**_DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( @@ -67,15 +68,19 @@ def generate( Args: n_events: Total number of events combined, for all the decays. - normalize_weights: Normalize weights according to all events generated. This also changes the return values. + normalize_weights: Normalize weights according to all events generated. + This also changes the return values. See the phasespace documentation for more details. kwargs: Additional parameters passed to all calls of GenParticle.generate Returns: - The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. - However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding - to the return arguments from the corresponding GenParticle instances in self.gen_particles. - Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. + The arguments returned by GenParticle.generate are returned. + See the phasespace documentation for details. + However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists + corresponding to the return arguments from the corresponding GenParticle instances in + self.gen_particles. + Note that when normalize_weights is True, the weights are normalized to the maximum of all + returned events. """ # Input to tf.random.categorical must be 2D rand_i = tf.random.categorical( @@ -109,8 +114,8 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: preexisting_particles: Names that the particle cannot have as name. Returns: - name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i begins at 0 - and increases until the name is not preexisting_particles. + name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i + begins at 0 and increases until the name is not preexisting_particles. """ if name not in preexisting_particles: preexisting_particles.add(name) diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py similarity index 100% rename from phasespace/fulldecay/mass_functions.py rename to phasespace/fromdecay/mass_functions.py diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py deleted file mode 100644 index 42f877ee..00000000 --- a/phasespace/fulldecay/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -import sys - -from .fulldecay import FullDecay - -try: - import zfit - import zfit_physics as zphys - from particle import Particle -except ModuleNotFoundError as error: - raise ModuleNotFoundError( - "The fulldecay functionality in phasespace requires particle and zfit-physics. " - "Either install phasespace[fulldecay] or particle and zfit-physics.", - file=sys.stderr, - ) from error diff --git a/phasespace/kinematics.py b/phasespace/kinematics.py index 460c127e..317651dd 100644 --- a/phasespace/kinematics.py +++ b/phasespace/kinematics.py @@ -1,152 +1,152 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file kinematics.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 12.02.2019 -# ============================================================================= -"""Basic kinematics.""" - -import tensorflow.experimental.numpy as tnp - -from phasespace.backend import function, function_jit - - -@function_jit -def scalar_product(vec1, vec2): - """Calculate scalar product of two 3-vectors. - - Args: - vec1: First vector. - vec2: Second vector. - """ - return tnp.sum(vec1 * vec2, axis=1) - - -@function_jit -def spatial_component(vector): - """Extract spatial components of the input Lorentz vector. - - Args: - vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). - """ - return tnp.take(vector, indices=[0, 1, 2], axis=-1) - - -@function_jit -def time_component(vector): - """Extract time component of the input Lorentz vector. - - Args: - vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). - """ - return tnp.take(vector, indices=[3], axis=-1) - - -@function -def x_component(vector): - """Extract spatial X component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[0], axis=-1) - - -@function_jit -def y_component(vector): - """Extract spatial Y component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[1], axis=-1) - - -@function_jit -def z_component(vector): - """Extract spatial Z component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[2], axis=-1) - - -@function_jit -def mass(vector): - """Calculate mass scalar for Lorentz 4-momentum. - - Args: - vector: Input Lorentz momentum vector. - """ - return tnp.sqrt( - tnp.sum(tnp.square(vector) * metric_tensor(), axis=-1, keepdims=True) - ) - - -@function_jit -def lorentz_vector(space, time): - """Make a Lorentz vector from spatial and time components. - - Args: - space: 3-vector of spatial components. - time: Time component. - """ - return tnp.concatenate([space, time], axis=-1) - - -@function_jit -def lorentz_boost(vector, boostvector): - """Perform Lorentz boost. - - Args: - vector: 4-vector to be boosted - boostvector: Boost vector. Can be either 3-vector or 4-vector, since - only spatial components are used. - """ - boost = spatial_component(boostvector) - b2 = tnp.expand_dims(scalar_product(boost, boost), axis=-1) - - def boost_fn(): - gamma = 1.0 / tnp.sqrt(1.0 - b2) - gamma2 = (gamma - 1.0) / b2 - ve = time_component(vector) - vp = spatial_component(vector) - bp = tnp.expand_dims(scalar_product(vp, boost), axis=-1) - vp2 = vp + (gamma2 * bp + gamma * ve) * boost - ve2 = gamma * (ve + bp) - return lorentz_vector(vp2, ve2) - - # if boost vector is zero, return the original vector - all_b2_zero = tnp.all(tnp.equal(b2, tnp.zeros_like(b2))) - boosted_vector = tnp.where(all_b2_zero, vector, boost_fn()) - return boosted_vector - - -@function_jit -def beta(vector): - """Calculate beta of a given 4-vector. - - Args: - vector: Input Lorentz momentum vector. - """ - return mass(vector) / time_component(vector) - - -@function_jit -def boost_components(vector): - """Get the boost components of a given 4-vector. - - Args: - vector: Input Lorentz momentum vector. - """ - return spatial_component(vector) / time_component(vector) - - -@function_jit -def metric_tensor(): - """Metric tensor for Lorentz space (constant).""" - return tnp.asarray([-1.0, -1.0, -1.0, 1.0], dtype=tnp.float64) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file kinematics.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 12.02.2019 +# ============================================================================= +"""Basic kinematics.""" + +import tensorflow.experimental.numpy as tnp + +from phasespace.backend import function, function_jit + + +@function_jit +def scalar_product(vec1, vec2): + """Calculate scalar product of two 3-vectors. + + Args: + vec1: First vector. + vec2: Second vector. + """ + return tnp.sum(vec1 * vec2, axis=1) + + +@function_jit +def spatial_component(vector): + """Extract spatial components of the input Lorentz vector. + + Args: + vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). + """ + return tnp.take(vector, indices=[0, 1, 2], axis=-1) + + +@function_jit +def time_component(vector): + """Extract time component of the input Lorentz vector. + + Args: + vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). + """ + return tnp.take(vector, indices=[3], axis=-1) + + +@function +def x_component(vector): + """Extract spatial X component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[0], axis=-1) + + +@function_jit +def y_component(vector): + """Extract spatial Y component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[1], axis=-1) + + +@function_jit +def z_component(vector): + """Extract spatial Z component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[2], axis=-1) + + +@function_jit +def mass(vector): + """Calculate mass scalar for Lorentz 4-momentum. + + Args: + vector: Input Lorentz momentum vector. + """ + return tnp.sqrt( + tnp.sum(tnp.square(vector) * metric_tensor(), axis=-1, keepdims=True) + ) + + +@function_jit +def lorentz_vector(space, time): + """Make a Lorentz vector from spatial and time components. + + Args: + space: 3-vector of spatial components. + time: Time component. + """ + return tnp.concatenate([space, time], axis=-1) + + +@function_jit +def lorentz_boost(vector, boostvector): + """Perform Lorentz boost. + + Args: + vector: 4-vector to be boosted + boostvector: Boost vector. Can be either 3-vector or 4-vector, since + only spatial components are used. + """ + boost = spatial_component(boostvector) + b2 = tnp.expand_dims(scalar_product(boost, boost), axis=-1) + + def boost_fn(): + gamma = 1.0 / tnp.sqrt(1.0 - b2) + gamma2 = (gamma - 1.0) / b2 + ve = time_component(vector) + vp = spatial_component(vector) + bp = tnp.expand_dims(scalar_product(vp, boost), axis=-1) + vp2 = vp + (gamma2 * bp + gamma * ve) * boost + ve2 = gamma * (ve + bp) + return lorentz_vector(vp2, ve2) + + # if boost vector is zero, return the original vector + all_b2_zero = tnp.all(tnp.equal(b2, tnp.zeros_like(b2))) + boosted_vector = tnp.where(all_b2_zero, vector, boost_fn()) + return boosted_vector + + +@function_jit +def beta(vector): + """Calculate beta of a given 4-vector. + + Args: + vector: Input Lorentz momentum vector. + """ + return mass(vector) / time_component(vector) + + +@function_jit +def boost_components(vector): + """Get the boost components of a given 4-vector. + + Args: + vector: Input Lorentz momentum vector. + """ + return spatial_component(vector) / time_component(vector) + + +@function_jit +def metric_tensor(): + """Metric tensor for Lorentz space (constant).""" + return tnp.asarray([-1.0, -1.0, -1.0, 1.0], dtype=tnp.float64) + + +# EOF diff --git a/phasespace/phasespace.py b/phasespace/phasespace.py index b5b821b9..70f09251 100644 --- a/phasespace/phasespace.py +++ b/phasespace/phasespace.py @@ -1,776 +1,776 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file phasespace.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 25.02.2019 -# ============================================================================= -"""Implementation of the Raubold and Lynch method to generate n-body events. - -The code is based on the GENBOD function (W515 from CERNLIB), documented in: - - F. James, Monte Carlo Phase Space, CERN 68-15 (1968) -""" -import inspect -import warnings -from math import pi -from typing import Callable, Dict, Optional, Tuple, Union - -import tensorflow as tf -import tensorflow.experimental.numpy as tnp - -from . import kinematics as kin -from .backend import function, function_jit_fixedshape -from .random import SeedLike, get_rng - -RELAX_SHAPES = False - - -def process_list_to_tensor(lst): - """Convert a list to a tensor. - - The list is converted to a tensor and transposed to get the proper shape. - - Note: - If `lst` is a tensor, nothing is done to it other than convert it to `tf.float64`. - - Arguments: - lst (list): List to convert. - - Return: - ~`tf.Tensor` - """ - if isinstance(lst, list): - lst = tnp.transpose(tnp.asarray(lst, dtype=tnp.float64)) - return tnp.asarray(lst, dtype=tnp.float64) - - -@function -def pdk(a, b, c): - """Calculate the PDK (2-body phase space) function. - - Based on Eq. (9.17) in CERN 68-15 (1968). - - Arguments: - a (~`tf.Tensor`): :math:`M_{i+1}` in Eq. (9.17). - b (~`tf.Tensor`): :math:`M_{i}` in Eq. (9.17). - c (~`tf.Tensor`): :math:`m_{i+1}` in Eq. (9.17). - - Return: - ~`tf.Tensor` - """ - x = (a - b - c) * (a + b + c) * (a - b + c) * (a + b - c) - return tnp.sqrt(x) / (tnp.asarray(2.0, dtype=tnp.float64) * a) - - -class GenParticle: - """Representation of a particle. - - Instances of this class can be combined with each other to build decay chains, - which can then be used to generate phase space events through the `generate` - or `generate_tensor` method. - - A `GenParticle` must have - + a `name`, which is ensured not to clash with any others in - the decay chain. - + a `mass`, which can be either a number or a function to generate it according to - a certain distribution. The returned ~`tf.Tensor` needs to have shape (nevents,). - In this case, the particle is not considered as having a - fixed mass and the `has_fixed_mass` method will return False. - - It may also have: - - + Children, ie, decay products, which are also `GenParticle` instances. - - - Arguments: - name (str): Name of the particle. - mass (float, ~`tf.Tensor`, callable): Mass of the particle. If it's a float, it get - converted to a `tf.constant`. - """ - - def __init__(self, name: str, mass: Union[Callable, int, float]) -> None: # noqa - self.name = name - self.children = [] - self._mass_val = mass - if not callable(mass) and not isinstance(mass, tf.Variable): - mass = tnp.asarray(mass, dtype=tnp.float64) - else: - mass = tf.function( - mass, autograph=False, experimental_relax_shapes=RELAX_SHAPES - ) - self._mass = mass - self._generate_called = False # not yet called, children can be set - - def __repr__(self): - return "".format( - self.name, - f"{self._mass_val:.2f}" if self.has_fixed_mass else "variable", - ", ".join(child.name for child in self.children), - ) - - def _do_names_clash(self, particles): - def get_list_of_names(part): - output = [part.name] - for child in part.children: - output.extend(get_list_of_names(child)) - return output - - names_to_check = [self.name] - for part in particles: - names_to_check.extend(get_list_of_names(part)) - # Find top - dup_names = {name for name in names_to_check if names_to_check.count(name) > 1} - if dup_names: - return dup_names - return None - - @function - def get_mass( - self, - min_mass: tf.Tensor = None, - max_mass: tf.Tensor = None, - n_events: Union[tf.Tensor, tf.Variable] = None, - seed: SeedLike = None, - ) -> tf.Tensor: - """Get the particle mass. - - If the particle is resonant, the mass function will be called with the - `min_mass`, `max_mass`, `n_events` and optionally a `seed` parameter. - - Arguments: - min_mass (tensor): Lower mass range. Defaults to None, which - is only valid in the case of fixed mass. - max_mass (tensor): Upper mass range. Defaults to None, which - is only valid in the case of fixed mass. - n_events (): Number of events to produce. Has to be specified if the particle is resonant. - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - ~`tf.Tensor`: Mass of the particles, either a scalar or shape (nevents,) - - Raise: - ValueError: If the mass is requested and has not been set. - """ - if self.has_fixed_mass: - mass = self._mass - else: - seed = get_rng(seed) - min_mass = tnp.reshape(min_mass, (n_events,)) - max_mass = tnp.reshape(max_mass, (n_events,)) - signature = inspect.signature(self._mass) - if "seed" in signature.parameters: - mass = self._mass(min_mass, max_mass, n_events, seed=seed) - else: - mass = self._mass(min_mass, max_mass, n_events) - return mass - - @property - def has_fixed_mass(self): - """bool: Is the mass a callable function?""" - return not callable(self._mass) - - def set_children(self, *children): - """Assign children. - - Arguments: - children (GenParticle): Two or more children to assign to the current particle. - - Return: - self - - Raise: - ValueError: If there is an inconsistency in the parent/children relationship, ie, - if children were already set, if their parent was or if less than two children were given. - KeyError: If there is a particle name clash. - RuntimeError: If `generate` was already called before. - """ - # self._set_cache_validity(False) - if self._generate_called: - raise RuntimeError( - "Cannot set children after the first call to `generate`." - ) - if self.children: - raise ValueError("Children already set!") - if len(children) <= 1: - raise ValueError( - f"Have to set at least 2 children, not {len(children)} for a particle to decay" - ) - # Check name clashes - name_clash = self._do_names_clash(children) - if name_clash: - raise KeyError(f"Particle name {name_clash} already used") - self.children = children - return self - - @property - def has_children(self): - """bool: Does the particle have children?""" - return bool(self.children) - - @property - def has_grandchildren(self): - """bool: Does the particle have grandchildren?""" - if not self.children: - return False - return any(child.has_children for child in self.children) - - @staticmethod - def _preprocess(momentum, n_events): - """Preprocess momentum input and determine number of events to generate. - - Both `momentum` and `n_events` are converted to tensors if they - are not already. - - Arguments: - `momentum`: Momentum vector, of shape (4, x), where x is optional. - `n_events`: Number of events to generate. If `None`, the number of events - to generate is calculated from the shape of `momentum`. - - Return: - tuple: Processed `momentum` and `n_events`. - - Raise: - tf.errors.InvalidArgumentError: If the number of events deduced from the - shape of `momentum` is inconsistent with `n_events`. - """ - momentum = process_list_to_tensor(momentum) - - # Check sanity of inputs - if len(momentum.shape) not in (1, 2): - raise ValueError(f"Bad shape for momentum -> {list(momentum.shape)}") - # Check compatibility of inputs - if len(momentum.shape) == 2: - if n_events is not None: - momentum_shape = momentum.shape[0] - if momentum_shape is None: - momentum_shape = tf.shape(momentum)[0] - momentum_shape = tnp.asarray(momentum_shape, tnp.int64) - else: - momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) - tf.assert_equal( - n_events, - momentum_shape, - message="Conflicting inputs -> momentum_shape and n_events", - ) - - if n_events is None: - if len(momentum.shape) == 2: - n_events = momentum.shape[0] - if n_events is None: # dynamic shape - n_events = tf.shape(momentum)[0] - n_events = tnp.asarray(n_events, dtype=tnp.int64) - else: - n_events = tnp.asarray(1, dtype=tnp.int64) - n_events = tnp.asarray(n_events, dtype=tnp.int64) - # Now preparation of tensors - if len(momentum.shape) == 1: - momentum = tnp.expand_dims(momentum, axis=0) - return momentum, n_events - - @staticmethod - @function_jit_fixedshape - def _get_w_max(available_mass, masses): - emmax = available_mass + tnp.take(masses, indices=[0], axis=1) - emmin = tnp.zeros_like(emmax, dtype=tnp.float64) - w_max = tnp.ones_like(emmax, dtype=tnp.float64) - for i in range(1, masses.shape[1]): - emmin += tnp.take(masses, [i - 1], axis=1) - emmax += tnp.take(masses, [i], axis=1) - w_max *= pdk(emmax, emmin, tnp.take(masses, [i], axis=1)) - return w_max - - def _generate(self, momentum, n_events, rng): - """Generate a n-body decay according to the Raubold and Lynch method. - - The number and mass of the children particles are taken from self.children. - - Note: - This method generates the same results as the GENBOD routine. - - Arguments: - momentum (tensor): Momentum of the parent particle. All generated particles - will be boosted to that momentum. - n_events (int): Number of events to generate. - - Return: - tuple: Result of the generation (per-event weights, maximum weights, output particles - and their output masses). - """ - self._generate_called = True - if not self.children: - raise ValueError("No children have been configured") - p_top, n_events = self._preprocess(momentum, n_events) - top_mass = tnp.broadcast_to(kin.mass(p_top), (n_events, 1)) - n_particles = len(self.children) - - # Prepare masses - def recurse_stable(part): - output_mass = 0 - for child in part.children: - if child.has_fixed_mass: - output_mass += child.get_mass() - else: - output_mass += recurse_stable(child) - return output_mass - - mass_from_stable = tnp.broadcast_to( - tnp.sum( - [child.get_mass() for child in self.children if child.has_fixed_mass], - axis=0, - ), - (n_events, 1), - ) - max_mass = top_mass - mass_from_stable - masses = [] - for child in self.children: - if child.has_fixed_mass: - masses.append(tnp.broadcast_to(child.get_mass(), (n_events, 1))) - else: - # Recurse that particle to know the minimum mass we need to generate - min_mass = tnp.broadcast_to(recurse_stable(child), (n_events, 1)) - mass = child.get_mass(min_mass, max_mass, n_events) - mass = tnp.reshape(mass, (n_events, 1)) - max_mass -= mass - masses.append(mass) - masses = tnp.concatenate(masses, axis=-1) - # if len(masses.shape) == 1: - # masses = tnp.expand_dims(masses, axis=0) - available_mass = top_mass - tnp.sum(masses, axis=1, keepdims=True) - tf.debugging.assert_greater_equal( - available_mass, - tnp.zeros_like(available_mass, dtype=tnp.float64), - message="Forbidden decay", - name="mass_check", - ) # Calculate the max weight, initial beta, etc - w_max = self._get_w_max(available_mass, masses) - p_top_boost = kin.boost_components(p_top) - # Start the generation - random_numbers = rng.uniform((n_events, n_particles - 2), dtype=tnp.float64) - random = tnp.concatenate( - [ - tnp.zeros((n_events, 1), dtype=tnp.float64), - tnp.sort(random_numbers, axis=1), - tnp.ones((n_events, 1), dtype=tnp.float64), - ], - axis=1, - ) - if random.shape[1] is None: - random.set_shape((None, n_particles)) - # random = tnp.expand_dims(random, axis=-1) - sum_ = tnp.zeros((n_events, 1), dtype=tnp.float64) - inv_masses = [] - # TODO(Mayou36): rewrite with cumsum? - for i in range(n_particles): - sum_ += tnp.take(masses, [i], axis=1) - inv_masses.append(tnp.take(random, [i], axis=1) * available_mass + sum_) - generated_particles, weights = self._generate_part2( - inv_masses, masses, n_events, n_particles, rng=rng - ) - # Final boost of all particles - generated_particles = [ - kin.lorentz_boost(part, p_top_boost) for part in generated_particles - ] - return ( - tnp.reshape(weights, (n_events,)), - tnp.reshape(w_max, (n_events,)), - generated_particles, - masses, - ) - - @staticmethod - @function - def _generate_part2(inv_masses, masses, n_events, n_particles, rng): - pds = [] - # Calculate weights of the events - for i in range(n_particles - 1): - pds.append( - pdk( - inv_masses[i + 1], - inv_masses[i], - tnp.take(masses, [i + 1], axis=1), - ) - ) - weights = tnp.prod(pds, axis=0) - zero_component = tnp.zeros_like(pds[0], dtype=tnp.float64) - generated_particles = [ - tnp.concatenate( - [ - zero_component, - pds[0], - zero_component, - tnp.sqrt( - tnp.square(pds[0]) + tnp.square(tnp.take(masses, [0], axis=1)) - ), - ], - axis=1, - ) - ] - part_num = 1 - while True: - generated_particles.append( - tnp.concatenate( - [ - zero_component, - -pds[part_num - 1], - zero_component, - tnp.sqrt( - tnp.square(pds[part_num - 1]) - + tnp.square(tnp.take(masses, [part_num], axis=1)) - ), - ], - axis=1, - ) - ) - - cos_z = tnp.asarray(2.0, dtype=tnp.float64) * rng.uniform( - (n_events, 1), dtype=tnp.float64 - ) - tnp.asarray(1.0, dtype=tnp.float64) - sin_z = tnp.sqrt(tnp.asarray(1.0, dtype=tnp.float64) - cos_z * cos_z) - ang_y = ( - tnp.asarray(2.0, dtype=tnp.float64) - * tnp.asarray(pi, dtype=tnp.float64) - * rng.uniform((n_events, 1), dtype=tnp.float64) - ) - cos_y = tnp.cos(ang_y) - sin_y = tnp.sin(ang_y) - # Do the rotations - for j in range(part_num + 1): - px = kin.x_component(generated_particles[j]) - py = kin.y_component(generated_particles[j]) - # Rotate about z - # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it? - generated_particles[j] = tnp.concatenate( - [ - cos_z * px - sin_z * py, - sin_z * px + cos_z * py, - kin.z_component(generated_particles[j]), - kin.time_component(generated_particles[j]), - ], - axis=1, - ) - # Rotate about y - px = kin.x_component(generated_particles[j]) - pz = kin.z_component(generated_particles[j]) - generated_particles[j] = tnp.concatenate( - [ - cos_y * px - sin_y * pz, - kin.y_component(generated_particles[j]), - sin_y * px + cos_y * pz, - kin.time_component(generated_particles[j]), - ], - axis=1, - ) - if part_num == (n_particles - 1): - break - betas = pds[part_num] / tnp.sqrt( - tnp.square(pds[part_num]) + tnp.square(inv_masses[part_num]) - ) - generated_particles = [ - kin.lorentz_boost( - part, - tnp.concatenate([zero_component, betas, zero_component], axis=1), - ) - for part in generated_particles - ] - part_num += 1 - return generated_particles, weights - - @function - def _recursive_generate( - self, - n_events, - boost_to=None, - recalculate_max_weights=False, - rng: SeedLike = None, - ): - """Recursively generate normalized n-body phase space as tensorflow tensors. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (tensor, optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - recalculate_max_weights (bool, optional): Recalculate the maximum weight of the event - using all the particles of the tree? This is necessary for the top particle of a decay, - otherwise the maximum weight calculation is going to be wrong (particles from subdecays - would not be taken into account). Defaults to False. - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - tuple: Result of the generation (per-event weights, maximum weights, output particles - and their output masses). - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - if boost_to is not None: - momentum = boost_to - else: - if self.has_fixed_mass: - momentum = tnp.broadcast_to( - tnp.stack((0.0, 0.0, 0.0, self.get_mass()), axis=-1), (n_events, 4) - ) - else: - raise ValueError("Cannot use resonance as top particle") - weights, weights_max, parts, children_masses = self._generate( - momentum, n_events, rng=rng - ) - output_particles = { - child.name: parts[child_num] - for child_num, child in enumerate(self.children) - } - output_masses = { - child.name: tnp.take(children_masses, [child_num], axis=1) - for child_num, child in enumerate(self.children) - } - for child_num, child in enumerate(self.children): - if child.has_children: - ( - child_weights, - _, - child_gen_particles, - child_masses, - ) = child._recursive_generate( - n_events=n_events, - boost_to=parts[child_num], - recalculate_max_weights=False, - rng=rng, - ) - weights *= child_weights - output_particles.update(child_gen_particles) - output_masses.update(child_masses) - if recalculate_max_weights: - - def build_mass_tree(particle, leaf): - if particle.has_children: - leaf[particle.name] = {} - for child in particle.children: - build_mass_tree(child, leaf[particle.name]) - else: - leaf[particle.name] = output_masses[particle.name] - - def get_flattened_values(dict_): - output = [] - for val in dict_.values(): - if isinstance(val, dict): - output.extend(get_flattened_values(val)) - else: - output.append(val) - return output - - def recurse_w_max(parent_mass, current_mass_tree): - available_mass = parent_mass - sum( - get_flattened_values(current_mass_tree) - ) - masses = [] - w_max = tnp.ones_like(available_mass) - for child, child_mass in current_mass_tree.items(): - if isinstance(child_mass, dict): - w_max *= recurse_w_max( - parent_mass - - sum( - get_flattened_values( - { - ch_it: ch_m_it - for ch_it, ch_m_it in current_mass_tree.items() - if ch_it != child - } - ) - ), - child_mass, - ) - masses.append(sum(get_flattened_values(child_mass))) - else: - masses.append(child_mass) - masses = tnp.concatenate(masses, axis=1) - w_max *= self._get_w_max(available_mass, masses) - return w_max - - mass_tree = {} - build_mass_tree(self, mass_tree) - momentum = process_list_to_tensor(momentum) - if len(momentum.shape) == 1: - momentum = tnp.expand_dims(momentum, axis=-1) - weights_max = tnp.reshape( - recurse_w_max(kin.mass(momentum), mass_tree[self.name]), (n_events,) - ) - return weights, weights_max, output_particles, output_masses - - def generate( - self, - n_events: Union[int, tf.Tensor, tf.Variable], - boost_to: Optional[tf.Tensor] = None, - normalize_weights: bool = True, - seed: SeedLike = None, - ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]: - """Generate normalized n-body phase space as tensorflow tensors. - - Any TensorFlow tensor can always be converted to a numpy array with the method `numpy()`. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - normalize_weights (bool, optional): Normalize the event weight to its max? - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - tuple: Result of the generation, which varies with the value of `normalize_weights`: - - + If True, the tuple elements are the normalized event weights as a tensor of shape - (n_events, ), and the momenta generated particles as a dictionary of tensors of shape - (4, n_events) with particle names as keys. - - + If False, the tuple weights are the unnormalized event weights as a tensor of shape - (n_events, ), the maximum per-event weights as a tensor of shape (n_events, ) and the - momenta generated particles as a dictionary of tensors of shape (4, n_events) with particle - names as keys. - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - rng = get_rng(seed) - if boost_to is not None: - message = ( - f"The number of events requested ({n_events}) doesn't match the boost_to input size " - f"of {boost_to.shape}" - ) - tf.assert_equal(tf.shape(boost_to)[0], tf.shape(n_events), message=message) - if not isinstance(n_events, tf.Variable): - n_events = tnp.asarray(n_events, dtype=tnp.int64) - weights, weights_max, parts, _ = self._recursive_generate( - n_events=n_events, - boost_to=boost_to, - recalculate_max_weights=self.has_grandchildren, - rng=rng, - ) - return ( - (weights / weights_max, parts) - if normalize_weights - else (weights, weights_max, parts) - ) - - def generate_tensor( - self, - n_events: int, - boost_to=None, - normalize_weights: bool = True, - ): - """Generate normalized n-body phase space as numpy arrays. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - normalize_weights (bool, optional): Normalize the event weight to its max - - - Return: - tuple: Result of the generation, which varies with the value of `normalize_weights`: - - + If True, the tuple elements are the normalized event weights as an array of shape - (n_events, ), and the momenta generated particles as a dictionary of arrays of shape - (4, n_events) with particle names as keys. - - + If False, the tuple weights are the unnormalized event weights as an array of shape - (n_events, ), the maximum per-event weights as an array of shape (n_events, ) and the - momenta generated particles as a dictionary of arrays of shape (4, n_events) with particle - names as keys. - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - - # Run generation - warnings.warn( - "This function is depreceated. Use `generate` which does not return a Tensor as well." - ) - generate_tf = self.generate(n_events, boost_to, normalize_weights) - # self._cache = generate_tf - # self._set_cache_validity(True, propagate=True) - return generate_tf - - -# legacy class to warn user about name change -class Particle: - """Deprecated Particle class. - - Renamed to GenParticle. - """ - - def __init__(self): - raise NameError( - "'Particle' has been renamed to 'GenParticle'. Please update your code accordingly." - "For more information, see: https://github.com/zfit/phasespace/issues/22" - ) - - -def nbody_decay(mass_top: float, masses: list, top_name: str = "", names: list = None): - """Shortcut to build an n-body decay of a GenParticle. - - If the particle names are not given, the top particle is called 'top' and the - children 'p_{i}', where i corresponds to their position in the `masses` sequence. - - Arguments: - mass_top (tensor, list): Mass of the top particle. Can be a list of 4-vectors. - masses (list): Masses of the child particles. - name_top (str, optional): Name of the top particle. If not given, the top particle is - named top. - names (list, optional): Names of the child particles. If not given, they are build as - 'p_{i}', where i is given by their ordering in the `masses` list. - - Return: - `GenParticle`: Particle decay. - - Raise: - ValueError: If the length of `masses` and `names` doesn't match. - """ - if not top_name: - top_name = "top" - if not names: - names = [f"p_{num}" for num in range(len(masses))] - if len(names) != len(masses): - raise ValueError("Mismatch in length between children masses and their names.") - return GenParticle(top_name, mass_top).set_children( - *(GenParticle(names[num], mass=mass) for num, mass in enumerate(masses)) - ) - - -def generate_decay(*args, **kwargs): - """Deprecated.""" - raise NameError( - "'generate_decay' has been removed. A similar behavior can be accomplished with 'nbody_decay'. " - "For more information see https://github.com/zfit/phasespace/issues/22" - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file phasespace.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 25.02.2019 +# ============================================================================= +"""Implementation of the Raubold and Lynch method to generate n-body events. + +The code is based on the GENBOD function (W515 from CERNLIB), documented in: + + F. James, Monte Carlo Phase Space, CERN 68-15 (1968) +""" +import inspect +import warnings +from math import pi +from typing import Callable, Dict, Optional, Tuple, Union + +import tensorflow as tf +import tensorflow.experimental.numpy as tnp + +from . import kinematics as kin +from .backend import function, function_jit_fixedshape +from .random import SeedLike, get_rng + +RELAX_SHAPES = False + + +def process_list_to_tensor(lst): + """Convert a list to a tensor. + + The list is converted to a tensor and transposed to get the proper shape. + + Note: + If `lst` is a tensor, nothing is done to it other than convert it to `tf.float64`. + + Arguments: + lst (list): List to convert. + + Return: + ~`tf.Tensor` + """ + if isinstance(lst, list): + lst = tnp.transpose(tnp.asarray(lst, dtype=tnp.float64)) + return tnp.asarray(lst, dtype=tnp.float64) + + +@function +def pdk(a, b, c): + """Calculate the PDK (2-body phase space) function. + + Based on Eq. (9.17) in CERN 68-15 (1968). + + Arguments: + a (~`tf.Tensor`): :math:`M_{i+1}` in Eq. (9.17). + b (~`tf.Tensor`): :math:`M_{i}` in Eq. (9.17). + c (~`tf.Tensor`): :math:`m_{i+1}` in Eq. (9.17). + + Return: + ~`tf.Tensor` + """ + x = (a - b - c) * (a + b + c) * (a - b + c) * (a + b - c) + return tnp.sqrt(x) / (tnp.asarray(2.0, dtype=tnp.float64) * a) + + +class GenParticle: + """Representation of a particle. + + Instances of this class can be combined with each other to build decay chains, + which can then be used to generate phase space events through the `generate` + or `generate_tensor` method. + + A `GenParticle` must have + + a `name`, which is ensured not to clash with any others in + the decay chain. + + a `mass`, which can be either a number or a function to generate it according to + a certain distribution. The returned ~`tf.Tensor` needs to have shape (nevents,). + In this case, the particle is not considered as having a + fixed mass and the `has_fixed_mass` method will return False. + + It may also have: + + + Children, ie, decay products, which are also `GenParticle` instances. + + + Arguments: + name (str): Name of the particle. + mass (float, ~`tf.Tensor`, callable): Mass of the particle. If it's a float, it get + converted to a `tf.constant`. + """ + + def __init__(self, name: str, mass: Union[Callable, int, float]) -> None: # noqa + self.name = name + self.children = [] + self._mass_val = mass + if not callable(mass) and not isinstance(mass, tf.Variable): + mass = tnp.asarray(mass, dtype=tnp.float64) + else: + mass = tf.function( + mass, autograph=False, experimental_relax_shapes=RELAX_SHAPES + ) + self._mass = mass + self._generate_called = False # not yet called, children can be set + + def __repr__(self): + return "".format( + self.name, + f"{self._mass_val:.2f}" if self.has_fixed_mass else "variable", + ", ".join(child.name for child in self.children), + ) + + def _do_names_clash(self, particles): + def get_list_of_names(part): + output = [part.name] + for child in part.children: + output.extend(get_list_of_names(child)) + return output + + names_to_check = [self.name] + for part in particles: + names_to_check.extend(get_list_of_names(part)) + # Find top + dup_names = {name for name in names_to_check if names_to_check.count(name) > 1} + if dup_names: + return dup_names + return None + + @function + def get_mass( + self, + min_mass: tf.Tensor = None, + max_mass: tf.Tensor = None, + n_events: Union[tf.Tensor, tf.Variable] = None, + seed: SeedLike = None, + ) -> tf.Tensor: + """Get the particle mass. + + If the particle is resonant, the mass function will be called with the + `min_mass`, `max_mass`, `n_events` and optionally a `seed` parameter. + + Arguments: + min_mass (tensor): Lower mass range. Defaults to None, which + is only valid in the case of fixed mass. + max_mass (tensor): Upper mass range. Defaults to None, which + is only valid in the case of fixed mass. + n_events (): Number of events to produce. Has to be specified if the particle is resonant. + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + ~`tf.Tensor`: Mass of the particles, either a scalar or shape (nevents,) + + Raise: + ValueError: If the mass is requested and has not been set. + """ + if self.has_fixed_mass: + mass = self._mass + else: + seed = get_rng(seed) + min_mass = tnp.reshape(min_mass, (n_events,)) + max_mass = tnp.reshape(max_mass, (n_events,)) + signature = inspect.signature(self._mass) + if "seed" in signature.parameters: + mass = self._mass(min_mass, max_mass, n_events, seed=seed) + else: + mass = self._mass(min_mass, max_mass, n_events) + return mass + + @property + def has_fixed_mass(self): + """bool: Is the mass a callable function?""" + return not callable(self._mass) + + def set_children(self, *children): + """Assign children. + + Arguments: + children (GenParticle): Two or more children to assign to the current particle. + + Return: + self + + Raise: + ValueError: If there is an inconsistency in the parent/children relationship, ie, + if children were already set, if their parent was or if less than two children were given. + KeyError: If there is a particle name clash. + RuntimeError: If `generate` was already called before. + """ + # self._set_cache_validity(False) + if self._generate_called: + raise RuntimeError( + "Cannot set children after the first call to `generate`." + ) + if self.children: + raise ValueError("Children already set!") + if len(children) <= 1: + raise ValueError( + f"Have to set at least 2 children, not {len(children)} for a particle to decay" + ) + # Check name clashes + name_clash = self._do_names_clash(children) + if name_clash: + raise KeyError(f"Particle name {name_clash} already used") + self.children = children + return self + + @property + def has_children(self): + """bool: Does the particle have children?""" + return bool(self.children) + + @property + def has_grandchildren(self): + """bool: Does the particle have grandchildren?""" + if not self.children: + return False + return any(child.has_children for child in self.children) + + @staticmethod + def _preprocess(momentum, n_events): + """Preprocess momentum input and determine number of events to generate. + + Both `momentum` and `n_events` are converted to tensors if they + are not already. + + Arguments: + `momentum`: Momentum vector, of shape (4, x), where x is optional. + `n_events`: Number of events to generate. If `None`, the number of events + to generate is calculated from the shape of `momentum`. + + Return: + tuple: Processed `momentum` and `n_events`. + + Raise: + tf.errors.InvalidArgumentError: If the number of events deduced from the + shape of `momentum` is inconsistent with `n_events`. + """ + momentum = process_list_to_tensor(momentum) + + # Check sanity of inputs + if len(momentum.shape) not in (1, 2): + raise ValueError(f"Bad shape for momentum -> {list(momentum.shape)}") + # Check compatibility of inputs + if len(momentum.shape) == 2: + if n_events is not None: + momentum_shape = momentum.shape[0] + if momentum_shape is None: + momentum_shape = tf.shape(momentum)[0] + momentum_shape = tnp.asarray(momentum_shape, tnp.int64) + else: + momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) + tf.assert_equal( + n_events, + momentum_shape, + message="Conflicting inputs -> momentum_shape and n_events", + ) + + if n_events is None: + if len(momentum.shape) == 2: + n_events = momentum.shape[0] + if n_events is None: # dynamic shape + n_events = tf.shape(momentum)[0] + n_events = tnp.asarray(n_events, dtype=tnp.int64) + else: + n_events = tnp.asarray(1, dtype=tnp.int64) + n_events = tnp.asarray(n_events, dtype=tnp.int64) + # Now preparation of tensors + if len(momentum.shape) == 1: + momentum = tnp.expand_dims(momentum, axis=0) + return momentum, n_events + + @staticmethod + @function_jit_fixedshape + def _get_w_max(available_mass, masses): + emmax = available_mass + tnp.take(masses, indices=[0], axis=1) + emmin = tnp.zeros_like(emmax, dtype=tnp.float64) + w_max = tnp.ones_like(emmax, dtype=tnp.float64) + for i in range(1, masses.shape[1]): + emmin += tnp.take(masses, [i - 1], axis=1) + emmax += tnp.take(masses, [i], axis=1) + w_max *= pdk(emmax, emmin, tnp.take(masses, [i], axis=1)) + return w_max + + def _generate(self, momentum, n_events, rng): + """Generate a n-body decay according to the Raubold and Lynch method. + + The number and mass of the children particles are taken from self.children. + + Note: + This method generates the same results as the GENBOD routine. + + Arguments: + momentum (tensor): Momentum of the parent particle. All generated particles + will be boosted to that momentum. + n_events (int): Number of events to generate. + + Return: + tuple: Result of the generation (per-event weights, maximum weights, output particles + and their output masses). + """ + self._generate_called = True + if not self.children: + raise ValueError("No children have been configured") + p_top, n_events = self._preprocess(momentum, n_events) + top_mass = tnp.broadcast_to(kin.mass(p_top), (n_events, 1)) + n_particles = len(self.children) + + # Prepare masses + def recurse_stable(part): + output_mass = 0 + for child in part.children: + if child.has_fixed_mass: + output_mass += child.get_mass() + else: + output_mass += recurse_stable(child) + return output_mass + + mass_from_stable = tnp.broadcast_to( + tnp.sum( + [child.get_mass() for child in self.children if child.has_fixed_mass], + axis=0, + ), + (n_events, 1), + ) + max_mass = top_mass - mass_from_stable + masses = [] + for child in self.children: + if child.has_fixed_mass: + masses.append(tnp.broadcast_to(child.get_mass(), (n_events, 1))) + else: + # Recurse that particle to know the minimum mass we need to generate + min_mass = tnp.broadcast_to(recurse_stable(child), (n_events, 1)) + mass = child.get_mass(min_mass, max_mass, n_events) + mass = tnp.reshape(mass, (n_events, 1)) + max_mass -= mass + masses.append(mass) + masses = tnp.concatenate(masses, axis=-1) + # if len(masses.shape) == 1: + # masses = tnp.expand_dims(masses, axis=0) + available_mass = top_mass - tnp.sum(masses, axis=1, keepdims=True) + tf.debugging.assert_greater_equal( + available_mass, + tnp.zeros_like(available_mass, dtype=tnp.float64), + message="Forbidden decay", + name="mass_check", + ) # Calculate the max weight, initial beta, etc + w_max = self._get_w_max(available_mass, masses) + p_top_boost = kin.boost_components(p_top) + # Start the generation + random_numbers = rng.uniform((n_events, n_particles - 2), dtype=tnp.float64) + random = tnp.concatenate( + [ + tnp.zeros((n_events, 1), dtype=tnp.float64), + tnp.sort(random_numbers, axis=1), + tnp.ones((n_events, 1), dtype=tnp.float64), + ], + axis=1, + ) + if random.shape[1] is None: + random.set_shape((None, n_particles)) + # random = tnp.expand_dims(random, axis=-1) + sum_ = tnp.zeros((n_events, 1), dtype=tnp.float64) + inv_masses = [] + # TODO(Mayou36): rewrite with cumsum? + for i in range(n_particles): + sum_ += tnp.take(masses, [i], axis=1) + inv_masses.append(tnp.take(random, [i], axis=1) * available_mass + sum_) + generated_particles, weights = self._generate_part2( + inv_masses, masses, n_events, n_particles, rng=rng + ) + # Final boost of all particles + generated_particles = [ + kin.lorentz_boost(part, p_top_boost) for part in generated_particles + ] + return ( + tnp.reshape(weights, (n_events,)), + tnp.reshape(w_max, (n_events,)), + generated_particles, + masses, + ) + + @staticmethod + @function + def _generate_part2(inv_masses, masses, n_events, n_particles, rng): + pds = [] + # Calculate weights of the events + for i in range(n_particles - 1): + pds.append( + pdk( + inv_masses[i + 1], + inv_masses[i], + tnp.take(masses, [i + 1], axis=1), + ) + ) + weights = tnp.prod(pds, axis=0) + zero_component = tnp.zeros_like(pds[0], dtype=tnp.float64) + generated_particles = [ + tnp.concatenate( + [ + zero_component, + pds[0], + zero_component, + tnp.sqrt( + tnp.square(pds[0]) + tnp.square(tnp.take(masses, [0], axis=1)) + ), + ], + axis=1, + ) + ] + part_num = 1 + while True: + generated_particles.append( + tnp.concatenate( + [ + zero_component, + -pds[part_num - 1], + zero_component, + tnp.sqrt( + tnp.square(pds[part_num - 1]) + + tnp.square(tnp.take(masses, [part_num], axis=1)) + ), + ], + axis=1, + ) + ) + + cos_z = tnp.asarray(2.0, dtype=tnp.float64) * rng.uniform( + (n_events, 1), dtype=tnp.float64 + ) - tnp.asarray(1.0, dtype=tnp.float64) + sin_z = tnp.sqrt(tnp.asarray(1.0, dtype=tnp.float64) - cos_z * cos_z) + ang_y = ( + tnp.asarray(2.0, dtype=tnp.float64) + * tnp.asarray(pi, dtype=tnp.float64) + * rng.uniform((n_events, 1), dtype=tnp.float64) + ) + cos_y = tnp.cos(ang_y) + sin_y = tnp.sin(ang_y) + # Do the rotations + for j in range(part_num + 1): + px = kin.x_component(generated_particles[j]) + py = kin.y_component(generated_particles[j]) + # Rotate about z + # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it? + generated_particles[j] = tnp.concatenate( + [ + cos_z * px - sin_z * py, + sin_z * px + cos_z * py, + kin.z_component(generated_particles[j]), + kin.time_component(generated_particles[j]), + ], + axis=1, + ) + # Rotate about y + px = kin.x_component(generated_particles[j]) + pz = kin.z_component(generated_particles[j]) + generated_particles[j] = tnp.concatenate( + [ + cos_y * px - sin_y * pz, + kin.y_component(generated_particles[j]), + sin_y * px + cos_y * pz, + kin.time_component(generated_particles[j]), + ], + axis=1, + ) + if part_num == (n_particles - 1): + break + betas = pds[part_num] / tnp.sqrt( + tnp.square(pds[part_num]) + tnp.square(inv_masses[part_num]) + ) + generated_particles = [ + kin.lorentz_boost( + part, + tnp.concatenate([zero_component, betas, zero_component], axis=1), + ) + for part in generated_particles + ] + part_num += 1 + return generated_particles, weights + + @function + def _recursive_generate( + self, + n_events, + boost_to=None, + recalculate_max_weights=False, + rng: SeedLike = None, + ): + """Recursively generate normalized n-body phase space as tensorflow tensors. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (tensor, optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + recalculate_max_weights (bool, optional): Recalculate the maximum weight of the event + using all the particles of the tree? This is necessary for the top particle of a decay, + otherwise the maximum weight calculation is going to be wrong (particles from subdecays + would not be taken into account). Defaults to False. + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + tuple: Result of the generation (per-event weights, maximum weights, output particles + and their output masses). + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + if boost_to is not None: + momentum = boost_to + else: + if self.has_fixed_mass: + momentum = tnp.broadcast_to( + tnp.stack((0.0, 0.0, 0.0, self.get_mass()), axis=-1), (n_events, 4) + ) + else: + raise ValueError("Cannot use resonance as top particle") + weights, weights_max, parts, children_masses = self._generate( + momentum, n_events, rng=rng + ) + output_particles = { + child.name: parts[child_num] + for child_num, child in enumerate(self.children) + } + output_masses = { + child.name: tnp.take(children_masses, [child_num], axis=1) + for child_num, child in enumerate(self.children) + } + for child_num, child in enumerate(self.children): + if child.has_children: + ( + child_weights, + _, + child_gen_particles, + child_masses, + ) = child._recursive_generate( + n_events=n_events, + boost_to=parts[child_num], + recalculate_max_weights=False, + rng=rng, + ) + weights *= child_weights + output_particles.update(child_gen_particles) + output_masses.update(child_masses) + if recalculate_max_weights: + + def build_mass_tree(particle, leaf): + if particle.has_children: + leaf[particle.name] = {} + for child in particle.children: + build_mass_tree(child, leaf[particle.name]) + else: + leaf[particle.name] = output_masses[particle.name] + + def get_flattened_values(dict_): + output = [] + for val in dict_.values(): + if isinstance(val, dict): + output.extend(get_flattened_values(val)) + else: + output.append(val) + return output + + def recurse_w_max(parent_mass, current_mass_tree): + available_mass = parent_mass - sum( + get_flattened_values(current_mass_tree) + ) + masses = [] + w_max = tnp.ones_like(available_mass) + for child, child_mass in current_mass_tree.items(): + if isinstance(child_mass, dict): + w_max *= recurse_w_max( + parent_mass + - sum( + get_flattened_values( + { + ch_it: ch_m_it + for ch_it, ch_m_it in current_mass_tree.items() + if ch_it != child + } + ) + ), + child_mass, + ) + masses.append(sum(get_flattened_values(child_mass))) + else: + masses.append(child_mass) + masses = tnp.concatenate(masses, axis=1) + w_max *= self._get_w_max(available_mass, masses) + return w_max + + mass_tree = {} + build_mass_tree(self, mass_tree) + momentum = process_list_to_tensor(momentum) + if len(momentum.shape) == 1: + momentum = tnp.expand_dims(momentum, axis=-1) + weights_max = tnp.reshape( + recurse_w_max(kin.mass(momentum), mass_tree[self.name]), (n_events,) + ) + return weights, weights_max, output_particles, output_masses + + def generate( + self, + n_events: Union[int, tf.Tensor, tf.Variable], + boost_to: Optional[tf.Tensor] = None, + normalize_weights: bool = True, + seed: SeedLike = None, + ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]: + """Generate normalized n-body phase space as tensorflow tensors. + + Any TensorFlow tensor can always be converted to a numpy array with the method `numpy()`. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + normalize_weights (bool, optional): Normalize the event weight to its max? + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + tuple: Result of the generation, which varies with the value of `normalize_weights`: + + + If True, the tuple elements are the normalized event weights as a tensor of shape + (n_events, ), and the momenta generated particles as a dictionary of tensors of shape + (4, n_events) with particle names as keys. + + + If False, the tuple weights are the unnormalized event weights as a tensor of shape + (n_events, ), the maximum per-event weights as a tensor of shape (n_events, ) and the + momenta generated particles as a dictionary of tensors of shape (4, n_events) with particle + names as keys. + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + rng = get_rng(seed) + if boost_to is not None: + message = ( + f"The number of events requested ({n_events}) doesn't match the boost_to input size " + f"of {boost_to.shape}" + ) + tf.assert_equal(tf.shape(boost_to)[0], tf.shape(n_events), message=message) + if not isinstance(n_events, tf.Variable): + n_events = tnp.asarray(n_events, dtype=tnp.int64) + weights, weights_max, parts, _ = self._recursive_generate( + n_events=n_events, + boost_to=boost_to, + recalculate_max_weights=self.has_grandchildren, + rng=rng, + ) + return ( + (weights / weights_max, parts) + if normalize_weights + else (weights, weights_max, parts) + ) + + def generate_tensor( + self, + n_events: int, + boost_to=None, + normalize_weights: bool = True, + ): + """Generate normalized n-body phase space as numpy arrays. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + normalize_weights (bool, optional): Normalize the event weight to its max + + + Return: + tuple: Result of the generation, which varies with the value of `normalize_weights`: + + + If True, the tuple elements are the normalized event weights as an array of shape + (n_events, ), and the momenta generated particles as a dictionary of arrays of shape + (4, n_events) with particle names as keys. + + + If False, the tuple weights are the unnormalized event weights as an array of shape + (n_events, ), the maximum per-event weights as an array of shape (n_events, ) and the + momenta generated particles as a dictionary of arrays of shape (4, n_events) with particle + names as keys. + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + + # Run generation + warnings.warn( + "This function is depreceated. Use `generate` which does not return a Tensor as well." + ) + generate_tf = self.generate(n_events, boost_to, normalize_weights) + # self._cache = generate_tf + # self._set_cache_validity(True, propagate=True) + return generate_tf + + +# legacy class to warn user about name change +class Particle: + """Deprecated Particle class. + + Renamed to GenParticle. + """ + + def __init__(self): + raise NameError( + "'Particle' has been renamed to 'GenParticle'. Please update your code accordingly." + "For more information, see: https://github.com/zfit/phasespace/issues/22" + ) + + +def nbody_decay(mass_top: float, masses: list, top_name: str = "", names: list = None): + """Shortcut to build an n-body decay of a GenParticle. + + If the particle names are not given, the top particle is called 'top' and the + children 'p_{i}', where i corresponds to their position in the `masses` sequence. + + Arguments: + mass_top (tensor, list): Mass of the top particle. Can be a list of 4-vectors. + masses (list): Masses of the child particles. + name_top (str, optional): Name of the top particle. If not given, the top particle is + named top. + names (list, optional): Names of the child particles. If not given, they are build as + 'p_{i}', where i is given by their ordering in the `masses` list. + + Return: + `GenParticle`: Particle decay. + + Raise: + ValueError: If the length of `masses` and `names` doesn't match. + """ + if not top_name: + top_name = "top" + if not names: + names = [f"p_{num}" for num in range(len(masses))] + if len(names) != len(masses): + raise ValueError("Mismatch in length between children masses and their names.") + return GenParticle(top_name, mass_top).set_children( + *(GenParticle(names[num], mass=mass) for num, mass in enumerate(masses)) + ) + + +def generate_decay(*args, **kwargs): + """Deprecated.""" + raise NameError( + "'generate_decay' has been removed. A similar behavior can be accomplished with 'nbody_decay'. " + "For more information see https://github.com/zfit/phasespace/issues/22" + ) + + +# EOF diff --git a/phasespace/random.py b/phasespace/random.py index 57c90aee..92fedde5 100644 --- a/phasespace/random.py +++ b/phasespace/random.py @@ -1,41 +1,41 @@ -"""Random number generation. - -As the random number generation is not a trivial thing, this module handles it uniformly. - -It mimicks the TensorFlows API on random generators and relies (currently) in global states on the TF states. -Especially on the global random number generator which will be used to get new generators. -""" -from typing import Optional, Union - -import tensorflow as tf - -SeedLike = Optional[Union[int, tf.random.Generator]] - - -def get_rng(seed: SeedLike = None) -> tf.random.Generator: - """Get or create a random number generators of type `tf.random.Generator`. - - This can be used to either retrieve random number generators deterministically from them - - global random number generator from TensorFlow, - - from a random number generator generated from the seed or - - from the random number generator passed. - - Both when using either the global generator or a random number generator is passed, they advance - by exactly one step as `split` is called on them. - - Args: - seed: This can be - - `None` to get the global random number generator - - a numerical seed to create a random number generator - - a `tf.random.Generator`. - - Returns: - A list of `tf.random.Generator` - """ - if seed is None: - rng = tf.random.get_global_generator() - elif not isinstance(seed, tf.random.Generator): # it's a seed, not an rng - rng = tf.random.Generator.from_seed(seed=seed) - else: - rng = seed - return rng +"""Random number generation. + +As the random number generation is not a trivial thing, this module handles it uniformly. + +It mimicks the TensorFlows API on random generators and relies (currently) in global states on the TF states. +Especially on the global random number generator which will be used to get new generators. +""" +from typing import Optional, Union + +import tensorflow as tf + +SeedLike = Optional[Union[int, tf.random.Generator]] + + +def get_rng(seed: SeedLike = None) -> tf.random.Generator: + """Get or create a random number generators of type `tf.random.Generator`. + + This can be used to either retrieve random number generators deterministically from them + - global random number generator from TensorFlow, + - from a random number generator generated from the seed or + - from the random number generator passed. + + Both when using either the global generator or a random number generator is passed, they advance + by exactly one step as `split` is called on them. + + Args: + seed: This can be + - `None` to get the global random number generator + - a numerical seed to create a random number generator + - a `tf.random.Generator`. + + Returns: + A list of `tf.random.Generator` + """ + if seed is None: + rng = tf.random.get_global_generator() + elif not isinstance(seed, tf.random.Generator): # it's a seed, not an rng + rng = tf.random.Generator.from_seed(seed=seed) + else: + rng = seed + return rng diff --git a/pyproject.toml b/pyproject.toml index affcbfb2..f7d864aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ -[build-system] -requires = [ - "setuptools>=42", - "setuptools_scm[toml]>=3.4", - "setuptools_scm_git_archive", - "wheel" -] - -build-backend = "setuptools.build_meta" - -[tool.isort] -profile = "black" -src_paths = ["phasespace", "tests"] +[build-system] +requires = [ + "setuptools>=42", + "setuptools_scm[toml]>=3.4", + "setuptools_scm_git_archive", + "wheel" +] + +build-backend = "setuptools.build_meta" + +[tool.isort] +profile = "black" +src_paths = ["phasespace", "tests"] diff --git a/scripts/prepare_test_samples.cxx b/scripts/prepare_test_samples.cxx index 24448d14..d5efd34b 100644 --- a/scripts/prepare_test_samples.cxx +++ b/scripts/prepare_test_samples.cxx @@ -1,123 +1,123 @@ -#include "TROOT.h" -#include "TSystem.h" -#include "TFile.h" -#include "TTree.h" -#include "TLorentzVector.h" -#include "TGenPhaseSpace.h" - -Int_t N_EVENTS = 100000; -Double_t B0_MASS = 5279.58; -Double_t PION_MASS = 139.57018; - - -int prepare_two_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[2] = {PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 2, masses); - - TLorentzVector pion_1, pion_2; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_three_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[3] = {PION_MASS, PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 3, masses); - - TLorentzVector pion_1, pion_2, pion_3; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("pion_3", "TLorentzVector", &pion_3); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - TLorentzVector *p_3 = event.GetDecay(2); - pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_four_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[4] = {PION_MASS, PION_MASS, PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 4, masses); - - TLorentzVector pion_1, pion_2, pion_3, pion_4; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("pion_3", "TLorentzVector", &pion_3); - out_tree.Branch("pion_4", "TLorentzVector", &pion_4); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - TLorentzVector *p_3 = event.GetDecay(2); - pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); - TLorentzVector *p_4 = event.GetDecay(3); - pion_4.SetPxPyPzE(p_4->Px(), p_4->Py(), p_4->Pz(), p_4->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_test_samples(TString two_body_file, TString three_body_file, TString four_body_file){ - if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); - - prepare_two_body(two_body_file); - prepare_three_body(three_body_file); - prepare_four_body(four_body_file); - - return 0; -} +#include "TROOT.h" +#include "TSystem.h" +#include "TFile.h" +#include "TTree.h" +#include "TLorentzVector.h" +#include "TGenPhaseSpace.h" + +Int_t N_EVENTS = 100000; +Double_t B0_MASS = 5279.58; +Double_t PION_MASS = 139.57018; + + +int prepare_two_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[2] = {PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 2, masses); + + TLorentzVector pion_1, pion_2; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_three_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[3] = {PION_MASS, PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 3, masses); + + TLorentzVector pion_1, pion_2, pion_3; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("pion_3", "TLorentzVector", &pion_3); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + TLorentzVector *p_3 = event.GetDecay(2); + pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_four_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[4] = {PION_MASS, PION_MASS, PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 4, masses); + + TLorentzVector pion_1, pion_2, pion_3, pion_4; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("pion_3", "TLorentzVector", &pion_3); + out_tree.Branch("pion_4", "TLorentzVector", &pion_4); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + TLorentzVector *p_3 = event.GetDecay(2); + pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); + TLorentzVector *p_4 = event.GetDecay(3); + pion_4.SetPxPyPzE(p_4->Px(), p_4->Py(), p_4->Pz(), p_4->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_test_samples(TString two_body_file, TString three_body_file, TString four_body_file){ + if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); + + prepare_two_body(two_body_file); + prepare_three_body(three_body_file); + prepare_four_body(four_body_file); + + return 0; +} diff --git a/setup.cfg b/setup.cfg index 3a5d7a3b..3d0cf47a 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,95 +1,95 @@ -[metadata] -name = phasespace -description = TensorFlow implementation of the Raubold and Lynch method for n-body events -long_description = file: README.rst -long_description_content_type = text/x-rst -url = https://github.com/zfit/phasespace -author = Albert Puig Navarro -author_email = apuignav@gmail.com -maintainer = zfit -maintainer_email = zfit@physik.uzh.ch -license = BSD-3-Clause -license_file = LICENSE -classifiers = - Development Status :: 5 - Production/Stable - Intended Audience :: Science/Research - License :: OSI Approved :: BSD License - Natural Language :: English - Operating System :: MacOS - Operating System :: Unix - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.6 - Programming Language :: Python :: 3.7 - Programming Language :: Python :: 3.8 - Programming Language :: Python :: 3.9 - Topic :: Scientific/Engineering :: Physics -keywords = TensorFlow, phasespace, HEP - -[options] -packages = find: -setup_requires = - setuptools_scm -install_requires = - tensorflow>=2.5,<2.7 # tensorflow.experimental.numpy - tensorflow_probability>=0.11 -python_requires = >=3.6 -include_package_data = True -testpaths = tests -zip_safe = False - -[options.extras_require] -test = - awkward - coverage - flaky - matplotlib - numpy - pytest - pytest-cov - pytest-xdist - scipy - uproot4 - wget -doc = - Sphinx - sphinx_bootstrap_theme - jupyter_sphinx - sphinx-math-dollar -dev = - %(doc)s - %(test)s - bumpversion - pre-commit - twine - watchdog - -[bdist_wheel] -universal = 1 - -[flake8] -exclude = - benchmark, - data, - dist, - docs, - paper, - scripts, - utils -max-line-length = 110 -statistics = True -max-complexity = 30 - -[coverage:run] -branch = True -include = - phasespace/* - -[tool:pytest] -addopts = - --color=yes - --ignore=setup.py -filterwarnings = - ignore:.*the imp module is deprecated in favour of importlib.*:DeprecationWarning -norecursedirs = - tests/helpers +[metadata] +name = phasespace +description = TensorFlow implementation of the Raubold and Lynch method for n-body events +long_description = file: README.rst +long_description_content_type = text/x-rst +url = https://github.com/zfit/phasespace +author = Albert Puig Navarro +author_email = apuignav@gmail.com +maintainer = zfit +maintainer_email = zfit@physik.uzh.ch +license = BSD-3-Clause +license_file = LICENSE +classifiers = + Development Status :: 5 - Production/Stable + Intended Audience :: Science/Research + License :: OSI Approved :: BSD License + Natural Language :: English + Operating System :: MacOS + Operating System :: Unix + Programming Language :: Python :: 3 + Programming Language :: Python :: 3 :: Only + Programming Language :: Python :: 3.6 + Programming Language :: Python :: 3.7 + Programming Language :: Python :: 3.8 + Programming Language :: Python :: 3.9 + Topic :: Scientific/Engineering :: Physics +keywords = TensorFlow, phasespace, HEP + +[options] +packages = find: +setup_requires = + setuptools_scm +install_requires = + tensorflow>=2.5,<2.7 # tensorflow.experimental.numpy + tensorflow_probability>=0.11 +python_requires = >=3.6 +include_package_data = True +testpaths = tests +zip_safe = False + +[options.extras_require] +test = + awkward + coverage + flaky + matplotlib + numpy + pytest + pytest-cov + pytest-xdist + scipy + uproot4 + wget +doc = + Sphinx + sphinx_bootstrap_theme + jupyter_sphinx + sphinx-math-dollar +dev = + %(doc)s + %(test)s + bumpversion + pre-commit + twine + watchdog + +[bdist_wheel] +universal = 1 + +[flake8] +exclude = + benchmark, + data, + dist, + docs, + paper, + scripts, + utils +max-line-length = 110 +statistics = True +max-complexity = 30 + +[coverage:run] +branch = True +include = + phasespace/* + +[tool:pytest] +addopts = + --color=yes + --ignore=setup.py +filterwarnings = + ignore:.*the imp module is deprecated in favour of importlib.*:DeprecationWarning +norecursedirs = + tests/helpers diff --git a/setup.py b/setup.py index edbdeba7..bdfb6d20 100644 --- a/setup.py +++ b/setup.py @@ -1,20 +1,20 @@ -#!/usr/bin/env python - -"""The setup script.""" - -import os - -from setuptools import setup - -here = os.path.abspath(os.path.dirname(__file__)) - -with open(os.path.join(here, "README.rst"), encoding="utf-8") as readme_file: - readme = readme_file.read() - -with open(os.path.join(here, "CHANGELOG.rst"), encoding="utf-8") as history_file: - history = history_file.read() - -setup( - long_description=readme.replace(":math:", "") + "\n\n" + history, - use_scm_version=True, -) +#!/usr/bin/env python + +"""The setup script.""" + +import os + +from setuptools import setup + +here = os.path.abspath(os.path.dirname(__file__)) + +with open(os.path.join(here, "README.rst"), encoding="utf-8") as readme_file: + readme = readme_file.read() + +with open(os.path.join(here, "CHANGELOG.rst"), encoding="utf-8") as history_file: + history = history_file.read() + +setup( + long_description=readme.replace(":math:", "") + "\n\n" + history, + use_scm_version=True, +) diff --git a/tests/conftest.py b/tests/conftest.py index f0de9679..c9c23692 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,4 @@ -import os -import sys - -sys.path.append(os.path.join(os.path.dirname(__file__), "helpers")) +import os +import sys + +sys.path.append(os.path.join(os.path.dirname(__file__), "helpers")) diff --git a/tests/fulldecay/__init__.py b/tests/fromdecay/__init__.py similarity index 63% rename from tests/fulldecay/__init__.py rename to tests/fromdecay/__init__.py index b98e9386..87d1a928 100644 --- a/tests/fulldecay/__init__.py +++ b/tests/fromdecay/__init__.py @@ -1,4 +1,4 @@ -import pytest - -# This makes it so that assert errors are more helpful for e.g., the check_norm helper function -pytest.register_assert_rewrite("fulldecay.test_fulldecay") +import pytest + +# This makes it so that assert errors are more helpful for e.g., the check_norm helper function +pytest.register_assert_rewrite("fromdecay.test_fulldecay") diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py similarity index 87% rename from tests/fulldecay/example_decay_chains.py rename to tests/fromdecay/example_decay_chains.py index ef71fda6..c04d867d 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -25,5 +25,5 @@ ): decay_mode["zfit"] = mass_function -# D*+ particle that has multiple child particles, grandchild particles, many of which can decay in multiple ways. +# 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/fulldecay/example_decays.dec b/tests/fromdecay/example_decays.dec similarity index 96% rename from tests/fulldecay/example_decays.dec rename to tests/fromdecay/example_decays.dec index c1eae76e..f5bed887 100644 --- a/tests/fulldecay/example_decays.dec +++ b/tests/fromdecay/example_decays.dec @@ -1,31 +1,31 @@ -# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec -# Example decay chain for testing purposes -# Considered by itself, this file in in fact incomplete, -# as there are no instructions on how to decay the anti-D0 and the D-! - -Decay D*+ -0.6770 D0 pi+ VSS; -0.3070 D+ pi0 VSS; -0.0160 D+ gamma VSP_PWAVE; -Enddecay - -Decay D*- -0.6770 anti-D0 pi- VSS; -0.3070 D- pi0 VSS; -0.0160 D- gamma VSP_PWAVE; -Enddecay - -Decay D0 -1.0 K- pi+ PHSP; -Enddecay - -Decay D+ -1.0 K- pi+ pi+ pi0 PHSP; -Enddecay - -Decay pi0 -0.988228297 gamma gamma PHSP; -0.011738247 e+ e- gamma PI0_DALITZ; -0.000033392 e+ e+ e- e- PHSP; -0.000000065 e+ e- PHSP; -Enddecay +# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec +# Example decay chain for testing purposes +# Considered by itself, this file in in fact incomplete, +# as there are no instructions on how to decay the anti-D0 and the D-! + +Decay D*+ +0.6770 D0 pi+ VSS; +0.3070 D+ pi0 VSS; +0.0160 D+ gamma VSP_PWAVE; +Enddecay + +Decay D*- +0.6770 anti-D0 pi- VSS; +0.3070 D- pi0 VSS; +0.0160 D- gamma VSP_PWAVE; +Enddecay + +Decay D0 +1.0 K- pi+ PHSP; +Enddecay + +Decay D+ +1.0 K- pi+ pi+ pi0 PHSP; +Enddecay + +Decay pi0 +0.988228297 gamma gamma PHSP; +0.011738247 e+ e- gamma PI0_DALITZ; +0.000033392 e+ e+ e- e- PHSP; +0.000000065 e+ e- PHSP; +Enddecay diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fromdecay/test_fulldecay.py similarity index 84% rename from tests/fulldecay/test_fulldecay.py rename to tests/fromdecay/test_fulldecay.py index 4baa7568..ad04cac1 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fromdecay/test_fulldecay.py @@ -1,9 +1,9 @@ from numpy.testing import assert_almost_equal -from phasespace.fulldecay import FullDecay -from phasespace.fulldecay.mass_functions import _DEFAULT_CONVERTER +from phasespace.fromdecay import FullDecay +from phasespace.fromdecay.mass_functions import _DEFAULT_CONVERTER -from .example_decay_chains import * # TODO remove * since it is bad practice? +from . import example_decay_chains def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: @@ -33,7 +33,7 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: def test_single_chain(): """Test converting a decaylanguage dict with only one possible decay.""" - container = FullDecay.from_dict(dplus_single, tolerance=1e-10) + container = FullDecay.from_dict(example_decay_chains.dplus_single, tolerance=1e-10) output_decay = container.gen_particles assert len(output_decay) == 1 prob, gen = output_decay[0] @@ -56,7 +56,7 @@ def test_single_chain(): def test_branching_children(): """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" - container = FullDecay.from_dict(pi0_4branches, tolerance=1e-10) + container = FullDecay.from_dict(example_decay_chains.pi0_4branches, 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) @@ -66,7 +66,7 @@ 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 = FullDecay.from_dict(dplus_4grandbranches) + container = FullDecay.from_dict(example_decay_chains.dplus_4grandbranches) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) @@ -77,7 +77,7 @@ def test_branching_grandchilden(): def test_mass_converter(): """Test that the mass_converter parameter works as intended.""" - dplus_4grandbranches_massfunc = dplus_4grandbranches.copy() + dplus_4grandbranches_massfunc = example_decay_chains.dplus_4grandbranches.copy() dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" container = FullDecay.from_dict( dplus_4grandbranches_massfunc, @@ -98,7 +98,7 @@ def test_mass_converter(): def test_big_decay(): - container = FullDecay.from_dict(dstarplus_big_decay) + container = FullDecay.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) diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py similarity index 94% rename from tests/fulldecay/test_mass_functions.py rename to tests/fromdecay/test_mass_functions.py index a7c12d3e..e8c4eac0 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fromdecay/test_mass_functions.py @@ -5,7 +5,7 @@ import tensorflow_probability as tfp from particle import Particle -import phasespace.fulldecay.mass_functions as mf +import phasespace.fromdecay.mass_functions as mf _kstarz = Particle.from_evtgen_name("K*0") KSTARZ_MASS = _kstarz.mass diff --git a/tests/helpers/decays.py b/tests/helpers/decays.py index 6ff09099..313b2e11 100644 --- a/tests/helpers/decays.py +++ b/tests/helpers/decays.py @@ -1,80 +1,80 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file decays.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Some physics models to test with.""" - -import tensorflow as tf -import tensorflow_probability as tfp - -from phasespace import GenParticle - -# Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat) -B0_MASS = 5279.58 -PION_MASS = 139.57018 -KAON_MASS = 493.677 -K1_MASS = 1272.0 -K1_WIDTH = 90.0 -KSTARZ_MASS = 895.81 -KSTARZ_WIDTH = 47.4 - - -def b0_to_kstar_gamma(kstar_width=KSTARZ_WIDTH): - """Generate B0 -> K*gamma.""" - - def kstar_mass(min_mass, max_mass, n_events): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(kstar_width, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - if kstar_width > 0: - kstar_mass = tfp.distributions.TruncatedNormal( - loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass - ).sample() - return kstar_mass - - return GenParticle("B0", B0_MASS).set_children( - GenParticle("K*0", mass=kstar_mass).set_children( - GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) - ), - GenParticle("gamma", mass=0.0), - ) - - -def bp_to_k1_kstar_pi_gamma(k1_width=K1_WIDTH, kstar_width=KSTARZ_WIDTH): - """Generate B+ -> K1 (-> K* (->K pi) pi) gamma.""" - - def res_mass(mass, width, min_mass, max_mass, n_events): - mass = tf.cast(mass, tf.float64) - width = tf.cast(width, tf.float64) - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - masses = tf.broadcast_to(mass, shape=(n_events,)) - if kstar_width > 0: - masses = tfp.distributions.TruncatedNormal( - loc=masses, scale=width, low=min_mass, high=max_mass - ).sample() - return masses - - def k1_mass(min_mass, max_mass, n_events): - return res_mass(K1_MASS, k1_width, min_mass, max_mass, n_events) - - def kstar_mass(min_mass, max_mass, n_events): - return res_mass(KSTARZ_MASS, kstar_width, min_mass, max_mass, n_events) - - return GenParticle("B+", B0_MASS).set_children( - GenParticle("K1+", mass=k1_mass).set_children( - GenParticle("K*0", mass=kstar_mass).set_children( - GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) - ), - GenParticle("pi+", mass=PION_MASS), - ), - GenParticle("gamma", mass=0.0), - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file decays.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Some physics models to test with.""" + +import tensorflow as tf +import tensorflow_probability as tfp + +from phasespace import GenParticle + +# Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat) +B0_MASS = 5279.58 +PION_MASS = 139.57018 +KAON_MASS = 493.677 +K1_MASS = 1272.0 +K1_WIDTH = 90.0 +KSTARZ_MASS = 895.81 +KSTARZ_WIDTH = 47.4 + + +def b0_to_kstar_gamma(kstar_width=KSTARZ_WIDTH): + """Generate B0 -> K*gamma.""" + + def kstar_mass(min_mass, max_mass, n_events): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(kstar_width, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + if kstar_width > 0: + kstar_mass = tfp.distributions.TruncatedNormal( + loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass + ).sample() + return kstar_mass + + return GenParticle("B0", B0_MASS).set_children( + GenParticle("K*0", mass=kstar_mass).set_children( + GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) + ), + GenParticle("gamma", mass=0.0), + ) + + +def bp_to_k1_kstar_pi_gamma(k1_width=K1_WIDTH, kstar_width=KSTARZ_WIDTH): + """Generate B+ -> K1 (-> K* (->K pi) pi) gamma.""" + + def res_mass(mass, width, min_mass, max_mass, n_events): + mass = tf.cast(mass, tf.float64) + width = tf.cast(width, tf.float64) + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + masses = tf.broadcast_to(mass, shape=(n_events,)) + if kstar_width > 0: + masses = tfp.distributions.TruncatedNormal( + loc=masses, scale=width, low=min_mass, high=max_mass + ).sample() + return masses + + def k1_mass(min_mass, max_mass, n_events): + return res_mass(K1_MASS, k1_width, min_mass, max_mass, n_events) + + def kstar_mass(min_mass, max_mass, n_events): + return res_mass(KSTARZ_MASS, kstar_width, min_mass, max_mass, n_events) + + return GenParticle("B+", B0_MASS).set_children( + GenParticle("K1+", mass=k1_mass).set_children( + GenParticle("K*0", mass=kstar_mass).set_children( + GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) + ), + GenParticle("pi+", mass=PION_MASS), + ), + GenParticle("gamma", mass=0.0), + ) + + +# EOF diff --git a/tests/helpers/plotting.py b/tests/helpers/plotting.py index 3051c214..1e222b5a 100644 --- a/tests/helpers/plotting.py +++ b/tests/helpers/plotting.py @@ -1,28 +1,28 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file plotting.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Plotting helpers for tests.""" - -import numpy as np - - -def make_norm_histo(array, range_, weights=None): - """Make histo and modify dimensions.""" - histo = np.histogram(array, 100, range=range_, weights=weights)[0] - return histo / np.sum(histo) - - -def mass(vector): - """Calculate mass scalar for Lorentz 4-momentum.""" - return np.sqrt( - np.sum( - vector * vector * np.reshape(np.array([-1.0, -1.0, -1.0, 1.0]), (1, 4)), - axis=1, - ) - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file plotting.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Plotting helpers for tests.""" + +import numpy as np + + +def make_norm_histo(array, range_, weights=None): + """Make histo and modify dimensions.""" + histo = np.histogram(array, 100, range=range_, weights=weights)[0] + return histo / np.sum(histo) + + +def mass(vector): + """Calculate mass scalar for Lorentz 4-momentum.""" + return np.sqrt( + np.sum( + vector * vector * np.reshape(np.array([-1.0, -1.0, -1.0, 1.0]), (1, 4)), + axis=1, + ) + ) + + +# EOF diff --git a/tests/helpers/rapidsim.py b/tests/helpers/rapidsim.py index 94821a5c..d6d94cc0 100644 --- a/tests/helpers/rapidsim.py +++ b/tests/helpers/rapidsim.py @@ -1,111 +1,111 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file rapidsim.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Utils to crossheck against RapidSim.""" - -import os - -import numpy as np -import uproot4 - -BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) -FONLL_FILE = os.path.join(BASE_PATH, "data", "fonll", "LHC{}{}.root") - - -def get_fonll_histos(energy, quark): - with uproot4.open(FONLL_FILE.format(quark, int(energy))) as histo_file: - return histo_file["pT"], histo_file["eta"] - - -def generate_fonll(mass, beam_energy, quark, n_events): - def analyze_histo(histo): - x_axis = histo.axis(0) - x_bins = x_axis.edges() - bin_width = x_axis.width - bin_centers = x_bins[:-1] + bin_width / 2 - normalized_values = histo.values() / np.sum(histo.values()) - return bin_centers, normalized_values - - pt_histo, eta_histo = get_fonll_histos(beam_energy, quark) - pt_bin_centers, pt_normalized_values = analyze_histo(pt_histo) - eta_bin_centers, eta_normalized_values = analyze_histo(eta_histo) - pt_rand = np.random.choice( - pt_bin_centers, - size=n_events, - p=pt_normalized_values, - ) - pt_rand = 1_000 * np.abs(pt_rand) - eta_rand = np.random.choice( - eta_bin_centers, - size=n_events, - p=eta_normalized_values, - ) - phi_rand = np.random.uniform(0, 2 * np.pi, size=n_events) - px = pt_rand * np.cos(phi_rand) - py = pt_rand * np.sin(phi_rand) - pz = pt_rand * np.sinh(eta_rand) - e = np.sqrt(px * px + py * py + pz * pz + mass * mass) - return np.stack([px, py, pz, e]) - - -def load_generated_histos(file_name, particles): - with uproot4.open(file_name) as rapidsim_file: - return { - particle: [ - rapidsim_file.get(f"{particle}_{coord}_TRUE").array(library="np") - for coord in ("PX", "PY", "PZ", "E") - ] - for particle in particles - } - - -def get_tree(file_name, top_particle, particles): - """Load a RapidSim tree.""" - with uproot4.open(file_name) as rapidsim_file: - tree = rapidsim_file["DecayTree"] - return { - particle: np.stack( - [ - 1000.0 * tree[f"{particle}_{coord}_TRUE"].array(library="np") - for coord in ("PX", "PY", "PZ", "E") - ] - ) - for particle in particles - } - - -def get_tree_in_b_rest_frame(file_name, top_particle, particles): - def lorentz_boost(part_to_boost, boost): - """ - Perform Lorentz boost - vector : 4-vector to be boosted - boostvector: boost vector. Can be either 3-vector or 4-vector (only spatial components - are used) - """ - boost_vec = -boost[:3, :] / boost[3, :] - b2 = np.sum(boost_vec * boost_vec, axis=0) - gamma = 1.0 / np.sqrt(1.0 - b2) - gamma2 = (gamma - 1.0) / b2 - part_time = part_to_boost[3, :] - part_space = part_to_boost[:3, :] - bp = np.sum(part_space * boost_vec, axis=0) - return np.concatenate( - [ - part_space + (gamma2 * bp + gamma * part_time) * boost_vec, - np.expand_dims(gamma * (part_time + bp), axis=0), - ], - axis=0, - ) - - part_dict = get_tree(file_name, top_particle, list(particles) + [top_particle]) - top_parts = part_dict.pop(top_particle) - return { - part_name: lorentz_boost(part, top_parts) - for part_name, part in part_dict.items() - } - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file rapidsim.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Utils to crossheck against RapidSim.""" + +import os + +import numpy as np +import uproot4 + +BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +FONLL_FILE = os.path.join(BASE_PATH, "data", "fonll", "LHC{}{}.root") + + +def get_fonll_histos(energy, quark): + with uproot4.open(FONLL_FILE.format(quark, int(energy))) as histo_file: + return histo_file["pT"], histo_file["eta"] + + +def generate_fonll(mass, beam_energy, quark, n_events): + def analyze_histo(histo): + x_axis = histo.axis(0) + x_bins = x_axis.edges() + bin_width = x_axis.width + bin_centers = x_bins[:-1] + bin_width / 2 + normalized_values = histo.values() / np.sum(histo.values()) + return bin_centers, normalized_values + + pt_histo, eta_histo = get_fonll_histos(beam_energy, quark) + pt_bin_centers, pt_normalized_values = analyze_histo(pt_histo) + eta_bin_centers, eta_normalized_values = analyze_histo(eta_histo) + pt_rand = np.random.choice( + pt_bin_centers, + size=n_events, + p=pt_normalized_values, + ) + pt_rand = 1_000 * np.abs(pt_rand) + eta_rand = np.random.choice( + eta_bin_centers, + size=n_events, + p=eta_normalized_values, + ) + phi_rand = np.random.uniform(0, 2 * np.pi, size=n_events) + px = pt_rand * np.cos(phi_rand) + py = pt_rand * np.sin(phi_rand) + pz = pt_rand * np.sinh(eta_rand) + e = np.sqrt(px * px + py * py + pz * pz + mass * mass) + return np.stack([px, py, pz, e]) + + +def load_generated_histos(file_name, particles): + with uproot4.open(file_name) as rapidsim_file: + return { + particle: [ + rapidsim_file.get(f"{particle}_{coord}_TRUE").array(library="np") + for coord in ("PX", "PY", "PZ", "E") + ] + for particle in particles + } + + +def get_tree(file_name, top_particle, particles): + """Load a RapidSim tree.""" + with uproot4.open(file_name) as rapidsim_file: + tree = rapidsim_file["DecayTree"] + return { + particle: np.stack( + [ + 1000.0 * tree[f"{particle}_{coord}_TRUE"].array(library="np") + for coord in ("PX", "PY", "PZ", "E") + ] + ) + for particle in particles + } + + +def get_tree_in_b_rest_frame(file_name, top_particle, particles): + def lorentz_boost(part_to_boost, boost): + """ + Perform Lorentz boost + vector : 4-vector to be boosted + boostvector: boost vector. Can be either 3-vector or 4-vector (only spatial components + are used) + """ + boost_vec = -boost[:3, :] / boost[3, :] + b2 = np.sum(boost_vec * boost_vec, axis=0) + gamma = 1.0 / np.sqrt(1.0 - b2) + gamma2 = (gamma - 1.0) / b2 + part_time = part_to_boost[3, :] + part_space = part_to_boost[:3, :] + bp = np.sum(part_space * boost_vec, axis=0) + return np.concatenate( + [ + part_space + (gamma2 * bp + gamma * part_time) * boost_vec, + np.expand_dims(gamma * (part_time + bp), axis=0), + ], + axis=0, + ) + + part_dict = get_tree(file_name, top_particle, list(particles) + [top_particle]) + top_parts = part_dict.pop(top_particle) + return { + part_name: lorentz_boost(part, top_parts) + for part_name, part in part_dict.items() + } + + +# EOF diff --git a/tests/test_chain.py b/tests/test_chain.py index 1bcf201e..555c8cf9 100644 --- a/tests/test_chain.py +++ b/tests/test_chain.py @@ -1,137 +1,137 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_chain.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 01.03.2019 -# ============================================================================= -"""Test decay chain tools.""" - -import os -import sys - -import numpy as np -import pytest - -from phasespace import GenParticle - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays # noqa: E402 - - -def test_name_clashes(): - """Test clashes in particle naming.""" - # In children - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - ) - # With itself - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Top", mass=decays.KSTARZ_MASS), - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - ) - # In grandchildren - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( - GenParticle("K+", mass=decays.KAON_MASS), - GenParticle("pi-", mass=decays.PION_MASS), - ), - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( - GenParticle("K+", mass=decays.KAON_MASS), - GenParticle("pi-_1", mass=decays.PION_MASS), - ), - ) - - -def test_wrong_children(): - """Test wrong number of children.""" - with pytest.raises(ValueError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS) - ) - - -def test_grandchildren(): - """Test that grandchildren detection is correct.""" - top = GenParticle("Top", 0) - assert not top.has_children - assert not top.has_grandchildren - assert not top.set_children( - GenParticle("Child1", mass=decays.KSTARZ_MASS), - GenParticle("Child2", mass=decays.KSTARZ_MASS), - ).has_grandchildren - - -def test_reset_children(): - """Test when children are set twice.""" - top = GenParticle("Top", 0).set_children( - GenParticle("Child1", mass=decays.KSTARZ_MASS), - GenParticle("Child2", mass=decays.KSTARZ_MASS), - ) - with pytest.raises(ValueError): - top.set_children( - GenParticle("Child3", mass=decays.KSTARZ_MASS), - GenParticle("Child4", mass=decays.KSTARZ_MASS), - ) - - -def test_no_children(): - """Test when no children have been configured.""" - top = GenParticle("Top", 0) - with pytest.raises(ValueError): - top.generate(n_events=1) - - -def test_resonance_top(): - """Test when a resonance is used as the top particle.""" - kstar = decays.b0_to_kstar_gamma().children[0] - with pytest.raises(ValueError): - kstar.generate(n_events=1) - - -def test_kstargamma(): - """Test B0 -> K*gamma.""" - decay = decays.b0_to_kstar_gamma() - norm_weights, particles = decay.generate(n_events=1000) - assert norm_weights.shape[0] == 1000 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 4 - assert set(particles.keys()) == {"K*0", "gamma", "K+", "pi-"} - assert all(part.shape == (1000, 4) for part in particles.values()) - - -def test_k1gamma(): - """Test B+ -> K1 (K*pi) gamma.""" - decay = decays.bp_to_k1_kstar_pi_gamma() - norm_weights, particles = decay.generate(n_events=1000) - assert norm_weights.shape[0] == 1000 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 6 - assert set(particles.keys()) == {"K1+", "K*0", "gamma", "K+", "pi-", "pi+"} - assert all(part.shape == (1000, 4) for part in particles.values()) - - -def test_repr(): - """Test string representation.""" - b0 = decays.b0_to_kstar_gamma() - kst = b0.children[0] - assert ( - str(b0) - == "" - ) - assert ( - str(kst) - == "" - ) - - -if __name__ == "__main__": - test_name_clashes() - test_kstargamma() - test_k1gamma() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_chain.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 01.03.2019 +# ============================================================================= +"""Test decay chain tools.""" + +import os +import sys + +import numpy as np +import pytest + +from phasespace import GenParticle + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays # noqa: E402 + + +def test_name_clashes(): + """Test clashes in particle naming.""" + # In children + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + ) + # With itself + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Top", mass=decays.KSTARZ_MASS), + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + ) + # In grandchildren + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( + GenParticle("K+", mass=decays.KAON_MASS), + GenParticle("pi-", mass=decays.PION_MASS), + ), + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( + GenParticle("K+", mass=decays.KAON_MASS), + GenParticle("pi-_1", mass=decays.PION_MASS), + ), + ) + + +def test_wrong_children(): + """Test wrong number of children.""" + with pytest.raises(ValueError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS) + ) + + +def test_grandchildren(): + """Test that grandchildren detection is correct.""" + top = GenParticle("Top", 0) + assert not top.has_children + assert not top.has_grandchildren + assert not top.set_children( + GenParticle("Child1", mass=decays.KSTARZ_MASS), + GenParticle("Child2", mass=decays.KSTARZ_MASS), + ).has_grandchildren + + +def test_reset_children(): + """Test when children are set twice.""" + top = GenParticle("Top", 0).set_children( + GenParticle("Child1", mass=decays.KSTARZ_MASS), + GenParticle("Child2", mass=decays.KSTARZ_MASS), + ) + with pytest.raises(ValueError): + top.set_children( + GenParticle("Child3", mass=decays.KSTARZ_MASS), + GenParticle("Child4", mass=decays.KSTARZ_MASS), + ) + + +def test_no_children(): + """Test when no children have been configured.""" + top = GenParticle("Top", 0) + with pytest.raises(ValueError): + top.generate(n_events=1) + + +def test_resonance_top(): + """Test when a resonance is used as the top particle.""" + kstar = decays.b0_to_kstar_gamma().children[0] + with pytest.raises(ValueError): + kstar.generate(n_events=1) + + +def test_kstargamma(): + """Test B0 -> K*gamma.""" + decay = decays.b0_to_kstar_gamma() + norm_weights, particles = decay.generate(n_events=1000) + assert norm_weights.shape[0] == 1000 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 4 + assert set(particles.keys()) == {"K*0", "gamma", "K+", "pi-"} + assert all(part.shape == (1000, 4) for part in particles.values()) + + +def test_k1gamma(): + """Test B+ -> K1 (K*pi) gamma.""" + decay = decays.bp_to_k1_kstar_pi_gamma() + norm_weights, particles = decay.generate(n_events=1000) + assert norm_weights.shape[0] == 1000 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 6 + assert set(particles.keys()) == {"K1+", "K*0", "gamma", "K+", "pi-", "pi+"} + assert all(part.shape == (1000, 4) for part in particles.values()) + + +def test_repr(): + """Test string representation.""" + b0 = decays.b0_to_kstar_gamma() + kst = b0.children[0] + assert ( + str(b0) + == "" + ) + assert ( + str(kst) + == "" + ) + + +if __name__ == "__main__": + test_name_clashes() + test_kstargamma() + test_k1gamma() + +# EOF diff --git a/tests/test_generate.py b/tests/test_generate.py index 260918cd..548b6a97 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -1,86 +1,86 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_generate.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Basic dimensionality tests.""" - -import os -import sys - -import numpy as np -import pytest - -import phasespace - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays # noqa: E402 - -B0_MASS = decays.B0_MASS -PION_MASS = decays.PION_MASS - - -def test_one_event(): - """Test B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=1) - assert norm_weights.shape[0] == 1 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (1, 4) for part in particles.values()) - - -def test_one_event_tf(): - """Test B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=1) - - assert norm_weights.shape[0] == 1 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (1, 4) for part in particles.values()) - - -@pytest.mark.parametrize("n_events", argvalues=[5, 523]) -def test_n_events(n_events): - """Test 5 B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=n_events) - assert norm_weights.shape[0] == n_events - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (n_events, 4) for part in particles.values()) - - -def test_deterministic_events(): - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - common_seed = 36 - norm_weights_seeded1, particles_seeded1 = decay.generate( - n_events=100, seed=common_seed - ) - norm_weights_global, particles_global = decay.generate(n_events=100) - norm_weights_rnd, particles_rnd = decay.generate(n_events=100, seed=152) - norm_weights_seeded2, particles_seeded2 = decay.generate( - n_events=100, seed=common_seed - ) - - np.testing.assert_allclose(norm_weights_seeded1, norm_weights_seeded2) - for part1, part2 in zip(particles_seeded1.values(), particles_seeded2.values()): - np.testing.assert_allclose(part1, part2) - - assert not np.allclose(norm_weights_seeded1, norm_weights_rnd) - for part1, part2 in zip(particles_seeded1.values(), particles_rnd.values()): - assert not np.allclose(part1, part2) - - assert not np.allclose(norm_weights_global, norm_weights_rnd) - for part1, part2 in zip(particles_global.values(), particles_rnd.values()): - assert not np.allclose(part1, part2) - - -if __name__ == "__main__": - test_one_event() - test_n_events(5) - - # EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_generate.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Basic dimensionality tests.""" + +import os +import sys + +import numpy as np +import pytest + +import phasespace + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays # noqa: E402 + +B0_MASS = decays.B0_MASS +PION_MASS = decays.PION_MASS + + +def test_one_event(): + """Test B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=1) + assert norm_weights.shape[0] == 1 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (1, 4) for part in particles.values()) + + +def test_one_event_tf(): + """Test B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=1) + + assert norm_weights.shape[0] == 1 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (1, 4) for part in particles.values()) + + +@pytest.mark.parametrize("n_events", argvalues=[5, 523]) +def test_n_events(n_events): + """Test 5 B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=n_events) + assert norm_weights.shape[0] == n_events + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (n_events, 4) for part in particles.values()) + + +def test_deterministic_events(): + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + common_seed = 36 + norm_weights_seeded1, particles_seeded1 = decay.generate( + n_events=100, seed=common_seed + ) + norm_weights_global, particles_global = decay.generate(n_events=100) + norm_weights_rnd, particles_rnd = decay.generate(n_events=100, seed=152) + norm_weights_seeded2, particles_seeded2 = decay.generate( + n_events=100, seed=common_seed + ) + + np.testing.assert_allclose(norm_weights_seeded1, norm_weights_seeded2) + for part1, part2 in zip(particles_seeded1.values(), particles_seeded2.values()): + np.testing.assert_allclose(part1, part2) + + assert not np.allclose(norm_weights_seeded1, norm_weights_rnd) + for part1, part2 in zip(particles_seeded1.values(), particles_rnd.values()): + assert not np.allclose(part1, part2) + + assert not np.allclose(norm_weights_global, norm_weights_rnd) + for part1, part2 in zip(particles_global.values(), particles_rnd.values()): + assert not np.allclose(part1, part2) + + +if __name__ == "__main__": + test_one_event() + test_n_events(5) + + # EOF diff --git a/tests/test_nbody_decay.py b/tests/test_nbody_decay.py index efd11df9..461beae9 100644 --- a/tests/test_nbody_decay.py +++ b/tests/test_nbody_decay.py @@ -1,64 +1,64 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_nbody_decay.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 14.06.2019 -# ============================================================================= -"""Test n-body decay generator.""" - -import pytest - -from phasespace import nbody_decay - -from .helpers import decays - -B0_MASS = decays.B0_MASS -PION_MASS = decays.PION_MASS - - -def test_no_names(): - """Test particle naming when no name is given.""" - decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - assert decay.name == "top" - assert all( - part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) - ) - - -def test_top_name(): - """Test particle naming when only top name is given.""" - decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0") - assert decay.name == "B0" - assert all( - part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) - ) - - -def test_children_names(): - """Test particle naming when only children names are given.""" - children_names = [f"pion_{i}" for i in range(3)] - decay = nbody_decay( - B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names - ) - assert decay.name == "top" - assert children_names == [part.name for part in decay.children] - - -def test_all_names(): - """Test particle naming when all names are given.""" - children_names = [f"pion_{i}" for i in range(3)] - decay = nbody_decay( - B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0", names=children_names - ) - assert decay.name == "B0" - assert children_names == [part.name for part in decay.children] - - -def test_mismatching_names(): - """Test wrong number of names given for children.""" - children_names = [f"pion_{i}" for i in range(4)] - with pytest.raises(ValueError): - nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_nbody_decay.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 14.06.2019 +# ============================================================================= +"""Test n-body decay generator.""" + +import pytest + +from phasespace import nbody_decay + +from .helpers import decays + +B0_MASS = decays.B0_MASS +PION_MASS = decays.PION_MASS + + +def test_no_names(): + """Test particle naming when no name is given.""" + decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + assert decay.name == "top" + assert all( + part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) + ) + + +def test_top_name(): + """Test particle naming when only top name is given.""" + decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0") + assert decay.name == "B0" + assert all( + part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) + ) + + +def test_children_names(): + """Test particle naming when only children names are given.""" + children_names = [f"pion_{i}" for i in range(3)] + decay = nbody_decay( + B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names + ) + assert decay.name == "top" + assert children_names == [part.name for part in decay.children] + + +def test_all_names(): + """Test particle naming when all names are given.""" + children_names = [f"pion_{i}" for i in range(3)] + decay = nbody_decay( + B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0", names=children_names + ) + assert decay.name == "B0" + assert children_names == [part.name for part in decay.children] + + +def test_mismatching_names(): + """Test wrong number of names given for children.""" + children_names = [f"pion_{i}" for i in range(4)] + with pytest.raises(ValueError): + nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names) + + +# EOF diff --git a/tests/test_physics.py b/tests/test_physics.py index 4cfa98c9..025d9312 100644 --- a/tests/test_physics.py +++ b/tests/test_physics.py @@ -1,411 +1,411 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_physics.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Test physics output.""" - -import platform -import subprocess - -import numpy as np -import pytest -from scipy.stats import ks_2samp - -if platform.system() == "Darwin": - import matplotlib - - matplotlib.use("TkAgg") - -import os -import sys - -import matplotlib.pyplot as plt -import tensorflow as tf -import uproot4 - -from phasespace import phasespace - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays, rapidsim # noqa: E402 -from .helpers.plotting import make_norm_histo # noqa: E402 - -BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -PLOT_DIR = os.path.join(BASE_PATH, "tests", "plots") - - -def setup_method(): - phasespace.GenParticle._sess.close() - tf.compat.v1.reset_default_graph() - - -def create_ref_histos(n_pions): - """Load reference histogram data.""" - ref_dir = os.path.join(BASE_PATH, "data") - if not os.path.exists(ref_dir): - os.mkdir(ref_dir) - ref_file = os.path.join(ref_dir, f"bto{n_pions}pi.root") - if not os.path.exists(ref_file): - script = os.path.join( - BASE_PATH, - "scripts", - "prepare_test_samples.cxx+({})".format( - ",".join( - '"{}"'.format(os.path.join(BASE_PATH, "data", f"bto{i + 1}pi.root")) - for i in range(1, 4) - ) - ), - ) - subprocess.call(f"root -qb '{script}'", shell=True) - events = uproot4.open(ref_file)["events"] - pion_names = [f"pion_{pion + 1}" for pion in range(n_pions)] - pions = {pion_name: events[pion_name] for pion_name in pion_names} - weights = events["weight"] - normalized_histograms = [] - for pion in pions.values(): - pion_array = pion.array() - energy = pion_array.fE - momentum = pion_array.fP - for coord, array in enumerate([momentum.fX, momentum.fY, momentum.fZ, energy]): - numpy_array = np.array(array) - histogram = make_norm_histo( - numpy_array, - range_=(-3000 if coord % 4 != 3 else 0, 3000), - weights=weights, - ) - normalized_histograms.append(histogram) - - return normalized_histograms, make_norm_histo(weights, range_=(0, 1 + 1e-8)) - - -def run_test(n_particles, test_prefix): - first_run_n_events = 100 - main_run_n_events = 100000 - n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64) - - decay = phasespace.nbody_decay(decays.B0_MASS, [decays.PION_MASS] * n_particles) - generate = decay.generate(n_events) - weights1, _ = generate # only generate to test change in n_events - assert len(weights1) == first_run_n_events - - # change n_events and run again - n_events.assign(main_run_n_events) - weights, particles = decay.generate(n_events) - parts = np.concatenate( - [particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1 - ) - histos = [ - make_norm_histo( - parts[:, coord], - range_=(-3000 if coord % 4 != 3 else 0, 3000), - weights=weights, - ) - for coord in range(parts.shape[1]) - ] - weight_histos = make_norm_histo(weights, range_=(0, 1 + 1e-8)) - ref_histos, ref_weights = create_ref_histos(n_particles) - p_values = np.array( - [ - ks_2samp(histos[coord], ref_histos[coord])[1] - for coord, _ in enumerate(histos) - ] - + [ks_2samp(weight_histos, ref_weights)[1]] - ) - # Let's plot - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - for coord, _ in enumerate(histos): - plt.hist( - x if coord % 4 != 3 else e, - weights=histos[coord], - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histos[coord], - alpha=0.5, - label="TGenPhasespace", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "{}_pion_{}_{}.png".format( - test_prefix, int(coord / 4) + 1, ["px", "py", "pz", "e"][coord % 4] - ), - ) - ) - plt.clf() - plt.hist( - np.linspace(0, 1, 100), - weights=weight_histos, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - np.linspace(0, 1, 100), - weights=ref_weights, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"{test_prefix}_weights.png")) - plt.clf() - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_two_body(): - """Test B->pipi decay.""" - run_test(2, "two_body") - - -@pytest.mark.flaky(3) # Stats are limited -def test_three_body(): - """Test B -> pi pi pi decay.""" - run_test(3, "three_body") - - -@pytest.mark.flaky(3) # Stats are limited -def test_four_body(): - """Test B -> pi pi pi pi decay.""" - run_test(4, "four_body") - - -def run_kstargamma(input_file, kstar_width, b_at_rest, suffix): - """Run B0->K*gamma test.""" - n_events = 1000000 - if b_at_rest: - booster = None - rapidsim_getter = rapidsim.get_tree_in_b_rest_frame - else: - booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) - booster = booster.transpose() - rapidsim_getter = rapidsim.get_tree - decay = decays.b0_to_kstar_gamma(kstar_width=kstar_width) - norm_weights, particles = decay.generate(n_events=n_events, boost_to=booster) - rapidsim_parts = rapidsim_getter( - os.path.join(BASE_PATH, "data", input_file), - "B0_0", - ("Kst0_0", "gamma_0", "Kp_0", "pim_0"), - ) - name_matching = {"Kst0_0": "K*0", "gamma_0": "gamma", "Kp_0": "K+", "pim_0": "pi-"} - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - p_values = {} - for ref_name, ref_part in rapidsim_parts.items(): - tf_part = name_matching[ref_name] - ref_part = ref_part.transpose() # for consistency - for coord, coord_name in enumerate(("px", "py", "pz", "e")): - range_ = (-3000 if coord % 4 != 3 else 0, 3000) - ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) - tf_histo = make_norm_histo( - particles[tf_part][:, coord], range_=range_, weights=norm_weights - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=tf_histo, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histo, - alpha=0.5, - label="RapidSim", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "B0_Kstar_gamma_Kstar{}_{}_{}.png".format( - suffix, tf_part.replace("*", "star"), coord_name - ), - ) - ) - plt.clf() - p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] - plt.hist( - np.linspace(0, 1, 100), - weights=make_norm_histo(norm_weights, range_=(0, 1)), - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"B0_Kstar_gamma_Kstar{suffix}_weights.png")) - plt.clf() - return np.array(list(p_values.values())) - - -@pytest.mark.flaky(3) # Stats are limited -def test_kstargamma_kstarnonresonant_at_rest(): - """Test B0 -> K* gamma physics with fixed mass for K*.""" - p_values = run_kstargamma( - "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", 0, True, "NonResonant" - ) - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_kstargamma_kstarnonresonant_lhc(): - """Test B0 -> K* gamma physics with fixed mass for K* with LHC kinematics.""" - p_values = run_kstargamma( - "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", - 0, - False, - "NonResonant_LHC", - ) - assert np.all(p_values > 0.05) - - -def test_kstargamma_resonant_at_rest(): - """Test B0 -> K* gamma physics with Gaussian mass for K*. - - Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov - test wrt to RapidSim, so plots are generated and can be inspected by the user. However, small differences - are expected in the tails of the energy distributions of the kaon and the pion. - """ - run_kstargamma( - "B2KstGamma_RapidSim_7TeV_Tree.root", decays.KSTARZ_WIDTH, True, "Gaussian" - ) - - -def run_k1_gamma(input_file, k1_width, kstar_width, b_at_rest, suffix): - """Run B+ -> K1gamma test.""" - n_events = 1000000 - if b_at_rest: - booster = None - rapidsim_getter = rapidsim.get_tree_in_b_rest_frame - else: - booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) - booster = booster.transpose() - rapidsim_getter = rapidsim.get_tree - gamma = decays.bp_to_k1_kstar_pi_gamma(k1_width=k1_width, kstar_width=kstar_width) - norm_weights, particles = gamma.generate(n_events=n_events, boost_to=booster) - rapidsim_parts = rapidsim_getter( - os.path.join(BASE_PATH, "data", input_file), - "Bp_0", - ("K1_1270_p_0", "Kst0_0", "gamma_0", "Kp_0", "pim_0", "pip_0"), - ) - name_matching = { - "K1_1270_p_0": "K1+", - "Kst0_0": "K*0", - "gamma_0": "gamma", - "Kp_0": "K+", - "pim_0": "pi-", - "pip_0": "pi+", - } - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - p_values = {} - for ref_name, ref_part in rapidsim_parts.items(): - tf_part = name_matching[ref_name] - ref_part = ( - ref_part.transpose() - ) # to be consistent with internal shape (nevents, nobs) - for coord, coord_name in enumerate(("px", "py", "pz", "e")): - range_ = (-3000 if coord % 4 != 3 else 0, 3000) - ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) - tf_histo = make_norm_histo( - particles[tf_part][:, coord], range_=range_, weights=norm_weights - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=tf_histo, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histo, - alpha=0.5, - label="RapidSim", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "Bp_K1_gamma_K1Kstar{}_{}_{}.png".format( - suffix, tf_part.replace("*", "star"), coord_name - ), - ) - ) - plt.clf() - p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] - plt.hist( - np.linspace(0, 1, 100), - weights=make_norm_histo(norm_weights, range_=(0, 1)), - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"Bp_K1_gamma_K1Kstar{suffix}_weights.png")) - plt.clf() - return np.array(list(p_values.values())) - - -@pytest.mark.flaky(3) # Stats are limited -def test_k1gamma_kstarnonresonant_at_rest(): - """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances.""" - p_values = run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", - 0, - 0, - True, - "NonResonant", - ) - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_k1gamma_kstarnonresonant_lhc(): - """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances with LHC kinematics.""" - p_values = run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", - 0, - 0, - False, - "NonResonant_LHC", - ) - assert np.all(p_values > 0.05) - - -def test_k1gamma_resonant_at_rest(): - """Test B0 -> K1 (->K*pi) gamma physics. - - Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov - test wrt to RapidSim, so plots are generated and can be inspected by the user. - """ - run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_Tree.root", - decays.K1_WIDTH, - decays.KSTARZ_WIDTH, - True, - "Gaussian", - ) - - -if __name__ == "__main__": - test_two_body() - test_three_body() - test_four_body() - test_kstargamma_kstarnonresonant_at_rest() - test_kstargamma_kstarnonresonant_lhc() - test_kstargamma_resonant_at_rest() - test_k1gamma_kstarnonresonant_at_rest() - test_k1gamma_kstarnonresonant_lhc() - test_k1gamma_resonant_at_rest() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_physics.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Test physics output.""" + +import platform +import subprocess + +import numpy as np +import pytest +from scipy.stats import ks_2samp + +if platform.system() == "Darwin": + import matplotlib + + matplotlib.use("TkAgg") + +import os +import sys + +import matplotlib.pyplot as plt +import tensorflow as tf +import uproot4 + +from phasespace import phasespace + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays, rapidsim # noqa: E402 +from .helpers.plotting import make_norm_histo # noqa: E402 + +BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +PLOT_DIR = os.path.join(BASE_PATH, "tests", "plots") + + +def setup_method(): + phasespace.GenParticle._sess.close() + tf.compat.v1.reset_default_graph() + + +def create_ref_histos(n_pions): + """Load reference histogram data.""" + ref_dir = os.path.join(BASE_PATH, "data") + if not os.path.exists(ref_dir): + os.mkdir(ref_dir) + ref_file = os.path.join(ref_dir, f"bto{n_pions}pi.root") + if not os.path.exists(ref_file): + script = os.path.join( + BASE_PATH, + "scripts", + "prepare_test_samples.cxx+({})".format( + ",".join( + '"{}"'.format(os.path.join(BASE_PATH, "data", f"bto{i + 1}pi.root")) + for i in range(1, 4) + ) + ), + ) + subprocess.call(f"root -qb '{script}'", shell=True) + events = uproot4.open(ref_file)["events"] + pion_names = [f"pion_{pion + 1}" for pion in range(n_pions)] + pions = {pion_name: events[pion_name] for pion_name in pion_names} + weights = events["weight"] + normalized_histograms = [] + for pion in pions.values(): + pion_array = pion.array() + energy = pion_array.fE + momentum = pion_array.fP + for coord, array in enumerate([momentum.fX, momentum.fY, momentum.fZ, energy]): + numpy_array = np.array(array) + histogram = make_norm_histo( + numpy_array, + range_=(-3000 if coord % 4 != 3 else 0, 3000), + weights=weights, + ) + normalized_histograms.append(histogram) + + return normalized_histograms, make_norm_histo(weights, range_=(0, 1 + 1e-8)) + + +def run_test(n_particles, test_prefix): + first_run_n_events = 100 + main_run_n_events = 100000 + n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64) + + decay = phasespace.nbody_decay(decays.B0_MASS, [decays.PION_MASS] * n_particles) + generate = decay.generate(n_events) + weights1, _ = generate # only generate to test change in n_events + assert len(weights1) == first_run_n_events + + # change n_events and run again + n_events.assign(main_run_n_events) + weights, particles = decay.generate(n_events) + parts = np.concatenate( + [particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1 + ) + histos = [ + make_norm_histo( + parts[:, coord], + range_=(-3000 if coord % 4 != 3 else 0, 3000), + weights=weights, + ) + for coord in range(parts.shape[1]) + ] + weight_histos = make_norm_histo(weights, range_=(0, 1 + 1e-8)) + ref_histos, ref_weights = create_ref_histos(n_particles) + p_values = np.array( + [ + ks_2samp(histos[coord], ref_histos[coord])[1] + for coord, _ in enumerate(histos) + ] + + [ks_2samp(weight_histos, ref_weights)[1]] + ) + # Let's plot + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + for coord, _ in enumerate(histos): + plt.hist( + x if coord % 4 != 3 else e, + weights=histos[coord], + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histos[coord], + alpha=0.5, + label="TGenPhasespace", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "{}_pion_{}_{}.png".format( + test_prefix, int(coord / 4) + 1, ["px", "py", "pz", "e"][coord % 4] + ), + ) + ) + plt.clf() + plt.hist( + np.linspace(0, 1, 100), + weights=weight_histos, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + np.linspace(0, 1, 100), + weights=ref_weights, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"{test_prefix}_weights.png")) + plt.clf() + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_two_body(): + """Test B->pipi decay.""" + run_test(2, "two_body") + + +@pytest.mark.flaky(3) # Stats are limited +def test_three_body(): + """Test B -> pi pi pi decay.""" + run_test(3, "three_body") + + +@pytest.mark.flaky(3) # Stats are limited +def test_four_body(): + """Test B -> pi pi pi pi decay.""" + run_test(4, "four_body") + + +def run_kstargamma(input_file, kstar_width, b_at_rest, suffix): + """Run B0->K*gamma test.""" + n_events = 1000000 + if b_at_rest: + booster = None + rapidsim_getter = rapidsim.get_tree_in_b_rest_frame + else: + booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) + booster = booster.transpose() + rapidsim_getter = rapidsim.get_tree + decay = decays.b0_to_kstar_gamma(kstar_width=kstar_width) + norm_weights, particles = decay.generate(n_events=n_events, boost_to=booster) + rapidsim_parts = rapidsim_getter( + os.path.join(BASE_PATH, "data", input_file), + "B0_0", + ("Kst0_0", "gamma_0", "Kp_0", "pim_0"), + ) + name_matching = {"Kst0_0": "K*0", "gamma_0": "gamma", "Kp_0": "K+", "pim_0": "pi-"} + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + p_values = {} + for ref_name, ref_part in rapidsim_parts.items(): + tf_part = name_matching[ref_name] + ref_part = ref_part.transpose() # for consistency + for coord, coord_name in enumerate(("px", "py", "pz", "e")): + range_ = (-3000 if coord % 4 != 3 else 0, 3000) + ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) + tf_histo = make_norm_histo( + particles[tf_part][:, coord], range_=range_, weights=norm_weights + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=tf_histo, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histo, + alpha=0.5, + label="RapidSim", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "B0_Kstar_gamma_Kstar{}_{}_{}.png".format( + suffix, tf_part.replace("*", "star"), coord_name + ), + ) + ) + plt.clf() + p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] + plt.hist( + np.linspace(0, 1, 100), + weights=make_norm_histo(norm_weights, range_=(0, 1)), + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"B0_Kstar_gamma_Kstar{suffix}_weights.png")) + plt.clf() + return np.array(list(p_values.values())) + + +@pytest.mark.flaky(3) # Stats are limited +def test_kstargamma_kstarnonresonant_at_rest(): + """Test B0 -> K* gamma physics with fixed mass for K*.""" + p_values = run_kstargamma( + "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", 0, True, "NonResonant" + ) + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_kstargamma_kstarnonresonant_lhc(): + """Test B0 -> K* gamma physics with fixed mass for K* with LHC kinematics.""" + p_values = run_kstargamma( + "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", + 0, + False, + "NonResonant_LHC", + ) + assert np.all(p_values > 0.05) + + +def test_kstargamma_resonant_at_rest(): + """Test B0 -> K* gamma physics with Gaussian mass for K*. + + Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov + test wrt to RapidSim, so plots are generated and can be inspected by the user. However, small differences + are expected in the tails of the energy distributions of the kaon and the pion. + """ + run_kstargamma( + "B2KstGamma_RapidSim_7TeV_Tree.root", decays.KSTARZ_WIDTH, True, "Gaussian" + ) + + +def run_k1_gamma(input_file, k1_width, kstar_width, b_at_rest, suffix): + """Run B+ -> K1gamma test.""" + n_events = 1000000 + if b_at_rest: + booster = None + rapidsim_getter = rapidsim.get_tree_in_b_rest_frame + else: + booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) + booster = booster.transpose() + rapidsim_getter = rapidsim.get_tree + gamma = decays.bp_to_k1_kstar_pi_gamma(k1_width=k1_width, kstar_width=kstar_width) + norm_weights, particles = gamma.generate(n_events=n_events, boost_to=booster) + rapidsim_parts = rapidsim_getter( + os.path.join(BASE_PATH, "data", input_file), + "Bp_0", + ("K1_1270_p_0", "Kst0_0", "gamma_0", "Kp_0", "pim_0", "pip_0"), + ) + name_matching = { + "K1_1270_p_0": "K1+", + "Kst0_0": "K*0", + "gamma_0": "gamma", + "Kp_0": "K+", + "pim_0": "pi-", + "pip_0": "pi+", + } + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + p_values = {} + for ref_name, ref_part in rapidsim_parts.items(): + tf_part = name_matching[ref_name] + ref_part = ( + ref_part.transpose() + ) # to be consistent with internal shape (nevents, nobs) + for coord, coord_name in enumerate(("px", "py", "pz", "e")): + range_ = (-3000 if coord % 4 != 3 else 0, 3000) + ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) + tf_histo = make_norm_histo( + particles[tf_part][:, coord], range_=range_, weights=norm_weights + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=tf_histo, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histo, + alpha=0.5, + label="RapidSim", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "Bp_K1_gamma_K1Kstar{}_{}_{}.png".format( + suffix, tf_part.replace("*", "star"), coord_name + ), + ) + ) + plt.clf() + p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] + plt.hist( + np.linspace(0, 1, 100), + weights=make_norm_histo(norm_weights, range_=(0, 1)), + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"Bp_K1_gamma_K1Kstar{suffix}_weights.png")) + plt.clf() + return np.array(list(p_values.values())) + + +@pytest.mark.flaky(3) # Stats are limited +def test_k1gamma_kstarnonresonant_at_rest(): + """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances.""" + p_values = run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", + 0, + 0, + True, + "NonResonant", + ) + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_k1gamma_kstarnonresonant_lhc(): + """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances with LHC kinematics.""" + p_values = run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", + 0, + 0, + False, + "NonResonant_LHC", + ) + assert np.all(p_values > 0.05) + + +def test_k1gamma_resonant_at_rest(): + """Test B0 -> K1 (->K*pi) gamma physics. + + Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov + test wrt to RapidSim, so plots are generated and can be inspected by the user. + """ + run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_Tree.root", + decays.K1_WIDTH, + decays.KSTARZ_WIDTH, + True, + "Gaussian", + ) + + +if __name__ == "__main__": + test_two_body() + test_three_body() + test_four_body() + test_kstargamma_kstarnonresonant_at_rest() + test_kstargamma_kstarnonresonant_lhc() + test_kstargamma_resonant_at_rest() + test_k1gamma_kstarnonresonant_at_rest() + test_k1gamma_kstarnonresonant_lhc() + test_k1gamma_resonant_at_rest() + +# EOF diff --git a/tests/test_random.py b/tests/test_random.py index a0148d29..5581bbb8 100644 --- a/tests/test_random.py +++ b/tests/test_random.py @@ -1,27 +1,27 @@ -import numpy as np -import pytest -import tensorflow as tf - -import phasespace as phsp - - -@pytest.mark.parametrize( - "seed", [lambda: 15, lambda: tf.random.Generator.from_seed(15)] -) -def test_get_rng(seed): - rng1 = phsp.random.get_rng(seed()) - rng2 = phsp.random.get_rng(seed()) - rnd1_seeded = rng1.uniform_full_int(shape=(100,)) - rnd2_seeded = rng2.uniform_full_int(shape=(100,)) - - rng3 = phsp.random.get_rng() - rng4 = phsp.random.get_rng(seed()) - # advance rng4 by one step - _ = rng4.split(1) - - rnd3 = rng3.uniform_full_int(shape=(100,)) - rnd4 = rng4.uniform_full_int(shape=(100,)) - - np.testing.assert_array_equal(rnd1_seeded, rnd2_seeded) - assert not np.array_equal(rnd1_seeded, rnd3) - assert not np.array_equal(rnd4, rnd3) +import numpy as np +import pytest +import tensorflow as tf + +import phasespace as phsp + + +@pytest.mark.parametrize( + "seed", [lambda: 15, lambda: tf.random.Generator.from_seed(15)] +) +def test_get_rng(seed): + rng1 = phsp.random.get_rng(seed()) + rng2 = phsp.random.get_rng(seed()) + rnd1_seeded = rng1.uniform_full_int(shape=(100,)) + rnd2_seeded = rng2.uniform_full_int(shape=(100,)) + + rng3 = phsp.random.get_rng() + rng4 = phsp.random.get_rng(seed()) + # advance rng4 by one step + _ = rng4.split(1) + + rnd3 = rng3.uniform_full_int(shape=(100,)) + rnd4 = rng4.uniform_full_int(shape=(100,)) + + np.testing.assert_array_equal(rnd1_seeded, rnd2_seeded) + assert not np.array_equal(rnd1_seeded, rnd3) + assert not np.array_equal(rnd4, rnd3) From dec6a2393e0a6d9dc36d6f5ca5dd65f4a298d969 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sun, 10 Oct 2021 17:41:12 +0200 Subject: [PATCH 31/71] Rename fulldecay to fromdecay --- docs/fromdecay.ipynb | 716 +++++++++++++++++++++---------------------- 1 file changed, 358 insertions(+), 358 deletions(-) diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index ee248d7d..4eb738b1 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -1,358 +1,358 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "collapsed": true, - "pycharm": { - "is_executing": true, - "name": "#%% md\n" - } - }, - "source": [ - "# Tutorial for `fromdecay` functionality\n", - "This tutorial shows how `phasespace.fromdecay` can be used.\n", - "\n", - "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", - "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "# Import libraries\n", - "from pprint import pprint\n", - "\n", - "import zfit\n", - "from particle import Particle\n", - "from decaylanguage import DecFileParser, DecayChainViewer\n", - "# TODO rename\n", - "from phasespace.fulldecay import FullDecay" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Intro to DecayLanguage\n", - "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", - "\n", - "We will begin by parsing a .dec file using DecayLanguage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "parser = DecFileParser('../tests/fulldecay/example_decays.dec')\n", - "parser.parse()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_chain = parser.build_decay_chains(\"pi0\")\n", - "pprint(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DecayChainViewer(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Creating a FullDecay object\n", - "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_decay = FullDecay.from_dict(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", - "\n", - "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for probability, particle in pi0_decay.gen_particles:\n", - " print(f\"There is a probability of {probability} \"\n", - " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", - "\n", - "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "weights, events = pi0_decay.generate(n_events=10_000)\n", - "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can confirm that the counts above are close to the expected counts based on the probabilities. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Changing mass settings FullDecay\n", - "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", - "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", - "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", - "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", - "\n", - "### Constant vs variable mass\n", - "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", - "This will be demonsttrated with the decay below:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", - "DecayChainViewer(dsplus_chain)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", - " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$\\pi^0$ has a greater width than $D^0$. \n", - "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", - "# Loop over D0 and pi+ particles, see graph above\n", - "for particle in dstar_decay.gen_particles[0][1].children:\n", - " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", - " assert particle.has_fixed_mass\n", - " \n", - "# Loop over D+ and pi0. See above.\n", - "for particle in dstar_decay.gen_particles[1][1].children:\n", - " if particle.name == \"pi0\":\n", - " assert not particle.has_fixed_mass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Configuring mass fucntions\n", - "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_custom_mass_func = dsplus_chain.copy()\n", - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", - "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def custom_gauss(mass, width):\n", - " particle_mass = tf.cast(mass, tf.float64)\n", - " particle_width = tf.cast(width, tf.float64)\n", - " \n", - " # This is the actual mass function that will be returned\n", - " def mass_func(min_mass, max_mass, n_events):\n", - " min_mass = tf.cast(min_mass, tf.float64)\n", - " max_mass = tf.cast(max_mass, tf.float64)\n", - " # Use a zfit PDF\n", - " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", - " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", - " return tf.vectorized_map(\n", - " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", - " )\n", - "\n", - " return mass_func" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "\n", - "# Set the mass function of pi0 to the custom gaussian distribution \n", - "# when it decays into an electron-positron pair and a photon (gamma)\n", - "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.7" - } - }, - "nbformat": 4, - "nbformat_minor": 1 -} +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": true, + "pycharm": { + "is_executing": true, + "name": "#%% md\n" + } + }, + "source": [ + "# Tutorial for `fromdecay` functionality\n", + "This tutorial shows how `phasespace.fromdecay` can be used.\n", + "\n", + "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", + "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Import libraries\n", + "from pprint import pprint\n", + "\n", + "import zfit\n", + "from particle import Particle\n", + "from decaylanguage import DecFileParser, DecayChainViewer\n", + "\n", + "from phasespace.fromdecay import FullDecay" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Intro to DecayLanguage\n", + "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", + "\n", + "We will begin by parsing a .dec file using DecayLanguage:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "parser = DecFileParser('../tests/fromdecay/example_decays.dec')\n", + "parser.parse()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_chain = parser.build_decay_chains(\"pi0\")\n", + "pprint(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DecayChainViewer(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating a FullDecay object\n", + "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_decay = FullDecay.from_dict(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for probability, particle in pi0_decay.gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", + "\n", + "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "weights, events = pi0_decay.generate(n_events=10_000)\n", + "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can confirm that the counts above are close to the expected counts based on the probabilities. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Changing mass settings FullDecay\n", + "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", + "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", + "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", + "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", + "\n", + "### Constant vs variable mass\n", + "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", + "This will be demonsttrated with the decay below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", + "DecayChainViewer(dsplus_chain)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", + " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$\\pi^0$ has a greater width than $D^0$. \n", + "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", + "# Loop over D0 and pi+ particles, see graph above\n", + "for particle in dstar_decay.gen_particles[0][1].children:\n", + " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", + " assert particle.has_fixed_mass\n", + " \n", + "# Loop over D+ and pi0. See above.\n", + "for particle in dstar_decay.gen_particles[1][1].children:\n", + " if particle.name == \"pi0\":\n", + " assert not particle.has_fixed_mass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Configuring mass fucntions\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_custom_mass_func = dsplus_chain.copy()\n", + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def custom_gauss(mass, width):\n", + " particle_mass = tf.cast(mass, tf.float64)\n", + " particle_width = tf.cast(width, tf.float64)\n", + " \n", + " # This is the actual mass function that will be returned\n", + " def mass_func(min_mass, max_mass, n_events):\n", + " min_mass = tf.cast(min_mass, tf.float64)\n", + " max_mass = tf.cast(max_mass, tf.float64)\n", + " # Use a zfit PDF\n", + " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", + " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", + " return tf.vectorized_map(\n", + " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", + " )\n", + "\n", + " return mass_func" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "\n", + "# Set the mass function of pi0 to the custom gaussian distribution \n", + "# when it decays into an electron-positron pair and a photon (gamma)\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} From d14d69e1db80b6991be7c8a907bf18e4b8c66499 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 12 Oct 2021 11:51:15 +0200 Subject: [PATCH 32/71] Revert "Rename fulldecay to fromdecay" This reverts commit dec6a2393e0a6d9dc36d6f5ca5dd65f4a298d969. --- docs/fromdecay.ipynb | 716 +++++++++++++++++++++---------------------- 1 file changed, 358 insertions(+), 358 deletions(-) diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index 4eb738b1..ee248d7d 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -1,358 +1,358 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "collapsed": true, - "pycharm": { - "is_executing": true, - "name": "#%% md\n" - } - }, - "source": [ - "# Tutorial for `fromdecay` functionality\n", - "This tutorial shows how `phasespace.fromdecay` can be used.\n", - "\n", - "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", - "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "# Import libraries\n", - "from pprint import pprint\n", - "\n", - "import zfit\n", - "from particle import Particle\n", - "from decaylanguage import DecFileParser, DecayChainViewer\n", - "\n", - "from phasespace.fromdecay import FullDecay" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Intro to DecayLanguage\n", - "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", - "\n", - "We will begin by parsing a .dec file using DecayLanguage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "parser = DecFileParser('../tests/fromdecay/example_decays.dec')\n", - "parser.parse()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_chain = parser.build_decay_chains(\"pi0\")\n", - "pprint(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DecayChainViewer(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Creating a FullDecay object\n", - "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_decay = FullDecay.from_dict(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", - "\n", - "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for probability, particle in pi0_decay.gen_particles:\n", - " print(f\"There is a probability of {probability} \"\n", - " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", - "\n", - "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "weights, events = pi0_decay.generate(n_events=10_000)\n", - "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can confirm that the counts above are close to the expected counts based on the probabilities. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Changing mass settings FullDecay\n", - "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", - "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", - "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", - "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", - "\n", - "### Constant vs variable mass\n", - "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", - "This will be demonsttrated with the decay below:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", - "DecayChainViewer(dsplus_chain)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", - " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$\\pi^0$ has a greater width than $D^0$. \n", - "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", - "# Loop over D0 and pi+ particles, see graph above\n", - "for particle in dstar_decay.gen_particles[0][1].children:\n", - " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", - " assert particle.has_fixed_mass\n", - " \n", - "# Loop over D+ and pi0. See above.\n", - "for particle in dstar_decay.gen_particles[1][1].children:\n", - " if particle.name == \"pi0\":\n", - " assert not particle.has_fixed_mass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Configuring mass fucntions\n", - "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_custom_mass_func = dsplus_chain.copy()\n", - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", - "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def custom_gauss(mass, width):\n", - " particle_mass = tf.cast(mass, tf.float64)\n", - " particle_width = tf.cast(width, tf.float64)\n", - " \n", - " # This is the actual mass function that will be returned\n", - " def mass_func(min_mass, max_mass, n_events):\n", - " min_mass = tf.cast(min_mass, tf.float64)\n", - " max_mass = tf.cast(max_mass, tf.float64)\n", - " # Use a zfit PDF\n", - " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", - " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", - " return tf.vectorized_map(\n", - " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", - " )\n", - "\n", - " return mass_func" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "\n", - "# Set the mass function of pi0 to the custom gaussian distribution \n", - "# when it decays into an electron-positron pair and a photon (gamma)\n", - "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.7" - } - }, - "nbformat": 4, - "nbformat_minor": 1 -} +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": true, + "pycharm": { + "is_executing": true, + "name": "#%% md\n" + } + }, + "source": [ + "# Tutorial for `fromdecay` functionality\n", + "This tutorial shows how `phasespace.fromdecay` can be used.\n", + "\n", + "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", + "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Import libraries\n", + "from pprint import pprint\n", + "\n", + "import zfit\n", + "from particle import Particle\n", + "from decaylanguage import DecFileParser, DecayChainViewer\n", + "# TODO rename\n", + "from phasespace.fulldecay import FullDecay" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Intro to DecayLanguage\n", + "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", + "\n", + "We will begin by parsing a .dec file using DecayLanguage:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "parser = DecFileParser('../tests/fulldecay/example_decays.dec')\n", + "parser.parse()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_chain = parser.build_decay_chains(\"pi0\")\n", + "pprint(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DecayChainViewer(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating a FullDecay object\n", + "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_decay = FullDecay.from_dict(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for probability, particle in pi0_decay.gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", + "\n", + "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "weights, events = pi0_decay.generate(n_events=10_000)\n", + "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can confirm that the counts above are close to the expected counts based on the probabilities. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Changing mass settings FullDecay\n", + "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", + "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", + "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", + "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", + "\n", + "### Constant vs variable mass\n", + "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", + "This will be demonsttrated with the decay below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", + "DecayChainViewer(dsplus_chain)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", + " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$\\pi^0$ has a greater width than $D^0$. \n", + "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", + "# Loop over D0 and pi+ particles, see graph above\n", + "for particle in dstar_decay.gen_particles[0][1].children:\n", + " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", + " assert particle.has_fixed_mass\n", + " \n", + "# Loop over D+ and pi0. See above.\n", + "for particle in dstar_decay.gen_particles[1][1].children:\n", + " if particle.name == \"pi0\":\n", + " assert not particle.has_fixed_mass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Configuring mass fucntions\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_custom_mass_func = dsplus_chain.copy()\n", + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def custom_gauss(mass, width):\n", + " particle_mass = tf.cast(mass, tf.float64)\n", + " particle_width = tf.cast(width, tf.float64)\n", + " \n", + " # This is the actual mass function that will be returned\n", + " def mass_func(min_mass, max_mass, n_events):\n", + " min_mass = tf.cast(min_mass, tf.float64)\n", + " max_mass = tf.cast(max_mass, tf.float64)\n", + " # Use a zfit PDF\n", + " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", + " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", + " return tf.vectorized_map(\n", + " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", + " )\n", + "\n", + " return mass_func" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "\n", + "# Set the mass function of pi0 to the custom gaussian distribution \n", + "# when it decays into an electron-positron pair and a photon (gamma)\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} From bba9a5cd7c41dbad24dba810fb9156a8a237afae Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 12 Oct 2021 11:53:25 +0200 Subject: [PATCH 33/71] Revert "Rename fuldecay to fromdecay." This reverts commit 02549c5f31750d16c43768ffcb8b2f4a0a8e0ee8. --- .git_archival.txt | 2 +- .gitattributes | 10 +- .gitignore | 224 +-- AUTHORS.rst | 36 +- CHANGELOG.rst | 248 +-- CONTRIBUTING.rst | 250 +-- LICENSE | 48 +- MANIFEST.in | 28 +- README.rst | 466 ++--- benchmark/bench_phasespace.py | 280 +-- benchmark/bench_tgenphasespace.cxx | 42 +- benchmark/monitoring.py | 152 +- data/download_test_files.py | 98 +- docs/Makefile | 40 +- docs/authors.rst | 2 +- docs/conf.py | 510 +++--- docs/contributing.rst | 2 +- docs/history.rst | 2 +- docs/index.rst | 34 +- docs/make.bat | 72 +- docs/phasespace.rst | 36 +- docs/usage.rst | 406 ++--- paper/paper.bib | 230 +-- phasespace/__init__.py | 52 +- phasespace/backend.py | 42 +- phasespace/fromdecay/__init__.py | 14 - phasespace/fulldecay/__init__.py | 14 + .../{fromdecay => fulldecay}/fulldecay.py | 23 +- .../mass_functions.py | 0 phasespace/kinematics.py | 304 ++-- phasespace/phasespace.py | 1552 ++++++++--------- phasespace/random.py | 82 +- pyproject.toml | 26 +- scripts/prepare_test_samples.cxx | 246 +-- setup.cfg | 190 +- setup.py | 40 +- tests/conftest.py | 8 +- tests/{fromdecay => fulldecay}/__init__.py | 8 +- .../example_decay_chains.py | 2 +- .../example_decays.dec | 62 +- .../test_fulldecay.py | 16 +- .../test_mass_functions.py | 2 +- tests/helpers/decays.py | 160 +- tests/helpers/plotting.py | 56 +- tests/helpers/rapidsim.py | 222 +-- tests/test_chain.py | 274 +-- tests/test_generate.py | 172 +- tests/test_nbody_decay.py | 128 +- tests/test_physics.py | 822 ++++----- tests/test_random.py | 54 +- 50 files changed, 3892 insertions(+), 3897 deletions(-) delete mode 100644 phasespace/fromdecay/__init__.py create mode 100644 phasespace/fulldecay/__init__.py rename phasespace/{fromdecay => fulldecay}/fulldecay.py (90%) rename phasespace/{fromdecay => fulldecay}/mass_functions.py (100%) rename tests/{fromdecay => fulldecay}/__init__.py (63%) rename tests/{fromdecay => fulldecay}/example_decay_chains.py (87%) rename tests/{fromdecay => fulldecay}/example_decays.dec (96%) rename tests/{fromdecay => fulldecay}/test_fulldecay.py (84%) rename tests/{fromdecay => fulldecay}/test_mass_functions.py (94%) diff --git a/.git_archival.txt b/.git_archival.txt index 2e6e85ba..95cb3eea 100644 --- a/.git_archival.txt +++ b/.git_archival.txt @@ -1 +1 @@ -ref-names: $Format:%D$ +ref-names: $Format:%D$ diff --git a/.gitattributes b/.gitattributes index 4663edcf..7f892985 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,5 +1,5 @@ -data/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2K1Gamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text -data/B2KstGamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text -.git_archival.txt export-subst +data/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2K1Gamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text +data/B2KstGamma_RapidSim_7TeV_Tree.root filter=lfs diff=lfs merge=lfs -text +.git_archival.txt export-subst diff --git a/.gitignore b/.gitignore index 47e05415..363a46df 100644 --- a/.gitignore +++ b/.gitignore @@ -1,112 +1,112 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -env/ -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# pyenv -.python-version - -# celery beat schedule file -celerybeat-schedule - -# SageMath parsed files -*.sage.py - -# dotenv -.env - -# virtualenv -.venv -venv/ -ENV/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy and pyright -.mypy_cache/ -/typings/ -pyrightconfig.json - -# specific -*/*_cxx* -tests/plots -/data/backup/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root -/data/backup/B2K1Gamma_RapidSim_7TeV_Tree.root -/data/backup/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root -/data/backup/B2KstGamma_RapidSim_7TeV_Tree.root +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# dotenv +.env + +# virtualenv +.venv +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy and pyright +.mypy_cache/ +/typings/ +pyrightconfig.json + +# specific +*/*_cxx* +tests/plots +/data/backup/B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root +/data/backup/B2K1Gamma_RapidSim_7TeV_Tree.root +/data/backup/B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root +/data/backup/B2KstGamma_RapidSim_7TeV_Tree.root diff --git a/AUTHORS.rst b/AUTHORS.rst index e98f898f..11a8d85e 100644 --- a/AUTHORS.rst +++ b/AUTHORS.rst @@ -1,18 +1,18 @@ -========== -Credits -========== - -Development Lead ----------------- - -* Albert Puig Navarro - -Core Developers ---------------- - -* Jonas Eschle - -Contributors ------------- - -None yet. Why not be the first? +========== +Credits +========== + +Development Lead +---------------- + +* Albert Puig Navarro + +Core Developers +--------------- + +* Jonas Eschle + +Contributors +------------ + +None yet. Why not be the first? diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f0339867..3dbc8e5c 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,124 +1,124 @@ -********* -Changelog -********* - -Develop -========== - - -Major Features and Improvements -------------------------------- - -Behavioral changes ------------------- - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - - -Thanks ------- - -1.4.1 (27.08.2021) -================== - -Requirement changes -------------------- -- Losen restriction on TensorFlow, allow version 2.6 (and 2.5) - -1.4.0 (11.06.2021) -================== - -Requirement changes -------------------- -- require TensorFlow 2.5 as 2.4 breaks some functionality - -1.3.0 (28.05.2021) -=================== - - -Major Features and Improvements -------------------------------- - -- Support Python 3.9 -- Support TensorFlow 2.5 -- improved compilation in tf.functions, use of XLA where applicable -- developer: modernization of setup, CI and more - -Thanks ------- - -- Remco de Boer for many commits and cleanups - -1.2.0 (17.12.20) -================ - - -Major Features and Improvements -------------------------------- - -- Python 3.8 support -- Allow eager execution by setting with `tf.config.run_functions_eagerly(True)` - or the environment variable "PHASESPACE_EAGER" -- Deterministic random number generation via seed - or `tf.random.Generator` instance - -Behavioral changes ------------------- - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - -- tighten TensorFlow to 2.3/2.4 -- tighten TensorFlow Probability to 0.11/0.12 - -Thanks ------- -- Remco de Boer and Stefan Pflüger for discussions on random number genration - -1.1.0 (27.1.2020) -================= - -This release switched to TensorFlow 2.0 eager mode. Please upgrade your TensorFlow installation if possible and change -your code (minimal changes) as described under "Behavioral changes". -In case this is currently impossible to do, please downgrade to < 1.1.0. - -Major Features and Improvements -------------------------------- - - full TF2 compatibility - -Behavioral changes ------------------- - - `generate` now returns an eager Tensor. This is basically a numpy array wrapped by TensorFlow. - To explicitly convert it to a numpy array, use the `numpy()` method of the eager Tensor. - - `generate_tensor` is now depreceated, `generate` can directly be used instead. - - -Bug fixes and small changes ---------------------------- - -Requirement changes -------------------- - - requires now TensorFlow >= 2.0.0 - - -Thanks ------- - - -1.0.4 (13-10-2019) -========================== - - -Major Features and Improvements -------------------------------- - -Release to conda-forge, thanks to Chris Burr +********* +Changelog +********* + +Develop +========== + + +Major Features and Improvements +------------------------------- + +Behavioral changes +------------------ + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + + +Thanks +------ + +1.4.1 (27.08.2021) +================== + +Requirement changes +------------------- +- Losen restriction on TensorFlow, allow version 2.6 (and 2.5) + +1.4.0 (11.06.2021) +================== + +Requirement changes +------------------- +- require TensorFlow 2.5 as 2.4 breaks some functionality + +1.3.0 (28.05.2021) +=================== + + +Major Features and Improvements +------------------------------- + +- Support Python 3.9 +- Support TensorFlow 2.5 +- improved compilation in tf.functions, use of XLA where applicable +- developer: modernization of setup, CI and more + +Thanks +------ + +- Remco de Boer for many commits and cleanups + +1.2.0 (17.12.20) +================ + + +Major Features and Improvements +------------------------------- + +- Python 3.8 support +- Allow eager execution by setting with `tf.config.run_functions_eagerly(True)` + or the environment variable "PHASESPACE_EAGER" +- Deterministic random number generation via seed + or `tf.random.Generator` instance + +Behavioral changes +------------------ + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + +- tighten TensorFlow to 2.3/2.4 +- tighten TensorFlow Probability to 0.11/0.12 + +Thanks +------ +- Remco de Boer and Stefan Pflüger for discussions on random number genration + +1.1.0 (27.1.2020) +================= + +This release switched to TensorFlow 2.0 eager mode. Please upgrade your TensorFlow installation if possible and change +your code (minimal changes) as described under "Behavioral changes". +In case this is currently impossible to do, please downgrade to < 1.1.0. + +Major Features and Improvements +------------------------------- + - full TF2 compatibility + +Behavioral changes +------------------ + - `generate` now returns an eager Tensor. This is basically a numpy array wrapped by TensorFlow. + To explicitly convert it to a numpy array, use the `numpy()` method of the eager Tensor. + - `generate_tensor` is now depreceated, `generate` can directly be used instead. + + +Bug fixes and small changes +--------------------------- + +Requirement changes +------------------- + - requires now TensorFlow >= 2.0.0 + + +Thanks +------ + + +1.0.4 (13-10-2019) +========================== + + +Major Features and Improvements +------------------------------- + +Release to conda-forge, thanks to Chris Burr diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index 88e9e2ca..88d4d412 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -1,125 +1,125 @@ -.. highlight:: shell - -============ -Contributing -============ - -Contributions are welcome, and they are greatly appreciated! Every little bit -helps, and credit will always be given. - -You can contribute in many ways: - -Types of Contributions ----------------------- - -Report Bugs -~~~~~~~~~~~ - -Report bugs at https://github.com/zfit/phasespace/issues. - -If you are reporting a bug, please include: - -* Your operating system name and version. -* Any details about your local setup that might be helpful in troubleshooting. -* Detailed steps to reproduce the bug. - -Fix Bugs -~~~~~~~~ - -Look through the GitHub issues for bugs. Anything tagged with "bug" and "help -wanted" is open to whoever wants to implement it. - -Implement Features -~~~~~~~~~~~~~~~~~~ - -Look through the GitHub issues for features. Anything tagged with "enhancement" -and "help wanted" is open to whoever wants to implement it. - -Write Documentation -~~~~~~~~~~~~~~~~~~~ - -TensorFlow PhaseSpace could always use more documentation, whether as part of the -official TensorFlow PhaseSpace docs, in docstrings, or even on the web in blog posts, -articles, and such. - -Submit Feedback -~~~~~~~~~~~~~~~ - -The best way to send feedback is to file an issue at https://github.com/zfit/phasespace/issues. - -If you are proposing a feature: - -* Explain in detail how it would work. -* Keep the scope as narrow as possible, to make it easier to implement. -* Remember that this is a volunteer-driven project, and that contributions - are welcome :) - -Get Started! ------------- - -Ready to contribute? Here's how to set up `phasespace` for local development. - -1. Fork the `phasespace` repo on GitHub. -2. Clone your fork locally:: - - $ git clone git@github.com:your_name_here/phasespace.git - -3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: - - $ mkvirtualenv phasespace - $ cd phasespace/ - $ pip install -e .[dev] - -4. Create a branch for local development:: - - $ git checkout -b name-of-your-bugfix-or-feature - - Now you can make your changes locally. - -5. When you're done making changes, check that your changes pass pre-commit and the - tests: - - $ pytest - $ pre-commit run -a - -6. Commit your changes and push your branch to GitHub:: - - $ git add . - $ git commit -m "Your detailed description of your changes." - $ git push origin name-of-your-bugfix-or-feature - -7. Submit a pull request through the GitHub website. - -Pull Request Guidelines ------------------------ - -Before you submit a pull request, check that it meets these guidelines: - -1. The pull request should include tests. -2. If the pull request adds functionality, the docs should be updated. Put - your new functionality into a function with a docstring, and add the - feature to the list in README.rst. -3. The pull request should work for Python 2.7, 3.4, 3.5 and 3.6, and for PyPy. Check - https://github.com/zfit/phasespace/actions/workflows/ci.yml - and make sure that the tests pass for all supported Python versions. - -Tips ----- - -To run a subset of tests (for example those in `tests/test_generate.py`):: - - - $ pytest -k test_generate - -Deploying ---------- - -A reminder for the maintainers on how to deploy. -Make sure all your changes are committed (including an entry in HISTORY.rst). -Then run:: - -$ bumpversion patch # possible: major / minor / patch -$ git push -$ git push --tags - -GitHub Actions will then deploy to PyPI if tests pass. +.. highlight:: shell + +============ +Contributing +============ + +Contributions are welcome, and they are greatly appreciated! Every little bit +helps, and credit will always be given. + +You can contribute in many ways: + +Types of Contributions +---------------------- + +Report Bugs +~~~~~~~~~~~ + +Report bugs at https://github.com/zfit/phasespace/issues. + +If you are reporting a bug, please include: + +* Your operating system name and version. +* Any details about your local setup that might be helpful in troubleshooting. +* Detailed steps to reproduce the bug. + +Fix Bugs +~~~~~~~~ + +Look through the GitHub issues for bugs. Anything tagged with "bug" and "help +wanted" is open to whoever wants to implement it. + +Implement Features +~~~~~~~~~~~~~~~~~~ + +Look through the GitHub issues for features. Anything tagged with "enhancement" +and "help wanted" is open to whoever wants to implement it. + +Write Documentation +~~~~~~~~~~~~~~~~~~~ + +TensorFlow PhaseSpace could always use more documentation, whether as part of the +official TensorFlow PhaseSpace docs, in docstrings, or even on the web in blog posts, +articles, and such. + +Submit Feedback +~~~~~~~~~~~~~~~ + +The best way to send feedback is to file an issue at https://github.com/zfit/phasespace/issues. + +If you are proposing a feature: + +* Explain in detail how it would work. +* Keep the scope as narrow as possible, to make it easier to implement. +* Remember that this is a volunteer-driven project, and that contributions + are welcome :) + +Get Started! +------------ + +Ready to contribute? Here's how to set up `phasespace` for local development. + +1. Fork the `phasespace` repo on GitHub. +2. Clone your fork locally:: + + $ git clone git@github.com:your_name_here/phasespace.git + +3. Install your local copy into a virtualenv. Assuming you have virtualenvwrapper installed, this is how you set up your fork for local development:: + + $ mkvirtualenv phasespace + $ cd phasespace/ + $ pip install -e .[dev] + +4. Create a branch for local development:: + + $ git checkout -b name-of-your-bugfix-or-feature + + Now you can make your changes locally. + +5. When you're done making changes, check that your changes pass pre-commit and the + tests: + + $ pytest + $ pre-commit run -a + +6. Commit your changes and push your branch to GitHub:: + + $ git add . + $ git commit -m "Your detailed description of your changes." + $ git push origin name-of-your-bugfix-or-feature + +7. Submit a pull request through the GitHub website. + +Pull Request Guidelines +----------------------- + +Before you submit a pull request, check that it meets these guidelines: + +1. The pull request should include tests. +2. If the pull request adds functionality, the docs should be updated. Put + your new functionality into a function with a docstring, and add the + feature to the list in README.rst. +3. The pull request should work for Python 2.7, 3.4, 3.5 and 3.6, and for PyPy. Check + https://github.com/zfit/phasespace/actions/workflows/ci.yml + and make sure that the tests pass for all supported Python versions. + +Tips +---- + +To run a subset of tests (for example those in `tests/test_generate.py`):: + + + $ pytest -k test_generate + +Deploying +--------- + +A reminder for the maintainers on how to deploy. +Make sure all your changes are committed (including an entry in HISTORY.rst). +Then run:: + +$ bumpversion patch # possible: major / minor / patch +$ git push +$ git push --tags + +GitHub Actions will then deploy to PyPI if tests pass. diff --git a/LICENSE b/LICENSE index d3037b76..5b001265 100644 --- a/LICENSE +++ b/LICENSE @@ -1,24 +1,24 @@ -Copyright (c) 2019, zfit -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - * Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - * Neither the name of the nor the - names of its contributors may be used to endorse or promote products - derived from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL BE LIABLE FOR ANY -DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND -ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +Copyright (c) 2019, zfit +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/MANIFEST.in b/MANIFEST.in index a383b4ba..534b0b22 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,14 +1,14 @@ -include AUTHORS.rst -include CONTRIBUTING.rst -include CHANGELOG.rst -include LICENSE -include README.rst -include pyproject.toml -include *.rst -include *.txt - -recursive-include tests * -recursive-exclude * __pycache__ -recursive-exclude * *.py[co] - -recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif +include AUTHORS.rst +include CONTRIBUTING.rst +include CHANGELOG.rst +include LICENSE +include README.rst +include pyproject.toml +include *.rst +include *.txt + +recursive-include tests * +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] + +recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif diff --git a/README.rst b/README.rst index b7b4fe72..1079a8e8 100644 --- a/README.rst +++ b/README.rst @@ -1,233 +1,233 @@ -******************************* -PhaseSpace -******************************* - -.. image:: https://joss.theoj.org/papers/10.21105/joss.01570/status.svg - :target: https://doi.org/10.21105/joss.01570 -.. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.2591993.svg - :target: https://doi.org/10.5281/zenodo.2591993 -.. image:: https://img.shields.io/pypi/status/phasespace.svg - :target: https://pypi.org/project/phasespace/ -.. image:: https://img.shields.io/pypi/pyversions/phasespace.svg - :target: https://pypi.org/project/phasespace/ -.. image:: https://github.com/zfit/phasespace/workflows/tests/badge.svg - :target: https://github.com/zfit/phasespace/actions/workflows/ci.yml?query=branch%3Amaster -.. image:: https://codecov.io/gh/zfit/phasespace/branch/master/graph/badge.svg - :target: https://codecov.io/gh/zfit/phasespace -.. image:: https://readthedocs.org/projects/phasespace/badge/?version=stable - :target: https://phasespace.readthedocs.io/en/latest/?badge=stable - :alt: Documentation Status -.. image:: https://badges.gitter.im/zfit/phasespace.svg - :target: https://gitter.im/zfit/phasespace?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge - :alt: Gitter chat - -Python implementation of the Raubold and Lynch method for `n`-body events using -TensorFlow as a backend. - -The code is based on the GENBOD function (W515 from CERNLIB), documented in [1] -and tries to follow it as closely as possible. - -Detailed documentation, including the API, can be found in https://phasespace.readthedocs.io. -Don't hesitate to join our `gitter`_ channel for questions and comments. - -If you use phasespace in a scientific publication we would appreciate citations to the `JOSS`_ publication: - -.. code-block:: bibtex - - @article{puig_eschle_phasespace-2019, - title = {phasespace: n-body phase space generation in Python}, - doi = {10.21105/joss.01570}, - url = {https://doi.org/10.21105/joss.01570}, - year = {2019}, - month = {oct}, - publisher = {The Open Journal}, - author = {Albert Puig and Jonas Eschle}, - journal = {Journal of Open Source Software} - } - -Free software: BSD-3-Clause. - -[1] F. James, Monte Carlo Phase Space, CERN 68-15 (1968) - -.. _JOSS: https://joss.theoj.org/papers/10.21105/joss.01570 -.. _Gitter: https://gitter.im/zfit/phasespace - - -Why? -==== -Lately, data analysis in High Energy Physics (HEP), traditionally performed within the `ROOT`_ ecosystem, -has been moving more and more towards Python. -The possibility of carrying out purely Python-based analyses has become real thanks to the -development of many open source Python packages, -which have allowed to replace most ROOT functionality with Python-based packages. - -One of the aspects where this is still not possible is in the random generation of `n`-body phase space events, -which are widely used in the field, for example to study kinematics -of the particle decays of interest, or to perform importance sampling in the case of complex amplitude models. -This has been traditionally done with the `TGenPhaseSpace`_ class, which is based of the GENBOD function of the -CERNLIB FORTRAN libraries and which requires a full working ROOT installation. - -This package aims to address this issue by providing a TensorFlow-based implementation of such a function -to generate `n`-body decays without requiring a ROOT installation. -Additionally, an oft-needed functionality to generate complex decay chains, not included in ``TGenPhaseSpace``, -is also offered, leaving room for decaying resonances (which don't have a fixed mass, but can be seen as a -broad peak). - -.. _ROOT: https://root.cern.ch -.. _TGenPhaseSpace: https://root.cern.ch/doc/master/classTGenPhaseSpace.html - -Installing -========== - -``phasespace`` is available on conda-forge and pip. - -To install ``phasespace`` with conda, run: - - -.. code-block:: console - - $ conda install phasespace -c conda-forge - -To install with pip: - -.. code-block:: console - - $ pip install phasespace - -This is the preferred method to install ``phasespace``, as it will always install the most recent stable release. - -For the newest development version, which may be unstable, you can install the version from git with - -.. code-block:: console - - $ pip install git+https://github.com/zfit/phasespace - - -How to use -========== - -The generation of simple `n`-body decays can be done using the ``nbody_decay`` shortcut to create a decay chain -with a very simple interface: one needs to pass the mass of the top particle and the masses of the children particle as -a list, optionally giving the names of the particles. Then, the `generate` method can be used to produce the desired sample. -For example, to generate :math:`B^0\to K\pi`, we would do: - -.. code-block:: python - - import phasespace - - B0_MASS = 5279.65 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - weights, particles = phasespace.nbody_decay(B0_MASS, - [PION_MASS, KAON_MASS]).generate(n_events=1000) - -Behind the scenes, this function runs the TensorFlow graph. It returns `tf.Tensor`, which, as TensorFlow 2.x is in eager mode, -is basically a numpy array. Any `tf.Tensor` can be explicitly converted to a numpy array by calling `tf.Tensor.numpy()` on it. -The `generate` function returns a `tf.Tensor` of 1000 elements in the case of ``weights`` and a list of -``n particles`` (2) arrays of (1000, 4) shape, -where each of the 4-dimensions corresponds to one of the components of the generated Lorentz 4-vector. -All particles are generated in the rest frame of the top particle; boosting to a certain momentum (or list of momenta) can be -achieved by passing the momenta to the ``boost_to`` argument. - -Sequential decays can be handled with the ``GenParticle`` class (used internally by ``generate``) and its ``set_children`` method. -As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which :math:`K^*\to K\pi`, we would write: - -.. code-block:: python - - from phasespace import GenParticle - - B0_MASS = 5279.65 - KSTARZ_MASS = 895.55 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - kaon = GenParticle('K+', KAON_MASS) - pion = GenParticle('pi-', PION_MASS) - kstar = GenParticle('K*', KSTARZ_MASS).set_children(kaon, pion) - gamma = GenParticle('gamma', 0) - bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) - - weights, particles = bz.generate(n_events=1000) - -Where we have used the fact that ``set_children`` returns the parent particle. -In this case, ``particles`` is a ``dict`` with the particle names as keys: - -.. code-block:: pycon - - >>> particles - {'K*': array([[ 1732.79325872, -1632.88873127, 950.85807735, 2715.78804872], - [-1633.95329448, 239.88921123, -1961.0402768 , 2715.78804872], - [ 407.15613764, -2236.6569286 , -1185.16616251, 2715.78804872], - ..., - [ 1091.64603395, -1301.78721269, 1920.07503991, 2715.78804872], - [ -517.3125083 , 1901.39296899, 1640.15905194, 2715.78804872], - [ 656.56413668, -804.76922982, 2343.99214816, 2715.78804872]]), - 'K+': array([[ 750.08077976, -547.22569019, 224.6920906 , 1075.30490935], - [-1499.90049089, 289.19714633, -1935.27960292, 2514.43047106], - [ 97.64746732, -1236.68112923, -381.09526192, 1388.47607911], - ..., - [ 508.66157459, -917.93523639, 1474.7064148 , 1876.11771642], - [ -212.28646168, 540.26381432, 610.86656669, 976.63988936], - [ 177.16656666, -535.98777569, 946.12636904, 1207.28744488]]), - 'gamma': array([[-1732.79325872, 1632.88873127, -950.85807735, 2563.79195128], - [ 1633.95329448, -239.88921123, 1961.0402768 , 2563.79195128], - [ -407.15613764, 2236.6569286 , 1185.16616251, 2563.79195128], - ..., - [-1091.64603395, 1301.78721269, -1920.07503991, 2563.79195128], - [ 517.3125083 , -1901.39296899, -1640.15905194, 2563.79195128], - [ -656.56413668, 804.76922982, -2343.99214816, 2563.79195128]]), - 'pi-': array([[ 982.71247896, -1085.66304109, 726.16598675, 1640.48313937], - [ -134.0528036 , -49.3079351 , -25.76067389, 201.35757766], - [ 309.50867032, -999.97579937, -804.0709006 , 1327.31196961], - ..., - [ 582.98445936, -383.85197629, 445.36862511, 839.6703323 ], - [ -305.02604662, 1361.12915468, 1029.29248526, 1739.14815935], - [ 479.39757002, -268.78145413, 1397.86577911, 1508.50060384]])} - -The `GenParticle` class is able to cache the graphs so it is possible to generate in a loop -without overhead: - -.. code-block:: pycon - - for i in range(10): - weights, particles = bz.generate(n_events=1000) - ... - (do something with weights and particles) - ... - -This way of generating is recommended in the case of large samples, as it allows to benefit from -parallelisation while at the same time keep the memory usage low. - -If we want to operate with the TensorFlow graph instead, we can use the `generate_tensor` method -of `GenParticle`, which has the same signature as `generate`. - -More examples can be found in the ``tests`` folder and in the `documentation`_. - -.. _documentation: https://phasespace.readthedocs.io/en/latest/usage.html - - -Physics validation -================== - -Physics validation is performed continuously in the included tests (``tests/test_physics.py``), run through GitHub Actions. -This validation is performed at two levels: - -- In simple `n`-body decays, the results of ``phasespace`` are checked against ``TGenPhaseSpace``. -- For sequential decays, the results of ``phasespace`` are checked against `RapidSim`_, a "fast Monte Carlo generator - for simulation of heavy-quark hadron decays". - In the case of resonances, differences are expected because our tests don't include proper modelling of their - mass shape, as it would require the introduction of - further dependencies. However, the results of the comparison can be expected visually. - -The results of all physics validation performed by the ``tests_physics.py`` test are written in ``tests/plots``. - -.. _RapidSim: https://github.com/gcowan/RapidSim/ - - -Contributing -============ - -Contributions are always welcome, please have a look at the `Contributing guide`_. - -.. _Contributing guide: CONTRIBUTING.rst +******************************* +PhaseSpace +******************************* + +.. image:: https://joss.theoj.org/papers/10.21105/joss.01570/status.svg + :target: https://doi.org/10.21105/joss.01570 +.. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.2591993.svg + :target: https://doi.org/10.5281/zenodo.2591993 +.. image:: https://img.shields.io/pypi/status/phasespace.svg + :target: https://pypi.org/project/phasespace/ +.. image:: https://img.shields.io/pypi/pyversions/phasespace.svg + :target: https://pypi.org/project/phasespace/ +.. image:: https://github.com/zfit/phasespace/workflows/tests/badge.svg + :target: https://github.com/zfit/phasespace/actions/workflows/ci.yml?query=branch%3Amaster +.. image:: https://codecov.io/gh/zfit/phasespace/branch/master/graph/badge.svg + :target: https://codecov.io/gh/zfit/phasespace +.. image:: https://readthedocs.org/projects/phasespace/badge/?version=stable + :target: https://phasespace.readthedocs.io/en/latest/?badge=stable + :alt: Documentation Status +.. image:: https://badges.gitter.im/zfit/phasespace.svg + :target: https://gitter.im/zfit/phasespace?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge + :alt: Gitter chat + +Python implementation of the Raubold and Lynch method for `n`-body events using +TensorFlow as a backend. + +The code is based on the GENBOD function (W515 from CERNLIB), documented in [1] +and tries to follow it as closely as possible. + +Detailed documentation, including the API, can be found in https://phasespace.readthedocs.io. +Don't hesitate to join our `gitter`_ channel for questions and comments. + +If you use phasespace in a scientific publication we would appreciate citations to the `JOSS`_ publication: + +.. code-block:: bibtex + + @article{puig_eschle_phasespace-2019, + title = {phasespace: n-body phase space generation in Python}, + doi = {10.21105/joss.01570}, + url = {https://doi.org/10.21105/joss.01570}, + year = {2019}, + month = {oct}, + publisher = {The Open Journal}, + author = {Albert Puig and Jonas Eschle}, + journal = {Journal of Open Source Software} + } + +Free software: BSD-3-Clause. + +[1] F. James, Monte Carlo Phase Space, CERN 68-15 (1968) + +.. _JOSS: https://joss.theoj.org/papers/10.21105/joss.01570 +.. _Gitter: https://gitter.im/zfit/phasespace + + +Why? +==== +Lately, data analysis in High Energy Physics (HEP), traditionally performed within the `ROOT`_ ecosystem, +has been moving more and more towards Python. +The possibility of carrying out purely Python-based analyses has become real thanks to the +development of many open source Python packages, +which have allowed to replace most ROOT functionality with Python-based packages. + +One of the aspects where this is still not possible is in the random generation of `n`-body phase space events, +which are widely used in the field, for example to study kinematics +of the particle decays of interest, or to perform importance sampling in the case of complex amplitude models. +This has been traditionally done with the `TGenPhaseSpace`_ class, which is based of the GENBOD function of the +CERNLIB FORTRAN libraries and which requires a full working ROOT installation. + +This package aims to address this issue by providing a TensorFlow-based implementation of such a function +to generate `n`-body decays without requiring a ROOT installation. +Additionally, an oft-needed functionality to generate complex decay chains, not included in ``TGenPhaseSpace``, +is also offered, leaving room for decaying resonances (which don't have a fixed mass, but can be seen as a +broad peak). + +.. _ROOT: https://root.cern.ch +.. _TGenPhaseSpace: https://root.cern.ch/doc/master/classTGenPhaseSpace.html + +Installing +========== + +``phasespace`` is available on conda-forge and pip. + +To install ``phasespace`` with conda, run: + + +.. code-block:: console + + $ conda install phasespace -c conda-forge + +To install with pip: + +.. code-block:: console + + $ pip install phasespace + +This is the preferred method to install ``phasespace``, as it will always install the most recent stable release. + +For the newest development version, which may be unstable, you can install the version from git with + +.. code-block:: console + + $ pip install git+https://github.com/zfit/phasespace + + +How to use +========== + +The generation of simple `n`-body decays can be done using the ``nbody_decay`` shortcut to create a decay chain +with a very simple interface: one needs to pass the mass of the top particle and the masses of the children particle as +a list, optionally giving the names of the particles. Then, the `generate` method can be used to produce the desired sample. +For example, to generate :math:`B^0\to K\pi`, we would do: + +.. code-block:: python + + import phasespace + + B0_MASS = 5279.65 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + weights, particles = phasespace.nbody_decay(B0_MASS, + [PION_MASS, KAON_MASS]).generate(n_events=1000) + +Behind the scenes, this function runs the TensorFlow graph. It returns `tf.Tensor`, which, as TensorFlow 2.x is in eager mode, +is basically a numpy array. Any `tf.Tensor` can be explicitly converted to a numpy array by calling `tf.Tensor.numpy()` on it. +The `generate` function returns a `tf.Tensor` of 1000 elements in the case of ``weights`` and a list of +``n particles`` (2) arrays of (1000, 4) shape, +where each of the 4-dimensions corresponds to one of the components of the generated Lorentz 4-vector. +All particles are generated in the rest frame of the top particle; boosting to a certain momentum (or list of momenta) can be +achieved by passing the momenta to the ``boost_to`` argument. + +Sequential decays can be handled with the ``GenParticle`` class (used internally by ``generate``) and its ``set_children`` method. +As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which :math:`K^*\to K\pi`, we would write: + +.. code-block:: python + + from phasespace import GenParticle + + B0_MASS = 5279.65 + KSTARZ_MASS = 895.55 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + kaon = GenParticle('K+', KAON_MASS) + pion = GenParticle('pi-', PION_MASS) + kstar = GenParticle('K*', KSTARZ_MASS).set_children(kaon, pion) + gamma = GenParticle('gamma', 0) + bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) + + weights, particles = bz.generate(n_events=1000) + +Where we have used the fact that ``set_children`` returns the parent particle. +In this case, ``particles`` is a ``dict`` with the particle names as keys: + +.. code-block:: pycon + + >>> particles + {'K*': array([[ 1732.79325872, -1632.88873127, 950.85807735, 2715.78804872], + [-1633.95329448, 239.88921123, -1961.0402768 , 2715.78804872], + [ 407.15613764, -2236.6569286 , -1185.16616251, 2715.78804872], + ..., + [ 1091.64603395, -1301.78721269, 1920.07503991, 2715.78804872], + [ -517.3125083 , 1901.39296899, 1640.15905194, 2715.78804872], + [ 656.56413668, -804.76922982, 2343.99214816, 2715.78804872]]), + 'K+': array([[ 750.08077976, -547.22569019, 224.6920906 , 1075.30490935], + [-1499.90049089, 289.19714633, -1935.27960292, 2514.43047106], + [ 97.64746732, -1236.68112923, -381.09526192, 1388.47607911], + ..., + [ 508.66157459, -917.93523639, 1474.7064148 , 1876.11771642], + [ -212.28646168, 540.26381432, 610.86656669, 976.63988936], + [ 177.16656666, -535.98777569, 946.12636904, 1207.28744488]]), + 'gamma': array([[-1732.79325872, 1632.88873127, -950.85807735, 2563.79195128], + [ 1633.95329448, -239.88921123, 1961.0402768 , 2563.79195128], + [ -407.15613764, 2236.6569286 , 1185.16616251, 2563.79195128], + ..., + [-1091.64603395, 1301.78721269, -1920.07503991, 2563.79195128], + [ 517.3125083 , -1901.39296899, -1640.15905194, 2563.79195128], + [ -656.56413668, 804.76922982, -2343.99214816, 2563.79195128]]), + 'pi-': array([[ 982.71247896, -1085.66304109, 726.16598675, 1640.48313937], + [ -134.0528036 , -49.3079351 , -25.76067389, 201.35757766], + [ 309.50867032, -999.97579937, -804.0709006 , 1327.31196961], + ..., + [ 582.98445936, -383.85197629, 445.36862511, 839.6703323 ], + [ -305.02604662, 1361.12915468, 1029.29248526, 1739.14815935], + [ 479.39757002, -268.78145413, 1397.86577911, 1508.50060384]])} + +The `GenParticle` class is able to cache the graphs so it is possible to generate in a loop +without overhead: + +.. code-block:: pycon + + for i in range(10): + weights, particles = bz.generate(n_events=1000) + ... + (do something with weights and particles) + ... + +This way of generating is recommended in the case of large samples, as it allows to benefit from +parallelisation while at the same time keep the memory usage low. + +If we want to operate with the TensorFlow graph instead, we can use the `generate_tensor` method +of `GenParticle`, which has the same signature as `generate`. + +More examples can be found in the ``tests`` folder and in the `documentation`_. + +.. _documentation: https://phasespace.readthedocs.io/en/latest/usage.html + + +Physics validation +================== + +Physics validation is performed continuously in the included tests (``tests/test_physics.py``), run through GitHub Actions. +This validation is performed at two levels: + +- In simple `n`-body decays, the results of ``phasespace`` are checked against ``TGenPhaseSpace``. +- For sequential decays, the results of ``phasespace`` are checked against `RapidSim`_, a "fast Monte Carlo generator + for simulation of heavy-quark hadron decays". + In the case of resonances, differences are expected because our tests don't include proper modelling of their + mass shape, as it would require the introduction of + further dependencies. However, the results of the comparison can be expected visually. + +The results of all physics validation performed by the ``tests_physics.py`` test are written in ``tests/plots``. + +.. _RapidSim: https://github.com/gcowan/RapidSim/ + + +Contributing +============ + +Contributions are always welcome, please have a look at the `Contributing guide`_. + +.. _Contributing guide: CONTRIBUTING.rst diff --git a/benchmark/bench_phasespace.py b/benchmark/bench_phasespace.py index 52bcc4f2..d9cae3b0 100644 --- a/benchmark/bench_phasespace.py +++ b/benchmark/bench_phasespace.py @@ -1,140 +1,140 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file bench_phasespace.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Benchmark phasespace.""" - -import os -import sys -from timeit import default_timer - -import tensorflow as tf - -from phasespace import phasespace - -sys.path.append(os.path.dirname(__file__)) - - -def memory_usage(): - """Get memory usage of current process in MiB. - - Tries to use :mod:`psutil`, if possible, otherwise fallback to calling - ``ps`` directly. - - Return: - float: Memory usage of the current process. - """ - pid = os.getpid() - try: - import psutil - - process = psutil.Process(pid) - mem = process.memory_info()[0] / float(2 ** 20) - except ImportError: - import subprocess - - out = ( - subprocess.Popen(["ps", "v", "-p", str(pid)], stdout=subprocess.PIPE) - .communicate()[0] - .split(b"\n") - ) - vsz_index = out[0].split().index(b"RSS") - mem = float(out[1].split()[vsz_index]) / 1024 - return mem - - -# pylint: disable=too-few-public-methods -class Timer: - """Time the code placed inside its context. - - Taken from http://coreygoldberg.blogspot.ch/2012/06/python-timer-class-context-manager-for.html - - Attributes: - verbose (bool): Print the elapsed time at context exit? - start (float): Start time in seconds since Epoch Time. Value set - to 0 if not run. - elapsed (float): Elapsed seconds in the timer. Value set to - 0 if not run. - - Arguments: - verbose (bool, optional): Print the elapsed time at - context exit? Defaults to False. - """ - - def __init__(self, verbose=False, n=1): - """Initialize the timer.""" - self.verbose = verbose - self.n = n - self._timer = default_timer - self.start = 0 - self.elapsed = 0 - - def __enter__(self): - self.start = self._timer() - return self - - def __exit__(self, *args): - self.elapsed = self._timer() - self.start - if self.verbose: - print(f"Elapsed time: {self.elapsed * 1000.0 / self.n} ms") - - -# EOF - - -# to play around with optimization, no big effect though -NUM_PARALLEL_EXEC_UNITS = 1 -# config = tf.ConfigProto( -# intra_op_parallelism_threads=NUM_PARALLEL_EXEC_UNITS, -# inter_op_parallelism_threads=1, -# allow_soft_placement=True, -# device_count={"CPU": NUM_PARALLEL_EXEC_UNITS}, -# ) - -B_MASS = 5279.0 -B_AT_REST = tf.stack((0.0, 0.0, 0.0, B_MASS), axis=-1) -PION_MASS = 139.6 - -N_EVENTS = 1000000 -CHUNK_SIZE = int(N_EVENTS) - -n_runs = 10 - - -# N_EVENTS_VAR = tf.Variable(initial_value=N_EVENTS) -# CHUNK_SIZE_VAR = tf.Variable(initial_value=CHUNK_SIZE) - - -def test_three_body(): - """Test B -> pi pi pi decay.""" - with Timer(verbose=True): - print("Initial run (may takes more time than consequent runs)") - do_run() # to get rid of initial overhead - print("starting benchmark") - with Timer(verbose=True, n=n_runs): - for _ in range(n_runs): - # CHUNK_SIZE_VAR.assign(CHUNK_SIZE + 1) # +1 to make sure we're not using any trivial caching - samples = do_run() - - print(f"nevents produced {samples[0][0].shape}") - print("Shape of one particle momentum", samples[0][1]["p_0"].shape) - - -decay = phasespace.nbody_decay( - B_MASS, - [PION_MASS, PION_MASS, PION_MASS], -) - - -# tf.config.run_functions_eagerly(True) -@tf.function(autograph=False) -def do_run(): - return [decay.generate(N_EVENTS) for _ in range(0, N_EVENTS, CHUNK_SIZE)] - - -if __name__ == "__main__": - test_three_body() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file bench_phasespace.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Benchmark phasespace.""" + +import os +import sys +from timeit import default_timer + +import tensorflow as tf + +from phasespace import phasespace + +sys.path.append(os.path.dirname(__file__)) + + +def memory_usage(): + """Get memory usage of current process in MiB. + + Tries to use :mod:`psutil`, if possible, otherwise fallback to calling + ``ps`` directly. + + Return: + float: Memory usage of the current process. + """ + pid = os.getpid() + try: + import psutil + + process = psutil.Process(pid) + mem = process.memory_info()[0] / float(2 ** 20) + except ImportError: + import subprocess + + out = ( + subprocess.Popen(["ps", "v", "-p", str(pid)], stdout=subprocess.PIPE) + .communicate()[0] + .split(b"\n") + ) + vsz_index = out[0].split().index(b"RSS") + mem = float(out[1].split()[vsz_index]) / 1024 + return mem + + +# pylint: disable=too-few-public-methods +class Timer: + """Time the code placed inside its context. + + Taken from http://coreygoldberg.blogspot.ch/2012/06/python-timer-class-context-manager-for.html + + Attributes: + verbose (bool): Print the elapsed time at context exit? + start (float): Start time in seconds since Epoch Time. Value set + to 0 if not run. + elapsed (float): Elapsed seconds in the timer. Value set to + 0 if not run. + + Arguments: + verbose (bool, optional): Print the elapsed time at + context exit? Defaults to False. + """ + + def __init__(self, verbose=False, n=1): + """Initialize the timer.""" + self.verbose = verbose + self.n = n + self._timer = default_timer + self.start = 0 + self.elapsed = 0 + + def __enter__(self): + self.start = self._timer() + return self + + def __exit__(self, *args): + self.elapsed = self._timer() - self.start + if self.verbose: + print(f"Elapsed time: {self.elapsed * 1000.0 / self.n} ms") + + +# EOF + + +# to play around with optimization, no big effect though +NUM_PARALLEL_EXEC_UNITS = 1 +# config = tf.ConfigProto( +# intra_op_parallelism_threads=NUM_PARALLEL_EXEC_UNITS, +# inter_op_parallelism_threads=1, +# allow_soft_placement=True, +# device_count={"CPU": NUM_PARALLEL_EXEC_UNITS}, +# ) + +B_MASS = 5279.0 +B_AT_REST = tf.stack((0.0, 0.0, 0.0, B_MASS), axis=-1) +PION_MASS = 139.6 + +N_EVENTS = 1000000 +CHUNK_SIZE = int(N_EVENTS) + +n_runs = 10 + + +# N_EVENTS_VAR = tf.Variable(initial_value=N_EVENTS) +# CHUNK_SIZE_VAR = tf.Variable(initial_value=CHUNK_SIZE) + + +def test_three_body(): + """Test B -> pi pi pi decay.""" + with Timer(verbose=True): + print("Initial run (may takes more time than consequent runs)") + do_run() # to get rid of initial overhead + print("starting benchmark") + with Timer(verbose=True, n=n_runs): + for _ in range(n_runs): + # CHUNK_SIZE_VAR.assign(CHUNK_SIZE + 1) # +1 to make sure we're not using any trivial caching + samples = do_run() + + print(f"nevents produced {samples[0][0].shape}") + print("Shape of one particle momentum", samples[0][1]["p_0"].shape) + + +decay = phasespace.nbody_decay( + B_MASS, + [PION_MASS, PION_MASS, PION_MASS], +) + + +# tf.config.run_functions_eagerly(True) +@tf.function(autograph=False) +def do_run(): + return [decay.generate(N_EVENTS) for _ in range(0, N_EVENTS, CHUNK_SIZE)] + + +if __name__ == "__main__": + test_three_body() + +# EOF diff --git a/benchmark/bench_tgenphasespace.cxx b/benchmark/bench_tgenphasespace.cxx index 4001a587..5f6211da 100644 --- a/benchmark/bench_tgenphasespace.cxx +++ b/benchmark/bench_tgenphasespace.cxx @@ -1,21 +1,21 @@ -#include "TROOT.h" -#include "TSystem.h" -#include "TFile.h" -#include "TTree.h" -#include "TLorentzVector.h" -#include "TGenPhaseSpace.h" - -Int_t N_EVENTS = 1000000; - -int bench_tgenphasespace() -{ - TLorentzVector B(0.0, 0.0, 0.0, 5279.0); - Double_t masses[3] = {139.6, 139.6, 139.6}; - - if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); - TGenPhaseSpace event; - event.SetDecay(B, 3, masses); - - for (Int_t n=0; nGetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); + TGenPhaseSpace event; + event.SetDecay(B, 3, masses); + + for (Int_t n=0; n google style - -# Napoleon settings (convert numpy/google docstrings to proper ReST -napoleon_google_docstring = not using_numpy_style -napoleon_numpy_docstring = using_numpy_style -napoleon_include_init_with_doc = False -napoleon_include_private_with_doc = False -napoleon_include_special_with_doc = True -napoleon_use_admonition_for_examples = False -napoleon_use_admonition_for_notes = False -napoleon_use_admonition_for_references = False -napoleon_use_ivar = False -napoleon_use_param = True -napoleon_use_rtype = True - -# Add any paths that contain templates here, relative to this directory. -templates_path = ["_templates"] - -# The suffix(es) of source filenames. -# You can specify multiple suffix as a list of string: -# -# source_suffix = ['.rst', '.md'] -source_suffix = ".rst" - -# The master toctree document. -master_doc = "index" - -# General information about the project. -project = "TensorFlow PhaseSpace" -copyright = "2019, Albert Puig Navarro" -author = "Albert Puig Navarro" - -# The version info for the project you're documenting, acts as replacement -# for |version| and |release|, also used in various other places throughout -# the built documents. -# -# The short X.Y version. -version = phasespace.__version__ -# The full version, including alpha/beta/rc tags. -release = phasespace.__version__ - -# The language for content autogenerated by Sphinx. Refer to documentation -# for a list of supported languages. -# -# This is also used if you do content translation via gettext catalogs. -# Usually you set "language" from the command line for these cases. -language = "en" - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - -# The name of the Pygments (syntax highlighting) style to use. -pygments_style = "sphinx" - -# If true, `todo` and `todoList` produce output, else they produce nothing. -todo_include_todos = False - -# makes the jupyter extension executable -jupyter_sphinx_thebelab_config = { - "requestKernel": True, - "binderOptions": { - "repo": "zfit/phasespace", - "binderUrl": "https://mybinder.org", - "repoProvider": "github", - }, -} - -# -- Options for HTML output ------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "bootstrap" -html_theme_path = sphinx_bootstrap_theme.get_html_theme_path() -html_theme_options = { - # Navigation bar title. (Default: ``project`` value) - "navbar_title": "phasespace", - # Tab name for entire site. (Default: "Site") - # 'navbar_site_name': "Docs", - # 'navbar_site_name': "Overview", - # A list of tuples containing pages or urls to link to. - # Valid tuples should be in the following forms: - # (name, page) # a link to a page - # (name, "/aa/bb", 1) # a link to an arbitrary relative url - # (name, "http://example.com", True) # arbitrary absolute url - # Note the "1" or "True" value above as the third argument to indicate - # an arbitrary url. - "navbar_links": [ - ("Phasespace", "index"), - ("Usage", "usage"), - ("API", "phasespace"), - ("Contributing", "contributing"), - # ("Link", "http://example.com", True), - ], - # Render the next and previous page links in navbar. (Default: true) - "navbar_sidebarrel": False, - # Render the current pages TOC in the navbar. (Default: true) - "navbar_pagenav": False, - # Tab name for the current pages TOC. (Default: "Page") - # 'navbar_pagenav_name': "Page", - # Global TOC depth for "site" navbar tab. (Default: 1) - # Switching to -1 shows all levels. - "globaltoc_depth": 1, - # Include hidden TOCs in Site navbar? - # - # Note: If this is "false", you cannot have mixed ``:hidden:`` and - # non-hidden ``toctree`` directives in the same page, or else the build - # will break. - # - # Values: "true" (default) or "false" - "globaltoc_includehidden": "true", - # HTML navbar class (Default: "navbar") to attach to
element. - # For black navbar, do "navbar navbar-inverse" - # 'navbar_class': "navbar navbar-inverse", - "navbar_class": "navbar", - # Fix navigation bar to top of page? - # Values: "true" (default) or "false" - "navbar_fixed_top": "true", - # Location of link to source. - # Options are "nav" (default), "footer" or anything else to exclude. - # 'source_link_position': "nav", - "source_link_position": False, - # Bootswatch (http://bootswatch.com/) theme. - # - # Options are nothing (default) or the name of a valid theme - # such as "cosmo" or "sandstone". - # - # The set of valid themes depend on the version of Bootstrap - # that's used (the next config option). - # - # Currently, the supported themes are: - # - Bootstrap 2: https://bootswatch.com/2 - # - Bootstrap 3: https://bootswatch.com/3 - "bootswatch_theme": "flatly", - # Choose Bootstrap version. - # Values: "3" (default) or "2" (in quotes) - "bootstrap_version": "4", -} -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = [] - -# -- Options for HTMLHelp output --------------------------------------- - -# Output file base name for HTML help builder. -htmlhelp_basename = "phasespacedoc" - -# -- Options for LaTeX output ------------------------------------------ - -latex_elements = { - # The paper size ('letterpaper' or 'a4paper'). - # - # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). - # - # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. - # - # 'preamble': '', - # Latex figure (float) alignment - # - # 'figure_align': 'htbp', -} - -# Grouping the document tree into LaTeX files. List of tuples -# (source start file, target name, title, author, documentclass -# [howto, manual, or own class]). -latex_documents = [ - ( - master_doc, - "phasespace.tex", - "TensorFlow PhaseSpace Documentation", - "Albert Puig Navarro", - "manual", - ), -] - -# -- Options for manual page output ------------------------------------ - -# One entry per manual page. List of tuples -# (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, "phasespace", "TensorFlow PhaseSpace Documentation", [author], 1) -] - -# -- Options for Texinfo output ---------------------------------------- - -# Grouping the document tree into Texinfo files. List of tuples -# (source start file, target name, title, author, -# dir menu entry, description, category) -texinfo_documents = [ - ( - master_doc, - "phasespace", - "TensorFlow PhaseSpace Documentation", - author, - "phasespace", - "One line description of project.", - "Miscellaneous", - ), -] +#!/usr/bin/env python +# +# phasespace documentation build configuration file, created by +# sphinx-quickstart on Fri Jun 9 13:47:02 2017. +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +# If extensions (or modules to document with autodoc) are in another +# directory, add these directories to sys.path here. If the directory is +# relative to the documentation root, use os.path.abspath to make it +# absolute, like shown here. +# + +import sphinx_bootstrap_theme + +import phasespace + +# -- General configuration --------------------------------------------- + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# use for classes the class and the __init__ docs combined +autoclass_content = "both" + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones. +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.viewcode", + "sphinx.ext.napoleon", + "sphinx.ext.mathjax", + "sphinx_math_dollar", + "jupyter_sphinx", +] + + +mathjax_config = { + "tex2jax": { + "inlineMath": [["\\(", "\\)"]], + "displayMath": [["\\[", "\\]"]], + }, +} +using_numpy_style = False # False -> google style + +# Napoleon settings (convert numpy/google docstrings to proper ReST +napoleon_google_docstring = not using_numpy_style +napoleon_numpy_docstring = using_numpy_style +napoleon_include_init_with_doc = False +napoleon_include_private_with_doc = False +napoleon_include_special_with_doc = True +napoleon_use_admonition_for_examples = False +napoleon_use_admonition_for_notes = False +napoleon_use_admonition_for_references = False +napoleon_use_ivar = False +napoleon_use_param = True +napoleon_use_rtype = True + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# +# source_suffix = ['.rst', '.md'] +source_suffix = ".rst" + +# The master toctree document. +master_doc = "index" + +# General information about the project. +project = "TensorFlow PhaseSpace" +copyright = "2019, Albert Puig Navarro" +author = "Albert Puig Navarro" + +# The version info for the project you're documenting, acts as replacement +# for |version| and |release|, also used in various other places throughout +# the built documents. +# +# The short X.Y version. +version = phasespace.__version__ +# The full version, including alpha/beta/rc tags. +release = phasespace.__version__ + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This patterns also effect to html_static_path and html_extra_path +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = False + +# makes the jupyter extension executable +jupyter_sphinx_thebelab_config = { + "requestKernel": True, + "binderOptions": { + "repo": "zfit/phasespace", + "binderUrl": "https://mybinder.org", + "repoProvider": "github", + }, +} + +# -- Options for HTML output ------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "bootstrap" +html_theme_path = sphinx_bootstrap_theme.get_html_theme_path() +html_theme_options = { + # Navigation bar title. (Default: ``project`` value) + "navbar_title": "phasespace", + # Tab name for entire site. (Default: "Site") + # 'navbar_site_name': "Docs", + # 'navbar_site_name': "Overview", + # A list of tuples containing pages or urls to link to. + # Valid tuples should be in the following forms: + # (name, page) # a link to a page + # (name, "/aa/bb", 1) # a link to an arbitrary relative url + # (name, "http://example.com", True) # arbitrary absolute url + # Note the "1" or "True" value above as the third argument to indicate + # an arbitrary url. + "navbar_links": [ + ("Phasespace", "index"), + ("Usage", "usage"), + ("API", "phasespace"), + ("Contributing", "contributing"), + # ("Link", "http://example.com", True), + ], + # Render the next and previous page links in navbar. (Default: true) + "navbar_sidebarrel": False, + # Render the current pages TOC in the navbar. (Default: true) + "navbar_pagenav": False, + # Tab name for the current pages TOC. (Default: "Page") + # 'navbar_pagenav_name': "Page", + # Global TOC depth for "site" navbar tab. (Default: 1) + # Switching to -1 shows all levels. + "globaltoc_depth": 1, + # Include hidden TOCs in Site navbar? + # + # Note: If this is "false", you cannot have mixed ``:hidden:`` and + # non-hidden ``toctree`` directives in the same page, or else the build + # will break. + # + # Values: "true" (default) or "false" + "globaltoc_includehidden": "true", + # HTML navbar class (Default: "navbar") to attach to
element. + # For black navbar, do "navbar navbar-inverse" + # 'navbar_class': "navbar navbar-inverse", + "navbar_class": "navbar", + # Fix navigation bar to top of page? + # Values: "true" (default) or "false" + "navbar_fixed_top": "true", + # Location of link to source. + # Options are "nav" (default), "footer" or anything else to exclude. + # 'source_link_position': "nav", + "source_link_position": False, + # Bootswatch (http://bootswatch.com/) theme. + # + # Options are nothing (default) or the name of a valid theme + # such as "cosmo" or "sandstone". + # + # The set of valid themes depend on the version of Bootstrap + # that's used (the next config option). + # + # Currently, the supported themes are: + # - Bootstrap 2: https://bootswatch.com/2 + # - Bootstrap 3: https://bootswatch.com/3 + "bootswatch_theme": "flatly", + # Choose Bootstrap version. + # Values: "3" (default) or "2" (in quotes) + "bootstrap_version": "4", +} +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = [] + +# -- Options for HTMLHelp output --------------------------------------- + +# Output file base name for HTML help builder. +htmlhelp_basename = "phasespacedoc" + +# -- Options for LaTeX output ------------------------------------------ + +latex_elements = { + # The paper size ('letterpaper' or 'a4paper'). + # + # 'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + # + # 'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + # + # 'preamble': '', + # Latex figure (float) alignment + # + # 'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, author, documentclass +# [howto, manual, or own class]). +latex_documents = [ + ( + master_doc, + "phasespace.tex", + "TensorFlow PhaseSpace Documentation", + "Albert Puig Navarro", + "manual", + ), +] + +# -- Options for manual page output ------------------------------------ + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + (master_doc, "phasespace", "TensorFlow PhaseSpace Documentation", [author], 1) +] + +# -- Options for Texinfo output ---------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + ( + master_doc, + "phasespace", + "TensorFlow PhaseSpace Documentation", + author, + "phasespace", + "One line description of project.", + "Miscellaneous", + ), +] diff --git a/docs/contributing.rst b/docs/contributing.rst index 819f45e0..e582053e 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -1 +1 @@ -.. include:: ../CONTRIBUTING.rst +.. include:: ../CONTRIBUTING.rst diff --git a/docs/history.rst b/docs/history.rst index 9ad7a341..565b0521 100644 --- a/docs/history.rst +++ b/docs/history.rst @@ -1 +1 @@ -.. include:: ../CHANGELOG.rst +.. include:: ../CHANGELOG.rst diff --git a/docs/index.rst b/docs/index.rst index 232f91e1..24a715e8 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,17 +1,17 @@ - -.. include:: ../README.rst - -.. include:: authors.rst - -================= -Table of Contents -================= - -.. toctree:: - :maxdepth: 1 - - usage - API - contributing - authors - history + +.. include:: ../README.rst + +.. include:: authors.rst + +================= +Table of Contents +================= + +.. toctree:: + :maxdepth: 1 + + usage + API + contributing + authors + history diff --git a/docs/make.bat b/docs/make.bat index 25284887..31e560f0 100644 --- a/docs/make.bat +++ b/docs/make.bat @@ -1,36 +1,36 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=python -msphinx -) -set SOURCEDIR=. -set BUILDDIR=_build -set SPHINXPROJ=phasespace - -if "%1" == "" goto help - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The Sphinx module was not found. Make sure you have Sphinx installed, - echo.then set the SPHINXBUILD environment variable to point to the full - echo.path of the 'sphinx-build' executable. Alternatively you may add the - echo.Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.http://sphinx-doc.org/ - exit /b 1 -) - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% - -:end -popd +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=python -msphinx +) +set SOURCEDIR=. +set BUILDDIR=_build +set SPHINXPROJ=phasespace + +if "%1" == "" goto help + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The Sphinx module was not found. Make sure you have Sphinx installed, + echo.then set the SPHINXBUILD environment variable to point to the full + echo.path of the 'sphinx-build' executable. Alternatively you may add the + echo.Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 +) + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% + +:end +popd diff --git a/docs/phasespace.rst b/docs/phasespace.rst index 02104012..ab905f15 100644 --- a/docs/phasespace.rst +++ b/docs/phasespace.rst @@ -1,18 +1,18 @@ -phasespace package -==================== - -phasespace.phasespace module --------------------------------- - -.. automodule:: phasespace.phasespace - :members: - :undoc-members: - :show-inheritance: - -phasespace.kinematics module ------------------------------- - -.. automodule:: phasespace.kinematics - :members: - :undoc-members: - :show-inheritance: +phasespace package +==================== + +phasespace.phasespace module +-------------------------------- + +.. automodule:: phasespace.phasespace + :members: + :undoc-members: + :show-inheritance: + +phasespace.kinematics module +------------------------------ + +.. automodule:: phasespace.kinematics + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/usage.rst b/docs/usage.rst index 11232063..ee6f0f38 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -1,203 +1,203 @@ -===== -Usage -===== - -The base of ``phasespace`` is the ``GenParticle`` object. -This object, which represents a particle, either stable or decaying, has only one mandatory argument, its name. - -In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either -a number or ``tf.constant``, or a function. -Functions are used to specify the mass of particles such as resonances, which are not fixed but vary according to -a broad distribution. -These mass functions get three arguments, and must return a ``TensorFlow`` Tensor: - -- The minimum mass allowed by the decay chain, which will be of shape `(n_events,)`. -- The maximum mass available, which will be of shape `(n_events,)`. -- The number of events to generate. - -This function signature allows to handle threshold effects cleanly, giving enough information to produce kinematically -allowed decays (NB: ``phasespace`` will throw an error if a kinematically forbidden decay is requested). - -A simple example --------------------------- - - -With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the ``GenParticle`` -class (which returns the class itself). As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which -:math:`K^*\to K\pi` with a fixed mass, we would write: - -.. jupyter-execute:: - - from phasespace import GenParticle - - B0_MASS = 5279.58 - KSTARZ_MASS = 895.81 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - pion = GenParticle('pi-', PION_MASS) - kaon = GenParticle('K+', KAON_MASS) - kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon) - gamma = GenParticle('gamma', 0) - bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) - -.. thebe-button:: Run this interactively - - -Phasespace events can be generated using the ``generate`` method, which gets the number of events to generate as input. -The method returns: - -- The normalized weights of each event, as an array of dimension (n_events,). -- The 4-momenta of the generated particles as values of a dictionary with the particle name as key. These momenta - are expressed as arrays of dimension (n_events, 4). - -.. jupyter-execute:: - - N_EVENTS = 1000 - - weights, particles = bz.generate(n_events=N_EVENTS) - -The ``generate`` method return an eager ``Tensor``: this is basically a wrapped numpy array. With ``Tensor.numpy()``, -it can always directly be converted to a numpy array (if really needed). - -Boosting the particles --------------------------- - - -The particles are generated in the rest frame of the top particle. -To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in both -``generate`` and ~`tf.Tensor`. This latter approach can be useful if the momentum of the top particle -is generated according to some distribution, for example the kinematics of the LHC (see ``test_kstargamma_kstarnonresonant_lhc`` -and ``test_k1gamma_kstarnonresonant_lhc`` in ``tests/test_physics.py`` to see how this could be done). - -Weights --------------------------- - - -Additionally, it is possible to obtain the unnormalized weights by using the ``generate_unnormalized`` flag in -``generate``. In this case, the method returns the unnormalized weights, the per-event maximum weight -and the particle dictionary. - -.. jupyter-execute:: - - print(particles) - - -It is worth noting that the graph generation is cached even when using ``generate``, so iterative generation -can be performed using normal python loops without loss in performance: - -.. jupyter-execute:: - - for i in range(5): - weights, particles = bz.generate(n_events=100) - # ... - # (do something with weights and particles) - # ... - - - -Resonances with variable mass ------------------------------- - - -To generate the mass of a resonance, we need to give a function as its mass instead of a floating number. -This function should take as input the per-event lower mass allowed, per-event upper mass allowed and the number of -events, and should return a ~`tf.Tensor` with the generated masses and shape (nevents,). Well suited for this task -are the `TensorFlow Probability distributions `_ -or, for more customized mass shapes, the -`zfit pdfs `_ (currently an -*experimental feature* is needed, contact the `zfit developers `_ to learn more). - -Following with the same example as above, and approximating the resonance shape by a gaussian, we could -write the :math:`B^{0}\to K^{*}\gamma` decay chain as (more details can be found in ``tests/helpers/decays.py``): - -.. jupyter-execute:: - :hide-output: - - import tensorflow as tf - import tensorflow_probability as tfp - from phasespace import GenParticle - - KSTARZ_MASS = 895.81 - KSTARZ_WIDTH = 47.4 - - def kstar_mass(min_mass, max_mass, n_events): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - if KSTARZ_WIDTH > 0: - kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, - scale=kstar_width_cast, - low=min_mass, - high=max_mass).sample() - return kstar_mass - - bz = GenParticle('B0', B0_MASS).set_children(GenParticle('K*0', mass=kstar_mass) - .set_children(GenParticle('K+', mass=KAON_MASS), - GenParticle('pi-', mass=PION_MASS)), - GenParticle('gamma', mass=0.0)) - - bz.generate(n_events=500) - - -Shortcut for simple decays --------------------------- - -The generation of simple `n`-body decay chains can be done using the ``nbody_decay`` function of ``phasespace``, which takes - -- The mass of the top particle. -- The mass of children particles as a list. -- The name of the top particle (optional). -- The names of the children particles (optional). - -If the names are not given, `top` and `p_{i}` are assigned. For example, to generate :math:`B^0\to K\pi`, one would do: - -.. jupyter-execute:: - - import phasespace - - N_EVENTS = 1000 - - B0_MASS = 5279.58 - PION_MASS = 139.57018 - KAON_MASS = 493.677 - - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS], - top_name="B0", names=["pi", "K"]) - weights, particles = decay.generate(n_events=N_EVENTS) - -In this example, ``decay`` is simply a ``GenParticle`` with the corresponding children. - - -Eager execution ---------------- - -By default, `phasespace` uses TensorFlow to build a graph of the computations. This is usually more -performant, especially if used multiple times. However, this has the disadvantage that _inside_ -`phasespac`, the actual values are not computed on Python runtime, e.g. if a breakpoint is set -the values of a `tf.Tensor` won't be available. - -TensorFlow (since version 2.0) however can easily switch to so called "eager execution": in this -mode, it behaves the same as Numpy; values are computed instantly and the Python code is not only -executed once but every time. - -To switch this on or off, the global flag in TensorFlow `tf.config.run_functions_eagerly(True)` or -the enviroment variable "PHASESPACE_EAGER" (which switches this flag) can be used. - -Random numbers --------------- - -The random number generation inside `phasespace` is transparent in order to allow for deterministic -behavior if desired. A function that uses random number generation inside always takes a `seed` (or `rng`) -argument. The behavior is as follows - -- if no seed is given, the global random number generator of TensorFlow will be used. Setting this - instance explicitly or by setting the seed via `tf.random.set_seed` allows for a deterministic - execution of a whole _script_. -- if the seed is a number it will be used to create a random number generator from this. Using the - same seed again will result in the same output. -- if the seed is an instance of :py:class:`tf.random.Generator`, this instance will directly be used - and advances an undefined number of steps. +===== +Usage +===== + +The base of ``phasespace`` is the ``GenParticle`` object. +This object, which represents a particle, either stable or decaying, has only one mandatory argument, its name. + +In most cases (except for the top particle of a decay), one wants to also specify its mass, which can be either +a number or ``tf.constant``, or a function. +Functions are used to specify the mass of particles such as resonances, which are not fixed but vary according to +a broad distribution. +These mass functions get three arguments, and must return a ``TensorFlow`` Tensor: + +- The minimum mass allowed by the decay chain, which will be of shape `(n_events,)`. +- The maximum mass available, which will be of shape `(n_events,)`. +- The number of events to generate. + +This function signature allows to handle threshold effects cleanly, giving enough information to produce kinematically +allowed decays (NB: ``phasespace`` will throw an error if a kinematically forbidden decay is requested). + +A simple example +-------------------------- + + +With these considerations in mind, one can build a decay chain by using the ``set_children`` method of the ``GenParticle`` +class (which returns the class itself). As an example, to build the :math:`B^{0}\to K^{*}\gamma` decay in which +:math:`K^*\to K\pi` with a fixed mass, we would write: + +.. jupyter-execute:: + + from phasespace import GenParticle + + B0_MASS = 5279.58 + KSTARZ_MASS = 895.81 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + pion = GenParticle('pi-', PION_MASS) + kaon = GenParticle('K+', KAON_MASS) + kstar = GenParticle('K*', KSTARZ_MASS).set_children(pion, kaon) + gamma = GenParticle('gamma', 0) + bz = GenParticle('B0', B0_MASS).set_children(kstar, gamma) + +.. thebe-button:: Run this interactively + + +Phasespace events can be generated using the ``generate`` method, which gets the number of events to generate as input. +The method returns: + +- The normalized weights of each event, as an array of dimension (n_events,). +- The 4-momenta of the generated particles as values of a dictionary with the particle name as key. These momenta + are expressed as arrays of dimension (n_events, 4). + +.. jupyter-execute:: + + N_EVENTS = 1000 + + weights, particles = bz.generate(n_events=N_EVENTS) + +The ``generate`` method return an eager ``Tensor``: this is basically a wrapped numpy array. With ``Tensor.numpy()``, +it can always directly be converted to a numpy array (if really needed). + +Boosting the particles +-------------------------- + + +The particles are generated in the rest frame of the top particle. +To produce them at a given momentum of the top particle, one can pass these momenta with the ``boost_to`` argument in both +``generate`` and ~`tf.Tensor`. This latter approach can be useful if the momentum of the top particle +is generated according to some distribution, for example the kinematics of the LHC (see ``test_kstargamma_kstarnonresonant_lhc`` +and ``test_k1gamma_kstarnonresonant_lhc`` in ``tests/test_physics.py`` to see how this could be done). + +Weights +-------------------------- + + +Additionally, it is possible to obtain the unnormalized weights by using the ``generate_unnormalized`` flag in +``generate``. In this case, the method returns the unnormalized weights, the per-event maximum weight +and the particle dictionary. + +.. jupyter-execute:: + + print(particles) + + +It is worth noting that the graph generation is cached even when using ``generate``, so iterative generation +can be performed using normal python loops without loss in performance: + +.. jupyter-execute:: + + for i in range(5): + weights, particles = bz.generate(n_events=100) + # ... + # (do something with weights and particles) + # ... + + + +Resonances with variable mass +------------------------------ + + +To generate the mass of a resonance, we need to give a function as its mass instead of a floating number. +This function should take as input the per-event lower mass allowed, per-event upper mass allowed and the number of +events, and should return a ~`tf.Tensor` with the generated masses and shape (nevents,). Well suited for this task +are the `TensorFlow Probability distributions `_ +or, for more customized mass shapes, the +`zfit pdfs `_ (currently an +*experimental feature* is needed, contact the `zfit developers `_ to learn more). + +Following with the same example as above, and approximating the resonance shape by a gaussian, we could +write the :math:`B^{0}\to K^{*}\gamma` decay chain as (more details can be found in ``tests/helpers/decays.py``): + +.. jupyter-execute:: + :hide-output: + + import tensorflow as tf + import tensorflow_probability as tfp + from phasespace import GenParticle + + KSTARZ_MASS = 895.81 + KSTARZ_WIDTH = 47.4 + + def kstar_mass(min_mass, max_mass, n_events): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + if KSTARZ_WIDTH > 0: + kstar_mass = tfp.distributions.TruncatedNormal(loc=kstar_mass, + scale=kstar_width_cast, + low=min_mass, + high=max_mass).sample() + return kstar_mass + + bz = GenParticle('B0', B0_MASS).set_children(GenParticle('K*0', mass=kstar_mass) + .set_children(GenParticle('K+', mass=KAON_MASS), + GenParticle('pi-', mass=PION_MASS)), + GenParticle('gamma', mass=0.0)) + + bz.generate(n_events=500) + + +Shortcut for simple decays +-------------------------- + +The generation of simple `n`-body decay chains can be done using the ``nbody_decay`` function of ``phasespace``, which takes + +- The mass of the top particle. +- The mass of children particles as a list. +- The name of the top particle (optional). +- The names of the children particles (optional). + +If the names are not given, `top` and `p_{i}` are assigned. For example, to generate :math:`B^0\to K\pi`, one would do: + +.. jupyter-execute:: + + import phasespace + + N_EVENTS = 1000 + + B0_MASS = 5279.58 + PION_MASS = 139.57018 + KAON_MASS = 493.677 + + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS], + top_name="B0", names=["pi", "K"]) + weights, particles = decay.generate(n_events=N_EVENTS) + +In this example, ``decay`` is simply a ``GenParticle`` with the corresponding children. + + +Eager execution +--------------- + +By default, `phasespace` uses TensorFlow to build a graph of the computations. This is usually more +performant, especially if used multiple times. However, this has the disadvantage that _inside_ +`phasespac`, the actual values are not computed on Python runtime, e.g. if a breakpoint is set +the values of a `tf.Tensor` won't be available. + +TensorFlow (since version 2.0) however can easily switch to so called "eager execution": in this +mode, it behaves the same as Numpy; values are computed instantly and the Python code is not only +executed once but every time. + +To switch this on or off, the global flag in TensorFlow `tf.config.run_functions_eagerly(True)` or +the enviroment variable "PHASESPACE_EAGER" (which switches this flag) can be used. + +Random numbers +-------------- + +The random number generation inside `phasespace` is transparent in order to allow for deterministic +behavior if desired. A function that uses random number generation inside always takes a `seed` (or `rng`) +argument. The behavior is as follows + +- if no seed is given, the global random number generator of TensorFlow will be used. Setting this + instance explicitly or by setting the seed via `tf.random.set_seed` allows for a deterministic + execution of a whole _script_. +- if the seed is a number it will be used to create a random number generator from this. Using the + same seed again will result in the same output. +- if the seed is an instance of :py:class:`tf.random.Generator`, this instance will directly be used + and advances an undefined number of steps. diff --git a/paper/paper.bib b/paper/paper.bib index 88950958..7dd5b40c 100644 --- a/paper/paper.bib +++ b/paper/paper.bib @@ -1,115 +1,115 @@ -@article{Brun:1997pa, - author = "Brun, R. and Rademakers, F.", - title = "{ROOT: An object oriented data analysis framework}", - booktitle = "{New computing techniques in physics research V. - Proceedings, 5th International Workshop, AIHENP '96, - Lausanne, Switzerland, September 2-6, 1996}", - journal = "Nucl. Instrum. Meth.", - volume = "A389", - year = "1997", - pages = "81-86", - doi = "10.1016/S0168-9002(97)00048-X", - notes = "{See also \url{http://root.cern.ch/}}", - SLACcitation = "%%CITATION = NUIMA,A389,81;%%" -} - - -@article{Cowan:2016tnm, - author = "Cowan, G. A. and Craik, D. C. and Needham, M. D.", - title = "{RapidSim: an application for the fast simulation of - heavy-quark hadron decays}", - journal = "Comput. Phys. Commun.", - volume = "214", - year = "2017", - pages = "239-246", - doi = "10.1016/j.cpc.2017.01.029", - eprint = "1612.07489", - archivePrefix = "arXiv", - primaryClass = "hep-ex", - SLACcitation = "%%CITATION = ARXIV:1612.07489;%%" -} - -@article{James:1968gu, - author = "James, F.", - title = "{Monte-Carlo phase space}", - year = "1968", - reportNumber = "CERN-68-15", - SLACcitation = "%%CITATION = CERN-68-15;%%" -} - -@article{Poluektov:2266468, - author = "Poluektov, Anton", - title = "{Performing amplitude fits with TensorFlow: LHCb - experience. HEP analysis ecosystem workshop}", - month = "May", - year = "2017", - reportNumber = "LHCb-TALK-2017-134", - url = "https://cds.cern.ch/record/2266468", -} - -@misc{tensorflow2015-whitepaper, -title={ {TensorFlow}: Large-Scale Machine Learning on Heterogeneous Systems}, -url={https://www.tensorflow.org/}, -note={Software available from tensorflow.org}, -author={ - Mart\'{\i}n~Abadi and - Ashish~Agarwal and - Paul~Barham and - Eugene~Brevdo and - Zhifeng~Chen and - Craig~Citro and - Greg~S.~Corrado and - Andy~Davis and - Jeffrey~Dean and - Matthieu~Devin and - Sanjay~Ghemawat and - Ian~Goodfellow and - Andrew~Harp and - Geoffrey~Irving and - Michael~Isard and - Yangqing Jia and - Rafal~Jozefowicz and - Lukasz~Kaiser and - Manjunath~Kudlur and - Josh~Levenberg and - Dandelion~Man\'{e} and - Rajat~Monga and - Sherry~Moore and - Derek~Murray and - Chris~Olah and - Mike~Schuster and - Jonathon~Shlens and - Benoit~Steiner and - Ilya~Sutskever and - Kunal~Talwar and - Paul~Tucker and - Vincent~Vanhoucke and - Vijay~Vasudevan and - Fernanda~Vi\'{e}gas and - Oriol~Vinyals and - Pete~Warden and - Martin~Wattenberg and - Martin~Wicke and - Yuan~Yu and - Xiaoqiang~Zheng}, - year={2015}, -} - - - -@misc{zfit, - Author = {Jonas Eschle and Albert {Puig Navarro} and Rafael {Silva Coutinho}}, - Date-Added = {2019-03-12 14:59:41 +0100}, - Date-Modified = {2019-03-21 14:39:21 -0400}, - Doi = {10.5281/zenodo.2602043}, - Title = {{zfit: scalable pythonic fitting}}, - Year = {2019}} - - -@article{zenodo, - title={phasespace: n-body phase space generation in Python}, - DOI={10.5281/zenodo.2591993}, - publisher={Zenodo}, - author={Albert {Puig Navarro} and Jonas Eschle}, - year={2019}, - month={Mar}} +@article{Brun:1997pa, + author = "Brun, R. and Rademakers, F.", + title = "{ROOT: An object oriented data analysis framework}", + booktitle = "{New computing techniques in physics research V. + Proceedings, 5th International Workshop, AIHENP '96, + Lausanne, Switzerland, September 2-6, 1996}", + journal = "Nucl. Instrum. Meth.", + volume = "A389", + year = "1997", + pages = "81-86", + doi = "10.1016/S0168-9002(97)00048-X", + notes = "{See also \url{http://root.cern.ch/}}", + SLACcitation = "%%CITATION = NUIMA,A389,81;%%" +} + + +@article{Cowan:2016tnm, + author = "Cowan, G. A. and Craik, D. C. and Needham, M. D.", + title = "{RapidSim: an application for the fast simulation of + heavy-quark hadron decays}", + journal = "Comput. Phys. Commun.", + volume = "214", + year = "2017", + pages = "239-246", + doi = "10.1016/j.cpc.2017.01.029", + eprint = "1612.07489", + archivePrefix = "arXiv", + primaryClass = "hep-ex", + SLACcitation = "%%CITATION = ARXIV:1612.07489;%%" +} + +@article{James:1968gu, + author = "James, F.", + title = "{Monte-Carlo phase space}", + year = "1968", + reportNumber = "CERN-68-15", + SLACcitation = "%%CITATION = CERN-68-15;%%" +} + +@article{Poluektov:2266468, + author = "Poluektov, Anton", + title = "{Performing amplitude fits with TensorFlow: LHCb + experience. HEP analysis ecosystem workshop}", + month = "May", + year = "2017", + reportNumber = "LHCb-TALK-2017-134", + url = "https://cds.cern.ch/record/2266468", +} + +@misc{tensorflow2015-whitepaper, +title={ {TensorFlow}: Large-Scale Machine Learning on Heterogeneous Systems}, +url={https://www.tensorflow.org/}, +note={Software available from tensorflow.org}, +author={ + Mart\'{\i}n~Abadi and + Ashish~Agarwal and + Paul~Barham and + Eugene~Brevdo and + Zhifeng~Chen and + Craig~Citro and + Greg~S.~Corrado and + Andy~Davis and + Jeffrey~Dean and + Matthieu~Devin and + Sanjay~Ghemawat and + Ian~Goodfellow and + Andrew~Harp and + Geoffrey~Irving and + Michael~Isard and + Yangqing Jia and + Rafal~Jozefowicz and + Lukasz~Kaiser and + Manjunath~Kudlur and + Josh~Levenberg and + Dandelion~Man\'{e} and + Rajat~Monga and + Sherry~Moore and + Derek~Murray and + Chris~Olah and + Mike~Schuster and + Jonathon~Shlens and + Benoit~Steiner and + Ilya~Sutskever and + Kunal~Talwar and + Paul~Tucker and + Vincent~Vanhoucke and + Vijay~Vasudevan and + Fernanda~Vi\'{e}gas and + Oriol~Vinyals and + Pete~Warden and + Martin~Wattenberg and + Martin~Wicke and + Yuan~Yu and + Xiaoqiang~Zheng}, + year={2015}, +} + + + +@misc{zfit, + Author = {Jonas Eschle and Albert {Puig Navarro} and Rafael {Silva Coutinho}}, + Date-Added = {2019-03-12 14:59:41 +0100}, + Date-Modified = {2019-03-21 14:39:21 -0400}, + Doi = {10.5281/zenodo.2602043}, + Title = {{zfit: scalable pythonic fitting}}, + Year = {2019}} + + +@article{zenodo, + title={phasespace: n-body phase space generation in Python}, + DOI={10.5281/zenodo.2591993}, + publisher={Zenodo}, + author={Albert {Puig Navarro} and Jonas Eschle}, + year={2019}, + month={Mar}} diff --git a/phasespace/__init__.py b/phasespace/__init__.py index 5c190557..5132c914 100644 --- a/phasespace/__init__.py +++ b/phasespace/__init__.py @@ -1,26 +1,26 @@ -"""Top-level package for TensorFlow PhaseSpace.""" -from pkg_resources import get_distribution - -__author__ = """Albert Puig Navarro""" -__email__ = "apuignav@gmail.com" -__version__ = get_distribution(__name__).version -__maintainer__ = "zfit" - -__credits__ = ["Jonas Eschle "] - -__all__ = ["nbody_decay", "GenParticle", "random"] - -import tensorflow as tf - -from . import random -from .phasespace import GenParticle, nbody_decay - - -def _set_eager_mode(): - import os - - is_eager = bool(os.environ.get("PHASESPACE_EAGER")) - tf.config.run_functions_eagerly(is_eager) - - -_set_eager_mode() +"""Top-level package for TensorFlow PhaseSpace.""" +from pkg_resources import get_distribution + +__author__ = """Albert Puig Navarro""" +__email__ = "apuignav@gmail.com" +__version__ = get_distribution(__name__).version +__maintainer__ = "zfit" + +__credits__ = ["Jonas Eschle "] + +__all__ = ["nbody_decay", "GenParticle", "random"] + +import tensorflow as tf + +from . import random +from .phasespace import GenParticle, nbody_decay + + +def _set_eager_mode(): + import os + + is_eager = bool(os.environ.get("PHASESPACE_EAGER")) + tf.config.run_functions_eagerly(is_eager) + + +_set_eager_mode() diff --git a/phasespace/backend.py b/phasespace/backend.py index afb55167..25e9498b 100644 --- a/phasespace/backend.py +++ b/phasespace/backend.py @@ -1,21 +1,21 @@ -import tensorflow as tf - -RELAX_SHAPES = True -if int(tf.__version__.split(".")[1]) < 5: # smaller than 2.5 - jit_compile_argname = "experimental_compile" -else: - jit_compile_argname = "jit_compile" -function = tf.function( - autograph=False, - experimental_relax_shapes=RELAX_SHAPES, - **{jit_compile_argname: False} -) -function_jit = tf.function( - autograph=False, - experimental_relax_shapes=RELAX_SHAPES, - **{jit_compile_argname: True} -) - -function_jit_fixedshape = tf.function( - autograph=False, experimental_relax_shapes=False, **{jit_compile_argname: True} -) +import tensorflow as tf + +RELAX_SHAPES = True +if int(tf.__version__.split(".")[1]) < 5: # smaller than 2.5 + jit_compile_argname = "experimental_compile" +else: + jit_compile_argname = "jit_compile" +function = tf.function( + autograph=False, + experimental_relax_shapes=RELAX_SHAPES, + **{jit_compile_argname: False} +) +function_jit = tf.function( + autograph=False, + experimental_relax_shapes=RELAX_SHAPES, + **{jit_compile_argname: True} +) + +function_jit_fixedshape = tf.function( + autograph=False, experimental_relax_shapes=False, **{jit_compile_argname: True} +) diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py deleted file mode 100644 index 38788a14..00000000 --- a/phasespace/fromdecay/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -import sys - -from .fulldecay import FullDecay # noqa: F401 - -try: - import zfit # noqa: F401 - import zfit_physics as zphys # noqa: F401 - from particle import Particle # noqa: F401 -except ModuleNotFoundError as error: - raise ModuleNotFoundError( - "The fromdecay functionality in phasespace requires particle and zfit-physics. " - "Either install phasespace[fromdecay] or particle and zfit-physics.", - file=sys.stderr, - ) from error diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py new file mode 100644 index 00000000..42f877ee --- /dev/null +++ b/phasespace/fulldecay/__init__.py @@ -0,0 +1,14 @@ +import sys + +from .fulldecay import FullDecay + +try: + import zfit + import zfit_physics as zphys + from particle import Particle +except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "The fulldecay functionality in phasespace requires particle and zfit-physics. " + "Either install phasespace[fulldecay] or particle and zfit-physics.", + file=sys.stderr, + ) from error diff --git a/phasespace/fromdecay/fulldecay.py b/phasespace/fulldecay/fulldecay.py similarity index 90% rename from phasespace/fromdecay/fulldecay.py rename to phasespace/fulldecay/fulldecay.py index 2334d8be..76e755af 100644 --- a/phasespace/fromdecay/fulldecay.py +++ b/phasespace/fulldecay/fulldecay.py @@ -41,8 +41,7 @@ def from_dict( These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. - tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass - to be constant. + tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. Returns: The created FullDecay object. @@ -50,7 +49,7 @@ def from_dict( if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER else: - # Combine the default mass functions with the mass functions specified from the input. + # Combine the mass functions specified by the package to the mass functions specified from the input. total_mass_converter = {**_DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( @@ -68,19 +67,15 @@ def generate( Args: n_events: Total number of events combined, for all the decays. - normalize_weights: Normalize weights according to all events generated. - This also changes the return values. + normalize_weights: Normalize weights according to all events generated. This also changes the return values. See the phasespace documentation for more details. kwargs: Additional parameters passed to all calls of GenParticle.generate Returns: - The arguments returned by GenParticle.generate are returned. - See the phasespace documentation for details. - However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists - corresponding to the return arguments from the corresponding GenParticle instances in - self.gen_particles. - Note that when normalize_weights is True, the weights are normalized to the maximum of all - returned events. + The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. + However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding + to the return arguments from the corresponding GenParticle instances in self.gen_particles. + Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. """ # Input to tf.random.categorical must be 2D rand_i = tf.random.categorical( @@ -114,8 +109,8 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: preexisting_particles: Names that the particle cannot have as name. Returns: - name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i - begins at 0 and increases until the name is not preexisting_particles. + name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i begins at 0 + and increases until the name is not preexisting_particles. """ if name not in preexisting_particles: preexisting_particles.add(name) diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fulldecay/mass_functions.py similarity index 100% rename from phasespace/fromdecay/mass_functions.py rename to phasespace/fulldecay/mass_functions.py diff --git a/phasespace/kinematics.py b/phasespace/kinematics.py index 317651dd..460c127e 100644 --- a/phasespace/kinematics.py +++ b/phasespace/kinematics.py @@ -1,152 +1,152 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file kinematics.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 12.02.2019 -# ============================================================================= -"""Basic kinematics.""" - -import tensorflow.experimental.numpy as tnp - -from phasespace.backend import function, function_jit - - -@function_jit -def scalar_product(vec1, vec2): - """Calculate scalar product of two 3-vectors. - - Args: - vec1: First vector. - vec2: Second vector. - """ - return tnp.sum(vec1 * vec2, axis=1) - - -@function_jit -def spatial_component(vector): - """Extract spatial components of the input Lorentz vector. - - Args: - vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). - """ - return tnp.take(vector, indices=[0, 1, 2], axis=-1) - - -@function_jit -def time_component(vector): - """Extract time component of the input Lorentz vector. - - Args: - vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). - """ - return tnp.take(vector, indices=[3], axis=-1) - - -@function -def x_component(vector): - """Extract spatial X component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[0], axis=-1) - - -@function_jit -def y_component(vector): - """Extract spatial Y component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[1], axis=-1) - - -@function_jit -def z_component(vector): - """Extract spatial Z component of the input Lorentz or 3-vector. - - Args: - vector: Input vector. - """ - return tnp.take(vector, indices=[2], axis=-1) - - -@function_jit -def mass(vector): - """Calculate mass scalar for Lorentz 4-momentum. - - Args: - vector: Input Lorentz momentum vector. - """ - return tnp.sqrt( - tnp.sum(tnp.square(vector) * metric_tensor(), axis=-1, keepdims=True) - ) - - -@function_jit -def lorentz_vector(space, time): - """Make a Lorentz vector from spatial and time components. - - Args: - space: 3-vector of spatial components. - time: Time component. - """ - return tnp.concatenate([space, time], axis=-1) - - -@function_jit -def lorentz_boost(vector, boostvector): - """Perform Lorentz boost. - - Args: - vector: 4-vector to be boosted - boostvector: Boost vector. Can be either 3-vector or 4-vector, since - only spatial components are used. - """ - boost = spatial_component(boostvector) - b2 = tnp.expand_dims(scalar_product(boost, boost), axis=-1) - - def boost_fn(): - gamma = 1.0 / tnp.sqrt(1.0 - b2) - gamma2 = (gamma - 1.0) / b2 - ve = time_component(vector) - vp = spatial_component(vector) - bp = tnp.expand_dims(scalar_product(vp, boost), axis=-1) - vp2 = vp + (gamma2 * bp + gamma * ve) * boost - ve2 = gamma * (ve + bp) - return lorentz_vector(vp2, ve2) - - # if boost vector is zero, return the original vector - all_b2_zero = tnp.all(tnp.equal(b2, tnp.zeros_like(b2))) - boosted_vector = tnp.where(all_b2_zero, vector, boost_fn()) - return boosted_vector - - -@function_jit -def beta(vector): - """Calculate beta of a given 4-vector. - - Args: - vector: Input Lorentz momentum vector. - """ - return mass(vector) / time_component(vector) - - -@function_jit -def boost_components(vector): - """Get the boost components of a given 4-vector. - - Args: - vector: Input Lorentz momentum vector. - """ - return spatial_component(vector) / time_component(vector) - - -@function_jit -def metric_tensor(): - """Metric tensor for Lorentz space (constant).""" - return tnp.asarray([-1.0, -1.0, -1.0, 1.0], dtype=tnp.float64) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file kinematics.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 12.02.2019 +# ============================================================================= +"""Basic kinematics.""" + +import tensorflow.experimental.numpy as tnp + +from phasespace.backend import function, function_jit + + +@function_jit +def scalar_product(vec1, vec2): + """Calculate scalar product of two 3-vectors. + + Args: + vec1: First vector. + vec2: Second vector. + """ + return tnp.sum(vec1 * vec2, axis=1) + + +@function_jit +def spatial_component(vector): + """Extract spatial components of the input Lorentz vector. + + Args: + vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). + """ + return tnp.take(vector, indices=[0, 1, 2], axis=-1) + + +@function_jit +def time_component(vector): + """Extract time component of the input Lorentz vector. + + Args: + vector: Input Lorentz vector (where indexes 0-2 are space, index 3 is time). + """ + return tnp.take(vector, indices=[3], axis=-1) + + +@function +def x_component(vector): + """Extract spatial X component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[0], axis=-1) + + +@function_jit +def y_component(vector): + """Extract spatial Y component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[1], axis=-1) + + +@function_jit +def z_component(vector): + """Extract spatial Z component of the input Lorentz or 3-vector. + + Args: + vector: Input vector. + """ + return tnp.take(vector, indices=[2], axis=-1) + + +@function_jit +def mass(vector): + """Calculate mass scalar for Lorentz 4-momentum. + + Args: + vector: Input Lorentz momentum vector. + """ + return tnp.sqrt( + tnp.sum(tnp.square(vector) * metric_tensor(), axis=-1, keepdims=True) + ) + + +@function_jit +def lorentz_vector(space, time): + """Make a Lorentz vector from spatial and time components. + + Args: + space: 3-vector of spatial components. + time: Time component. + """ + return tnp.concatenate([space, time], axis=-1) + + +@function_jit +def lorentz_boost(vector, boostvector): + """Perform Lorentz boost. + + Args: + vector: 4-vector to be boosted + boostvector: Boost vector. Can be either 3-vector or 4-vector, since + only spatial components are used. + """ + boost = spatial_component(boostvector) + b2 = tnp.expand_dims(scalar_product(boost, boost), axis=-1) + + def boost_fn(): + gamma = 1.0 / tnp.sqrt(1.0 - b2) + gamma2 = (gamma - 1.0) / b2 + ve = time_component(vector) + vp = spatial_component(vector) + bp = tnp.expand_dims(scalar_product(vp, boost), axis=-1) + vp2 = vp + (gamma2 * bp + gamma * ve) * boost + ve2 = gamma * (ve + bp) + return lorentz_vector(vp2, ve2) + + # if boost vector is zero, return the original vector + all_b2_zero = tnp.all(tnp.equal(b2, tnp.zeros_like(b2))) + boosted_vector = tnp.where(all_b2_zero, vector, boost_fn()) + return boosted_vector + + +@function_jit +def beta(vector): + """Calculate beta of a given 4-vector. + + Args: + vector: Input Lorentz momentum vector. + """ + return mass(vector) / time_component(vector) + + +@function_jit +def boost_components(vector): + """Get the boost components of a given 4-vector. + + Args: + vector: Input Lorentz momentum vector. + """ + return spatial_component(vector) / time_component(vector) + + +@function_jit +def metric_tensor(): + """Metric tensor for Lorentz space (constant).""" + return tnp.asarray([-1.0, -1.0, -1.0, 1.0], dtype=tnp.float64) + + +# EOF diff --git a/phasespace/phasespace.py b/phasespace/phasespace.py index 70f09251..b5b821b9 100644 --- a/phasespace/phasespace.py +++ b/phasespace/phasespace.py @@ -1,776 +1,776 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file phasespace.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 25.02.2019 -# ============================================================================= -"""Implementation of the Raubold and Lynch method to generate n-body events. - -The code is based on the GENBOD function (W515 from CERNLIB), documented in: - - F. James, Monte Carlo Phase Space, CERN 68-15 (1968) -""" -import inspect -import warnings -from math import pi -from typing import Callable, Dict, Optional, Tuple, Union - -import tensorflow as tf -import tensorflow.experimental.numpy as tnp - -from . import kinematics as kin -from .backend import function, function_jit_fixedshape -from .random import SeedLike, get_rng - -RELAX_SHAPES = False - - -def process_list_to_tensor(lst): - """Convert a list to a tensor. - - The list is converted to a tensor and transposed to get the proper shape. - - Note: - If `lst` is a tensor, nothing is done to it other than convert it to `tf.float64`. - - Arguments: - lst (list): List to convert. - - Return: - ~`tf.Tensor` - """ - if isinstance(lst, list): - lst = tnp.transpose(tnp.asarray(lst, dtype=tnp.float64)) - return tnp.asarray(lst, dtype=tnp.float64) - - -@function -def pdk(a, b, c): - """Calculate the PDK (2-body phase space) function. - - Based on Eq. (9.17) in CERN 68-15 (1968). - - Arguments: - a (~`tf.Tensor`): :math:`M_{i+1}` in Eq. (9.17). - b (~`tf.Tensor`): :math:`M_{i}` in Eq. (9.17). - c (~`tf.Tensor`): :math:`m_{i+1}` in Eq. (9.17). - - Return: - ~`tf.Tensor` - """ - x = (a - b - c) * (a + b + c) * (a - b + c) * (a + b - c) - return tnp.sqrt(x) / (tnp.asarray(2.0, dtype=tnp.float64) * a) - - -class GenParticle: - """Representation of a particle. - - Instances of this class can be combined with each other to build decay chains, - which can then be used to generate phase space events through the `generate` - or `generate_tensor` method. - - A `GenParticle` must have - + a `name`, which is ensured not to clash with any others in - the decay chain. - + a `mass`, which can be either a number or a function to generate it according to - a certain distribution. The returned ~`tf.Tensor` needs to have shape (nevents,). - In this case, the particle is not considered as having a - fixed mass and the `has_fixed_mass` method will return False. - - It may also have: - - + Children, ie, decay products, which are also `GenParticle` instances. - - - Arguments: - name (str): Name of the particle. - mass (float, ~`tf.Tensor`, callable): Mass of the particle. If it's a float, it get - converted to a `tf.constant`. - """ - - def __init__(self, name: str, mass: Union[Callable, int, float]) -> None: # noqa - self.name = name - self.children = [] - self._mass_val = mass - if not callable(mass) and not isinstance(mass, tf.Variable): - mass = tnp.asarray(mass, dtype=tnp.float64) - else: - mass = tf.function( - mass, autograph=False, experimental_relax_shapes=RELAX_SHAPES - ) - self._mass = mass - self._generate_called = False # not yet called, children can be set - - def __repr__(self): - return "".format( - self.name, - f"{self._mass_val:.2f}" if self.has_fixed_mass else "variable", - ", ".join(child.name for child in self.children), - ) - - def _do_names_clash(self, particles): - def get_list_of_names(part): - output = [part.name] - for child in part.children: - output.extend(get_list_of_names(child)) - return output - - names_to_check = [self.name] - for part in particles: - names_to_check.extend(get_list_of_names(part)) - # Find top - dup_names = {name for name in names_to_check if names_to_check.count(name) > 1} - if dup_names: - return dup_names - return None - - @function - def get_mass( - self, - min_mass: tf.Tensor = None, - max_mass: tf.Tensor = None, - n_events: Union[tf.Tensor, tf.Variable] = None, - seed: SeedLike = None, - ) -> tf.Tensor: - """Get the particle mass. - - If the particle is resonant, the mass function will be called with the - `min_mass`, `max_mass`, `n_events` and optionally a `seed` parameter. - - Arguments: - min_mass (tensor): Lower mass range. Defaults to None, which - is only valid in the case of fixed mass. - max_mass (tensor): Upper mass range. Defaults to None, which - is only valid in the case of fixed mass. - n_events (): Number of events to produce. Has to be specified if the particle is resonant. - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - ~`tf.Tensor`: Mass of the particles, either a scalar or shape (nevents,) - - Raise: - ValueError: If the mass is requested and has not been set. - """ - if self.has_fixed_mass: - mass = self._mass - else: - seed = get_rng(seed) - min_mass = tnp.reshape(min_mass, (n_events,)) - max_mass = tnp.reshape(max_mass, (n_events,)) - signature = inspect.signature(self._mass) - if "seed" in signature.parameters: - mass = self._mass(min_mass, max_mass, n_events, seed=seed) - else: - mass = self._mass(min_mass, max_mass, n_events) - return mass - - @property - def has_fixed_mass(self): - """bool: Is the mass a callable function?""" - return not callable(self._mass) - - def set_children(self, *children): - """Assign children. - - Arguments: - children (GenParticle): Two or more children to assign to the current particle. - - Return: - self - - Raise: - ValueError: If there is an inconsistency in the parent/children relationship, ie, - if children were already set, if their parent was or if less than two children were given. - KeyError: If there is a particle name clash. - RuntimeError: If `generate` was already called before. - """ - # self._set_cache_validity(False) - if self._generate_called: - raise RuntimeError( - "Cannot set children after the first call to `generate`." - ) - if self.children: - raise ValueError("Children already set!") - if len(children) <= 1: - raise ValueError( - f"Have to set at least 2 children, not {len(children)} for a particle to decay" - ) - # Check name clashes - name_clash = self._do_names_clash(children) - if name_clash: - raise KeyError(f"Particle name {name_clash} already used") - self.children = children - return self - - @property - def has_children(self): - """bool: Does the particle have children?""" - return bool(self.children) - - @property - def has_grandchildren(self): - """bool: Does the particle have grandchildren?""" - if not self.children: - return False - return any(child.has_children for child in self.children) - - @staticmethod - def _preprocess(momentum, n_events): - """Preprocess momentum input and determine number of events to generate. - - Both `momentum` and `n_events` are converted to tensors if they - are not already. - - Arguments: - `momentum`: Momentum vector, of shape (4, x), where x is optional. - `n_events`: Number of events to generate. If `None`, the number of events - to generate is calculated from the shape of `momentum`. - - Return: - tuple: Processed `momentum` and `n_events`. - - Raise: - tf.errors.InvalidArgumentError: If the number of events deduced from the - shape of `momentum` is inconsistent with `n_events`. - """ - momentum = process_list_to_tensor(momentum) - - # Check sanity of inputs - if len(momentum.shape) not in (1, 2): - raise ValueError(f"Bad shape for momentum -> {list(momentum.shape)}") - # Check compatibility of inputs - if len(momentum.shape) == 2: - if n_events is not None: - momentum_shape = momentum.shape[0] - if momentum_shape is None: - momentum_shape = tf.shape(momentum)[0] - momentum_shape = tnp.asarray(momentum_shape, tnp.int64) - else: - momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) - tf.assert_equal( - n_events, - momentum_shape, - message="Conflicting inputs -> momentum_shape and n_events", - ) - - if n_events is None: - if len(momentum.shape) == 2: - n_events = momentum.shape[0] - if n_events is None: # dynamic shape - n_events = tf.shape(momentum)[0] - n_events = tnp.asarray(n_events, dtype=tnp.int64) - else: - n_events = tnp.asarray(1, dtype=tnp.int64) - n_events = tnp.asarray(n_events, dtype=tnp.int64) - # Now preparation of tensors - if len(momentum.shape) == 1: - momentum = tnp.expand_dims(momentum, axis=0) - return momentum, n_events - - @staticmethod - @function_jit_fixedshape - def _get_w_max(available_mass, masses): - emmax = available_mass + tnp.take(masses, indices=[0], axis=1) - emmin = tnp.zeros_like(emmax, dtype=tnp.float64) - w_max = tnp.ones_like(emmax, dtype=tnp.float64) - for i in range(1, masses.shape[1]): - emmin += tnp.take(masses, [i - 1], axis=1) - emmax += tnp.take(masses, [i], axis=1) - w_max *= pdk(emmax, emmin, tnp.take(masses, [i], axis=1)) - return w_max - - def _generate(self, momentum, n_events, rng): - """Generate a n-body decay according to the Raubold and Lynch method. - - The number and mass of the children particles are taken from self.children. - - Note: - This method generates the same results as the GENBOD routine. - - Arguments: - momentum (tensor): Momentum of the parent particle. All generated particles - will be boosted to that momentum. - n_events (int): Number of events to generate. - - Return: - tuple: Result of the generation (per-event weights, maximum weights, output particles - and their output masses). - """ - self._generate_called = True - if not self.children: - raise ValueError("No children have been configured") - p_top, n_events = self._preprocess(momentum, n_events) - top_mass = tnp.broadcast_to(kin.mass(p_top), (n_events, 1)) - n_particles = len(self.children) - - # Prepare masses - def recurse_stable(part): - output_mass = 0 - for child in part.children: - if child.has_fixed_mass: - output_mass += child.get_mass() - else: - output_mass += recurse_stable(child) - return output_mass - - mass_from_stable = tnp.broadcast_to( - tnp.sum( - [child.get_mass() for child in self.children if child.has_fixed_mass], - axis=0, - ), - (n_events, 1), - ) - max_mass = top_mass - mass_from_stable - masses = [] - for child in self.children: - if child.has_fixed_mass: - masses.append(tnp.broadcast_to(child.get_mass(), (n_events, 1))) - else: - # Recurse that particle to know the minimum mass we need to generate - min_mass = tnp.broadcast_to(recurse_stable(child), (n_events, 1)) - mass = child.get_mass(min_mass, max_mass, n_events) - mass = tnp.reshape(mass, (n_events, 1)) - max_mass -= mass - masses.append(mass) - masses = tnp.concatenate(masses, axis=-1) - # if len(masses.shape) == 1: - # masses = tnp.expand_dims(masses, axis=0) - available_mass = top_mass - tnp.sum(masses, axis=1, keepdims=True) - tf.debugging.assert_greater_equal( - available_mass, - tnp.zeros_like(available_mass, dtype=tnp.float64), - message="Forbidden decay", - name="mass_check", - ) # Calculate the max weight, initial beta, etc - w_max = self._get_w_max(available_mass, masses) - p_top_boost = kin.boost_components(p_top) - # Start the generation - random_numbers = rng.uniform((n_events, n_particles - 2), dtype=tnp.float64) - random = tnp.concatenate( - [ - tnp.zeros((n_events, 1), dtype=tnp.float64), - tnp.sort(random_numbers, axis=1), - tnp.ones((n_events, 1), dtype=tnp.float64), - ], - axis=1, - ) - if random.shape[1] is None: - random.set_shape((None, n_particles)) - # random = tnp.expand_dims(random, axis=-1) - sum_ = tnp.zeros((n_events, 1), dtype=tnp.float64) - inv_masses = [] - # TODO(Mayou36): rewrite with cumsum? - for i in range(n_particles): - sum_ += tnp.take(masses, [i], axis=1) - inv_masses.append(tnp.take(random, [i], axis=1) * available_mass + sum_) - generated_particles, weights = self._generate_part2( - inv_masses, masses, n_events, n_particles, rng=rng - ) - # Final boost of all particles - generated_particles = [ - kin.lorentz_boost(part, p_top_boost) for part in generated_particles - ] - return ( - tnp.reshape(weights, (n_events,)), - tnp.reshape(w_max, (n_events,)), - generated_particles, - masses, - ) - - @staticmethod - @function - def _generate_part2(inv_masses, masses, n_events, n_particles, rng): - pds = [] - # Calculate weights of the events - for i in range(n_particles - 1): - pds.append( - pdk( - inv_masses[i + 1], - inv_masses[i], - tnp.take(masses, [i + 1], axis=1), - ) - ) - weights = tnp.prod(pds, axis=0) - zero_component = tnp.zeros_like(pds[0], dtype=tnp.float64) - generated_particles = [ - tnp.concatenate( - [ - zero_component, - pds[0], - zero_component, - tnp.sqrt( - tnp.square(pds[0]) + tnp.square(tnp.take(masses, [0], axis=1)) - ), - ], - axis=1, - ) - ] - part_num = 1 - while True: - generated_particles.append( - tnp.concatenate( - [ - zero_component, - -pds[part_num - 1], - zero_component, - tnp.sqrt( - tnp.square(pds[part_num - 1]) - + tnp.square(tnp.take(masses, [part_num], axis=1)) - ), - ], - axis=1, - ) - ) - - cos_z = tnp.asarray(2.0, dtype=tnp.float64) * rng.uniform( - (n_events, 1), dtype=tnp.float64 - ) - tnp.asarray(1.0, dtype=tnp.float64) - sin_z = tnp.sqrt(tnp.asarray(1.0, dtype=tnp.float64) - cos_z * cos_z) - ang_y = ( - tnp.asarray(2.0, dtype=tnp.float64) - * tnp.asarray(pi, dtype=tnp.float64) - * rng.uniform((n_events, 1), dtype=tnp.float64) - ) - cos_y = tnp.cos(ang_y) - sin_y = tnp.sin(ang_y) - # Do the rotations - for j in range(part_num + 1): - px = kin.x_component(generated_particles[j]) - py = kin.y_component(generated_particles[j]) - # Rotate about z - # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it? - generated_particles[j] = tnp.concatenate( - [ - cos_z * px - sin_z * py, - sin_z * px + cos_z * py, - kin.z_component(generated_particles[j]), - kin.time_component(generated_particles[j]), - ], - axis=1, - ) - # Rotate about y - px = kin.x_component(generated_particles[j]) - pz = kin.z_component(generated_particles[j]) - generated_particles[j] = tnp.concatenate( - [ - cos_y * px - sin_y * pz, - kin.y_component(generated_particles[j]), - sin_y * px + cos_y * pz, - kin.time_component(generated_particles[j]), - ], - axis=1, - ) - if part_num == (n_particles - 1): - break - betas = pds[part_num] / tnp.sqrt( - tnp.square(pds[part_num]) + tnp.square(inv_masses[part_num]) - ) - generated_particles = [ - kin.lorentz_boost( - part, - tnp.concatenate([zero_component, betas, zero_component], axis=1), - ) - for part in generated_particles - ] - part_num += 1 - return generated_particles, weights - - @function - def _recursive_generate( - self, - n_events, - boost_to=None, - recalculate_max_weights=False, - rng: SeedLike = None, - ): - """Recursively generate normalized n-body phase space as tensorflow tensors. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (tensor, optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - recalculate_max_weights (bool, optional): Recalculate the maximum weight of the event - using all the particles of the tree? This is necessary for the top particle of a decay, - otherwise the maximum weight calculation is going to be wrong (particles from subdecays - would not be taken into account). Defaults to False. - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - tuple: Result of the generation (per-event weights, maximum weights, output particles - and their output masses). - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - if boost_to is not None: - momentum = boost_to - else: - if self.has_fixed_mass: - momentum = tnp.broadcast_to( - tnp.stack((0.0, 0.0, 0.0, self.get_mass()), axis=-1), (n_events, 4) - ) - else: - raise ValueError("Cannot use resonance as top particle") - weights, weights_max, parts, children_masses = self._generate( - momentum, n_events, rng=rng - ) - output_particles = { - child.name: parts[child_num] - for child_num, child in enumerate(self.children) - } - output_masses = { - child.name: tnp.take(children_masses, [child_num], axis=1) - for child_num, child in enumerate(self.children) - } - for child_num, child in enumerate(self.children): - if child.has_children: - ( - child_weights, - _, - child_gen_particles, - child_masses, - ) = child._recursive_generate( - n_events=n_events, - boost_to=parts[child_num], - recalculate_max_weights=False, - rng=rng, - ) - weights *= child_weights - output_particles.update(child_gen_particles) - output_masses.update(child_masses) - if recalculate_max_weights: - - def build_mass_tree(particle, leaf): - if particle.has_children: - leaf[particle.name] = {} - for child in particle.children: - build_mass_tree(child, leaf[particle.name]) - else: - leaf[particle.name] = output_masses[particle.name] - - def get_flattened_values(dict_): - output = [] - for val in dict_.values(): - if isinstance(val, dict): - output.extend(get_flattened_values(val)) - else: - output.append(val) - return output - - def recurse_w_max(parent_mass, current_mass_tree): - available_mass = parent_mass - sum( - get_flattened_values(current_mass_tree) - ) - masses = [] - w_max = tnp.ones_like(available_mass) - for child, child_mass in current_mass_tree.items(): - if isinstance(child_mass, dict): - w_max *= recurse_w_max( - parent_mass - - sum( - get_flattened_values( - { - ch_it: ch_m_it - for ch_it, ch_m_it in current_mass_tree.items() - if ch_it != child - } - ) - ), - child_mass, - ) - masses.append(sum(get_flattened_values(child_mass))) - else: - masses.append(child_mass) - masses = tnp.concatenate(masses, axis=1) - w_max *= self._get_w_max(available_mass, masses) - return w_max - - mass_tree = {} - build_mass_tree(self, mass_tree) - momentum = process_list_to_tensor(momentum) - if len(momentum.shape) == 1: - momentum = tnp.expand_dims(momentum, axis=-1) - weights_max = tnp.reshape( - recurse_w_max(kin.mass(momentum), mass_tree[self.name]), (n_events,) - ) - return weights, weights_max, output_particles, output_masses - - def generate( - self, - n_events: Union[int, tf.Tensor, tf.Variable], - boost_to: Optional[tf.Tensor] = None, - normalize_weights: bool = True, - seed: SeedLike = None, - ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]: - """Generate normalized n-body phase space as tensorflow tensors. - - Any TensorFlow tensor can always be converted to a numpy array with the method `numpy()`. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - normalize_weights (bool, optional): Normalize the event weight to its max? - seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used - as a seed to create a random number generator inside the function or directly as - the random number generator instance, respectively. - - Return: - tuple: Result of the generation, which varies with the value of `normalize_weights`: - - + If True, the tuple elements are the normalized event weights as a tensor of shape - (n_events, ), and the momenta generated particles as a dictionary of tensors of shape - (4, n_events) with particle names as keys. - - + If False, the tuple weights are the unnormalized event weights as a tensor of shape - (n_events, ), the maximum per-event weights as a tensor of shape (n_events, ) and the - momenta generated particles as a dictionary of tensors of shape (4, n_events) with particle - names as keys. - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - rng = get_rng(seed) - if boost_to is not None: - message = ( - f"The number of events requested ({n_events}) doesn't match the boost_to input size " - f"of {boost_to.shape}" - ) - tf.assert_equal(tf.shape(boost_to)[0], tf.shape(n_events), message=message) - if not isinstance(n_events, tf.Variable): - n_events = tnp.asarray(n_events, dtype=tnp.int64) - weights, weights_max, parts, _ = self._recursive_generate( - n_events=n_events, - boost_to=boost_to, - recalculate_max_weights=self.has_grandchildren, - rng=rng, - ) - return ( - (weights / weights_max, parts) - if normalize_weights - else (weights, weights_max, parts) - ) - - def generate_tensor( - self, - n_events: int, - boost_to=None, - normalize_weights: bool = True, - ): - """Generate normalized n-body phase space as numpy arrays. - - Events are generated in the rest frame of the particle, unless `boost_to` is given. - - Note: - In this method, the event weights are returned normalized to their maximum. - - Arguments: - n_events (int): Number of events to generate. - boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where - the resulting events will be boosted. If not specified, events are generated - in the rest frame of the particle. - normalize_weights (bool, optional): Normalize the event weight to its max - - - Return: - tuple: Result of the generation, which varies with the value of `normalize_weights`: - - + If True, the tuple elements are the normalized event weights as an array of shape - (n_events, ), and the momenta generated particles as a dictionary of arrays of shape - (4, n_events) with particle names as keys. - - + If False, the tuple weights are the unnormalized event weights as an array of shape - (n_events, ), the maximum per-event weights as an array of shape (n_events, ) and the - momenta generated particles as a dictionary of arrays of shape (4, n_events) with particle - names as keys. - - Raise: - tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. - ValueError: If `n_events` and the size of `boost_to` don't match. - See `GenParticle.generate_unnormalized`. - """ - - # Run generation - warnings.warn( - "This function is depreceated. Use `generate` which does not return a Tensor as well." - ) - generate_tf = self.generate(n_events, boost_to, normalize_weights) - # self._cache = generate_tf - # self._set_cache_validity(True, propagate=True) - return generate_tf - - -# legacy class to warn user about name change -class Particle: - """Deprecated Particle class. - - Renamed to GenParticle. - """ - - def __init__(self): - raise NameError( - "'Particle' has been renamed to 'GenParticle'. Please update your code accordingly." - "For more information, see: https://github.com/zfit/phasespace/issues/22" - ) - - -def nbody_decay(mass_top: float, masses: list, top_name: str = "", names: list = None): - """Shortcut to build an n-body decay of a GenParticle. - - If the particle names are not given, the top particle is called 'top' and the - children 'p_{i}', where i corresponds to their position in the `masses` sequence. - - Arguments: - mass_top (tensor, list): Mass of the top particle. Can be a list of 4-vectors. - masses (list): Masses of the child particles. - name_top (str, optional): Name of the top particle. If not given, the top particle is - named top. - names (list, optional): Names of the child particles. If not given, they are build as - 'p_{i}', where i is given by their ordering in the `masses` list. - - Return: - `GenParticle`: Particle decay. - - Raise: - ValueError: If the length of `masses` and `names` doesn't match. - """ - if not top_name: - top_name = "top" - if not names: - names = [f"p_{num}" for num in range(len(masses))] - if len(names) != len(masses): - raise ValueError("Mismatch in length between children masses and their names.") - return GenParticle(top_name, mass_top).set_children( - *(GenParticle(names[num], mass=mass) for num, mass in enumerate(masses)) - ) - - -def generate_decay(*args, **kwargs): - """Deprecated.""" - raise NameError( - "'generate_decay' has been removed. A similar behavior can be accomplished with 'nbody_decay'. " - "For more information see https://github.com/zfit/phasespace/issues/22" - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file phasespace.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 25.02.2019 +# ============================================================================= +"""Implementation of the Raubold and Lynch method to generate n-body events. + +The code is based on the GENBOD function (W515 from CERNLIB), documented in: + + F. James, Monte Carlo Phase Space, CERN 68-15 (1968) +""" +import inspect +import warnings +from math import pi +from typing import Callable, Dict, Optional, Tuple, Union + +import tensorflow as tf +import tensorflow.experimental.numpy as tnp + +from . import kinematics as kin +from .backend import function, function_jit_fixedshape +from .random import SeedLike, get_rng + +RELAX_SHAPES = False + + +def process_list_to_tensor(lst): + """Convert a list to a tensor. + + The list is converted to a tensor and transposed to get the proper shape. + + Note: + If `lst` is a tensor, nothing is done to it other than convert it to `tf.float64`. + + Arguments: + lst (list): List to convert. + + Return: + ~`tf.Tensor` + """ + if isinstance(lst, list): + lst = tnp.transpose(tnp.asarray(lst, dtype=tnp.float64)) + return tnp.asarray(lst, dtype=tnp.float64) + + +@function +def pdk(a, b, c): + """Calculate the PDK (2-body phase space) function. + + Based on Eq. (9.17) in CERN 68-15 (1968). + + Arguments: + a (~`tf.Tensor`): :math:`M_{i+1}` in Eq. (9.17). + b (~`tf.Tensor`): :math:`M_{i}` in Eq. (9.17). + c (~`tf.Tensor`): :math:`m_{i+1}` in Eq. (9.17). + + Return: + ~`tf.Tensor` + """ + x = (a - b - c) * (a + b + c) * (a - b + c) * (a + b - c) + return tnp.sqrt(x) / (tnp.asarray(2.0, dtype=tnp.float64) * a) + + +class GenParticle: + """Representation of a particle. + + Instances of this class can be combined with each other to build decay chains, + which can then be used to generate phase space events through the `generate` + or `generate_tensor` method. + + A `GenParticle` must have + + a `name`, which is ensured not to clash with any others in + the decay chain. + + a `mass`, which can be either a number or a function to generate it according to + a certain distribution. The returned ~`tf.Tensor` needs to have shape (nevents,). + In this case, the particle is not considered as having a + fixed mass and the `has_fixed_mass` method will return False. + + It may also have: + + + Children, ie, decay products, which are also `GenParticle` instances. + + + Arguments: + name (str): Name of the particle. + mass (float, ~`tf.Tensor`, callable): Mass of the particle. If it's a float, it get + converted to a `tf.constant`. + """ + + def __init__(self, name: str, mass: Union[Callable, int, float]) -> None: # noqa + self.name = name + self.children = [] + self._mass_val = mass + if not callable(mass) and not isinstance(mass, tf.Variable): + mass = tnp.asarray(mass, dtype=tnp.float64) + else: + mass = tf.function( + mass, autograph=False, experimental_relax_shapes=RELAX_SHAPES + ) + self._mass = mass + self._generate_called = False # not yet called, children can be set + + def __repr__(self): + return "".format( + self.name, + f"{self._mass_val:.2f}" if self.has_fixed_mass else "variable", + ", ".join(child.name for child in self.children), + ) + + def _do_names_clash(self, particles): + def get_list_of_names(part): + output = [part.name] + for child in part.children: + output.extend(get_list_of_names(child)) + return output + + names_to_check = [self.name] + for part in particles: + names_to_check.extend(get_list_of_names(part)) + # Find top + dup_names = {name for name in names_to_check if names_to_check.count(name) > 1} + if dup_names: + return dup_names + return None + + @function + def get_mass( + self, + min_mass: tf.Tensor = None, + max_mass: tf.Tensor = None, + n_events: Union[tf.Tensor, tf.Variable] = None, + seed: SeedLike = None, + ) -> tf.Tensor: + """Get the particle mass. + + If the particle is resonant, the mass function will be called with the + `min_mass`, `max_mass`, `n_events` and optionally a `seed` parameter. + + Arguments: + min_mass (tensor): Lower mass range. Defaults to None, which + is only valid in the case of fixed mass. + max_mass (tensor): Upper mass range. Defaults to None, which + is only valid in the case of fixed mass. + n_events (): Number of events to produce. Has to be specified if the particle is resonant. + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + ~`tf.Tensor`: Mass of the particles, either a scalar or shape (nevents,) + + Raise: + ValueError: If the mass is requested and has not been set. + """ + if self.has_fixed_mass: + mass = self._mass + else: + seed = get_rng(seed) + min_mass = tnp.reshape(min_mass, (n_events,)) + max_mass = tnp.reshape(max_mass, (n_events,)) + signature = inspect.signature(self._mass) + if "seed" in signature.parameters: + mass = self._mass(min_mass, max_mass, n_events, seed=seed) + else: + mass = self._mass(min_mass, max_mass, n_events) + return mass + + @property + def has_fixed_mass(self): + """bool: Is the mass a callable function?""" + return not callable(self._mass) + + def set_children(self, *children): + """Assign children. + + Arguments: + children (GenParticle): Two or more children to assign to the current particle. + + Return: + self + + Raise: + ValueError: If there is an inconsistency in the parent/children relationship, ie, + if children were already set, if their parent was or if less than two children were given. + KeyError: If there is a particle name clash. + RuntimeError: If `generate` was already called before. + """ + # self._set_cache_validity(False) + if self._generate_called: + raise RuntimeError( + "Cannot set children after the first call to `generate`." + ) + if self.children: + raise ValueError("Children already set!") + if len(children) <= 1: + raise ValueError( + f"Have to set at least 2 children, not {len(children)} for a particle to decay" + ) + # Check name clashes + name_clash = self._do_names_clash(children) + if name_clash: + raise KeyError(f"Particle name {name_clash} already used") + self.children = children + return self + + @property + def has_children(self): + """bool: Does the particle have children?""" + return bool(self.children) + + @property + def has_grandchildren(self): + """bool: Does the particle have grandchildren?""" + if not self.children: + return False + return any(child.has_children for child in self.children) + + @staticmethod + def _preprocess(momentum, n_events): + """Preprocess momentum input and determine number of events to generate. + + Both `momentum` and `n_events` are converted to tensors if they + are not already. + + Arguments: + `momentum`: Momentum vector, of shape (4, x), where x is optional. + `n_events`: Number of events to generate. If `None`, the number of events + to generate is calculated from the shape of `momentum`. + + Return: + tuple: Processed `momentum` and `n_events`. + + Raise: + tf.errors.InvalidArgumentError: If the number of events deduced from the + shape of `momentum` is inconsistent with `n_events`. + """ + momentum = process_list_to_tensor(momentum) + + # Check sanity of inputs + if len(momentum.shape) not in (1, 2): + raise ValueError(f"Bad shape for momentum -> {list(momentum.shape)}") + # Check compatibility of inputs + if len(momentum.shape) == 2: + if n_events is not None: + momentum_shape = momentum.shape[0] + if momentum_shape is None: + momentum_shape = tf.shape(momentum)[0] + momentum_shape = tnp.asarray(momentum_shape, tnp.int64) + else: + momentum_shape = tnp.asarray(momentum_shape, dtype=tnp.int64) + tf.assert_equal( + n_events, + momentum_shape, + message="Conflicting inputs -> momentum_shape and n_events", + ) + + if n_events is None: + if len(momentum.shape) == 2: + n_events = momentum.shape[0] + if n_events is None: # dynamic shape + n_events = tf.shape(momentum)[0] + n_events = tnp.asarray(n_events, dtype=tnp.int64) + else: + n_events = tnp.asarray(1, dtype=tnp.int64) + n_events = tnp.asarray(n_events, dtype=tnp.int64) + # Now preparation of tensors + if len(momentum.shape) == 1: + momentum = tnp.expand_dims(momentum, axis=0) + return momentum, n_events + + @staticmethod + @function_jit_fixedshape + def _get_w_max(available_mass, masses): + emmax = available_mass + tnp.take(masses, indices=[0], axis=1) + emmin = tnp.zeros_like(emmax, dtype=tnp.float64) + w_max = tnp.ones_like(emmax, dtype=tnp.float64) + for i in range(1, masses.shape[1]): + emmin += tnp.take(masses, [i - 1], axis=1) + emmax += tnp.take(masses, [i], axis=1) + w_max *= pdk(emmax, emmin, tnp.take(masses, [i], axis=1)) + return w_max + + def _generate(self, momentum, n_events, rng): + """Generate a n-body decay according to the Raubold and Lynch method. + + The number and mass of the children particles are taken from self.children. + + Note: + This method generates the same results as the GENBOD routine. + + Arguments: + momentum (tensor): Momentum of the parent particle. All generated particles + will be boosted to that momentum. + n_events (int): Number of events to generate. + + Return: + tuple: Result of the generation (per-event weights, maximum weights, output particles + and their output masses). + """ + self._generate_called = True + if not self.children: + raise ValueError("No children have been configured") + p_top, n_events = self._preprocess(momentum, n_events) + top_mass = tnp.broadcast_to(kin.mass(p_top), (n_events, 1)) + n_particles = len(self.children) + + # Prepare masses + def recurse_stable(part): + output_mass = 0 + for child in part.children: + if child.has_fixed_mass: + output_mass += child.get_mass() + else: + output_mass += recurse_stable(child) + return output_mass + + mass_from_stable = tnp.broadcast_to( + tnp.sum( + [child.get_mass() for child in self.children if child.has_fixed_mass], + axis=0, + ), + (n_events, 1), + ) + max_mass = top_mass - mass_from_stable + masses = [] + for child in self.children: + if child.has_fixed_mass: + masses.append(tnp.broadcast_to(child.get_mass(), (n_events, 1))) + else: + # Recurse that particle to know the minimum mass we need to generate + min_mass = tnp.broadcast_to(recurse_stable(child), (n_events, 1)) + mass = child.get_mass(min_mass, max_mass, n_events) + mass = tnp.reshape(mass, (n_events, 1)) + max_mass -= mass + masses.append(mass) + masses = tnp.concatenate(masses, axis=-1) + # if len(masses.shape) == 1: + # masses = tnp.expand_dims(masses, axis=0) + available_mass = top_mass - tnp.sum(masses, axis=1, keepdims=True) + tf.debugging.assert_greater_equal( + available_mass, + tnp.zeros_like(available_mass, dtype=tnp.float64), + message="Forbidden decay", + name="mass_check", + ) # Calculate the max weight, initial beta, etc + w_max = self._get_w_max(available_mass, masses) + p_top_boost = kin.boost_components(p_top) + # Start the generation + random_numbers = rng.uniform((n_events, n_particles - 2), dtype=tnp.float64) + random = tnp.concatenate( + [ + tnp.zeros((n_events, 1), dtype=tnp.float64), + tnp.sort(random_numbers, axis=1), + tnp.ones((n_events, 1), dtype=tnp.float64), + ], + axis=1, + ) + if random.shape[1] is None: + random.set_shape((None, n_particles)) + # random = tnp.expand_dims(random, axis=-1) + sum_ = tnp.zeros((n_events, 1), dtype=tnp.float64) + inv_masses = [] + # TODO(Mayou36): rewrite with cumsum? + for i in range(n_particles): + sum_ += tnp.take(masses, [i], axis=1) + inv_masses.append(tnp.take(random, [i], axis=1) * available_mass + sum_) + generated_particles, weights = self._generate_part2( + inv_masses, masses, n_events, n_particles, rng=rng + ) + # Final boost of all particles + generated_particles = [ + kin.lorentz_boost(part, p_top_boost) for part in generated_particles + ] + return ( + tnp.reshape(weights, (n_events,)), + tnp.reshape(w_max, (n_events,)), + generated_particles, + masses, + ) + + @staticmethod + @function + def _generate_part2(inv_masses, masses, n_events, n_particles, rng): + pds = [] + # Calculate weights of the events + for i in range(n_particles - 1): + pds.append( + pdk( + inv_masses[i + 1], + inv_masses[i], + tnp.take(masses, [i + 1], axis=1), + ) + ) + weights = tnp.prod(pds, axis=0) + zero_component = tnp.zeros_like(pds[0], dtype=tnp.float64) + generated_particles = [ + tnp.concatenate( + [ + zero_component, + pds[0], + zero_component, + tnp.sqrt( + tnp.square(pds[0]) + tnp.square(tnp.take(masses, [0], axis=1)) + ), + ], + axis=1, + ) + ] + part_num = 1 + while True: + generated_particles.append( + tnp.concatenate( + [ + zero_component, + -pds[part_num - 1], + zero_component, + tnp.sqrt( + tnp.square(pds[part_num - 1]) + + tnp.square(tnp.take(masses, [part_num], axis=1)) + ), + ], + axis=1, + ) + ) + + cos_z = tnp.asarray(2.0, dtype=tnp.float64) * rng.uniform( + (n_events, 1), dtype=tnp.float64 + ) - tnp.asarray(1.0, dtype=tnp.float64) + sin_z = tnp.sqrt(tnp.asarray(1.0, dtype=tnp.float64) - cos_z * cos_z) + ang_y = ( + tnp.asarray(2.0, dtype=tnp.float64) + * tnp.asarray(pi, dtype=tnp.float64) + * rng.uniform((n_events, 1), dtype=tnp.float64) + ) + cos_y = tnp.cos(ang_y) + sin_y = tnp.sin(ang_y) + # Do the rotations + for j in range(part_num + 1): + px = kin.x_component(generated_particles[j]) + py = kin.y_component(generated_particles[j]) + # Rotate about z + # TODO(Mayou36): only list? will be overwritten below anyway, but can `*_component` handle it? + generated_particles[j] = tnp.concatenate( + [ + cos_z * px - sin_z * py, + sin_z * px + cos_z * py, + kin.z_component(generated_particles[j]), + kin.time_component(generated_particles[j]), + ], + axis=1, + ) + # Rotate about y + px = kin.x_component(generated_particles[j]) + pz = kin.z_component(generated_particles[j]) + generated_particles[j] = tnp.concatenate( + [ + cos_y * px - sin_y * pz, + kin.y_component(generated_particles[j]), + sin_y * px + cos_y * pz, + kin.time_component(generated_particles[j]), + ], + axis=1, + ) + if part_num == (n_particles - 1): + break + betas = pds[part_num] / tnp.sqrt( + tnp.square(pds[part_num]) + tnp.square(inv_masses[part_num]) + ) + generated_particles = [ + kin.lorentz_boost( + part, + tnp.concatenate([zero_component, betas, zero_component], axis=1), + ) + for part in generated_particles + ] + part_num += 1 + return generated_particles, weights + + @function + def _recursive_generate( + self, + n_events, + boost_to=None, + recalculate_max_weights=False, + rng: SeedLike = None, + ): + """Recursively generate normalized n-body phase space as tensorflow tensors. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (tensor, optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + recalculate_max_weights (bool, optional): Recalculate the maximum weight of the event + using all the particles of the tree? This is necessary for the top particle of a decay, + otherwise the maximum weight calculation is going to be wrong (particles from subdecays + would not be taken into account). Defaults to False. + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + tuple: Result of the generation (per-event weights, maximum weights, output particles + and their output masses). + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + if boost_to is not None: + momentum = boost_to + else: + if self.has_fixed_mass: + momentum = tnp.broadcast_to( + tnp.stack((0.0, 0.0, 0.0, self.get_mass()), axis=-1), (n_events, 4) + ) + else: + raise ValueError("Cannot use resonance as top particle") + weights, weights_max, parts, children_masses = self._generate( + momentum, n_events, rng=rng + ) + output_particles = { + child.name: parts[child_num] + for child_num, child in enumerate(self.children) + } + output_masses = { + child.name: tnp.take(children_masses, [child_num], axis=1) + for child_num, child in enumerate(self.children) + } + for child_num, child in enumerate(self.children): + if child.has_children: + ( + child_weights, + _, + child_gen_particles, + child_masses, + ) = child._recursive_generate( + n_events=n_events, + boost_to=parts[child_num], + recalculate_max_weights=False, + rng=rng, + ) + weights *= child_weights + output_particles.update(child_gen_particles) + output_masses.update(child_masses) + if recalculate_max_weights: + + def build_mass_tree(particle, leaf): + if particle.has_children: + leaf[particle.name] = {} + for child in particle.children: + build_mass_tree(child, leaf[particle.name]) + else: + leaf[particle.name] = output_masses[particle.name] + + def get_flattened_values(dict_): + output = [] + for val in dict_.values(): + if isinstance(val, dict): + output.extend(get_flattened_values(val)) + else: + output.append(val) + return output + + def recurse_w_max(parent_mass, current_mass_tree): + available_mass = parent_mass - sum( + get_flattened_values(current_mass_tree) + ) + masses = [] + w_max = tnp.ones_like(available_mass) + for child, child_mass in current_mass_tree.items(): + if isinstance(child_mass, dict): + w_max *= recurse_w_max( + parent_mass + - sum( + get_flattened_values( + { + ch_it: ch_m_it + for ch_it, ch_m_it in current_mass_tree.items() + if ch_it != child + } + ) + ), + child_mass, + ) + masses.append(sum(get_flattened_values(child_mass))) + else: + masses.append(child_mass) + masses = tnp.concatenate(masses, axis=1) + w_max *= self._get_w_max(available_mass, masses) + return w_max + + mass_tree = {} + build_mass_tree(self, mass_tree) + momentum = process_list_to_tensor(momentum) + if len(momentum.shape) == 1: + momentum = tnp.expand_dims(momentum, axis=-1) + weights_max = tnp.reshape( + recurse_w_max(kin.mass(momentum), mass_tree[self.name]), (n_events,) + ) + return weights, weights_max, output_particles, output_masses + + def generate( + self, + n_events: Union[int, tf.Tensor, tf.Variable], + boost_to: Optional[tf.Tensor] = None, + normalize_weights: bool = True, + seed: SeedLike = None, + ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]: + """Generate normalized n-body phase space as tensorflow tensors. + + Any TensorFlow tensor can always be converted to a numpy array with the method `numpy()`. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + normalize_weights (bool, optional): Normalize the event weight to its max? + seed (`SeedLike`): The seed can be a number or a `tf.random.Generator` that are used + as a seed to create a random number generator inside the function or directly as + the random number generator instance, respectively. + + Return: + tuple: Result of the generation, which varies with the value of `normalize_weights`: + + + If True, the tuple elements are the normalized event weights as a tensor of shape + (n_events, ), and the momenta generated particles as a dictionary of tensors of shape + (4, n_events) with particle names as keys. + + + If False, the tuple weights are the unnormalized event weights as a tensor of shape + (n_events, ), the maximum per-event weights as a tensor of shape (n_events, ) and the + momenta generated particles as a dictionary of tensors of shape (4, n_events) with particle + names as keys. + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + rng = get_rng(seed) + if boost_to is not None: + message = ( + f"The number of events requested ({n_events}) doesn't match the boost_to input size " + f"of {boost_to.shape}" + ) + tf.assert_equal(tf.shape(boost_to)[0], tf.shape(n_events), message=message) + if not isinstance(n_events, tf.Variable): + n_events = tnp.asarray(n_events, dtype=tnp.int64) + weights, weights_max, parts, _ = self._recursive_generate( + n_events=n_events, + boost_to=boost_to, + recalculate_max_weights=self.has_grandchildren, + rng=rng, + ) + return ( + (weights / weights_max, parts) + if normalize_weights + else (weights, weights_max, parts) + ) + + def generate_tensor( + self, + n_events: int, + boost_to=None, + normalize_weights: bool = True, + ): + """Generate normalized n-body phase space as numpy arrays. + + Events are generated in the rest frame of the particle, unless `boost_to` is given. + + Note: + In this method, the event weights are returned normalized to their maximum. + + Arguments: + n_events (int): Number of events to generate. + boost_to (optional): Momentum vector of shape (x, 4), where x is optional, to where + the resulting events will be boosted. If not specified, events are generated + in the rest frame of the particle. + normalize_weights (bool, optional): Normalize the event weight to its max + + + Return: + tuple: Result of the generation, which varies with the value of `normalize_weights`: + + + If True, the tuple elements are the normalized event weights as an array of shape + (n_events, ), and the momenta generated particles as a dictionary of arrays of shape + (4, n_events) with particle names as keys. + + + If False, the tuple weights are the unnormalized event weights as an array of shape + (n_events, ), the maximum per-event weights as an array of shape (n_events, ) and the + momenta generated particles as a dictionary of arrays of shape (4, n_events) with particle + names as keys. + + Raise: + tf.errors.InvalidArgumentError: If the the decay is kinematically forbidden. + ValueError: If `n_events` and the size of `boost_to` don't match. + See `GenParticle.generate_unnormalized`. + """ + + # Run generation + warnings.warn( + "This function is depreceated. Use `generate` which does not return a Tensor as well." + ) + generate_tf = self.generate(n_events, boost_to, normalize_weights) + # self._cache = generate_tf + # self._set_cache_validity(True, propagate=True) + return generate_tf + + +# legacy class to warn user about name change +class Particle: + """Deprecated Particle class. + + Renamed to GenParticle. + """ + + def __init__(self): + raise NameError( + "'Particle' has been renamed to 'GenParticle'. Please update your code accordingly." + "For more information, see: https://github.com/zfit/phasespace/issues/22" + ) + + +def nbody_decay(mass_top: float, masses: list, top_name: str = "", names: list = None): + """Shortcut to build an n-body decay of a GenParticle. + + If the particle names are not given, the top particle is called 'top' and the + children 'p_{i}', where i corresponds to their position in the `masses` sequence. + + Arguments: + mass_top (tensor, list): Mass of the top particle. Can be a list of 4-vectors. + masses (list): Masses of the child particles. + name_top (str, optional): Name of the top particle. If not given, the top particle is + named top. + names (list, optional): Names of the child particles. If not given, they are build as + 'p_{i}', where i is given by their ordering in the `masses` list. + + Return: + `GenParticle`: Particle decay. + + Raise: + ValueError: If the length of `masses` and `names` doesn't match. + """ + if not top_name: + top_name = "top" + if not names: + names = [f"p_{num}" for num in range(len(masses))] + if len(names) != len(masses): + raise ValueError("Mismatch in length between children masses and their names.") + return GenParticle(top_name, mass_top).set_children( + *(GenParticle(names[num], mass=mass) for num, mass in enumerate(masses)) + ) + + +def generate_decay(*args, **kwargs): + """Deprecated.""" + raise NameError( + "'generate_decay' has been removed. A similar behavior can be accomplished with 'nbody_decay'. " + "For more information see https://github.com/zfit/phasespace/issues/22" + ) + + +# EOF diff --git a/phasespace/random.py b/phasespace/random.py index 92fedde5..57c90aee 100644 --- a/phasespace/random.py +++ b/phasespace/random.py @@ -1,41 +1,41 @@ -"""Random number generation. - -As the random number generation is not a trivial thing, this module handles it uniformly. - -It mimicks the TensorFlows API on random generators and relies (currently) in global states on the TF states. -Especially on the global random number generator which will be used to get new generators. -""" -from typing import Optional, Union - -import tensorflow as tf - -SeedLike = Optional[Union[int, tf.random.Generator]] - - -def get_rng(seed: SeedLike = None) -> tf.random.Generator: - """Get or create a random number generators of type `tf.random.Generator`. - - This can be used to either retrieve random number generators deterministically from them - - global random number generator from TensorFlow, - - from a random number generator generated from the seed or - - from the random number generator passed. - - Both when using either the global generator or a random number generator is passed, they advance - by exactly one step as `split` is called on them. - - Args: - seed: This can be - - `None` to get the global random number generator - - a numerical seed to create a random number generator - - a `tf.random.Generator`. - - Returns: - A list of `tf.random.Generator` - """ - if seed is None: - rng = tf.random.get_global_generator() - elif not isinstance(seed, tf.random.Generator): # it's a seed, not an rng - rng = tf.random.Generator.from_seed(seed=seed) - else: - rng = seed - return rng +"""Random number generation. + +As the random number generation is not a trivial thing, this module handles it uniformly. + +It mimicks the TensorFlows API on random generators and relies (currently) in global states on the TF states. +Especially on the global random number generator which will be used to get new generators. +""" +from typing import Optional, Union + +import tensorflow as tf + +SeedLike = Optional[Union[int, tf.random.Generator]] + + +def get_rng(seed: SeedLike = None) -> tf.random.Generator: + """Get or create a random number generators of type `tf.random.Generator`. + + This can be used to either retrieve random number generators deterministically from them + - global random number generator from TensorFlow, + - from a random number generator generated from the seed or + - from the random number generator passed. + + Both when using either the global generator or a random number generator is passed, they advance + by exactly one step as `split` is called on them. + + Args: + seed: This can be + - `None` to get the global random number generator + - a numerical seed to create a random number generator + - a `tf.random.Generator`. + + Returns: + A list of `tf.random.Generator` + """ + if seed is None: + rng = tf.random.get_global_generator() + elif not isinstance(seed, tf.random.Generator): # it's a seed, not an rng + rng = tf.random.Generator.from_seed(seed=seed) + else: + rng = seed + return rng diff --git a/pyproject.toml b/pyproject.toml index f7d864aa..affcbfb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ -[build-system] -requires = [ - "setuptools>=42", - "setuptools_scm[toml]>=3.4", - "setuptools_scm_git_archive", - "wheel" -] - -build-backend = "setuptools.build_meta" - -[tool.isort] -profile = "black" -src_paths = ["phasespace", "tests"] +[build-system] +requires = [ + "setuptools>=42", + "setuptools_scm[toml]>=3.4", + "setuptools_scm_git_archive", + "wheel" +] + +build-backend = "setuptools.build_meta" + +[tool.isort] +profile = "black" +src_paths = ["phasespace", "tests"] diff --git a/scripts/prepare_test_samples.cxx b/scripts/prepare_test_samples.cxx index d5efd34b..24448d14 100644 --- a/scripts/prepare_test_samples.cxx +++ b/scripts/prepare_test_samples.cxx @@ -1,123 +1,123 @@ -#include "TROOT.h" -#include "TSystem.h" -#include "TFile.h" -#include "TTree.h" -#include "TLorentzVector.h" -#include "TGenPhaseSpace.h" - -Int_t N_EVENTS = 100000; -Double_t B0_MASS = 5279.58; -Double_t PION_MASS = 139.57018; - - -int prepare_two_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[2] = {PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 2, masses); - - TLorentzVector pion_1, pion_2; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_three_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[3] = {PION_MASS, PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 3, masses); - - TLorentzVector pion_1, pion_2, pion_3; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("pion_3", "TLorentzVector", &pion_3); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - TLorentzVector *p_3 = event.GetDecay(2); - pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_four_body(TString filename) -{ - TFile out_file(filename, "RECREATE"); - TTree out_tree("events", "Generated events"); - - TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); - Double_t masses[4] = {PION_MASS, PION_MASS, PION_MASS, PION_MASS}; - - TGenPhaseSpace event; - event.SetDecay(B, 4, masses); - - TLorentzVector pion_1, pion_2, pion_3, pion_4; - Double_t weight; - out_tree.Branch("pion_1", "TLorentzVector", &pion_1); - out_tree.Branch("pion_2", "TLorentzVector", &pion_2); - out_tree.Branch("pion_3", "TLorentzVector", &pion_3); - out_tree.Branch("pion_4", "TLorentzVector", &pion_4); - out_tree.Branch("weight", &weight, "weight/D"); - - - for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); - TLorentzVector *p_2 = event.GetDecay(1); - pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); - TLorentzVector *p_3 = event.GetDecay(2); - pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); - TLorentzVector *p_4 = event.GetDecay(3); - pion_4.SetPxPyPzE(p_4->Px(), p_4->Py(), p_4->Pz(), p_4->E()); - out_tree.Fill(); - } - out_file.Write(); - return 0; -} - - -int prepare_test_samples(TString two_body_file, TString three_body_file, TString four_body_file){ - if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); - - prepare_two_body(two_body_file); - prepare_three_body(three_body_file); - prepare_four_body(four_body_file); - - return 0; -} +#include "TROOT.h" +#include "TSystem.h" +#include "TFile.h" +#include "TTree.h" +#include "TLorentzVector.h" +#include "TGenPhaseSpace.h" + +Int_t N_EVENTS = 100000; +Double_t B0_MASS = 5279.58; +Double_t PION_MASS = 139.57018; + + +int prepare_two_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[2] = {PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 2, masses); + + TLorentzVector pion_1, pion_2; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_three_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[3] = {PION_MASS, PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 3, masses); + + TLorentzVector pion_1, pion_2, pion_3; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("pion_3", "TLorentzVector", &pion_3); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + TLorentzVector *p_3 = event.GetDecay(2); + pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_four_body(TString filename) +{ + TFile out_file(filename, "RECREATE"); + TTree out_tree("events", "Generated events"); + + TLorentzVector B(0.0, 0.0, 0.0, B0_MASS); + Double_t masses[4] = {PION_MASS, PION_MASS, PION_MASS, PION_MASS}; + + TGenPhaseSpace event; + event.SetDecay(B, 4, masses); + + TLorentzVector pion_1, pion_2, pion_3, pion_4; + Double_t weight; + out_tree.Branch("pion_1", "TLorentzVector", &pion_1); + out_tree.Branch("pion_2", "TLorentzVector", &pion_2); + out_tree.Branch("pion_3", "TLorentzVector", &pion_3); + out_tree.Branch("pion_4", "TLorentzVector", &pion_4); + out_tree.Branch("weight", &weight, "weight/D"); + + + for (Int_t n=0; nPx(), p_1->Py(), p_1->Pz(), p_1->E()); + TLorentzVector *p_2 = event.GetDecay(1); + pion_2.SetPxPyPzE(p_2->Px(), p_2->Py(), p_2->Pz(), p_2->E()); + TLorentzVector *p_3 = event.GetDecay(2); + pion_3.SetPxPyPzE(p_3->Px(), p_3->Py(), p_3->Pz(), p_3->E()); + TLorentzVector *p_4 = event.GetDecay(3); + pion_4.SetPxPyPzE(p_4->Px(), p_4->Py(), p_4->Pz(), p_4->E()); + out_tree.Fill(); + } + out_file.Write(); + return 0; +} + + +int prepare_test_samples(TString two_body_file, TString three_body_file, TString four_body_file){ + if (!gROOT->GetClass("TGenPhaseSpace")) gSystem->Load("libPhysics"); + + prepare_two_body(two_body_file); + prepare_three_body(three_body_file); + prepare_four_body(four_body_file); + + return 0; +} diff --git a/setup.cfg b/setup.cfg index 3d0cf47a..3a5d7a3b 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,95 +1,95 @@ -[metadata] -name = phasespace -description = TensorFlow implementation of the Raubold and Lynch method for n-body events -long_description = file: README.rst -long_description_content_type = text/x-rst -url = https://github.com/zfit/phasespace -author = Albert Puig Navarro -author_email = apuignav@gmail.com -maintainer = zfit -maintainer_email = zfit@physik.uzh.ch -license = BSD-3-Clause -license_file = LICENSE -classifiers = - Development Status :: 5 - Production/Stable - Intended Audience :: Science/Research - License :: OSI Approved :: BSD License - Natural Language :: English - Operating System :: MacOS - Operating System :: Unix - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.6 - Programming Language :: Python :: 3.7 - Programming Language :: Python :: 3.8 - Programming Language :: Python :: 3.9 - Topic :: Scientific/Engineering :: Physics -keywords = TensorFlow, phasespace, HEP - -[options] -packages = find: -setup_requires = - setuptools_scm -install_requires = - tensorflow>=2.5,<2.7 # tensorflow.experimental.numpy - tensorflow_probability>=0.11 -python_requires = >=3.6 -include_package_data = True -testpaths = tests -zip_safe = False - -[options.extras_require] -test = - awkward - coverage - flaky - matplotlib - numpy - pytest - pytest-cov - pytest-xdist - scipy - uproot4 - wget -doc = - Sphinx - sphinx_bootstrap_theme - jupyter_sphinx - sphinx-math-dollar -dev = - %(doc)s - %(test)s - bumpversion - pre-commit - twine - watchdog - -[bdist_wheel] -universal = 1 - -[flake8] -exclude = - benchmark, - data, - dist, - docs, - paper, - scripts, - utils -max-line-length = 110 -statistics = True -max-complexity = 30 - -[coverage:run] -branch = True -include = - phasespace/* - -[tool:pytest] -addopts = - --color=yes - --ignore=setup.py -filterwarnings = - ignore:.*the imp module is deprecated in favour of importlib.*:DeprecationWarning -norecursedirs = - tests/helpers +[metadata] +name = phasespace +description = TensorFlow implementation of the Raubold and Lynch method for n-body events +long_description = file: README.rst +long_description_content_type = text/x-rst +url = https://github.com/zfit/phasespace +author = Albert Puig Navarro +author_email = apuignav@gmail.com +maintainer = zfit +maintainer_email = zfit@physik.uzh.ch +license = BSD-3-Clause +license_file = LICENSE +classifiers = + Development Status :: 5 - Production/Stable + Intended Audience :: Science/Research + License :: OSI Approved :: BSD License + Natural Language :: English + Operating System :: MacOS + Operating System :: Unix + Programming Language :: Python :: 3 + Programming Language :: Python :: 3 :: Only + Programming Language :: Python :: 3.6 + Programming Language :: Python :: 3.7 + Programming Language :: Python :: 3.8 + Programming Language :: Python :: 3.9 + Topic :: Scientific/Engineering :: Physics +keywords = TensorFlow, phasespace, HEP + +[options] +packages = find: +setup_requires = + setuptools_scm +install_requires = + tensorflow>=2.5,<2.7 # tensorflow.experimental.numpy + tensorflow_probability>=0.11 +python_requires = >=3.6 +include_package_data = True +testpaths = tests +zip_safe = False + +[options.extras_require] +test = + awkward + coverage + flaky + matplotlib + numpy + pytest + pytest-cov + pytest-xdist + scipy + uproot4 + wget +doc = + Sphinx + sphinx_bootstrap_theme + jupyter_sphinx + sphinx-math-dollar +dev = + %(doc)s + %(test)s + bumpversion + pre-commit + twine + watchdog + +[bdist_wheel] +universal = 1 + +[flake8] +exclude = + benchmark, + data, + dist, + docs, + paper, + scripts, + utils +max-line-length = 110 +statistics = True +max-complexity = 30 + +[coverage:run] +branch = True +include = + phasespace/* + +[tool:pytest] +addopts = + --color=yes + --ignore=setup.py +filterwarnings = + ignore:.*the imp module is deprecated in favour of importlib.*:DeprecationWarning +norecursedirs = + tests/helpers diff --git a/setup.py b/setup.py index bdfb6d20..edbdeba7 100644 --- a/setup.py +++ b/setup.py @@ -1,20 +1,20 @@ -#!/usr/bin/env python - -"""The setup script.""" - -import os - -from setuptools import setup - -here = os.path.abspath(os.path.dirname(__file__)) - -with open(os.path.join(here, "README.rst"), encoding="utf-8") as readme_file: - readme = readme_file.read() - -with open(os.path.join(here, "CHANGELOG.rst"), encoding="utf-8") as history_file: - history = history_file.read() - -setup( - long_description=readme.replace(":math:", "") + "\n\n" + history, - use_scm_version=True, -) +#!/usr/bin/env python + +"""The setup script.""" + +import os + +from setuptools import setup + +here = os.path.abspath(os.path.dirname(__file__)) + +with open(os.path.join(here, "README.rst"), encoding="utf-8") as readme_file: + readme = readme_file.read() + +with open(os.path.join(here, "CHANGELOG.rst"), encoding="utf-8") as history_file: + history = history_file.read() + +setup( + long_description=readme.replace(":math:", "") + "\n\n" + history, + use_scm_version=True, +) diff --git a/tests/conftest.py b/tests/conftest.py index c9c23692..f0de9679 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,4 @@ -import os -import sys - -sys.path.append(os.path.join(os.path.dirname(__file__), "helpers")) +import os +import sys + +sys.path.append(os.path.join(os.path.dirname(__file__), "helpers")) diff --git a/tests/fromdecay/__init__.py b/tests/fulldecay/__init__.py similarity index 63% rename from tests/fromdecay/__init__.py rename to tests/fulldecay/__init__.py index 87d1a928..b98e9386 100644 --- a/tests/fromdecay/__init__.py +++ b/tests/fulldecay/__init__.py @@ -1,4 +1,4 @@ -import pytest - -# This makes it so that assert errors are more helpful for e.g., the check_norm helper function -pytest.register_assert_rewrite("fromdecay.test_fulldecay") +import pytest + +# This makes it so that assert errors are more helpful for e.g., the check_norm helper function +pytest.register_assert_rewrite("fulldecay.test_fulldecay") diff --git a/tests/fromdecay/example_decay_chains.py b/tests/fulldecay/example_decay_chains.py similarity index 87% rename from tests/fromdecay/example_decay_chains.py rename to tests/fulldecay/example_decay_chains.py index c04d867d..ef71fda6 100644 --- a/tests/fromdecay/example_decay_chains.py +++ b/tests/fulldecay/example_decay_chains.py @@ -25,5 +25,5 @@ ): decay_mode["zfit"] = mass_function -# D*+ particle that has multiple children, grandchild particles, many of which can decay in multiple ways. +# D*+ particle that has multiple child particles, grandchild particles, many of which can decay in multiple ways. dstarplus_big_decay = dfp.build_decay_chains("D*+") diff --git a/tests/fromdecay/example_decays.dec b/tests/fulldecay/example_decays.dec similarity index 96% rename from tests/fromdecay/example_decays.dec rename to tests/fulldecay/example_decays.dec index f5bed887..c1eae76e 100644 --- a/tests/fromdecay/example_decays.dec +++ b/tests/fulldecay/example_decays.dec @@ -1,31 +1,31 @@ -# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec -# Example decay chain for testing purposes -# Considered by itself, this file in in fact incomplete, -# as there are no instructions on how to decay the anti-D0 and the D-! - -Decay D*+ -0.6770 D0 pi+ VSS; -0.3070 D+ pi0 VSS; -0.0160 D+ gamma VSP_PWAVE; -Enddecay - -Decay D*- -0.6770 anti-D0 pi- VSS; -0.3070 D- pi0 VSS; -0.0160 D- gamma VSP_PWAVE; -Enddecay - -Decay D0 -1.0 K- pi+ PHSP; -Enddecay - -Decay D+ -1.0 K- pi+ pi+ pi0 PHSP; -Enddecay - -Decay pi0 -0.988228297 gamma gamma PHSP; -0.011738247 e+ e- gamma PI0_DALITZ; -0.000033392 e+ e+ e- e- PHSP; -0.000000065 e+ e- PHSP; -Enddecay +# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec +# Example decay chain for testing purposes +# Considered by itself, this file in in fact incomplete, +# as there are no instructions on how to decay the anti-D0 and the D-! + +Decay D*+ +0.6770 D0 pi+ VSS; +0.3070 D+ pi0 VSS; +0.0160 D+ gamma VSP_PWAVE; +Enddecay + +Decay D*- +0.6770 anti-D0 pi- VSS; +0.3070 D- pi0 VSS; +0.0160 D- gamma VSP_PWAVE; +Enddecay + +Decay D0 +1.0 K- pi+ PHSP; +Enddecay + +Decay D+ +1.0 K- pi+ pi+ pi0 PHSP; +Enddecay + +Decay pi0 +0.988228297 gamma gamma PHSP; +0.011738247 e+ e- gamma PI0_DALITZ; +0.000033392 e+ e+ e- e- PHSP; +0.000000065 e+ e- PHSP; +Enddecay diff --git a/tests/fromdecay/test_fulldecay.py b/tests/fulldecay/test_fulldecay.py similarity index 84% rename from tests/fromdecay/test_fulldecay.py rename to tests/fulldecay/test_fulldecay.py index ad04cac1..4baa7568 100644 --- a/tests/fromdecay/test_fulldecay.py +++ b/tests/fulldecay/test_fulldecay.py @@ -1,9 +1,9 @@ from numpy.testing import assert_almost_equal -from phasespace.fromdecay import FullDecay -from phasespace.fromdecay.mass_functions import _DEFAULT_CONVERTER +from phasespace.fulldecay import FullDecay +from phasespace.fulldecay.mass_functions import _DEFAULT_CONVERTER -from . import example_decay_chains +from .example_decay_chains import * # TODO remove * since it is bad practice? def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: @@ -33,7 +33,7 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: def test_single_chain(): """Test converting a decaylanguage dict with only one possible decay.""" - container = FullDecay.from_dict(example_decay_chains.dplus_single, tolerance=1e-10) + container = FullDecay.from_dict(dplus_single, tolerance=1e-10) output_decay = container.gen_particles assert len(output_decay) == 1 prob, gen = output_decay[0] @@ -56,7 +56,7 @@ def test_single_chain(): def test_branching_children(): """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" - container = FullDecay.from_dict(example_decay_chains.pi0_4branches, tolerance=1e-10) + container = FullDecay.from_dict(pi0_4branches, 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) @@ -66,7 +66,7 @@ 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 = FullDecay.from_dict(example_decay_chains.dplus_4grandbranches) + container = FullDecay.from_dict(dplus_4grandbranches) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) @@ -77,7 +77,7 @@ def test_branching_grandchilden(): 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 = dplus_4grandbranches.copy() dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" container = FullDecay.from_dict( dplus_4grandbranches_massfunc, @@ -98,7 +98,7 @@ def test_mass_converter(): def test_big_decay(): - container = FullDecay.from_dict(example_decay_chains.dstarplus_big_decay) + container = FullDecay.from_dict(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) diff --git a/tests/fromdecay/test_mass_functions.py b/tests/fulldecay/test_mass_functions.py similarity index 94% rename from tests/fromdecay/test_mass_functions.py rename to tests/fulldecay/test_mass_functions.py index e8c4eac0..a7c12d3e 100644 --- a/tests/fromdecay/test_mass_functions.py +++ b/tests/fulldecay/test_mass_functions.py @@ -5,7 +5,7 @@ import tensorflow_probability as tfp from particle import Particle -import phasespace.fromdecay.mass_functions as mf +import phasespace.fulldecay.mass_functions as mf _kstarz = Particle.from_evtgen_name("K*0") KSTARZ_MASS = _kstarz.mass diff --git a/tests/helpers/decays.py b/tests/helpers/decays.py index 313b2e11..6ff09099 100644 --- a/tests/helpers/decays.py +++ b/tests/helpers/decays.py @@ -1,80 +1,80 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file decays.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Some physics models to test with.""" - -import tensorflow as tf -import tensorflow_probability as tfp - -from phasespace import GenParticle - -# Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat) -B0_MASS = 5279.58 -PION_MASS = 139.57018 -KAON_MASS = 493.677 -K1_MASS = 1272.0 -K1_WIDTH = 90.0 -KSTARZ_MASS = 895.81 -KSTARZ_WIDTH = 47.4 - - -def b0_to_kstar_gamma(kstar_width=KSTARZ_WIDTH): - """Generate B0 -> K*gamma.""" - - def kstar_mass(min_mass, max_mass, n_events): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(kstar_width, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - if kstar_width > 0: - kstar_mass = tfp.distributions.TruncatedNormal( - loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass - ).sample() - return kstar_mass - - return GenParticle("B0", B0_MASS).set_children( - GenParticle("K*0", mass=kstar_mass).set_children( - GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) - ), - GenParticle("gamma", mass=0.0), - ) - - -def bp_to_k1_kstar_pi_gamma(k1_width=K1_WIDTH, kstar_width=KSTARZ_WIDTH): - """Generate B+ -> K1 (-> K* (->K pi) pi) gamma.""" - - def res_mass(mass, width, min_mass, max_mass, n_events): - mass = tf.cast(mass, tf.float64) - width = tf.cast(width, tf.float64) - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - masses = tf.broadcast_to(mass, shape=(n_events,)) - if kstar_width > 0: - masses = tfp.distributions.TruncatedNormal( - loc=masses, scale=width, low=min_mass, high=max_mass - ).sample() - return masses - - def k1_mass(min_mass, max_mass, n_events): - return res_mass(K1_MASS, k1_width, min_mass, max_mass, n_events) - - def kstar_mass(min_mass, max_mass, n_events): - return res_mass(KSTARZ_MASS, kstar_width, min_mass, max_mass, n_events) - - return GenParticle("B+", B0_MASS).set_children( - GenParticle("K1+", mass=k1_mass).set_children( - GenParticle("K*0", mass=kstar_mass).set_children( - GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) - ), - GenParticle("pi+", mass=PION_MASS), - ), - GenParticle("gamma", mass=0.0), - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file decays.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Some physics models to test with.""" + +import tensorflow as tf +import tensorflow_probability as tfp + +from phasespace import GenParticle + +# Use RapidSim values (https://github.com/gcowan/RapidSim/blob/master/config/particles.dat) +B0_MASS = 5279.58 +PION_MASS = 139.57018 +KAON_MASS = 493.677 +K1_MASS = 1272.0 +K1_WIDTH = 90.0 +KSTARZ_MASS = 895.81 +KSTARZ_WIDTH = 47.4 + + +def b0_to_kstar_gamma(kstar_width=KSTARZ_WIDTH): + """Generate B0 -> K*gamma.""" + + def kstar_mass(min_mass, max_mass, n_events): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(kstar_width, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + if kstar_width > 0: + kstar_mass = tfp.distributions.TruncatedNormal( + loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass + ).sample() + return kstar_mass + + return GenParticle("B0", B0_MASS).set_children( + GenParticle("K*0", mass=kstar_mass).set_children( + GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) + ), + GenParticle("gamma", mass=0.0), + ) + + +def bp_to_k1_kstar_pi_gamma(k1_width=K1_WIDTH, kstar_width=KSTARZ_WIDTH): + """Generate B+ -> K1 (-> K* (->K pi) pi) gamma.""" + + def res_mass(mass, width, min_mass, max_mass, n_events): + mass = tf.cast(mass, tf.float64) + width = tf.cast(width, tf.float64) + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + masses = tf.broadcast_to(mass, shape=(n_events,)) + if kstar_width > 0: + masses = tfp.distributions.TruncatedNormal( + loc=masses, scale=width, low=min_mass, high=max_mass + ).sample() + return masses + + def k1_mass(min_mass, max_mass, n_events): + return res_mass(K1_MASS, k1_width, min_mass, max_mass, n_events) + + def kstar_mass(min_mass, max_mass, n_events): + return res_mass(KSTARZ_MASS, kstar_width, min_mass, max_mass, n_events) + + return GenParticle("B+", B0_MASS).set_children( + GenParticle("K1+", mass=k1_mass).set_children( + GenParticle("K*0", mass=kstar_mass).set_children( + GenParticle("K+", mass=KAON_MASS), GenParticle("pi-", mass=PION_MASS) + ), + GenParticle("pi+", mass=PION_MASS), + ), + GenParticle("gamma", mass=0.0), + ) + + +# EOF diff --git a/tests/helpers/plotting.py b/tests/helpers/plotting.py index 1e222b5a..3051c214 100644 --- a/tests/helpers/plotting.py +++ b/tests/helpers/plotting.py @@ -1,28 +1,28 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file plotting.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Plotting helpers for tests.""" - -import numpy as np - - -def make_norm_histo(array, range_, weights=None): - """Make histo and modify dimensions.""" - histo = np.histogram(array, 100, range=range_, weights=weights)[0] - return histo / np.sum(histo) - - -def mass(vector): - """Calculate mass scalar for Lorentz 4-momentum.""" - return np.sqrt( - np.sum( - vector * vector * np.reshape(np.array([-1.0, -1.0, -1.0, 1.0]), (1, 4)), - axis=1, - ) - ) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file plotting.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Plotting helpers for tests.""" + +import numpy as np + + +def make_norm_histo(array, range_, weights=None): + """Make histo and modify dimensions.""" + histo = np.histogram(array, 100, range=range_, weights=weights)[0] + return histo / np.sum(histo) + + +def mass(vector): + """Calculate mass scalar for Lorentz 4-momentum.""" + return np.sqrt( + np.sum( + vector * vector * np.reshape(np.array([-1.0, -1.0, -1.0, 1.0]), (1, 4)), + axis=1, + ) + ) + + +# EOF diff --git a/tests/helpers/rapidsim.py b/tests/helpers/rapidsim.py index d6d94cc0..94821a5c 100644 --- a/tests/helpers/rapidsim.py +++ b/tests/helpers/rapidsim.py @@ -1,111 +1,111 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file rapidsim.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 07.03.2019 -# ============================================================================= -"""Utils to crossheck against RapidSim.""" - -import os - -import numpy as np -import uproot4 - -BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) -FONLL_FILE = os.path.join(BASE_PATH, "data", "fonll", "LHC{}{}.root") - - -def get_fonll_histos(energy, quark): - with uproot4.open(FONLL_FILE.format(quark, int(energy))) as histo_file: - return histo_file["pT"], histo_file["eta"] - - -def generate_fonll(mass, beam_energy, quark, n_events): - def analyze_histo(histo): - x_axis = histo.axis(0) - x_bins = x_axis.edges() - bin_width = x_axis.width - bin_centers = x_bins[:-1] + bin_width / 2 - normalized_values = histo.values() / np.sum(histo.values()) - return bin_centers, normalized_values - - pt_histo, eta_histo = get_fonll_histos(beam_energy, quark) - pt_bin_centers, pt_normalized_values = analyze_histo(pt_histo) - eta_bin_centers, eta_normalized_values = analyze_histo(eta_histo) - pt_rand = np.random.choice( - pt_bin_centers, - size=n_events, - p=pt_normalized_values, - ) - pt_rand = 1_000 * np.abs(pt_rand) - eta_rand = np.random.choice( - eta_bin_centers, - size=n_events, - p=eta_normalized_values, - ) - phi_rand = np.random.uniform(0, 2 * np.pi, size=n_events) - px = pt_rand * np.cos(phi_rand) - py = pt_rand * np.sin(phi_rand) - pz = pt_rand * np.sinh(eta_rand) - e = np.sqrt(px * px + py * py + pz * pz + mass * mass) - return np.stack([px, py, pz, e]) - - -def load_generated_histos(file_name, particles): - with uproot4.open(file_name) as rapidsim_file: - return { - particle: [ - rapidsim_file.get(f"{particle}_{coord}_TRUE").array(library="np") - for coord in ("PX", "PY", "PZ", "E") - ] - for particle in particles - } - - -def get_tree(file_name, top_particle, particles): - """Load a RapidSim tree.""" - with uproot4.open(file_name) as rapidsim_file: - tree = rapidsim_file["DecayTree"] - return { - particle: np.stack( - [ - 1000.0 * tree[f"{particle}_{coord}_TRUE"].array(library="np") - for coord in ("PX", "PY", "PZ", "E") - ] - ) - for particle in particles - } - - -def get_tree_in_b_rest_frame(file_name, top_particle, particles): - def lorentz_boost(part_to_boost, boost): - """ - Perform Lorentz boost - vector : 4-vector to be boosted - boostvector: boost vector. Can be either 3-vector or 4-vector (only spatial components - are used) - """ - boost_vec = -boost[:3, :] / boost[3, :] - b2 = np.sum(boost_vec * boost_vec, axis=0) - gamma = 1.0 / np.sqrt(1.0 - b2) - gamma2 = (gamma - 1.0) / b2 - part_time = part_to_boost[3, :] - part_space = part_to_boost[:3, :] - bp = np.sum(part_space * boost_vec, axis=0) - return np.concatenate( - [ - part_space + (gamma2 * bp + gamma * part_time) * boost_vec, - np.expand_dims(gamma * (part_time + bp), axis=0), - ], - axis=0, - ) - - part_dict = get_tree(file_name, top_particle, list(particles) + [top_particle]) - top_parts = part_dict.pop(top_particle) - return { - part_name: lorentz_boost(part, top_parts) - for part_name, part in part_dict.items() - } - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file rapidsim.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 07.03.2019 +# ============================================================================= +"""Utils to crossheck against RapidSim.""" + +import os + +import numpy as np +import uproot4 + +BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +FONLL_FILE = os.path.join(BASE_PATH, "data", "fonll", "LHC{}{}.root") + + +def get_fonll_histos(energy, quark): + with uproot4.open(FONLL_FILE.format(quark, int(energy))) as histo_file: + return histo_file["pT"], histo_file["eta"] + + +def generate_fonll(mass, beam_energy, quark, n_events): + def analyze_histo(histo): + x_axis = histo.axis(0) + x_bins = x_axis.edges() + bin_width = x_axis.width + bin_centers = x_bins[:-1] + bin_width / 2 + normalized_values = histo.values() / np.sum(histo.values()) + return bin_centers, normalized_values + + pt_histo, eta_histo = get_fonll_histos(beam_energy, quark) + pt_bin_centers, pt_normalized_values = analyze_histo(pt_histo) + eta_bin_centers, eta_normalized_values = analyze_histo(eta_histo) + pt_rand = np.random.choice( + pt_bin_centers, + size=n_events, + p=pt_normalized_values, + ) + pt_rand = 1_000 * np.abs(pt_rand) + eta_rand = np.random.choice( + eta_bin_centers, + size=n_events, + p=eta_normalized_values, + ) + phi_rand = np.random.uniform(0, 2 * np.pi, size=n_events) + px = pt_rand * np.cos(phi_rand) + py = pt_rand * np.sin(phi_rand) + pz = pt_rand * np.sinh(eta_rand) + e = np.sqrt(px * px + py * py + pz * pz + mass * mass) + return np.stack([px, py, pz, e]) + + +def load_generated_histos(file_name, particles): + with uproot4.open(file_name) as rapidsim_file: + return { + particle: [ + rapidsim_file.get(f"{particle}_{coord}_TRUE").array(library="np") + for coord in ("PX", "PY", "PZ", "E") + ] + for particle in particles + } + + +def get_tree(file_name, top_particle, particles): + """Load a RapidSim tree.""" + with uproot4.open(file_name) as rapidsim_file: + tree = rapidsim_file["DecayTree"] + return { + particle: np.stack( + [ + 1000.0 * tree[f"{particle}_{coord}_TRUE"].array(library="np") + for coord in ("PX", "PY", "PZ", "E") + ] + ) + for particle in particles + } + + +def get_tree_in_b_rest_frame(file_name, top_particle, particles): + def lorentz_boost(part_to_boost, boost): + """ + Perform Lorentz boost + vector : 4-vector to be boosted + boostvector: boost vector. Can be either 3-vector or 4-vector (only spatial components + are used) + """ + boost_vec = -boost[:3, :] / boost[3, :] + b2 = np.sum(boost_vec * boost_vec, axis=0) + gamma = 1.0 / np.sqrt(1.0 - b2) + gamma2 = (gamma - 1.0) / b2 + part_time = part_to_boost[3, :] + part_space = part_to_boost[:3, :] + bp = np.sum(part_space * boost_vec, axis=0) + return np.concatenate( + [ + part_space + (gamma2 * bp + gamma * part_time) * boost_vec, + np.expand_dims(gamma * (part_time + bp), axis=0), + ], + axis=0, + ) + + part_dict = get_tree(file_name, top_particle, list(particles) + [top_particle]) + top_parts = part_dict.pop(top_particle) + return { + part_name: lorentz_boost(part, top_parts) + for part_name, part in part_dict.items() + } + + +# EOF diff --git a/tests/test_chain.py b/tests/test_chain.py index 555c8cf9..1bcf201e 100644 --- a/tests/test_chain.py +++ b/tests/test_chain.py @@ -1,137 +1,137 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_chain.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 01.03.2019 -# ============================================================================= -"""Test decay chain tools.""" - -import os -import sys - -import numpy as np -import pytest - -from phasespace import GenParticle - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays # noqa: E402 - - -def test_name_clashes(): - """Test clashes in particle naming.""" - # In children - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - ) - # With itself - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Top", mass=decays.KSTARZ_MASS), - GenParticle("Kstarz", mass=decays.KSTARZ_MASS), - ) - # In grandchildren - with pytest.raises(KeyError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( - GenParticle("K+", mass=decays.KAON_MASS), - GenParticle("pi-", mass=decays.PION_MASS), - ), - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( - GenParticle("K+", mass=decays.KAON_MASS), - GenParticle("pi-_1", mass=decays.PION_MASS), - ), - ) - - -def test_wrong_children(): - """Test wrong number of children.""" - with pytest.raises(ValueError): - GenParticle("Top", 0).set_children( - GenParticle("Kstarz0", mass=decays.KSTARZ_MASS) - ) - - -def test_grandchildren(): - """Test that grandchildren detection is correct.""" - top = GenParticle("Top", 0) - assert not top.has_children - assert not top.has_grandchildren - assert not top.set_children( - GenParticle("Child1", mass=decays.KSTARZ_MASS), - GenParticle("Child2", mass=decays.KSTARZ_MASS), - ).has_grandchildren - - -def test_reset_children(): - """Test when children are set twice.""" - top = GenParticle("Top", 0).set_children( - GenParticle("Child1", mass=decays.KSTARZ_MASS), - GenParticle("Child2", mass=decays.KSTARZ_MASS), - ) - with pytest.raises(ValueError): - top.set_children( - GenParticle("Child3", mass=decays.KSTARZ_MASS), - GenParticle("Child4", mass=decays.KSTARZ_MASS), - ) - - -def test_no_children(): - """Test when no children have been configured.""" - top = GenParticle("Top", 0) - with pytest.raises(ValueError): - top.generate(n_events=1) - - -def test_resonance_top(): - """Test when a resonance is used as the top particle.""" - kstar = decays.b0_to_kstar_gamma().children[0] - with pytest.raises(ValueError): - kstar.generate(n_events=1) - - -def test_kstargamma(): - """Test B0 -> K*gamma.""" - decay = decays.b0_to_kstar_gamma() - norm_weights, particles = decay.generate(n_events=1000) - assert norm_weights.shape[0] == 1000 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 4 - assert set(particles.keys()) == {"K*0", "gamma", "K+", "pi-"} - assert all(part.shape == (1000, 4) for part in particles.values()) - - -def test_k1gamma(): - """Test B+ -> K1 (K*pi) gamma.""" - decay = decays.bp_to_k1_kstar_pi_gamma() - norm_weights, particles = decay.generate(n_events=1000) - assert norm_weights.shape[0] == 1000 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 6 - assert set(particles.keys()) == {"K1+", "K*0", "gamma", "K+", "pi-", "pi+"} - assert all(part.shape == (1000, 4) for part in particles.values()) - - -def test_repr(): - """Test string representation.""" - b0 = decays.b0_to_kstar_gamma() - kst = b0.children[0] - assert ( - str(b0) - == "" - ) - assert ( - str(kst) - == "" - ) - - -if __name__ == "__main__": - test_name_clashes() - test_kstargamma() - test_k1gamma() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_chain.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 01.03.2019 +# ============================================================================= +"""Test decay chain tools.""" + +import os +import sys + +import numpy as np +import pytest + +from phasespace import GenParticle + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays # noqa: E402 + + +def test_name_clashes(): + """Test clashes in particle naming.""" + # In children + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + ) + # With itself + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Top", mass=decays.KSTARZ_MASS), + GenParticle("Kstarz", mass=decays.KSTARZ_MASS), + ) + # In grandchildren + with pytest.raises(KeyError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( + GenParticle("K+", mass=decays.KAON_MASS), + GenParticle("pi-", mass=decays.PION_MASS), + ), + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS).set_children( + GenParticle("K+", mass=decays.KAON_MASS), + GenParticle("pi-_1", mass=decays.PION_MASS), + ), + ) + + +def test_wrong_children(): + """Test wrong number of children.""" + with pytest.raises(ValueError): + GenParticle("Top", 0).set_children( + GenParticle("Kstarz0", mass=decays.KSTARZ_MASS) + ) + + +def test_grandchildren(): + """Test that grandchildren detection is correct.""" + top = GenParticle("Top", 0) + assert not top.has_children + assert not top.has_grandchildren + assert not top.set_children( + GenParticle("Child1", mass=decays.KSTARZ_MASS), + GenParticle("Child2", mass=decays.KSTARZ_MASS), + ).has_grandchildren + + +def test_reset_children(): + """Test when children are set twice.""" + top = GenParticle("Top", 0).set_children( + GenParticle("Child1", mass=decays.KSTARZ_MASS), + GenParticle("Child2", mass=decays.KSTARZ_MASS), + ) + with pytest.raises(ValueError): + top.set_children( + GenParticle("Child3", mass=decays.KSTARZ_MASS), + GenParticle("Child4", mass=decays.KSTARZ_MASS), + ) + + +def test_no_children(): + """Test when no children have been configured.""" + top = GenParticle("Top", 0) + with pytest.raises(ValueError): + top.generate(n_events=1) + + +def test_resonance_top(): + """Test when a resonance is used as the top particle.""" + kstar = decays.b0_to_kstar_gamma().children[0] + with pytest.raises(ValueError): + kstar.generate(n_events=1) + + +def test_kstargamma(): + """Test B0 -> K*gamma.""" + decay = decays.b0_to_kstar_gamma() + norm_weights, particles = decay.generate(n_events=1000) + assert norm_weights.shape[0] == 1000 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 4 + assert set(particles.keys()) == {"K*0", "gamma", "K+", "pi-"} + assert all(part.shape == (1000, 4) for part in particles.values()) + + +def test_k1gamma(): + """Test B+ -> K1 (K*pi) gamma.""" + decay = decays.bp_to_k1_kstar_pi_gamma() + norm_weights, particles = decay.generate(n_events=1000) + assert norm_weights.shape[0] == 1000 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 6 + assert set(particles.keys()) == {"K1+", "K*0", "gamma", "K+", "pi-", "pi+"} + assert all(part.shape == (1000, 4) for part in particles.values()) + + +def test_repr(): + """Test string representation.""" + b0 = decays.b0_to_kstar_gamma() + kst = b0.children[0] + assert ( + str(b0) + == "" + ) + assert ( + str(kst) + == "" + ) + + +if __name__ == "__main__": + test_name_clashes() + test_kstargamma() + test_k1gamma() + +# EOF diff --git a/tests/test_generate.py b/tests/test_generate.py index 548b6a97..260918cd 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -1,86 +1,86 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_generate.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Basic dimensionality tests.""" - -import os -import sys - -import numpy as np -import pytest - -import phasespace - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays # noqa: E402 - -B0_MASS = decays.B0_MASS -PION_MASS = decays.PION_MASS - - -def test_one_event(): - """Test B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=1) - assert norm_weights.shape[0] == 1 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (1, 4) for part in particles.values()) - - -def test_one_event_tf(): - """Test B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=1) - - assert norm_weights.shape[0] == 1 - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (1, 4) for part in particles.values()) - - -@pytest.mark.parametrize("n_events", argvalues=[5, 523]) -def test_n_events(n_events): - """Test 5 B->pi pi pi.""" - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - norm_weights, particles = decay.generate(n_events=n_events) - assert norm_weights.shape[0] == n_events - assert np.alltrue(norm_weights < 1) - assert len(particles) == 3 - assert all(part.shape == (n_events, 4) for part in particles.values()) - - -def test_deterministic_events(): - decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - common_seed = 36 - norm_weights_seeded1, particles_seeded1 = decay.generate( - n_events=100, seed=common_seed - ) - norm_weights_global, particles_global = decay.generate(n_events=100) - norm_weights_rnd, particles_rnd = decay.generate(n_events=100, seed=152) - norm_weights_seeded2, particles_seeded2 = decay.generate( - n_events=100, seed=common_seed - ) - - np.testing.assert_allclose(norm_weights_seeded1, norm_weights_seeded2) - for part1, part2 in zip(particles_seeded1.values(), particles_seeded2.values()): - np.testing.assert_allclose(part1, part2) - - assert not np.allclose(norm_weights_seeded1, norm_weights_rnd) - for part1, part2 in zip(particles_seeded1.values(), particles_rnd.values()): - assert not np.allclose(part1, part2) - - assert not np.allclose(norm_weights_global, norm_weights_rnd) - for part1, part2 in zip(particles_global.values(), particles_rnd.values()): - assert not np.allclose(part1, part2) - - -if __name__ == "__main__": - test_one_event() - test_n_events(5) - - # EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_generate.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Basic dimensionality tests.""" + +import os +import sys + +import numpy as np +import pytest + +import phasespace + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays # noqa: E402 + +B0_MASS = decays.B0_MASS +PION_MASS = decays.PION_MASS + + +def test_one_event(): + """Test B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=1) + assert norm_weights.shape[0] == 1 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (1, 4) for part in particles.values()) + + +def test_one_event_tf(): + """Test B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=1) + + assert norm_weights.shape[0] == 1 + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (1, 4) for part in particles.values()) + + +@pytest.mark.parametrize("n_events", argvalues=[5, 523]) +def test_n_events(n_events): + """Test 5 B->pi pi pi.""" + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + norm_weights, particles = decay.generate(n_events=n_events) + assert norm_weights.shape[0] == n_events + assert np.alltrue(norm_weights < 1) + assert len(particles) == 3 + assert all(part.shape == (n_events, 4) for part in particles.values()) + + +def test_deterministic_events(): + decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + common_seed = 36 + norm_weights_seeded1, particles_seeded1 = decay.generate( + n_events=100, seed=common_seed + ) + norm_weights_global, particles_global = decay.generate(n_events=100) + norm_weights_rnd, particles_rnd = decay.generate(n_events=100, seed=152) + norm_weights_seeded2, particles_seeded2 = decay.generate( + n_events=100, seed=common_seed + ) + + np.testing.assert_allclose(norm_weights_seeded1, norm_weights_seeded2) + for part1, part2 in zip(particles_seeded1.values(), particles_seeded2.values()): + np.testing.assert_allclose(part1, part2) + + assert not np.allclose(norm_weights_seeded1, norm_weights_rnd) + for part1, part2 in zip(particles_seeded1.values(), particles_rnd.values()): + assert not np.allclose(part1, part2) + + assert not np.allclose(norm_weights_global, norm_weights_rnd) + for part1, part2 in zip(particles_global.values(), particles_rnd.values()): + assert not np.allclose(part1, part2) + + +if __name__ == "__main__": + test_one_event() + test_n_events(5) + + # EOF diff --git a/tests/test_nbody_decay.py b/tests/test_nbody_decay.py index 461beae9..efd11df9 100644 --- a/tests/test_nbody_decay.py +++ b/tests/test_nbody_decay.py @@ -1,64 +1,64 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_nbody_decay.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 14.06.2019 -# ============================================================================= -"""Test n-body decay generator.""" - -import pytest - -from phasespace import nbody_decay - -from .helpers import decays - -B0_MASS = decays.B0_MASS -PION_MASS = decays.PION_MASS - - -def test_no_names(): - """Test particle naming when no name is given.""" - decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) - assert decay.name == "top" - assert all( - part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) - ) - - -def test_top_name(): - """Test particle naming when only top name is given.""" - decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0") - assert decay.name == "B0" - assert all( - part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) - ) - - -def test_children_names(): - """Test particle naming when only children names are given.""" - children_names = [f"pion_{i}" for i in range(3)] - decay = nbody_decay( - B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names - ) - assert decay.name == "top" - assert children_names == [part.name for part in decay.children] - - -def test_all_names(): - """Test particle naming when all names are given.""" - children_names = [f"pion_{i}" for i in range(3)] - decay = nbody_decay( - B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0", names=children_names - ) - assert decay.name == "B0" - assert children_names == [part.name for part in decay.children] - - -def test_mismatching_names(): - """Test wrong number of names given for children.""" - children_names = [f"pion_{i}" for i in range(4)] - with pytest.raises(ValueError): - nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names) - - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_nbody_decay.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 14.06.2019 +# ============================================================================= +"""Test n-body decay generator.""" + +import pytest + +from phasespace import nbody_decay + +from .helpers import decays + +B0_MASS = decays.B0_MASS +PION_MASS = decays.PION_MASS + + +def test_no_names(): + """Test particle naming when no name is given.""" + decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS]) + assert decay.name == "top" + assert all( + part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) + ) + + +def test_top_name(): + """Test particle naming when only top name is given.""" + decay = nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0") + assert decay.name == "B0" + assert all( + part.name == f"p_{part_num}" for part_num, part in enumerate(decay.children) + ) + + +def test_children_names(): + """Test particle naming when only children names are given.""" + children_names = [f"pion_{i}" for i in range(3)] + decay = nbody_decay( + B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names + ) + assert decay.name == "top" + assert children_names == [part.name for part in decay.children] + + +def test_all_names(): + """Test particle naming when all names are given.""" + children_names = [f"pion_{i}" for i in range(3)] + decay = nbody_decay( + B0_MASS, [PION_MASS, PION_MASS, PION_MASS], top_name="B0", names=children_names + ) + assert decay.name == "B0" + assert children_names == [part.name for part in decay.children] + + +def test_mismatching_names(): + """Test wrong number of names given for children.""" + children_names = [f"pion_{i}" for i in range(4)] + with pytest.raises(ValueError): + nbody_decay(B0_MASS, [PION_MASS, PION_MASS, PION_MASS], names=children_names) + + +# EOF diff --git a/tests/test_physics.py b/tests/test_physics.py index 025d9312..4cfa98c9 100644 --- a/tests/test_physics.py +++ b/tests/test_physics.py @@ -1,411 +1,411 @@ -#!/usr/bin/env python3 -# ============================================================================= -# @file test_physics.py -# @author Albert Puig (albert.puig@cern.ch) -# @date 27.02.2019 -# ============================================================================= -"""Test physics output.""" - -import platform -import subprocess - -import numpy as np -import pytest -from scipy.stats import ks_2samp - -if platform.system() == "Darwin": - import matplotlib - - matplotlib.use("TkAgg") - -import os -import sys - -import matplotlib.pyplot as plt -import tensorflow as tf -import uproot4 - -from phasespace import phasespace - -sys.path.append(os.path.dirname(__file__)) - -from .helpers import decays, rapidsim # noqa: E402 -from .helpers.plotting import make_norm_histo # noqa: E402 - -BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -PLOT_DIR = os.path.join(BASE_PATH, "tests", "plots") - - -def setup_method(): - phasespace.GenParticle._sess.close() - tf.compat.v1.reset_default_graph() - - -def create_ref_histos(n_pions): - """Load reference histogram data.""" - ref_dir = os.path.join(BASE_PATH, "data") - if not os.path.exists(ref_dir): - os.mkdir(ref_dir) - ref_file = os.path.join(ref_dir, f"bto{n_pions}pi.root") - if not os.path.exists(ref_file): - script = os.path.join( - BASE_PATH, - "scripts", - "prepare_test_samples.cxx+({})".format( - ",".join( - '"{}"'.format(os.path.join(BASE_PATH, "data", f"bto{i + 1}pi.root")) - for i in range(1, 4) - ) - ), - ) - subprocess.call(f"root -qb '{script}'", shell=True) - events = uproot4.open(ref_file)["events"] - pion_names = [f"pion_{pion + 1}" for pion in range(n_pions)] - pions = {pion_name: events[pion_name] for pion_name in pion_names} - weights = events["weight"] - normalized_histograms = [] - for pion in pions.values(): - pion_array = pion.array() - energy = pion_array.fE - momentum = pion_array.fP - for coord, array in enumerate([momentum.fX, momentum.fY, momentum.fZ, energy]): - numpy_array = np.array(array) - histogram = make_norm_histo( - numpy_array, - range_=(-3000 if coord % 4 != 3 else 0, 3000), - weights=weights, - ) - normalized_histograms.append(histogram) - - return normalized_histograms, make_norm_histo(weights, range_=(0, 1 + 1e-8)) - - -def run_test(n_particles, test_prefix): - first_run_n_events = 100 - main_run_n_events = 100000 - n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64) - - decay = phasespace.nbody_decay(decays.B0_MASS, [decays.PION_MASS] * n_particles) - generate = decay.generate(n_events) - weights1, _ = generate # only generate to test change in n_events - assert len(weights1) == first_run_n_events - - # change n_events and run again - n_events.assign(main_run_n_events) - weights, particles = decay.generate(n_events) - parts = np.concatenate( - [particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1 - ) - histos = [ - make_norm_histo( - parts[:, coord], - range_=(-3000 if coord % 4 != 3 else 0, 3000), - weights=weights, - ) - for coord in range(parts.shape[1]) - ] - weight_histos = make_norm_histo(weights, range_=(0, 1 + 1e-8)) - ref_histos, ref_weights = create_ref_histos(n_particles) - p_values = np.array( - [ - ks_2samp(histos[coord], ref_histos[coord])[1] - for coord, _ in enumerate(histos) - ] - + [ks_2samp(weight_histos, ref_weights)[1]] - ) - # Let's plot - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - for coord, _ in enumerate(histos): - plt.hist( - x if coord % 4 != 3 else e, - weights=histos[coord], - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histos[coord], - alpha=0.5, - label="TGenPhasespace", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "{}_pion_{}_{}.png".format( - test_prefix, int(coord / 4) + 1, ["px", "py", "pz", "e"][coord % 4] - ), - ) - ) - plt.clf() - plt.hist( - np.linspace(0, 1, 100), - weights=weight_histos, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - np.linspace(0, 1, 100), - weights=ref_weights, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"{test_prefix}_weights.png")) - plt.clf() - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_two_body(): - """Test B->pipi decay.""" - run_test(2, "two_body") - - -@pytest.mark.flaky(3) # Stats are limited -def test_three_body(): - """Test B -> pi pi pi decay.""" - run_test(3, "three_body") - - -@pytest.mark.flaky(3) # Stats are limited -def test_four_body(): - """Test B -> pi pi pi pi decay.""" - run_test(4, "four_body") - - -def run_kstargamma(input_file, kstar_width, b_at_rest, suffix): - """Run B0->K*gamma test.""" - n_events = 1000000 - if b_at_rest: - booster = None - rapidsim_getter = rapidsim.get_tree_in_b_rest_frame - else: - booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) - booster = booster.transpose() - rapidsim_getter = rapidsim.get_tree - decay = decays.b0_to_kstar_gamma(kstar_width=kstar_width) - norm_weights, particles = decay.generate(n_events=n_events, boost_to=booster) - rapidsim_parts = rapidsim_getter( - os.path.join(BASE_PATH, "data", input_file), - "B0_0", - ("Kst0_0", "gamma_0", "Kp_0", "pim_0"), - ) - name_matching = {"Kst0_0": "K*0", "gamma_0": "gamma", "Kp_0": "K+", "pim_0": "pi-"} - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - p_values = {} - for ref_name, ref_part in rapidsim_parts.items(): - tf_part = name_matching[ref_name] - ref_part = ref_part.transpose() # for consistency - for coord, coord_name in enumerate(("px", "py", "pz", "e")): - range_ = (-3000 if coord % 4 != 3 else 0, 3000) - ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) - tf_histo = make_norm_histo( - particles[tf_part][:, coord], range_=range_, weights=norm_weights - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=tf_histo, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histo, - alpha=0.5, - label="RapidSim", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "B0_Kstar_gamma_Kstar{}_{}_{}.png".format( - suffix, tf_part.replace("*", "star"), coord_name - ), - ) - ) - plt.clf() - p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] - plt.hist( - np.linspace(0, 1, 100), - weights=make_norm_histo(norm_weights, range_=(0, 1)), - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"B0_Kstar_gamma_Kstar{suffix}_weights.png")) - plt.clf() - return np.array(list(p_values.values())) - - -@pytest.mark.flaky(3) # Stats are limited -def test_kstargamma_kstarnonresonant_at_rest(): - """Test B0 -> K* gamma physics with fixed mass for K*.""" - p_values = run_kstargamma( - "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", 0, True, "NonResonant" - ) - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_kstargamma_kstarnonresonant_lhc(): - """Test B0 -> K* gamma physics with fixed mass for K* with LHC kinematics.""" - p_values = run_kstargamma( - "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", - 0, - False, - "NonResonant_LHC", - ) - assert np.all(p_values > 0.05) - - -def test_kstargamma_resonant_at_rest(): - """Test B0 -> K* gamma physics with Gaussian mass for K*. - - Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov - test wrt to RapidSim, so plots are generated and can be inspected by the user. However, small differences - are expected in the tails of the energy distributions of the kaon and the pion. - """ - run_kstargamma( - "B2KstGamma_RapidSim_7TeV_Tree.root", decays.KSTARZ_WIDTH, True, "Gaussian" - ) - - -def run_k1_gamma(input_file, k1_width, kstar_width, b_at_rest, suffix): - """Run B+ -> K1gamma test.""" - n_events = 1000000 - if b_at_rest: - booster = None - rapidsim_getter = rapidsim.get_tree_in_b_rest_frame - else: - booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) - booster = booster.transpose() - rapidsim_getter = rapidsim.get_tree - gamma = decays.bp_to_k1_kstar_pi_gamma(k1_width=k1_width, kstar_width=kstar_width) - norm_weights, particles = gamma.generate(n_events=n_events, boost_to=booster) - rapidsim_parts = rapidsim_getter( - os.path.join(BASE_PATH, "data", input_file), - "Bp_0", - ("K1_1270_p_0", "Kst0_0", "gamma_0", "Kp_0", "pim_0", "pip_0"), - ) - name_matching = { - "K1_1270_p_0": "K1+", - "Kst0_0": "K*0", - "gamma_0": "gamma", - "Kp_0": "K+", - "pim_0": "pi-", - "pip_0": "pi+", - } - if not os.path.exists(PLOT_DIR): - os.mkdir(PLOT_DIR) - x = np.linspace(-3000, 3000, 100) - e = np.linspace(0, 3000, 100) - p_values = {} - for ref_name, ref_part in rapidsim_parts.items(): - tf_part = name_matching[ref_name] - ref_part = ( - ref_part.transpose() - ) # to be consistent with internal shape (nevents, nobs) - for coord, coord_name in enumerate(("px", "py", "pz", "e")): - range_ = (-3000 if coord % 4 != 3 else 0, 3000) - ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) - tf_histo = make_norm_histo( - particles[tf_part][:, coord], range_=range_, weights=norm_weights - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=tf_histo, - alpha=0.5, - label="phasespace", - bins=100, - ) - plt.hist( - x if coord % 4 != 3 else e, - weights=ref_histo, - alpha=0.5, - label="RapidSim", - bins=100, - ) - plt.legend(loc="upper right") - plt.savefig( - os.path.join( - PLOT_DIR, - "Bp_K1_gamma_K1Kstar{}_{}_{}.png".format( - suffix, tf_part.replace("*", "star"), coord_name - ), - ) - ) - plt.clf() - p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] - plt.hist( - np.linspace(0, 1, 100), - weights=make_norm_histo(norm_weights, range_=(0, 1)), - bins=100, - ) - plt.savefig(os.path.join(PLOT_DIR, f"Bp_K1_gamma_K1Kstar{suffix}_weights.png")) - plt.clf() - return np.array(list(p_values.values())) - - -@pytest.mark.flaky(3) # Stats are limited -def test_k1gamma_kstarnonresonant_at_rest(): - """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances.""" - p_values = run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", - 0, - 0, - True, - "NonResonant", - ) - assert np.all(p_values > 0.05) - - -@pytest.mark.flaky(3) # Stats are limited -def test_k1gamma_kstarnonresonant_lhc(): - """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances with LHC kinematics.""" - p_values = run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", - 0, - 0, - False, - "NonResonant_LHC", - ) - assert np.all(p_values > 0.05) - - -def test_k1gamma_resonant_at_rest(): - """Test B0 -> K1 (->K*pi) gamma physics. - - Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov - test wrt to RapidSim, so plots are generated and can be inspected by the user. - """ - run_k1_gamma( - "B2K1Gamma_RapidSim_7TeV_Tree.root", - decays.K1_WIDTH, - decays.KSTARZ_WIDTH, - True, - "Gaussian", - ) - - -if __name__ == "__main__": - test_two_body() - test_three_body() - test_four_body() - test_kstargamma_kstarnonresonant_at_rest() - test_kstargamma_kstarnonresonant_lhc() - test_kstargamma_resonant_at_rest() - test_k1gamma_kstarnonresonant_at_rest() - test_k1gamma_kstarnonresonant_lhc() - test_k1gamma_resonant_at_rest() - -# EOF +#!/usr/bin/env python3 +# ============================================================================= +# @file test_physics.py +# @author Albert Puig (albert.puig@cern.ch) +# @date 27.02.2019 +# ============================================================================= +"""Test physics output.""" + +import platform +import subprocess + +import numpy as np +import pytest +from scipy.stats import ks_2samp + +if platform.system() == "Darwin": + import matplotlib + + matplotlib.use("TkAgg") + +import os +import sys + +import matplotlib.pyplot as plt +import tensorflow as tf +import uproot4 + +from phasespace import phasespace + +sys.path.append(os.path.dirname(__file__)) + +from .helpers import decays, rapidsim # noqa: E402 +from .helpers.plotting import make_norm_histo # noqa: E402 + +BASE_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +PLOT_DIR = os.path.join(BASE_PATH, "tests", "plots") + + +def setup_method(): + phasespace.GenParticle._sess.close() + tf.compat.v1.reset_default_graph() + + +def create_ref_histos(n_pions): + """Load reference histogram data.""" + ref_dir = os.path.join(BASE_PATH, "data") + if not os.path.exists(ref_dir): + os.mkdir(ref_dir) + ref_file = os.path.join(ref_dir, f"bto{n_pions}pi.root") + if not os.path.exists(ref_file): + script = os.path.join( + BASE_PATH, + "scripts", + "prepare_test_samples.cxx+({})".format( + ",".join( + '"{}"'.format(os.path.join(BASE_PATH, "data", f"bto{i + 1}pi.root")) + for i in range(1, 4) + ) + ), + ) + subprocess.call(f"root -qb '{script}'", shell=True) + events = uproot4.open(ref_file)["events"] + pion_names = [f"pion_{pion + 1}" for pion in range(n_pions)] + pions = {pion_name: events[pion_name] for pion_name in pion_names} + weights = events["weight"] + normalized_histograms = [] + for pion in pions.values(): + pion_array = pion.array() + energy = pion_array.fE + momentum = pion_array.fP + for coord, array in enumerate([momentum.fX, momentum.fY, momentum.fZ, energy]): + numpy_array = np.array(array) + histogram = make_norm_histo( + numpy_array, + range_=(-3000 if coord % 4 != 3 else 0, 3000), + weights=weights, + ) + normalized_histograms.append(histogram) + + return normalized_histograms, make_norm_histo(weights, range_=(0, 1 + 1e-8)) + + +def run_test(n_particles, test_prefix): + first_run_n_events = 100 + main_run_n_events = 100000 + n_events = tf.Variable(initial_value=first_run_n_events, dtype=tf.int64) + + decay = phasespace.nbody_decay(decays.B0_MASS, [decays.PION_MASS] * n_particles) + generate = decay.generate(n_events) + weights1, _ = generate # only generate to test change in n_events + assert len(weights1) == first_run_n_events + + # change n_events and run again + n_events.assign(main_run_n_events) + weights, particles = decay.generate(n_events) + parts = np.concatenate( + [particles[f"p_{part_num}"] for part_num in range(n_particles)], axis=1 + ) + histos = [ + make_norm_histo( + parts[:, coord], + range_=(-3000 if coord % 4 != 3 else 0, 3000), + weights=weights, + ) + for coord in range(parts.shape[1]) + ] + weight_histos = make_norm_histo(weights, range_=(0, 1 + 1e-8)) + ref_histos, ref_weights = create_ref_histos(n_particles) + p_values = np.array( + [ + ks_2samp(histos[coord], ref_histos[coord])[1] + for coord, _ in enumerate(histos) + ] + + [ks_2samp(weight_histos, ref_weights)[1]] + ) + # Let's plot + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + for coord, _ in enumerate(histos): + plt.hist( + x if coord % 4 != 3 else e, + weights=histos[coord], + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histos[coord], + alpha=0.5, + label="TGenPhasespace", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "{}_pion_{}_{}.png".format( + test_prefix, int(coord / 4) + 1, ["px", "py", "pz", "e"][coord % 4] + ), + ) + ) + plt.clf() + plt.hist( + np.linspace(0, 1, 100), + weights=weight_histos, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + np.linspace(0, 1, 100), + weights=ref_weights, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"{test_prefix}_weights.png")) + plt.clf() + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_two_body(): + """Test B->pipi decay.""" + run_test(2, "two_body") + + +@pytest.mark.flaky(3) # Stats are limited +def test_three_body(): + """Test B -> pi pi pi decay.""" + run_test(3, "three_body") + + +@pytest.mark.flaky(3) # Stats are limited +def test_four_body(): + """Test B -> pi pi pi pi decay.""" + run_test(4, "four_body") + + +def run_kstargamma(input_file, kstar_width, b_at_rest, suffix): + """Run B0->K*gamma test.""" + n_events = 1000000 + if b_at_rest: + booster = None + rapidsim_getter = rapidsim.get_tree_in_b_rest_frame + else: + booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) + booster = booster.transpose() + rapidsim_getter = rapidsim.get_tree + decay = decays.b0_to_kstar_gamma(kstar_width=kstar_width) + norm_weights, particles = decay.generate(n_events=n_events, boost_to=booster) + rapidsim_parts = rapidsim_getter( + os.path.join(BASE_PATH, "data", input_file), + "B0_0", + ("Kst0_0", "gamma_0", "Kp_0", "pim_0"), + ) + name_matching = {"Kst0_0": "K*0", "gamma_0": "gamma", "Kp_0": "K+", "pim_0": "pi-"} + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + p_values = {} + for ref_name, ref_part in rapidsim_parts.items(): + tf_part = name_matching[ref_name] + ref_part = ref_part.transpose() # for consistency + for coord, coord_name in enumerate(("px", "py", "pz", "e")): + range_ = (-3000 if coord % 4 != 3 else 0, 3000) + ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) + tf_histo = make_norm_histo( + particles[tf_part][:, coord], range_=range_, weights=norm_weights + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=tf_histo, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histo, + alpha=0.5, + label="RapidSim", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "B0_Kstar_gamma_Kstar{}_{}_{}.png".format( + suffix, tf_part.replace("*", "star"), coord_name + ), + ) + ) + plt.clf() + p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] + plt.hist( + np.linspace(0, 1, 100), + weights=make_norm_histo(norm_weights, range_=(0, 1)), + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"B0_Kstar_gamma_Kstar{suffix}_weights.png")) + plt.clf() + return np.array(list(p_values.values())) + + +@pytest.mark.flaky(3) # Stats are limited +def test_kstargamma_kstarnonresonant_at_rest(): + """Test B0 -> K* gamma physics with fixed mass for K*.""" + p_values = run_kstargamma( + "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", 0, True, "NonResonant" + ) + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_kstargamma_kstarnonresonant_lhc(): + """Test B0 -> K* gamma physics with fixed mass for K* with LHC kinematics.""" + p_values = run_kstargamma( + "B2KstGamma_RapidSim_7TeV_KstarNonResonant_Tree.root", + 0, + False, + "NonResonant_LHC", + ) + assert np.all(p_values > 0.05) + + +def test_kstargamma_resonant_at_rest(): + """Test B0 -> K* gamma physics with Gaussian mass for K*. + + Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov + test wrt to RapidSim, so plots are generated and can be inspected by the user. However, small differences + are expected in the tails of the energy distributions of the kaon and the pion. + """ + run_kstargamma( + "B2KstGamma_RapidSim_7TeV_Tree.root", decays.KSTARZ_WIDTH, True, "Gaussian" + ) + + +def run_k1_gamma(input_file, k1_width, kstar_width, b_at_rest, suffix): + """Run B+ -> K1gamma test.""" + n_events = 1000000 + if b_at_rest: + booster = None + rapidsim_getter = rapidsim.get_tree_in_b_rest_frame + else: + booster = rapidsim.generate_fonll(decays.B0_MASS, 7, "b", n_events) + booster = booster.transpose() + rapidsim_getter = rapidsim.get_tree + gamma = decays.bp_to_k1_kstar_pi_gamma(k1_width=k1_width, kstar_width=kstar_width) + norm_weights, particles = gamma.generate(n_events=n_events, boost_to=booster) + rapidsim_parts = rapidsim_getter( + os.path.join(BASE_PATH, "data", input_file), + "Bp_0", + ("K1_1270_p_0", "Kst0_0", "gamma_0", "Kp_0", "pim_0", "pip_0"), + ) + name_matching = { + "K1_1270_p_0": "K1+", + "Kst0_0": "K*0", + "gamma_0": "gamma", + "Kp_0": "K+", + "pim_0": "pi-", + "pip_0": "pi+", + } + if not os.path.exists(PLOT_DIR): + os.mkdir(PLOT_DIR) + x = np.linspace(-3000, 3000, 100) + e = np.linspace(0, 3000, 100) + p_values = {} + for ref_name, ref_part in rapidsim_parts.items(): + tf_part = name_matching[ref_name] + ref_part = ( + ref_part.transpose() + ) # to be consistent with internal shape (nevents, nobs) + for coord, coord_name in enumerate(("px", "py", "pz", "e")): + range_ = (-3000 if coord % 4 != 3 else 0, 3000) + ref_histo = make_norm_histo(ref_part[:, coord], range_=range_) + tf_histo = make_norm_histo( + particles[tf_part][:, coord], range_=range_, weights=norm_weights + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=tf_histo, + alpha=0.5, + label="phasespace", + bins=100, + ) + plt.hist( + x if coord % 4 != 3 else e, + weights=ref_histo, + alpha=0.5, + label="RapidSim", + bins=100, + ) + plt.legend(loc="upper right") + plt.savefig( + os.path.join( + PLOT_DIR, + "Bp_K1_gamma_K1Kstar{}_{}_{}.png".format( + suffix, tf_part.replace("*", "star"), coord_name + ), + ) + ) + plt.clf() + p_values[(tf_part, coord_name)] = ks_2samp(tf_histo, ref_histo)[1] + plt.hist( + np.linspace(0, 1, 100), + weights=make_norm_histo(norm_weights, range_=(0, 1)), + bins=100, + ) + plt.savefig(os.path.join(PLOT_DIR, f"Bp_K1_gamma_K1Kstar{suffix}_weights.png")) + plt.clf() + return np.array(list(p_values.values())) + + +@pytest.mark.flaky(3) # Stats are limited +def test_k1gamma_kstarnonresonant_at_rest(): + """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances.""" + p_values = run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", + 0, + 0, + True, + "NonResonant", + ) + assert np.all(p_values > 0.05) + + +@pytest.mark.flaky(3) # Stats are limited +def test_k1gamma_kstarnonresonant_lhc(): + """Test B0 -> K1 (->K*pi) gamma physics with fixed-mass resonances with LHC kinematics.""" + p_values = run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_K1KstarNonResonant_Tree.root", + 0, + 0, + False, + "NonResonant_LHC", + ) + assert np.all(p_values > 0.05) + + +def test_k1gamma_resonant_at_rest(): + """Test B0 -> K1 (->K*pi) gamma physics. + + Since we don't have BW and we model the resonances with Gaussians, we can't really perform the Kolmogorov + test wrt to RapidSim, so plots are generated and can be inspected by the user. + """ + run_k1_gamma( + "B2K1Gamma_RapidSim_7TeV_Tree.root", + decays.K1_WIDTH, + decays.KSTARZ_WIDTH, + True, + "Gaussian", + ) + + +if __name__ == "__main__": + test_two_body() + test_three_body() + test_four_body() + test_kstargamma_kstarnonresonant_at_rest() + test_kstargamma_kstarnonresonant_lhc() + test_kstargamma_resonant_at_rest() + test_k1gamma_kstarnonresonant_at_rest() + test_k1gamma_kstarnonresonant_lhc() + test_k1gamma_resonant_at_rest() + +# EOF diff --git a/tests/test_random.py b/tests/test_random.py index 5581bbb8..a0148d29 100644 --- a/tests/test_random.py +++ b/tests/test_random.py @@ -1,27 +1,27 @@ -import numpy as np -import pytest -import tensorflow as tf - -import phasespace as phsp - - -@pytest.mark.parametrize( - "seed", [lambda: 15, lambda: tf.random.Generator.from_seed(15)] -) -def test_get_rng(seed): - rng1 = phsp.random.get_rng(seed()) - rng2 = phsp.random.get_rng(seed()) - rnd1_seeded = rng1.uniform_full_int(shape=(100,)) - rnd2_seeded = rng2.uniform_full_int(shape=(100,)) - - rng3 = phsp.random.get_rng() - rng4 = phsp.random.get_rng(seed()) - # advance rng4 by one step - _ = rng4.split(1) - - rnd3 = rng3.uniform_full_int(shape=(100,)) - rnd4 = rng4.uniform_full_int(shape=(100,)) - - np.testing.assert_array_equal(rnd1_seeded, rnd2_seeded) - assert not np.array_equal(rnd1_seeded, rnd3) - assert not np.array_equal(rnd4, rnd3) +import numpy as np +import pytest +import tensorflow as tf + +import phasespace as phsp + + +@pytest.mark.parametrize( + "seed", [lambda: 15, lambda: tf.random.Generator.from_seed(15)] +) +def test_get_rng(seed): + rng1 = phsp.random.get_rng(seed()) + rng2 = phsp.random.get_rng(seed()) + rnd1_seeded = rng1.uniform_full_int(shape=(100,)) + rnd2_seeded = rng2.uniform_full_int(shape=(100,)) + + rng3 = phsp.random.get_rng() + rng4 = phsp.random.get_rng(seed()) + # advance rng4 by one step + _ = rng4.split(1) + + rnd3 = rng3.uniform_full_int(shape=(100,)) + rnd4 = rng4.uniform_full_int(shape=(100,)) + + np.testing.assert_array_equal(rnd1_seeded, rnd2_seeded) + assert not np.array_equal(rnd1_seeded, rnd3) + assert not np.array_equal(rnd4, rnd3) From 32cc7819c2c2438a1b9dd2eaf768ba880283961e Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Tue, 12 Oct 2021 12:35:02 +0200 Subject: [PATCH 34/71] Rename fulldecay to fromdecay and FullDecay to MultiDecayChain --- docs/fromdecay.ipynb | 716 +++++++++--------- phasespace/fromdecay/__init__.py | 14 + .../mass_functions.py | 128 ++-- .../multidecaychain.py} | 30 +- phasespace/fulldecay/__init__.py | 14 - tests/{fulldecay => fromdecay}/__init__.py | 2 +- .../example_decay_chains.py | 2 +- .../example_decays.dec | 0 .../test_fromdecay.py} | 26 +- .../test_mass_functions.py | 108 +-- 10 files changed, 523 insertions(+), 517 deletions(-) create mode 100644 phasespace/fromdecay/__init__.py rename phasespace/{fulldecay => fromdecay}/mass_functions.py (97%) rename phasespace/{fulldecay/fulldecay.py => fromdecay/multidecaychain.py} (85%) delete mode 100644 phasespace/fulldecay/__init__.py rename tests/{fulldecay => fromdecay}/__init__.py (65%) rename tests/{fulldecay => fromdecay}/example_decay_chains.py (87%) rename tests/{fulldecay => fromdecay}/example_decays.dec (100%) rename tests/{fulldecay/test_fulldecay.py => fromdecay/test_fromdecay.py} (79%) rename tests/{fulldecay => fromdecay}/test_mass_functions.py (94%) diff --git a/docs/fromdecay.ipynb b/docs/fromdecay.ipynb index ee248d7d..430d05f1 100644 --- a/docs/fromdecay.ipynb +++ b/docs/fromdecay.ipynb @@ -1,358 +1,358 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "collapsed": true, - "pycharm": { - "is_executing": true, - "name": "#%% md\n" - } - }, - "source": [ - "# Tutorial for `fromdecay` functionality\n", - "This tutorial shows how `phasespace.fromdecay` can be used.\n", - "\n", - "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", - "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "# Import libraries\n", - "from pprint import pprint\n", - "\n", - "import zfit\n", - "from particle import Particle\n", - "from decaylanguage import DecFileParser, DecayChainViewer\n", - "# TODO rename\n", - "from phasespace.fulldecay import FullDecay" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Intro to DecayLanguage\n", - "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", - "\n", - "We will begin by parsing a .dec file using DecayLanguage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "parser = DecFileParser('../tests/fulldecay/example_decays.dec')\n", - "parser.parse()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_chain = parser.build_decay_chains(\"pi0\")\n", - "pprint(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DecayChainViewer(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Creating a FullDecay object\n", - "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `FullDecay` object can be created directly from a DecayLanguage dict:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_decay = FullDecay.from_dict(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When creating a `FullDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", - "\n", - "These GenParticle instances and the probabilities of that decay mode can be accessed via `FullDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for probability, particle in pi0_decay.gen_particles:\n", - " print(f\"There is a probability of {probability} \"\n", - " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", - "\n", - "When calling the `FullDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `FullDecay.gen_particles`. The outputs are placed in a list, which is returned." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "weights, events = pi0_decay.generate(n_events=10_000)\n", - "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can confirm that the counts above are close to the expected counts based on the probabilities. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Changing mass settings FullDecay\n", - "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", - "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", - "These settings can be modified by passing in additional parameters to `FullDecay.from_dict`.\n", - "There are two optional parameters that can be passed to `FullDecay.from_dict`: `tolerance` and `mass_converter`.\n", - "\n", - "### Constant vs variable mass\n", - "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", - "This will be demonsttrated with the decay below:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", - "DecayChainViewer(dsplus_chain)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", - " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$\\pi^0$ has a greater width than $D^0$. \n", - "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dstar_decay = FullDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", - "# Loop over D0 and pi+ particles, see graph above\n", - "for particle in dstar_decay.gen_particles[0][1].children:\n", - " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", - " assert particle.has_fixed_mass\n", - " \n", - "# Loop over D+ and pi0. See above.\n", - "for particle in dstar_decay.gen_particles[1][1].children:\n", - " if particle.name == \"pi0\":\n", - " assert not particle.has_fixed_mass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Configuring mass fucntions\n", - "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_custom_mass_func = dsplus_chain.copy()\n", - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", - "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `FullDecay.from_dict`, like before." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def custom_gauss(mass, width):\n", - " particle_mass = tf.cast(mass, tf.float64)\n", - " particle_width = tf.cast(width, tf.float64)\n", - " \n", - " # This is the actual mass function that will be returned\n", - " def mass_func(min_mass, max_mass, n_events):\n", - " min_mass = tf.cast(min_mass, tf.float64)\n", - " max_mass = tf.cast(max_mass, tf.float64)\n", - " # Use a zfit PDF\n", - " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", - " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", - " return tf.vectorized_map(\n", - " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", - " )\n", - "\n", - " return mass_func" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This function can then be passed to `FullDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "\n", - "# Set the mass function of pi0 to the custom gaussian distribution \n", - "# when it decays into an electron-positron pair and a photon (gamma)\n", - "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FullDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.7" - } - }, - "nbformat": 4, - "nbformat_minor": 1 -} +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": true, + "pycharm": { + "is_executing": true, + "name": "#%% md\n" + } + }, + "source": [ + "# Tutorial for `fromdecay` functionality\n", + "This tutorial shows how `phasespace.fromdecay` can be used.\n", + "\n", + "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", + "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Import libraries\n", + "from pprint import pprint\n", + "\n", + "import zfit\n", + "from particle import Particle\n", + "from decaylanguage import DecFileParser, DecayChainViewer\n", + "\n", + "from phasespace.fromdecay import MultiDecayChain" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Intro to DecayLanguage\n", + "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", + "\n", + "We will begin by parsing a .dec file using DecayLanguage:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "parser = DecFileParser('../tests/fromdecay/example_decays.dec')\n", + "parser.parse()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_chain = parser.build_decay_chains(\"pi0\")\n", + "pprint(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DecayChainViewer(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating a MultiDecayChain object\n", + "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `MultiDecayChain` object can be created directly from a DecayLanguage dict:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_decay = MultiDecayChain.from_dict(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When creating a `MultiDecayChain` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `MultiDecayChain.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for probability, particle in pi0_decay.gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", + "\n", + "When calling the `MultiDecayChain.generate` method, it internally calls the generate method on the of the GenParticle instances in `MultiDecayChain.gen_particles`. The outputs are placed in a list, which is returned." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "weights, events = pi0_decay.generate(n_events=10_000)\n", + "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can confirm that the counts above are close to the expected counts based on the probabilities. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Changing mass settings MultiDecayChain\n", + "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", + "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", + "These settings can be modified by passing in additional parameters to `MultiDecayChain.from_dict`.\n", + "There are two optional parameters that can be passed to `MultiDecayChain.from_dict`: `tolerance` and `mass_converter`.\n", + "\n", + "### Constant vs variable mass\n", + "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", + "This will be demonsttrated with the decay below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", + "DecayChainViewer(dsplus_chain)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", + " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$\\pi^0$ has a greater width than $D^0$. \n", + "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dstar_decay = MultiDecayChain.from_dict(dsplus_chain, tolerance=1e-8)\n", + "# Loop over D0 and pi+ particles, see graph above\n", + "for particle in dstar_decay.gen_particles[0][1].children:\n", + " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", + " assert particle.has_fixed_mass\n", + " \n", + "# Loop over D+ and pi0. See above.\n", + "for particle in dstar_decay.gen_particles[1][1].children:\n", + " if particle.name == \"pi0\":\n", + " assert not particle.has_fixed_mass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Configuring mass fucntions\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_custom_mass_func = dsplus_chain.copy()\n", + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `MultiDecayChain.from_dict`, like before." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "MultiDecayChain.from_dict(dsplus_custom_mass_func)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def custom_gauss(mass, width):\n", + " particle_mass = tf.cast(mass, tf.float64)\n", + " particle_width = tf.cast(width, tf.float64)\n", + " \n", + " # This is the actual mass function that will be returned\n", + " def mass_func(min_mass, max_mass, n_events):\n", + " min_mass = tf.cast(min_mass, tf.float64)\n", + " max_mass = tf.cast(max_mass, tf.float64)\n", + " # Use a zfit PDF\n", + " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", + " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", + " return tf.vectorized_map(\n", + " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", + " )\n", + "\n", + " return mass_func" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This function can then be passed to `MultiDecayChain.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "\n", + "# Set the mass function of pi0 to the custom gaussian distribution \n", + "# when it decays into an electron-positron pair and a photon (gamma)\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "MultiDecayChain.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py new file mode 100644 index 00000000..a93c7240 --- /dev/null +++ b/phasespace/fromdecay/__init__.py @@ -0,0 +1,14 @@ +import sys + +from .multidecaychain import MultiDecayChain # noqa: F401 + +try: + import zfit # noqa: F401 + import zfit_physics as zphys # noqa: F401 + from particle import Particle # noqa: F401 +except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "The fromdecay functionality in phasespace requires particle and zfit-physics. " + "Either install phasespace[fromdecay] or particle and zfit-physics.", + file=sys.stderr, + ) from error diff --git a/phasespace/fulldecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py similarity index 97% rename from phasespace/fulldecay/mass_functions.py rename to phasespace/fromdecay/mass_functions.py index 80738fb9..23abd45a 100644 --- a/phasespace/fulldecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -1,64 +1,64 @@ -import tensorflow as tf -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) - - def mass_func(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="") - iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator - ) - - return mass_func - - -def breitwigner(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): - 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="") - iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator - ) - - return mass_func - - -def relativistic_breitwigner(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): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - pdf = zphys.pdf.RelativisticBreitWigner( - m=particle_mass, gamma=particle_width, obs="" - ) - iterator = tf.stack([min_mass, max_mass], axis=-1) - - # this works with map_fn but not with vectorized_map as no analytic sampling is available. - return tf.map_fn( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator - ) - - return mass_func - - -_DEFAULT_CONVERTER = { - "gauss": gauss, - "bw": breitwigner, - "relbw": relativistic_breitwigner, -} +import tensorflow as tf +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) + + def mass_func(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="") + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + + return mass_func + + +def breitwigner(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): + 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="") + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + + return mass_func + + +def relativistic_breitwigner(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): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + pdf = zphys.pdf.RelativisticBreitWigner( + m=particle_mass, gamma=particle_width, obs="" + ) + iterator = tf.stack([min_mass, max_mass], axis=-1) + + # this works with map_fn but not with vectorized_map as no analytic sampling is available. + return tf.map_fn( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator + ) + + return mass_func + + +_DEFAULT_CONVERTER = { + "gauss": gauss, + "bw": breitwigner, + "relbw": relativistic_breitwigner, +} diff --git a/phasespace/fulldecay/fulldecay.py b/phasespace/fromdecay/multidecaychain.py similarity index 85% rename from phasespace/fulldecay/fulldecay.py rename to phasespace/fromdecay/multidecaychain.py index 76e755af..0f58781a 100644 --- a/phasespace/fulldecay/fulldecay.py +++ b/phasespace/fromdecay/multidecaychain.py @@ -13,7 +13,7 @@ _DEFAULT_MASS_FUNC = "relbw" -class FullDecay: +class MultiDecayChain: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): """A container that works like GenParticle that can handle multiple decays. Can be created from. @@ -33,23 +33,24 @@ def from_dict( mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE, ): - """Create a FullDecay instance from a dict in the decaylanguage format. + """Create a MultiDecayChain instance from a dict in the decaylanguage format. Args: - dec_dict: The input dict from which the FullDecay object will be created from. + dec_dict: The input dict from which the MultiDecayChain object will be created from. mass_converter: A dict with mass function names and their corresponding mass functions. These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. - tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. + tolerance: Minimum mass width of the particle to use a mass function instead of + assuming the mass to be constant. Returns: - The created FullDecay object. + The created MultiDecayChain object. """ if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER else: - # Combine the mass functions specified by the package to the mass functions specified from the input. + # Combine the default mass functions specified with the mass functions from input. total_mass_converter = {**_DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( @@ -67,15 +68,16 @@ def generate( Args: n_events: Total number of events combined, for all the decays. - normalize_weights: Normalize weights according to all events generated. This also changes the return values. - See the phasespace documentation for more details. + normalize_weights: Normalize weights according to all events generated. + This also changes the return values. See the phasespace documentation for more details. kwargs: Additional parameters passed to all calls of GenParticle.generate Returns: - The arguments returned by GenParticle.generate are returned. See the phasespace documentation for details. - However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, each entry in the lists corresponding - to the return arguments from the corresponding GenParticle instances in self.gen_particles. - Note that when normalize_weights is True, the weights are normalized to the maximum of all returned events. + The arguments returned by GenParticle.generate are returned. See the phasespace documentation for + details. However, instead of being 2 or 3 tensors, it is 2 or 3 lists of tensors, + each entry in the lists corresponding to the return arguments from the corresponding GenParticle + instances in self.gen_particles. Note that when normalize_weights is True, + the weights are normalized to the maximum of all returned events. """ # Input to tf.random.categorical must be 2D rand_i = tf.random.categorical( @@ -109,8 +111,8 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: preexisting_particles: Names that the particle cannot have as name. Returns: - name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i begins at 0 - and increases until the name is not preexisting_particles. + name: Will be `name` if `name` is not in preexisting_particles or of the format "name [i]" where i + begins at 0 and increases until the name is not preexisting_particles. """ if name not in preexisting_particles: preexisting_particles.add(name) diff --git a/phasespace/fulldecay/__init__.py b/phasespace/fulldecay/__init__.py deleted file mode 100644 index 42f877ee..00000000 --- a/phasespace/fulldecay/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -import sys - -from .fulldecay import FullDecay - -try: - import zfit - import zfit_physics as zphys - from particle import Particle -except ModuleNotFoundError as error: - raise ModuleNotFoundError( - "The fulldecay functionality in phasespace requires particle and zfit-physics. " - "Either install phasespace[fulldecay] or particle and zfit-physics.", - file=sys.stderr, - ) from error diff --git a/tests/fulldecay/__init__.py b/tests/fromdecay/__init__.py similarity index 65% rename from tests/fulldecay/__init__.py rename to tests/fromdecay/__init__.py index b98e9386..0234b711 100644 --- a/tests/fulldecay/__init__.py +++ b/tests/fromdecay/__init__.py @@ -1,4 +1,4 @@ import pytest # This makes it so that assert errors are more helpful for e.g., the check_norm helper function -pytest.register_assert_rewrite("fulldecay.test_fulldecay") +pytest.register_assert_rewrite("fromdecay.test_fulldecay") diff --git a/tests/fulldecay/example_decay_chains.py b/tests/fromdecay/example_decay_chains.py similarity index 87% rename from tests/fulldecay/example_decay_chains.py rename to tests/fromdecay/example_decay_chains.py index ef71fda6..c04d867d 100644 --- a/tests/fulldecay/example_decay_chains.py +++ b/tests/fromdecay/example_decay_chains.py @@ -25,5 +25,5 @@ ): decay_mode["zfit"] = mass_function -# D*+ particle that has multiple child particles, grandchild particles, many of which can decay in multiple ways. +# 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/fulldecay/example_decays.dec b/tests/fromdecay/example_decays.dec similarity index 100% rename from tests/fulldecay/example_decays.dec rename to tests/fromdecay/example_decays.dec diff --git a/tests/fulldecay/test_fulldecay.py b/tests/fromdecay/test_fromdecay.py similarity index 79% rename from tests/fulldecay/test_fulldecay.py rename to tests/fromdecay/test_fromdecay.py index 4baa7568..8751f3ac 100644 --- a/tests/fulldecay/test_fulldecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,13 +1,13 @@ from numpy.testing import assert_almost_equal -from phasespace.fulldecay import FullDecay -from phasespace.fulldecay.mass_functions import _DEFAULT_CONVERTER +from phasespace.fromdecay import MultiDecayChain +from phasespace.fromdecay.mass_functions import _DEFAULT_CONVERTER -from .example_decay_chains import * # TODO remove * since it is bad practice? +from . import example_decay_chains -def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: - """Helper function that checks whether the normalize_weights argument works for FullDecay.generate. +def check_norm(full_decay: MultiDecayChain, **kwargs) -> list[tuple]: + """Helper function that checks whether the normalize_weights argument works for MultiDecayChain.generate. Args: full_decay: full_decay.generate will be called. kwargs: Additional parameters passed to generate. @@ -33,7 +33,9 @@ def check_norm(full_decay: FullDecay, **kwargs) -> list[tuple]: def test_single_chain(): """Test converting a decaylanguage dict with only one possible decay.""" - container = FullDecay.from_dict(dplus_single, tolerance=1e-10) + container = MultiDecayChain.from_dict( + example_decay_chains.dplus_single, tolerance=1e-10 + ) output_decay = container.gen_particles assert len(output_decay) == 1 prob, gen = output_decay[0] @@ -56,7 +58,9 @@ def test_single_chain(): def test_branching_children(): """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" - container = FullDecay.from_dict(pi0_4branches, tolerance=1e-10) + container = MultiDecayChain.from_dict( + example_decay_chains.pi0_4branches, 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) @@ -66,7 +70,7 @@ 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 = FullDecay.from_dict(dplus_4grandbranches) + container = MultiDecayChain.from_dict(example_decay_chains.dplus_4grandbranches) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) @@ -77,9 +81,9 @@ def test_branching_grandchilden(): def test_mass_converter(): """Test that the mass_converter parameter works as intended.""" - dplus_4grandbranches_massfunc = dplus_4grandbranches.copy() + dplus_4grandbranches_massfunc = example_decay_chains.dplus_4grandbranches.copy() dplus_4grandbranches_massfunc["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" - container = FullDecay.from_dict( + container = MultiDecayChain.from_dict( dplus_4grandbranches_massfunc, tolerance=1e-10, mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}, @@ -98,7 +102,7 @@ def test_mass_converter(): def test_big_decay(): - container = FullDecay.from_dict(dstarplus_big_decay) + container = MultiDecayChain.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) diff --git a/tests/fulldecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py similarity index 94% rename from tests/fulldecay/test_mass_functions.py rename to tests/fromdecay/test_mass_functions.py index a7c12d3e..43e75b48 100644 --- a/tests/fulldecay/test_mass_functions.py +++ b/tests/fromdecay/test_mass_functions.py @@ -1,54 +1,54 @@ -from typing import Callable - -import pytest -import tensorflow as tf -import tensorflow_probability as tfp -from particle import Particle - -import phasespace.fulldecay.mass_functions as mf - -_kstarz = Particle.from_evtgen_name("K*0") -KSTARZ_MASS = _kstarz.mass -KSTARZ_WIDTH = _kstarz.width - - -def ref_mass_func(min_mass, max_mass, n_events): - """Reference mass function used to compare the behavior of the actual mass functions. - - Args: - min_mass: lower limit of mass. - max_mass: upper limit of mass. - n_events: number of mass values that should be generated. - - Returns: - kstar_mass: Generated mass. - - Notes: - Code taken from phasespace documentation. - """ - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - kstar_mass = tfp.distributions.TruncatedNormal( - loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass - ).sample() - return kstar_mass - - -@pytest.mark.parametrize( - "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) -) -@pytest.mark.parametrize("size", (1, 10)) -def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): - rng = tf.random.Generator.from_seed(1234) - min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) - min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) - assert tf.reduce_all(min_mass <= max_mass) - ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) - sample = function(*params)(min_mass, max_mass, len(min_mass)) - assert sample.shape[0] == ref_sample.shape[0] - assert all( - i <= 1 for i in sample.shape[1:] - ) # Sample.shape have extra dimensions with just 1 or 0, which can be ignored +from typing import Callable + +import pytest +import tensorflow as tf +import tensorflow_probability as tfp +from particle import Particle + +import phasespace.fromdecay.mass_functions as mf + +_kstarz = Particle.from_evtgen_name("K*0") +KSTARZ_MASS = _kstarz.mass +KSTARZ_WIDTH = _kstarz.width + + +def ref_mass_func(min_mass, max_mass, n_events): + """Reference mass function used to compare the behavior of the actual mass functions. + + Args: + min_mass: lower limit of mass. + max_mass: upper limit of mass. + n_events: number of mass values that should be generated. + + Returns: + kstar_mass: Generated mass. + + Notes: + Code taken from phasespace documentation. + """ + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + kstar_mass = tfp.distributions.TruncatedNormal( + loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass + ).sample() + return kstar_mass + + +@pytest.mark.parametrize( + "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) +) +@pytest.mark.parametrize("size", (1, 10)) +def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): + rng = tf.random.Generator.from_seed(1234) + min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) + min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) + assert tf.reduce_all(min_mass <= max_mass) + ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) + sample = function(*params)(min_mass, max_mass, len(min_mass)) + assert sample.shape[0] == ref_sample.shape[0] + assert all( + i <= 1 for i in sample.shape[1:] + ) # Sample.shape have extra dimensions with just 1 or 0, which can be ignored From f5da87c67ee150347fe7b3064b02d48733b0cfc1 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 13 Oct 2021 20:15:29 +0200 Subject: [PATCH 35/71] Update setup.cfg to work with fromdecay --- setup.cfg | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/setup.cfg b/setup.cfg index 3a5d7a3b..091fad12 100644 --- a/setup.cfg +++ b/setup.cfg @@ -39,6 +39,10 @@ testpaths = tests zip_safe = False [options.extras_require] +fromdecay = + particle + zfit + zfit-physics >= 0.2 test = awkward coverage @@ -51,6 +55,7 @@ test = scipy uproot4 wget + decaylanguage doc = Sphinx sphinx_bootstrap_theme From 75c4a87bc964d254385e7e2e94d31f6972c49d8e Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Wed, 3 Nov 2021 22:58:33 +0100 Subject: [PATCH 36/71] Update particle package version requirement Co-authored-by: Eduardo Rodrigues --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 091fad12..c10769f7 100644 --- a/setup.cfg +++ b/setup.cfg @@ -40,7 +40,7 @@ zip_safe = False [options.extras_require] fromdecay = - particle + particle >= 0.16.0 zfit zfit-physics >= 0.2 test = From 128a577cf3391fdea908b75ac6c4ce27ad15b562 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Wed, 3 Nov 2021 23:02:47 +0100 Subject: [PATCH 37/71] Add decaylanguage to package requirements --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index c10769f7..79988259 100644 --- a/setup.cfg +++ b/setup.cfg @@ -43,6 +43,7 @@ fromdecay = particle >= 0.16.0 zfit zfit-physics >= 0.2 + decaylanguage >= 0.12.0 # not required but everyone using this feature will likely use decaylanguage test = awkward coverage From 58f1215e82b0c391dcf43d0d734c8c3c47ed65ec Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 6 Nov 2021 17:50:33 +0100 Subject: [PATCH 38/71] Rename class to GenMultiDecay --- ...cay.ipynb => GenMultiDecay_Tutorial.ipynb} | 716 +++++++++--------- phasespace/fromdecay/__init__.py | 2 +- .../{multidecaychain.py => genmultidecay.py} | 14 +- phasespace/fromdecay/mass_functions.py | 128 ++-- tests/fromdecay/test_fromdecay.py | 16 +- tests/fromdecay/test_mass_functions.py | 108 +-- 6 files changed, 495 insertions(+), 489 deletions(-) rename docs/{fromdecay.ipynb => GenMultiDecay_Tutorial.ipynb} (82%) rename phasespace/fromdecay/{multidecaychain.py => genmultidecay.py} (92%) diff --git a/docs/fromdecay.ipynb b/docs/GenMultiDecay_Tutorial.ipynb similarity index 82% rename from docs/fromdecay.ipynb rename to docs/GenMultiDecay_Tutorial.ipynb index 430d05f1..cc289df2 100644 --- a/docs/fromdecay.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -1,358 +1,358 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "collapsed": true, - "pycharm": { - "is_executing": true, - "name": "#%% md\n" - } - }, - "source": [ - "# Tutorial for `fromdecay` functionality\n", - "This tutorial shows how `phasespace.fromdecay` can be used.\n", - "\n", - "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\n", - "More generally, `fromdecay` can also be used as a high-level interface for simulating particles that can decay in multiple different ways." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "# Import libraries\n", - "from pprint import pprint\n", - "\n", - "import zfit\n", - "from particle import Particle\n", - "from decaylanguage import DecFileParser, DecayChainViewer\n", - "\n", - "from phasespace.fromdecay import MultiDecayChain" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Intro to DecayLanguage\n", - "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", - "\n", - "We will begin by parsing a .dec file using DecayLanguage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "pycharm": { - "name": "#%%\n" - } - }, - "outputs": [], - "source": [ - "parser = DecFileParser('../tests/fromdecay/example_decays.dec')\n", - "parser.parse()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_chain = parser.build_decay_chains(\"pi0\")\n", - "pprint(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DecayChainViewer(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Creating a MultiDecayChain object\n", - "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `MultiDecayChain` object can be created directly from a DecayLanguage dict:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pi0_decay = MultiDecayChain.from_dict(pi0_chain)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When creating a `MultiDecayChain` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", - "\n", - "These GenParticle instances and the probabilities of that decay mode can be accessed via `MultiDecayChain.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for probability, particle in pi0_decay.gen_particles:\n", - " print(f\"There is a probability of {probability} \"\n", - " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", - "\n", - "When calling the `MultiDecayChain.generate` method, it internally calls the generate method on the of the GenParticle instances in `MultiDecayChain.gen_particles`. The outputs are placed in a list, which is returned." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "weights, events = pi0_decay.generate(n_events=10_000)\n", - "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can confirm that the counts above are close to the expected counts based on the probabilities. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Changing mass settings MultiDecayChain\n", - "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", - "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", - "These settings can be modified by passing in additional parameters to `MultiDecayChain.from_dict`.\n", - "There are two optional parameters that can be passed to `MultiDecayChain.from_dict`: `tolerance` and `mass_converter`.\n", - "\n", - "### Constant vs variable mass\n", - "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", - "This will be demonsttrated with the decay below:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", - "DecayChainViewer(dsplus_chain)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", - " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$\\pi^0$ has a greater width than $D^0$. \n", - "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dstar_decay = MultiDecayChain.from_dict(dsplus_chain, tolerance=1e-8)\n", - "# Loop over D0 and pi+ particles, see graph above\n", - "for particle in dstar_decay.gen_particles[0][1].children:\n", - " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", - " assert particle.has_fixed_mass\n", - " \n", - "# Loop over D+ and pi0. See above.\n", - "for particle in dstar_decay.gen_particles[1][1].children:\n", - " if particle.name == \"pi0\":\n", - " assert not particle.has_fixed_mass" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Configuring mass fucntions\n", - "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_custom_mass_func = dsplus_chain.copy()\n", - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", - "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `MultiDecayChain.from_dict`, like before." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MultiDecayChain.from_dict(dsplus_custom_mass_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def custom_gauss(mass, width):\n", - " particle_mass = tf.cast(mass, tf.float64)\n", - " particle_width = tf.cast(width, tf.float64)\n", - " \n", - " # This is the actual mass function that will be returned\n", - " def mass_func(min_mass, max_mass, n_events):\n", - " min_mass = tf.cast(min_mass, tf.float64)\n", - " max_mass = tf.cast(max_mass, tf.float64)\n", - " # Use a zfit PDF\n", - " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", - " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", - " return tf.vectorized_map(\n", - " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", - " )\n", - "\n", - " return mass_func" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This function can then be passed to `MultiDecayChain.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", - "print(\"Before:\")\n", - "pprint(dsplus_chain_subset)\n", - "\n", - "# Set the mass function of pi0 to the custom gaussian distribution \n", - "# when it decays into an electron-positron pair and a photon (gamma)\n", - "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", - "print(\"After:\")\n", - "pprint(dsplus_chain_subset)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MultiDecayChain.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.9.7" - } - }, - "nbformat": 4, - "nbformat_minor": 1 -} +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": true, + "pycharm": { + "is_executing": true, + "name": "#%% md\n" + } + }, + "source": [ + "# Tutorial for `GenMultiDecay` class\n", + "This tutorial shows how `phasespace.fromdecay.GenMultiDecay` can be used.\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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Import libraries\n", + "from pprint import pprint\n", + "\n", + "import zfit\n", + "from particle import Particle\n", + "from decaylanguage import DecFileParser, DecayChainViewer\n", + "\n", + "from phasespace.fromdecay import GenMultiDecay" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Intro to DecayLanguage\n", + "DecayLanguage can be used to parse and view .dec files. These files contain information about how a particle decays and with which probability. For more information about DecayLanguage and .dec files, see the [DecayLanguage](https://github.com/scikit-hep/decaylanguage) documentation.\n", + "\n", + "We will begin by parsing a .dec file using DecayLanguage:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "parser = DecFileParser('../tests/fromdecay/example_decays.dec')\n", + "parser.parse()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "From the `parser` variable, one can access a certain decay for a particle using `parser.build_decay_chains`. This will be a `dict` that contains all information about how the mother particle, daughter particles etc. decay." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_chain = parser.build_decay_chains(\"pi0\")\n", + "pprint(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This `dict` can also be displayed in a more human-readable way using `DecayChainViewer`: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DecayChainViewer(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating a GenMultiDecay object\n", + "A regular `phasespace.GenParticle` instance would not be able to simulate this decay, since the $\\pi^0$ particle can decay in four different ways. However, a `GenMultiDecay` object can be created directly from a DecayLanguage dict:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "pi0_decay = GenMultiDecay.from_dict(pi0_chain)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When creating a `GenMultiDecay` object, the DecayLanguage dict is \"unpacked\" into separate GenParticle instances, where each GenParticle instance corresponds to one way that the particle can decay.\n", + "\n", + "These GenParticle instances and the probabilities of that decay mode can be accessed via `GenMultiDecay.gen_particles`. This is a list of tuples, where the first element in the tuple is the probability and the second element is the GenParticle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for probability, particle in pi0_decay.gen_particles:\n", + " print(f\"There is a probability of {probability} \"\n", + " f\"that pi0 decays into {', '.join(child.name for child in particle.children)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One can simulate this decay using the `.generate` method, which works the same as the `GenParticle.generate` method.\n", + "\n", + "When calling the `GenMultiDecay.generate` method, it internally calls the generate method on the of the GenParticle instances in `GenMultiDecay.gen_particles`. The outputs are placed in a list, which is returned." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "weights, events = pi0_decay.generate(n_events=10_000)\n", + "print(\"Number of events for each decay mode:\", \", \".join(str(len(w)) for w in weights))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can confirm that the counts above are close to the expected counts based on the probabilities. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Changing mass settings\n", + "Since DecayLanguage dicts do not contain any information about the mass of a particle, the `fromdecay` submodule uses the [particle](https://github.com/scikit-hep/particle) package to find the mass of a particle based on its name. \n", + "The mass can either be a constant value or a function (besides the top particle, which is always a constant). \n", + "These settings can be modified by passing in additional parameters to `GenMultiDecay.from_dict`.\n", + "There are two optional parameters that can be passed to `GenMultiDecay.from_dict`: `tolerance` and `mass_converter`.\n", + "\n", + "### Constant vs variable mass\n", + "If a particle has a width less than `tolerance`, its mass is set to a constant value.\n", + "This will be demonsttrated with the decay below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain = parser.build_decay_chains(\"D*+\", stable_particles=[\"D+\"])\n", + "DecayChainViewer(dsplus_chain)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(f\"pi0 width = {Particle.from_evtgen_name('pi0').width}\\n\"\n", + " f\"D0 width = {Particle.from_evtgen_name('D0').width}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "$\\pi^0$ has a greater width than $D^0$. \n", + "If the tolerance is set to a value between their widths, the $D^0$ particle will have a constant mass while $\\pi^0$ will not. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dstar_decay = GenMultiDecay.from_dict(dsplus_chain, tolerance=1e-8)\n", + "# Loop over D0 and pi+ particles, see graph above\n", + "for particle in dstar_decay.gen_particles[0][1].children:\n", + " # If a particle width is less than tolerance or if it does not have any children, its mass will be fixed.\n", + " assert particle.has_fixed_mass\n", + " \n", + "# Loop over D+ and pi0. See above.\n", + "for particle in dstar_decay.gen_particles[1][1].children:\n", + " if particle.name == \"pi0\":\n", + " assert not particle.has_fixed_mass" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Configuring mass fucntions\n", + "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_custom_mass_func = dsplus_chain.copy()\n", + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Notice the added `zfit` field to the first decay mode of the $\\pi^0$ particle. This dict can then be passed to `GenMultiDecay.from_dict`, like before." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "GenMultiDecay.from_dict(dsplus_custom_mass_func)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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, it will automatically use the relativistic Breit-Wigner distribution.\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:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def custom_gauss(mass, width):\n", + " particle_mass = tf.cast(mass, tf.float64)\n", + " particle_width = tf.cast(width, tf.float64)\n", + " \n", + " # This is the actual mass function that will be returned\n", + " def mass_func(min_mass, max_mass, n_events):\n", + " min_mass = tf.cast(min_mass, tf.float64)\n", + " max_mass = tf.cast(max_mass, tf.float64)\n", + " # Use a zfit PDF\n", + " pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs=\"\")\n", + " iterator = tf.stack([min_mass, max_mass], axis=-1)\n", + " return tf.vectorized_map(\n", + " lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator\n", + " )\n", + "\n", + " return mass_func" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This function can then be passed to `GenMultiDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", + "print(\"Before:\")\n", + "pprint(dsplus_chain_subset)\n", + "\n", + "# Set the mass function of pi0 to the custom gaussian distribution \n", + "# when it decays into an electron-positron pair and a photon (gamma)\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "print(\"After:\")\n", + "pprint(dsplus_chain_subset)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "GenMultiDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 1 +} diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py index a93c7240..45536e6f 100644 --- a/phasespace/fromdecay/__init__.py +++ b/phasespace/fromdecay/__init__.py @@ -1,6 +1,6 @@ import sys -from .multidecaychain import MultiDecayChain # noqa: F401 +from .genmultidecay import GenMultiDecay # noqa: F401 try: import zfit # noqa: F401 diff --git a/phasespace/fromdecay/multidecaychain.py b/phasespace/fromdecay/genmultidecay.py similarity index 92% rename from phasespace/fromdecay/multidecaychain.py rename to phasespace/fromdecay/genmultidecay.py index 0f58781a..9018ac79 100644 --- a/phasespace/fromdecay/multidecaychain.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -13,7 +13,7 @@ _DEFAULT_MASS_FUNC = "relbw" -class MultiDecayChain: +class GenMultiDecay: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): """A container that works like GenParticle that can handle multiple decays. Can be created from. @@ -33,19 +33,25 @@ def from_dict( mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE, ): - """Create a MultiDecayChain instance from a dict in the decaylanguage format. + """Create a GenMultiDecay instance from a dict in the decaylanguage format. Args: - dec_dict: The input dict from which the MultiDecayChain object will be created from. + dec_dict: + The input dict from which the GenMultiDecay object will be created from. + A dec_dict has the same structure as the dicts used in DecayLanguage. + The structure of these dicts are: + >>> {particle_name: [{TODO}]} + mass_converter: A dict with mass function names and their corresponding mass functions. These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. + TODO more docs here tolerance: Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. Returns: - The created MultiDecayChain object. + The created GenMultiDecay object. """ if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index 23abd45a..80738fb9 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -1,64 +1,64 @@ -import tensorflow as tf -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) - - def mass_func(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="") - iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator - ) - - return mass_func - - -def breitwigner(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): - 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="") - iterator = tf.stack([min_mass, max_mass], axis=-1) - return tf.vectorized_map( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator - ) - - return mass_func - - -def relativistic_breitwigner(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): - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - pdf = zphys.pdf.RelativisticBreitWigner( - m=particle_mass, gamma=particle_width, obs="" - ) - iterator = tf.stack([min_mass, max_mass], axis=-1) - - # this works with map_fn but not with vectorized_map as no analytic sampling is available. - return tf.map_fn( - lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator - ) - - return mass_func - - -_DEFAULT_CONVERTER = { - "gauss": gauss, - "bw": breitwigner, - "relbw": relativistic_breitwigner, -} +import tensorflow as tf +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) + + def mass_func(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="") + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + + return mass_func + + +def breitwigner(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): + 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="") + iterator = tf.stack([min_mass, max_mass], axis=-1) + return tf.vectorized_map( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + ) + + return mass_func + + +def relativistic_breitwigner(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): + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + pdf = zphys.pdf.RelativisticBreitWigner( + m=particle_mass, gamma=particle_width, obs="" + ) + iterator = tf.stack([min_mass, max_mass], axis=-1) + + # this works with map_fn but not with vectorized_map as no analytic sampling is available. + return tf.map_fn( + lambda lim: pdf.sample(1, limits=(lim[0], lim[1])).unstack_x(), iterator + ) + + return mass_func + + +_DEFAULT_CONVERTER = { + "gauss": gauss, + "bw": breitwigner, + "relbw": relativistic_breitwigner, +} diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 8751f3ac..867e566f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,13 +1,13 @@ from numpy.testing import assert_almost_equal -from phasespace.fromdecay import MultiDecayChain +from phasespace.fromdecay import GenMultiDecay from phasespace.fromdecay.mass_functions import _DEFAULT_CONVERTER from . import example_decay_chains -def check_norm(full_decay: MultiDecayChain, **kwargs) -> list[tuple]: - """Helper function that checks whether the normalize_weights argument works for MultiDecayChain.generate. +def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: + """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate. Args: full_decay: full_decay.generate will be called. kwargs: Additional parameters passed to generate. @@ -33,7 +33,7 @@ def check_norm(full_decay: MultiDecayChain, **kwargs) -> list[tuple]: def test_single_chain(): """Test converting a decaylanguage dict with only one possible decay.""" - container = MultiDecayChain.from_dict( + container = GenMultiDecay.from_dict( example_decay_chains.dplus_single, tolerance=1e-10 ) output_decay = container.gen_particles @@ -58,7 +58,7 @@ def test_single_chain(): def test_branching_children(): """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" - container = MultiDecayChain.from_dict( + container = GenMultiDecay.from_dict( example_decay_chains.pi0_4branches, tolerance=1e-10 ) output_decays = container.gen_particles @@ -70,7 +70,7 @@ 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 = MultiDecayChain.from_dict(example_decay_chains.dplus_4grandbranches) + container = GenMultiDecay.from_dict(example_decay_chains.dplus_4grandbranches) output_decays = container.gen_particles assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) @@ -83,7 +83,7 @@ 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["D+"][0]["fs"][-1]["pi0"][-1]["zfit"] = "rel-BW" - container = MultiDecayChain.from_dict( + container = GenMultiDecay.from_dict( dplus_4grandbranches_massfunc, tolerance=1e-10, mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}, @@ -102,7 +102,7 @@ def test_mass_converter(): def test_big_decay(): - container = MultiDecayChain.from_dict(example_decay_chains.dstarplus_big_decay) + 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) diff --git a/tests/fromdecay/test_mass_functions.py b/tests/fromdecay/test_mass_functions.py index 43e75b48..e8c4eac0 100644 --- a/tests/fromdecay/test_mass_functions.py +++ b/tests/fromdecay/test_mass_functions.py @@ -1,54 +1,54 @@ -from typing import Callable - -import pytest -import tensorflow as tf -import tensorflow_probability as tfp -from particle import Particle - -import phasespace.fromdecay.mass_functions as mf - -_kstarz = Particle.from_evtgen_name("K*0") -KSTARZ_MASS = _kstarz.mass -KSTARZ_WIDTH = _kstarz.width - - -def ref_mass_func(min_mass, max_mass, n_events): - """Reference mass function used to compare the behavior of the actual mass functions. - - Args: - min_mass: lower limit of mass. - max_mass: upper limit of mass. - n_events: number of mass values that should be generated. - - Returns: - kstar_mass: Generated mass. - - Notes: - Code taken from phasespace documentation. - """ - min_mass = tf.cast(min_mass, tf.float64) - max_mass = tf.cast(max_mass, tf.float64) - kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) - kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) - kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) - kstar_mass = tfp.distributions.TruncatedNormal( - loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass - ).sample() - return kstar_mass - - -@pytest.mark.parametrize( - "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) -) -@pytest.mark.parametrize("size", (1, 10)) -def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): - rng = tf.random.Generator.from_seed(1234) - min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) - min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) - assert tf.reduce_all(min_mass <= max_mass) - ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) - sample = function(*params)(min_mass, max_mass, len(min_mass)) - assert sample.shape[0] == ref_sample.shape[0] - assert all( - i <= 1 for i in sample.shape[1:] - ) # Sample.shape have extra dimensions with just 1 or 0, which can be ignored +from typing import Callable + +import pytest +import tensorflow as tf +import tensorflow_probability as tfp +from particle import Particle + +import phasespace.fromdecay.mass_functions as mf + +_kstarz = Particle.from_evtgen_name("K*0") +KSTARZ_MASS = _kstarz.mass +KSTARZ_WIDTH = _kstarz.width + + +def ref_mass_func(min_mass, max_mass, n_events): + """Reference mass function used to compare the behavior of the actual mass functions. + + Args: + min_mass: lower limit of mass. + max_mass: upper limit of mass. + n_events: number of mass values that should be generated. + + Returns: + kstar_mass: Generated mass. + + Notes: + Code taken from phasespace documentation. + """ + min_mass = tf.cast(min_mass, tf.float64) + max_mass = tf.cast(max_mass, tf.float64) + kstar_width_cast = tf.cast(KSTARZ_WIDTH, tf.float64) + kstar_mass_cast = tf.cast(KSTARZ_MASS, dtype=tf.float64) + kstar_mass = tf.broadcast_to(kstar_mass_cast, shape=(n_events,)) + kstar_mass = tfp.distributions.TruncatedNormal( + loc=kstar_mass, scale=kstar_width_cast, low=min_mass, high=max_mass + ).sample() + return kstar_mass + + +@pytest.mark.parametrize( + "function", (mf.gauss, mf.breitwigner, mf.relativistic_breitwigner) +) +@pytest.mark.parametrize("size", (1, 10)) +def test_shape(function: Callable, size: int, params: tuple = (1.0, 1.0)): + rng = tf.random.Generator.from_seed(1234) + min_max_mass = rng.uniform(minval=0, maxval=1000, shape=(2, size), dtype=tf.float64) + min_mass, max_mass = tf.unstack(tf.sort(min_max_mass, axis=0), axis=0) + assert tf.reduce_all(min_mass <= max_mass) + ref_sample = ref_mass_func(min_mass, max_mass, len(min_mass)) + sample = function(*params)(min_mass, max_mass, len(min_mass)) + assert sample.shape[0] == ref_sample.shape[0] + assert all( + i <= 1 for i in sample.shape[1:] + ) # Sample.shape have extra dimensions with just 1 or 0, which can be ignored From b6b0e3365f368e9986eed6ac2d7f79343c3d1313 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 6 Nov 2021 18:48:03 +0100 Subject: [PATCH 39/71] Add examples to docstring of GenMultiDecay.from_dict --- docs/GenMultiDecay_Tutorial.ipynb | 6 ++- phasespace/fromdecay/genmultidecay.py | 68 ++++++++++++++++++++++++--- 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index cc289df2..3560c5ab 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -335,9 +335,11 @@ } ], "metadata": { + "interpreter": { + "hash": "bb65b8ec85759003b5d99658bc6210aa8fd7b9c8f144db79d452bd242727ce5f" + }, "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", + "display_name": "Python 3.9.7 64-bit ('phasespace': conda)", "name": "python3" }, "language_info": { diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 9018ac79..1808e56c 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -20,9 +20,6 @@ def __init__(self, gen_particles: list[tuple[float, GenParticle]]): Args: gen_particles: All the GenParticles and their corresponding probabilities. The list must be of the format [[probability, GenParticle instance], [probability, ... - - Notes: - Input format might change """ self.gen_particles = gen_particles @@ -39,9 +36,6 @@ def from_dict( dec_dict: The input dict from which the GenMultiDecay object will be created from. A dec_dict has the same structure as the dicts used in DecayLanguage. - The structure of these dicts are: - >>> {particle_name: [{TODO}]} - mass_converter: A dict with mass function names and their corresponding mass functions. These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. @@ -52,6 +46,68 @@ def from_dict( Returns: The created GenMultiDecay object. + + Examples: + DecayLanguage is usually used to create a dict that can be understood by GenMultiDecay. + >>> from decaylanguage import DecFileParser + >>> from phasespace.fromdecay import GenMultiDecay + + Parse a .dec file to create a DecayLanguage dict describing a D*+ particle + that can decay in 2 different ways: D*+ -> D0(->K- pi+) pi+ or D*+ -> D+ gamma. + + >>> parser = DecFileParser('example_decays.dec') + >>> parser.parse() + >>> dst_chain = parser.build_decay_chains("D*+") + >>> dst_chain + {'D*+': [{'bf': 0.984, + 'fs': [{'D0': [{'bf': 1.0, + 'fs': ['K-', 'pi+']}]}, + 'pi+']}, + {'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: + >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "gauss" + >>> dst_chain + {'D*+': [{'bf': 0.984, + 'fs': [{'D0': [{'bf': 1.0, + 'fs': ['K-', 'pi+'], + 'zfit': 'gauss'}]}, + 'pi+']}, + {'bf': 0.016, + 'fs': ['D+', 'gamma']}]} + + This dict can then be passed to `GenMultiDecay.from_dict`: + >>> dst_gen = GenMultiDecay.from_dict(dst_chain) + + If the decay of the D0 particle instead 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) + >>> particle_width = tf.cast(width, tf.float64) + >>> def mass_func(min_mass, max_mass, n_events): + >>> min_mass = tf.cast(min_mass, tf.float64) + >>> max_mass = tf.cast(max_mass, tf.float64) + >>> # Use a zfit PDF + >>> pdf = zfit.pdf.Gauss(mu=particle_mass, sigma=particle_width, obs="") + >>> iterator = tf.stack([min_mass, max_mass], axis=-1) + >>> return tf.vectorized_map( + >>> lambda lim: pdf.sample(1, limits=(lim[0], lim[1])), iterator + >>> ) + >>> return mass_func + + Once again change the distribution in the `dst_chain` dict. Here, it is changed to "custom gauss" + but any name can be used. + >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "custom gauss" + + 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}) + + Notes: + For a more in-depth tutorial, see the tutorial on GenMultiDecay in the + [documentation](https://phasespace.readthedocs.io/en/stable/tutorials/GenMultiDecay_Tutorial). """ if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER From 30d8ab0befbb612dcbaa63bef8e0dfde1221f45a Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Mon, 8 Nov 2021 23:17:42 +0100 Subject: [PATCH 40/71] Format classes and packages in monospaced font Co-authored-by: Eduardo Rodrigues --- phasespace/fromdecay/genmultidecay.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 1808e56c..d485bd1c 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -30,7 +30,9 @@ def from_dict( mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE, ): - """Create a GenMultiDecay instance from a dict in the decaylanguage format. + """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, + which is typically the result of parsing a `.dec` decay file. + Args: dec_dict: From e1aead91018dddaadd33841ec6107cdf9a179545 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 8 Nov 2021 22:18:07 +0000 Subject: [PATCH 41/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- phasespace/fromdecay/genmultidecay.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index d485bd1c..4e51c212 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -30,9 +30,8 @@ def from_dict( mass_converter: dict[str, Callable] = None, tolerance: float = _MASS_WIDTH_TOLERANCE, ): - """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, - which is typically the result of parsing a `.dec` decay file. - + """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, which is + typically the result of parsing a `.dec` decay file. Args: dec_dict: From 434cdb875a995318e557f75c8c94a6efdeaa9488 Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Fri, 12 Nov 2021 19:08:45 +0100 Subject: [PATCH 42/71] Clarify docstring Co-authored-by: Eduardo Rodrigues --- phasespace/fromdecay/genmultidecay.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 4e51c212..54e7ff89 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -15,7 +15,8 @@ class GenMultiDecay: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): - """A container that works like GenParticle that can handle multiple decays. Can be created from. + """A `GenParticle`-type container that can handle multiple decays. + Args: gen_particles: All the GenParticles and their corresponding probabilities. From 0e41fc3b6c1bfd00f6461ec969adc4671f0b6c42 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 12 Nov 2021 18:09:04 +0000 Subject: [PATCH 43/71] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- phasespace/fromdecay/genmultidecay.py | 1 - 1 file changed, 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 54e7ff89..671cf235 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -17,7 +17,6 @@ class GenMultiDecay: def __init__(self, gen_particles: list[tuple[float, GenParticle]]): """A `GenParticle`-type container that can handle multiple decays. - Args: gen_particles: All the GenParticles and their corresponding probabilities. The list must be of the format [[probability, GenParticle instance], [probability, ... From d8062e7a0fe80a90fc60057331c75dd326747354 Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Fri, 12 Nov 2021 19:11:03 +0100 Subject: [PATCH 44/71] Reference examples in dec_dict docstring Co-authored-by: Eduardo Rodrigues --- phasespace/fromdecay/genmultidecay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 671cf235..b0a0adbb 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -36,7 +36,7 @@ def from_dict( Args: dec_dict: The input dict from which the GenMultiDecay object will be created from. - A dec_dict has the same structure as the dicts used in DecayLanguage. + A dec_dict has the same structure as the dicts used in DecayLanguage, see the examples below. mass_converter: A dict with mass function names and their corresponding mass functions. These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. From 36ae8b13b307ae36630978ba731c5193c402a323 Mon Sep 17 00:00:00 2001 From: simonthor <45770021+simonthor@users.noreply.github.com> Date: Fri, 12 Nov 2021 19:12:48 +0100 Subject: [PATCH 45/71] decaylanguage to DecayLanguage Co-authored-by: Eduardo Rodrigues --- docs/GenMultiDecay_Tutorial.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index 3560c5ab..b26d5116 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -13,7 +13,7 @@ "# Tutorial for `GenMultiDecay` class\n", "This tutorial shows how `phasespace.fromdecay.GenMultiDecay` can be used.\n", "\n", - "This submodule makes it possible for `phasespace` and [`decaylanguage`](https://github.com/scikit-hep/decaylanguage/) to work together.\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." ] }, From bb80371975aa2764273b20c79f1a30b56e226a87 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 19:17:37 +0100 Subject: [PATCH 46/71] decaylanguage to DecayLanguage --- phasespace/fromdecay/genmultidecay.py | 8 ++++---- tests/fromdecay/example_decays.dec | 2 +- tests/fromdecay/test_fromdecay.py | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index b0a0adbb..5ff951c1 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -220,12 +220,12 @@ def _recursively_traverse( preexisting_particles: set[str] = None, tolerance: float = _MASS_WIDTH_TOLERANCE, ) -> list[tuple[float, GenParticle]]: - """Create all possible GenParticles by recursively traversing a dict from decaylanguage. + """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. Args: - 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. + 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. Returns: The generated GenParticle instances, one for each possible way of the decay. diff --git a/tests/fromdecay/example_decays.dec b/tests/fromdecay/example_decays.dec index c1eae76e..3b65a042 100644 --- a/tests/fromdecay/example_decays.dec +++ b/tests/fromdecay/example_decays.dec @@ -1,4 +1,4 @@ -# File originally from decaylanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec +# File originally from DecayLanguage tests: https://github.com/scikit-hep/decaylanguage/blob/master/tests/data/test_example_Dst.dec # Example decay chain for testing purposes # Considered by itself, this file in in fact incomplete, # as there are no instructions on how to decay the anti-D0 and the D-! diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 867e566f..2cf10996 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -32,7 +32,7 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: def test_single_chain(): - """Test converting a decaylanguage dict with only one possible decay.""" + """Test converting a DecayLanguage dict with only one possible decay.""" container = GenMultiDecay.from_dict( example_decay_chains.dplus_single, tolerance=1e-10 ) @@ -57,7 +57,7 @@ def test_single_chain(): def test_branching_children(): - """Test converting a decaylanguage dict where the mother particle can decay in many ways.""" + """Test converting a DecayLanguage dict where the mother particle can decay in many ways.""" container = GenMultiDecay.from_dict( example_decay_chains.pi0_4branches, tolerance=1e-10 ) @@ -69,7 +69,7 @@ def test_branching_children(): def test_branching_grandchilden(): - """Test converting a decaylanguage dict where children to the mother particle can decay in many ways.""" + """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) output_decays = container.gen_particles assert len(output_decays) == 4 From 632e44329710271df92c98f7765464f92aee7537 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 19:20:29 +0100 Subject: [PATCH 47/71] Same as last commit --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 79988259..380271a4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -43,7 +43,7 @@ fromdecay = particle >= 0.16.0 zfit zfit-physics >= 0.2 - decaylanguage >= 0.12.0 # not required but everyone using this feature will likely use decaylanguage + decaylanguage >= 0.12.0 # not required but everyone using this feature will likely use DecayLanguage test = awkward coverage From 0b6ec078e05213984811275d063028b07a477c1d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 20:37:46 +0100 Subject: [PATCH 48/71] Make constants public --- phasespace/fromdecay/__init__.py | 8 ++++++++ phasespace/fromdecay/genmultidecay.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py index 45536e6f..eb35172c 100644 --- a/phasespace/fromdecay/__init__.py +++ b/phasespace/fromdecay/__init__.py @@ -1,4 +1,5 @@ import sys +from typing import Tuple from .genmultidecay import GenMultiDecay # noqa: F401 @@ -12,3 +13,10 @@ "Either install phasespace[fromdecay] or particle and zfit-physics.", file=sys.stderr, ) from error + + +__all__ = ("GenMultiDecay", "MASS_WIDTH_TOLERANCE", "DEFAULT_MASS_FUNC") + + +def __dir__() -> Tuple[str, ...]: + return __all__ diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 5ff951c1..183b62b2 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -9,8 +9,8 @@ from .mass_functions import _DEFAULT_CONVERTER -_MASS_WIDTH_TOLERANCE = 0.01 -_DEFAULT_MASS_FUNC = "relbw" +MASS_WIDTH_TOLERANCE = 0.01 +DEFAULT_MASS_FUNC = "relbw" class GenMultiDecay: @@ -18,8 +18,8 @@ def __init__(self, gen_particles: list[tuple[float, GenParticle]]): """A `GenParticle`-type container that can handle multiple decays. Args: - gen_particles: All the GenParticles and their corresponding probabilities. - The list must be of the format [[probability, GenParticle instance], [probability, ... + gen_particles: All the GenParticles and their corresponding probabilities. + The list must be of the format [[probability, GenParticle instance], [probability, ... """ self.gen_particles = gen_particles @@ -28,7 +28,7 @@ def from_dict( cls, dec_dict: dict, mass_converter: dict[str, Callable] = None, - tolerance: float = _MASS_WIDTH_TOLERANCE, + tolerance: float = MASS_WIDTH_TOLERANCE, ): """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, which is typically the result of parsing a `.dec` decay file. @@ -194,7 +194,7 @@ def _get_particle_mass( name: str, mass_converter: dict[str, Callable], mass_func: str, - tolerance: float = _MASS_WIDTH_TOLERANCE, + tolerance: float = MASS_WIDTH_TOLERANCE, ) -> Union[Callable, float]: """ Get mass or mass function of particle using the particle package. @@ -218,7 +218,7 @@ def _recursively_traverse( decaychain: dict, mass_converter: dict[str, Callable], preexisting_particles: set[str] = None, - tolerance: float = _MASS_WIDTH_TOLERANCE, + tolerance: float = MASS_WIDTH_TOLERANCE, ) -> list[tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -279,7 +279,7 @@ def _recursively_traverse( mother_mass = _get_particle_mass( original_mother_name, mass_converter=mass_converter, - mass_func=dm.get("zfit", _DEFAULT_MASS_FUNC), + mass_func=dm.get("zfit", DEFAULT_MASS_FUNC), tolerance=tolerance, ) From bbdf8901368310f4702d2b193df069ea642e53ea Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 20:53:24 +0100 Subject: [PATCH 49/71] Minor clarifications Made the docs more clear, changed the link to the tutorial, and added some info about the constants. --- docs/GenMultiDecay_Tutorial.ipynb | 9 +++++---- phasespace/fromdecay/genmultidecay.py | 15 ++++++++------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index b26d5116..b8019bb2 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -33,6 +33,7 @@ "import zfit\n", "from particle import Particle\n", "from decaylanguage import DecFileParser, DecayChainViewer\n", + "import tensorflow as tf\n", "\n", "from phasespace.fromdecay import GenMultiDecay" ] @@ -264,7 +265,7 @@ "source": [ "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, it will automatically use the relativistic Breit-Wigner distribution.\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 `phasespace.fromdecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`.\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:" ] @@ -297,7 +298,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "This function can then be passed to `GenMultiDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom gauss\"`." + "This function can then be passed to `GenMultiDecay.from_dict` as a dict, where the key specifies the `zfit` parameter name. In the example below, it is set to `\"custom_gauss\"`. However, this name can be chosen arbitrarily and does not need to be the same as the function name." ] }, { @@ -312,7 +313,7 @@ "\n", "# Set the mass function of pi0 to the custom gaussian distribution \n", "# when it decays into an electron-positron pair and a photon (gamma)\n", - "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom gauss\"\n", + "dsplus_chain_subset[\"pi0\"][1][\"zfit\"] = \"custom_gauss\"\n", "print(\"After:\")\n", "pprint(dsplus_chain_subset)" ] @@ -323,7 +324,7 @@ "metadata": {}, "outputs": [], "source": [ - "GenMultiDecay.from_dict(dsplus_custom_mass_func, {\"custom gauss\": custom_gauss})" + "GenMultiDecay.from_dict(dsplus_custom_mass_func, {\"custom_gauss\": custom_gauss})" ] }, { diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 183b62b2..31908055 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -41,9 +41,10 @@ def from_dict( These functions should take the particle mass and the mass width as inputs and return a mass function that phasespace can understand. This dict will be combined with the predefined mass functions in this package. - TODO more docs here + 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. + assuming the mass to be constant. The default value is defined by MASS_WIDTH_TOLERANCE and + can be customized if desired. Returns: The created GenMultiDecay object. @@ -98,17 +99,17 @@ def from_dict( >>> ) >>> return mass_func - Once again change the distribution in the `dst_chain` dict. Here, it is changed to "custom gauss" + Once again change the distribution in the `dst_chain` dict. Here, it is changed to "custom_gauss" but any name can be used. - >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "custom gauss" + >>> dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "custom_gauss" - One can then pass the `custom_gauss` function and its name (in this case "custom gauss") as a + 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}) + >>> dst_gen = GenMultiDecay.from_dict(dst_chain, mass_converter={"custom_gauss": custom_gauss}) Notes: For a more in-depth tutorial, see the tutorial on GenMultiDecay in the - [documentation](https://phasespace.readthedocs.io/en/stable/tutorials/GenMultiDecay_Tutorial). + [documentation](https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html). """ if mass_converter is None: total_mass_converter = _DEFAULT_CONVERTER From 2a299dd7ab29e53526b1f8a64244bf8ca157ccd8 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 21:07:24 +0100 Subject: [PATCH 50/71] Make type hints complaiant with python < 3.9. Also changed _DEFAULT_CONVERTER to DEFAULT_CONVERTER --- phasespace/fromdecay/genmultidecay.py | 26 +++++++++++++------------- phasespace/fromdecay/mass_functions.py | 2 +- tests/fromdecay/test_fromdecay.py | 4 +++- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 31908055..06045d64 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -1,5 +1,5 @@ import itertools -from typing import Callable, Union +from typing import Callable, Dict, List, Set, Tuple, Union import tensorflow as tf import tensorflow.experimental.numpy as tnp @@ -7,14 +7,14 @@ from phasespace import GenParticle -from .mass_functions import _DEFAULT_CONVERTER +from .mass_functions import DEFAULT_CONVERTER MASS_WIDTH_TOLERANCE = 0.01 DEFAULT_MASS_FUNC = "relbw" class GenMultiDecay: - def __init__(self, gen_particles: list[tuple[float, GenParticle]]): + def __init__(self, gen_particles: List[Tuple[float, GenParticle]]): """A `GenParticle`-type container that can handle multiple decays. Args: @@ -27,7 +27,7 @@ def __init__(self, gen_particles: list[tuple[float, GenParticle]]): def from_dict( cls, dec_dict: dict, - mass_converter: dict[str, Callable] = None, + mass_converter: Dict[str, Callable] = None, tolerance: float = MASS_WIDTH_TOLERANCE, ): """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, which is @@ -112,10 +112,10 @@ def from_dict( [documentation](https://phasespace.readthedocs.io/en/stable/GenMultiDecay_Tutorial.html). """ if mass_converter is None: - total_mass_converter = _DEFAULT_CONVERTER + 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} + total_mass_converter = {**DEFAULT_CONVERTER, **mass_converter} gen_particles = _recursively_traverse( dec_dict, total_mass_converter, tolerance=tolerance @@ -125,8 +125,8 @@ def from_dict( def generate( self, n_events: int, normalize_weights: bool = True, **kwargs ) -> Union[ - 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). @@ -167,7 +167,7 @@ def generate( return weights, max_weights, events -def _unique_name(name: str, preexisting_particles: set[str]) -> str: +def _unique_name(name: str, preexisting_particles: Set[str]) -> str: """Create a string that does not exist in preexisting_particles based on name. Args: @@ -193,7 +193,7 @@ def _unique_name(name: str, preexisting_particles: set[str]) -> str: def _get_particle_mass( name: str, - mass_converter: dict[str, Callable], + mass_converter: Dict[str, Callable], mass_func: str, tolerance: float = MASS_WIDTH_TOLERANCE, ) -> Union[Callable, float]: @@ -217,10 +217,10 @@ def _get_particle_mass( def _recursively_traverse( decaychain: dict, - mass_converter: dict[str, Callable], - preexisting_particles: set[str] = None, + mass_converter: Dict[str, Callable], + preexisting_particles: Set[str] = None, tolerance: float = MASS_WIDTH_TOLERANCE, -) -> list[tuple[float, GenParticle]]: +) -> List[Tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. Args: diff --git a/phasespace/fromdecay/mass_functions.py b/phasespace/fromdecay/mass_functions.py index 80738fb9..f4ea260b 100644 --- a/phasespace/fromdecay/mass_functions.py +++ b/phasespace/fromdecay/mass_functions.py @@ -57,7 +57,7 @@ def mass_func(min_mass, max_mass, n_events): return mass_func -_DEFAULT_CONVERTER = { +DEFAULT_CONVERTER = { "gauss": gauss, "bw": breitwigner, "relbw": relativistic_breitwigner, diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 2cf10996..86a9ac96 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -1,3 +1,5 @@ +from typing import List + from numpy.testing import assert_almost_equal from phasespace.fromdecay import GenMultiDecay @@ -6,7 +8,7 @@ from . import example_decay_chains -def check_norm(full_decay: GenMultiDecay, **kwargs) -> list[tuple]: +def check_norm(full_decay: GenMultiDecay, **kwargs) -> List[tuple]: """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate. Args: full_decay: full_decay.generate will be called. From 512294227a5201e1a52ec77439c7ca8c4e8faa0d Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 21:33:23 +0100 Subject: [PATCH 51/71] Improve code quality --- phasespace/fromdecay/__init__.py | 12 +++++++++++- phasespace/fromdecay/genmultidecay.py | 4 ++-- tests/fromdecay/__init__.py | 1 + tests/fromdecay/test_fromdecay.py | 5 +++-- 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py index eb35172c..8cd67533 100644 --- a/phasespace/fromdecay/__init__.py +++ b/phasespace/fromdecay/__init__.py @@ -1,7 +1,17 @@ +"""This submodule makes it possible for `phasespace` and. + +[`DecayLanguage`](https://github.com/scikit-hep/decaylanguage/) to work together. +More generally, the `GenMultiDecay` object can also be used as a high-level interface for simulating particles +that can decay in multiple different ways. +""" import sys from typing import Tuple -from .genmultidecay import GenMultiDecay # noqa: F401 +from .genmultidecay import ( # noqa: F401 + DEFAULT_MASS_FUNC, + MASS_WIDTH_TOLERANCE, + GenMultiDecay, +) try: import zfit # noqa: F401 diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 06045d64..139f2163 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -197,8 +197,8 @@ def _get_particle_mass( mass_func: str, tolerance: float = MASS_WIDTH_TOLERANCE, ) -> Union[Callable, float]: - """ - Get mass or mass function of particle using the particle package. + """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 diff --git a/tests/fromdecay/__init__.py b/tests/fromdecay/__init__.py index 0234b711..87cb446b 100644 --- a/tests/fromdecay/__init__.py +++ b/tests/fromdecay/__init__.py @@ -1,3 +1,4 @@ +"""Tests for the fromdecay submodule.""" import pytest # This makes it so that assert errors are more helpful for e.g., the check_norm helper function diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 86a9ac96..2db85d5f 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -10,6 +10,7 @@ def check_norm(full_decay: GenMultiDecay, **kwargs) -> List[tuple]: """Helper function that checks whether the normalize_weights argument works for GenMultiDecay.generate. + Args: full_decay: full_decay.generate will be called. kwargs: Additional parameters passed to generate. @@ -51,7 +52,7 @@ def test_single_chain(): assert p.has_fixed_mass check_norm(container, n_events=1) - (normed_weights, decay_list), _ = check_norm(container, n_events=100) + (_, decay_list), _ = check_norm(container, n_events=100) assert len(decay_list) == 1 events = decay_list[0] assert set(events.keys()) == {"K-", "pi+", "pi+ [0]", "pi0", "gamma", "gamma [0]"} @@ -67,7 +68,7 @@ def test_branching_children(): assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) check_norm(container, n_events=1) - (normed_weights, events), _ = check_norm(container, n_events=100) + check_norm(container, n_events=100) def test_branching_grandchilden(): From 66d3ff3af68d6fb0804a756771c9ef0f225042ea Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 21:45:14 +0100 Subject: [PATCH 52/71] Improve code quality (2) --- phasespace/fromdecay/__init__.py | 3 +-- tests/fromdecay/test_fromdecay.py | 6 +++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py index 8cd67533..70932ea8 100644 --- a/phasespace/fromdecay/__init__.py +++ b/phasespace/fromdecay/__init__.py @@ -1,6 +1,5 @@ -"""This submodule makes it possible for `phasespace` and. +"""This submodule makes it possible for `phasespace` and `DecayLanguage` to work together. -[`DecayLanguage`](https://github.com/scikit-hep/decaylanguage/) to work together. More generally, the `GenMultiDecay` object can also be used as a high-level interface for simulating particles that can decay in multiple different ways. """ diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 2db85d5f..2efebc7b 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -78,7 +78,7 @@ def test_branching_grandchilden(): assert len(output_decays) == 4 assert_almost_equal(sum(d[0] for d in output_decays), 1) check_norm(container, n_events=1) - (normed_weights, events), _ = check_norm(container, n_events=100) + check_norm(container, n_events=100) # TODO add more asserts here @@ -101,7 +101,7 @@ def test_mass_converter(): assert not child.has_fixed_mass check_norm(container, n_events=1) - (normed_weights, events), _ = check_norm(container, n_events=100) + check_norm(container, n_events=100) def test_big_decay(): @@ -109,5 +109,5 @@ def test_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) - (normed_weights, events), _ = check_norm(container, n_events=100) + check_norm(container, n_events=100) # TODO add more asserts here From 4a47428fa6414816778d77377a1dd29bc6832e7f Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Fri, 12 Nov 2021 21:52:37 +0100 Subject: [PATCH 53/71] Shorten docstring summary to improve code quality --- phasespace/fromdecay/genmultidecay.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 139f2163..1cb2a613 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -30,8 +30,7 @@ def from_dict( mass_converter: Dict[str, Callable] = None, tolerance: float = MASS_WIDTH_TOLERANCE, ): - """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format, which is - typically the result of parsing a `.dec` decay file. + """Create a `GenMultiDecay` instance from a dict in the `DecayLanguage` package format. Args: dec_dict: From decf916103ad7ea600beeb1436f1812557de6a7e Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 13 Nov 2021 16:40:15 +0100 Subject: [PATCH 54/71] Change MASS_WIDTH_TOLERANCE and DEFAULT_MASS_FUNC to class variables. --- docs/GenMultiDecay_Tutorial.ipynb | 2 +- phasespace/fromdecay/__init__.py | 8 ++------ phasespace/fromdecay/genmultidecay.py | 12 ++++++------ tests/fromdecay/test_fromdecay.py | 4 ++-- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index b8019bb2..f9c9f6da 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -265,7 +265,7 @@ "source": [ "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 `phasespace.fromdecay.DEFAULT_MASS_FUNC` to a different string, e.g., `\"gauss\"`.\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", "\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/phasespace/fromdecay/__init__.py b/phasespace/fromdecay/__init__.py index 70932ea8..da9b6387 100644 --- a/phasespace/fromdecay/__init__.py +++ b/phasespace/fromdecay/__init__.py @@ -6,11 +6,7 @@ import sys from typing import Tuple -from .genmultidecay import ( # noqa: F401 - DEFAULT_MASS_FUNC, - MASS_WIDTH_TOLERANCE, - GenMultiDecay, -) +from .genmultidecay import GenMultiDecay # noqa: F401 try: import zfit # noqa: F401 @@ -24,7 +20,7 @@ ) from error -__all__ = ("GenMultiDecay", "MASS_WIDTH_TOLERANCE", "DEFAULT_MASS_FUNC") +__all__ = ("GenMultiDecay",) def __dir__() -> Tuple[str, ...]: diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index 1cb2a613..fccc3566 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -9,11 +9,11 @@ from .mass_functions import DEFAULT_CONVERTER -MASS_WIDTH_TOLERANCE = 0.01 -DEFAULT_MASS_FUNC = "relbw" - class GenMultiDecay: + MASS_WIDTH_TOLERANCE = 0.01 + DEFAULT_MASS_FUNC = "relbw" + def __init__(self, gen_particles: List[Tuple[float, GenParticle]]): """A `GenParticle`-type container that can handle multiple decays. @@ -194,7 +194,7 @@ def _get_particle_mass( name: str, mass_converter: Dict[str, Callable], mass_func: str, - tolerance: float = MASS_WIDTH_TOLERANCE, + tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE, ) -> Union[Callable, float]: """Get mass or mass function of particle using the particle package. @@ -218,7 +218,7 @@ def _recursively_traverse( decaychain: dict, mass_converter: Dict[str, Callable], preexisting_particles: Set[str] = None, - tolerance: float = MASS_WIDTH_TOLERANCE, + tolerance: float = GenMultiDecay.MASS_WIDTH_TOLERANCE, ) -> List[Tuple[float, GenParticle]]: """Create all possible GenParticles by recursively traversing a dict from DecayLanguage, see Examples. @@ -279,7 +279,7 @@ def _recursively_traverse( mother_mass = _get_particle_mass( original_mother_name, mass_converter=mass_converter, - mass_func=dm.get("zfit", DEFAULT_MASS_FUNC), + mass_func=dm.get("zfit", GenMultiDecay.DEFAULT_MASS_FUNC), tolerance=tolerance, ) diff --git a/tests/fromdecay/test_fromdecay.py b/tests/fromdecay/test_fromdecay.py index 2efebc7b..76c14ea7 100644 --- a/tests/fromdecay/test_fromdecay.py +++ b/tests/fromdecay/test_fromdecay.py @@ -3,7 +3,7 @@ from numpy.testing import assert_almost_equal from phasespace.fromdecay import GenMultiDecay -from phasespace.fromdecay.mass_functions import _DEFAULT_CONVERTER +from phasespace.fromdecay.mass_functions import DEFAULT_CONVERTER from . import example_decay_chains @@ -89,7 +89,7 @@ def test_mass_converter(): container = GenMultiDecay.from_dict( dplus_4grandbranches_massfunc, tolerance=1e-10, - mass_converter={"rel-BW": _DEFAULT_CONVERTER["relbw"]}, + mass_converter={"rel-BW": DEFAULT_CONVERTER["relbw"]}, ) output_decays = container.gen_particles assert len(output_decays) == 4 From 03f4deffd77c9fbd9b826099ad87dc4f7fabb997 Mon Sep 17 00:00:00 2001 From: Simon Thor Date: Sat, 13 Nov 2021 19:50:26 +0100 Subject: [PATCH 55/71] Change docstring to reference class variable --- phasespace/fromdecay/genmultidecay.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/phasespace/fromdecay/genmultidecay.py b/phasespace/fromdecay/genmultidecay.py index fccc3566..d790b246 100644 --- a/phasespace/fromdecay/genmultidecay.py +++ b/phasespace/fromdecay/genmultidecay.py @@ -42,8 +42,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 MASS_WIDTH_TOLERANCE and - can be customized if desired. + assuming the mass to be constant. The default value is defined by the class variable + MASS_WIDTH_TOLERANCE and can be customized if desired. Returns: The created GenMultiDecay object. From 3374b3fd5111eb3cebb4909bf957a0f47c998da3 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 19:38:08 +0100 Subject: [PATCH 56/71] Add nbval to CI --- .github/workflows/ci.yml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0c95dfdd..f94fde50 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,6 +52,10 @@ jobs: - name: Download test data run: python data/download_test_files.py > /dev/null - name: Test with pytest - run: PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} + run: | + PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} + PHASESPACE_EAGER=0 pytest --nbval-lax docs - name: Test with pytest (eager mode) - run: PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} + run: | + PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} + PHASESPACE_EAGER=1 pytest --nbval-lax docs From 2cf9c8cf93a12b8500e3d92bef39024d112d197f Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 19:42:55 +0100 Subject: [PATCH 57/71] Update ci.yml --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f94fde50..1de9a7a6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,7 +3,7 @@ name: tests -on: [push] +on: [push, pull_request] jobs: codecov: From 690e16d23b6709caf62508877c411396ed8179e7 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 19:50:58 +0100 Subject: [PATCH 58/71] Add fromdecay to test requirements --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 380271a4..75a3a034 100644 --- a/setup.cfg +++ b/setup.cfg @@ -45,6 +45,7 @@ fromdecay = zfit-physics >= 0.2 decaylanguage >= 0.12.0 # not required but everyone using this feature will likely use DecayLanguage test = + %(fromdecay)s awkward coverage flaky From dfe560421338d75e495be6d968874d43e82ee900 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 19:58:33 +0100 Subject: [PATCH 59/71] Update GenMultiDecay_Tutorial.ipynb --- docs/GenMultiDecay_Tutorial.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index f9c9f6da..17bf6227 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -223,7 +223,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### Configuring mass fucntions\n", + "### Configuring mass functions\n", "By default, the mass function used for variable mass is the relativistic Breit-Wigner distribution. This can however be changed. If you want the mother particle to have a specific mass function for a specific decay, you can add a `zfit` parameter to the DecayLanguage dict. Consider for example the previous $D^{*+}$ example:" ] }, @@ -237,7 +237,7 @@ "dsplus_chain_subset = dsplus_custom_mass_func[\"D*+\"][1][\"fs\"][1]\n", "print(\"Before:\")\n", "pprint(dsplus_chain_subset)\n", - "# Set the mass function of pi0 to a gaussian distribution when it decays into 2 photons (gamma)\n", + "# Set the mass function of pi0 to a gaussian distribution when it decays into two photons (gamma)\n", "dsplus_chain_subset[\"pi0\"][0][\"zfit\"] = \"gauss\"\n", "print(\"After:\")\n", "pprint(dsplus_chain_subset)" From 58538110a90a9c9d5d3ec5a883745617942ed51a Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 19:59:22 +0100 Subject: [PATCH 60/71] Update setup.cfg --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 75a3a034..4db9bd33 100644 --- a/setup.cfg +++ b/setup.cfg @@ -50,6 +50,7 @@ test = coverage flaky matplotlib + nbval numpy pytest pytest-cov From e5c9f8a2571670be3b2f7602a65bb6db8addaf71 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Sun, 14 Nov 2021 20:12:16 +0100 Subject: [PATCH 61/71] CI: correct execution place for notebook --- .github/workflows/ci.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1de9a7a6..30f36be2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,8 +54,12 @@ jobs: - name: Test with pytest run: | PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} - PHASESPACE_EAGER=0 pytest --nbval-lax docs + cd docs + PHASESPACE_EAGER=0 pytest --nbval-lax + cd .. - name: Test with pytest (eager mode) run: | PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} + cd docs PHASESPACE_EAGER=1 pytest --nbval-lax docs + cd .. From eae7abb5f14197f844db18b424d403b5802e4c2b Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Sun, 14 Nov 2021 21:09:49 +0100 Subject: [PATCH 62/71] ci: explicit test directory --- .github/workflows/ci.yml | 4 ++-- docs/GenMultiDecay_Tutorial.ipynb | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 30f36be2..4c1dda73 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,13 +53,13 @@ jobs: run: python data/download_test_files.py > /dev/null - name: Test with pytest run: | - PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} + PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} tests cd docs PHASESPACE_EAGER=0 pytest --nbval-lax cd .. - name: Test with pytest (eager mode) run: | - PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} + PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} tests cd docs PHASESPACE_EAGER=1 pytest --nbval-lax docs cd .. diff --git a/docs/GenMultiDecay_Tutorial.ipynb b/docs/GenMultiDecay_Tutorial.ipynb index 17bf6227..642e1adc 100644 --- a/docs/GenMultiDecay_Tutorial.ipynb +++ b/docs/GenMultiDecay_Tutorial.ipynb @@ -340,7 +340,8 @@ "hash": "bb65b8ec85759003b5d99658bc6210aa8fd7b9c8f144db79d452bd242727ce5f" }, "kernelspec": { - "display_name": "Python 3.9.7 64-bit ('phasespace': conda)", + "display_name": "Python 3 (ipykernel)", + "language": "python", "name": "python3" }, "language_info": { @@ -358,4 +359,4 @@ }, "nbformat": 4, "nbformat_minor": 1 -} +} \ No newline at end of file From d446b7cc9bdea023157049210381e149bc365e14 Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Sun, 14 Nov 2021 21:33:31 +0100 Subject: [PATCH 63/71] chore: add graphviz to dependencies for notebook --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 4db9bd33..d6947b01 100644 --- a/setup.cfg +++ b/setup.cfg @@ -60,6 +60,7 @@ test = wget decaylanguage doc = + graphviz Sphinx sphinx_bootstrap_theme jupyter_sphinx From 649d86db49fb52cbc7aa5de805b430b860919425 Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Mon, 15 Nov 2021 10:23:09 +0100 Subject: [PATCH 64/71] chore: remove duplicated requirements --- setup.cfg | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index d6947b01..389b5534 100644 --- a/setup.cfg +++ b/setup.cfg @@ -58,7 +58,6 @@ test = scipy uproot4 wget - decaylanguage doc = graphviz Sphinx From b46c0643ffe8444dcdc2cc39844eb87cba9966b4 Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Mon, 15 Nov 2021 14:52:06 +0100 Subject: [PATCH 65/71] ci: add dev installation --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4c1dda73..19e3fc95 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,7 +21,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install --use-feature=2020-resolver -e .[test] + pip install --use-feature=2020-resolver -e .[dev] pip install coverage - name: Download test data run: python data/download_test_files.py > /dev/null From c88957b5b24b295924cf48410811677289e9974a Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Mon, 15 Nov 2021 16:32:40 +0100 Subject: [PATCH 66/71] ci: ignore ipynb checkpoints --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 19e3fc95..930933f3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -55,11 +55,11 @@ jobs: run: | PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} tests cd docs - PHASESPACE_EAGER=0 pytest --nbval-lax + PHASESPACE_EAGER=0 pytest --nbval-lax --ignore=.ipynb_checkpoints cd .. - name: Test with pytest (eager mode) run: | PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} tests cd docs - PHASESPACE_EAGER=1 pytest --nbval-lax docs + PHASESPACE_EAGER=1 pytest --nbval-lax --ignore=.ipynb_checkpoints cd .. From f8d0341b7ea4db6c90a6d0d62637c0d0661da1ab Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Mon, 15 Nov 2021 18:13:02 +0100 Subject: [PATCH 67/71] ci: ignore notebook --- .github/workflows/ci.yml | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 930933f3..3516d3e4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,12 +54,14 @@ jobs: - name: Test with pytest run: | PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} tests - cd docs - PHASESPACE_EAGER=0 pytest --nbval-lax --ignore=.ipynb_checkpoints - cd .. + # TODO: activate tutorial again. Fails because graphviz is not installed + # cd docs + # PHASESPACE_EAGER=0 pytest --nbval-lax --ignore=.ipynb_checkpoints + # cd .. - name: Test with pytest (eager mode) run: | PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} tests - cd docs - PHASESPACE_EAGER=1 pytest --nbval-lax --ignore=.ipynb_checkpoints - cd .. + +# cd docs +# PHASESPACE_EAGER=1 pytest --nbval-lax --ignore=.ipynb_checkpoints +# cd .. From c529263651124380066d5019abacbeeb3f4fab37 Mon Sep 17 00:00:00 2001 From: Jonas Eschle Date: Mon, 22 Nov 2021 14:26:21 +0100 Subject: [PATCH 68/71] Update setup.cfg --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 389b5534..8cf7e88d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -31,7 +31,7 @@ packages = find: setup_requires = setuptools_scm install_requires = - tensorflow>=2.5,<2.7 # tensorflow.experimental.numpy + tensorflow>=2.5,<2.8 # tensorflow.experimental.numpy tensorflow_probability>=0.11 python_requires = >=3.6 include_package_data = True From c1ef3b3f51caad5554043e3efc183d5638e31156 Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Wed, 24 Nov 2021 10:22:06 +0100 Subject: [PATCH 69/71] fix: error with new tf version, requiring keras < 2.7 --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index 8cf7e88d..1b4a1579 100644 --- a/setup.cfg +++ b/setup.cfg @@ -33,6 +33,7 @@ setup_requires = install_requires = tensorflow>=2.5,<2.8 # tensorflow.experimental.numpy tensorflow_probability>=0.11 + keras <2.7 python_requires = >=3.6 include_package_data = True testpaths = tests From 2bd2d2f2dabe0033c1dae19aab6d237bfae3e547 Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Wed, 24 Nov 2021 15:20:18 +0100 Subject: [PATCH 70/71] ci: debug notebook error 1 --- .github/workflows/ci.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3516d3e4..83e599d0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,10 +54,9 @@ jobs: - name: Test with pytest run: | PHASESPACE_EAGER=0 pytest --basetemp={envtmpdir} tests - # TODO: activate tutorial again. Fails because graphviz is not installed - # cd docs - # PHASESPACE_EAGER=0 pytest --nbval-lax --ignore=.ipynb_checkpoints - # cd .. + cd docs + PHASESPACE_EAGER=0 pytest --nbval-lax --ignore=.ipynb_checkpoints + cd .. - name: Test with pytest (eager mode) run: | PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} tests From 80f812150cb3015ada92c3a87de0214cd5d337fb Mon Sep 17 00:00:00 2001 From: Jonas Eschle 'Mayou36 Date: Wed, 24 Nov 2021 15:30:03 +0100 Subject: [PATCH 71/71] ci: debug notebook error 2, adding graphviz system package --- .github/workflows/ci.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 83e599d0..ab2cec3b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,6 +48,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip + sudo apt-get install graphviz -y pip install --use-feature=2020-resolver -e .[dev] - name: Download test data run: python data/download_test_files.py > /dev/null @@ -60,7 +61,6 @@ jobs: - name: Test with pytest (eager mode) run: | PHASESPACE_EAGER=1 pytest --basetemp={envtmpdir} tests - -# cd docs -# PHASESPACE_EAGER=1 pytest --nbval-lax --ignore=.ipynb_checkpoints -# cd .. + cd docs + PHASESPACE_EAGER=1 pytest --nbval-lax --ignore=.ipynb_checkpoints + cd ..