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

Hide code cell content

import matplotlib.pyplot as plt
import numpy as np
import phasespace
import tensorflow as tf
from pylorentz import Momentum4

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.

Show code cell source

Hide 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

Hide 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([[  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

Show code cell source

Hide 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

Hide 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()

---

[1]

Do you remember the difference between a BW width and a Lorentzian FWHM? The widths you get out of the PDG booklet are FWHM’s.