Construct three-body decay

\(\Lambda_b^0 \to p K^- \gamma\)

Model definition: lb2pkg-lhcb-2765817.json.

This page demonstrates deserialization and evaluation of an amplitude model for the decay \(\Lambda_b^0 \to p K^- \gamma\). The resonant structure was studied using proton-proton collision data recorded at centre-of-mass energies of \(7\), \(8\), and \(13\) TeV with the LHCb detector, INSPIRE-HEP 2765817.

Import Python libraries 
from copy import deepcopy
from pathlib import Path
import json
import logging

import jax.numpy as jnp
import matplotlib.pyplot as plt
import pandas as pd
import sympy as sp
from ampform.dynamics.form_factor import BreakupMomentumSquared
from ampform_dpd import DefinedExpression
from ampform_dpd.decay import FinalStateID, State, ThreeBodyDecay
from ampform_dpd.io import aslatex, cached, perform_cached_lambdify
from ampform_dpd.io.serialization import amplitude as serialization_amplitude
from ampform_dpd.io.serialization.amplitude import formulate
from ampform_dpd.io.serialization.decay import get_final_state, to_decay
from ampform_dpd.io.serialization.dynamics import (
    formulate_dynamics,
    to_mandelstam_symbol,
    to_mass_symbol,
)
from ampform_dpd.io.serialization.format import (
    get_decay_chains,
    get_function_definition,
)
from IPython.display import Math
from matplotlib_inline.backend_inline import set_matplotlib_formats
from mpl_toolkits.axes_grid1 import make_axes_locatable  # cspell:ignore mpl_toolkits

THIS_DIR = Path(".").absolute()
logging.getLogger("ampform.sympy").setLevel(logging.ERROR)
set_matplotlib_formats("svg")
with open(THIS_DIR.parent.parent / "models" / "lb2pkg-lhcb-2765817.json") as f:
    MODEL_DEFINITION = json.load(f)
Name-to-LaTeX converter 
def to_latex(name: str) -> str:
    latex = {
        "Lb": R"\Lambda_b^0",
        "g": R"\gamma",
        "K": "K^-",
        "p": "p",
    }.get(name)
    if latex is not None:
        return latex
    if name.startswith("LNR"):
        return R"\Lambda_\mathrm{NR}"
    if name.startswith("L"):
        return Rf"\Lambda({name[1:]})"
    return name


DECAY = to_decay(MODEL_DEFINITION, to_latex=to_latex)
Math(aslatex(DECAY, with_jp=True))

\(\displaystyle \begin{array}{c} \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1405)\left[J=\frac{1}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1520)\left[J=\frac{3}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1600)\left[J=\frac{1}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1670)\left[J=\frac{1}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1690)\left[J=\frac{3}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1800)\left[J=\frac{1}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1810)\left[J=\frac{1}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1820)\left[J=\frac{5}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1830)\left[J=\frac{5}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(1890)\left[J=\frac{3}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(2100)\left[J=\frac{7}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(2110)\left[J=\frac{5}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda(2350)\left[J=\frac{9}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \Lambda_b^0\left[J=\frac{1}{2}\right] \to \left(\Lambda_\mathrm{NR}\left[J=\frac{3}{2}\right] \to p\left[J=\frac{1}{2}\right] K^-\left[J=0\right]\right) \gamma\left[J=1\right] \\ \end{array}\)

Dynamics

The model contains Breit-Wigner and multichannel Breit-Wigner propagators, as well as a constant non-resonant term. The latter is serialized as a zeroth-order polynomial.

CHAIN_DEFS = get_decay_chains(MODEL_DEFINITION)
chains_by_propagator = {
    propagator["parametrization"]: chain
    for chain in CHAIN_DEFS
    for propagator in chain["propagators"]
}
L1520_BW = formulate_dynamics(chains_by_propagator["L1520_BW"], MODEL_DEFINITION)
Math(aslatex(L1520_BW))

\(\displaystyle \begin{array}{rcl} \mathcal{R}^\mathrm{BW}_{L=2}\left(\sigma_{3}; m_{L1520}, \Gamma_{L1520}\right) &=& \mathcal{R}^\mathrm{BW}\left(\sigma_{3}; m_{L1520}, \Gamma_{1520}\left(\sigma_{3}\right)\right) \\ m_{L1520} &=& 1.519 \\ \Gamma_{L1520} &=& 0.016 \\ m_{1} &=& 0.938 \\ m_{2} &=& 0.493 \\ R_\mathrm{res} &=& 1.5 \\ \end{array}\)

L1405_Flatte = formulate_dynamics(
    chains_by_propagator["L1405_Flatte"],
    MODEL_DEFINITION,
)
Math(aslatex(L1405_Flatte))

\(\displaystyle \begin{array}{rcl} \mathcal{R}^\mathrm{BW}_\mathrm{multi}\left(\sigma_{3}; \Gamma_{L1405}, \Gamma_{L1405}^\text{ch. 2}\right) &=& \mathcal{R}^\mathrm{BW}_{L=0}\left(\sigma_{3}; m_{L1405}, \Gamma^\text{ch}\left(\sigma_{3}; m_{L1405}, \Gamma_{L1405}\right) + \Gamma^\text{ch}\left(\sigma_{3}; m_{L1405}, \Gamma_{L1405}^\text{ch. 2}\right)\right) \\ m_{L1405} &=& 1.405 \\ \Gamma_{L1405} &=& 0.24941478752959237 \\ m_{1} &=& 0.938 \\ m_{2} &=& 0.493 \\ R_{L1405} &=& 0 \\ m_{a,2} &=& 1.197 \\ m_{b,2} &=& 0.14 \\ \Gamma_{L1405}^\text{ch. 2} &=& 0.24941478752959237 \\ \end{array}\)

Construct amplitude model

The current Python preview API needs two small compatibility adapters for this older LS-coupled model. They preserve a unique symbol for every serialized LS coupling and implement the serialized MomentumPower and constant Polynomial functions.

Compatibility adapters for LS-coupled models 
def prepare_ls_model(model_definition: dict) -> dict:
    model_definition = deepcopy(model_definition)
    chains = get_decay_chains(model_definition)
    for chain_idx, chain in enumerate(chains):
        for vertex in chain["vertices"]:
            if "helicities" in vertex:
                continue
            has_isobar = any(isinstance(item, list) for item in vertex["node"])
            vertex["helicities"] = [str(chain_idx), "0"] if has_isobar else ["0", "0"]
    return model_definition


def get_ls_child_spins(model_definition, chain_idx, vertex_idx):
    chain = get_decay_chains(model_definition)[chain_idx]
    vertex = chain["vertices"][vertex_idx]
    final_state = get_final_state(model_definition)
    resonance_spin = sp.Rational(chain["propagators"][0]["spin"])
    return tuple(
        final_state[item].spin if isinstance(item, int) else resonance_spin
        for item in vertex["node"]
    )


def formulate_polynomial(propagator, resonance, model):
    definition = get_function_definition(propagator["parametrization"], model)
    variable = to_mandelstam_symbol(propagator["node"])
    expression = sum(
        coefficient * variable**power
        for power, coefficient in enumerate(definition["coefficients"])
    )
    return DefinedExpression(expression=expression)


_original_formulate_form_factor = serialization_amplitude.formulate_form_factor


def formulate_form_factor(vertex, model):
    function_name = vertex.get("formfactor")
    if not function_name:
        return DefinedExpression()
    definition = get_function_definition(function_name, model)
    if definition["type"] != "MomentumPower":
        return _original_formulate_form_factor(vertex, model)
    node = vertex["node"]
    if all(isinstance(item, int) for item in node):
        s = to_mandelstam_symbol(node)
        m1, m2 = (to_mass_symbol(item) for item in node)
    else:
        s = to_mandelstam_symbol(node) ** 2
        m1, m2 = (to_mass_symbol(item) for item in node)
        m1 = sp.sqrt(m1)
    power = sp.Rational(definition["l"], 2)
    return DefinedExpression(expression=BreakupMomentumSquared(s, m1, m2) ** power)


serialization_amplitude._get_child_spins = get_ls_child_spins
serialization_amplitude.formulate_form_factor = formulate_form_factor
LS_MODEL_DEFINITION = prepare_ls_model(MODEL_DEFINITION)
ADDITIONAL_BUILDERS = {"Polynomial": formulate_polynomial}
model = formulate(
    LS_MODEL_DEFINITION,
    cleanup_summations=True,
    to_latex=to_latex,
    additional_builders=ADDITIONAL_BUILDERS,
)

Validation

The table compares all serialized checkpoints with the Python implementation. The marks โ€œ๐ŸŸขโ€, โ€œ๐ŸŸกโ€, and โ€œ๐Ÿ”ดโ€ indicate an absolute difference of \(<10^{-10}\), \(<10^{-2}\), or \(\ge10^{-2}\), respectively.

Validation helpers 
checksum_points = {
    point["name"]: {parameter["name"]: parameter["value"] for parameter in point["parameters"]}
    for point in MODEL_DEFINITION["parameter_points"]
}


def to_number(value: float | str) -> float | complex:
    if isinstance(value, str):
        value = complex(value.replace(" ", "").replace("i", "j"))
    number = complex(value)
    return number.real if number.imag == 0 else number


def round_number(value: float | complex, digits: int = 6) -> float | complex:
    number = complex(value)
    if number.imag == 0:
        return round(number.real, digits)
    return complex(round(number.real, digits), round(number.imag, digits))


def label_diff(difference: complex) -> str:
    absolute_difference = abs(difference)
    if absolute_difference < 1e-10:
        return "๐ŸŸข"
    if absolute_difference < 1e-2:
        return "๐ŸŸก"
    return "๐Ÿ”ด"


def create_intensity_function(
    model,
    subsystem: int,
    resonance_name: str | None = None,
):
    invariant = next(s for s in model.invariants if str(s) == f"sigma{subsystem}")
    intensity_expr = cached.xreplace(cached.unfold(model), model.variables)
    parameter_defaults = select_resonance_parameters(model, resonance_name)
    intensity_expr = cached.xreplace(intensity_expr, parameter_defaults)
    invariant_expr = model.invariants[invariant].xreplace(model.masses).doit()
    intensity_expr = cached.doit(intensity_expr.xreplace({invariant: invariant_expr}))
    return cached.lambdify(intensity_expr, backend="jax")


def select_resonance_parameters(model, resonance_name: str | None):
    if resonance_name is None:
        return model.parameter_defaults
    return {
        symbol: value
        if not str(symbol).startswith("c^") or resonance_name in str(symbol)
        else 0
        for symbol, value in model.parameter_defaults.items()
    }


def evaluate_intensity(model, point: dict[str, float], pair: tuple[int, int]) -> float:
    i, j = pair
    k, *_ = {1, 2, 3} - set(pair)
    invariants = {str(s): s for s in model.invariants}
    sigma_j = invariants[f"sigma{j}"]
    sigma_k = invariants[f"sigma{k}"]
    angle_expr = next(
        expression
        for angle, expression in model.variables.items()
        if str(angle) == f"theta_{i}{j}"
    )
    sigma_k_value = point[f"m_{i}{j}"] ** 2
    equation = sp.Eq(
        sp.cos(angle_expr).xreplace(model.parameter_defaults),
        point[f"cos_theta_{i}{j}"],
    )
    sigma_j_value = sp.solve(equation.subs(sigma_k, sigma_k_value), sigma_j)[0]
    function = create_intensity_function(model, subsystem=i)
    value = function({str(sigma_j): float(sigma_j_value), str(sigma_k): sigma_k_value})
    return float(jnp.real(value))
Compute every serialized checkpoint 
validation_results = []
for checksum in MODEL_DEFINITION["misc"]["amplitude_model_checksums"]:
    distribution_name = checksum["distribution"]
    point = checksum_points[checksum["point"]]
    reference_value = to_number(checksum["value"])
    chain = chains_by_propagator.get(distribution_name)
    if chain is None:
        computed_value = evaluate_intensity(model, point, pair=(1, 2))
    else:
        dynamics = formulate_dynamics(
            chain,
            MODEL_DEFINITION,
            additional_definitions=ADDITIONAL_BUILDERS,
        )
        expression = dynamics.expression.doit()
        variables = expression.free_symbols - dynamics.parameters.keys()
        if variables:
            function = perform_cached_lambdify(
                expression,
                parameters=dynamics.parameters,
            )
            variable = next(iter(variables))
            computed_value = function({str(variable): next(iter(point.values()))})
        else:
            computed_value = complex(expression.xreplace(dynamics.parameters))
    difference = abs(reference_value - computed_value)
    validation_results.append({
        "Distribution": distribution_name,
        "Point": checksum["point"],
        "Reference": round_number(reference_value),
        "Computed": round_number(computed_value),
        "Difference": difference,
        "Status": label_diff(reference_value - computed_value),
    })

pd.DataFrame(validation_results, dtype=object)
Distribution Point Reference Computed Difference Status
0 default_model validation_point1 645.687778 644.103426 1.584352 ๐Ÿ”ด
1 default_model validation_point2 33.313132 33.314165 0.001032 ๐ŸŸก
2 default_model validation_point3 38.286729 38.287889 0.00116 ๐ŸŸก
3 L1405_Flatte validation_point_m12sq (-0.740198+0.236368j) (-0.775611+0.176283j) 0.06974548265443277 ๐Ÿ”ด
4 L1890_BW validation_point_m12sq (2.198693+1.036641j) (2.198693+1.036641j) 0.0 ๐ŸŸข
5 L1800_BW validation_point_m12sq (0.311734+2.7742j) (0.311734+2.7742j) 4.577566798522237e-16 ๐ŸŸข
6 L1520_BW validation_point_m12sq (-0.643171+0.553948j) (-0.643171+0.553948j) 2.482534153247273e-16 ๐ŸŸข
7 L2110_BW validation_point_m12sq (0.854068+0.041564j) (0.854068+0.041564j) 1.2738115944464653e-16 ๐ŸŸข
8 L1680_BW validation_point_m12sq (-2.46133+0.361098j) (-2.46133+0.361098j) 1.1102230246251565e-16 ๐ŸŸข
9 L2350_BW validation_point_m12sq (0.43057+0.000197j) (0.43057+0.000197j) 1.0028870095490916e-18 ๐ŸŸข
10 L1600_BW validation_point_m12sq (-0.629331+0.766337j) (-0.629331+0.766337j) 2.220446049250313e-16 ๐ŸŸข
11 L1810_BW validation_point_m12sq (0.106436+5.093986j) (0.106436+5.093986j) 8.899114524108741e-16 ๐ŸŸข
12 L1820_BW validation_point_m12sq (4.647408+4.443945j) (4.647408+4.443945j) 4.440892098500626e-15 ๐ŸŸข
13 L1830_BW validation_point_m12sq (3.824528+3.82775j) (3.824528+3.82775j) 1.7763568394002505e-15 ๐ŸŸข
14 L2100_BW validation_point_m12sq (0.826225+0.013532j) (0.826225+0.013532j) 1.5612511283791264e-17 ๐ŸŸข
15 L1690_BW validation_point_m12sq (-1.859729+1.396126j) (-1.859729+1.396126j) 1.4043333874306805e-15 ๐ŸŸข
16 LNR30_NR validation_point_m12sq 1.0 1.0 0.0 ๐ŸŸข

Visualization

Dalitz plot

The Dalitz plot shows the intensity up to \(m_{pK^-}=2.5\ \mathrm{GeV}\). The plotting range includes a margin below the physical threshold so that the left phase-space boundary remains visible, as in the Julia implementation.

Configure the Dalitz plot 
x_subsystem, y_subsystem = 3, 2
eliminated_subsystem, *_ = {1, 2, 3} - {x_subsystem, y_subsystem}
resolution = 250
x_min, x_max = 1.9, 2.5**2
y_min, y_max = 5, 27
X, Y = jnp.meshgrid(
    jnp.linspace(x_min, x_max, resolution),
    jnp.linspace(y_min, y_max, resolution),
)
intensity_function = create_intensity_function(model, subsystem=eliminated_subsystem)
intensities = jnp.real(intensity_function({
    f"sigma{x_subsystem}": X,
    f"sigma{y_subsystem}": Y,
}))
normalized_intensities = intensities / jnp.nansum(intensities)


def get_decay_products(
    decay: ThreeBodyDecay,
    subsystem: FinalStateID,
) -> tuple[State, State]:
    return tuple(state for state in decay.final_state.values() if state.index != subsystem)


sigma_labels = {
    subsystem: Rf"$\sigma_{subsystem} = M^2\left({' '.join(p.latex for p in get_decay_products(DECAY, subsystem))}\right)$"
    for subsystem in (1, 2, 3)
}
# cspell:ignore bbox edgecolor facecolor fontset fontsize labelcolor mathtext ncol
# cspell:ignore savefig sharey startswith STIX wspace xtick ytick
plt.rcParams.update({
    "axes.edgecolor": "#808080",
    "axes.facecolor": "none",
    "axes.labelcolor": "#808080",
    "figure.facecolor": "none",
    "font.family": "serif",
    "font.serif": ["STIXGeneral", "DejaVu Serif"],
    "font.size": 18,
    "mathtext.fontset": "stix",
    "savefig.transparent": True,
    "text.color": "#808080",
    "xtick.color": "#808080",
    "ytick.color": "#808080",
})
Render the Dalitz plot 
fig, ax = plt.subplots(figsize=(10, 8))
mesh = ax.pcolormesh(X, Y, normalized_intensities, rasterized=True)
ax.contour(
    X,
    Y,
    jnp.isfinite(intensities),
    levels=[0.5],
    colors="black",
    linewidths=3,
)
ax.set_xlabel(sigma_labels[x_subsystem])
ax.set_ylabel(sigma_labels[y_subsystem])
divider = make_axes_locatable(ax)
colorbar_axes = divider.append_axes("right", size="5%", pad=0.1)
colorbar = fig.colorbar(mesh, cax=colorbar_axes)
colorbar.ax.set_ylabel("Normalized intensity (a.u.)")
plt.show()

Mass projection

The projection onto the \(pK^-\) mass is shown as a black line. The colored curves show the incoherent contribution from each \(\Lambda\) resonance.

Render the proton-kaon mass projection 
resonance_names = sorted({chain["name"] for chain in CHAIN_DEFS})
mass_values = jnp.sqrt(X[0])
minimum_mass = float(
    (
        sp.Symbol("m1", nonnegative=True) + sp.Symbol("m2", nonnegative=True)
    ).xreplace(model.masses)
)
physical_region = mass_values >= minimum_mass
mass_values = mass_values[physical_region]


def project_to_pk(intensity_values):
    projection = jnp.nansum(intensity_values, axis=0)
    return mass_values * projection[physical_region]


total_projection = project_to_pk(intensities)
normalization = jnp.trapezoid(total_projection, mass_values)

fig, ax = plt.subplots(figsize=(12, 7))
ax.plot(
    mass_values,
    total_projection / normalization,
    color="#202020",
    lw=3,
    label="Total",
)
for resonance_name in resonance_names:
    component_function = create_intensity_function(
        model,
        subsystem=eliminated_subsystem,
        resonance_name=to_latex(resonance_name),
    )
    component = jnp.real(component_function({"sigma2": Y, "sigma3": X}))
    component = jnp.where(jnp.isnan(intensities), jnp.nan, component)
    label = Rf"${to_latex(resonance_name)}$"
    ax.plot(
        mass_values,
        project_to_pk(component) / normalization,
        label=label,
    )
ax.set_xlabel(R"$m_{pK^-}$ [GeV]")
ax.set_ylabel("Normalized intensity (a.u.)")
ax.set_xlim(minimum_mass, 2.5)
ax.set_ylim(bottom=0)
ax.legend(fontsize="small", loc="upper right", ncol=3)
plt.show()