Sub-intensities of P vector#

Hide code cell source
from __future__ import annotations

import logging
import os
import re
from collections import defaultdict
from functools import lru_cache
from pathlib import Path
from typing import Any

import ampform
import attrs
import graphviz
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import qrules
import sympy as sp
from ampform.dynamics.builder import TwoBodyKinematicVariableSet
from ampform.helicity import ParameterValues
from ampform.io import aslatex, improve_latex_rendering
from ampform.kinematics.phasespace import Kallen
from ampform.sympy import perform_cached_doit, unevaluated
from attrs import define, field
from IPython.display import Math
from qrules.particle import Particle, ParticleCollection
from sympy import Abs
from tensorwaves.data import (
    SympyDataTransformer,
    TFPhaseSpaceGenerator,
    TFUniformRealNumberGenerator,
)
from tensorwaves.function.sympy import create_parametrized_function
from tensorwaves.interface import DataSample, Function, ParametrizedFunction

improve_latex_rendering()
logging.getLogger("absl").setLevel(logging.ERROR)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
plt.rc("font", size=12)

Studied decay#

Hide code cell source
@lru_cache(maxsize=1)
def create_particle_database() -> ParticleCollection:
    particles = qrules.load_default_particles()
    for nstar in particles.filter(lambda p: p.name.startswith("N")):
        particles.remove(nstar)
    particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)
    particles += create_nstar(mass=1.92, width=0.6, parity=+1, spin=1.5, idx=2)
    return particles


def create_nstar(
    mass: float, width: float, parity: int, spin: float, idx: int
) -> Particle:
    spin = sp.Rational(spin)
    parity_symbol = "⁺" if parity > 0 else "⁻"
    unicode_subscripts = list("₀₁₂₃₄₅₆₇₈₉")
    return Particle(
        name=f"N{unicode_subscripts[idx]}({spin}{parity_symbol})",
        latex=Rf"N_{idx}({spin.numerator}/{spin.denominator}^-)",
        pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,
        mass=mass,
        width=width,
        baryon_number=1,
        charge=+1,
        isospin=(0.5, +0.5),
        parity=parity,
        spin=1.5,
    )
reaction = qrules.generate_transitions(
    initial_state="J/psi(1S)",
    final_state=["eta", "p", "p~"],
    allowed_intermediate_particles=["N"],
    allowed_interaction_types=["strong"],
    formalism="helicity",
    particle_db=create_particle_database(),
)
Hide code cell source
dot = qrules.io.asdot(reaction, collapse_graphs=True)
graph = graphviz.Source(dot)
output_file = Path("qrules-output")
graph.render(output_file, format="svg")
output_file.unlink()
graph

Image

Amplitude builder#

Hide code cell source
@define
class DynamicsSymbolBuilder:
    collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (
        field(factory=lambda: defaultdict(set))
    )

    def __call__(
        self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet
    ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:
        jp = render_jp(resonance)
        charge = resonance.charge
        if variable_pool.angular_momentum is not None:
            L = sp.Rational(variable_pool.angular_momentum)
            X = sp.Symbol(Rf"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}")
        else:
            X = sp.Symbol(Rf"X_{{{jp}, Q={charge:+d}}}")
        self.collected_symbols[X].add((resonance, variable_pool))
        parameter_defaults = {}
        return X, parameter_defaults


def render_jp(particle: Particle) -> str:
    spin = sp.Rational(particle.spin)
    j = (
        str(spin)
        if spin.denominator == 1
        else Rf"\frac{{{spin.numerator}}}{{{spin.denominator}}}"
    )
    if particle.parity is None:
        return f"J={j}"
    p = "-" if particle.parity < 0 else "+"
    return f"J^P={{{j}}}^{{{p}}}"
model_builder = ampform.get_builder(reaction)
model_builder.adapter.permutate_registered_topologies()
model_builder.config.scalar_initial_state_mass = True
model_builder.config.stable_final_state_ids = [0, 1, 2]
create_dynamics_symbol = DynamicsSymbolBuilder()
for resonance in reaction.get_intermediate_particles():
    model_builder.set_dynamics(resonance.name, create_dynamics_symbol)
model = model_builder.formulate()
model.intensity.cleanup()
\[\displaystyle \sum_{m_{A}=-1}^{1} \sum_{m_{1}=-1/2}^{1/2} \sum_{m_{2}=-1/2}^{1/2}{\left|{A^{01}_{m_{A}, 0, m_{1}, m_{2}}}\right|^{2}}\]
Hide code cell source
selected_amplitudes = {
    k: v for i, (k, v) in enumerate(model.amplitudes.items()) if i == 0
}
Math(aslatex(selected_amplitudes, terms_per_line=1))
\[\begin{split}\displaystyle \begin{array}{rcl} A^{01}_{0, 0, - \frac{1}{2}, - \frac{1}{2}} &=& - C_{J/\psi(1S) \to {N_1(3/2^-)}_{+1/2} \overline{p}_{+1/2}; N_1(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,0}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{- \frac{1}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ &+& - C_{J/\psi(1S) \to {N_1(3/2^-)}_{+1/2} \overline{p}_{-1/2}; N_1(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,1}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{\frac{1}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ &+& - C_{J/\psi(1S) \to {N_1(3/2^-)}_{+3/2} \overline{p}_{+1/2}; N_1(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,-1}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{- \frac{3}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ &+& - C_{J/\psi(1S) \to {N_2(3/2^-)}_{+1/2} \overline{p}_{+1/2}; N_2(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,0}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{- \frac{1}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ &+& - C_{J/\psi(1S) \to {N_2(3/2^-)}_{+1/2} \overline{p}_{-1/2}; N_2(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,1}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{\frac{1}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ &+& - C_{J/\psi(1S) \to {N_2(3/2^-)}_{+3/2} \overline{p}_{+1/2}; N_2(3/2^-) \to \eta_{0} p_{+1/2}} X_{J^P={\frac{3}{2}}^{+}, Q=+1} D^{1}_{0,-1}\left(- \phi_{01},\theta_{01},0\right) D^{\frac{3}{2}}_{- \frac{3}{2},\frac{1}{2}}\left(- \phi^{01}_{0},\theta^{01}_{0},0\right) \\ \end{array}\end{split}\]
Hide code cell source
src = R"\begin{array}{cll}" "\n"
for symbol, resonances in create_dynamics_symbol.collected_symbols.items():
    src += Rf"  {symbol} \\" "\n"
    for p, _ in resonances:
        src += Rf"  {p.latex} & m={p.mass:g}\text{{ GeV}} & \Gamma={p.width:g}\text{{ GeV}} \\"
        src += "\n"
src += R"\end{array}"
Math(src)
\[\begin{split}\displaystyle \begin{array}{cll} X_{J^P={\frac{3}{2}}^{+}, Q=+1} \\ N_2(3/2^-) & m=1.92\text{ GeV} & \Gamma=0.6\text{ GeV} \\ N_1(3/2^-) & m=1.82\text{ GeV} & \Gamma=0.6\text{ GeV} \\ \end{array}\end{split}\]

Dynamics parametrization#

Phasespace factor#

See also

TR-026 and TR-027 on analyticity and Riemann sheets.

Hide code cell source
@unevaluated(real=False)
class PhaseSpaceCM(sp.Expr):
    s: Any
    m1: Any
    m2: Any
    _latex_repr_ = R"\rho^\mathrm{{CM}}_{{{m1},{m2}}}\left({s}\right)"

    def evaluate(self) -> sp.Expr:
        s, m1, m2 = self.args
        return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)


@unevaluated(real=False)
class ChewMandelstam(sp.Expr):
    s: Any
    m1: Any
    m2: Any
    _latex_repr_ = R"\Sigma\left({s}\right)"

    def evaluate(self) -> sp.Expr:
        s, m1, m2 = self.args
        q = BreakupMomentum(s, m1, m2)
        return (
            (2 * q / sp.sqrt(s))
            * sp.log(Abs((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)))
            - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)
        ) / (16 * sp.pi**2)


@unevaluated(real=False)
class BreakupMomentum(sp.Expr):
    s: Any
    m1: Any
    m2: Any
    _latex_repr_ = R"q\left({s}\right)"

    def evaluate(self) -> sp.Expr:
        s, m1, m2 = self.args
        return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))
Hide code cell source
s, m1, m2 = sp.symbols("s m1 m2", nonnegative=True)
exprs = [
    PhaseSpaceCM(s, m1, m2),
    ChewMandelstam(s, m1, m2),
    BreakupMomentum(s, m1, m2),
]
Math(aslatex({e: e.doit(deep=False) for e in exprs}))
\[\begin{split}\displaystyle \begin{array}{rcl} \rho^\mathrm{CM}_{m_{1},m_{2}}\left(s\right) &=& - 16 i \pi \Sigma\left(s\right) \\ \Sigma\left(s\right) &=& \frac{- \left(m_{1}^{2} - m_{2}^{2}\right) \left(- \frac{1}{\left(m_{1} + m_{2}\right)^{2}} + \frac{1}{s}\right) \log{\left(\frac{m_{1}}{m_{2}} \right)} + \frac{2 \log{\left(\frac{\left|{m_{1}^{2} + m_{2}^{2} + 2 \sqrt{s} q\left(s\right) - s}\right|}{2 m_{1} m_{2}} \right)} q\left(s\right)}{\sqrt{s}}}{16 \pi^{2}} \\ q\left(s\right) &=& \frac{\sqrt{\lambda\left(s, m_{1}^{2}, m_{2}^{2}\right)}}{2 \sqrt{s}} \\ \end{array}\end{split}\]

Relativistic Breit-Wigner#

PARAMETERS_BW = dict(model.parameter_defaults)
def formulate_breit_wigner(
    resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],
) -> sp.Expr:
    (_, variables), *_ = resonances
    s = variables.incoming_state_mass**2
    m1 = variables.outgoing_state_mass1
    m2 = variables.outgoing_state_mass2
    ρ = PhaseSpaceCM(s, m1, m2)
    m = [sp.Symbol(Rf"m_{{{p.latex}}}") for p, _ in resonances]
    Γ0 = [sp.Symbol(Rf"\Gamma_{{{p.latex}}}") for p, _ in resonances]
    β = [sp.Symbol(Rf"\beta_{{{p.latex}}}") for p, _ in resonances]
    expr = sum(
        (β_ * m_ * Γ0_) / (m_**2 - s - m_ * Γ0_ * ρ)
        for m_, Γ0_, β_ in zip(m, Γ0, β, strict=True)
    )
    for i, (resonance, _) in enumerate(resonances):
        PARAMETERS_BW[β[i]] = 1 + 0j
        PARAMETERS_BW[m[i]] = resonance.mass
        PARAMETERS_BW[Γ0[i]] = resonance.width
    return expr
Hide code cell source
dynamics_expressions_bw = {
    symbol: formulate_breit_wigner(resonances)
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
}
model_bw = attrs.evolve(
    model,
    parameter_defaults=ParameterValues({
        **model.parameter_defaults,
        **PARAMETERS_BW,
    }),
)
Math(aslatex(dynamics_expressions_bw))
\[\begin{split}\displaystyle \begin{array}{rcl} X_{J^P={\frac{3}{2}}^{+}, Q=+1} &=& \frac{\Gamma_{N_1(3/2^-)} \beta_{N_1(3/2^-)} m_{N_1(3/2^-)}}{- \Gamma_{N_1(3/2^-)} m_{N_1(3/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_1(3/2^-)}\right)^{2}} + \frac{\Gamma_{N_2(3/2^-)} \beta_{N_2(3/2^-)} m_{N_2(3/2^-)}}{- \Gamma_{N_2(3/2^-)} m_{N_2(3/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_2(3/2^-)}\right)^{2}} \\ \end{array}\end{split}\]

\(P\) vector#

PARAMETERS_F = dict(model.parameter_defaults)
def formulate_k_matrix(
    resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],
) -> sp.Expr:
    (_, variables), *_ = resonances
    s = variables.incoming_state_mass**2
    m = [sp.Symbol(Rf"m_{{{p.latex}}}") for p, _ in resonances]
    g = [sp.Symbol(Rf"g_{{{p.latex}}}") for p, _ in resonances]
    expr = sum((g_**2) / (m_**2 - s) for m_, g_ in zip(m, g, strict=True))
    for i, (resonance, _) in enumerate(resonances):
        PARAMETERS_F[m[i]] = resonance.mass
        PARAMETERS_F[g[i]] = 1
    return expr
def formulate_p_vector(
    resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],
) -> sp.Expr:
    (_, variables), *_ = resonances
    s = variables.incoming_state_mass**2
    g = [sp.Symbol(Rf"g_{{{p.latex}}}") for p, _ in resonances]
    m = [sp.Symbol(Rf"m_{{{p.latex}}}") for p, _ in resonances]
    β = [sp.Symbol(Rf"\beta_{{{p.latex}}}") for p, _ in resonances]
    expr = sum((g_ * β_) / (m_**2 - s) for m_, g_, β_ in zip(m, g, β, strict=True))
    for i, (resonance, _) in enumerate(resonances):
        PARAMETERS_F[β[i]] = 1 + 0j
        PARAMETERS_F[m[i]] = resonance.mass
        PARAMETERS_F[g[i]] = 1
    return expr
def formulate_f_vector(
    resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],
) -> sp.Expr:
    (_, variables), *_ = resonances
    s = variables.incoming_state_mass**2
    m1 = variables.outgoing_state_mass1
    m2 = variables.outgoing_state_mass2
    rho = PhaseSpaceCM(s, m1, m2)
    K = formulate_k_matrix(resonances)
    P = formulate_p_vector(resonances)
    return (1 / (1 - rho * K)) * P
Hide code cell source
dynamics_expressions_fvector = {
    symbol: formulate_f_vector(resonances)
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
}
model_fvector = attrs.evolve(
    model,
    parameter_defaults=ParameterValues({
        **model.parameter_defaults,
        **PARAMETERS_F,
    }),
)
Math(aslatex(dynamics_expressions_fvector))
\[\begin{split}\displaystyle \begin{array}{rcl} X_{J^P={\frac{3}{2}}^{+}, Q=+1} &=& \frac{\frac{\beta_{N_1(3/2^-)} g_{N_1(3/2^-)}}{- m_{01}^{2} + \left(m_{N_1(3/2^-)}\right)^{2}} + \frac{\beta_{N_2(3/2^-)} g_{N_2(3/2^-)}}{- m_{01}^{2} + \left(m_{N_2(3/2^-)}\right)^{2}}}{- \left(\frac{\left(g_{N_1(3/2^-)}\right)^{2}}{- m_{01}^{2} + \left(m_{N_1(3/2^-)}\right)^{2}} + \frac{\left(g_{N_2(3/2^-)}\right)^{2}}{- m_{01}^{2} + \left(m_{N_2(3/2^-)}\right)^{2}}\right) \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) + 1} \\ \end{array}\end{split}\]

Create numerical functions#

Amplitude model function#

full_expression_bw = perform_cached_doit(model_bw.expression).xreplace(
    dynamics_expressions_bw
)
intensity_func_bw = create_parametrized_function(
    expression=perform_cached_doit(full_expression_bw),
    backend="jax",
    parameters=PARAMETERS_BW,
)
full_expression_fvector = perform_cached_doit(model_fvector.expression).xreplace(
    dynamics_expressions_fvector
)
intensity_func_fvector = create_parametrized_function(
    expression=perform_cached_doit(full_expression_fvector),
    backend="jax",
    parameters=PARAMETERS_F,
)

Dynamics function#

Hide code cell source
dynamics_expr_bw, *_ = dynamics_expressions_bw.values()
dynamics_expr_bw
\[\displaystyle \frac{\Gamma_{N_1(3/2^-)} \beta_{N_1(3/2^-)} m_{N_1(3/2^-)}}{- \Gamma_{N_1(3/2^-)} m_{N_1(3/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_1(3/2^-)}\right)^{2}} + \frac{\Gamma_{N_2(3/2^-)} \beta_{N_2(3/2^-)} m_{N_2(3/2^-)}}{- \Gamma_{N_2(3/2^-)} m_{N_2(3/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_2(3/2^-)}\right)^{2}}\]
Hide code cell source
dynamics_expr_fvector, *_ = dynamics_expressions_fvector.values()
dynamics_expr_fvector.simplify(doit=False)
\[\displaystyle \frac{\beta_{N_1(3/2^-)} g_{N_1(3/2^-)} \left(m_{01}^{2} - \left(m_{N_2(3/2^-)}\right)^{2}\right) + \beta_{N_2(3/2^-)} g_{N_2(3/2^-)} \left(m_{01}^{2} - \left(m_{N_1(3/2^-)}\right)^{2}\right)}{- \left(m_{01}^{2} - \left(m_{N_1(3/2^-)}\right)^{2}\right) \left(m_{01}^{2} - \left(m_{N_2(3/2^-)}\right)^{2}\right) + \left(- \left(g_{N_1(3/2^-)}\right)^{2} \left(m_{01}^{2} - \left(m_{N_2(3/2^-)}\right)^{2}\right) - \left(g_{N_2(3/2^-)}\right)^{2} \left(m_{01}^{2} - \left(m_{N_1(3/2^-)}\right)^{2}\right)\right) \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right)}\]
dynamics_func_bw = create_parametrized_function(
    expression=perform_cached_doit(dynamics_expr_bw),
    backend="jax",
    parameters=model_bw.parameter_defaults,
)
dynamics_func_fvector = create_parametrized_function(
    expression=perform_cached_doit(dynamics_expr_fvector),
    backend="jax",
    parameters=model_fvector.parameter_defaults,
)

Generate data#

Generate phase space sample#

rng = TFUniformRealNumberGenerator(seed=0)
phsp_generator = TFPhaseSpaceGenerator(
    initial_state_mass=reaction.initial_state[-1].mass,
    final_state_masses={i: p.mass for i, p in reaction.final_state.items()},
)
phsp_momenta = phsp_generator.generate(500_000, rng)

ε = 1e-8
transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend="jax")
phsp = transformer(phsp_momenta)
phsp = {k: v + ε * 1j if re.match(r"^m_\d\d$", k) else v for k, v in phsp.items()}

Update function parameters#

m_res1 = 1.82
m_res2 = 1.92
g_res1 = 1
g_res2 = 1
toy_parameters_bw = {
    R"m_{N_1(3/2^-)}": m_res1,
    R"m_{N_2(3/2^-)}": m_res2,
    R"\Gamma_{N_1(3/2^-)}": g_res1 / m_res1,
    R"\Gamma_{N_2(3/2^-)}": g_res2 / m_res2,
}
dynamics_func_bw.update_parameters(toy_parameters_bw)
intensity_func_bw.update_parameters(toy_parameters_bw)
toy_parameters_fvector = {
    R"\beta_{N_1(3/2^-)}": 1 + 0j,
    R"\beta_{N_2(3/2^-)}": 1 + 0j,
    R"m_{N_1(3/2^-)}": m_res1,
    R"m_{N_2(3/2^-)}": m_res2,
    R"g_{N_1(3/2^-)}": g_res1,
    R"g_{N_2(3/2^-)}": g_res2,
}
dynamics_func_fvector.update_parameters(toy_parameters_fvector)
intensity_func_fvector.update_parameters(toy_parameters_fvector)

Plots#

Sub-intensities#

Hide code cell source
def compute_sub_intensity(
    func: ParametrizedFunction,
    input_data: DataSample,
    resonances: list[str],
    coupling_pattern: str,
):
    original_parameters = dict(func.parameters)
    negative_lookahead = f"(?!{'|'.join(map(re.escape, resonances))})"
    # https://regex101.com/r/WrgGyD/1
    pattern = rf"^{coupling_pattern}({negative_lookahead}.)*$"
    set_parameters_to_zero(func, pattern)
    array = func(input_data)
    func.update_parameters(original_parameters)
    return array


def set_parameters_to_zero(func: ParametrizedFunction, name_pattern: str) -> None:
    toy_parameters = dict(func.parameters)
    for par_name in func.parameters:
        if re.match(name_pattern, par_name) is not None:
            toy_parameters[par_name] = 0
    func.update_parameters(toy_parameters)
total_intensities_bw = intensity_func_bw(phsp)
sub_intensities_bw = {
    p: compute_sub_intensity(
        intensity_func_bw,
        phsp,
        resonances=[p.latex],
        coupling_pattern=r"\\beta",
    )
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
    for p, _ in resonances
}
total_intensities_fvector = intensity_func_fvector(phsp)
sub_intensities_fvector = {
    p: compute_sub_intensity(
        intensity_func_fvector,
        phsp,
        resonances=[p.latex],
        coupling_pattern=r"\\beta",
    )
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
    for p, _ in resonances
}
Hide code cell source
def fast_histogram(
    data: jnp.ndarray,
    weights: jnp.ndarray | None = None,
    bins: int = 100,
    density: bool | None = None,
    fill: bool = True,
    ax=plt,
    **plot_kwargs,
) -> None:
    bin_values, bin_edges = jnp.histogram(
        data,
        bins=bins,
        density=density,
        weights=weights,
    )
    if fill:
        bin_rights = bin_edges[1:]
        ax.fill_between(bin_rights, bin_values, step="pre", **plot_kwargs)
    else:
        bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2
        ax.step(bin_mids, bin_values, **plot_kwargs)
Hide code cell source
fig, ax = plt.subplots(figsize=(8, 5))
ax.set_xlim(2, 5)
ax.set_xlabel(R"$m_{p\eta}^{2}$ [GeV$^2$]")
ax.set_ylabel(R"Intensity [a. u.]")
ax.set_yticks([])

bins = 150
phsp_projection = np.real(phsp["m_01"]) ** 2
fast_histogram(
    phsp_projection,
    weights=total_intensities_fvector,
    alpha=0.2,
    bins=bins,
    color="hotpink",
    label="Full intensity $F$ vector",
    ax=ax,
)
fast_histogram(
    phsp_projection,
    weights=total_intensities_bw,
    alpha=0.2,
    bins=bins,
    color="grey",
    label="Full intensity Breit-Wigner",
    ax=ax,
)
for i, (p, v) in enumerate(sub_intensities_fvector.items()):
    fast_histogram(
        phsp_projection,
        weights=v,
        alpha=0.6,
        bins=bins,
        color=f"C{i}",
        fill=False,
        label=Rf"Resonance at ${p.mass}\,\mathrm{{GeV}}$ $F$ vector",
        linewidth=2,
        ax=ax,
    )
for i, (p, v) in enumerate(sub_intensities_bw.items()):
    fast_histogram(
        phsp_projection,
        weights=v,
        alpha=0.6,
        bins=bins,
        color=f"C{i}",
        fill=False,
        label=Rf"Resonance at ${p.mass}\,\mathrm{{GeV^2}}$ Breit-Wigner",
        linestyle="dashed",
        ax=ax,
    )

ax.set_ylim(0, None)
fig.legend(loc="upper right")
fig.tight_layout()

fig.savefig("sub-intensities.svg", bbox_inches="tight")
plt.show()

Image

Argand plots#

ε = 1e-8
x = np.linspace(2, 5, num=400)
plot_data = {"m_01": np.sqrt(x) + ε * 1j}
total_dynamics_bw = dynamics_func_bw(plot_data)
sub_dynamics_bw = {
    p: compute_sub_intensity(
        dynamics_func_bw,
        plot_data,
        resonances=[p.latex],
        coupling_pattern=r"\\beta",
    )
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
    for p, _ in resonances
}
total_dynamics_fvector = dynamics_func_fvector(plot_data)
sub_dynamics_fvector = {
    p: compute_sub_intensity(
        dynamics_func_fvector,
        plot_data,
        resonances=[p.latex],
        coupling_pattern=r"\\beta",
    )
    for symbol, resonances in create_dynamics_symbol.collected_symbols.items()
    for p, _ in resonances
}
x1 = np.linspace(2.0, (m_res1**2 + m_res2**2) / 2, num=500)
x2 = np.linspace((m_res1**2 + m_res2**2) / 2, 5.0, num=500)
plot_data1 = {"m_01": np.sqrt(x1) + ε * 1j}
plot_data2 = {"m_01": np.sqrt(x2) + ε * 1j}
Hide code cell source
def plot_argand(
    total_func: Function, sub_funcs: dict[Particle, Function], title: str
) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharey=True)
    fig.subplots_adjust(wspace=0.05)
    fig.suptitle(title, y=0.99)
    ax1, ax2 = axes
    ax1.set_title("Total amplitude")
    ax2.set_title("Amplitude for resonance only")
    ax1.set_ylabel(R"$\text{Im}\,F$")
    for ax in axes:
        ax.axhline(0, color="black", linewidth=0.5)
        ax.axvline(0, color="black", linewidth=0.5)
        ax.set_xlabel(R"$\text{Re}\,F$")

    y1 = total_func(plot_data1)
    ax1.plot(
        y1.real,
        y1.imag,
        label=f"Domain of {m_res1}-GeV resonance ",
        color="C0",
    )
    y2 = total_func(plot_data2)
    ax1.plot(
        y2.real,
        y2.imag,
        label=f"Domain of {m_res2}-GeV resonance ",
        color="C1",
    )
    for i, (k, v) in enumerate(sub_funcs.items()):
        ax2.plot(
            v.real,
            v.imag,
            color=f"C{i}",
            label=f"Resonance at {k.mass} GeV $F$-vector",
        )

    ax1.legend(loc="upper left")

    output_file = f"argand-{title.lower().replace(' ', '-')}.svg"
    fig.savefig(output_file, bbox_inches="tight")
    fig.show()
Hide code cell source
plot_argand(
    dynamics_func_fvector,
    sub_dynamics_fvector,
    title="F vector",
)

Image

Hide code cell source
plot_argand(
    dynamics_func_bw,
    sub_dynamics_bw,
    title="Breit-Wigner",
)

Image

Phase#

total_phase_bw = np.angle(total_dynamics_bw)
total_phase_fvector = np.angle(total_dynamics_fvector)
sub_phase_bw = {p: np.angle(v) for p, v in sub_dynamics_bw.items()}
sub_phase_fvector = {p: np.angle(v) for p, v in sub_dynamics_fvector.items()}
Hide code cell source
def plot_phases(
    total_intensity_array: np.ndarray,
    sub_intensity_arrays: dict[Particle, np.ndarray],
    total_phase_array: np.ndarray,
    sub_phase_arrays: dict[Particle, np.ndarray],
    title: str,
) -> None:
    fig, ax1 = plt.subplots(figsize=(10, 6))
    ax1.set_title(title)
    ax2 = ax1.twinx()
    ax1.set_xlim(2.0, 5.0)
    ax1.set_xlabel(R"$m_{p\eta}^{2}$ [GeV$^{2}$]")
    ax1.set_ylabel("Intensity [a. u.]")
    ax2.set_ylabel("Angle")
    ax1.set_yticks([])
    ax2.set_ylim([-np.pi, +np.pi])
    ax2.set_yticks([
        -np.pi,
        -np.pi / 2,
        0,
        +np.pi / 2,
        +np.pi,
    ])
    ax2.set_yticklabels([
        R"$-\pi$",
        R"$-\frac{\pi}{2}$",
        "0",
        R"$+\frac{\pi}{2}$",
        R"$+\pi$",
    ])
    ax2.axhline(0, c="black", lw=0.5)

    # Plot background histograms
    phsp_projection = np.real(phsp["m_01"]) ** 2
    fast_histogram(
        phsp_projection,
        weights=total_intensity_array,
        bins=bins,
        alpha=0.2,
        color="gray",
        label="Full intensity",
        ax=ax1,
    )
    for i, (k, v) in enumerate(sub_intensity_arrays.items()):
        fast_histogram(
            phsp_projection,
            weights=v,
            bins=bins,
            alpha=0.2,
            color=f"C{i}",
            label=Rf"Resonance at ${k.mass}\,\mathrm{{GeV}}$",
            ax=ax1,
        )
    ax1.set_ylim(0, None)

    # Plot phases
    ax2.scatter(
        x,
        total_phase_array,
        color="gray",
        label="Total Phase",
        s=18,
    )
    for i, (k, v) in enumerate(sub_phase_arrays.items()):
        ax2.scatter(
            x,
            v,
            alpha=0.5,
            color=f"C{i}",
            label=f"Resonance at {k.mass} GeV",
            s=8,
        )
        ax2.axvline(k.mass**2, linestyle="dotted", color=f"C{i}")

    # Add legends
    fig.legend(bbox_to_anchor=(0.1, 0.9), loc="upper left")
    fig.tight_layout()

    output_file = f"phase-{title.lower().replace(' ', '-')}.svg"
    fig.savefig(output_file, bbox_inches="tight")
    fig.show()
Hide code cell source
plot_phases(
    total_intensity_array=total_intensities_fvector,
    sub_intensity_arrays=sub_intensities_fvector,
    total_phase_array=total_phase_fvector,
    sub_phase_arrays=sub_phase_fvector,
    title="F vector",
)

Image

Hide code cell source
plot_phases(
    total_intensity_array=total_intensities_bw,
    sub_intensity_arrays=sub_intensities_bw,
    total_phase_array=total_phase_bw,
    sub_phase_arrays=sub_phase_bw,
    title="Breit-Wigner",
)

Image