B-matrix extension of polarimeter#
Import python libraries
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 matplotlib import cm
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:
Formulate amplitude models
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}\]
Symbol definitions
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#
Unfold symbolic expressions
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)
Lambdify to JAX
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)
Generate data grid for Dalitz plot
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
Plots#
Show 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("022/b00-is-intensity.png")
plt.show()
Show 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_{{{τ}{p}}}$")
cmap = cm.viridis
else:
Z = B_norm[τ, ρ].real
ax.set_title(f"$B_{{{τ}{p}}} / B_{{00}}$")
cmap = 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("022/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}\]
Function definitions for quiver plot
def plot_field(vx, vy, v_abs, ax, strides=12, cmap=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
Show 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("022/alpha-field.svg")
plt.show()
Show 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("022/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)
Show 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=cm.coolwarm).set_clim(vmin=-1, vmax=+1)
ax2.pcolormesh(X, Y, β_abs.real, cmap=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("022/alpha-beta-comparison.png")
plt.show()