7.7. Amplitude model with LS-couplings#
Import python libraries
from __future__ import annotations
import logging
import os
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import sympy as sp
from IPython.display import Latex
from sympy.core.symbol import Str
from tensorwaves.interface import Function
from tqdm.auto import tqdm
from polarimetry.amplitude import (
AmplitudeModel,
get_indexed_base,
simplify_latex_rendering,
)
from polarimetry.data import (
create_data_transformer,
generate_meshgrid_sample,
generate_phasespace_sample,
)
from polarimetry.decay import Particle
from polarimetry.function import integrate_intensity, sub_intensity
from polarimetry.io import (
as_latex,
display_latex,
mute_jax_warnings,
perform_cached_doit,
perform_cached_lambdify,
)
from polarimetry.lhcb import (
get_conversion_factor_ls,
load_model_builder,
load_model_parameters,
)
from polarimetry.lhcb.particle import load_particles
from polarimetry.plot import use_mpl_latex_fonts
mute_jax_warnings()
simplify_latex_rendering()
MODEL_FILE = "../../data/model-definitions.yaml"
PARTICLES = load_particles("../../data/particle-definitions.yaml")
NO_TQDM = "EXECUTE_NB" in os.environ
if NO_TQDM:
logging.getLogger().setLevel(logging.ERROR)
logging.getLogger("polarimetry.io").setLevel(logging.ERROR)
logging.getLogger("tensorwaves.data").setLevel(logging.ERROR)
7.7.1. Model inspection#
Show code cell source
def formulate_model(title: str) -> AmplitudeModel:
builder = load_model_builder(MODEL_FILE, PARTICLES, title)
imported_parameters = load_model_parameters(
MODEL_FILE, builder.decay, title, PARTICLES
)
model = builder.formulate()
model.parameter_defaults.update(imported_parameters)
return model
def simplify_notation(expr: sp.Expr) -> sp.Expr:
def substitute_node(node):
if isinstance(node, sp.Indexed) and node.indices[2:] == (0, 0):
return sp.Indexed(node.base, *node.indices[:2])
return node
for node in sp.preorder_traversal(expr):
new_node = substitute_node(node)
expr = expr.xreplace({node: new_node})
return expr
LS_MODEL = formulate_model("Alternative amplitude model obtained using LS couplings")
simplify_notation(LS_MODEL.intensity.args[0].args[0].args[0].cleanup())
Show code cell source
display_latex({simplify_notation(k): v for k, v in LS_MODEL.amplitudes.items()})
Show conversion factors for LHCb-PAPER-2022-002
H_prod = get_indexed_base("production", min_ls=False)
latex = R"""
\begin{array}{c|ccc|c}
\textbf{Decay} & \textbf{coupling} & & & \textbf{factor} \\
\hline
"""
for chain in LS_MODEL.decay.chains:
R = Str(chain.resonance.name)
L = chain.incoming_ls.L
S = chain.incoming_ls.S
symbol = H_prod[R, L, S]
value = sp.sympify(LS_MODEL.parameter_defaults[symbol])
factor = get_conversion_factor_ls(chain.resonance, L, S)
coupling_value = f"{as_latex(symbol)} &=& {as_latex(value.n(3))}"
latex += Rf" {as_latex(chain)} & {coupling_value} & {factor:+d} \\" "\n"
latex += R"\end{array}"
Latex(f"{latex}")
Show conversion factors for LHCb-PAPER-2022-002
It is asserted that these amplitude expressions to not evaluate to
Show code cell content
def assert_non_zero_amplitudes(model: AmplitudeModel) -> None:
for amplitude in tqdm(model.amplitudes.values(), disable=NO_TQDM):
assert amplitude.doit() != 0
assert_non_zero_amplitudes(LS_MODEL)
See also
See Resonances and LS-scheme for the allowed
7.7.2. Distribution#
Convert expressions to numerical functions
def lambdify(model: AmplitudeModel) -> sp.Expr:
intensity_expr = unfold_intensity(model)
pars = model.parameter_defaults
free_parameters = {s: v for s, v in pars.items() if "production" in str(s)}
fixed_parameters = {s: v for s, v in pars.items() if s not in free_parameters}
subs_intensity_expr = intensity_expr.xreplace(fixed_parameters)
return perform_cached_lambdify(subs_intensity_expr, free_parameters)
def unfold_intensity(model: AmplitudeModel) -> sp.Expr:
unfolded_intensity = perform_cached_doit(model.intensity)
return perform_cached_doit(unfolded_intensity.xreplace(model.amplitudes))
NOMINAL_MODEL = formulate_model("Default amplitude model")
NOMINAL_INTENSITY_FUNC = lambdify(NOMINAL_MODEL)
LS_INTENSITY_FUNC = lambdify(LS_MODEL)
Create phase space grid
GRID = generate_meshgrid_sample(NOMINAL_MODEL.decay, resolution=300)
TRANSFORMER = create_data_transformer(NOMINAL_MODEL)
GRID.update(TRANSFORMER(GRID))
Show code cell source
def compare_2d_distributions() -> None:
NOMINAL_INTENSITIES = compute_normalized_intensity(NOMINAL_INTENSITY_FUNC)
LS_INTENSITIES = compute_normalized_intensity(LS_INTENSITY_FUNC)
max_intensity = max(
jnp.nanmax(NOMINAL_INTENSITIES),
jnp.nanmax(LS_INTENSITIES),
)
use_mpl_latex_fonts()
fig, axes = plt.subplots(
dpi=200,
figsize=(12, 5),
ncols=2,
)
for ax in axes:
ax.set_box_aspect(1)
ax1, ax2 = axes
ax1.set_title("Nominal model")
ax2.set_title("LS-model")
ax1.pcolormesh(
GRID["sigma1"],
GRID["sigma2"],
NOMINAL_INTENSITIES,
vmax=max_intensity,
)
ax2.pcolormesh(
GRID["sigma1"],
GRID["sigma2"],
LS_INTENSITIES,
vmax=max_intensity,
)
plt.show()
def compute_normalized_intensity(func: Function) -> jax.Array:
intensities = func(GRID)
integral = jnp.nansum(intensities)
return intensities / integral
compare_2d_distributions()

7.7.3. Decay rates#
Show code cell source
def to_regex(text: str) -> str:
text = text.replace("(", r"\(")
return text.replace(")", r"\)")
def compute_decay_rates() -> dict[Particle, tuple[float, float]]:
decay_rates = {}
nominal_I_tot = integrate_intensity(NOMINAL_INTENSITY_FUNC(PHSP))
LS_I_tot = integrate_intensity(LS_INTENSITY_FUNC(PHSP))
for chain in tqdm(NOMINAL_MODEL.decay.chains, disable=NO_TQDM):
filter_ = [to_regex(chain.resonance.name)]
LS_I_sub = sub_intensity(LS_INTENSITY_FUNC, PHSP, filter_)
nominal_I_sub = sub_intensity(NOMINAL_INTENSITY_FUNC, PHSP, filter_)
decay_rates[chain.resonance] = (
float(nominal_I_sub / nominal_I_tot),
float(LS_I_sub / LS_I_tot),
)
return decay_rates
PHSP = generate_phasespace_sample(NOMINAL_MODEL.decay, n_events=100_000, seed=0)
PHSP = TRANSFORMER(PHSP)
DECAY_RATES = compute_decay_rates()
src = R"""
\begin{array}{l|rr|r}
\textbf{Resonance} & \textbf{Nominal} & \textbf{LS-model} & \textbf{Difference}\\
\hline
"""
for res, (nominal_rate, ls_rate) in DECAY_RATES.items():
nominal_rate *= 100
ls_rate *= 100
src += (
Rf" {res.latex} & {nominal_rate:.2f} & {ls_rate:.2f} &"
rf" {ls_rate - nominal_rate:+.2f} \\"
"\n"
)
del res, nominal_rate, ls_rate
src += R"\end{array}"
Latex(src)
Tip
Compare with the values with uncertainties as reported in Decay rates.