Amplitude analysis with zfit

Amplitude analysis with zfit#

Hide code cell content 
%config InlineBackend.figure_formats = ['svg']

import logging
import os
import warnings

JAX_LOGGER = logging.getLogger("absl")
JAX_LOGGER.setLevel(logging.ERROR)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
warnings.simplefilter("ignore", UserWarning)
warnings.filterwarnings("ignore")

Formulating the model#

import qrules

reaction = qrules.generate_transitions(
    initial_state=("J/psi(1S)", [-1, +1]),
    final_state=["gamma", "pi0", "pi0"],
    allowed_intermediate_particles=["f(0)"],
    allowed_interaction_types=["strong", "EM"],
    formalism="helicity",
)
Hide code cell source 
import graphviz

dot = qrules.io.asdot(reaction, collapse_graphs=True)
graphviz.Source(dot)

image

import ampform
from ampform.dynamics.builder import (
    create_non_dynamic_with_ff,
    create_relativistic_breit_wigner_with_ff,
)

model_builder = ampform.get_builder(reaction)
model_builder.scalar_initial_state_mass = True
model_builder.stable_final_state_ids = [0, 1, 2]
model_builder.set_dynamics("J/psi(1S)", create_non_dynamic_with_ff)
for name in reaction.get_intermediate_particles().names:
    model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)
model = model_builder.formulate()

Generate data#

Phase space sample#

from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator

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)

Intensity-based sample#

from tensorwaves.function.sympy import create_function

unfolded_expression = model.expression.doit()
fixed_intensity_func = create_function(
    unfolded_expression.xreplace(model.parameter_defaults),
    backend="jax",
)

from tensorwaves.data import SympyDataTransformer

transform_momenta = SympyDataTransformer.from_sympy(
    model.kinematic_variables, backend="jax"
)

from tensorwaves.data import (
    IntensityDistributionGenerator,
    TFWeightedPhaseSpaceGenerator,
)

weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(
    initial_state_mass=reaction.initial_state[-1].mass,
    final_state_masses={i: p.mass for i, p in reaction.final_state.items()},
)
data_generator = IntensityDistributionGenerator(
    domain_generator=weighted_phsp_generator,
    function=fixed_intensity_func,
    domain_transformer=transform_momenta,
)
data_momenta = data_generator.generate(10_000, rng)

import pandas as pd

phsp = transform_momenta(phsp_momenta)
data = transform_momenta(data_momenta)
pd.DataFrame(data)
m_12 phi_0 phi_1^12 theta_0 theta_1^12
0 1.499845+0.000000j 2.941350 -0.984419 2.344617 1.064114
1 0.580070+0.000000j 1.422127 0.183725 1.086667 1.535691
2 1.495937+0.000000j 2.695585 3.063622 0.777978 1.730394
3 1.172263+0.000000j 0.527850 1.515685 1.343530 0.602596
4 1.581282+0.000000j -0.678981 -2.951556 2.987470 1.959462
... ... ... ... ... ...
9995 1.486016+0.000000j -1.271331 -1.387495 2.792571 2.565453
9996 0.584599+0.000000j -2.452912 -1.957086 1.070889 2.313677
9997 1.956302+0.000000j 0.378314 2.711496 0.588987 1.551541
9998 1.585024+0.000000j -0.816920 -1.166315 2.076068 1.807813
9999 1.712966+0.000000j 0.604657 0.553347 1.264140 2.079405

10000 rows × 5 columns

Hide code cell source 
import matplotlib.pyplot as plt
import numpy as np

resonances = sorted(
    reaction.get_intermediate_particles(),
    key=lambda p: p.mass,
)
evenly_spaced_interval = np.linspace(0, 1, len(resonances))
colors = [plt.cm.rainbow(x) for x in evenly_spaced_interval]
fig, ax = plt.subplots(figsize=(9, 4))
ax.hist(
    np.real(data["m_12"]),
    bins=100,
    alpha=0.5,
    density=True,
)
ax.set_xlabel("$m$ [GeV]")
for p, color in zip(resonances, colors, strict=True):
    ax.axvline(x=p.mass, linestyle="dotted", label=p.name, color=color)
ax.legend()
plt.show()

Fit#

Determine free parameters#

initial_parameters = {
    R"C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}": (
        1.0 + 0.0j
    ),
    "m_{f_{0}(500)}": 0.4,
    "m_{f_{0}(980)}": 0.88,
    "m_{f_{0}(1370)}": 1.22,
    "m_{f_{0}(1500)}": 1.45,
    "m_{f_{0}(1710)}": 1.83,
    R"\Gamma_{f_{0}(500)}": 0.3,
    R"\Gamma_{f_{0}(980)}": 0.1,
    R"\Gamma_{f_{0}(1710)}": 0.3,
}

Parametrized function and caching#

from tensorwaves.function.sympy import create_parametrized_function

intensity_func = create_parametrized_function(
    expression=unfolded_expression,
    parameters=model.parameter_defaults,
    backend="jax",
)

from tensorwaves.estimator import create_cached_function

free_parameter_symbols = [
    symbol
    for symbol in model.parameter_defaults
    if symbol.name in set(initial_parameters)
]
cached_intensity_func, transform_to_cache = create_cached_function(
    unfolded_expression,
    parameters=model.parameter_defaults,
    free_parameters=free_parameter_symbols,
    backend="jax",
)
cached_data = transform_to_cache(data)
cached_phsp = transform_to_cache(phsp)

Estimator#

from tensorwaves.estimator import UnbinnedNLL

estimator = UnbinnedNLL(
    intensity_func,
    data=data,
    phsp=phsp,
    backend="jax",
)
estimator_with_caching = UnbinnedNLL(
    cached_intensity_func,
    data=cached_data,
    phsp=cached_phsp,
    backend="jax",
)

Optimize fit parameters#

Hide code cell content 
import matplotlib.pyplot as plt
import numpy as np

reaction_info = model.reaction_info
resonances = sorted(
    reaction_info.get_intermediate_particles(),
    key=lambda p: p.mass,
)

evenly_spaced_interval = np.linspace(0, 1, len(resonances))
colors = [plt.cm.rainbow(x) for x in evenly_spaced_interval]


def indicate_masses(ax):
    ax.set_xlabel("$m$ [GeV]")
    for color, resonance in zip(colors, resonances, strict=True):
        ax.axvline(
            x=resonance.mass,
            linestyle="dotted",
            label=resonance.name,
            color=color,
        )


def compare_model(
    variable_name,
    data,
    phsp,
    function,
    bins=100,
):
    intensities = function(phsp)
    _, ax = plt.subplots(figsize=(9, 4))
    data_projection = np.real(data[variable_name])
    ax = plt.gca()
    ax.hist(
        data_projection,
        bins=bins,
        alpha=0.5,
        label="data",
        density=True,
    )
    phsp_projection = np.real(phsp[variable_name])
    ax.hist(
        phsp_projection,
        weights=np.array(intensities),
        bins=bins,
        histtype="step",
        color="red",
        label="fit model",
        density=True,
    )
    indicate_masses(ax)
    ax.legend()

original_parameters = intensity_func.parameters
intensity_func.update_parameters(initial_parameters)
compare_model("m_12", data, phsp, intensity_func)


from tensorwaves.optimizer import Minuit2
from tensorwaves.optimizer.callbacks import CSVSummary

minuit2 = Minuit2(
    callback=CSVSummary("fit_traceback.csv"),
    use_analytic_gradient=False,
)
fit_result = minuit2.optimize(estimator, initial_parameters)
fit_result

FitResult(
 minimum_valid=True,
 execution_time=7.060763359069824,
 function_calls=539,
 estimator_value=-4891.01730754809,
 parameter_values={
  'm_{f_{0}(500)}': 0.6102707294724865,
  'm_{f_{0}(980)}': 0.9902119846615327,
  'm_{f_{0}(1370)}': 1.3456300421915652,
  'm_{f_{0}(1500)}': 1.50502995100389,
  'm_{f_{0}(1710)}': 1.7096496843682751,
  '\\Gamma_{f_{0}(500)}': 0.4226040807774344,
  '\\Gamma_{f_{0}(980)}': 0.06479339507889993,
  '\\Gamma_{f_{0}(1710)}': 0.13301019075808046,
  'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.0699249014701417-0.018664035501929042j),
 },
 parameter_errors={
  'm_{f_{0}(500)}': 0.006168655466103817,
  'm_{f_{0}(980)}': 0.0016283609785222876,
  'm_{f_{0}(1370)}': 0.005122588422790316,
  'm_{f_{0}(1500)}': 0.0033157863330869892,
  'm_{f_{0}(1710)}': 0.0025660827305775034,
  '\\Gamma_{f_{0}(500)}': 0.023838186430050128,
  '\\Gamma_{f_{0}(980)}': 0.003556673018336295,
  '\\Gamma_{f_{0}(1710)}': 0.007573518980113613,
  'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (0.04106392764099969+0.07043808181098646j),
 },
)

minuit2 = Minuit2()
fit_result_with_caching = minuit2.optimize(estimator_with_caching, initial_parameters)
fit_result_with_caching

FitResult(
 minimum_valid=True,
 execution_time=3.6658225059509277,
 function_calls=539,
 estimator_value=-4891.01730754809,
 parameter_values={
  'm_{f_{0}(500)}': 0.6102707294731716,
  'm_{f_{0}(980)}': 0.9902119846618569,
  'm_{f_{0}(1370)}': 1.3456300421927978,
  'm_{f_{0}(1500)}': 1.5050299510041418,
  'm_{f_{0}(1710)}': 1.7096496843680975,
  '\\Gamma_{f_{0}(500)}': 0.42260408077678696,
  '\\Gamma_{f_{0}(980)}': 0.06479339507977673,
  '\\Gamma_{f_{0}(1710)}': 0.13301019075895135,
  'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.069924901473717-0.018664035486070114j),
 },
 parameter_errors={
  'm_{f_{0}(500)}': 0.006168655451483166,
  'm_{f_{0}(980)}': 0.0016283609759060128,
  'm_{f_{0}(1370)}': 0.005122588414282541,
  'm_{f_{0}(1500)}': 0.0033157863009583644,
  'm_{f_{0}(1710)}': 0.0025660827200538303,
  '\\Gamma_{f_{0}(500)}': 0.023838186345858253,
  '\\Gamma_{f_{0}(980)}': 0.00355667300785808,
  '\\Gamma_{f_{0}(1710)}': 0.007573518972833387,
  'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (0.04106392765352627+0.07043808113241967j),
 },
)

Fit result analysis#

Hide code cell source 
intensity_func.update_parameters(fit_result.parameter_values)
compare_model("m_12", data, phsp, intensity_func)

Hide code cell source 
fit_traceback = pd.read_csv("fit_traceback.csv")
fig, (ax1, ax2) = plt.subplots(
    2, figsize=(7, 9), sharex=True, gridspec_kw={"height_ratios": [1, 2]}
)
fit_traceback.plot("function_call", "estimator_value", ax=ax1)
fit_traceback.plot("function_call", sorted(initial_parameters), ax=ax2)
fig.tight_layout()
ax2.set_xlabel("function call")
plt.show()

Zfit#

PDF definition#

import jax.numpy as jnp
import zfit  # suppress tf warnings
import zfit.z.numpy as znp
from zfit import supports, z

zfit.run.set_graph_mode(False)  # We cannot (yet) compile through the function
zfit.run.set_autograd_mode(False)


class TensorWavesPDF(zfit.pdf.BasePDF):
    def __init__(self, intensity, norm, obs, params=None, name="tensorwaves"):
        """tensorwaves intensity normalized over the *norm* dataset."""
        super().__init__(obs, params, name)
        self.intensity = intensity
        norm = {
            ob: jnp.asarray(ar)
            for ob, ar in zip(self.obs, z.unstack_x(norm), strict=True)
        }
        self.norm_sample = norm

    @supports(norm=True)
    def _pdf(self, x, norm):
        # we can also use better mechanics, where it automatically normalizes or not
        # this here is rather to take full control, it is always possible

        # updating the parameters of the model. This seems not very TF compatible?
        self.intensity.update_parameters({
            p.name: float(p) for p in self.params.values()
        })

        # converting the data to a dict for tensorwaves
        data = {
            ob: jnp.asarray(ar) for ob, ar in zip(self.obs, z.unstack_x(x), strict=True)
        }

        non_normalized_pdf = self.intensity(data)
        # this is not really needed, but can be useful for e.g. sampling with `pdf(..., norm_range=False)`
        if norm is False:
            out = non_normalized_pdf
        else:
            out = non_normalized_pdf / jnp.mean(self.intensity(self.norm_sample))
        return znp.asarray(out)

params = [
    zfit.param.convert_to_parameter(val, name, prefer_constant=False)
    for name, val in model.parameter_defaults.items()
]

def reset_parameters():
    for p in params_fit:
        if p.name in initial_parameters:
            p.set_value(initial_parameters[p.name])

obs = [
    zfit.Space(ob, limits=(np.min(data[ob]) - 1, np.max(data[ob]) + 1))
    for ob in pd.DataFrame(phsp)
]
obs_all = zfit.dimension.combine_spaces(*obs)

Data conversion#

phsp_zfit = zfit.Data.from_pandas(pd.DataFrame(phsp), obs=obs_all)
data_zfit = zfit.Data.from_pandas(pd.DataFrame(data), obs=obs_all)

Perform fit#

complex parameters need to be removed first:

params_fit = [p for p in params if p.name in initial_parameters if p.independent]

jax_intensity_func = create_parametrized_function(
    expression=unfolded_expression,
    parameters=model.parameter_defaults,
    backend="jax",
)

reset_parameters()

pdf = TensorWavesPDF(
    obs=obs_all,
    intensity=jax_intensity_func,
    norm=phsp_zfit,
    params={f"{p.name}": p for i, p in enumerate(params_fit)},
)
loss = zfit.loss.UnbinnedNLL(pdf, data_zfit)

minimizer = zfit.minimize.Minuit(gradient=True, mode=0)

Note

You can also try different minimizers, like ScipyTrustConstrV1, but Minuit seems to perform best.

%%time

result = minimizer.minimize(loss)
result

CPU times: user 22 s, sys: 188 ms, total: 22.2 s
Wall time: 8.56 s

FitResult of
<UnbinnedNLL model=[<zfit.<class '__main__.TensorWavesPDF'>  params=[\Gamma_{f_{0}(1710)}, \Gamma_{f_{0}(500)}, \Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[<zfit.core.data.Data object at 0x7fdc203d0430>] constraints=[]> 
with
<Minuit Minuit tol=0.001>

╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕
│ valid   │ converged   │ param at limit   │ edm     │ min value   │
╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡
│ True    │ True        │ False            │ 0.00041 │ -1871.035   │
╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛

Parameters
name                    value  (rounded)    at limit
--------------------  ------------------  ----------
m_{f_{0}(500)}                  0.608864       False
\Gamma_{f_{0}(500)}             0.419716       False
m_{f_{0}(980)}                  0.990038       False
\Gamma_{f_{0}(980)}            0.0643328       False
m_{f_{0}(1370)}                  1.35137       False
m_{f_{0}(1500)}                  1.50627       False
m_{f_{0}(1710)}                  1.70956       False
\Gamma_{f_{0}(1710)}            0.132484       False

%%time

result.hesse(name="hesse")
result

CPU times: user 2.5 s, sys: 12.6 ms, total: 2.51 s
Wall time: 953 ms

FitResult of
<UnbinnedNLL model=[<zfit.<class '__main__.TensorWavesPDF'>  params=[\Gamma_{f_{0}(1710)}, \Gamma_{f_{0}(500)}, \Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[<zfit.core.data.Data object at 0x7fdc203d0430>] constraints=[]> 
with
<Minuit Minuit tol=0.001>

╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕
│ valid   │ converged   │ param at limit   │ edm     │ min value   │
╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡
│ True    │ True        │ False            │ 0.00041 │ -1871.035   │
╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛

Parameters
name                    value  (rounded)        hesse    at limit
--------------------  ------------------  -----------  ----------
m_{f_{0}(500)}                  0.608864  +/-  0.0061       False
\Gamma_{f_{0}(500)}             0.419716  +/-   0.024       False
m_{f_{0}(980)}                  0.990038  +/-  0.0016       False
\Gamma_{f_{0}(980)}            0.0643328  +/-  0.0035       False
m_{f_{0}(1370)}                  1.35137  +/-  0.0039       False
m_{f_{0}(1500)}                  1.50627  +/-   0.002       False
m_{f_{0}(1710)}                  1.70956  +/-  0.0023       False
\Gamma_{f_{0}(1710)}            0.132484  +/-   0.007       False

%%time

result.errors(name="errors")
result

CPU times: user 45.3 s, sys: 393 ms, total: 45.7 s
Wall time: 17.2 s

FitResult of
<UnbinnedNLL model=[<zfit.<class '__main__.TensorWavesPDF'>  params=[\Gamma_{f_{0}(1710)}, \Gamma_{f_{0}(500)}, \Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[<zfit.core.data.Data object at 0x7fdc203d0430>] constraints=[]> 
with
<Minuit Minuit tol=0.001>

╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕
│ valid   │ converged   │ param at limit   │ edm     │ min value   │
╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡
│ True    │ True        │ False            │ 0.00041 │ -1871.035   │
╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛

Parameters
name                    value  (rounded)        hesse               errors    at limit
--------------------  ------------------  -----------  -------------------  ----------
m_{f_{0}(500)}                  0.608864  +/-  0.0061  -  0.006   + 0.0063       False
\Gamma_{f_{0}(500)}             0.419716  +/-   0.024  -  0.024   +  0.023       False
m_{f_{0}(980)}                  0.990038  +/-  0.0016  - 0.0016   + 0.0016       False
\Gamma_{f_{0}(980)}            0.0643328  +/-  0.0035  - 0.0034   + 0.0036       False
m_{f_{0}(1370)}                  1.35137  +/-  0.0039  - 0.0039   + 0.0039       False
m_{f_{0}(1500)}                  1.50627  +/-   0.002  -  0.002   +  0.002       False
m_{f_{0}(1710)}                  1.70956  +/-  0.0023  - 0.0024   + 0.0024       False
\Gamma_{f_{0}(1710)}            0.132484  +/-   0.007  - 0.0068   + 0.0073       False

Statistical inference using the hepstats library#

hepstats is built on top of zfit-interface:

from hepstats.hypotests import ConfidenceInterval
from hepstats.hypotests.calculators import AsymptoticCalculator
from hepstats.hypotests.parameters import POIarray

calculator = AsymptoticCalculator(result, minimizer)

We take one of the parameters as POI:

poi = pdf.params[r"\Gamma_{f_{0}(500)}"]
poi

<zfit.Parameter '\Gamma_{f_{0}(500)}' floating=True value=0.4197>

poi_null = POIarray(poi, np.linspace(poi - 0.1, poi + 0.1, 50))
ci = ConfidenceInterval(calculator, poi_null)
alpha = 0.328
ci.interval(alpha=alpha);

Confidence interval on \Gamma_{f_{0}(500)}:
	0.3964206394323228 < \Gamma_{f_{0}(500)} < 0.44257337109434974 at 67.2% C.L.

A helper function to plot the result:

def one_minus_cl_plot(x, pvalues, alpha=None, ax=None):
    if alpha is None:
        alpha = [0.32]
    if isinstance(alpha, float | int):
        alpha = [alpha]
    if ax is None:
        ax = plt.gca()

    ax.plot(x, pvalues, ".--")
    for a in alpha:
        ax.axhline(a, color="red", label="$\\alpha = " + str(a) + "$")
    ax.set_ylabel("1-CL")

    return ax

plt.figure(figsize=(9, 8))
one_minus_cl_plot(poi_null.values, ci.pvalues(), alpha=alpha)
plt.xlabel(f"${poi.name}$")
plt.show()