Skip to content

API reference#

phasespace#

phasespace.phasespace #

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)

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 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 jax.Array 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.

Parameters:

Name Type Description Default
name str

Name of the particle.

 required
mass (float, array - like, callable)

Mass of the particle. If it's a float, it gets converted to array-like. If it is a callable, it is called as mass(min_mass, max_mass, n_events, key) and has to be jit-compatible.

 required

has_fixed_mass property #

bool: Is the mass a callable function?

has_children property #

bool: Does the particle have children?

has_grandchildren property #

bool: Does the particle have grandchildren?

get_mass(min_mass=None, max_mass=None, n_events=None, key=None) #

Get the particle mass.

If the particle is resonant, the mass function is called as mass(min_mass, max_mass, n_events, key).

Parameters:

Name Type Description Default
min_mass array

Lower mass range. Defaults to None, which is only valid in the case of fixed mass.

 None
max_mass array

Upper mass range. Defaults to None, which is only valid in the case of fixed mass.

 None
n_events int

Number of events to produce. Has to be specified if the particle is resonant.

 None
key ``jax.Array``

JAX PRNG key, handed to the mass function. Has to be specified if the particle is resonant.

 None

Returns:

Type Description
Array

jax.Array: Mass of the particles, either a scalar or shape (nevents,)

Raises:

Type Description
ValueError

If the mass is requested and has not been set.

set_children(*children) #

Assign children.

Parameters:

Name Type Description Default
children GenParticle

Two or more children to assign to the current particle.

 ()

Returns:

Type Description

self

Raises:

Type Description
ValueError

If there is an inconsistency in the parent/children relationship, ie,
KeyError

If there is a particle name clash.
RuntimeError

If generate was already called before.

generate(n_events, boost_to=None, normalize_weights=True, key=None, *, as_vectors=None, chunk_size=None) #

Generate normalized n-body phase space as JAX arrays.

Events are generated in the rest frame of the particle, unless boost_to is given.

Notes

In this method, the event weights are returned normalized to their maximum.

The generation is jit-compiled with n_events as a static argument: calling this with a new value of n_events triggers a recompilation, while repeated calls with the same value reuse the compiled function.

Chunking changes which events are drawn, as every chunk consumes its own split of key: generate(n, key=k) and generate(n, key=k, chunk_size=c) give different but equally valid samples, each of them reproducible. It bounds the memory of the generation itself, not of the returned arrays, and keeps the number of compilations at two (a full chunk and the remainder) whatever n_events is.

Parameters:

Name Type Description Default
n_events int

Number of events to generate.

 required
boost_to optional

Momentum vector of shape (x, 4), where x is optional, to where the resulting events will be boosted in the (px, py, pz, E) format. Can also be a vector momentum Lorentz vector. If not specified, events are generated in the rest frame of the particle.

 None
normalize_weights bool

Normalize the event weight to its max?

 True
key ``KeyLike``

Either an integer seed, a JAX PRNG key or None, in which case a new key is created from OS entropy (and the generation is not reproducible).

 None
as_vectors bool

If True, the output momenta are returned as vector objects.

 None
chunk_size int

Generate the events in chunks of at most this many rather than all at once, which bounds the peak memory of the generation. Defaults to None, which generates everything in one go.

 None

Returns:

Name Type Description
tuple tuple[Array, dict[str, Array]] | tuple[Array, Array, dict[str, Array]]

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 of the generated particles as a dictionary of arrays of shape (n_events, 4) with particle names as keys.

  • If False, the tuple elements 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 of the generated particles as a dictionary of arrays of shape (n_events, 4) with particle names as keys.

Raises:

Type Description
ValueError

If the decay is kinematically forbidden, if n_events and the size of boost_to don't match or if chunk_size is not positive.

nbody_decay(mass_top, masses, top_name='', names=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.

Parameters:

Name Type Description Default
mass_top (array, list)

Mass of the top particle. Can be a list of 4-vectors.

 required
masses list

Masses of the child particles.

 required
top_name str

Name of the top particle. If not given, the top particle is named top.

 ''
names list

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.

 None

Returns:

Type Description

GenParticle: Particle decay.

Raises:

Type Description
ValueError

If the length of masses and names doesn't match.

to_vectors(particles) #

Convert a dictionary of particles to a dictionary of vector.Momentum instances.

Parameters:

Name Type Description Default
particles dict

Dictionary of particles, with the keys being the particle names and the values being the momenta.

 required

Returns:

Name Type Description
dict dict[str, Momentum]

Dictionary of vector.Momentum instances with numpy arrays

phasespace.kinematics#

phasespace.kinematics #

Basic kinematics.

scalar_product(vec1, vec2) #

Calculate scalar product of two 3-vectors.

Parameters:

Name Type Description Default
vec1

First vector.

 required
vec2

Second vector.

 required

Returns:

Type Description

Scalar product of the two vectors.

spatial_component(vector) #

Extract spatial components of the input Lorentz vector.

Parameters:

Name Type Description Default
vector

Input Lorentz vector (where indexes 0-2 are space, index 3 is time).

 required

Returns:

Type Description

Spatial components (3-vector) of the input Lorentz vector.

time_component(vector) #

Extract time component of the input Lorentz vector.

Parameters:

Name Type Description Default
vector

Input Lorentz vector (where indexes 0-2 are space, index 3 is time).

 required

Returns:

Type Description

Time component of the input Lorentz vector.

x_component(vector) #

Extract spatial X component of the input Lorentz or 3-vector.

Parameters:

Name Type Description Default
vector

Input vector.

 required

Returns:

Type Description

X component of the input vector.

y_component(vector) #

Extract spatial Y component of the input Lorentz or 3-vector.

Parameters:

Name Type Description Default
vector

Input vector.

 required

Returns:

Type Description

Y component of the input vector.

z_component(vector) #

Extract spatial Z component of the input Lorentz or 3-vector.

Parameters:

Name Type Description Default
vector

Input vector.

 required

Returns:

Type Description

Z component of the input vector.

mass(vector) #

Calculate mass scalar for Lorentz 4-momentum.

Parameters:

Name Type Description Default
vector

Input Lorentz momentum vector.

 required

Returns:

Type Description

Mass of the Lorentz 4-momentum vector.

lorentz_vector(space, time) #

Make a Lorentz vector from spatial and time components.

Parameters:

Name Type Description Default
space

3-vector of spatial components.

 required
time

Time component.

 required

Returns:

Type Description

Lorentz 4-vector combining spatial and time components.

lorentz_boost(vector, boostvector) #

Perform Lorentz boost.

Parameters:

Name Type Description Default
vector

4-vector to be boosted

 required
boostvector

Boost vector. Can be either 3-vector or 4-vector, since only spatial components are used.

 required

Returns:

Type Description

Boosted 4-vector.

beta(vector) #

Calculate beta of a given 4-vector.

Parameters:

Name Type Description Default
vector

Input Lorentz momentum vector.

 required

Returns:

Type Description

Beta (v/c) of the Lorentz momentum vector.

boost_components(vector) #

Get the boost components of a given 4-vector.

Parameters:

Name Type Description Default
vector

Input Lorentz momentum vector.

 required

Returns:

Type Description

Boost components (3-vector) of the Lorentz momentum vector.

metric_tensor() #

Metric tensor for Lorentz space (constant).

Returns:

Type Description

Metric tensor for Lorentz space with signature (-1, -1, -1, 1).

phasespace.random#

phasespace.random #

Random number generation.

JAX random number generation is purely functional: every draw is an explicit function of a PRNG key. This module only normalizes what users may pass as a key.

ensure_key(key=None) #

Normalize a user-supplied key into a JAX PRNG key.

Parameters:

Name Type Description Default
key KeyLike

This can be - None to create a new, non-reproducible key from OS entropy, - an integer seed to create a key deterministically, - a JAX PRNG key, which is returned unchanged.

 None

Returns:

Type Description
Array

A JAX PRNG key.
Notes

Never call this with None inside a jitted function: the key would be created once at trace time and every call would then reuse the very same random numbers.

phasespace.precision#

phasespace.precision #

Double precision handling.

The phase space computation asks for float64 explicitly everywhere, but JAX only honours that while its x64 mode is enabled: with the mode off, every request is truncated to float32 and warns once per call. Single precision is not an option here, as pdk suffers catastrophic cancellation close to threshold, which degrades energy-momentum conservation from ~1e-15 to ~1e-6 relative.

The mode is therefore enabled per call rather than at import time, which keeps the dtype defaults of the calling program untouched.

with_float64(func) #

Run func with JAX's double precision mode enabled.

Enabling the mode is scoped to the call, so it does not change the dtype defaults of the calling program. The returned arrays are float64.

phasespace.fromdecay#

phasespace.fromdecay.genmultidecay #

GenMultiDecay #

__init__(gen_particles) #

A GenParticle-type container that can handle multiple decays.

Parameters:

Name Type Description Default
gen_particles list[tuple[float, GenParticle]]

All the GenParticles and their corresponding probabilities. The list must be of the format [[probability, GenParticle instance], [probability, ...

 required

from_dict(dec_dict, mass_converter=None, tolerance=None, particle_model_map=None) classmethod #

Create a GenMultiDecay instance from a dict in the DecayLanguage package format.

Parameters:

Name Type Description Default
dec_dict 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, see the examples below.

 required
mass_converter dict[str, Callable] | None

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. See the Example below or the tutorial for how to use this parameter.

 None
tolerance float | None

Minimum mass width of the particle to use a mass function instead of assuming the mass to be constant. If None, the default value, defined by the class variable MASS_WIDTH_TOLERANCE, will be used. This value can be customized if desired.

 None
particle_model_map dict[str, str] | None

A dict where the key is a particle name and the value is a mass function name. If a particle is specified in the particle_model_map, then all appearances of this particle in dec_dict will get the same mass function. This way, one doesn't have to manually add the zfit parameter to every place where this particle appears in dec_dict. If the zfit parameter is specified for a particle which is also included in particle_model_map, the zfit parameter mass function name will be prioritized.

 None

Returns:

Type Description

The created GenMultiDecay object.

Examples:

Basic Usage

DecayLanguage is usually used to create a dict that can be understood by GenMultiDecay:

.. code-block:: python

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 will be:
# {'D*+': [{'bf': 0.984,
#     'fs': [{'D0': [{'bf': 1.0,
#             'fs': ['K-', 'pi+']}]},
#         'pi+']},
#     {'bf': 0.016,
#      'fs': ['D+', 'gamma']}]}

Using particle_model_map

If the D0 particle should have a mass distribution of a gaussian when it decays, one can pass the particle_model_map parameter to from_dict:

.. code-block:: python

dst_gen = GenMultiDecay.from_dict(dst_chain, particle_model_map={"D0": "gauss"})

This will then set the mass function of D0 to a gaussian for all its decays.

Using the zfit parameter

If more custom control is required, e.g., if D0 can decay in multiple ways and one of the decays should have a specific mass function, a zfit parameter can be added to its decay dict:

.. code-block:: python

dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "gauss"
dst_gen = GenMultiDecay.from_dict(dst_chain)

This will now make the D0 particle have a gaussian mass function, only when it decays into K- and pi+. In this case, there are no other ways that D0 can decay, so using particle_model_map is a cleaner and easier option.

Custom Mass Functions

If the decay of the D0 particle should be modelled by a mass distribution that does not come with the package, a custom distribution can be created:

.. code-block:: python

def custom_gauss(mass, width):
    def mass_func(min_mass, max_mass, n_events, key):
        # a normal distribution truncated to the kinematic limits
        standard = jax.random.truncated_normal(
            key,
            lower=(min_mass - mass) / width,
            upper=(max_mass - mass) / width,
            shape=(n_events,),
            dtype=jnp.float64,
        )
        return mass + width * standard
    return mass_func

# Change the distribution in the dst_chain dict
dst_chain["D*+"][0]["fs"][0]["D0"][0]["zfit"] = "custom_gauss"

# Pass the custom_gauss function to from_dict
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://cirkiters.github.io/phasespace-jax/GenMultiDecay_Tutorial/>__.

generate(n_events, normalize_weights=True, key=None, **kwargs) #

Generate four-momentum vectors from the decay(s).

Parameters:

Name Type Description Default
n_events int

Total number of events combined, for all the decays.

 required
normalize_weights bool

Normalize weights according to all events generated. This also changes the return values. See the phasespace documentation for more details.

 True
key KeyLike

Either an integer seed, a JAX PRNG key or None, in which case a new key is created from OS entropy (and the generation is not reproducible).

 None
kwargs

Additional parameters passed to all calls of GenParticle.generate

 {}

Returns:

Type Description
tuple[list[Array], list[Array]] | tuple[list[Array], list[Array], list[Array]]

The arguments returned by GenParticle.generate are returned. See the phasespace documentation for
tuple[list[Array], list[Array]] | tuple[list[Array], list[Array], list[Array]]

details. However, instead of being 2 or 3 arrays, it is 2 or 3 lists of arrays,
tuple[list[Array], list[Array]] | tuple[list[Array], list[Array], list[Array]]

each entry in the lists corresponding to the return arguments from the corresponding GenParticle
tuple[list[Array], list[Array]] | tuple[list[Array], list[Array], list[Array]]

instances in self.gen_particles. Note that when normalize_weights is True,
tuple[list[Array], list[Array]] | tuple[list[Array], list[Array], list[Array]]

the weights are normalized to the maximum of all returned events.
Notes

The number of events per decay mode is drawn at random, so each call generally requires a recompilation of the underlying GenParticle.generate calls.

phasespace.fromdecay.mass_functions #

Mass distribution functions for resonant particles.

This module provides factory functions that create mass distribution functions for resonant particles. Each factory returns a callable with the signature (min_mass, max_mass, n_events, key) that samples masses truncated to [min_mass, max_mass] and is usable inside jitted code.

Sampling is done by inverse transform sampling (see e.g. L. Devroye, Non-Uniform Random Variate Generation, Springer 1986, Ch. II), analytically where a closed-form quantile function exists and on a precomputed grid for the relativistic Breit-Wigner, which has none.

gauss_factory(mass, width) #

Create a Gaussian mass distribution function.

Parameters:

Name Type Description Default
mass

Mean mass of the particle.

 required
width

Width (sigma) of the Gaussian distribution.

 required

Returns:

Type Description

Callable that generates masses from a Gaussian distribution truncated to

[min_mass, max_mass], with signature (min_mass, max_mass, n_events, key) and

returning an array of shape (n_events,).

breitwigner_factory(mass, width) #

Create a Breit-Wigner (Cauchy) mass distribution function.

Parameters:

Name Type Description Default
mass

Central mass (m) of the particle.

 required
width

Width (gamma) of the Breit-Wigner distribution.

 required

Returns:

Type Description

Callable that generates masses from a Breit-Wigner distribution truncated to

[min_mass, max_mass], with signature (min_mass, max_mass, n_events, key) and

returning an array of shape (n_events,).
Notes

The Cauchy CDF is :math:F(x) = 1/2 + \arctan((x - m) / \gamma) / \pi, which is inverted analytically to sample within the limits.

relativistic_breitwigner_factory(mass, width) #

Create a relativistic Breit-Wigner mass distribution function.

Parameters:

Name Type Description Default
mass

Central mass (m) of the particle.

 required
width

Width (gamma) of the relativistic Breit-Wigner distribution.

 required

Returns:

Type Description

Callable that generates masses from a relativistic Breit-Wigner distribution truncated to

[min_mass, max_mass], with signature (min_mass, max_mass, n_events, key) and

returning an array of shape (n_events,).
Notes

The density is the constant-width relativistic Breit-Wigner :math:f(m) \propto 1 / ((m^2 - m_0^2)^2 + m_0^2 \Gamma^2) (PDG, Review of Particle Physics, resonance section), matching zfit_physics.pdf.RelativisticBreitWigner.

It has no closed-form quantile function, so the CDF is tabulated once here and inverted by interpolation. The grid is placed at the quantiles of the Cauchy distribution that the density takes in :math:s = m^2, which makes it dense across the peak while still reaching far into the tails.