Usage#

TensorWaves is a package for fitting general mathematical expressions to data distributions. The fundamentals behind the package are illustrated in Core ideas illustrated.

While general in design, the package is intended for doing Partial Wave Analysis. First, the ampform package determines which transitions are allowed from some initial state to a final state. It then formulates those transitions mathematically as an amplitude model. TensorWaves can then lambdify() this expression to some computational backend. Finally, TensorWaves ‘fits’ this model to some data sample. Optionally, a data sample can be generated from the model.

This page shows a brief overview of the complete workflow. More info about each step can be found under Step-by-step workflow.

Overview#

import logging
import warnings

import ampform
import graphviz
import matplotlib.pyplot as plt
import pandas as pd
import qrules as q
from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff

from tensorwaves.data import generate_data, generate_phsp
from tensorwaves.data.transform import HelicityTransformer
from tensorwaves.estimator import UnbinnedNLL
from tensorwaves.model import LambdifiedFunction, SympyModel
from tensorwaves.optimizer.callbacks import CSVSummary
from tensorwaves.optimizer.minuit import Minuit2

logger = logging.getLogger()
logger.setLevel(logging.ERROR)
warnings.filterwarnings("ignore")

Construct a model#

See also

Step 1: Create amplitude model

result = q.generate_transitions(
    initial_state=("J/psi(1S)", [-1, +1]),
    final_state=["gamma", "pi0", "pi0"],
    allowed_intermediate_particles=["f(0)(980)", "f(0)(1500)", "f(0)(1710)"],
    allowed_interaction_types=["strong", "EM"],
    formalism_type="canonical-helicity",
)
dot = q.io.asdot(result, collapse_graphs=True)
graphviz.Source(dot)

model_builder = ampform.get_builder(result)
for name in result.get_intermediate_particles().names:
    model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)
model = model_builder.generate()

model.components[
    R"A[J/\psi(1S)_{-1} \to f_{0}(980)_{0} \gamma_{+1,L=2,S=1}; f_{0}(980)_{0}"
    R" \to \pi^{0}_{0} \pi^{0}_{0,L=0,S=0}]"
].doit()

Generate data sample#

See also

Step 2: Generate data samples

sympy_model = SympyModel(
    expression=model.expression,
    parameters=model.parameter_defaults,
)
intensity = LambdifiedFunction(sympy_model, backend="jax")
data_converter = HelicityTransformer(model.adapter)
phsp_sample = generate_phsp(100_000, model.adapter.reaction_info)
data_sample = generate_data(
    10_000, model.adapter.reaction_info, data_converter, intensity
)

import numpy as np
from matplotlib import cm

reaction_info = model.adapter.reaction_info
intermediate_states = sorted(
    (
        p
        for p in model.particles
        if p not in reaction_info.final_state.values()
        and p not in reaction_info.initial_state.values()
    ),
    key=lambda p: p.mass,
)

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


def indicate_masses():
    plt.xlabel("$m$ [GeV]")
    for i, p in enumerate(intermediate_states):
        plt.gca().axvline(
            x=p.mass, linestyle="dotted", label=p.name, color=colors[i]
        )

phsp_set = data_converter.transform(phsp_sample)
data_set = data_converter.transform(data_sample)
data_frame = pd.DataFrame(data_set)
data_frame["m_12"].hist(bins=100, alpha=0.5, density=True)
indicate_masses()
plt.legend();

Optimize the model#

See also

Step 3: Perform fit

import matplotlib.pyplot as plt
import numpy as np


def compare_model(
    variable_name,
    data_set,
    phsp_set,
    intensity_model,
    bins=150,
):
    data = data_set[variable_name]
    phsp = phsp_set[variable_name]
    intensities = intensity_model(phsp_set)
    plt.hist(data, bins=bins, alpha=0.5, label="data", density=True)
    plt.hist(
        phsp,
        weights=intensities,
        bins=bins,
        histtype="step",
        color="red",
        label="initial fit model",
        density=True,
    )
    indicate_masses()
    plt.legend()

estimator = UnbinnedNLL(
    intensity,
    data_set,
    phsp_set,
    backend="jax",
)
initial_parameters = {
    "m_f(0)(980)": 0.93,
    "m_f(0)(1500)": 1.45,
    "m_f(0)(1710)": 1.8,
    "Gamma_f(0)(980)": 0.1,
    "Gamma_f(0)(1710)": 0.2,
}
intensity.update_parameters(initial_parameters)

compare_model("m_12", data_set, phsp_set, intensity)
print("Number of free parameters:", len(initial_parameters))

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

optimized_parameters = fit_result.parameter_values
intensity.update_parameters(optimized_parameters)
compare_model("m_12", data_set, phsp_set, intensity)

import pandas as pd
import sympy as sp

converters = {p: lambda s: complex(s).real for p in initial_parameters}
fit_traceback = pd.read_csv("fit_traceback.csv", converters=converters)
fig, (ax1, ax2) = plt.subplots(
    2, figsize=(7, 8), gridspec_kw={"height_ratios": [1, 1.8]}
)
fit_traceback.plot("function_call", "estimator_value", ax=ax1)
fit_traceback.plot("function_call", sorted(initial_parameters), ax=ax2)
ax1.set_title("Negative log likelihood")
ax2.set_title("Parameter values")
ax1.set_xlabel("function call")
ax2.set_xlabel("function call")
fig.tight_layout()
ax1.legend().remove()
legend_texts = ax2.legend().get_texts()
for text in legend_texts:
    latex = f"${sp.latex(sp.Symbol(text.get_text()))}$"
    latex = latex.replace("\\\\", "\\")
    if latex[2] == "C":
        latex = Rf"\left|{latex}\right|"
    text.set_text(latex)
for line in ax2.get_lines():
    label = line.get_label()
    color = line.get_color()
    ax2.axhline(
        y=complex(sympy_model.parameters[label]).real,
        color=color,
        alpha=0.5,
        linestyle="dotted",
    )

Step-by-step workflow#

The following pages go through the usual workflow when using tensorwaves. The output in each of these steps is written to disk, so that each notebook can be run separately.