3. Intensity distribution#
Import Python libraries
from __future__ import annotations
import logging
import os
from itertools import product
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from ampform.helicity.naming import natural_sorting
from ampform_dpd.decay import Particle
from ampform_dpd.io import cached
from IPython.display import HTML, Markdown
from matplotlib.patches import Rectangle
from tensorwaves.interface import DataSample
from tqdm.auto import tqdm
from polarimetry.data import (
create_data_transformer,
generate_meshgrid_sample,
generate_phasespace_sample,
generate_sub_meshgrid_sample,
)
from polarimetry.function import (
compute_sub_function,
integrate_intensity,
interference_intensity,
sub_intensity,
)
from polarimetry.io import mute_jax_warnings
from polarimetry.lhcb import load_model
from polarimetry.lhcb.particle import load_particles
from polarimetry.plot import (
add_watermark,
get_contour_line,
reduce_svg_size,
use_mpl_latex_fonts,
)
mute_jax_warnings()
particles = load_particles("../data/particle-definitions.yaml")
model = load_model("../data/model-definitions.yaml", particles, model_id=0)
NO_LOG = "EXECUTE_NB" in os.environ
if NO_LOG:
logging.disable(logging.CRITICAL)
unfolded_intensity_expr = cached.unfold(model)
The complete intensity expression contains 43,198 mathematical operations.
3.1. Definition of free parameters#
free_parameters = {
symbol: value
for symbol, value in model.parameter_defaults.items()
if isinstance(symbol, sp.Indexed)
if "production" in str(symbol)
}
fixed_parameters = {
symbol: value
for symbol, value in model.parameter_defaults.items()
if symbol not in free_parameters
}
subs_intensity_expr = cached.xreplace(unfolded_intensity_expr, fixed_parameters)
After substituting the parameters that are not production couplings, the total intensity expression contains 9,516 operations.
3.2. Distribution#
intensity_func = cached.lambdify(subs_intensity_expr, parameters=free_parameters)
transformer = create_data_transformer(model)
grid_sample = generate_meshgrid_sample(model.decay, resolution=1_000)
grid_sample = transformer(grid_sample)
X = grid_sample["sigma1"]
Y = grid_sample["sigma2"]
Show code cell source
s_labels = {
1: R"$K^{**}: \; \sigma_1=m^2\left(K^-\pi^+\right)$ [GeV$^2$]",
2: R"$\Lambda^{**}: \; \sigma_2=m^2\left(pK^-\right)$ [GeV$^2$]",
3: R"$\Delta^{**}: \; \sigma_3=m^2\left(p\pi^+\right)$ [GeV$^2$]",
}
s1_label, s2_label, s3_label = s_labels.values()
plt.rcdefaults()
use_mpl_latex_fonts()
plt.rc("font", size=21)
fig, ax = plt.subplots(dpi=200, figsize=(8.22, 7), tight_layout=True)
fig.patch.set_facecolor("none")
ax.patch.set_facecolor("none")
ax.set_xlabel(s1_label)
ax.set_ylabel(s2_label)
INTENSITIES = intensity_func(grid_sample)
INTENSITY_INTEGRAL = jnp.nansum(INTENSITIES)
NORMALIZED_INTENSITIES = INTENSITIES / INTENSITY_INTEGRAL
np.testing.assert_almost_equal(jnp.nansum(NORMALIZED_INTENSITIES), 1.0)
mesh = ax.pcolormesh(X, Y, NORMALIZED_INTENSITIES, rasterized=True)
c_bar = fig.colorbar(mesh, ax=ax, pad=0.02)
c_bar.ax.set_ylabel("Normalized intensity")
add_watermark(ax, 0.7, 0.82, fontsize=24)
output_path = "_static/images/intensity-distribution.svg"
fig.savefig(output_path, bbox_inches="tight", dpi=1000)
reduce_svg_size(output_path)
plt.show(fig)
High-resolution image can be downloaded here: intensity-distribution.svg
Comparison with Figure 2 from the original LHCb study [1]:
Show code cell source
def plot_horizontal_intensity(ax) -> None:
ax.set_xlabel("$" + s2_label[s2_label.find("m^2") :])
ax.set_ylabel("$" + s1_label[s1_label.find("m^2") :])
ax.set_xlim(1.79, 4.95)
ax.set_ylim(0.18, 2.05)
add_watermark(ax, 0.7, 0.78, fontsize=18)
mesh = ax.pcolormesh(
grid_sample["sigma2"],
grid_sample["sigma1"],
NORMALIZED_INTENSITIES,
rasterized=True,
)
c_bar = fig.colorbar(mesh, ax=ax, pad=0.02)
c_bar.ax.set_ylabel("Normalized intensity")
plt.ioff()
plt.rcdefaults()
use_mpl_latex_fonts()
plt.rc("font", size=20)
fig, ax = plt.subplots(dpi=200, figsize=(7, 4.9))
fig.patch.set_color("none")
plot_horizontal_intensity(ax)
lhcb_fig2_path = "_static/images/LHCb-PAPER-2022-002-Fig2.svg"
output_path = "_static/images/intensity-distribution-low-res.svg"
left_plot_high_res_path = "_static/images/intensity-distribution.svg"
fig.savefig(output_path, bbox_inches="tight")
fig.savefig(left_plot_high_res_path, bbox_inches="tight", dpi=2000)
reduce_svg_size(output_path)
reduce_svg_size(left_plot_high_res_path)
plt.ion()
plt.close(fig)
HTML(f"""
<table style="width:100%; border-collapse:collapse;">
<tr>
<td style="width:50%; padding-right:5px;">
<img src="{output_path}" style="width:100%; height:auto;">
</td>
<td style="width:50%; padding-left:5px;">
<img src="{lhcb_fig2_path}" style="width:100%; height:auto;">
</td>
</tr>
</table>
""")
|
|
σ₃ vs σ₁ plot for σ₃ projection
grid_sample_31 = generate_meshgrid_sample(
model.decay, resolution=1_000, x_mandelstam=3, y_mandelstam=1
)
grid_sample_31 = transformer(grid_sample_31)
INTENSITIES_31 = intensity_func(grid_sample_31)
INTENSITY_INTEGRAL_31 = jnp.nansum(INTENSITIES_31)
NORMALIZED_INTENSITIES_31 = INTENSITIES_31 / INTENSITY_INTEGRAL_31
np.testing.assert_almost_equal(jnp.nansum(NORMALIZED_INTENSITIES_31), 1.0)
fig, ax = plt.subplots()
fig.patch.set_color("none")
ax.set_xlabel(s3_label)
ax.set_ylabel(s1_label)
mesh = ax.pcolormesh(
grid_sample_31["sigma3"],
grid_sample_31["sigma1"],
NORMALIZED_INTENSITIES_31,
rasterized=True,
)
c_bar = fig.colorbar(mesh, ax=ax)
c_bar.ax.set_ylabel("Normalized intensity")
plt.show(fig)
σ₃ vs σ₁ plot for σ₃ projection
Show code cell source
plt.rcdefaults()
use_mpl_latex_fonts()
plt.rc("font", size=16)
fig, axes = plt.subplots(
ncols=3,
figsize=(12, 3.5),
tight_layout=True,
sharey=True,
)
fig.patch.set_color("none")
ax1, ax2, ax3 = axes
for ax, label in zip(axes, s_labels.values()):
ax.patch.set_color("none")
ax.set_xlabel(label)
ax1.set_ylabel("Normalized intensity (a.u.)")
ax1.set_yticks([])
subsystem_identifiers = ["K", "L", "D"]
subsystem_labels = ["K^{**}", R"\Lambda^{**}", R"\Delta^{**}"]
x, y = X[0], Y[:, 0]
ax1.fill(x, jnp.nansum(NORMALIZED_INTENSITIES, axis=0), alpha=0.3)
ax2.fill(y, jnp.nansum(NORMALIZED_INTENSITIES, axis=1), alpha=0.3)
ax3.fill(y, jnp.nansum(NORMALIZED_INTENSITIES_31, axis=0), alpha=0.3)
original_parameters = dict(intensity_func.parameters)
for label, identifier in zip(subsystem_labels, subsystem_identifiers):
label = f"${label}$"
sub_intensities = (
compute_sub_function(intensity_func, grid_sample, [identifier])
/ INTENSITY_INTEGRAL
)
sub_intensities_31 = (
compute_sub_function(intensity_func, grid_sample_31, [identifier])
/ INTENSITY_INTEGRAL_31
)
ax1.plot(x, jnp.nansum(sub_intensities, axis=0), label=label)
ax2.plot(y, jnp.nansum(sub_intensities, axis=1), label=label)
ax3.plot(y, jnp.nansum(sub_intensities_31, axis=0), label=label)
del sub_intensities
intensity_func.update_parameters(original_parameters)
ax1.set_ylim(0, None)
ax2.set_ylim(0, None)
ax3.set_ylim(0, None)
ax3.legend()
output_path = "_images/intensity-distributions-1D.svg"
fig.savefig(output_path, bbox_inches="tight")
reduce_svg_size(output_path)
plt.show(fig)
3.3. Decay rates#
integration_sample = generate_phasespace_sample(model.decay, n_events=100_000, seed=0)
integration_sample = transformer(integration_sample)
I_tot = integrate_intensity(intensity_func(integration_sample))
I_K = sub_intensity(intensity_func, integration_sample, non_zero_couplings=["K"])
I_Λ = sub_intensity(intensity_func, integration_sample, non_zero_couplings=["L"])
I_Δ = sub_intensity(intensity_func, integration_sample, non_zero_couplings=["D"])
I_ΛΔ = interference_intensity(intensity_func, integration_sample, ["L"], ["D"])
I_KΔ = interference_intensity(intensity_func, integration_sample, ["K"], ["D"])
I_KΛ = interference_intensity(intensity_func, integration_sample, ["K"], ["L"])
np.testing.assert_allclose(I_tot, I_K + I_Λ + I_Δ + I_ΛΔ + I_KΔ + I_KΛ)
Functions for computing the decay rate
def compute_decay_rates(func, integration_sample: DataSample):
decay_rates = np.zeros(shape=(n_resonances, n_resonances))
combinations = list(product(enumerate(resonances), enumerate(resonances)))
progress_bar = tqdm(
desc="Calculating rate matrix",
disable=NO_LOG,
total=(len(combinations) + n_resonances) // 2,
)
I_tot = integrate_intensity(intensity_func(integration_sample))
for (i, resonance1), (j, resonance2) in combinations:
if j < i:
continue
progress_bar.postfix = f"{resonance1.name} × {resonance2.name}"
res1 = resonance1.latex
res2 = resonance2.latex
if res1 == res2:
I_sub = sub_intensity(func, integration_sample, non_zero_couplings=[res1])
else:
I_sub = interference_intensity(func, integration_sample, [res1], [res2])
decay_rates[i, j] = I_sub / I_tot
if i != j:
decay_rates[j, i] = decay_rates[i, j]
progress_bar.update()
progress_bar.close()
return decay_rates
def sort_resonances(resonance: Particle):
KDL = {"L": 1, "D": 2, "K": 3}
return KDL[resonance.name[0]], natural_sorting(resonance.name)
resonances = sorted(
(chain.resonance for chain in model.decay.chains),
key=sort_resonances,
reverse=True,
)
n_resonances = len(resonances)
Show code cell source
def visualize_decay_rates(decay_rates, title=R"Rate matrix for isobars (\%)"):
vmax = jnp.max(jnp.abs(decay_rates))
plt.rcdefaults()
use_mpl_latex_fonts()
plt.rc("font", size=14)
plt.rc("axes", titlesize=24)
plt.rc("xtick", labelsize=16)
plt.rc("ytick", labelsize=16)
fig, ax = plt.subplots(figsize=(9, 9))
fig.patch.set_color("none")
ax.set_title(title)
ax.matshow(jnp.rot90(decay_rates).T, cmap=plt.cm.coolwarm, vmin=-vmax, vmax=+vmax)
resonance_latex = [f"${p.latex}$" for p in resonances]
ax.set_xticks(range(n_resonances))
ax.set_xticklabels(reversed(resonance_latex), rotation=90)
ax.set_yticks(range(n_resonances))
ax.set_yticklabels(resonance_latex)
for i in range(n_resonances):
for j in range(n_resonances):
if j < i:
continue
rate = decay_rates[i, j]
ax.text(
n_resonances - j - 1,
i,
f"{100 * rate:.2f}",
va="center",
ha="center",
)
fig.tight_layout()
return fig
decay_rates = compute_decay_rates(intensity_func, integration_sample)
fig = visualize_decay_rates(decay_rates)
output_path = "_images/rate-matrix.svg"
fig.savefig(output_path, bbox_inches="tight")
reduce_svg_size(output_path)
plt.show(fig)
np.testing.assert_almost_equal(compute_sum_over_decay_rates(decay_rates), 1.0)
3.4. Dominant decays#
Show code cell source
threshold = 0.5
percentage = int(100 * threshold)
I_tot = intensity_func(grid_sample)
plt.rcdefaults()
use_mpl_latex_fonts()
plt.rc("font", size=18)
fig, ax = plt.subplots(figsize=(9.1, 7), sharey=True, tight_layout=True)
ax.set_ylabel(s2_label)
ax.set_xlabel(s1_label)
fig.suptitle(
Rf"Regions where the resonance has a decay ratio of $\geq {percentage}$\%",
y=0.95,
)
fig.patch.set_color("none")
phsp_region = jnp.select(
[I_tot > 0, True],
(1, 0),
)
ax.contour(X, Y, phsp_region, colors=["black"], levels=[0], linewidths=[0.2])
resonances = [c.resonance for c in model.decay.chains]
colors = [plt.cm.rainbow(x) for x in np.linspace(0, 1, len(resonances))]
linestyles = {
"K": "dotted",
"L": "dashed",
"D": "solid",
}
items = list(zip(resonances, colors)) # tqdm requires len
progress_bar = tqdm(
desc="Computing dominant region contours",
disable=NO_LOG,
total=len(items),
)
legend_elements = {}
for resonance, color in items:
progress_bar.postfix = resonance.name
I_sub = compute_sub_function(intensity_func, grid_sample, [resonance.latex])
ratio = I_sub / I_tot
selection = jnp.select(
[jnp.isnan(ratio), ratio < threshold, True],
[0, 0, 1],
)
progress_bar.update()
if jnp.all(selection == 0):
continue
contour_set = ax.contour(
*(X, Y, selection),
colors=[color],
levels=[0],
linestyles=linestyles[resonance.name[0]],
)
contour_set.set_clim(vmin=1, vmax=len(model.decay.chains))
line_collection = get_contour_line(contour_set)
legend_elements[f"${resonance.latex}$"] = line_collection
progress_bar.close()
sub_region_x_range = (0.68, 0.72)
sub_region_y_range = (2.52, 2.60)
region_indicator = Rectangle(
xy=(
sub_region_x_range[0],
sub_region_y_range[0],
),
width=sub_region_x_range[1] - sub_region_x_range[0],
height=sub_region_y_range[1] - sub_region_y_range[0],
edgecolor="black",
facecolor="none",
label="Sub-region",
linewidth=0.5,
zorder=10,
)
ax.add_patch(region_indicator)
legend_elements[region_indicator.get_label()] = region_indicator
leg = ax.legend(
handles=legend_elements.values(),
labels=legend_elements.keys(),
bbox_to_anchor=(1.38, 1),
framealpha=1,
loc="upper right",
)
output_path = "_images/sub-regions.svg"
fig.savefig(output_path, bbox_inches="tight")
reduce_svg_size(output_path)
plt.show(fig)
Show code cell source
sub_sample = generate_sub_meshgrid_sample(
model.decay,
resolution=50,
x_range=sub_region_x_range,
y_range=sub_region_y_range,
)
sub_sample = transformer(sub_sample)
sub_decay_rates = compute_decay_rates(intensity_func, sub_sample)
fig = visualize_decay_rates(sub_decay_rates, title="Rate matrix over sub-region")
output_path = "_images/rate-matrix-sub-region.svg"
fig.savefig(output_path, bbox_inches="tight")
reduce_svg_size(output_path)
plt.show(fig)
np.testing.assert_almost_equal(compute_sum_over_decay_rates(sub_decay_rates), 1.0)