B-matrix extension of polarimeter

B-matrix extension of polarimeter#

Hide code cell content 
from __future__ import annotations

import logging
from pathlib import Path
from warnings import filterwarnings

import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import polarimetry
import sympy as sp
from ampform.sympy import PoolSum
from IPython.display import display
from polarimetry import _to_index
from polarimetry.data import create_data_transformer, generate_meshgrid_sample
from polarimetry.io import (
    mute_jax_warnings,
    perform_cached_doit,
    perform_cached_lambdify,
)
from polarimetry.lhcb import load_model_builder, load_model_parameters
from polarimetry.lhcb.particle import load_particles
from sympy.physics.matrices import msigma
from tqdm.auto import tqdm

filterwarnings("ignore")
logging.getLogger("polarimetry.function").setLevel(logging.INFO)
mute_jax_warnings()
POLARIMETRY_DIR = Path(polarimetry.__file__).parent

Formulate expressions#

Reference subsystem 1 is defined as:

orientation-K.svg

Hide code cell source 
MODEL_CHOICE = 0
MODEL_FILE = POLARIMETRY_DIR / "lhcb/model-definitions.yaml"
PARTICLES = load_particles(POLARIMETRY_DIR / "lhcb/particle-definitions.yaml")
BUILDER = load_model_builder(MODEL_FILE, PARTICLES, model_id=MODEL_CHOICE)
IMPORTED_PARAMETER_VALUES = load_model_parameters(
    MODEL_FILE, BUILDER.decay, MODEL_CHOICE, PARTICLES
)
REFERENCE_SUBSYSTEM = 1
MODEL = BUILDER.formulate(REFERENCE_SUBSYSTEM, cleanup_summations=True)
MODEL.parameter_defaults.update(IMPORTED_PARAMETER_VALUES)
\[\begin{split} \vec\alpha = \sum_{\nu',\nu,\lambda} A^*_{\nu',\lambda}\vec\sigma_{\nu',\nu} A_{\nu,\lambda} / I_0 \\ \vec\beta = \sum_{\nu,\lambda',\lambda} A^*_{\nu,\lambda'} \vec\sigma_{\lambda',\lambda} A^*_{\nu,\lambda} / I_0 \\ B_{\tau,\rho} = \sum_{\nu,\nu',\lambda',\lambda} A^*_{\nu',\lambda'} \sigma_{\nu',\nu}^\tau A_{\nu,\lambda} \sigma_{\lambda',\lambda}^\rho \end{split}\]
Hide code cell source 
half = sp.Rational(1, 2)
λ, λp = sp.symbols(R"lambda \lambda^{\prime}", rational=True)
v, vp = sp.symbols(R"nu \nu^{\prime}", rational=True)
σ = [sp.Matrix([[1, 0], [0, 1]])]
σ.extend(msigma(i) for i in (1, 2, 3))

ref = REFERENCE_SUBSYSTEM
B = tuple(
    tuple(
        PoolSum(
            BUILDER.formulate_aligned_amplitude(vp, λp, 0, 0, ref)[0].conjugate()
            * σ[τ][_to_index(vp), _to_index(v)]
            * BUILDER.formulate_aligned_amplitude(v, λ, 0, 0, ref)[0]
            * σ[ρ][_to_index(λp), _to_index(λ)],
            (v, [-half, +half]),
            (vp, [-half, +half]),
            (λ, [-half, +half]),
            (λp, [-half, +half]),
        ).cleanup()
        for ρ in range(4)
    )
    for τ in range(4)
)
del ref
B = sp.Matrix(B)

Functions and data#

Hide code cell content 
progress_bar = tqdm(desc="Unfolding expressions", total=16)
B_exprs = []
for τ in range(4):
    row = []
    for ρ in range(4):
        expr = perform_cached_doit(B[τ, ρ].doit().xreplace(MODEL.amplitudes))
        progress_bar.update()
        row.append(expr)
    B_exprs.append(row)
progress_bar.close()
B_exprs = np.array(B_exprs)
B_exprs.shape

(4, 4)
Hide code cell content 
progress_bar = tqdm(desc="Lambdifying", total=16)
B_funcs = []
for τ in range(4):
    row = []
    for ρ in range(4):
        func = perform_cached_lambdify(
            B_exprs[τ, ρ].xreplace(MODEL.parameter_defaults),
            backend="jax",
        )
        progress_bar.update()
        row.append(func)
    B_funcs.append(row)
progress_bar.close()
B_funcs = np.array(B_funcs)
Hide code cell content 
transformer = create_data_transformer(MODEL)
GRID_SAMPLE = generate_meshgrid_sample(MODEL.decay, resolution=400)
GRID_SAMPLE.update(transformer(GRID_SAMPLE))
X = GRID_SAMPLE["sigma1"]
Y = GRID_SAMPLE["sigma2"]
del transformer
Hide code cell content 
B_arrays = jnp.array([[B_funcs[τ, ρ](GRID_SAMPLE) for ρ in range(4)] for τ in range(4)])
B_norm = B_arrays / B_arrays[0, 0]
B_arrays.shape

(4, 4, 400, 400)

Plots#

Hide code cell source 
%config InlineBackend.figure_formats = ['png']

plt.rcdefaults()
plt.rc("font", size=16)
s1_label = R"$m^2\left(K^-\pi^+\right)$ [GeV$^2$]"
s2_label = R"$m^2\left(pK^-\right)$ [GeV$^2$]"
fig, ax = plt.subplots(figsize=(8, 6.8))
ax.set_title("$I_0 = B_{0, 0}$")
ax.set_xlabel(s1_label)
ax.set_ylabel(s2_label)
ax.set_box_aspect(1)
ax.pcolormesh(X, Y, B_arrays[0, 0].real)
fig.savefig("b00-is-intensity.png")
plt.show()

Hide code cell source 
%config InlineBackend.figure_formats = ['png']

plt.rcdefaults()
plt.rc("font", size=10)
fig, axes = plt.subplots(
    dpi=200,
    figsize=(11, 10),
    ncols=4,
    nrows=4,
    sharex=True,
    sharey=True,
)
fig.suptitle(
    R"$B_{\tau,\rho} = \sum_{\nu,\nu',\lambda',\lambda} A^*_{\nu',\lambda'}"
    R" \sigma_{\nu',\nu}^\tau A_{\nu,\lambda} \sigma_{\lambda',\lambda}^\rho$"
)
progress_bar = tqdm(total=16)
for ρ in range(4):
    for τ in range(4):
        ax = axes[τ, ρ]
        ax.set_box_aspect(1)
        if τ == 0 and ρ == 0:
            Z = B_arrays[τ, ρ].real
            ax.set_title(f"$B_{{{τ}{ρ}}}$")
            cmap = plt.cm.viridis
        else:
            Z = B_norm[τ, ρ].real
            ax.set_title(f"$B_{{{τ}{ρ}}} / B_{{00}}$")
            cmap = plt.cm.coolwarm
        mesh = ax.pcolormesh(X, Y, Z, cmap=cmap)
        cbar = fig.colorbar(mesh, ax=ax, fraction=0.047, pad=0.01)
        if τ != 0 or ρ != 0:
            mesh.set_clim(vmin=-1, vmax=+1)
            cbar.set_ticks([-1, 0, +1])
            cbar.set_ticklabels(["-1", "0", "+1"])
        if τ == 3:
            ax.set_xlabel(s1_label)
        if ρ == 0:
            ax.set_ylabel(s2_label)
        progress_bar.update()
progress_bar.close()
fig.tight_layout()
fig.savefig("b-matrix-elements.png")
plt.show()

Hypothesis:

\[\begin{split} B_{0,\rho} = \vec\beta B_{00} \\ B_{\tau,0} = \vec\alpha B_{00} \\ B_{00} = I_0 \end{split}\]
Hide code cell content 
def plot_field(vx, vy, v_abs, ax, strides=12, cmap=plt.cm.viridis_r):
    mesh = ax.quiver(
        X[::strides, ::strides],
        Y[::strides, ::strides],
        vx[::strides, ::strides].real,
        vy[::strides, ::strides].real,
        v_abs[::strides, ::strides],
        cmap=cmap,
    )
    mesh.set_clim(vmin=0, vmax=+1)
    return mesh


def plot(x, y, z, strides=14):
    plt.rcdefaults()
    plt.rc("font", size=16)
    fig, ax = plt.subplots(figsize=(8, 6.8), tight_layout=True)
    ax.set_box_aspect(1)
    v_abs = jnp.sqrt(x.real**2 + y.real**2 + z.real**2)
    mesh = plot_field(x, y, v_abs, ax, strides)
    color_bar = fig.colorbar(mesh, ax=ax, pad=0.01)
    return fig, ax, color_bar
Hide code cell source 
%config InlineBackend.figure_formats = ['svg']

fig, ax, cbar = plot(
    x=B_norm[3, 0],
    y=B_norm[1, 0],
    z=B_norm[2, 0],
    strides=10,
)
ax.set_title(
    R"$B_{\tau, 0} / B_{00} = \sum_{\nu',\nu,\lambda}"
    R" A^*_{\nu',\lambda}\vec\sigma_{\nu',\nu} A_{\nu,\lambda} / I_0$"
)
ax.set_xlabel(Rf"{s1_label}, $\quad\alpha_z$")
ax.set_ylabel(Rf"{s2_label}, $\quad\alpha_x$")
cbar.set_label(R"$\left|\vec{\alpha}\right|$")
fig.savefig("alpha-field.svg")
plt.show()
Hide code cell source 
%config InlineBackend.figure_formats = ['svg']

fig, ax, cbar = plot(
    x=B_norm[0, 3],
    y=B_norm[0, 1],
    z=B_norm[0, 2],
    strides=10,
)
ax.set_title(
    R"$B_{0,\rho} / B_{00} = \sum_{\nu,\lambda',\lambda} A^*_{\nu,\lambda'}"
    R" \vec\sigma_{\lambda',\lambda} A^*_{\nu,\lambda} / I_0$"
)
ax.set_xlabel(Rf"{s1_label}, $\quad \beta_z = B_{{03}}$")
ax.set_ylabel(Rf"{s2_label}, $\quad \beta_x = B_{{01}}$")
cbar.set_label(R"$\left|\vec{\beta}\right|$")
fig.savefig("beta-field.svg")
plt.show()

Note that \(|\alpha| = |\beta|\):

α_abs = jnp.sqrt(jnp.sum(B_norm[1:, 0] ** 2, axis=0))
β_abs = jnp.sqrt(jnp.sum(B_norm[0, 1:] ** 2, axis=0))
np.testing.assert_allclose(α_abs, β_abs, rtol=1e-14)
Hide code cell source 
%config InlineBackend.figure_formats = ['png']

fig, axes = plt.subplots(
    figsize=(11, 6),
    ncols=2,
    sharey=True,
    tight_layout=True,
)
for ax in axes:
    ax.set_box_aspect(1)
ax1, ax2 = axes
ax1.set_title(R"$\alpha$")
ax2.set_title(R"$\beta$")
ax1.pcolormesh(X, Y, α_abs.real, cmap=plt.cm.coolwarm).set_clim(vmin=-1, vmax=+1)
ax2.pcolormesh(X, Y, β_abs.real, cmap=plt.cm.coolwarm).set_clim(vmin=-1, vmax=+1)
ax1.set_xlabel(s1_label)
ax2.set_xlabel(s1_label)
ax1.set_ylabel(s2_label)
fig.savefig("alpha-beta-comparison.png")
plt.show()
https://github.com/ComPWA/compwa.github.io/assets/29308176/c7268301-11c9-45f2-a5ec-4c2928352a68