docs/comparison/jpsi2phipipi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# J/ψ → φ(1020) π⁺ π⁻\n",
    "\n",
    "```{autolink-concat}\n",
    "```\n",
    "\n",
    "The decay $J/\\psi \\to \\phi(1020) \\pi^+ \\pi^-$ has an initial state with spin and has one vector meson, $\\phi(1020)$, in the final state. This is a follow-up to {doc}`jpsi2pipipi`, where there is no spin in the final state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Import Python libraries"
    },
    "tags": [
     "hide-cell",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import logging\n",
    "import os\n",
    "import warnings\n",
    "from textwrap import dedent\n",
    "from typing import TYPE_CHECKING\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import ipywidgets\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import qrules\n",
    "import sympy as sp\n",
    "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n",
    "from ampform.sympy import perform_cached_doit\n",
    "from IPython.display import Latex, Markdown, clear_output, display\n",
    "from ipywidgets import (\n",
    "    Accordion,\n",
    "    Checkbox,\n",
    "    GridBox,\n",
    "    HBox,\n",
    "    Layout,\n",
    "    SelectMultiple,\n",
    "    Tab,\n",
    "    ToggleButtons,\n",
    "    VBox,\n",
    "    interactive_output,\n",
    ")\n",
    "from tensorwaves.data.phasespace import TFPhaseSpaceGenerator\n",
    "from tensorwaves.data.rng import TFUniformRealNumberGenerator\n",
    "from tensorwaves.data.transform import SympyDataTransformer\n",
    "\n",
    "from ampform_dpd import DalitzPlotDecompositionBuilder\n",
    "from ampform_dpd.adapter.qrules import normalize_state_ids, to_three_body_decay\n",
    "from ampform_dpd.decay import Particle\n",
    "from ampform_dpd.io import (\n",
    "    as_markdown_table,\n",
    "    aslatex,\n",
    "    get_readable_hash,\n",
    "    perform_cached_lambdify,\n",
    "    simplify_latex_rendering,\n",
    ")\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from ampform.helicity import HelicityModel\n",
    "    from qrules.transition import ReactionInfo\n",
    "    from tensorwaves.interface import DataSample, ParameterValue, ParametrizedFunction\n",
    "\n",
    "simplify_latex_rendering()\n",
    "logging.getLogger(\"jax\").setLevel(logging.ERROR)  # mute JAX\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"  # mute TF\n",
    "warnings.simplefilter(\"ignore\", category=RuntimeWarning)\n",
    "NO_TQDM = \"EXECUTE_NB\" in os.environ\n",
    "if NO_TQDM:\n",
    "    logging.getLogger(\"ampform.sympy\").setLevel(logging.ERROR)\n",
    "    logging.getLogger(\"ampform_dpd.io\").setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Decay definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "mystnb": {
     "code_prompt_show": "Generate transitions"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "REACTION = qrules.generate_transitions(\n",
    "    initial_state=\"J/psi(1S)\",\n",
    "    final_state=[\"phi(1020)\", \"pi-\", \"pi+\"],\n",
    "    allowed_intermediate_particles=[\"a(0)(1450\", \"rho(1450)\"],\n",
    "    mass_conservation_factor=0,\n",
    "    formalism=\"helicity\",\n",
    ")\n",
    "REACTION123 = normalize_state_ids(REACTION)\n",
    "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "DECAY = to_three_body_decay(REACTION123.transitions, min_ls=True)\n",
    "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "resonances = sorted(\n",
    "    {t.resonance for t in DECAY.chains},\n",
    "    key=lambda p: (p.name[0], p.mass),\n",
    ")\n",
    "resonance_names = [p.name for p in resonances]\n",
    "Markdown(as_markdown_table(resonances))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(DECAY, with_jp=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Model formulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### DPD model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=True)\n",
    "DPD_MODEL = model_builder.formulate(reference_subsystem=1)\n",
    "del model_builder\n",
    "DPD_MODEL.intensity.cleanup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(DPD_MODEL.amplitudes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "There is an isobar Wigner-$d$ function, which takes the following helicity angles as argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(DPD_MODEL.variables))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### AmpForm model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "model_builder = ampform.get_builder(REACTION)\n",
    "model_builder.use_helicity_couplings = False\n",
    "model_builder.config.scalar_initial_state_mass = True\n",
    "model_builder.config.stable_final_state_ids = [0, 1, 2]\n",
    "AMPFORM_MODEL = model_builder.formulate()\n",
    "AMPFORM_MODEL.intensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(AMPFORM_MODEL.amplitudes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(AMPFORM_MODEL.kinematic_variables))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Phase space sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Formulate kinematic variables in terms of four-momenta"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "p1, p2, p3 = tuple(FourMomentumSymbol(f\"p{i}\", shape=[]) for i in (0, 1, 2))\n",
    "s1, s2, s3 = sorted(DPD_MODEL.invariants, key=str)\n",
    "mass_definitions = {\n",
    "    **DPD_MODEL.masses,\n",
    "    s1: InvariantMass(p2 + p3) ** 2,\n",
    "    s2: InvariantMass(p1 + p3) ** 2,\n",
    "    s3: InvariantMass(p1 + p2) ** 2,\n",
    "    sp.Symbol(\"m_01\", nonnegative=True): InvariantMass(p1 + p2),\n",
    "    sp.Symbol(\"m_02\", nonnegative=True): InvariantMass(p1 + p3),\n",
    "    sp.Symbol(\"m_12\", nonnegative=True): InvariantMass(p2 + p3),\n",
    "}\n",
    "dpd_variables = {\n",
    "    symbol: expr.doit().xreplace(DPD_MODEL.variables).xreplace(mass_definitions)\n",
    "    for symbol, expr in DPD_MODEL.variables.items()\n",
    "}\n",
    "dpd_transformer = SympyDataTransformer.from_sympy(dpd_variables, backend=\"jax\")\n",
    "\n",
    "ampform_transformer = SympyDataTransformer.from_sympy(\n",
    "    AMPFORM_MODEL.kinematic_variables, backend=\"jax\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_phase_space(reaction: ReactionInfo, size: int) -> dict[str, jnp.ndarray]:\n",
    "    rng = TFUniformRealNumberGenerator(seed=0)\n",
    "    phsp_generator = TFPhaseSpaceGenerator(\n",
    "        initial_state_mass=reaction.initial_state[-1].mass,\n",
    "        final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n",
    "    )\n",
    "    return phsp_generator.generate(size, rng)\n",
    "\n",
    "\n",
    "phsp = generate_phase_space(AMPFORM_MODEL.reaction_info, size=100_000)\n",
    "ampform_phsp = ampform_transformer(phsp)\n",
    "dpd_phsp = dpd_transformer(phsp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Convert to numerical functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def unfold_intensity(model: HelicityModel) -> sp.Expr:\n",
    "    unfolded_intensity = perform_cached_doit(model.intensity)\n",
    "    unfolded_amplitudes = {\n",
    "        symbol: perform_cached_doit(expr) for symbol, expr in model.amplitudes.items()\n",
    "    }\n",
    "    return unfolded_intensity.xreplace(unfolded_amplitudes)\n",
    "\n",
    "\n",
    "ampform_intensity_expr = unfold_intensity(AMPFORM_MODEL)\n",
    "dpd_intensity_expr = unfold_intensity(DPD_MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ampform_func = perform_cached_lambdify(\n",
    "    ampform_intensity_expr,\n",
    "    parameters=AMPFORM_MODEL.parameter_defaults,\n",
    ")\n",
    "dpd_func = perform_cached_lambdify(\n",
    "    dpd_intensity_expr,\n",
    "    parameters=DPD_MODEL.parameter_defaults,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Functions for computing sub-intensities"
    },
    "tags": [
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def compute_sub_intensities(\n",
    "    func: ParametrizedFunction, phsp: DataSample, resonance_name: str\n",
    ") -> jnp.ndarray:\n",
    "    original_parameters = dict(func.parameters)\n",
    "    _set_couplings_to_zero(func, resonance_name)\n",
    "    intensity_array = func(phsp)\n",
    "    func.update_parameters(original_parameters)\n",
    "    return intensity_array\n",
    "\n",
    "\n",
    "def _set_couplings_to_zero(\n",
    "    func: ParametrizedFunction, resonance_names: list[str]\n",
    ") -> None:\n",
    "    couplings_to_zero = {\n",
    "        key: value if any(r in key for r in resonance_names) else 0\n",
    "        for key, value in _get_couplings(func).items()\n",
    "    }\n",
    "    func.update_parameters(couplings_to_zero)\n",
    "\n",
    "\n",
    "def _get_couplings(func: ParametrizedFunction) -> dict[str, ParameterValue]:\n",
    "    return {\n",
    "        key: value\n",
    "        for key, value in func.parameters.items()\n",
    "        if key.startswith(\"C\") or \"production\" in key\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Define widget UI"
    },
    "tags": [
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def create_sliders() -> dict[str, ToggleButtons]:\n",
    "    all_parameters = dict(AMPFORM_MODEL.parameter_defaults.items())\n",
    "    all_parameters.update(DPD_MODEL.parameter_defaults)\n",
    "    sliders = {}\n",
    "    for symbol, value in all_parameters.items():\n",
    "        value = \"+1\"\n",
    "        if (\n",
    "            symbol.name.startswith(R\"\\mathcal{H}^\\mathrm{decay}\") and \"+\" in symbol.name\n",
    "        ) and any(s in symbol.name for s in [\"{1}\", \"*\", \"rho\"]):\n",
    "            value = \"-1\"\n",
    "        sliders[symbol.name] = ToggleButtons(\n",
    "            description=Rf\"\\({sp.latex(symbol)}\\)\",\n",
    "            options=[\"-1\", \"0\", \"+1\"],\n",
    "            value=value,\n",
    "            continuous_update=False,\n",
    "        )\n",
    "    return sliders\n",
    "\n",
    "\n",
    "def to_unicode(particle: Particle) -> str:\n",
    "    unicode = particle.name\n",
    "    unicode = unicode.replace(\"pi\", \"π\")\n",
    "    unicode = unicode.replace(\"rho\", \"p\")\n",
    "    unicode = unicode.replace(\"Sigma\", \"Σ\")\n",
    "    unicode = unicode.replace(\"~\", \"\")\n",
    "    unicode = unicode.replace(\"Σ\", \"~Σ\")\n",
    "    unicode = unicode.replace(\"+\", \"⁺\")\n",
    "    unicode = unicode.replace(\"-\", \"⁻\")\n",
    "    unicode = unicode.replace(\"(0)\", \"₀\")\n",
    "    unicode = unicode.replace(\"(1)\", \"₁\")\n",
    "    return unicode.replace(\")0\", \")⁰\")\n",
    "\n",
    "\n",
    "sliders = create_sliders()\n",
    "resonance_selector = SelectMultiple(\n",
    "    description=\"Resonance\",\n",
    "    options={to_unicode(p): p.latex for p in resonances},\n",
    "    value=[resonances[0].latex, resonances[1].latex],\n",
    "    layout=Layout(\n",
    "        height=f\"{14 * (len(resonances) + 1)}pt\",\n",
    "        width=\"auto\",\n",
    "    ),\n",
    ")\n",
    "hide_expressions = Checkbox(description=\"Hide expressions\", value=True)\n",
    "simplify_expressions = Checkbox(description=\"Simplify\", value=True)\n",
    "ipywidgets.link((hide_expressions, \"value\"), (simplify_expressions, \"disabled\"))\n",
    "\n",
    "package_names = (\"AmpForm\", \"AmpForm-DPD\")\n",
    "ui = HBox([\n",
    "    VBox([resonance_selector, hide_expressions, simplify_expressions]),\n",
    "    Tab(\n",
    "        children=[\n",
    "            Accordion(\n",
    "                children=[\n",
    "                    GridBox([\n",
    "                        sliders[key]\n",
    "                        for key in sorted(sliders)\n",
    "                        if p.latex in key\n",
    "                        if (\n",
    "                            key[0] in {\"C\", \"H\"}\n",
    "                            if package == \"AmpForm\"\n",
    "                            else key.startswith(R\"\\mathcal{H}\")\n",
    "                        )\n",
    "                    ])\n",
    "                    for package in package_names\n",
    "                ],\n",
    "                selected_index=1,\n",
    "                titles=package_names,\n",
    "            )\n",
    "            for p in resonances\n",
    "        ],\n",
    "        titles=[to_unicode(p) for p in resonances],\n",
    "    ),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Generate comparison widget"
    },
    "tags": [
     "full-width",
     "hide-input",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "plt.rc(\"font\", size=12)\n",
    "fig, axes = plt.subplots(figsize=(16, 6), ncols=3, nrows=2)\n",
    "fig.canvas.toolbar_visible = False\n",
    "fig.canvas.header_visible = False\n",
    "fig.canvas.footer_visible = False\n",
    "(\n",
    "    (ax_s1, ax_s2, ax_s3),\n",
    "    (ax_t1, ax_t2, ax_t3),\n",
    ") = axes\n",
    "for ax in axes[:, 0].flatten():\n",
    "    ax.set_ylabel(\"Intensity (a.u.)\")\n",
    "for ax in axes[:, 1:].flatten():\n",
    "    ax.set_yticks([])\n",
    "\n",
    "final_state = DECAY.final_state\n",
    "ax_s1.set_xlabel(f\"$m({final_state[2].latex}, {final_state[3].latex})$\")\n",
    "ax_s2.set_xlabel(f\"$m({final_state[1].latex}, {final_state[3].latex})$\")\n",
    "ax_s3.set_xlabel(f\"$m({final_state[1].latex}, {final_state[2].latex})$\")\n",
    "ax_t1.set_xlabel(Rf\"$\\theta({final_state[2].latex}, {final_state[3].latex})$\")\n",
    "ax_t2.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[3].latex})$\")\n",
    "ax_t3.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[2].latex})$\")\n",
    "fig.suptitle(f'Selected resonances: ${\", \".join(resonance_selector.value)}$')\n",
    "fig.tight_layout()\n",
    "\n",
    "lines = None\n",
    "\n",
    "__EXPRS: dict[int, sp.Expr] = {}\n",
    "\n",
    "\n",
    "def _get_symbol_values(\n",
    "    expr: sp.Expr,\n",
    "    parameters: dict[str, ParameterValue],\n",
    "    selected_resonances: list[str],\n",
    ") -> dict[sp.Symbol, sp.Rational]:\n",
    "    parameters = {\n",
    "        key: value if any(r in key for r in selected_resonances) else 0\n",
    "        for key, value in parameters.items()\n",
    "    }\n",
    "    return {\n",
    "        s: sp.Rational(parameters[s.name])\n",
    "        for s in expr.free_symbols\n",
    "        if s.name in parameters\n",
    "    }\n",
    "\n",
    "\n",
    "def _simplify(\n",
    "    intensity_expr: sp.Expr,\n",
    "    parameter_defaults: dict[str, ParameterValue],\n",
    "    variables: dict[sp.Symbol, sp.Expr],\n",
    "    selected_resonances: list[str],\n",
    ") -> sp.Expr:\n",
    "    parameters = _get_symbol_values(\n",
    "        intensity_expr, parameter_defaults, selected_resonances\n",
    "    )\n",
    "    fixed_variables = {k: v for k, v in variables.items() if not v.free_symbols}\n",
    "    obj = (\n",
    "        intensity_expr,\n",
    "        tuple((k, fixed_variables[k]) for k in sorted(fixed_variables, key=str)),\n",
    "        tuple((k, parameters[k]) for k in sorted(parameters, key=str)),\n",
    "    )\n",
    "    h = get_readable_hash(obj)\n",
    "    if h in __EXPRS:\n",
    "        return __EXPRS[h]\n",
    "    expr = intensity_expr.xreplace(parameters).xreplace(fixed_variables)\n",
    "    expr = sp.trigsimp(expr)\n",
    "    __EXPRS[h] = expr\n",
    "    return expr\n",
    "\n",
    "\n",
    "def plot_contributions(**kwargs) -> None:\n",
    "    kwargs.pop(\"resonance_selector\")\n",
    "    kwargs.pop(\"hide_expressions\")\n",
    "    kwargs.pop(\"simplify_expressions\")\n",
    "    selected_resonances = list(resonance_selector.value)\n",
    "    fig.suptitle(f'Selected resonances: ${\", \".join(selected_resonances)}$')\n",
    "    dpd_pars = {k: int(v) for k, v in kwargs.items() if k in dpd_func.parameters}\n",
    "    ampform_pars = {\n",
    "        k: int(v) for k, v in kwargs.items() if k in ampform_func.parameters\n",
    "    }\n",
    "    ampform_func.update_parameters(ampform_pars)\n",
    "    dpd_func.update_parameters(dpd_pars)\n",
    "    ampform_intensities = compute_sub_intensities(\n",
    "        ampform_func, ampform_phsp, selected_resonances\n",
    "    )\n",
    "    dpd_intensities = compute_sub_intensities(dpd_func, dpd_phsp, selected_resonances)\n",
    "\n",
    "    s1_edges = jnp.linspace(0.2, 2.1, num=50)\n",
    "    s23_edges = jnp.linspace(1.1, 3.0, num=50)\n",
    "    amp_values_s1, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_12\"].real,\n",
    "        bins=s1_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_s1, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_12\"].real,\n",
    "        bins=s1_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "    amp_values_s2, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_02\"].real,\n",
    "        bins=s23_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_s2, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_02\"].real,\n",
    "        bins=s23_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "    amp_values_s3, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_01\"].real,\n",
    "        bins=s23_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_s3, _ = jnp.histogram(\n",
    "        ampform_phsp[\"m_01\"].real,\n",
    "        bins=s23_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "\n",
    "    t_edges = jnp.linspace(0, jnp.pi, num=50)\n",
    "    amp_values_t1, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_23\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_t1, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_23\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "    amp_values_t2, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_31\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_t2, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_31\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "    amp_values_t3, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_23\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=ampform_intensities,\n",
    "    )\n",
    "    dpd_values_t3, _ = jnp.histogram(\n",
    "        dpd_phsp[\"theta_23\"].real,\n",
    "        bins=t_edges,\n",
    "        weights=dpd_intensities,\n",
    "    )\n",
    "\n",
    "    global lines\n",
    "    amp_kwargs = dict(color=\"r\", label=\"ampform\", linestyle=\"solid\")\n",
    "    dpd_kwargs = dict(color=\"blue\", label=\"dpd\", linestyle=\"dotted\")\n",
    "    if lines is None:\n",
    "        sx1 = (s1_edges[:-1] + s1_edges[1:]) / 2\n",
    "        sx23 = (s23_edges[:-1] + s23_edges[1:]) / 2\n",
    "        tx = (t_edges[:-1] + t_edges[1:]) / 2\n",
    "        lines = [\n",
    "            ax_s1.step(sx1, amp_values_s1, **amp_kwargs)[0],\n",
    "            ax_s1.step(sx1, dpd_values_s1, **dpd_kwargs)[0],\n",
    "            ax_s2.step(sx23, amp_values_s2, **amp_kwargs)[0],\n",
    "            ax_s2.step(sx23, dpd_values_s2, **dpd_kwargs)[0],\n",
    "            ax_s3.step(sx23, amp_values_s3, **amp_kwargs)[0],\n",
    "            ax_s3.step(sx23, dpd_values_s3, **dpd_kwargs)[0],\n",
    "            ax_t1.step(tx, amp_values_t1, **amp_kwargs)[0],\n",
    "            ax_t1.step(tx, dpd_values_t1, **dpd_kwargs)[0],\n",
    "            ax_t2.step(tx, amp_values_t2, **amp_kwargs)[0],\n",
    "            ax_t2.step(tx, dpd_values_t2, **dpd_kwargs)[0],\n",
    "            ax_t3.step(tx, amp_values_t3, **amp_kwargs)[0],\n",
    "            ax_t3.step(tx, dpd_values_t3, **dpd_kwargs)[0],\n",
    "        ]\n",
    "        ax_s1.legend(loc=\"upper right\")\n",
    "    else:\n",
    "        lines[0].set_ydata(amp_values_s1)\n",
    "        lines[1].set_ydata(dpd_values_s1)\n",
    "        lines[2].set_ydata(amp_values_s2)\n",
    "        lines[3].set_ydata(dpd_values_s2)\n",
    "        lines[4].set_ydata(amp_values_s3)\n",
    "        lines[5].set_ydata(dpd_values_s3)\n",
    "        lines[6].set_ydata(amp_values_t1)\n",
    "        lines[7].set_ydata(dpd_values_t1)\n",
    "        lines[8].set_ydata(amp_values_t2)\n",
    "        lines[9].set_ydata(dpd_values_t2)\n",
    "        lines[10].set_ydata(amp_values_t3)\n",
    "        lines[11].set_ydata(dpd_values_t3)\n",
    "\n",
    "    sy_max = max(\n",
    "        jnp.nanmax(amp_values_s1),\n",
    "        jnp.nanmax(dpd_values_s1),\n",
    "        jnp.nanmax(amp_values_s2),\n",
    "        jnp.nanmax(dpd_values_s2),\n",
    "        jnp.nanmax(amp_values_s3),\n",
    "        jnp.nanmax(dpd_values_s3),\n",
    "    )\n",
    "    ty_max = max(\n",
    "        jnp.nanmax(amp_values_t1),\n",
    "        jnp.nanmax(dpd_values_t1),\n",
    "        jnp.nanmax(amp_values_t2),\n",
    "        jnp.nanmax(dpd_values_t2),\n",
    "        jnp.nanmax(amp_values_t3),\n",
    "        jnp.nanmax(dpd_values_t3),\n",
    "    )\n",
    "    for ax in axes[0]:\n",
    "        ax.set_ylim(0, 1.05 * sy_max)\n",
    "    for ax in axes[1]:\n",
    "        ax.set_ylim(0, 1.05 * ty_max)\n",
    "\n",
    "    fig.canvas.draw_idle()\n",
    "\n",
    "    if hide_expressions.value:\n",
    "        clear_output()\n",
    "    else:\n",
    "        if simplify_expressions.value:\n",
    "            ampform_expr = _simplify(\n",
    "                ampform_intensity_expr,\n",
    "                ampform_pars,\n",
    "                AMPFORM_MODEL.kinematic_variables,\n",
    "                selected_resonances,\n",
    "            )\n",
    "            dpd_expr = _simplify(\n",
    "                dpd_intensity_expr,\n",
    "                dpd_pars,\n",
    "                DPD_MODEL.variables,\n",
    "                selected_resonances,\n",
    "            )\n",
    "        else:\n",
    "            ampform_expr = ampform_intensity_expr\n",
    "            dpd_expr = dpd_intensity_expr\n",
    "        src = Rf\"\"\"\n",
    "        \\begin{{eqnarray}}\n",
    "          \\text{{AmpForm:}} && {sp.latex(ampform_expr)} \\\\\n",
    "          \\text{{DPD:}}     && {sp.latex(dpd_expr)} \\\\\n",
    "        \\end{{eqnarray}}\n",
    "        \"\"\"\n",
    "        src = dedent(src).strip()\n",
    "        display(Latex(src))\n",
    "\n",
    "\n",
    "output = interactive_output(\n",
    "    plot_contributions,\n",
    "    controls={\n",
    "        **sliders,\n",
    "        \"resonance_selector\": resonance_selector,\n",
    "        \"hide_expressions\": hide_expressions,\n",
    "        \"simplify_expressions\": simplify_expressions,\n",
    "    },\n",
    ")\n",
    "display(output, ui)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}