6. Average polarimeter per resonance

Import Python libraries Hide code cell content

from __future__ import annotations

import logging
import os
from collections import defaultdict
from math import ceil, sqrt
from textwrap import dedent, wrap

import cloudpickle
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
import sympy as sp
import yaml
from ampform.io import aslatex
from ampform.sympy._cache import cache_to_disk
from ampform_dpd import AmplitudeModel
from ampform_dpd.decay import FinalStateID, Particle
from ampform_dpd.io import cached, simplify_latex_rendering
from IPython.display import Latex, Math
from plotly.subplots import make_subplots
from tensorwaves.interface import ParametrizedFunction
from tqdm.auto import tqdm

from polarimetry import formulate_polarimetry
from polarimetry.data import create_data_transformer, generate_phasespace_sample
from polarimetry.function import compute_sub_function
from polarimetry.lhcb import (
    ParameterBootstrap,
    ParameterType,
    flip_production_coupling_signs,
    load_model_builder,
    load_model_parameters,
)
from polarimetry.lhcb.particle import load_particles
from polarimetry.plot import reduce_svg_size

simplify_latex_rendering()
FUNCTION_CACHE: dict[sp.Expr, ParametrizedFunction] = {}
MODEL_FILE = "../data/model-definitions.yaml"
PARTICLES = load_particles("../data/particle-definitions.yaml")

NO_LOG = "EXECUTE_NB" in os.environ
if NO_LOG:
    logging.disable(logging.CRITICAL)

6.1. Computations

Formulate models Hide code cell source

@cache_to_disk(dependencies=["ampform", "ampform-dpd", "polarimetry-lc2pkpi", "sympy"])
def formulate_all_models() -> dict[str, dict[FinalStateID, AmplitudeModel]]:
    with open(MODEL_FILE) as f:
        data = yaml.load(f, Loader=yaml.SafeLoader)
    allowed_model_titles = list(data)
    models = defaultdict(dict)
    for title in tqdm(
        allowed_model_titles,
        desc="Formulate models",
        disable=NO_LOG,
    ):
        builder = load_model_builder(MODEL_FILE, PARTICLES, title)
        imported_parameters = load_model_parameters(
            MODEL_FILE, builder.decay, title, PARTICLES
        )
        for reference_subsystem in (1, 2, 3):
            model = builder.formulate(reference_subsystem)
            model.parameter_defaults.update(imported_parameters)
            models[title][reference_subsystem] = model
        models[title][2] = flip_production_coupling_signs(models[title][2], ["K", "L"])
        models[title][3] = flip_production_coupling_signs(models[title][3], ["K", "D"])
    return {i: dict(dct.items()) for i, dct in models.items()}


MODELS = formulate_all_models()
DEFAULT_MODEL_TITLE = "Default amplitude model"
DEFAULT_MODEL = MODELS[DEFAULT_MODEL_TITLE]
DECAY = DEFAULT_MODEL[1].decay

Unfold symbolic expressions Hide code cell source

@cache_to_disk(dependencies=["ampform", "ampform-dpd", "polarimetry-lc2pkpi", "sympy"])
def unfold_expressions() -> tuple[
    dict[str, sp.Expr],
    dict[str, sp.Expr],
    dict[str, dict[int, sp.Expr]],
]:
    intensity_exprs = {}
    alpha_x_exprs = {}
    alpha_z_exprs = defaultdict(dict)
    for title, ref_models in tqdm(
        MODELS.items(),
        desc="Unfolding expressions",
        disable=NO_LOG,
    ):
        model = ref_models[1]
        intensity_exprs[title] = cached.unfold(model)
        for ref, model in tqdm(ref_models.items(), disable=NO_LOG, leave=False):
            builder = load_model_builder(MODEL_FILE, PARTICLES, model_id=title)
            alpha_x, _, alpha_z = formulate_polarimetry(builder, ref)
            if ref == 1:
                alpha_x_exprs[title] = cached.unfold(alpha_x, model.amplitudes)
            alpha_z_exprs[title][ref] = cached.unfold(alpha_z, model.amplitudes)
    return (
        intensity_exprs,
        alpha_x_exprs,
        {i: dict(dct) for i, dct in alpha_z_exprs.items()},
    )


INTENSITY_EXPRS, ALPHA_X_EXPRS, ALPHA_Z_EXPRS = unfold_expressions()

Convert to numerical functions Hide code cell source

@cache_to_disk(
    dependencies=[
        "ampform",
        "ampform-dpd",
        "jax",
        "numpy",
        "polarimetry-lc2pkpi",
        "sympy",
    ],
    dump_function=cloudpickle.dump,
)
def lambdify_expressions() -> tuple[
    dict[str, ParametrizedFunction],
    dict[str, ParametrizedFunction],
    dict[str, dict[int, ParametrizedFunction]],
    dict[str, dict[int, float]],
]:
    intensity_funcs = {}
    alpha_x_funcs = {}
    alpha_z_funcs = defaultdict(dict)
    original_parameters = defaultdict(dict)
    for title, ref_models in tqdm(
        MODELS.items(),
        desc="Lambdifying",
        disable=NO_LOG,
    ):
        reference_subsystem = 1
        model = ref_models[reference_subsystem]
        intensity_funcs[title] = cached_lambdify(INTENSITY_EXPRS[title], model)
        alpha_x_funcs[title] = cached_lambdify(ALPHA_X_EXPRS[title], model)
        for ref, model in ref_models.items():
            alpha_z_expr = ALPHA_Z_EXPRS[title][ref]
            alpha_z_funcs[title][ref] = cached_lambdify(alpha_z_expr, model)
            str_parameters = {str(k): v for k, v in model.parameter_defaults.items()}
            original_parameters[title][ref] = str_parameters
    return (
        intensity_funcs,
        alpha_x_funcs,
        {i: dict(dct) for i, dct in alpha_z_funcs.items()},
        {i: dict(dct) for i, dct in original_parameters.items()},
    )


def cached_lambdify(expr: sp.Expr, model: AmplitudeModel) -> ParametrizedFunction:
    func = FUNCTION_CACHE.get(expr)
    if func is None:
        func = cached.lambdify(
            expr,
            parameters=model.parameter_defaults,
            backend="jax",
        )
        FUNCTION_CACHE[expr] = func
    str_parameters = {str(k): v for k, v in model.parameter_defaults.items()}
    func.update_parameters(str_parameters)
    return func


(
    INTENSITY_FUNCS,
    ALPHA_X_FUNCS,
    ALPHA_Z_FUNCS,
    ORIGINAL_PARAMETERS,
) = lambdify_expressions()

Generate phase space sample Hide code cell source

N_EVENTS = 100_000
PHSP = generate_phasespace_sample(DECAY, N_EVENTS, seed=0)
for ref in tqdm((1, 2, 3), disable=NO_LOG, leave=False):
    transformer = create_data_transformer(DEFAULT_MODEL[ref], backend="jax")
    PHSP.update(transformer(PHSP))
    del transformer

Compute statistics with parameter bootstrap Hide code cell source

def create_bootstraps(
    n_bootstraps: int,
) -> dict[FinalStateID, dict[str, ParameterType]]:
    bootstraps = {
        ref: ParameterBootstrap(MODEL_FILE, DECAY, DEFAULT_MODEL_TITLE)
        for ref, model in DEFAULT_MODEL.items()
    }
    bootstraps[2] = flip_production_coupling_signs(bootstraps[2], ["K", "L"])
    bootstraps[3] = flip_production_coupling_signs(bootstraps[3], ["K", "D"])
    return {
        i: bootstrap.create_distribution(n_bootstraps, seed=0)
        for i, bootstrap in bootstraps.items()
    }


@cache_to_disk(
    dependencies=[
        "ampform",
        "ampform-dpd",
        "jax",
        "numpy",
        "polarimetry-lc2pkpi",
        "sympy",
    ],
    dump_function=cloudpickle.dump,
)
def compute_statistics_weighted_alpha() -> tuple[
    dict[str, jnp.ndarray], dict[str, jnp.ndarray]
]:
    weighted_αx: dict[str, jnp.ndarray] = {}
    weighted_αz: dict[str, jnp.ndarray] = {}
    resonances = get_resonance_to_reference()
    intensity_func = INTENSITY_FUNCS[DEFAULT_MODEL_TITLE]
    αx_ref1_func = ALPHA_X_FUNCS[DEFAULT_MODEL_TITLE]
    αz_ref1_func = ALPHA_Z_FUNCS[DEFAULT_MODEL_TITLE][1]
    original_parameters_ref1 = ORIGINAL_PARAMETERS[DEFAULT_MODEL_TITLE][1]
    for resonance, ref in tqdm(
        resonances.items(),
        desc="Computing statistics",
        disable=NO_LOG,
    ):
        filter_ = [resonance.latex]
        αz_func = ALPHA_Z_FUNCS[DEFAULT_MODEL_TITLE][ref]
        intensity_func.update_parameters(BOOTSTRAP_PARAMETERS_T[ref])
        αx_ref1_func.update_parameters(BOOTSTRAP_PARAMETERS_T[ref])
        αz_ref1_func.update_parameters(BOOTSTRAP_PARAMETERS_T[ref])
        αz_func.update_parameters(BOOTSTRAP_PARAMETERS_T[1])
        intensities = compute_sub_function(intensity_func, PHSP, filter_)
        αx_ref1_array = compute_sub_function(αx_ref1_func, PHSP, filter_).real
        αz_ref1_array = compute_sub_function(αz_ref1_func, PHSP, filter_).real
        αz_array = compute_sub_function(αz_func, PHSP, filter_).real
        αx_ref1 = compute_weighted_average(αx_ref1_array, intensities)
        αz_ref1 = compute_weighted_average(αz_ref1_array, intensities)
        αz = compute_weighted_average(αz_array, intensities)
        weighted_αx[resonance.name] = jnp.stack([αx_ref1, αz_ref1]).T
        weighted_αz[resonance.name] = αz
        original_parameters = ORIGINAL_PARAMETERS[DEFAULT_MODEL_TITLE][ref]
        intensity_func.update_parameters(original_parameters_ref1)
        αx_ref1_func.update_parameters(original_parameters_ref1)
        αz_ref1_func.update_parameters(original_parameters_ref1)
        αz_func.update_parameters(original_parameters)
    return weighted_αx, weighted_αz


def get_resonance_to_reference() -> dict[Particle, int]:
    subsystem_ids: dict[str, FinalStateID] = dict(K=1, L=2, D=3)
    resonances = [
        c.resonance
        for c in sorted(
            DECAY.chains,
            key=lambda p: (subsystem_ids[p.resonance.name[0]], p.resonance.mass),
        )
    ]
    return {p: subsystem_ids[p.name[0]] for p in resonances}


def compute_weighted_average(v: jnp.ndarray, weights: jnp.ndarray) -> jnp.ndarray:
    return jnp.nansum(v * weights, axis=-1) / jnp.nansum(weights, axis=-1)


N_BOOTSTRAPS = 100
BOOTSTRAP_PARAMETERS = create_bootstraps(N_BOOTSTRAPS)
BOOTSTRAP_PARAMETERS_T = {
    ref: {k: v[None].T for k, v in dct.items()}
    for ref, dct in BOOTSTRAP_PARAMETERS.items()
}
STAT_WEIGHTED_ALPHA_REF1, STAT_WEIGHTED_ALPHA_Z = compute_statistics_weighted_alpha()

Compute systematics from alternative models Hide code cell source

@cache_to_disk(
    dependencies=[
        "ampform",
        "ampform-dpd",
        "jax",
        "numpy",
        "polarimetry-lc2pkpi",
        "sympy",
    ],
    dump_function=cloudpickle.dump,
)
def compute_systematic_weighted_alpha() -> tuple[
    dict[str, jnp.ndarray], dict[str, jnp.ndarray]
]:
    weighted_αx = defaultdict(list)
    weighted_αz = defaultdict(list)
    resonances = get_resonance_to_reference()
    for title in tqdm(
        MODELS,
        disable=NO_LOG,
        desc="Computing systematics",
    ):
        for resonance, ref in tqdm(resonances.items(), disable=NO_LOG, leave=False):
            filter_ = [resonance.latex]
            intensity_func = INTENSITY_FUNCS[title]
            αx_func_ref1 = ALPHA_X_FUNCS[title]
            αz_func_ref1 = ALPHA_Z_FUNCS[title][1]
            αz_func = ALPHA_Z_FUNCS[title][ref]
            intensity_func.update_parameters(ORIGINAL_PARAMETERS[title][1])
            αx_func_ref1.update_parameters(ORIGINAL_PARAMETERS[title][1])
            αz_func_ref1.update_parameters(ORIGINAL_PARAMETERS[title][1])
            αz_func.update_parameters(ORIGINAL_PARAMETERS[title][ref])
            αx_ref1_array = compute_sub_function(αx_func_ref1, PHSP, filter_)
            αz_ref1_array = compute_sub_function(αz_func_ref1, PHSP, filter_)
            αz_array = compute_sub_function(αz_func, PHSP, filter_)
            intensities = compute_sub_function(intensity_func, PHSP, filter_)
            αx_ref1 = compute_weighted_average(αx_ref1_array, intensities).real
            αz_ref1 = compute_weighted_average(αz_ref1_array, intensities).real
            αz = compute_weighted_average(αz_array, intensities).real
            weighted_αx[resonance.name].append((αx_ref1, αz_ref1))
            weighted_αz[resonance.name].append(αz)
    return (
        {k: jnp.array(v) for k, v in weighted_αx.items()},
        {k: jnp.array(v) for k, v in weighted_αz.items()},
    )


SYST_WEIGHTED_ALPHA_REF1, SYST_WEIGHTED_ALPHA_Z = compute_systematic_weighted_alpha()

6.2. Result and comparison

LHCb values are taken from the original study [1]:

Show code cell source Hide code cell source

def tabulate_alpha_z():
    with open("../data/observable-references.yaml") as f:
        data = yaml.load(f, Loader=yaml.SafeLoader)
    lhcb_values = {
        k: tuple(float(v) for v in row.split(" ± "))
        for k, row in data[DEFAULT_MODEL_TITLE]["alpha_z"].items()
        if k != "K(892)"
    }
    model_ids = list(range(len(MODELS)))
    alignment = "r".join("" for _ in model_ids)
    header = " & ".join(Rf"\textbf{{{i}}}" for i in model_ids[1:])
    src = Rf"\begin{{array}}{{l|c|c|{alignment}}}" "\n"
    src += Rf" & \textbf{{this study}} & \textbf{{LHCb}} & {header} \\" "\n"
    src += R"\hline" "\n"
    src = dedent(src)
    for resonance in get_resonance_to_reference():
        src += f" {resonance.latex} & "
        stat_array = 1e3 * STAT_WEIGHTED_ALPHA_Z[resonance.name]
        syst_array = 1e3 * SYST_WEIGHTED_ALPHA_Z[resonance.name]
        syst_diff = syst_array[1:] - syst_array[0]
        value = syst_array[0]
        std = stat_array.std()
        min_ = syst_diff.min()  # LS-model excluded
        max_ = syst_diff.max()  # LS-model excluded
        src += Rf"{value:>+.0f} \pm {std:.0f}_{{{min_:+.0f}}}^{{{max_:+.0f}}}"
        src += " & "
        lhcb = lhcb_values.get(resonance.name)
        if lhcb is not None:
            val, stat, syst, _ = lhcb
            src += Rf"{val:+.0f} \pm {stat:.0f} \pm {syst:.0f}"
        for diff in syst_diff:
            diff_str = f"{diff:>+.0f}"
            if diff == syst_diff.max():
                src += Rf" & \color{{red}}{{{diff_str}}}"
            elif diff == syst_diff.min():
                src += Rf" & \color{{blue}}{{{diff_str}}}"
            else:
                src += f" & {diff_str}"
        src += R" \\" "\n"
    src += R"\end{array}"
    return Latex(src)


tabulate_alpha_z()

\[\begin{split}\begin{array}{l|c|c|rrrrrrrrrrrrrrrrr} & \textbf{this study} & \textbf{LHCb} & \textbf{1} & \textbf{2} & \textbf{3} & \textbf{4} & \textbf{5} & \textbf{6} & \textbf{7} & \textbf{8} & \textbf{9} & \textbf{10} & \textbf{11} & \textbf{12} & \textbf{13} & \textbf{14} & \textbf{15} & \textbf{16} & \textbf{17} \\ \hline K(700) & +63 \pm 78_{-235}^{+238} & +60 \pm 660 \pm 240 & -5 & -14 & -55 & -113 & -100 & +57 & -176 & \color{blue}{-235} & \color{red}{+238} & +96 & +51 & +211 & +52 & +11 & -221 & +12 & +1 \\ K(892) & +29 \pm 15_{-17}^{+31} & & +2 & -0 & +2 & -9 & \color{blue}{-17} & +2 & -5 & +23 & \color{red}{+31} & -8 & +5 & +8 & -3 & -2 & +13 & +1 & +10 \\ K(1430) & -339 \pm 28_{-102}^{+139} & -340 \pm 30 \pm 140 & +3 & +3 & -1 & -2 & +45 & +102 & +125 & -9 & -102 & \color{red}{+139} & -15 & \color{blue}{-102} & +7 & +4 & +6 & -1 & +1 \\ \Lambda(1405) & +580 \pm 31_{-122}^{+278} & -580 \pm 50 \pm 280 & +14 & -7 & +3 & +31 & -3 & -30 & \color{blue}{-122} & -22 & +124 & -64 & +31 & \color{red}{+278} & -17 & -8 & +51 & +0 & +7 \\ \Lambda(1520) & +925 \pm 8_{-84}^{+16} & -925 \pm 25 \pm 84 & +7 & +2 & +2 & \color{red}{+16} & -34 & +2 & +8 & +11 & +7 & -3 & +4 & \color{blue}{-84} & +2 & +1 & -6 & +2 & -10 \\ \Lambda(1600) & +199 \pm 51_{-428}^{+499} & -200 \pm 60 \pm 500 & +10 & -5 & +14 & -5 & +21 & +138 & +100 & \color{red}{+499} & \color{blue}{-428} & -140 & +12 & +72 & -21 & -12 & +13 & -19 & +11 \\ \Lambda(1670) & +817 \pm 16_{-46}^{+73} & -817 \pm 42 \pm 73 & +9 & -10 & +12 & +70 & -41 & -5 & \color{red}{+73} & +30 & +47 & \color{blue}{-46} & +12 & +29 & -5 & -3 & +17 & +3 & -18 \\ \Lambda(1690) & +958 \pm 8_{-35}^{+27} & -958 \pm 20 \pm 27 & -3 & +6 & -12 & \color{blue}{-35} & -14 & +22 & \color{red}{+27} & -20 & +3 & -4 & +5 & +18 & -4 & -1 & -1 & -0 & -9 \\ \Lambda(2000) & -573 \pm 9_{-191}^{+124} & +570 \pm 30 \pm 190 & +9 & -1 & +12 & +47 & -24 & -45 & \color{blue}{-191} & +58 & +85 & +78 & -19 & \color{red}{+124} & +6 & +3 & -9 & -2 & -23 \\ \Delta(1232) & +548 \pm 8_{-27}^{+36} & -548 \pm 14 \pm 36 & +9 & +0 & -9 & -14 & +17 & -1 & +10 & \color{red}{+36} & +5 & -11 & +2 & -8 & -2 & -1 & +12 & -0 & \color{blue}{-27} \\ \Delta(1600) & -502 \pm 9_{-112}^{+162} & +500 \pm 50 \pm 170 & +19 & +10 & +6 & +107 & \color{blue}{-112} & +115 & +88 & +49 & \color{red}{+162} & +5 & +51 & +97 & +16 & +9 & -27 & -3 & -53 \\ \Delta(1700) & +216 \pm 18_{-75}^{+42} & -216 \pm 36 \pm 75 & +40 & +4 & -0 & -19 & -2 & +23 & +16 & \color{red}{+42} & +23 & \color{blue}{-75} & -4 & -2 & +18 & +11 & -3 & +5 & -15 \\ \end{array}\end{split}\]

6.3. Distribution analysis

Show code cell source Hide code cell source

def plot_distributions():
    layout_kwargs = dict(
        font=dict(size=18),
        height=800,
        width=1000,
        xaxis_title="Resonance",
        yaxis_title="<i>ɑ̅<sub>z</sub></i>",
        showlegend=False,
    )
    wrapped_titles = ["<br>".join(wrap(t, width=60)) for t in MODELS]
    colors = dict(  # https://stackoverflow.com/a/44727682
        K="#d62728",  # red
        L="#1f77b4",  # blue
        D="#2ca02c",  # green
    )
    align_left = {
        "K(700)",
        "K(1430)",
        "L(1520)",
        "L(2000)",
        "L(1670)",
        "D(1700)",
    }

    fig = go.Figure()
    for i, (resonance, alpha_z) in enumerate(STAT_WEIGHTED_ALPHA_Z.items()):
        subsystem_id = resonance[0]
        fig.add_trace(
            go.Violin(
                y=alpha_z,
                hovertemplate="<i>ɑ̅<sub>z</sub></i> = %{y:.3f}",
                marker_color=colors[subsystem_id],
                meanline_visible=True,
                name=to_unicode(resonance),
                points="all",
                text=wrapped_titles,
            )
        )
        fig.add_annotation(
            x=i,
            y=jnp.median(alpha_z),
            xshift=-65 if resonance in align_left else +50,
            font_color=colors[subsystem_id],
            font_size=14,
            showarrow=False,
            text=format_average(resonance),
        )
    fig.update_layout(
        title="<b>Statistical</b> distribution of weighted <i>ɑ̅<sub>z</sub></i>",
        **layout_kwargs,
    )
    fig.update_xaxes(tickangle=45)
    plt.show(fig)
    fig.update_layout(font=dict(size=24))
    fig.write_image("_images/alpha-z-per-resonance-statistical.svg")

    fig = go.Figure()
    for i, (resonance, alpha_z) in enumerate(SYST_WEIGHTED_ALPHA_Z.items()):
        subsystem_id = resonance[0]
        fig.add_trace(
            go.Box(
                y=alpha_z,
                boxpoints="all",
                hovertemplate=(
                    "<b>%{text}</b><br>%{x}: <i>ɑ̅<sub>z</sub></i> = %{y:.3f}"
                ),
                marker_color=colors[subsystem_id],
                name=to_unicode(resonance),
                text=wrapped_titles,
                line=dict(color="rgba(0,0,0,0)"),
                fillcolor="rgba(0,0,0,0)",
            )
        )
        fig.add_annotation(
            x=i,
            y=jnp.median(alpha_z),
            xshift=-65 if resonance in align_left else +15,
            font_color=colors[subsystem_id],
            font_size=14,
            showarrow=False,
            text=format_average(resonance),
        )
    fig.update_layout(
        title="<b>Systematics</b> distribution of weighted <i>ɑ̅<sub>z</sub></i>",
        **layout_kwargs,
    )
    fig.update_xaxes(tickangle=45)
    plt.show(fig)
    fig.update_layout(font=dict(size=24))
    fig.write_image("_images/alpha-z-per-resonance-systematics.svg")


def to_unicode(resonance: str) -> str:
    return resonance.replace("L", "Λ").replace("D", "Δ")


def format_average(resonance: str) -> str:
    stat_alpha = 1e3 * STAT_WEIGHTED_ALPHA_Z[resonance]
    syst_alpha = 1e3 * SYST_WEIGHTED_ALPHA_Z[resonance]
    diff = syst_alpha[1:] - syst_alpha[0]
    mean = syst_alpha[0]
    std = stat_alpha.std()
    return f"{diff.max():+.0f}<br><b>{mean:+.0f}</b>±{std:.0f}<br>{diff.min():+.0f}"


plot_distributions()

6.3.1. XZ-correlations

It follows from the definition of \(\vec\alpha\) for a single resonance that:

\[\begin{split} \begin{array}{rcl} \alpha_x &=& \left|\vec\alpha\right| \int I_0 \sin\left(\zeta^0\right) \,\mathrm{d}\tau \big/ \int I_0 \,\mathrm{d}\tau \\ \alpha_z &=& \left|\vec\alpha\right| \int I_0 \cos\left(\zeta^0\right) \,\mathrm{d}\tau \big/ \int I_0 \,\mathrm{d}\tau \end{array} \end{split}\]

This means that the correlation if 100% if \(I_0\) does not change in the bootstrap. This may explain the \(xz\)-correlation observed for \(\overline{\alpha}\) over the complete decay as reported in Average polarimetry values.

Show code cell source Hide code cell source

def round_nested(expression: sp.Expr, n_decimals: int) -> sp.Expr:
    no_sqrt_expr = expression.xreplace({
        node: node.n()
        for node in sp.preorder_traversal(expression)
        if not node.free_symbols
        if isinstance(node, sp.Pow)
        if node.args[1] == 1 / 2
    })
    return no_sqrt_expr.xreplace({
        node: round(node, n_decimals)
        for node in sp.preorder_traversal(no_sqrt_expr)
        if isinstance(node, (float, sp.Float))
    })


resonance_latex = R"\Lambda(2000)"
resonance_couplings = {
    k: 0 if "production" in str(k) and resonance_latex not in str(k) else v
    for k, v in DEFAULT_MODEL[1].parameter_defaults.items()
}
intensity_symbol = sp.Symbol(f"I_{{{resonance_latex}}}")
intensity_expr = round_nested(
    cached.simplify(
        cached.xreplace(INTENSITY_EXPRS[DEFAULT_MODEL_TITLE], resonance_couplings)
    ),
    n_decimals=3,
)
Math(aslatex({intensity_symbol: intensity_expr}))

\[\begin{split}\displaystyle \begin{array}{rcl} I_{\Lambda(2000)} &=& \frac{155.425 \sigma_{2}^{2}}{\left|{\sigma_{2} \left(\sigma_{2} - 3.953\right) + 0.79 i \sqrt{0.445 \sigma_{2}^{2} - \sigma_{2} + 0.18}}\right|^{2}} \\ \end{array}\end{split}\]

Show code cell source Hide code cell source

alpha_symbol = sp.Symbol(Rf"\alpha_{{x,{resonance_latex}}}")
alpha_expr = cached.simplify(
    round_nested(
        cached.simplify(
            cached.xreplace(ALPHA_X_EXPRS[DEFAULT_MODEL_TITLE], resonance_couplings)
        ),
        n_decimals=3,
    ).rewrite(sp.conjugate)
)
alpha_str = str(alpha_expr)
expected_str = R"-0.572*sin(\zeta^0_{2(1)})"
assert alpha_str == expected_str, alpha_str
Math(aslatex({alpha_symbol: alpha_expr}))

\[\begin{split}\displaystyle \begin{array}{rcl} \alpha_{x,\Lambda(2000)} &=& - 0.572 \sin{\left(\zeta^0_{2(1)} \right)} \\ \end{array}\end{split}\]

Show code cell source Hide code cell source

alpha_symbol = sp.Symbol(Rf"\alpha_{{z,{resonance_latex}}}")
alpha_expr = cached.simplify(
    round_nested(
        cached.simplify(
            cached.xreplace(ALPHA_Z_EXPRS[DEFAULT_MODEL_TITLE][1], resonance_couplings)
        ),
        n_decimals=3,
    ).rewrite(sp.conjugate)
)
alpha_str = str(alpha_expr)
expected_str = R"-0.572*cos(\zeta^0_{2(1)})"
assert alpha_str == expected_str, alpha_str
Math(aslatex({alpha_symbol: alpha_expr}))

\[\begin{split}\displaystyle \begin{array}{rcl} \alpha_{z,\Lambda(2000)} &=& - 0.572 \cos{\left(\zeta^0_{2(1)} \right)} \\ \end{array}\end{split}\]

Show code cell source Hide code cell source

def plot_correlation_xz_mpl(stat_or_syst, typ: str) -> None:
    resonances = get_resonance_to_reference()
    n_resonances = len(resonances)
    n_cols = ceil(sqrt(n_resonances))
    n_rows = ceil(n_resonances / n_cols)
    fig, axes = plt.subplots(
        figsize=(12, 8),
        ncols=n_cols,
        nrows=n_rows,
    )
    fig.patch.set_facecolor("none")
    fig.suptitle(typ + R" $\overline{\alpha}_{xz}$-distribution")
    colors = dict(  # https://stackoverflow.com/a/44727682
        K="#d62728",  # red
        L="#1f77b4",  # blue
        D="#2ca02c",  # green
    )
    for ax, resonance in zip(axes.flatten(), resonances):
        subsystem = resonance.name[0]
        color = colors[subsystem]
        αx, αz = stat_or_syst[resonance.name].T
        if αx.std() == 0:
            correlation = 1
            slope = 0
        else:
            correlation = np.corrcoef(αx, αz)[0, 1]
            slope, _ = np.polyfit(αz, αx, deg=1)
        ax.scatter(αz, αx, c=color, s=1)
        ax.set_aspect("equal", adjustable="datalim")
        ax.set_title(f"${resonance.latex}$", y=0.85)
        kwargs = dict(c=color, size=8, transform=ax.transAxes)
        ax.text(0.03, 0.11, f"slope: {slope:+.3g}", **kwargs)
        ax.text(0.03, 0.03, f"correlation: {correlation:+.3g}", **kwargs)
        if ax in axes[-1, :]:
            ax.set_xlabel(R"$\overline{\alpha}_z$")
        if ax in axes[:, 0]:
            ax.set_ylabel(R"$\overline{\alpha}_x$")
    fig.tight_layout()
    output_path = f"_images/alpha-xz-{typ.lower()}.svg"
    fig.savefig(output_path, bbox_inches="tight")
    reduce_svg_size(output_path)
    plt.show(fig)
    plt.close(fig)


plot_correlation_xz_mpl(STAT_WEIGHTED_ALPHA_REF1, "Statistics")
plot_correlation_xz_mpl(SYST_WEIGHTED_ALPHA_REF1, "Systematics")

Show code cell source Hide code cell source

def plot_correlation_xz(stat_or_syst, typ: str) -> None:
    resonances = get_resonance_to_reference()
    n_resonances = len(resonances)
    n_cols = ceil(sqrt(n_resonances))
    n_rows = ceil(n_resonances / n_cols)
    fig = make_subplots(
        cols=n_cols,
        rows=n_rows,
        subplot_titles=[to_unicode(r.name) for r in resonances],
        horizontal_spacing=0.02,
        vertical_spacing=0.08,
    )
    fig.update_layout(
        title_text=(
            "<b>Systematics</b> distribution of weighted <i>ɑ̅<sub>xz</sub></i>"
        ),
        height=800,
        width=1000,
        showlegend=False,
    )
    colors = dict(  # https://stackoverflow.com/a/44727682
        K="#d62728",  # red
        L="#1f77b4",  # blue
        D="#2ca02c",  # green
    )
    wrapped_titles = ["<br>".join(wrap(t, width=60)) for t in MODELS]
    hovertemplate = "<i>ɑ̅<sub>x</sub></i> = %{y:.3f}, <i>ɑ̅<sub>z</sub></i> = %{x:.3f}"
    if "syst" in typ.lower():
        hovertemplate = "<b>%{text}</b><br>" + hovertemplate
    for i, resonance in enumerate(resonances):
        αx, αz = stat_or_syst[resonance.name].T
        subsystem_name = resonance.name[0]
        fig.add_trace(
            go.Scatter(
                x=αz,
                y=αx,
                hovertemplate=hovertemplate,
                marker_color=colors[subsystem_name],
                name=to_unicode(resonance.name),
                text=wrapped_titles,
                mode="markers",
            ),
            col=i % n_cols + 1,
            row=i // n_cols + 1,
        )
    plt.show(fig)
    fig.write_image(f"_images/alpha-xz-{typ.lower()}-plotly.svg")


plot_correlation_xz(STAT_WEIGHTED_ALPHA_REF1, "Statistics")
plot_correlation_xz(SYST_WEIGHTED_ALPHA_REF1, "Systematics")