Amplitude analysis#

While TensorWaves can handle general mathematical expressions, it was originally created to perform amplitude analysis / Partial Wave Analysis (PWA), that is, to fit amplitude models to four-momenta data samples.

This notebook shows how to do an amplitude analysis with the ComPWA packages QRules, AmpForm, and TensorWaves. The ComPWA workflow generally consists of three stages:

1. Generate hit-and-miss data samples with this amplitude model.

Note

This notebook shows several tricks that can be helpful when doing an amplitude analysis. Most steps are optional though—they only serve to illustrate some tips that can be adopted and worked out further for specific analyses.

Step 1: Formulate model#

Whether generating data or fitting a model, TensorWaves takes mathematical expressions as input. When that expression is an amplitude model, it is most convenient to formulate it with qrules and ampform.

1.1 Generate transitions#

The first step is to generate all allowed transitions with QRules. In this example, we use the helicity formalism, but you can also use formalism="canonical-helicity". As you can see, we analyze the decay $$J/\psi \to \gamma\pi^0\pi^0$$ here.

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


Tip

See more advanced examples on QRules’ usage page.

1.2 Build amplitude model#

Next, we use ampform to formulate the transitions as an amplitude model (here: HelicityModel). This can be done with get_builder() and formulate():

import ampform

model_builder = ampform.get_builder(reaction)
model_no_dynamics = model_builder.formulate()
model_no_dynamics.intensity

$\displaystyle \sum_{m_{A}\in\left\{1,-1\right\}} \sum_{m_{0}\in\left\{1,-1\right\}} \sum_{m_{1}=0} \sum_{m_{2}=0}{\left|{{A^{12}}_{m_{A},m_{0},m_{1},m_{2}}}\right|^{2}}$
Hide code cell source
from ampform.io import aslatex
from IPython.display import Math

Math(aslatex(model_no_dynamics.amplitudes))

$\begin{split}\displaystyle \begin{array}{rcl} {A^{12}}_{-1,-1,0,0} &=& C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,-1}\left(- \phi_{0},\theta_{0},0\right) \\ {A^{12}}_{-1,1,0,0} &=& C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{-1,1}\left(- \phi_{0},\theta_{0},0\right) \\ {A^{12}}_{1,-1,0,0} &=& C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,-1}\left(- \phi_{0},\theta_{0},0\right) \\ {A^{12}}_{1,1,0,0} &=& C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) + C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} D^{0}_{0,0}\left(- \phi^{12}_{1},\theta^{12}_{1},0\right) D^{1}_{1,1}\left(- \phi_{0},\theta_{0},0\right) \\ \end{array}\end{split}$
Hide code cell source
Math(aslatex(model_no_dynamics.parameter_defaults))

$\begin{split}\displaystyle \begin{array}{rcl} C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ \end{array}\end{split}$

The heart of the model is a sympy expression that contains the full description of the intensity model. Note two things:

1. The coefficients for the different amplitudes are complex valued.

2. By default there is no dynamics in the model, so it still has to be specified.

We choose to use relativistic_breit_wigner_with_ff() as the lineshape for all resonances and use a Blatt-Weisskopf form factor (create_non_dynamic_with_ff()) for the production decay. The set_dynamics() is a convenience interface for replacing the dynamics for intermediate states.

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

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


Now let’s take another look at the parameters of the model to see which new parameters are there:

Hide code cell source
sorted_parameter_defaults = {
symbol: model.parameter_defaults[symbol]
for symbol in sorted(model.parameter_defaults, key=str)
}
src = aslatex(sorted_parameter_defaults)
Math(src)

$\begin{split}\displaystyle \begin{array}{rcl} C_{J/\psi(1S) \to {f_{0}(1370)}_{0} \gamma_{+1}; f_{0}(1370) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(1710)}_{0} \gamma_{+1}; f_{0}(1710) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(2020)}_{0} \gamma_{+1}; f_{0}(2020) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(500)}_{0} \gamma_{+1}; f_{0}(500) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ C_{J/\psi(1S) \to {f_{0}(980)}_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} &=& 1+0i \\ \Gamma_{f_{0}(1370)} &=& 0.35 \\ \Gamma_{f_{0}(1500)} &=& 0.108 \\ \Gamma_{f_{0}(1710)} &=& 0.15 \\ \Gamma_{f_{0}(2020)} &=& 0.44 \\ \Gamma_{f_{0}(500)} &=& 0.45 \\ \Gamma_{f_{0}(980)} &=& 0.06 \\ d_{J/\psi(1S)} &=& 1 \\ d_{f_{0}(1370)} &=& 1 \\ d_{f_{0}(1500)} &=& 1 \\ d_{f_{0}(1710)} &=& 1 \\ d_{f_{0}(2020)} &=& 1 \\ d_{f_{0}(500)} &=& 1 \\ d_{f_{0}(980)} &=& 1 \\ m_{f_{0}(1370)} &=& 1.35 \\ m_{f_{0}(1500)} &=& 1.522 \\ m_{f_{0}(1710)} &=& 1.733 \\ m_{f_{0}(2020)} &=& 1.982 \\ m_{f_{0}(500)} &=& 0.6 \\ m_{f_{0}(980)} &=& 0.99 \\ \end{array}\end{split}$

Optionally, we can backup the HelicityModel to disk via pickle. The ReactionInfo object (which takes longest to generate) can also be pickled, or it can also be serialized to JSON:

import pickle

qrules.io.write(reaction, "transitions.json")
with open("helicity_model.pickle", "wb") as stream:
pickle.dump(model, stream)


In the next steps, we use this HelicityModel as a template for a computational function to generate data and to perform a fit.

Tip

See more advanced examples on AmpForm’s usage page.

Step 2: Generate data#

In this section, we use the HelicityModel that we created with ampform in the previous step to generate a data sample via hit & miss Monte Carlo. We do this with the data module.

Optionally, we can load() the HelicityModel that was created in the previous step. This does not have to be done if the model has been generated in the same script or notebook (like in this notebook), but can be useful if the model was generated elsewhere.

import pickle

from ampform.helicity import HelicityModel

with open("helicity_model.pickle", "rb") as model_file:

Hide code cell source
initial_state, *_ = imported_model.reaction_info.initial_state.values()
print("Initial state:")
print(" ", initial_state.name)
print("Final state:")
for i, p in imported_model.reaction_info.final_state.items():
print(f"  {i}: {p.name}")
del initial_state

Initial state:
J/psi(1S)
Final state:
0: gamma
1: pi0
2: pi0


2.1 Generate phase space sample#

The ReactionInfo class defines the constraints of the phase space. As such, we have enough information to generate a phase-space sample for this particle reaction. We do this with a TFPhaseSpaceGenerator class, which is a DataGenerator for a DataSample of four-momenta arrays (using tensorflow and the phasespace package as a back-end). We also need to construct a RealNumberGenerator that can generate random numbers. TFUniformRealNumberGenerator is the natural choice here.

As opposed to the main Step 2: Generate data of the main usage example page, we will generate a deterministic data sample. This can be done by feeding a RealNumberGenerator with a specific seed and giving that generator to the TFPhaseSpaceGenerator.generate() method:

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)

Hide code cell source
import numpy as np
import pandas as pd

pd.DataFrame(
{
(k, label): np.transpose(v)[i]
for k, v in phsp_momenta.items()
for i, label in enumerate(["E", "px", "py", "pz"])
}
)

p0 p1 p2
E px py pz E px py pz E px py pz
0 0.811317 0.787621 0.144327 0.130609 1.092959 -0.019148 -0.531113 -0.945459 1.192624 -0.768473 0.386786 0.814849
1 0.602303 -0.343019 -0.367933 0.331259 1.431153 1.173457 -0.078558 -0.804244 1.063444 -0.830437 0.446491 0.472984
2 1.165704 0.115778 0.666757 -0.949155 1.443427 0.192011 -0.593016 1.294884 0.487770 -0.307789 -0.073741 -0.345729
3 0.593783 0.392870 0.407210 -0.180034 1.246726 0.037716 0.106090 1.234273 1.256391 -0.430586 -0.513300 -1.054238
4 0.864272 -0.668511 0.516082 0.183623 0.754404 -0.716281 -0.135528 0.139574 1.478224 1.384792 -0.380554 -0.323197
... ... ... ... ... ... ... ... ... ... ... ... ...
99995 0.472639 0.250787 0.395665 0.062791 1.126885 0.041148 0.795625 0.785454 1.497377 -0.291935 -1.191290 -0.848245
99996 1.330392 1.185399 0.538998 -0.272492 0.419241 -0.033333 0.139318 0.370167 1.347267 -1.152066 -0.678316 -0.097675
99997 0.864709 0.258491 0.294117 -0.770973 1.444051 -0.804658 0.065531 1.189662 0.788140 0.546167 -0.359648 -0.418689
99998 0.607279 0.563725 0.219317 -0.053868 0.968601 0.899961 0.061514 0.325968 1.521020 -1.463686 -0.280831 -0.272100
99999 1.148659 -0.481566 0.715177 0.758969 1.291588 0.469209 -1.158960 -0.294342 0.656653 0.012357 0.443783 -0.464627

100000 rows × 12 columns

The resulting phase space sample is a dict of final state IDs to an array of four-momenta. In the last step, we converted this sample in such a way that it is rendered as an understandable pandas.DataFrame.

Four-momentum arrays

Kinematic expressions from AmpForm that involve four-momenta should be formatted as $$\left(E, p_x, p_y, p_z\right)$$. In addition, the shape of input arrays should be (n, 4) with n the number of events.

Warning

When using the helicity formalism, the sum of the four-momenta should be in the rest frame, that is $$\sum_i p_i = \left(m_A, 0, 0, 0\right)$$ with $$m_A$$ the mass of the decaying particle $$A$$. This is because the helicity formalisms boosts through the decay chain starting from particle $$A$$. Take care to boost your experimental data into the rest frame, optionally following the kinematics classes provided by AmpForm.

import numpy as np

p = np.array(list(phsp_momenta.values()))
p.shape

(3, 100000, 4)

p.sum(axis=0).round(decimals=14)

array([[ 3.0969,  0.    , -0.    , -0.    ],
[ 3.0969,  0.    , -0.    ,  0.    ],
[ 3.0969,  0.    , -0.    ,  0.    ],
...,
[ 3.0969, -0.    , -0.    ,  0.    ],
[ 3.0969, -0.    , -0.    , -0.    ],
[ 3.0969, -0.    , -0.    ,  0.    ]])


2.2 Generate intensity-based sample#

‘Data samples’ are more complicated than phase space samples in that they represent the intensity profile resulting from a reaction. You therefore need a Function object that expresses an intensity distribution as well as a phase space over which to generate that distribution. We call such a data sample an intensity-based sample.

An intensity-based sample is generated over a phase space sample using the IntensityDistributionGenerator. Its usage is similar to TFPhaseSpaceGenerator, but now you have to provide a Function as well as a DataTransformer that is used to transform the four-momentum phase space sample to a data sample that can be understood by the Function.

Now, recall that in Step 1: Formulate model, we used the helicity formalism to mathematically express the reaction in terms of an amplitude model. TensorWaves needs to convert this HelicityModel to a Function object that can perform fast computations. This can be done with create_parametrized_function():

from tensorwaves.function.sympy import create_parametrized_function

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


Tip

If create_parametrized_function() takes a long time, have a look at Speed up lambdifying.

Hit & miss

A problem is that ParametrizedBackendFunction takes a DataSample with kinematic variables for the helicity formalism as input, not a set of four-momenta. We therefore need to construct a DataTransformer to transform these four-momenta to function variables. In this case, we work with the helicity formalism, so we construct a SympyDataTransformer. This numerical data transformer is created from symbolic expressions, in this case, from expressions that relate four-momenta to helicity angles and invariant masses:

Hide code cell source
Math(aslatex(model.kinematic_variables))

$\begin{split}\displaystyle \begin{array}{rcl} m_{0} &=& m_{{p}_{0}} \\ m_{1} &=& m_{{p}_{1}} \\ m_{2} &=& m_{{p}_{2}} \\ m_{012} &=& m_{{p}_{012}} \\ m_{12} &=& m_{{p}_{12}} \\ \phi_{0} &=& \phi\left({p}_{12}\right) \\ \phi^{12}_{1} &=& \phi\left(\boldsymbol{B_z}\left(\frac{\left|\vec{{p}_{12}}\right|}{E\left({p}_{12}\right)}\right) \boldsymbol{R_y}\left(- \theta\left({p}_{12}\right)\right) \boldsymbol{R_z}\left(- \phi\left({p}_{12}\right)\right) p_{1}\right) \\ \theta_{0} &=& \theta\left({p}_{12}\right) \\ \theta^{12}_{1} &=& \theta\left(\boldsymbol{B_z}\left(\frac{\left|\vec{{p}_{12}}\right|}{E\left({p}_{12}\right)}\right) \boldsymbol{R_y}\left(- \theta\left({p}_{12}\right)\right) \boldsymbol{R_z}\left(- \phi\left({p}_{12}\right)\right) p_{1}\right) \\ \end{array}\end{split}$
from tensorwaves.data import SympyDataTransformer

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


That’s it, now we have enough info to create an intensity-based data sample. Notice how the structure of the output data is the same as the phase-space sample we generated previously:

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=intensity_func,
domain_transformer=helicity_transformer,
)
data_momenta = data_generator.generate(10_000, rng)
pd.DataFrame(
{
(k, label): np.transpose(v)[i]
for k, v in data_momenta.items()
for i, label in enumerate(["E", "px", "py", "pz"])
}
)

weights p0 p1 p2
E px py pz E px py pz E px py pz E px py pz
0 0.259658 0.366444 0.390128 0.400137 1.331059 0.410932 0.311819 -1.227038 0.331184 -0.249848 -0.039500 -0.165771 1.434657 -0.161084 -0.272319 1.392808
1 0.259658 0.366444 0.390128 0.400137 1.494125 -0.195880 -1.307833 -0.695422 0.824596 0.072240 0.798756 0.136049 0.778180 0.123640 0.509077 0.559374
2 0.259658 0.366444 0.390128 0.400137 1.392857 -0.802316 0.178701 -1.124459 0.804786 -0.013665 -0.063677 0.790709 0.899256 0.815980 -0.115023 0.333749
3 0.259658 0.366444 0.390128 0.400137 1.187149 0.751684 -0.359411 -0.845646 0.862093 0.159626 -0.139034 0.824727 1.047659 -0.911310 0.498444 0.020920
4 0.259658 0.366444 0.390128 0.400137 1.326583 -1.116556 -0.650988 -0.298899 1.416935 0.996462 0.954615 0.291996 0.353382 0.120094 -0.303627 0.006904
... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
9995 0.259658 0.366444 0.390128 0.400137 1.399972 -0.184332 -0.530811 1.282257 1.326253 0.478162 0.546242 -1.101685 0.370675 -0.293831 -0.015432 -0.180572
9996 0.259658 0.366444 0.390128 0.400137 1.219380 0.235560 0.130632 1.189258 0.993276 0.435025 -0.472702 -0.745443 0.884243 -0.670585 0.342070 -0.443815
9997 0.259658 0.366444 0.390128 0.400137 1.264992 0.790155 0.874115 -0.460200 1.174563 -0.802853 -0.318591 0.784415 0.657345 0.012698 -0.555524 -0.324215
9998 0.259658 0.366444 0.390128 0.400137 1.530396 1.331257 0.223952 -0.720912 1.109122 -0.931223 -0.128974 0.572820 0.457382 -0.400034 -0.094978 0.148092
9999 0.259658 0.366444 0.390128 0.400137 1.251299 0.216628 -1.176751 0.366166 0.580159 0.329024 0.297080 0.349073 1.265442 -0.545653 0.879671 -0.715239

10000 rows × 16 columns

As before, we use a RealNumberGenerator with a specific seed to ensure we get a deterministic data sample.

Note

Instead of constructing a TFWeightedPhaseSpaceGenerator, it’s also possible to just reuse the unweighted TFPhaseSpaceGenerator that we constructed previously. This is, however, a bit slower, because the TFPhaseSpaceGenerator also uses a hit-and-miss strategy.

2.3 Visualize kinematic variables#

We now have a phase space sample and an intensity-based sample. Their data structure isn’t the most informative though: it’s just a collection of four-momentum tuples. But we can again use the SympyDataTransformer to convert these four-momenta to (in the case of the helicity formalism) invariant masses and helicity angles:

phsp = helicity_transformer(phsp_momenta)
data = helicity_transformer(data_momenta)
list(data)

['m_0',
'm_1',
'm_2',
'm_012',
'm_12',
'phi_0',
'phi_1^12',
'theta_0',
'theta_1^12']


Note

Check the remark about four-momentum format and rest frame of the decaying particle here.

The DataSample is a mapping of kinematic variables names to a 1-dimensional array of values. The numbers you see here are final state IDs as defined in the HelicityModel member of the HelicityModel:

Hide code cell source
for state_id, particle in reaction.final_state.items():
print(f"ID {state_id}:", particle.name)

ID 0: gamma
ID 1: pi0
ID 2: pi0


The DataSample can easily be converted to a pandas.DataFrame:

import pandas as pd

data_frame = pd.DataFrame(data)
phsp_frame = pd.DataFrame(data)
data_frame

m_0 m_1 m_2 m_012 m_12 phi_0 phi_1^12 theta_0 theta_1^12
0 2.107342e-08+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.160377+0.000000j -2.492477 -0.417233 0.397967 2.591367
1 0.000000e+00+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 0.580070+0.000000j 1.422127 0.183725 1.086667 1.535691
2 0.000000e+00+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 0.981687+0.000000j -0.219154 -3.002802 0.631228 1.641399
3 0.000000e+00+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.495937+0.000000j 2.695585 3.063622 0.777978 1.730394
4 0.000000e+00+2.107342e-08j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.172263+0.000000j 0.527850 1.515685 1.343530 0.602596
... ... ... ... ... ... ... ... ... ...
9995 2.580957e-08+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 0.958980+0.000000j 1.236561 -2.139350 2.728581 0.779400
9996 1.490116e-08+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.427653+0.000000j -2.635255 1.107227 2.918859 1.479611
9997 0.000000e+00+1.490116e-08j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.325021+0.000000j -2.305789 -2.436526 1.198456 1.139961
9998 2.107342e-08+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 0.334396+0.000000j -2.974927 -2.730497 1.080301 0.764671
9999 1.490116e-08+0.000000e+00j 0.134977+0.000000j 0.134977+0.000000j 3.0969+0.0000j 1.356648+0.000000j 1.752848 -2.387607 1.867771 2.163774

10000 rows × 9 columns

This also means that we can use all kinds of fancy plotting functionality of for instance matplotlib.pyplot to see what’s going on. Here’s an example:

Hide code cell source
import matplotlib.pyplot as plt
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_frame["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()


2.4 Export data sets#

To export the generated data samples, simply pickle.dump() them as follows:

import pickle

with open("data.pickle", "wb") as stream:
pickle.dump(data, stream)
with open("phsp.pickle", "wb") as stream:
pickle.dump(phsp, stream)


In the next step, we illustrate how to optimize() the intensity model to these data samples.

Step 3: Perform fit#

As explained in the previous step, a ParametrizedFunction can compute a list of intensities (real numbers) for an input DataSample. At this stage, we want to optimize the parameters of this ParametrizedFunction, so that it matches the distribution of our data sample. This is what we call ‘fitting’.

Hide code cell source
reaction = qrules.io.load("transitions.json")
with open("helicity_model.pickle", "rb") as stream:
with open("data.pickle", "rb") as stream:
with open("phsp.pickle", "rb") as stream:


3.1 Prepare parametrized function#

In principle, we can use the same ParametrizedFunction as the one that we created in 2.2 Generate intensity-based sample. However, when fitting such a function to a data distribution, an Optimizer will evaluate this function numerous times, so it is smart to apply some optimizations to the underlying expression tree before-hand.

Tip

The sections below illustrate some tricks for how to simplify the expression tree underneath a ParametrizedFunction. Most of this can also be achieved with create_cached_function(), which is illustrated in Constant sub-expressions and under Simplified procedure: create_cached_function(). But note that it is still smart to cast complex-valued data.

Determine free parameters#

It often happens that not all parameters in a ParametrizedFunction have to be optimized. These parameters are called fixed parameters, while the ones that will be optimized are called free parameters. In the fit example here, we will optimize the following free parameters, starting with a certain initial value.

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


To make the fit more interesting, we give these initial values a small offset compared to the suggested parameter_defaults―after all, we are fitting to a data sample that was generated with this very same ParametrizedFunction.

Complex-valued parameters

If initial parameter values are complex, the parameter is split into a real and an imaginary part during the fit. See also Covariance matrix.

Collapsing constant sub-trees#

Having decided which parameters are to be optimized, we can apply some additional optimizations to the amplitude model expression, so that the computations during the fit are faster. The first of these optimizations is to substitute the parameters that remain fixed with their suggested parameter values. Note how this decreases the number of operations (nodes) in the expression tree:

import sympy as sp

sp.count_ops(unfolded_expression)

2023

free_parameters = {
p for p in model.parameter_defaults if p.name in initial_parameters
}
fixed_parameters = {
p: v for p, v in model.parameter_defaults.items() if p not in free_parameters
}
optimized_expression = unfolded_expression.subs(fixed_parameters)
sp.count_ops(optimized_expression)

1863


Note that there are a few irrational numbers (square roots) in the expression as well, for example in this amplitude:

optimized_expression.args[0].args[0].args[0].args[0]

$\displaystyle \frac{0.43604 \sqrt{2} \sqrt{\frac{\left(m_{012}^{2} - \left(m_{0} - m_{12}\right)^{2}\right) \left(m_{012}^{2} - \left(m_{0} + m_{12}\right)^{2}\right)}{m_{012}^{2} \cdot \left(1 + \frac{\left(m_{012}^{2} - \left(m_{0} - m_{12}\right)^{2}\right) \left(m_{012}^{2} - \left(m_{0} + m_{12}\right)^{2}\right)}{4 m_{012}^{2}}\right)}} \left(\frac{\cos{\left(\theta_{0} \right)}}{2} + \frac{1}{2}\right) e^{i \phi_{0}}}{m_{12}^{2} - 3.928324 + \frac{3.42581279392 i \sqrt{\frac{\left(m_{12}^{2} - \left(m_{1} - m_{2}\right)^{2}\right) \left(m_{12}^{2} - \left(m_{1} + m_{2}\right)^{2}\right)}{m_{12}^{2}}}}{m_{12} \sqrt{\left(3.928324 - \left(m_{1} - m_{2}\right)^{2}\right) \left(3.928324 - \left(m_{1} + m_{2}\right)^{2}\right)}}}$

When we substitute these values, the expression tree becomes even smaller:

for node in sp.preorder_traversal(optimized_expression):
if node.free_symbols:
continue
if isinstance(node, sp.Pow):
optimized_expression = optimized_expression.xreplace({node: node.n()})
sp.count_ops(optimized_expression)

1791


Substitute scalar masses#

Since the final state particles $$\gamma, \pi^0, \pi^0$$ are stable, we can substitute their invariant masses $$m_0, m_1, m_2$$ with scalar values. SymPy will then further simply the expression, so that the computation is less complex.

mass_substitutions = {
sp.Symbol(f"m_{i}", nonnegative=True): particle.mass
for i, particle in reaction.final_state.items()
}
optimized_expression = optimized_expression.subs(mass_substitutions)
sp.count_ops(optimized_expression)

1139


Cast complex-valued data#

Note that some of the computed kinematic variable arrays are complex-valued. This can be useful in come cases, but in this decay channel, all kinematic variable values lie on the real axis. The fit can therefore be further optimized by casting the data arrays to real values:

from tensorwaves.interface import DataSample

def safe_downcast_to_real(data: DataSample) -> DataSample:
# using isrealobj instead of real_if_close to keep same array backend
return {k: v.real if np.isrealobj(v) else v for k, v in data.items()}

data_real = safe_downcast_to_real(data)
phsp_real = safe_downcast_to_real(phsp)


Create computational function#

Finally, we again use create_parametrized_function() to create a ParametrizedFunction for this optimized expression:

optimized_function = create_parametrized_function(
optimized_expression, model.parameter_defaults, backend="jax"
)


Note that the intensities computed by the optimized function are indeed the same as the original intensity function that was created in 2.2 Generate intensity-based sample and that it is much faster!

# JIT-compile functions and test equality
np.testing.assert_array_almost_equal(
optimized_function(data_real),
intensity_func(data_real),
decimal=13,
)

%timeit -n1 intensity_func(data)
%timeit -n1 optimized_function(data_real)

15.6 ms ± 198 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
5.93 ms ± 40.2 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)


The reason is that several sub-trees in the original expression tree have been collapsed, which results in fewer computations in the backend. Here’s one of the sub-amplitudes, taken from the total expression:

Original expression tree#
Hide code cell source
dot = sp.dotprint(unfolded_expression.args[0].args[0].args[0].args[0])
graphviz.Source(dot)

Optimized expression tree#
Hide code cell source
dot = sp.dotprint(optimized_expression.args[0].args[0].args[0].args[0])
graphviz.Source(dot)


Simplified procedure: create_cached_function()#

As noted under 3.1 Prepare parametrized function, most of what is described in this section can be achieved with the function create_cached_function(). The idea is described on Constant sub-expressions, but in this section, we show how this translates to amplitude analysis.

First, note that create_cached_function() works with mappings and iterable of sympy.Symbol and not with the str mappings that we defined in Determine free parameters. We can convert the parameter names back to Symbols as follows:

free_parameter_symbols = [
symbol
for symbol in model.parameter_defaults
if symbol.name in set(initial_parameters)
]


This gives us all the information we need to create a cached function, convert the data samples to cached data samples and construct an Estimator with these transformed items (compare 3.2 Define estimator):

from tensorwaves.estimator import UnbinnedNLL, create_cached_function

cached_function, cache_transformer = create_cached_function(
unfolded_expression,
parameters=model.parameter_defaults,
free_parameters=free_parameter_symbols,
backend="jax",
)
cached_data = cache_transformer(data_real)
cached_phsp = cache_transformer(phsp_real)
estimator_with_caching = UnbinnedNLL(
cached_function,
data=cached_data,
phsp=cached_phsp,
backend="jax",
)


Note that, just like in Create computational function, the computed intensities of both the original intensity function and the cached function are indeed the same:

np.testing.assert_array_almost_equal(
cached_function(cached_data),
intensity_func(data_real),
decimal=13,
)

%timeit -n1 intensity_func(data_real)
%timeit -n1 cached_function(cached_data)

15.5 ms ± 195 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
5.45 ms ± 68.3 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)


3.2 Define estimator#

To perform a fit, you need to define an Estimator. This is a measure for the discrepancy between the ParametrizedFunction and the data distribution to which you fit it. In PWA, we usually use an unbinned negative log likelihood estimator (UnbinnedNLL).

Generally, the ParametrizedFunction is not normalized with regards to the data sample, while a log likelihood estimator requires a normalized function. This is where the phase space data comes into play again: the ParametrizedFunction is evaluated over the phase space data, so that its output can be used as a normalization factor.

from tensorwaves.estimator import UnbinnedNLL

estimator = UnbinnedNLL(
optimized_function,
data=data_real,
phsp=phsp_real,
backend="jax",
)


Note that the UnbinnedNLL can be expressed with different backends, because it uses statistical operations like log and mean. Here, we use jax, which turns out to be the fastest backend for this model.

3.3 Optimize fit parameters#

Starting the fit itself is quite simple: just create an optimizer instance of your choice and call its optimize() method to start the fitting process. The optimize() method requires a mapping of parameter names to their initial values. Only the parameters listed in the mapping are optimized.

Let’s have a look at our first guess for the parameter values. Recall that a ParametrizedFunction object computes the intensity for a certain DataSample. This can be seen nicely when we use these intensities as weights on the phase space sample and plot it together with the original data sample. Here, we look at the invariant mass distribution projection of the final states 1 and 2, which, as we saw before, is the final state particle pair $$\pi^0\pi^0$$.

Don’t forget to use update_parameters() first!

Hide code cell content
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]

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 = optimized_function.parameters
optimized_function.update_parameters(initial_parameters)
compare_model("m_12", data_real, phsp_real, optimized_function)


Finally, we are ready to create an Optimizer that can optimize() the parameters in the ParametrizedFunction (which is embedded in the Estimator). Here, we choose the Minuit2 optimizer, which is the most common optimizer in high-energy physics (see also Perform fit with different optimizers). Since we constructed the UnbinnedNLL with jax, can an optimize the model with analytic gradient over the ParametrizedBackendFunction:

Note

The computation time depends on the complexity of the model, the number of data events, the size of the phase space sample, and the number of free parameters. This model is rather small and has but a few free parameters, so the optimization shouldn’t take more than a minute.

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

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

FitResult(
minimum_valid=True,
execution_time=40.43044304847717,
function_calls=581,
estimator_value=-2818.4110168059096,
parameter_values={
'm_{f_{0}(500)}': 0.6190451065240363,
'm_{f_{0}(980)}': 0.9903165662098864,
'm_{f_{0}(1370)}': 1.3469935411809348,
'm_{f_{0}(1500)}': 1.5184796602617077,
'm_{f_{0}(1710)}': 1.7445194835233515,
'\\Gamma_{f_{0}(500)}': 0.41722348489436895,
'\\Gamma_{f_{0}(980)}': 0.06393661146054923,
'\\Gamma_{f_{0}(1710)}': 0.1522819977397501,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.0410382201695805-0.04219995344283251j),
},
parameter_errors={
'm_{f_{0}(500)}': 0.007700201525274844,
'm_{f_{0}(980)}': 0.001968628625848101,
'm_{f_{0}(1370)}': 0.005929322670604979,
'm_{f_{0}(1500)}': 0.003974046469601777,
'm_{f_{0}(1710)}': 0.003092343574647879,
'\\Gamma_{f_{0}(500)}': 0.030597276293004426,
'\\Gamma_{f_{0}(980)}': 0.004379012917021767,
'\\Gamma_{f_{0}(1710)}': 0.012186008408374435,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (0.04466380767510002+0.07610105377542067j),
},
)


See Minuit2 for how to tweak the internal iminuit.Minuit optimizer.

As can be seen, the values of the optimized parameters in the FitResult are again comparable to the original parameter values.

Hide code cell source
optimized_parameters = fit_result.parameter_values
for p in optimized_parameters:
print(p)
print(f"  initial:   {initial_parameters[p]:.3}")
print(f"  optimized: {optimized_parameters[p]:.3}")
print(f"  original:  {original_parameters[p]:.3}")

m_{f_{0}(500)}
initial:   0.4
optimized: 0.619
original:  0.6
m_{f_{0}(980)}
initial:   0.88
optimized: 0.99
original:  0.99
m_{f_{0}(1370)}
initial:   1.22
optimized: 1.35
original:  1.35
m_{f_{0}(1500)}
initial:   1.45
optimized: 1.52
original:  1.52
m_{f_{0}(1710)}
initial:   1.83
optimized: 1.74
original:  1.73
\Gamma_{f_{0}(500)}
initial:   0.3
optimized: 0.417
original:  0.45
\Gamma_{f_{0}(980)}
initial:   0.1
optimized: 0.0639
original:  0.06
\Gamma_{f_{0}(1710)}
initial:   0.3
optimized: 0.152
original:  0.15
C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}
initial:   (1+0j)
optimized: (1.04-0.0422j)
original:  (1+0j)


3.4 Export and import#

In 3.3 Optimize fit parameters, we initialized Minuit2 with a Loadable callback. Such callback classes offer the possibility to load_latest_parameters(), so you can pick up the optimize process in case it crashes or if you pause it. Loading the latest parameters goes as follows:

latest_parameters = CSVSummary.load_latest_parameters("fit_traceback.csv")
latest_parameters

{'iteration': '',
'function_call': 581,
'time': '2023-11-08 14:14:33.147425',
'optimizer': 'Minuit2',
'estimator_type': 'UnbinnedNLL',
'estimator_value': -2818.4110168059096,
'm_{f_{0}(500)}': 0.6190451065240363,
'm_{f_{0}(980)}': 0.9903165662098864,
'm_{f_{0}(1370)}': 1.3469935411809348,
'm_{f_{0}(1500)}': 1.5184796602617077,
'm_{f_{0}(1710)}': 1.7445194835233515,
'\\Gamma_{f_{0}(500)}': 0.41722348489436895,
'\\Gamma_{f_{0}(980)}': 0.06393661146054923,
'\\Gamma_{f_{0}(1710)}': 0.1522819977397501,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.0410382201695805-0.04219995344283251j)}


Callbacks and this example of a (custom) plotting callback.

3.5 Analyze fit result#

Plot optimized model#

Using the same method as above, we renew the parameters of the ParametrizedFunction and plot it again over the phase space sample.

optimized_function.update_parameters(fit_result.parameter_values)
compare_model("m_12", data_real, phsp_real, optimized_function)


Covariance matrix#

Each of the optimizers offer more specific information about the fit result. This information can be accessed with FitResult.specifics. A common example would be to get the covariance matrix and correlation matrix:

covariance_matrix = fit_result.specifics.covariance
covariance_matrix.correlation()

 real_C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} imag_C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} m_{f_{0}(500)} m_{f_{0}(980)} m_{f_{0}(1370)} m_{f_{0}(1500)} m_{f_{0}(1710)} \Gamma_{f_{0}(500)} \Gamma_{f_{0}(980)} \Gamma_{f_{0}(1710)} 1 -0.2 0.1 0 -0.6 -0.2 0.1 0.1 0 0.1 -0.2 1 -0.1 -0.1 0.1 0.8 -0.3 0.1 -0.1 0.6 0.1 -0.1 1 0.1 0 -0.1 0.1 -0.5 0 0 0 -0.1 0.1 1 0 0 0 -0.1 -0.1 0 -0.6 0.1 0 0 1 0.3 0.1 0 0.2 -0.2 -0.2 0.8 -0.1 0 0.3 1 -0.1 0 0 0.3 0.1 -0.3 0.1 0 0.1 -0.1 1 -0.1 0 -0.2 0.1 0.1 -0.5 -0.1 0 0 -0.1 1 -0.4 0.1 0 -0.1 0 -0.1 0.2 0 0 -0.4 1 0.1 0.1 0.6 0 0 -0.2 0.3 -0.2 0.1 0.1 1

Note

Optimizers can only work with real numbers. For this reason, complex-valued parameters are split into a real and an imaginary part. This becomes apparent if we look at the covariance matrix.

AIC and BIC#

The AIC and BIC give us another quantification of the quality of the fit. Here’s how to compute them from this FitResult:

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

AIC: -5616.822033611819
BIC: -5544.718629892058


Extract intensity components#

Notice that the HelicityModel contains components attribute. This is dict of component names to sympy.Exprs helps us to identify amplitudes and intensities in the total amplitude.

sorted(model.components)

['A_{J/\\psi(1S)_{+1} \\to {f_{0}(1370)}_{0} \\gamma_{+1}; {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(1370)}_{0} \\gamma_{-1}; {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(1500)}_{0} \\gamma_{+1}; {f_{0}(1500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(1500)}_{0} \\gamma_{-1}; {f_{0}(1500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(1710)}_{0} \\gamma_{+1}; {f_{0}(1710)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(1710)}_{0} \\gamma_{-1}; {f_{0}(1710)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(2020)}_{0} \\gamma_{+1}; {f_{0}(2020)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(2020)}_{0} \\gamma_{-1}; {f_{0}(2020)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(500)}_{0} \\gamma_{+1}; {f_{0}(500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(500)}_{0} \\gamma_{-1}; {f_{0}(500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(980)}_{0} \\gamma_{+1}; {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{+1} \\to {f_{0}(980)}_{0} \\gamma_{-1}; {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1370)}_{0} \\gamma_{+1}; {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1370)}_{0} \\gamma_{-1}; {f_{0}(1370)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1500)}_{0} \\gamma_{+1}; {f_{0}(1500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1500)}_{0} \\gamma_{-1}; {f_{0}(1500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1710)}_{0} \\gamma_{+1}; {f_{0}(1710)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(1710)}_{0} \\gamma_{-1}; {f_{0}(1710)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(2020)}_{0} \\gamma_{+1}; {f_{0}(2020)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(2020)}_{0} \\gamma_{-1}; {f_{0}(2020)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(500)}_{0} \\gamma_{+1}; {f_{0}(500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(500)}_{0} \\gamma_{-1}; {f_{0}(500)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(980)}_{0} \\gamma_{+1}; {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'A_{J/\\psi(1S)_{-1} \\to {f_{0}(980)}_{0} \\gamma_{-1}; {f_{0}(980)}_{0} \\to \\pi^{0}_{0} \\pi^{0}_{0}}',
'I_{J/\\psi(1S)_{+1} \\to \\gamma_{+1} \\pi^{0}_{0} \\pi^{0}_{0}}',
'I_{J/\\psi(1S)_{+1} \\to \\gamma_{-1} \\pi^{0}_{0} \\pi^{0}_{0}}',
'I_{J/\\psi(1S)_{-1} \\to \\gamma_{+1} \\pi^{0}_{0} \\pi^{0}_{0}}',
'I_{J/\\psi(1S)_{-1} \\to \\gamma_{-1} \\pi^{0}_{0} \\pi^{0}_{0}}']

model.components[
R"A_{J/\psi(1S)_{+1} \to {f_{0}(1370)}_{0} \gamma_{+1};"
R" {f_{0}(1370)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}"
].subs(model.parameter_defaults).doit()

$\displaystyle \frac{0.23625 \sqrt{2} \sqrt{\frac{\left(m_{012}^{2} - \left(m_{0} - m_{12}\right)^{2}\right) \left(m_{012}^{2} - \left(m_{0} + m_{12}\right)^{2}\right)}{m_{012}^{2} \cdot \left(1 + \frac{\left(m_{012}^{2} - \left(m_{0} - m_{12}\right)^{2}\right) \left(m_{012}^{2} - \left(m_{0} + m_{12}\right)^{2}\right)}{4 m_{012}^{2}}\right)}} \left(\frac{\cos{\left(\theta_{0} \right)}}{2} + \frac{1}{2}\right) e^{i \phi_{0}}}{- m_{12}^{2} + 1.8225 - \frac{0.86113125 i \sqrt{\frac{\left(m_{12}^{2} - \left(m_{1} - m_{2}\right)^{2}\right) \left(m_{12}^{2} - \left(m_{1} + m_{2}\right)^{2}\right)}{m_{12}^{2}}}}{m_{12} \sqrt{\left(1.8225 - \left(m_{1} - m_{2}\right)^{2}\right) \left(1.8225 - \left(m_{1} + m_{2}\right)^{2}\right)}}}$

We can also compute separate sums of these components with the method sum_components() from the original HelicityModel:

added_components = model.sum_components(
components=[
R"A_{J/\psi(1S)_{+1} \to {f_{0}(500)}_{0} \gamma_{+1};"
R" {f_{0}(500)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}",
R"A_{J/\psi(1S)_{+1} \to {f_{0}(980)}_{0} \gamma_{+1};"
R" {f_{0}(980)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}",
R"A_{J/\psi(1S)_{+1} \to {f_{0}(1370)}_{0} \gamma_{+1};"
R" {f_{0}(1370)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}",
R"A_{J/\psi(1S)_{+1} \to {f_{0}(1500)}_{0} \gamma_{+1};"
R" {f_{0}(1500)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}",
R"A_{J/\psi(1S)_{+1} \to {f_{0}(1710)}_{0} \gamma_{+1};"
R" {f_{0}(1710)}_{0} \to \pi^{0}_{0} \pi^{0}_{0}}",
]
)


Just like in 2.2 Generate intensity-based sample, these intensity components can each be expressed in a computational backend. We do not have to parametrize this function now that we have already optimized the parameters, so we can just substitute the Symbols in all expression beforehand and use create_function() instead:

parameter_substitutions = model.parameter_defaults
for symbol in unfolded_expression.free_symbols:
optimized_value = fit_result.parameter_values.get(symbol.name)
if optimized_value is not None:
parameter_substitutions[symbol] = optimized_value

from tensorwaves.function.sympy import create_function

function_from_amplitude_sum = create_function(
backend="numpy",
)
function_from_intensity = create_function(
model.components[R"I_{J/\psi(1S)_{+1} \to \gamma_{+1} \pi^{0}_{0} \pi^{0}_{0}}"]
.doit()
.subs(parameter_substitutions)
.subs(mass_substitutions),
backend="numpy",
)


The result is a PositionalArgumentFunction that can be plotted just like in Plot optimized model:

Hide code cell source
fig, ax = plt.subplots(1, figsize=(8, 5))
bins = 150
phsp_projection = np.real(phsp["m_12"])
ax.hist(
phsp_projection,
weights=np.array(optimized_function(phsp)),
bins=bins,
alpha=0.2,
label="full intensity",
)
ax.hist(
phsp_projection,
weights=np.array(function_from_intensity(phsp)),
bins=bins,
histtype="step",
label=R"$J/\psi(1S)_{-1} \to \gamma_{-1} \pi^0 \pi^0$",
)
ax.set_xlim(0.25, 2.5)
ax.set_xlabel(R"$m_{\pi^0\pi^0}$ [GeV]")
ax.set_yticks([])
plt.legend()
plt.show()


Or generically, so that we can stack all the sub-intensities:

Hide code cell source
import logging
import warnings

logging.getLogger("tensorwaves.data").setLevel(logging.ERROR)  # hide progress bars
warnings.filterwarnings("ignore")  # hide negative sqrt warning

sub_intensity_functions = {
name: create_function(
sub_expression.doit().subs(parameter_substitutions).subs(mass_substitutions),
backend="jax",
)
for name, sub_expression in model.components.items()
if name.startswith("I")
}
initial_state_mass = reaction.initial_state[-1].mass
final_state_masses = {i: p.mass for i, p in reaction.final_state.items()}

masses = []
for sub_function in sub_intensity_functions.values():
sub_data_generator = IntensityDistributionGenerator(
phsp_generator,
function=sub_function,
domain_transformer=helicity_transformer,
)
sub_events = sub_data_generator.generate(5_000, rng)
sub_dataset = helicity_transformer(sub_events)
masses.append(np.real(sub_dataset["m_12"]))

fig, ax = plt.subplots(1, figsize=(8, 5))
plt.hist(masses, bins=100, stacked=True, alpha=0.6)
ax.set_xlim(0.25, 2.5)
ax.set_xlabel(R"$m_{\pi^0\pi^0}$ [GeV]")
ax.set_yticks([])
ax.legend(labels=[f"${name[3:-1]}$" for name in sub_intensity_functions])
plt.show()