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.
Import Python libraries
import matplotlib.pyplot as plt
import numpy as np
import phasespace
import tensorflow as tf
from pylorentz import Momentum4
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1730235440.338701 3123 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1730235440.342879 3123 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
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:1730235446.760321 3147 service.cc:148] XLA service 0x7f7854005d70 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1730235446.760362 3147 service.cc:156] StreamExecutor device (0): Host, Default Version
I0000 00:00:1730235446.778005 3149 device_compiler.h:188] 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([[ -6.15677921, -2348.78915304, -1149.20850963, 2618.58901417],
[ 764.40172391, -1030.8068558 , 2278.30108206, 2618.58901417],
[ 1414.71229124, -2031.48127234, 842.14105892, 2618.58901417],
...,
[ 331.91798772, -1805.01740603, 1862.59797148, 2618.58901417],
[ 2020.95505265, -1650.26405055, -172.90988586, 2618.58901417],
[ 1502.44758564, 1358.45676915, -1653.71552836, 2618.58901417]])>,
'p_1': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 6.15677921, 2348.78915304, 1149.20850963, 2661.06098583],
[ -764.40172391, 1030.8068558 , -2278.30108206, 2661.06098583],
[-1414.71229124, 2031.48127234, -842.14105892, 2661.06098583],
...,
[ -331.91798772, 1805.01740603, -1862.59797148, 2661.06098583],
[-2020.95505265, 1650.26405055, 172.90988586, 2661.06098583],
[-1502.44758564, -1358.45676915, 1653.71552836, 2661.06098583]])>}
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.
Show code cell source
fig, ax = plt.subplots(1, 4, tight_layout=True, figsize=(11, 3.5))
ax[0].hist(B0.m, bins=100, color="CornFlowerBlue", range=(5279, 5280))
ax[0].set_xlabel(R"i.m.($\pi$K) [MeV/$c^2$]")
ax[0].set_title("(pion-kaon) i.m. \n")
ax[1].hist(kaon.p_x, bins=70, color="lightcoral", hatch="//")
ax[1].hist(
pion.p_x,
bins=70,
color="springgreen",
histtype="barstacked",
hatch="\\",
alpha=0.5,
)
ax[1].set_xlabel("$p_x$ [MeV/$c$]")
ax[1].set_title("x mom. component \n")
ax[2].hist(kaon.p_y, bins=70, color="lightcoral", hatch="//")
ax[2].hist(
pion.p_y,
bins=70,
color="springgreen",
histtype="barstacked",
hatch="\\",
alpha=0.5,
)
ax[2].set_xlabel("$p_y$ [MeV/$c$]")
ax[2].set_title("y mom. component \n")
ax[3].hist(kaon.p_z, bins=70, color="lightcoral", hatch="//")
ax[3].hist(
pion.p_z,
bins=70,
color="springgreen",
histtype="barstacked",
hatch="\\",
alpha=0.5,
)
ax[3].set_xlabel("$p_z$ [MeV/$c$]")
ax[3].set_title("z mom. component \n")
plt.show()
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
Show code cell source
fig, ax = plt.subplots(1, 3, tight_layout=True, figsize=(12, 3.5))
f0 = ax[0].hist2d(s1, s3, bins=70, cmap="turbo", cmin=1)
fig.colorbar(f0[3], ax=ax[0])
ax[0].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[0].set_ylabel(R"i.m.$^2(\pi^-\pi^0)$ [(MeV/$c^2)^2]$")
f1 = ax[1].hist2d(s2, s3, bins=70, cmap="turbo", cmin=1)
fig.colorbar(f1[3], ax=ax[1])
ax[1].set_xlabel(R"i.m.$^2(\pi^0K^+)$ [(MeV/$c^2)^2]$")
ax[1].set_ylabel(R"i.m.$^2(\pi^-\pi^0)$ [(MeV/$c^2)^2]$")
f2 = ax[2].hist2d(s1, s2, bins=70, cmap="turbo", cmin=1)
fig.colorbar(f2[3], ax=ax[2])
ax[2].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[2].set_ylabel(R"i.m.$^2(\pi^0K^+)$ [(MeV/$c^2)^2]$")
plt.show()
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([[ 571.11959178, -73.84525593, -2498.36131643, 2715.77793272],
[ 754.36468872, 2140.24604025, -1193.19771245, 2715.77793272],
[-2061.34852086, -163.18532728, -1515.80104297, 2715.77793272],
...,
[ 238.92357257, 801.27317841, -2423.69899068, 2715.77793272],
[ -722.4134245 , 1958.77567506, -1488.20585816, 2715.77793272],
[-1721.39819116, 1864.32856242, -366.75230903, 2715.77793272]])>,
'gamma': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ -571.11959178, 73.84525593, 2498.36131643, 2563.87206728],
[ -754.36468872, -2140.24604025, 1193.19771245, 2563.87206728],
[ 2061.34852086, 163.18532728, 1515.80104297, 2563.87206728],
...,
[ -238.92357257, -801.27317841, 2423.69899068, 2563.87206728],
[ 722.4134245 , -1958.77567506, 1488.20585816, 2563.87206728],
[ 1721.39819116, -1864.32856242, 366.75230903, 2563.87206728]])>,
'K+': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 267.43463379, -322.29703013, -1573.88186916, 1701.82777547],
[ 165.12291731, 1048.89816129, -753.72137463, 1392.57524805],
[ -661.78195588, -65.8862658 , -716.06221655, 1094.87827426],
...,
[ 43.91597478, 830.1436363 , -2100.1399382 , 2312.00601429],
[ -191.45672606, 679.16272784, -737.92546862, 1134.09376438],
[-1329.00400946, 1261.54316707, -519.03831292, 1967.45027133]])>,
'pi-': <tf.Tensor: shape=(100000, 4), dtype=float64, numpy=
array([[ 303.68495799, 248.4517742 , -924.47944727, 1013.95015725],
[ 589.24177141, 1091.34787896, -439.47633782, 1323.20268467],
[-1399.56656498, -97.29906148, -799.73882642, 1620.89965846],
...,
[ 195.00759779, -28.87045789, -323.55905247, 403.77191843],
[ -530.95669845, 1279.61294722, -750.28038955, 1581.68416834],
[ -392.3941817 , 602.78539535, 152.2860039 , 748.32766139]])>}
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
Show code cell source
fig, ax = plt.subplots(1, 3, tight_layout=True, figsize=(12, 3.5))
f0 = ax[0].hist2d(s1, s2, bins=70, cmap="turbo")
fig.colorbar(f0[3], ax=ax[0])
ax[0].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[0].set_ylabel(R"i.m.$^2(K^+\gamma)$ [(MeV/$c^2)^2]$")
f1 = ax[1].hist2d(s2, s3, bins=70, cmap="turbo")
fig.colorbar(f1[3], ax=ax[1])
ax[1].set_xlabel(R"i.m.$^2(K^+\gamma)$ [(MeV/$c^2)^2]$")
ax[1].set_ylabel(R"i.m.$^2(\pi^-\gamma)$ [(MeV/$c^2)^2]$")
f2 = ax[2].hist2d(s1, s3, bins=70, cmap="turbo")
fig.colorbar(f2[3], ax=ax[2])
ax[2].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[2].set_ylabel(R"i.m.$^2(\pi^-\gamma)$ [(MeV/$c^2)^2]$")
plt.show()
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
Show code cell source
fig, ax = plt.subplots(1, 3, tight_layout=True, figsize=(12, 3.5))
f0 = ax[0].hist2d(s1, s2, bins=70, cmap="rainbow")
fig.colorbar(f0[3], ax=ax[0])
ax[0].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[0].set_ylabel(R"i.m.$^2(K^+\gamma)$ [(MeV/$c^2)^2]$")
f1 = ax[1].hist2d(s2, s3, bins=70, cmap="rainbow")
fig.colorbar(f1[3], ax=ax[1])
ax[1].set_xlabel(R"i.m.$^2(K^+\gamma)$ [(MeV/$c^2)^2]$")
ax[1].set_ylabel(R"i.m.$^2(\pi^-\gamma)$ [(MeV/$c^2)^2]$")
f2 = ax[2].hist2d(s1, s3, bins=70, cmap="rainbow")
fig.colorbar(f2[3], ax=ax[2])
ax[2].set_xlabel(R"i.m.$^2(\pi^-K^+)$ [(MeV/$c^2)^2]$")
ax[2].set_ylabel(R"i.m.$^2(\pi^-\gamma)$ [(MeV/$c^2)^2]$")
plt.show()
