B-matrix extension of polarimeter

B-matrix extension of polarimeter#

Hide code cell content

 
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 matplotlib_inline.backend_inline import set_matplotlib_formats
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
set_matplotlib_formats("svg")

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

 
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, rasterized=True)
fig.savefig("b00-is-intensity.png")
plt.show()

Hide code cell source

 
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 = "coolwarm"
        mesh = ax.pcolormesh(X, Y, Z, cmap=cmap, rasterized=True)
        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

 
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

 
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

 
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$")
style = dict(cmap="coolwarm", rasterized=True)
ax1.pcolormesh(X, Y, α_abs.real, **style).set_clim(vmin=-1, vmax=+1)
ax2.pcolormesh(X, Y, β_abs.real, **style).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