3. Intensity distribution

Import Python libraries Hide code cell content

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

σ₃ 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()):
    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 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=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(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(fig)

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

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