%matplotlib widget
plt.rc("font", size=12)
fig, axes = plt.subplots(figsize=(16, 6), ncols=3, nrows=2)
fig.canvas.toolbar_visible = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
(
(ax_s1, ax_s2, ax_s3),
(ax_t1, ax_t2, ax_t3),
) = axes
for ax in axes[:, 0].flatten():
ax.set_ylabel("Intensity (a.u.)")
for ax in axes[:, 1:].flatten():
ax.set_yticks([])
final_state = DECAY.final_state
ax_s1.set_xlabel(f"$m({final_state[2].latex}, {final_state[3].latex})$")
ax_s2.set_xlabel(f"$m({final_state[1].latex}, {final_state[3].latex})$")
ax_s3.set_xlabel(f"$m({final_state[1].latex}, {final_state[2].latex})$")
ax_t1.set_xlabel(Rf"$\theta({final_state[2].latex}, {final_state[3].latex})$")
ax_t2.set_xlabel(Rf"$\theta({final_state[1].latex}, {final_state[3].latex})$")
ax_t3.set_xlabel(Rf"$\theta({final_state[1].latex}, {final_state[2].latex})$")
fig.suptitle(f"Selected resonances: ${', '.join(resonance_selector.value)}$")
fig.tight_layout()
lines = None
__EXPRS: dict[int, sp.Expr] = {}
def _get_symbol_values(
expr: sp.Expr,
parameters: dict[str, ParameterValue],
selected_resonances: list[str],
) -> dict[sp.Symbol, sp.Rational]:
parameters = {
key: value if any(r in key for r in selected_resonances) else 0
for key, value in parameters.items()
}
return {
s: sp.Rational(parameters[s.name])
for s in expr.free_symbols
if s.name in parameters
}
def _simplify(
intensity_expr: sp.Expr,
parameter_defaults: dict[str, ParameterValue],
variables: dict[sp.Symbol, sp.Expr],
selected_resonances: list[str],
) -> sp.Expr:
parameters = _get_symbol_values(
intensity_expr, parameter_defaults, selected_resonances
)
fixed_variables = {k: v for k, v in variables.items() if not v.free_symbols}
obj = (
intensity_expr,
tuple((k, fixed_variables[k]) for k in sorted(fixed_variables, key=str)),
tuple((k, parameters[k]) for k in sorted(parameters, key=str)),
)
h = get_readable_hash(obj)
if h in __EXPRS:
return __EXPRS[h]
expr = cached.xreplace(intensity_expr, parameters)
expr = cached.xreplace(expr, fixed_variables)
expr = sp.trigsimp(expr)
__EXPRS[h] = expr
return expr
def plot_contributions(**kwargs) -> None:
kwargs.pop("resonance_selector")
kwargs.pop("hide_expressions")
kwargs.pop("simplify_expressions")
selected_resonances = list(resonance_selector.value)
fig.suptitle(f"Selected resonances: ${', '.join(selected_resonances)}$")
dpd_pars = {k: int(v) for k, v in kwargs.items() if k in dpd_func.parameters}
ampform_pars = {
k: int(v) for k, v in kwargs.items() if k in ampform_func.parameters
}
ampform_func.update_parameters(ampform_pars)
dpd_func.update_parameters(dpd_pars)
ampform_intensities = compute_sub_intensities(
ampform_func, ampform_phsp, selected_resonances
)
dpd_intensities = compute_sub_intensities(dpd_func, dpd_phsp, selected_resonances)
s_edges = jnp.linspace(0.98, 1.38, num=50)
amp_values_s1, _ = jnp.histogram(
ampform_phsp["m_12"].real,
bins=s_edges,
weights=ampform_intensities,
)
dpd_values_s1, _ = jnp.histogram(
ampform_phsp["m_12"].real,
bins=s_edges,
weights=dpd_intensities,
)
amp_values_s2, _ = jnp.histogram(
ampform_phsp["m_02"].real,
bins=s_edges,
weights=ampform_intensities,
)
dpd_values_s2, _ = jnp.histogram(
ampform_phsp["m_02"].real,
bins=s_edges,
weights=dpd_intensities,
)
amp_values_s3, _ = jnp.histogram(
ampform_phsp["m_01"].real,
bins=s_edges,
weights=ampform_intensities,
)
dpd_values_s3, _ = jnp.histogram(
ampform_phsp["m_01"].real,
bins=s_edges,
weights=dpd_intensities,
)
t_edges = jnp.linspace(0, jnp.pi, num=50)
amp_values_t1, _ = jnp.histogram(
dpd_phsp["theta_23"].real,
bins=t_edges,
weights=ampform_intensities,
)
dpd_values_t1, _ = jnp.histogram(
dpd_phsp["theta_23"].real,
bins=t_edges,
weights=dpd_intensities,
)
amp_values_t2, _ = jnp.histogram(
dpd_phsp["theta_31"].real,
bins=t_edges,
weights=ampform_intensities,
)
dpd_values_t2, _ = jnp.histogram(
dpd_phsp["theta_31"].real,
bins=t_edges,
weights=dpd_intensities,
)
amp_values_t3, _ = jnp.histogram(
dpd_phsp["theta_23"].real,
bins=t_edges,
weights=ampform_intensities,
)
dpd_values_t3, _ = jnp.histogram(
dpd_phsp["theta_23"].real,
bins=t_edges,
weights=dpd_intensities,
)
global lines
amp_kwargs = dict(color="r", label="ampform", linestyle="solid")
dpd_kwargs = dict(color="blue", label="dpd", linestyle="dotted")
if lines is None:
sx = (s_edges[:-1] + s_edges[1:]) / 2
tx = (t_edges[:-1] + t_edges[1:]) / 2
lines = [
ax_s1.step(sx, amp_values_s1, **amp_kwargs)[0],
ax_s1.step(sx, dpd_values_s1, **dpd_kwargs)[0],
ax_s2.step(sx, amp_values_s2, **amp_kwargs)[0],
ax_s2.step(sx, dpd_values_s2, **dpd_kwargs)[0],
ax_s3.step(sx, amp_values_s3, **amp_kwargs)[0],
ax_s3.step(sx, dpd_values_s3, **dpd_kwargs)[0],
ax_t1.step(tx, amp_values_t1, **amp_kwargs)[0],
ax_t1.step(tx, dpd_values_t1, **dpd_kwargs)[0],
ax_t2.step(tx, amp_values_t2, **amp_kwargs)[0],
ax_t2.step(tx, dpd_values_t2, **dpd_kwargs)[0],
ax_t3.step(tx, amp_values_t3, **amp_kwargs)[0],
ax_t3.step(tx, dpd_values_t3, **dpd_kwargs)[0],
]
ax_s1.legend(loc="upper right")
else:
lines[0].set_ydata(amp_values_s1)
lines[1].set_ydata(dpd_values_s1)
lines[2].set_ydata(amp_values_s2)
lines[3].set_ydata(dpd_values_s2)
lines[4].set_ydata(amp_values_s3)
lines[5].set_ydata(dpd_values_s3)
lines[6].set_ydata(amp_values_t1)
lines[7].set_ydata(dpd_values_t1)
lines[8].set_ydata(amp_values_t2)
lines[9].set_ydata(dpd_values_t2)
lines[10].set_ydata(amp_values_t3)
lines[11].set_ydata(dpd_values_t3)
sy_max = max(
jnp.nanmax(amp_values_s1),
jnp.nanmax(dpd_values_s1),
jnp.nanmax(amp_values_s2),
jnp.nanmax(dpd_values_s2),
jnp.nanmax(amp_values_s3),
jnp.nanmax(dpd_values_s3),
)
ty_max = max(
jnp.nanmax(amp_values_t1),
jnp.nanmax(dpd_values_t1),
jnp.nanmax(amp_values_t2),
jnp.nanmax(dpd_values_t2),
jnp.nanmax(amp_values_t3),
jnp.nanmax(dpd_values_t3),
)
for ax in axes[0]:
ax.set_ylim(0, 1.05 * sy_max)
for ax in axes[1]:
ax.set_ylim(0, 1.05 * ty_max)
fig.canvas.draw_idle()
if hide_expressions.value:
clear_output()
else:
if simplify_expressions.value:
ampform_expr = _simplify(
ampform_intensity_expr,
ampform_pars,
AMPFORM_MODEL.kinematic_variables,
selected_resonances,
)
dpd_expr = _simplify(
dpd_intensity_expr,
dpd_pars,
DPD_MODEL.variables,
selected_resonances,
)
else:
ampform_expr = ampform_intensity_expr
dpd_expr = dpd_intensity_expr
src = Rf"""
\begin{{eqnarray}}
\text{{AmpForm:}} && {sp.latex(ampform_expr)} \\
\text{{DPD:}} && {sp.latex(dpd_expr)} \\
\end{{eqnarray}}
"""
src = dedent(src).strip()
display(Latex(src))
output = interactive_output(
plot_contributions,
controls={
**sliders,
"resonance_selector": resonance_selector,
"hide_expressions": hide_expressions,
"simplify_expressions": simplify_expressions,
},
)
display(output, ui)