3. Intensity distribution#

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"]
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)
_images/1cc5b197a02f262c5467fa6dc42830c04051299ec6db6e4b41383f2899e4eb2a.svg

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

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

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:

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)
Hide code cell output
_images/c6e3bd20f7103183eeca4d94ba4669386d3200322d892ad1cda74d2d284ed946.svg
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)
_images/08e54bc798a9627ce0f91472434ce5490572bccfa3d90fa5828c3badfd53b1da.svg

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Λ)
Hide code cell content
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)
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(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)
_images/ac788bec04fe0456a345e8c8fd5870b16721bed9fdc54431a3abd735295081e1.svg
Hide code cell content
def compute_sum_over_decay_rates(decay_rate_matrix) -> float:
    decay_rate_sum = 0.0
    for i in range(len(resonances)):
        for j in range(len(resonances)):
            if j < i:
                continue
            decay_rate_sum += decay_rate_matrix[i, j]
    return decay_rate_sum
np.testing.assert_almost_equal(compute_sum_over_decay_rates(decay_rates), 1.0)

3.4. Dominant decays#

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)
_images/37600392ba06c1de8f88f9be6fbd9d40e73cbef5c263f41285306acf02a75879.svg
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_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)
_images/b0579b9204c6457d104112a778d29d0df7d08c4a671db59bea0343865cf42ae3.svg
np.testing.assert_almost_equal(compute_sum_over_decay_rates(sub_decay_rates), 1.0)