3. Intensity distribution

Import Python libraries

Hide code cell content

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 matplotlib_inline.backend_inline import set_matplotlib_formats
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()
set_matplotlib_formats("svg")
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 28,430 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 10,252 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

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

W0302 14:43:37.771299    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

High-resolution image can be downloaded here: intensity-distribution.svg

Comparison with Figure 2 from the original LHCb study [1]:

Show code cell source

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

Tip

High-resolution versions of these plots:

• This study (left)

• Original amplitude analysis (right)

σ₃ vs σ₁ plot for σ₃ projection

Hide code cell source

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

σ₃ vs σ₁ plot for σ₃ projection

Hide code cell output

Show code cell source

Hide 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(), strict=True):
    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, strict=True):
    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()

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

W0302 14:44:08.052843    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

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

Hide code cell content

def compute_decay_rates(rates: np.ndarray) -> np.ndarray:
    m, n = rates.shape
    assert m == n
    d = rates.diagonal()
    D = d * np.identity(n)
    X = d[None] + d[None].T
    return rates - X + 2 * D


def compute_sub_intensities(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 sub-intensities",
        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
        I_sub = sub_intensity(func, integration_sample, non_zero_couplings=[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,
)
n_resonances = len(resonances)

Show code cell source

Hide 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(decay_rates, cmap="coolwarm", vmin=-vmax, vmax=+vmax)

    resonance_latex = [f"${p.latex}$" for p in resonances]
    ax.set_xticks(range(n_resonances))
    ax.set_xticklabels(resonance_latex, rotation=90)
    ax.xaxis.tick_bottom()
    ax.set_yticks(range(n_resonances))
    ax.set_yticklabels(resonance_latex)
    for i in range(n_resonances):
        for j in range(n_resonances):
            if i > j:
                continue
            rate = decay_rates[i, j]
            ax.text(i, j, f"{100 * rate:.2f}", ha="center", va="center")
    fig.tight_layout()
    return fig


sub_intensities = compute_sub_intensities(intensity_func, integration_sample)
decay_rates = compute_decay_rates(sub_intensities)
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()

W0302 14:44:19.536454    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:44:40.070241    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:44:58.914858    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:45:15.817522    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:45:31.070588    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:45:44.207164    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:45:55.490093    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:46:05.260914    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:46:12.923665    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:46:18.818706    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:46:22.739001    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

W0302 14:46:24.742268    8339 pjrt_executable.cc:638] Assume version compatibility. PjRt-IFRT does not track XLA executable versions.

Show code cell source

Hide code cell source

np.testing.assert_array_almost_equal(
    100 * decay_rates[-1],
    [4.78, 0.16, -1.68, 0.04, 0.03, -7.82, 5.09, 1.96, -1.15, -1.95, 0.04, 14.7],
    decimal=2,
)

def compute_sum_over_decay_rates(rate_matrix: np.ndarray) -> float:
    return rate_matrix.diagonal().sum() + np.tril(rate_matrix, k=-1).sum()


np.testing.assert_almost_equal(compute_sum_over_decay_rates(decay_rates), 1.0)

3.4. Dominant decays

Show code cell source

Hide 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, strict=True))  # 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()

Show code cell source

Hide 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_intensities = compute_sub_intensities(intensity_func, sub_sample)
sub_decay_rates = compute_decay_rates(sub_decay_intensities)
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()

np.testing.assert_almost_equal(compute_sum_over_decay_rates(sub_decay_rates), 1.0)