Determination of polarization¶

def compute_polarized_intensity(
    Px: float,
    Py: float,
    Pz: float,
    I0: jnp.ndarray,
    alpha: jnp.ndarray,
    phsp: DataSample,
) -> jnp.array:
    P = jnp.array([Px, Py, Pz])
    R = compute_rotation_matrix(phsp)
    return I0 * (1 + jnp.einsum("i,ij...,j...->...", P, R, alpha))


def compute_rotation_matrix(phsp: DataSample) -> jnp.ndarray:
    ϕ = phsp["phi"]
    θ = jnp.arccos(phsp["cos_theta"])
    χ = phsp["chi"]
    return jnp.einsum("ki...,ij...,j...->k...", Rz(ϕ), Ry(θ), Rz(χ))


def Rz(angle: jnp.ndarray) -> jnp.ndarray:
    n_events = len(angle)
    ones = jnp.ones(n_events)
    zeros = jnp.zeros(n_events)
    return jnp.array(
        [
            [+jnp.cos(angle), -jnp.sin(angle), zeros],
            [+jnp.sin(angle), +jnp.cos(angle), zeros],
            [zeros, zeros, ones],
        ]
    )


def Ry(angle: jnp.ndarray) -> jnp.ndarray:
    n_events = len(angle)
    ones = jnp.ones(n_events)
    zeros = jnp.zeros(n_events)
    return jnp.array(
        [
            [+jnp.cos(angle), zeros, +jnp.sin(angle)],
            [zeros, ones, zeros],
            [-jnp.sin(angle), zeros, +jnp.cos(angle)],
        ]
    )

def extract_polarizations(fit_results: list[iminuit.Minuit]) -> np.ndarray:
    return np.array([extract_polarization(fit) for fit in fit_results])


def extract_polarization(fit_result: iminuit.Minuit) -> tuple[float, float, float]:
    return tuple(p.value for p in fit_result.params)


def render_fit_results(
    fit_results: list[iminuit.Minuit],
    compare: bool = False,
) -> str:
    P_syst = 100 * extract_polarizations(fit_results)
    P_nominal = P_syst[0]
    P_max = (P_syst[1:] - P_nominal).max(axis=0)
    P_min = (P_syst[1:] - P_nominal).min(axis=0)
    if compare:
        np.testing.assert_array_almost_equal(P_nominal, 100 * np.array(P), decimal=2)

    def render_p(i: int) -> str:
        return f"{P_nominal[i]:+.2f}_{{{P_min[i]:+.2f}}}^{{{P_max[i]:+.2f}}}"

    src = Rf"""
    \begin{{array}}{{ccc}}
    P_x &=& {render_p(0)} \\
    P_y &=& {render_p(1)} \\
    P_z &=& {render_p(2)} \\
    \end{{array}}
    """
    return dedent(src).strip()


src = Rf"""
The polarization $\vec{{P}}$ is determined to be (in %):

$$
{render_fit_results(SYST_FIT_RESULTS_FIELD, compare=True)}
$$

with the upper and lower sign being the systematic extrema uncertainties as determined by
the alternative models.
"""
Markdown(src)

def get_averaged_polarimeters(polar: bool = False) -> jnp.ndarray:
    with open("_static/export/averaged-polarimeter-vectors.json") as f:
        json_data = json.load(f)
    data = json_data["systematics"]
    typ = "polar" if polar else "cartesian"
    items = list("xyz")
    if polar:
        items = ("norm", "theta", "phi")
    return jnp.array([data[typ][i] for i in items]).T


def compute_differential_decay_rate(
    Px: float,
    Py: float,
    Pz: float,
    averaged_alpha: jnp.array,
    phsp: DataSample,
) -> jnp.array:
    P = jnp.array([Px, Py, Pz])
    R = compute_rotation_matrix(phsp)
    return 1 + jnp.einsum("i,ij...,j...->...", P, R, averaged_alpha)


SYST_AVERAGED_POLARIMETERS = get_averaged_polarimeters()
SYST_POLAR_POLARIMETERS = get_averaged_polarimeters(polar=True)
assert SYST_AVERAGED_POLARIMETERS.shape == (18, 3)
assert SYST_POLAR_POLARIMETERS.shape == (18, 3)

diff_rate = compute_differential_decay_rate(*P, SYST_AVERAGED_POLARIMETERS[0], PHSP)
assert diff_rate.shape == (N_EVENTS,)
del diff_rate

def get_alpha_systematics(
    all_values: jnp.ndarray,
) -> tuple[tuple[float, float], tuple[float, float], tuple[float, float]]:
    central = all_values[0]
    syst = np.abs(all_values - central).max(axis=0)
    return tuple((μ, σ) for μ, σ in zip(central.tolist(), syst.tolist()))


def render_min_max_averaged_polarimeter() -> Latex:
    cartesian = get_alpha_systematics(SYST_AVERAGED_POLARIMETERS)
    polar = get_alpha_systematics(SYST_POLAR_POLARIMETERS)
    src = R"""
    \begin{array}{c|r|c}
      \textbf{observable} & \textbf{central} & \textbf{stat} + \textbf{syst} \\
      \hline
    """
    src = dedent(src)
    for xyz, (central, systematic) in zip("xyz", cartesian):
        src += Rf"  \overline{{\alpha}}_{xyz} \; \left[10^{{-3}}\right]"
        src += Rf"  & {1e3*central:+6.1f} & {1e3*systematic:4.1f}"
        src += R" \\" "\n"
    src += R"  \hline" "\n"
    polar_labels = [
        R"\left|\overline{\alpha}\right|",
        R"\theta(\overline{\alpha})",
        R"\phi(\overline{\alpha})",
    ]
    for label, (central, systematic) in zip(polar_labels, polar):
        factor = "10^{-3}" if "left" in label else R"\pi"
        src += Rf"  {label:30s} \; \left[{factor:7s}\right]"
        if "left" in label:
            src += Rf"  & {1e3*central:6.1f} & {1e3*systematic:5.1f}"
        else:
            src += Rf"  & {central/π:+6.3f} & {systematic/π:5.3f}"
        src += R" \\" "\n"
    src += R"\end{array}"
    return Latex(src.strip())


render_min_max_averaged_polarimeter()

def polar_to_cartesian(
    r: float, theta: float, phi: float
) -> tuple[float, float, float]:
    return (
        r * np.sin(theta) * np.cos(phi),
        r * np.sin(theta) * np.sin(phi),
        r * np.cos(theta),
    )


def perform_combinatorics_fit(
    alpha_array: jnp.ndarray, polar: bool
) -> tuple[list[tuple[float, float, float]], list[tuple[float, float, float]]]:
    alpha_with_syst = get_alpha_systematics(alpha_array)
    alpha_combinations = tuple((μ - σ, μ + σ) for μ, σ in alpha_with_syst)
    alphas = []
    polarizations = []
    items = list(itertools.product(*alpha_combinations))
    for averaged_alpha in tqdm(items):
        alphas.append(averaged_alpha)
        if polar:
            averaged_alpha = polar_to_cartesian(*averaged_alpha)
        fit_result = perform_averaged_fit(PHSP, averaged_alpha)
        polarizations.append(extract_polarization(fit_result))
    return alphas, polarizations


(
    PROPAGATED_POLARIMETERS_CARTESIAN,
    PROPAGATED_POLARIZATIONS_CARTESIAN,
) = perform_combinatorics_fit(SYST_AVERAGED_POLARIMETERS, polar=False)
(
    PROPAGATED_POLARIMETERS_POLAR,
    PROPAGATED_POLARIZATIONS_POLAR,
) = perform_combinatorics_fit(SYST_POLAR_POLARIMETERS, polar=True)

def render_propagated_polarization(
    polarizations: list[tuple[float, float, float]], polar: bool
) -> str:
    nominal_p = extract_polarization(SYST_FIT_RESULTS_AVERAGED[0])
    diff_combinatorics = np.abs(np.array(polarizations) - np.array(nominal_p))
    px, py, pz = 100 * np.array(nominal_p)
    σx, σy, σz = 100 * diff_combinatorics.max(axis=0)
    src = Rf"""
    \begin{{array}}{{ccrcr}}
      P_x &=& {px:+6.2f} &\pm& {σx:5.2f} \\
      P_y &=& {py:+6.2f} &\pm& {σy:5.2f} \\
      P_z &=& {pz:+6.2f} &\pm& {σz:5.2f} \\
    \end{{array}}
    """
    return dedent(src).strip()


src = Rf"""
Polarizations from $\overline{{\alpha}}$ in cartesian coordinates:

$$
{render_propagated_polarization(PROPAGATED_POLARIZATIONS_CARTESIAN, polar=False)}
$$

Polarizations from $\overline{{\alpha}}$ in polar coordinates:

$$
{render_propagated_polarization(PROPAGATED_POLARIZATIONS_POLAR, polar=True)}
$$
"""
Markdown(src)

def render_combinatorics_fit(
    alphas: list[tuple[float, float, float]],
    polarizations: list[tuple[float, float, float]],
    polar: bool = False,
) -> None:
    src = R"\begin{array}{rrr|rrr|rrr}" "\n  "
    if polar:
        src += R"|\alpha| & \theta\;[\pi] & \phi\;[\pi]"
    else:
        src += R"\alpha_x & \alpha_y & \alpha_z"
    src += R" & P_x & P_y & P_z & \Delta P_x & \Delta P_y & \Delta P_z \\ " "\n"
    src += R"  \hline" "\n  "
    if polar:
        r, θ, φ = SYST_POLAR_POLARIMETERS[0]
        nominal_values = (f"{1e3*r:.1f}", f"{θ/π:.3f}", f"{φ/π:.3f}")
    else:
        αx, αy, αz = 1e3 * SYST_AVERAGED_POLARIMETERS[0]
        nominal_values = (f"{αx:.1f}", f"{αy:.1f}", f"{αz:.1f}")
    src += " & ".join(Rf"\color{{gray}}{{{v}}}" for v in nominal_values) + " & "
    nominal_α = 1e3 * SYST_AVERAGED_POLARIMETERS[0]
    if polar:
        nominal_α = (nominal_α[0], 1e-3 * nominal_α[1] / π)
    nominal_p = extract_polarization(SYST_FIT_RESULTS_AVERAGED[0])
    nominal_p = 100 * np.array(nominal_p)
    src += " & ".join(Rf"\color{{gray}}{{{v:+.2f}}}" for v in nominal_p)
    src += R" \\" "\n"
    for alpha, polarization in zip(alphas, polarizations):
        polarization = 100 * np.array(polarization)
        px, py, pz = polarization
        dx, dy, dz = polarization - nominal_p
        if polar:
            r, θ, φ = np.array(alpha)
            src += Rf"  {1e3*r:4.1f} & {θ/π:+5.2f} & {φ/π:+6.2f} "
        else:
            αx, αy, αz = 1e3 * np.array(alpha)
            src += Rf"  {αx:+5.1f} & {αy:+5.1f} & {αz:+6.1f} "
        src += Rf"& {px:+5.1f} & {py:+5.2f} & {pz:+5.1f} "
        src += Rf"& {dx:+5.2f} & {dy:+5.2f} & {dz:+5.1f} \\" "\n"
    src += R"\end{array}"
    display(Latex(src))


render_combinatorics_fit(
    PROPAGATED_POLARIMETERS_CARTESIAN,
    PROPAGATED_POLARIZATIONS_CARTESIAN,
)
render_combinatorics_fit(
    PROPAGATED_POLARIMETERS_POLAR,
    PROPAGATED_POLARIZATIONS_POLAR,
    polar=True,
)