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:
Create an amplitude model with
qrules
andampform
.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.
Simplified model: \(J/\psi \to \gamma f_0\), \(f_0 \rightarrow \pi^0\pi^0\)
As Step 3: Perform fit serves to illustrate usage only, we make the amplitude model here a bit simpler by not allowing \(\omega\) resonances (which are narrow and therefore hard to fit). For this reason, we can also limit the InteractionType
to STRONG
.
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
Show code cell source
from ampform.io import aslatex
from IPython.display import Math
Math(aslatex(model_no_dynamics.amplitudes))
Show code cell source
Math(aslatex(model_no_dynamics.parameter_defaults))
The heart of the model is a sympy expression
that contains the full description of the intensity model. Note two things:
The coefficients for the different amplitudes are complex valued.
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 assign()
method of the dynamics
attribute 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.dynamics.assign("J/psi(1S)", create_non_dynamic_with_ff)
for name in reaction.get_intermediate_particles().names:
model_builder.dynamics.assign(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:
Show 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)
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:
imported_model: HelicityModel = pickle.load(model_file)
Show 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)
Show 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.
See also
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:
Show code cell source
Math(aslatex(model.kinematic_variables))
from tensorwaves.data import SympyDataTransformer
helicity_transformer = SympyDataTransformer.from_sympy(
model.kinematic_variables, backend="jax"
)
See also
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
:
Show 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
Available kinematic variables
By default, tensorwaves
only generates invariant masses of the Topologies
that are of relevance to the decay problem. In this case, we only have resonances \(f_0 \to \pi^0\pi^0\). If you are interested in more invariant mass combinations, you can do so with the method register_topology()
.
The DataSample
can easily be converted to a pandas.DataFrame
:
import pandas as pd
data_frame = pd.DataFrame(data)
phsp_frame = pd.DataFrame(phsp)
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:
Show code cell source
%config InlineBackend.figure_formats = ['svg']
import matplotlib.pyplot as plt
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_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()
See also
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’.
Show code cell source
reaction = qrules.io.load("transitions.json")
with open("helicity_model.pickle", "rb") as stream:
model: HelicityModel = pickle.load(stream)
with open("data.pickle", "rb") as stream:
data = pickle.load(stream)
with open("phsp.pickle", "rb") as stream:
phsp = pickle.load(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]
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.4 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
5.81 ms ± 117 µ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#
Show code cell source
dot = sp.dotprint(unfolded_expression.args[0].args[0].args[0].args[0])
graphviz.Source(dot)
Optimized expression tree#
Show 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 Symbol
s 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.6 ms ± 167 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
5.23 ms ± 61 µ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!
Show code cell content
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]
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"),
use_analytic_gradient=False,
)
fit_result = minuit2.optimize(estimator, initial_parameters)
fit_result
FitResult(
minimum_valid=True,
execution_time=38.70925211906433,
function_calls=581,
estimator_value=-2818.4110168058587,
parameter_values={
'm_{f_{0}(500)}': 0.6190451064818804,
'm_{f_{0}(980)}': 0.9903165661915163,
'm_{f_{0}(1370)}': 1.3469935411846858,
'm_{f_{0}(1500)}': 1.518479660244283,
'm_{f_{0}(1710)}': 1.74451948353133,
'\\Gamma_{f_{0}(500)}': 0.4172234849380949,
'\\Gamma_{f_{0}(980)}': 0.06393661144847929,
'\\Gamma_{f_{0}(1710)}': 0.15228199760447464,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.0410382201089454-0.04219995352600525j),
},
parameter_errors={
'm_{f_{0}(500)}': 0.007700201509428148,
'm_{f_{0}(980)}': 0.0019686286262519392,
'm_{f_{0}(1370)}': 0.005929322670516124,
'm_{f_{0}(1500)}': 0.00397404644783676,
'm_{f_{0}(1710)}': 0.0030923435725945386,
'\\Gamma_{f_{0}(500)}': 0.030597276131929575,
'\\Gamma_{f_{0}(980)}': 0.004379012901789155,
'\\Gamma_{f_{0}(1710)}': 0.012186008357817284,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (0.04466380763254209+0.07610105335200282j),
},
)
See also
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.
Show 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': '2024-03-07 20:53:26.169919',
'optimizer': 'Minuit2',
'estimator_type': 'UnbinnedNLL',
'estimator_value': -2818.4110168058587,
'm_{f_{0}(500)}': 0.6190451064818804,
'm_{f_{0}(980)}': 0.9903165661915163,
'm_{f_{0}(1370)}': 1.3469935411846858,
'm_{f_{0}(1500)}': 1.518479660244283,
'm_{f_{0}(1710)}': 1.74451948353133,
'\\Gamma_{f_{0}(500)}': 0.4172234849380949,
'\\Gamma_{f_{0}(980)}': 0.06393661144847929,
'\\Gamma_{f_{0}(1710)}': 0.15228199760447464,
'C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}': (1.0410382201089454-0.04219995352600525j)}
See also
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 optimizer
s 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)} | |
---|---|---|---|---|---|---|---|---|---|---|
real_C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} | 1 | -0.2 | 0.1 | 0 | -0.6 | -0.2 | 0.1 | 0.1 | 0 | 0.1 |
imag_C_{J/\psi(1S) \to {f_{0}(1500)}_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} | -0.2 | 1 | -0.1 | -0.1 | 0.1 | 0.8 | -0.3 | 0.1 | -0.1 | 0.6 |
m_{f_{0}(500)} | 0.1 | -0.1 | 1 | 0.1 | 0 | -0.1 | 0.1 | -0.5 | 0 | 0 |
m_{f_{0}(980)} | 0 | -0.1 | 0.1 | 1 | 0 | 0 | 0 | -0.1 | -0.1 | 0 |
m_{f_{0}(1370)} | -0.6 | 0.1 | 0 | 0 | 1 | 0.3 | 0.1 | 0 | 0.2 | -0.2 |
m_{f_{0}(1500)} | -0.2 | 0.8 | -0.1 | 0 | 0.3 | 1 | -0.1 | 0 | 0 | 0.3 |
m_{f_{0}(1710)} | 0.1 | -0.3 | 0.1 | 0 | 0.1 | -0.1 | 1 | -0.1 | 0 | -0.2 |
\Gamma_{f_{0}(500)} | 0.1 | 0.1 | -0.5 | -0.1 | 0 | 0 | -0.1 | 1 | -0.4 | 0.1 |
\Gamma_{f_{0}(980)} | 0 | -0.1 | 0 | -0.1 | 0.2 | 0 | 0 | -0.4 | 1 | 0.1 |
\Gamma_{f_{0}(1710)} | 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.822033611717
BIC: -5544.718629891956
Fit fractions#
As we have seen when formulating the amplitude model, the helicity model is built up of an incoherent sum of real-valued intensities (one for each spin combination), which are each built up of a coherent sum of complex-valued amplitudes (one for each resonance). If we want to compute what each of these amplitudes contribute to the main intensity, we should use the optimized intensity and set coefficients or helicity couplings of amplitudes that were are not interested in to zero.
Here’s an example function that can do this. Using regular expressions, we set all coefficients in the intensity function to zero if they do not contain a certain resonance \(\LaTeX\) name:
import re
from tensorwaves.interface import ParametrizedFunction
def compute_sub_intensity(
func: ParametrizedFunction,
input_data: DataSample,
resonances: list[str],
):
original_parameters = dict(func.parameters)
negative_lookahead = f"(?!{'|'.join(map(re.escape, resonances))})"
# https://regex101.com/r/WrgGyD/1
pattern = rf"^(\\mathcal{{H}}|C_)({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:
new_parameters = dict(func.parameters)
for par_name in func.parameters:
if re.match(name_pattern, par_name) is not None:
new_parameters[par_name] = 0
func.update_parameters(new_parameters)
These functions can be used to compute and visualize the sub-intensity distributions over the phase space.
total_intensities = intensity_func(phsp)
sub_intensities = {
p: compute_sub_intensity(intensity_func, phsp, resonances=[p.latex])
for p in resonances
}
Show code cell source
fig, ax = plt.subplots(figsize=(8, 5))
ax.set_xlim(0.25, 2.5)
ax.set_xlabel(R"$m_{\pi^0\pi^0}$ [GeV]")
ax.set_yticks([])
bins = 150
phsp_projection = phsp["m_12"].real
ax.hist(
phsp_projection,
weights=total_intensities,
bins=bins,
color="red",
histtype="step",
label="full intensity",
)
ax.hist(
len(sub_intensities) * [phsp_projection],
weights=list(sub_intensities.values()),
bins=bins,
alpha=0.6,
label=[Rf"${p.latex}$" for p in sub_intensities],
stacked=True,
)
fig.legend()
plt.tight_layout()
plt.show()
We can also use these functions to compute the decay rates for each resonance. Notice how the sum of the decay rates does not add up to a 100%. This is because of the strong constructive interference between the resonances in this model.
total_intensity = intensity_func(phsp).sum()
fit_fractions = {
resonance.name: f"{sub_intensity.sum() / total_intensity:.1%}"
for resonance, sub_intensity in sub_intensities.items()
}
fit_fractions
{'f(0)(500)': '10.1%',
'f(0)(980)': '2.9%',
'f(0)(1370)': '17.2%',
'f(0)(1500)': '6.0%',
'f(0)(1710)': '7.9%',
'f(0)(2020)': '18.2%'}