Lecture 12 – Phase space simulation#
How to perform a phase space simulation with Python?
Prepare the notebook by importing numpy
, matplotlib.pyplot
, and pylorentz
and download the data file (CSV format) from Google Drive.
Two-body decay#
Let’s start with a simple two body decay at rest: \(B^0\rightarrow K^+\pi^-\).
B0_MASS = 5279.65
PION_MASS = 139.57018
KAON_MASS = 493.677
n_events = 100_000
decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, KAON_MASS])
weights, four_momenta = decay.generate(n_events=n_events)
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1760542591.090437 2863 device_compiler.h:196] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process.
The simulation produces a dictionary (four_momenta
) of tf.Tensor
objects. Each object can be addressed with particles['p_i']
, where i
is the number of the \(i\)-th generated particle.
four_momenta
{'p_0': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[-1745.82622074, -221.9171955 , -1934.00417665, 2618.58901417],
[ -269.81856691, 1146.32870346, 2334.66422309, 2618.58901417],
[ 546.26283836, -110.36827271, -2554.78851293, 2618.58901417],
...,
[ 199.44140457, -2081.81330626, 1569.65125871, 2618.58901417],
[-2383.23234941, -980.76217707, 442.53554743, 2618.58901417],
[ -447.4219548 , -619.32727104, 2500.75506911, 2618.58901417]],
shape=(100000, 4))>,
'p_1': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 1745.82622074, 221.9171955 , 1934.00417665, 2661.06098583],
[ 269.81856691, -1146.32870346, -2334.66422309, 2661.06098583],
[ -546.26283836, 110.36827271, 2554.78851293, 2661.06098583],
...,
[ -199.44140457, 2081.81330626, -1569.65125871, 2661.06098583],
[ 2383.23234941, 980.76217707, -442.53554743, 2661.06098583],
[ 447.4219548 , 619.32727104, -2500.75506911, 2661.06098583]],
shape=(100000, 4))>}
Each tf.Tensor
can be converted to a NumPy array, which can then be converted to a pylorentz
.
def to_lorentz(p: tf.Tensor) -> Momentum4:
p = p.numpy().T
return Momentum4(p[3], *p[:3])
pion = to_lorentz(four_momenta["p_0"])
kaon = to_lorentz(four_momenta["p_1"])
These objects can be used to do kinematic computations. Let’s first verify that the invariant mass of the kaon+pion system corresponds to the mass of the mother \(B^0\):
B0 = pion + kaon
np.testing.assert_almost_equal(B0.m.mean(), B0_MASS)
Let’s also plot the momentum components of the two daugther particles.

But it’s monochromatic!! of course it is… it’s a decay at rest. The momentum components are uniformly distributed in the available phase space.
Three-body decay#
Let’s consider now a three body decay like \(B^0\rightarrow K^+\pi^-\pi^0\) and repeat the plot of the relevant kinematic variables. We can also make Dalitz plots this time.
n_events = 50_000
PION0_MASS = 134.9766
decay = phasespace.nbody_decay(B0_MASS, [PION_MASS, PION0_MASS, KAON_MASS])
weights, four_momenta = decay.generate(n_events=n_events)
pim = to_lorentz(four_momenta["p_0"])
pi0 = to_lorentz(four_momenta["p_1"])
kaon = to_lorentz(four_momenta["p_2"])
s1 = (kaon + pim).m2
s2 = (kaon + pi0).m2
s3 = (pim + pi0).m2

Decay chain#
The phasespace
package allows to treat also multiple decays. Let’s consider the \(B^0\rightarrow K^{\ast 0}\gamma\) decay, followed by \(K^{\ast 0}\rightarrow \pi^-K^+\). It can be simulated using the following procedure:
from phasespace import GenParticle
B0_MASS = 5279.65
K0STAR_MASS = 895.55
PION_MASS = 139.57018
KAON_MASS = 493.677
GAMMA_MASS = 0.0
Kp = GenParticle("K+", KAON_MASS)
pim = GenParticle("pi-", PION_MASS)
Kstar = GenParticle("KStar", K0STAR_MASS).set_children(Kp, pim)
gamma = GenParticle("gamma", GAMMA_MASS)
B0 = GenParticle("B0", B0_MASS).set_children(Kstar, gamma)
weights, four_momenta = B0.generate(n_events=100_000)
four_momenta
{'KStar': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 224.3457034 , 2303.56944837, 1103.03072453, 2715.77793272],
[-1458.09857778, 2028.67725078, -576.07041843, 2715.77793272],
[-1463.97556835, -1467.36579601, 1508.99076648, 2715.77793272],
...,
[ -848.46467185, -1733.63669988, 1687.6170984 , 2715.77793272],
[-1174.19436843, 679.73264833, 2175.47031451, 2715.77793272],
[-1124.4603133 , 2206.57977328, -663.35110259, 2715.77793272]],
shape=(100000, 4))>,
'gamma': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ -224.3457034 , -2303.56944837, -1103.03072453, 2563.87206728],
[ 1458.09857778, -2028.67725078, 576.07041843, 2563.87206728],
[ 1463.97556835, 1467.36579601, -1508.99076648, 2563.87206728],
...,
[ 848.46467185, 1733.63669988, -1687.6170984 , 2563.87206728],
[ 1174.19436843, -679.73264833, -2175.47031451, 2563.87206728],
[ 1124.4603133 , -2206.57977328, 663.35110259, 2563.87206728]],
shape=(100000, 4))>,
'K+': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ -88.0955056 , 1001.43028826, 651.73946427, 1295.80274347],
[-1427.41092591, 1602.21153518, -520.65570317, 2262.60537798],
[ -957.36077871, -800.95531732, 1234.90674262, 1823.94646968],
...,
[ -878.71780555, -1677.80691895, 1644.70285654, 2556.54953124],
[ -859.00314063, 736.36481648, 1407.79579792, 1872.35827969],
[-1077.14550204, 2137.89436766, -519.94303874, 2498.97829196]],
shape=(100000, 4))>,
'pi-': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 312.441209 , 1302.13916011, 451.29126026, 1419.97518925],
[ -30.68765187, 426.4657156 , -55.41471526, 453.17255475],
[-506.61478964, -666.41047869, 274.08402387, 891.83146305],
...,
[ 30.2531337 , -55.82978093, 42.91424186, 159.22840148],
[-315.1912278 , -56.63216816, 767.67451659, 843.41965303],
[ -47.31481126, 68.68540563, -143.40806386, 216.79964076]],
shape=(100000, 4))>}
gamma = to_lorentz(four_momenta["gamma"])
pion = to_lorentz(four_momenta["pi-"])
kaon = to_lorentz(four_momenta["K+"])
Kstar = to_lorentz(four_momenta["KStar"])
Let’s build the Dalitz plots matching particle pairs. The particles measured in the final state are \(K^-,\; \pi^-\) and \(\gamma\).
s1 = (pion + kaon).m2
s2 = (gamma + kaon).m2
s3 = (gamma + pion).m2

Width distribution#
These distributions aren’t so interesting, because the masses of each particle are one fixed value. So let’s simulate a more realistic \(K^\ast\) particle; not monochromatic, but with a width of 47 MeV.[1] The mass is extracted from a Gaussian distribution centered at the B0_MASS value and with \(\sigma = 47/2.36 \sim 20\) MeV. See more info on how to do this with the phasespace
package here.
import tensorflow as tf
import tensorflow_probability as tfp
K0STAR_WIDTH = 47 / 2.36
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_mass_cast = tf.cast(K0STAR_MASS, dtype=tf.float64)
tf.cast(K0STAR_WIDTH, tf.float64)
tf.broadcast_to(kstar_mass_cast, shape=(n_events,))
return tfp.distributions.TruncatedNormal(
loc=K0STAR_MASS,
scale=K0STAR_WIDTH,
low=min_mass,
high=max_mass,
).sample()
K = GenParticle("K+", KAON_MASS)
pion = GenParticle("pi-", PION_MASS)
Kstar = GenParticle("KStar", kstar_mass).set_children(K, pion)
gamma = GenParticle("gamma", GAMMA_MASS)
B0 = GenParticle("B0", B0_MASS).set_children(Kstar, gamma)
weights, four_momenta = B0.generate(n_events=100_000)
gamma = to_lorentz(four_momenta["gamma"])
pion = to_lorentz(four_momenta["pi-"])
kaon = to_lorentz(four_momenta["K+"])
Kstar = to_lorentz(four_momenta["KStar"])
Now you have all the 4-vectors to plot the invariant mass distributions for the different steps of the decay chains.
s1 = (pion + kaon).m2
s2 = (gamma + kaon).m2
s3 = (gamma + pion).m2
