Amplitude analysis with zfit#
Show 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",
)
Show code cell source
import graphviz
dot = qrules.io.asdot(reaction, collapse_graphs=True)
graphviz.Source(dot)
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
Show code cell source
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
resonances = sorted(
reaction.get_intermediate_particles(),
key=lambda p: p.mass,
)
evenly_spaced_interval = np.linspace(0, 1, len(resonances))
colors = [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):
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#
Show code cell content
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
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 = [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):
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#
Show code cell source
intensity_func.update_parameters(fit_result.parameter_values)
compare_model("m_12", data, phsp, intensity_func)
Show 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))}
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))}
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()
