P-vector model fit, single channel

P-vector model fit, single channel#

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 pandas as pd
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 matplotlib.figure import Figure
from qrules.particle import Particle, ParticleCollection
from sympy import Abs
from tensorwaves.data import (
    IntensityDistributionGenerator,
    SympyDataTransformer,
    TFPhaseSpaceGenerator,
    TFUniformRealNumberGenerator,
    TFWeightedPhaseSpaceGenerator,
)
from tensorwaves.estimator import UnbinnedNLL
from tensorwaves.function.sympy import create_parametrized_function
from tensorwaves.interface import DataSample, FitResult, ParametrizedFunction
from tensorwaves.optimizer import Minuit2

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.65, width=0.6, parity=-1, spin=0.5, idx=1)
    particles += create_nstar(mass=1.75, width=0.6, parity=-1, spin=0.5, idx=2)
    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

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(1/2^-)}_{+1/2} \overline{p}_{+1/2}; N_1(1/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(1/2^-)}_{+1/2} \overline{p}_{-1/2}; N_1(1/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(1/2^-)}_{+3/2} \overline{p}_{+1/2}; N_1(1/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_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(1/2^-)}_{+1/2} \overline{p}_{+1/2}; N_2(1/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(1/2^-)}_{+1/2} \overline{p}_{-1/2}; N_2(1/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(1/2^-)}_{+3/2} \overline{p}_{+1/2}; N_2(1/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_1(1/2^-) & m=1.65\text{ GeV} & \Gamma=0.6\text{ GeV} \\ N_2(1/2^-) & m=1.75\text{ GeV} & \Gamma=0.6\text{ GeV} \\ 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(1/2^-)} \beta_{N_1(1/2^-)} m_{N_1(1/2^-)}}{- \Gamma_{N_1(1/2^-)} m_{N_1(1/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_1(1/2^-)}\right)^{2}} + \frac{\Gamma_{N_2(1/2^-)} \beta_{N_2(1/2^-)} m_{N_2(1/2^-)}}{- \Gamma_{N_2(1/2^-)} m_{N_2(1/2^-)} \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) - m_{01}^{2} + \left(m_{N_2(1/2^-)}\right)^{2}} \\ 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(1/2^-)} g_{N_1(1/2^-)}}{- m_{01}^{2} + \left(m_{N_1(1/2^-)}\right)^{2}} + \frac{\beta_{N_2(1/2^-)} g_{N_2(1/2^-)}}{- m_{01}^{2} + \left(m_{N_2(1/2^-)}\right)^{2}}}{- \left(\frac{\left(g_{N_1(1/2^-)}\right)^{2}}{- m_{01}^{2} + \left(m_{N_1(1/2^-)}\right)^{2}} + \frac{\left(g_{N_2(1/2^-)}\right)^{2}}{- m_{01}^{2} + \left(m_{N_2(1/2^-)}\right)^{2}}\right) \rho^\mathrm{CM}_{m_{0},m_{1}}\left(m_{01}^{2}\right) + 1} \\ 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#

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,
)

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(100_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#

toy_parameters_bw = {
    R"m_{N_1(1/2^-)}": 1.65,
    R"m_{N_2(1/2^-)}": 1.75,
    R"m_{N_1(3/2^-)}": 1.85,
    R"m_{N_2(3/2^-)}": 1.9,
    R"\Gamma_{N_1(1/2^-)}": 1 / 1.65,
    R"\Gamma_{N_2(1/2^-)}": 1 / 1.75,
    R"\Gamma_{N_1(3/2^-)}": 1 / 1.85,
    R"\Gamma_{N_2(3/2^-)}": 1 / 1.9,
}
intensity_func_bw.update_parameters(toy_parameters_bw)

toy_parameters_fvector = {
    R"\beta_{N_1(1/2^-)}": 1 + 0j,
    R"\beta_{N_2(1/2^-)}": 1 + 0j,
    R"\beta_{N_1(3/2^-)}": 1 + 0j,
    R"\beta_{N_2(3/2^-)}": 1 + 0j,
    R"m_{N_1(1/2^-)}": 1.65,
    R"m_{N_2(1/2^-)}": 1.75,
    R"m_{N_1(3/2^-)}": 1.95,
    R"m_{N_2(3/2^-)}": 1.9,
    R"g_{N_1(1/2^-)}": 1.65,
    R"g_{N_2(1/2^-)}": 1,
    R"g_{N_1(3/2^-)}": 1,
    R"g_{N_2(3/2^-)}": 1,
}
intensity_func_fvector.update_parameters(toy_parameters_fvector)

Plot 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")
plt.tight_layout()

fig.savefig("contributions.svg", bbox_inches="tight")
plt.show()

Weighted data with \(F\) vector#

Hide code cell source 
fig, ax = plt.subplots(figsize=(9, 5))
fast_histogram(
    phsp["m_01"].real,
    bins=100,
    weights=np.real(intensity_func_fvector(phsp)),
    ax=ax,
)
ax.set_xlabel(R"$M^2\left(\eta p\right)\, \mathrm{[(GeV/c)^2]}$")
ax.set_ylabel(R"Intensity [a.u.]")
ax.set_ylim(0, None)
fig.tight_layout()

fig.savefig("weighted-phsp-f-vector.svg", bbox_inches="tight")
fig.show()


weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(
    initial_state_mass=model.reaction_info.initial_state[-1].mass,
    final_state_masses={i: p.mass for i, p in model.reaction_info.final_state.items()},
)
data_generator = IntensityDistributionGenerator(
    domain_generator=weighted_phsp_generator,
    function=intensity_func_fvector,
    domain_transformer=transformer,
)
data_momenta = data_generator.generate(50_000, rng)
data = transformer(data_momenta)
Hide code cell source 
fig, ax = plt.subplots(figsize=(9, 4))
fast_histogram(
    np.real(data["m_01"]),
    bins=200,
    alpha=0.5,
    density=True,
    ax=ax,
)
mass_parameters = {p: v for p, v in toy_parameters_bw.items() if p.startswith("m_{")}
evenly_spaced_interval = np.linspace(0, 1, num=len(mass_parameters))
colors = [plt.cm.rainbow(x) for x in evenly_spaced_interval]
ax.set_xlabel("$m$ [GeV]")
for (k, v), color in zip(mass_parameters.items(), colors, strict=True):
    ax.axvline(v, c=color, label=f"${k}$", ls="dotted")
ax.set_ylim(0, None)
ax.legend()

fig.savefig("data-f-vector.svg", bbox_inches="tight")
plt.show()

Perform fit#

Estimator definition#

estimator_bw = UnbinnedNLL(
    intensity_func_bw,
    data=data,
    phsp=phsp,
    backend="jax",
)
estimator_fvector = UnbinnedNLL(
    intensity_func_fvector,
    data=data,
    phsp=phsp,
    backend="jax",
)

Initial parameters#

Hide code cell source 
def indicate_masses(ax, intensity_func, ls: str, lw: float, typ: str):
    mass_pars = {
        k: v for k, v in intensity_func.parameters.items() if k.startswith("m_{")
    }
    for i, (k, v) in enumerate(mass_pars.items()):
        ax.axvline(v, c=f"C{i}", label=f"${k}$ ({typ})", ls=ls, lw=lw)


def compare_model(
    variable_name,
    data,
    phsp,
    function1,
    function2,
    bins=100,
) -> Figure:
    intensities1 = function1(phsp)
    intensities2 = function2(phsp)
    fig, ax = plt.subplots(figsize=(11, 4))
    fig.subplots_adjust(right=0.85, top=0.95)
    ax.set_xlabel(R"$m_{p\eta}$ [GeV]")
    ax.set_ylabel("Intensity [a. u.]")
    ax.set_yticks([])
    data_projection = np.real(data[variable_name])
    fast_histogram(
        data_projection,
        bins=bins,
        alpha=0.5,
        label="data",
        density=True,
        ax=ax,
    )
    phsp_projection = np.real(phsp[variable_name])
    fast_histogram(
        phsp_projection,
        weights=np.array(intensities1),
        bins=bins,
        fill=False,
        color="red",
        label="Fit model with F vector",
        density=True,
        ax=ax,
    )
    fast_histogram(
        phsp_projection,
        weights=np.array(intensities2),
        bins=bins,
        fill=False,
        color="blue",
        label="Fit model with Breit-Wigner",
        density=True,
        ax=ax,
    )
    ax.set_ylim(0, None)
    indicate_masses(ax, function1, ls="dashed", lw=1, typ="F vector")
    indicate_masses(ax, function2, ls="dotted", lw=1, typ="Breit-Wigner")
    fig.legend(loc="outside upper right")
    fig.show()
    return fig

initial_parameters_beta = {
    R"\beta_{N_2(1/2^-)}": 1 + 0j,
    R"\beta_{N_2(3/2^-)}": 1 + 0j,
}
initial_parameters_masses = {
    R"m_{N_1(1/2^-)}": 1.6,
    R"m_{N_2(1/2^-)}": 1.7,
    R"m_{N_1(3/2^-)}": 1.8,
    R"m_{N_2(3/2^-)}": 1.93,
}
initial_parameters_bw = {
    **initial_parameters_beta,
    **initial_parameters_masses,
    R"\Gamma_{N_1(1/2^-)}": 1 / 1.6,
    R"\Gamma_{N_2(1/2^-)}": 1 / 1.65,
    R"\Gamma_{N_1(3/2^-)}": 1 / 1.85,
    R"\Gamma_{N_2(3/2^-)}": 1 / 1.93,
}
initial_parameters_fvector = {
    **initial_parameters_beta,
    **initial_parameters_masses,
    R"g_{N_1(1/2^-)}": 1.6,
    R"g_{N_2(1/2^-)}": 1,
    R"g_{N_1(3/2^-)}": 1.0,
    R"g_{N_2(3/2^-)}": 1.0,
}
Hide code cell source 
original_parameters_bw = dict(intensity_func_bw.parameters)
intensity_func_bw.update_parameters(initial_parameters_bw)
original_parameters_fvector = dict(intensity_func_fvector.parameters)
intensity_func_fvector.update_parameters(initial_parameters_fvector)
fig = compare_model("m_01", data, phsp, intensity_func_fvector, intensity_func_bw)
fig.savefig("before-fit.svg", bbox_inches="tight")

Optimize parameters#

minuit2 = Minuit2()

fit_result_bw = minuit2.optimize(estimator_bw, initial_parameters_bw)
assert fit_result_bw.minimum_valid
fit_result_bw

FitResult(
 minimum_valid=True,
 execution_time=12.6900634765625,
 function_calls=1682,
 estimator_value=-7931.137816713118,
 parameter_values={
  'm_{N_1(1/2^-)}': 1.6539461175761394,
  'm_{N_2(1/2^-)}': 1.7264292670507233,
  'm_{N_1(3/2^-)}': 1.7328191451306136,
  'm_{N_2(3/2^-)}': 1.908580828366668,
  '\\Gamma_{N_1(1/2^-)}': 1.8524926041523075,
  '\\Gamma_{N_2(1/2^-)}': 0.13046881396159438,
  '\\Gamma_{N_1(3/2^-)}': 0.1573516791553556,
  '\\Gamma_{N_2(3/2^-)}': 1.324380387465878,
  '\\beta_{N_2(1/2^-)}': (-1.6897481365226579-1.1568947003976526j),
  '\\beta_{N_2(3/2^-)}': (1.5375186053600138-0.44760512526013624j),
 },
 parameter_errors={
  'm_{N_1(1/2^-)}': 0.004603319099772491,
  'm_{N_2(1/2^-)}': 0.0006070545961158388,
  'm_{N_1(3/2^-)}': 0.0015172566194625206,
  'm_{N_2(3/2^-)}': 0.0077272263352300054,
  '\\Gamma_{N_1(1/2^-)}': 0.06361609357610253,
  '\\Gamma_{N_2(1/2^-)}': 0.005879299479538631,
  '\\Gamma_{N_1(3/2^-)}': 0.019367706202832097,
  '\\Gamma_{N_2(3/2^-)}': 0.05320290566507636,
  '\\beta_{N_2(1/2^-)}': (0.1311293412261402+0.1278774780728592j),
  '\\beta_{N_2(3/2^-)}': (0.11272653975271772+0.22792481789700972j),
 },
)

fit_result_fvector = minuit2.optimize(estimator_fvector, initial_parameters_fvector)
assert fit_result_fvector.minimum_valid
fit_result_fvector

FitResult(
 minimum_valid=True,
 execution_time=13.932363510131836,
 function_calls=1011,
 estimator_value=-8199.598714182841,
 parameter_values={
  'm_{N_1(1/2^-)}': 1.6483798486877363,
  'm_{N_2(1/2^-)}': 1.7483200793152158,
  'm_{N_1(3/2^-)}': 1.8986142342242096,
  'm_{N_2(3/2^-)}': 1.9496754881738385,
  'g_{N_1(1/2^-)}': 1.6624278005128597,
  'g_{N_2(1/2^-)}': 0.9493267500468856,
  'g_{N_1(3/2^-)}': 1.0079832421160817,
  'g_{N_2(3/2^-)}': 0.9800998848231933,
  '\\beta_{N_2(1/2^-)}': (0.9606446724644881-0.00903038965870049j),
  '\\beta_{N_2(3/2^-)}': (0.9709322547787542+0.008602453798822136j),
 },
 parameter_errors={
  'm_{N_1(1/2^-)}': 0.001040462783404995,
  'm_{N_2(1/2^-)}': 0.0006658939683350329,
  'm_{N_1(3/2^-)}': 0.0011425836015682435,
  'm_{N_2(3/2^-)}': 0.0015413452101229357,
  'g_{N_1(1/2^-)}': 0.010998823395799513,
  'g_{N_2(1/2^-)}': 0.018365301903530223,
  'g_{N_1(3/2^-)}': 0.013236677643628186,
  'g_{N_2(3/2^-)}': 0.03222752091156066,
  '\\beta_{N_2(1/2^-)}': (0.015907378725063155+0.016622785545148663j),
  '\\beta_{N_2(3/2^-)}': (0.040758323198047446+0.01968492218927783j),
 },
)
Hide code cell source 
intensity_func_fvector.update_parameters(fit_result_fvector.parameter_values)
intensity_func_bw.update_parameters(fit_result_bw.parameter_values)
fig = compare_model("m_01", data, phsp, intensity_func_fvector, intensity_func_bw)
fig.savefig("after-fit.svg", bbox_inches="tight")

Fit result comparison#

Hide code cell source 
def compute_aic_bic(fit_result: FitResult) -> tuple[float, float]:
    n_real_par = fit_result.count_number_of_parameters(complex_twice=True)
    n_events = len(next(iter(data.values())))
    log_likelihood = -fit_result.estimator_value
    aic = 2 * n_real_par - 2 * log_likelihood
    bic = n_real_par * np.log(n_events) - 2 * log_likelihood
    return aic, bic


def compare_parameters(initial: dict, optimized: dict, expected: dict) -> pd.DataFrame:
    parameters = sorted(set(initial) | set(optimized))
    df = pd.DataFrame(
        {
            f"${p}$": (
                f"{initial[p]:.3g}",
                f"{optimized[p]:.3g}",
                f"{expected[p]:.3g}",
                f"{100 * abs((optimized[p] - expected[p]) / expected[p]):.1f}%",
            )
            for p in parameters
        },
    ).T
    df.columns = ("initial", "fit result", "expected", "deviation")
    return df

P vector#

compute_aic_bic(fit_result_fvector)

(-16375.197428365682, -16269.360088952759)
Hide code cell source 
compare_parameters(
    initial=initial_parameters_fvector,
    optimized=fit_result_fvector.parameter_values,
    expected=original_parameters_fvector,
)
initial fit result expected deviation
$\beta_{N_2(1/2^-)}$ 1+0j 0.961-0.00903j 1+0j 4.0%
$\beta_{N_2(3/2^-)}$ 1+0j 0.971+0.0086j 1+0j 3.0%
$g_{N_1(1/2^-)}$ 1.6 1.66 1.65 0.8%
$g_{N_1(3/2^-)}$ 1 1.01 1 0.8%
$g_{N_2(1/2^-)}$ 1 0.949 1 5.1%
$g_{N_2(3/2^-)}$ 1 0.98 1 2.0%
$m_{N_1(1/2^-)}$ 1.6 1.65 1.65 0.1%
$m_{N_1(3/2^-)}$ 1.8 1.9 1.95 2.6%
$m_{N_2(1/2^-)}$ 1.7 1.75 1.75 0.1%
$m_{N_2(3/2^-)}$ 1.93 1.95 1.9 2.6%

Breit–Wigner#

compute_aic_bic(fit_result_bw)

(-15838.275633426236, -15732.438294013313)
Hide code cell source 
compare_parameters(
    initial=initial_parameters_bw,
    optimized=fit_result_bw.parameter_values,
    expected=original_parameters_bw,
)
initial fit result expected deviation
$\Gamma_{N_1(1/2^-)}$ 0.625 1.85 0.606 205.7%
$\Gamma_{N_1(3/2^-)}$ 0.541 0.157 0.541 70.9%
$\Gamma_{N_2(1/2^-)}$ 0.606 0.13 0.571 77.2%
$\Gamma_{N_2(3/2^-)}$ 0.518 1.32 0.526 151.6%
$\beta_{N_2(1/2^-)}$ 1+0j -1.69-1.16j 1+0j 292.8%
$\beta_{N_2(3/2^-)}$ 1+0j 1.54-0.448j 1+0j 69.9%
$m_{N_1(1/2^-)}$ 1.6 1.65 1.65 0.2%
$m_{N_1(3/2^-)}$ 1.8 1.73 1.85 6.3%
$m_{N_2(1/2^-)}$ 1.7 1.73 1.75 1.3%
$m_{N_2(3/2^-)}$ 1.93 1.91 1.9 0.5%