From a96767b70991028530dce0abfba53cb843376580 Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Wed, 14 Aug 2024 10:44:12 +0200
Subject: [PATCH 1/8] starting draft doc

---
 docs/lambda-k-pi/ampform-dpd.ipynb            | 759 ++++++++++++++++++
 .../{automated.ipynb => ampform.ipynb}        |   0
 2 files changed, 759 insertions(+)
 create mode 100644 docs/lambda-k-pi/ampform-dpd.ipynb
 rename docs/lambda-k-pi/{automated.ipynb => ampform.ipynb} (100%)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
new file mode 100644
index 0000000..adbd634
--- /dev/null
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -0,0 +1,759 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Amplitude model with `ampform-dpd` (Draft)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "PWA study on $p \\gamma \\to \\Lambda K^+ \\pi^0$.\n",
+    "We formulate the helicity amplitude model symbolically using AmpForm-DPD here."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Import Python libraries"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "import logging\n",
+    "import os\n",
+    "import warnings\n",
+    "from collections import defaultdict\n",
+    "from fractions import Fraction\n",
+    "from textwrap import dedent\n",
+    "\n",
+    "import ampform\n",
+    "import graphviz\n",
+    "import ipywidgets as w\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import qrules\n",
+    "import sympy as sp\n",
+    "from ampform.dynamics.builder import RelativisticBreitWignerBuilder\n",
+    "from ampform.io import aslatex, improve_latex_rendering\n",
+    "from IPython.display import Markdown, Math, display\n",
+    "from qrules.particle import Particle, Spin, create_particle, load_pdg\n",
+    "from tensorwaves.data import (\n",
+    "    SympyDataTransformer,\n",
+    "    TFPhaseSpaceGenerator,\n",
+    "    TFUniformRealNumberGenerator,\n",
+    ")\n",
+    "from tensorwaves.function.sympy import create_parametrized_function\n",
+    "\n",
+    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
+    "logging.disable(logging.WARNING)\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "improve_latex_rendering()\n",
+    "particle_db = load_pdg()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Decay definition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Particle definitions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "scroll-output",
+     "full-width",
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "def generate_markdown_table(particles: list[str]):\n",
+    "    src = dedent(r\"\"\"\n",
+    "    | Particle | Name | PID | $J^{PC} (I^G)$ | $I_3$ | $M$ | $\\Gamma$ | $Q$ | $S$ | $B$ |\n",
+    "    | :------- |------|-----|----------------|-------|-----|----------|-----|-----|-----|\n",
+    "    \"\"\")\n",
+    "    for name in particles:\n",
+    "        p = particle_db[name]\n",
+    "        src += f\"| ${p.latex}$ | `{p.name}` | {p.pid} | {jpc_ig(p)} | {i_3(p)} |  {p.mass:.3g}| {p.width:g} | {p.charge} |{p.strangeness} | {p.baryon_number}|\\n\"\n",
+    "    return src\n",
+    "\n",
+    "\n",
+    "def jpc_ig(particle: Particle) -> str:\n",
+    "    j = format_fraction(particle.spin)\n",
+    "    p = format_parity(particle.parity)\n",
+    "    c = format_parity(particle.c_parity)\n",
+    "    if particle.isospin is None:\n",
+    "        return f\"${j}^{{{p}{c}}}$\"\n",
+    "    i = format_fraction(particle.isospin.magnitude)\n",
+    "    g = format_parity(particle.g_parity)\n",
+    "    return rf\"${j}^{{{p}{c}}} \\; ({i}^{{{g}}})$\"\n",
+    "\n",
+    "\n",
+    "def i_3(particle: Particle) -> str:\n",
+    "    if particle.isospin is None:\n",
+    "        return \"N/A\"\n",
+    "    return f\"${format_fraction(particle.isospin.projection)}$\"\n",
+    "\n",
+    "\n",
+    "def format_fraction(value: float) -> str:\n",
+    "    value = Fraction(value)\n",
+    "    if value.denominator == 1:\n",
+    "        return str(value.numerator)\n",
+    "    return rf\"\\frac{{{value.numerator}}}{{{value.denominator}}}\"\n",
+    "\n",
+    "\n",
+    "def format_parity(parity: int | None) -> str:\n",
+    "    if parity is None:\n",
+    "        return \" \"\n",
+    "    if parity == -1:\n",
+    "        return \"-\"\n",
+    "    if parity == 1:\n",
+    "        return \"+\"\n",
+    "    raise NotImplementedError\n",
+    "\n",
+    "\n",
+    "particles = [\"Lambda\", \"K+\", \"pi0\", \"gamma\", \"p\"]\n",
+    "src = generate_markdown_table(particles)\n",
+    "Markdown(src)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In the table above, PID is the PDG ID from PDG particle numbering scheme,\n",
+    "$J$ is the spin, $P$ is the parity, $C$ is the C parity, $I$ is the isospin (magnitude), $G$ is the G parity.\n",
+    "$I_3$ is the isospin projection (or the 3rd component),\n",
+    "$M$ is the mass,\n",
+    "$\\Gamma$ is the width,\n",
+    "$Q$ is the charge,\n",
+    "$S$ is the strangeness number,\n",
+    "and $B$ is the baryon number."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Initial state definition"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Mass for $p \\gamma$ system"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "E_lab_gamma = 8.5\n",
+    "m_proton = 0.938\n",
+    "m_0 = np.sqrt(2 * E_lab_gamma * m_proton + m_proton**2)\n",
+    "m_eta = 0.548\n",
+    "m_pi = 0.135\n",
+    "m_0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Add custom particle $p \\gamma$"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pgamma1 = Particle(\n",
+    "    name=\"pgamma1\",\n",
+    "    latex=r\"p\\gamma (s1/2)\",\n",
+    "    spin=0.5,\n",
+    "    mass=m_0,\n",
+    "    charge=1,\n",
+    "    isospin=Spin(1 / 2, +1 / 2),\n",
+    "    baryon_number=1,\n",
+    "    parity=-1,\n",
+    "    pid=99990,\n",
+    ")\n",
+    "pgamma2 = create_particle(\n",
+    "    template_particle=pgamma1,\n",
+    "    name=\"pgamma2\",\n",
+    "    latex=R\"p\\gamma (s3/2)\",\n",
+    "    spin=1.5,\n",
+    "    pid=pgamma1.pid + 1,\n",
+    ")\n",
+    "particle_db.update([pgamma1, pgamma2])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Generate transitions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For simplicity, we use the initial state $p \\gamma$ (with spin-$\\frac{1}{2}$), \n",
+    "and set the allowed interaction type to be strong only,\n",
+    "the formalism is selected to be helicity formalism instead of canonical."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{seealso}\n",
+    "[Helicity versus canonical](https://ampform.readthedocs.io/stable/usage/helicity/formalism.html)\n",
+    ":::"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "reaction = qrules.generate_transitions(\n",
+    "    initial_state=(\"pgamma1\"),\n",
+    "    final_state=[\"Lambda\", \"K+\", \"pi0\"],\n",
+    "    allowed_interaction_types=[\"strong\"],\n",
+    "    formalism=\"helicity\",\n",
+    "    particle_db=particle_db,\n",
+    "    max_angular_momentum=4,\n",
+    "    max_spin_magnitude=4,\n",
+    "    mass_conservation_factor=0,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
+    "graphviz.Source(dot)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Formulate amplitude model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_builder = ampform.get_builder(reaction)\n",
+    "model_builder.config.scalar_initial_state_mass = True\n",
+    "model_builder.config.stable_final_state_ids = 0, 1, 2\n",
+    "bw_builder = RelativisticBreitWignerBuilder(\n",
+    "    energy_dependent_width=False,\n",
+    "    form_factor=False,\n",
+    ")\n",
+    "for name in reaction.get_intermediate_particles().names:\n",
+    "    model_builder.dynamics.assign(name, bw_builder)\n",
+    "model = model_builder.formulate()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.intensity"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The first term in the amplitude model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input",
+     "scroll-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "(symbol, expr), *_ = model.amplitudes.items()\n",
+    "Math(aslatex({symbol: expr}, terms_per_line=1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "mystnb": {
+     "code_prompt_show": "Model parameters"
+    },
+    "tags": [
+     "scroll-output",
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "Math(aslatex(model.parameter_defaults))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "mystnb": {
+     "code_prompt_show": "Kinematic variable definitions"
+    },
+    "tags": [
+     "hide-input",
+     "scroll-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "Math(aslatex(model.kinematic_variables))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Visualization"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "unfolded_expression = model.expression.doit()\n",
+    "intensity_func = create_parametrized_function(\n",
+    "    expression=unfolded_expression,\n",
+    "    parameters=model.parameter_defaults,\n",
+    "    backend=\"jax\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "remove-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "phsp_event = 500_000\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",
+    "phsp_momenta = phsp_generator.generate(phsp_event, rng)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "helicity_transformer = SympyDataTransformer.from_sympy(\n",
+    "    model.kinematic_variables,\n",
+    "    backend=\"jax\",\n",
+    ")\n",
+    "phsp = helicity_transformer(phsp_momenta)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-output",
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "resonances = defaultdict(set)\n",
+    "for transition in reaction.transitions:\n",
+    "    topology = transition.topology\n",
+    "    top_decay_products = topology.get_edge_ids_outgoing_from_node(0)\n",
+    "    (resonance_id, resonance), *_ = transition.intermediate_states.items()\n",
+    "    recoil_id, *_ = top_decay_products - {resonance_id}\n",
+    "    resonances[recoil_id].add(resonance.particle)\n",
+    "resonances = {k: sorted(v, key=lambda p: p.mass) for k, v in resonances.items()}\n",
+    "{k: [p.name for p in v] for k, v in resonances.items()}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Design slider UI"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "sliders = {}\n",
+    "categorized_sliders_m = defaultdict(list)\n",
+    "categorized_sliders_gamma = defaultdict(list)\n",
+    "categorized_cphi_pair = defaultdict(list)\n",
+    "\n",
+    "for symbol, value in model.parameter_defaults.items():\n",
+    "    if symbol.name.startswith(R\"\\Gamma_{\"):\n",
+    "        slider = w.FloatSlider(\n",
+    "            description=Rf\"\\({sp.latex(symbol)}\\)\",\n",
+    "            min=0.0,\n",
+    "            max=1.0,\n",
+    "            value=value,\n",
+    "            continuous_update=False,\n",
+    "        )\n",
+    "        sliders[symbol.name] = slider\n",
+    "        if symbol.name.startswith(R\"\\Gamma_{K\"):\n",
+    "            categorized_sliders_gamma[0].append(slider)\n",
+    "        elif symbol.name.startswith(R\"\\Gamma_{\\S\"):\n",
+    "            categorized_sliders_gamma[1].append(slider)\n",
+    "        elif symbol.name.startswith(R\"\\Gamma_{N\"):\n",
+    "            categorized_sliders_gamma[2].append(slider)\n",
+    "\n",
+    "    if symbol.name.startswith(\"m_{\"):\n",
+    "        slider = w.FloatSlider(\n",
+    "            description=Rf\"\\({sp.latex(symbol)}\\)\",\n",
+    "            min=0.63,\n",
+    "            max=4,\n",
+    "            value=value,\n",
+    "            continuous_update=False,\n",
+    "        )\n",
+    "        sliders[symbol.name] = slider\n",
+    "        if symbol.name.startswith(\"m_{K\"):\n",
+    "            categorized_sliders_m[0].append(slider)\n",
+    "        elif symbol.name.startswith(R\"m_{\\S\"):\n",
+    "            categorized_sliders_m[1].append(slider)\n",
+    "        elif symbol.name.startswith(\"m_{N\"):\n",
+    "            categorized_sliders_m[2].append(slider)\n",
+    "\n",
+    "    if symbol.name.startswith(\"C_{\"):\n",
+    "        c_latex = sp.latex(symbol)\n",
+    "        phi_latex = c_latex.replace(\"C\", R\"\\phi\", 1)\n",
+    "\n",
+    "        slider_c = w.FloatSlider(\n",
+    "            description=Rf\"\\({c_latex}\\)\",\n",
+    "            min=0,\n",
+    "            max=10,\n",
+    "            value=abs(value),\n",
+    "            continuous_update=False,\n",
+    "        )\n",
+    "        slider_phi = w.FloatSlider(\n",
+    "            description=Rf\"\\({phi_latex}\\)\",\n",
+    "            min=-np.pi,\n",
+    "            max=+np.pi,\n",
+    "            value=np.angle(value),\n",
+    "            continuous_update=False,\n",
+    "        )\n",
+    "\n",
+    "        sliders[symbol.name] = slider_c\n",
+    "        sliders[symbol.name.replace(\"C\", \"phi\", 1)] = slider_phi\n",
+    "\n",
+    "        cphi_hbox = w.HBox([slider_c, slider_phi])\n",
+    "        if \"Sigma\" in symbol.name:\n",
+    "            categorized_cphi_pair[1].append(cphi_hbox)\n",
+    "        elif R\"\\to N\" in symbol.name:\n",
+    "            categorized_cphi_pair[2].append(cphi_hbox)\n",
+    "        else:\n",
+    "            categorized_cphi_pair[0].append(cphi_hbox)\n",
+    "\n",
+    "\n",
+    "assert len(categorized_sliders_gamma) == 3\n",
+    "assert len(categorized_sliders_m) == 3\n",
+    "assert len(categorized_cphi_pair) == 3\n",
+    "\n",
+    "subtabs = {}\n",
+    "for category, resonance_list in resonances.items():\n",
+    "    subtabs[category] = []\n",
+    "    for particle in resonance_list:\n",
+    "        m_sliders = [\n",
+    "            slider\n",
+    "            for slider in categorized_sliders_m[category]\n",
+    "            if particle.latex in slider.description\n",
+    "        ]\n",
+    "        gamma_sliders = [\n",
+    "            slider\n",
+    "            for slider in categorized_sliders_gamma[category]\n",
+    "            if particle.latex in slider.description\n",
+    "        ]\n",
+    "        cphi_pairs = [\n",
+    "            hbox\n",
+    "            for hbox in categorized_cphi_pair[category]\n",
+    "            if particle.latex in hbox.children[0].description\n",
+    "        ]\n",
+    "        pole_pair = w.HBox(m_sliders + gamma_sliders)\n",
+    "        resonance_tab = w.VBox([pole_pair, *cphi_pairs])\n",
+    "        subtabs[category].append(resonance_tab)\n",
+    "assert len(subtabs) == 3\n",
+    "\n",
+    "main_tabs = []\n",
+    "for category, slider_boxes in subtabs.items():\n",
+    "    sub_tab_widget = w.Tab(children=slider_boxes)\n",
+    "    for i, particle in enumerate(resonances[category]):\n",
+    "        sub_tab_widget.set_title(i, particle.name)\n",
+    "\n",
+    "    main_tabs.append(sub_tab_widget)\n",
+    "\n",
+    "UI = w.Tab(children=main_tabs, titles=(\"K*\", \"Σ*\", \"N*\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    }
+   },
+   "outputs": [],
+   "source": [
+    "def insert_phi(parameters: dict) -> dict:\n",
+    "    updated_parameters = {}\n",
+    "    for key, value in parameters.items():\n",
+    "        if key.startswith(\"phi_\"):\n",
+    "            continue\n",
+    "        if key.startswith(\"C_\"):\n",
+    "            phi_key = key.replace(\"C_\", \"phi_\")\n",
+    "            if phi_key in parameters:\n",
+    "                phi = parameters[phi_key]\n",
+    "                value *= np.exp(1j * phi)  # noqa:PLW2901\n",
+    "        updated_parameters[key] = value\n",
+    "    return updated_parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input",
+     "full-width"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "%matplotlib widget\n",
+    "%config InlineBackend.figure_formats = ['png']\n",
+    "fig_2d, ax_2d = plt.subplots(dpi=200)\n",
+    "ax_2d.set_title(\"Model-weighted Phase space Dalitz Plot\")\n",
+    "ax_2d.set_xlabel(R\"$m^2(\\Lambda K^+)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n",
+    "ax_2d.set_ylabel(R\"$m^2(K^+ \\pi^0)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n",
+    "\n",
+    "mesh = None\n",
+    "\n",
+    "\n",
+    "def update_histogram(**parameters):\n",
+    "    global mesh\n",
+    "    parameters = insert_phi(parameters)\n",
+    "    intensity_func.update_parameters(parameters)\n",
+    "    intensity_weights = intensity_func(phsp)\n",
+    "    bin_values, xedges, yedges = jnp.histogram2d(\n",
+    "        phsp[\"m_01\"].real ** 2,\n",
+    "        phsp[\"m_12\"].real ** 2,\n",
+    "        bins=200,\n",
+    "        weights=intensity_weights,\n",
+    "        density=True,\n",
+    "    )\n",
+    "    bin_values = jnp.where(bin_values < 1e-6, jnp.nan, bin_values)\n",
+    "    x, y = jnp.meshgrid(xedges[:-1], yedges[:-1])\n",
+    "    if mesh is None:\n",
+    "        mesh = ax_2d.pcolormesh(x, y, bin_values.T, cmap=\"jet\", vmax=0.15)\n",
+    "    else:\n",
+    "        mesh.set_array(bin_values.T)\n",
+    "    fig_2d.canvas.draw_idle()\n",
+    "\n",
+    "\n",
+    "interactive_plot = w.interactive_output(update_histogram, sliders)\n",
+    "fig_2d.tight_layout()\n",
+    "fig_2d.colorbar(mesh, ax=ax_2d)\n",
+    "display(UI, interactive_plot)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input",
+     "full-width",
+     "scroll-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "%matplotlib widget\n",
+    "%config InlineBackend.figure_formats = ['svg']\n",
+    "\n",
+    "fig, axes = plt.subplots(figsize=(11, 3.5), ncols=3, sharey=True)\n",
+    "fig.canvas.toolbar_visible = False\n",
+    "fig.canvas.header_visible = False\n",
+    "fig.canvas.footer_visible = False\n",
+    "ax1, ax2, ax3 = axes\n",
+    "\n",
+    "for recoil_id, ax in enumerate(axes):\n",
+    "    decay_products = sorted({0, 1, 2} - {recoil_id})\n",
+    "    product_latex = \" \".join([reaction.final_state[i].latex for i in decay_products])\n",
+    "    ax.set_xlabel(f\"$m({product_latex})$ [GeV]\")\n",
+    "\n",
+    "LINES = 3 * [None]\n",
+    "RESONANCE_LINE = [None] * len(reaction.get_intermediate_particles())\n",
+    "\n",
+    "\n",
+    "def update_plot(**parameters):\n",
+    "    parameters = insert_phi(parameters)\n",
+    "    intensity_func.update_parameters(parameters)\n",
+    "    intensities = intensity_func(phsp)\n",
+    "    max_value = 0\n",
+    "    color_id = 0\n",
+    "    for recoil_id, ax in enumerate(axes):\n",
+    "        decay_products = sorted({0, 1, 2} - {recoil_id})\n",
+    "        key = f\"m_{''.join(str(i) for i in decay_products)}\"\n",
+    "        bin_values, bin_edges = jax.numpy.histogram(\n",
+    "            phsp[key].real,\n",
+    "            bins=120,\n",
+    "            density=True,\n",
+    "            weights=intensities,\n",
+    "        )\n",
+    "        max_value = max(max_value, bin_values.max())\n",
+    "\n",
+    "        if LINES[recoil_id] is None:\n",
+    "            LINES[recoil_id] = ax.step(bin_edges[:-1], bin_values, alpha=0.5)[0]\n",
+    "        else:\n",
+    "            LINES[recoil_id].set_ydata(bin_values)\n",
+    "\n",
+    "        for resonance in resonances[recoil_id]:\n",
+    "            key = f\"m_{{{resonance.latex}}}\"\n",
+    "            val = parameters.get(key, resonance.mass)\n",
+    "            if RESONANCE_LINE[color_id] is None:\n",
+    "                RESONANCE_LINE[color_id] = ax.axvline(\n",
+    "                    val,\n",
+    "                    c=f\"C{color_id}\",\n",
+    "                    ls=\"dotted\",\n",
+    "                    label=resonance.name,\n",
+    "                )\n",
+    "            else:\n",
+    "                RESONANCE_LINE[color_id].set_xdata([val, val])\n",
+    "            color_id += 1\n",
+    "\n",
+    "    for ax in axes:\n",
+    "        ax.set_ylim(0, max_value * 1.1)\n",
+    "\n",
+    "\n",
+    "interactive_plot = w.interactive_output(update_plot, sliders)\n",
+    "for ax in axes:\n",
+    "    ax.legend(fontsize=\"small\")\n",
+    "\n",
+    "fig.tight_layout()\n",
+    "display(UI, interactive_plot)"
+   ]
+  }
+ ],
+ "metadata": {
+  "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.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/docs/lambda-k-pi/automated.ipynb b/docs/lambda-k-pi/ampform.ipynb
similarity index 100%
rename from docs/lambda-k-pi/automated.ipynb
rename to docs/lambda-k-pi/ampform.ipynb

From 1e27d49a6943bf6051c7a994353d054fd07f6236 Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Wed, 14 Aug 2024 12:04:39 +0200
Subject: [PATCH 2/8] name in toctree

---
 docs/lambda-k-pi/index.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/lambda-k-pi/index.md b/docs/lambda-k-pi/index.md
index a3afaa4..78a8251 100644
--- a/docs/lambda-k-pi/index.md
+++ b/docs/lambda-k-pi/index.md
@@ -1,5 +1,6 @@
 # pγ → ΛK⁺π⁰
 
 ```{toctree}
-automated
+ampform
+ampform-dpd
 ```

From 71bde2f93b7e0a77dc5451cbf117fa64805170e7 Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Wed, 14 Aug 2024 14:44:08 +0200
Subject: [PATCH 3/8] add ampform-dpd via PyPI

---
 pixi.lock      | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++
 pyproject.toml |  1 +
 2 files changed, 65 insertions(+)

diff --git a/pixi.lock b/pixi.lock
index 05a3069..7cbe53d 100644
--- a/pixi.lock
+++ b/pixi.lock
@@ -158,6 +158,7 @@ environments:
       - pypi: https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/7e/b3/6b4067be973ae96ba0d615946e314c5ae35f9f993eca561b356540bb0c2b/alabaster-1.0.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/15/90/60468aaed8bb12b5bb0e8e100c95ff33c9d097d309e20fbe14b26406307b/ampform-0.15.4-py3-none-any.whl
+      - pypi: https://files.pythonhosted.org/packages/9d/bd/8b662d1eabe1d920fc44db38551483ef450322ef0faefd413f9a0a094621/ampform_dpd-0.2.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/7b/a2/10639a79341f6c019dedc95bd48a4928eed9f1d1197f4c04f546fc7ae0ff/anyio-4.4.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/a4/6a/e8a041599e78b6b3752da48000b14c8d1e8a04ded09c88c714ba047f34f5/argon2_cffi-23.1.0-py3-none-any.whl
       - pypi: https://files.pythonhosted.org/packages/ec/f7/378254e6dd7ae6f31fe40c8649eea7d4832a42243acaf0f1fff9083b2bed/argon2_cffi_bindings-21.2.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
@@ -461,6 +462,69 @@ packages:
   - sphinx-api-relink>=0.0.3 ; extra == 'types'
   - graphviz ; extra == 'viz'
   requires_python: '>=3.7'
+- kind: pypi
+  name: ampform-dpd
+  version: 0.2.0
+  url: https://files.pythonhosted.org/packages/9d/bd/8b662d1eabe1d920fc44db38551483ef450322ef0faefd413f9a0a094621/ampform_dpd-0.2.0-py3-none-any.whl
+  sha256: b2e06a5efd900c2454168c2a24cde21dc319b95e7f77a6c74a24f48b0588bd65
+  requires_dist:
+  - ampform>=0.14.8
+  - attrs>=20.1.0
+  - cloudpickle
+  - qrules>=0.10.0
+  - sympy>=1.10
+  - tensorwaves[jax]
+  - ampform-dpd[doc] ; extra == 'dev'
+  - ampform-dpd[jupyter] ; extra == 'dev'
+  - ampform-dpd[sty] ; extra == 'dev'
+  - ampform-dpd[test] ; extra == 'dev'
+  - sphinx-autobuild ; extra == 'dev'
+  - tox>=1.9 ; extra == 'dev'
+  - graphviz ; extra == 'doc'
+  - ipympl ; extra == 'doc'
+  - ipywidgets ; extra == 'doc'
+  - matplotlib ; extra == 'doc'
+  - myst-nb>=0.14 ; extra == 'doc'
+  - sphinx-api-relink>=0.0.4 ; extra == 'doc'
+  - sphinx-book-theme ; extra == 'doc'
+  - sphinx-codeautolink[ipython] ; extra == 'doc'
+  - sphinx-copybutton ; extra == 'doc'
+  - sphinx-design ; extra == 'doc'
+  - sphinx-pybtex-etal-style ; extra == 'doc'
+  - sphinx-togglebutton ; extra == 'doc'
+  - sphinxcontrib-bibtex ; extra == 'doc'
+  - tensorwaves[phsp] ; extra == 'doc'
+  - tqdm ; extra == 'doc'
+  - tensorwaves[jax] ; extra == 'jax'
+  - ampform-dpd[doc] ; extra == 'jupyter'
+  - black ; extra == 'jupyter'
+  - isort ; extra == 'jupyter'
+  - jupyterlab>=3.0 ; extra == 'jupyter'
+  - jupyterlab ; extra == 'jupyter'
+  - jupyterlab-code-formatter ; extra == 'jupyter'
+  - jupyterlab-git ; extra == 'jupyter'
+  - jupyterlab-lsp ; extra == 'jupyter'
+  - jupyterlab-myst ; extra == 'jupyter'
+  - python-lsp-ruff ; extra == 'jupyter'
+  - python-lsp-server[rope] ; extra == 'jupyter'
+  - tensorwaves[numba] ; extra == 'numba'
+  - ampform-dpd[types] ; extra == 'sty'
+  - mypy ; extra == 'sty'
+  - pre-commit>=1.4.0 ; extra == 'sty'
+  - ruff ; extra == 'sty'
+  - ampform-dpd[tf] ; extra == 'tensorflow'
+  - nbmake ; extra == 'test'
+  - numpy ; extra == 'test'
+  - pytest>=6.0 ; extra == 'test'
+  - pytest-cov ; extra == 'test'
+  - pytest-xdist ; extra == 'test'
+  - tensorwaves[tf] ; extra == 'tf'
+  - docutils ; extra == 'types'
+  - pybtex ; extra == 'types'
+  - pytest ; extra == 'types'
+  - sphinx ; extra == 'types'
+  - sphinx-api-relink>=0.0.4 ; extra == 'types'
+  requires_python: '>=3.8'
 - kind: pypi
   name: anyio
   version: 4.4.0
diff --git a/pyproject.toml b/pyproject.toml
index ea1a57a..1e0ef22 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,6 +6,7 @@ requires = ["setuptools"]
 authors = [{name = "Vitor Shen"}]
 dependencies = [
     "ampform",
+    "ampform-dpd>=0.2.0,<0.3",
     "graphviz",
     "ipympl",
     "ipywidgets",

From d36d52e473e49009ba2d9ddd99763342189fd6fc Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Fri, 16 Aug 2024 11:23:36 +0200
Subject: [PATCH 4/8] model formulation

---
 docs/lambda-k-pi/ampform-dpd.ipynb | 91 +++++++++++++++++++++++++-----
 pyproject.toml                     |  1 +
 2 files changed, 78 insertions(+), 14 deletions(-)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
index adbd634..d6c6a08 100644
--- a/docs/lambda-k-pi/ampform-dpd.ipynb
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -19,9 +19,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "mystnb": {
      "code_prompt_show": "Import Python libraries"
     },
@@ -50,8 +47,15 @@
     "import qrules\n",
     "import sympy as sp\n",
     "from ampform.dynamics.builder import RelativisticBreitWignerBuilder\n",
-    "from ampform.io import aslatex, improve_latex_rendering\n",
-    "from IPython.display import Markdown, Math, display\n",
+    "from ampform.io import improve_latex_rendering\n",
+    "from ampform_dpd import DalitzPlotDecompositionBuilder\n",
+    "from ampform_dpd.adapter.qrules import (\n",
+    "    to_three_body_decay,\n",
+    ")\n",
+    "from ampform_dpd.decay import DecayNode, ThreeBodyDecayChain\n",
+    "from ampform_dpd.dynamics.builder import create_mass_symbol, get_mandelstam_s\n",
+    "from ampform_dpd.io import aslatex\n",
+    "from IPython.display import Latex, Markdown, Math, display\n",
     "from qrules.particle import Particle, Spin, create_particle, load_pdg\n",
     "from tensorwaves.data import (\n",
     "    SympyDataTransformer,\n",
@@ -60,6 +64,8 @@
     ")\n",
     "from tensorwaves.function.sympy import create_parametrized_function\n",
     "\n",
+    "STATIC_PAGE = \"EXECUTE_NB\" in os.environ\n",
+    "\n",
     "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
     "logging.disable(logging.WARNING)\n",
     "warnings.filterwarnings(\"ignore\")\n",
@@ -283,6 +289,16 @@
     "graphviz.Source(dot)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "decay = to_three_body_decay(reaction.transitions)\n",
+    "Latex(aslatex(decay, with_jp=True))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -317,6 +333,60 @@
     "model.intensity"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def formulate_breit_wigner(\n",
+    "    decay: ThreeBodyDecayChain,\n",
+    ") -> tuple[sp.Expr, dict[sp.Symbol, complex | float]]:\n",
+    "    decay_node = decay.decay_node\n",
+    "    s = get_mandelstam_s(decay_node)\n",
+    "    parameter_defaults = {}\n",
+    "    breit_wigner, new_pars = _create_breit_wigner(s, decay_node)\n",
+    "    parameter_defaults.update(new_pars)\n",
+    "    return (\n",
+    "        breit_wigner,\n",
+    "        parameter_defaults,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def _create_breit_wigner(\n",
+    "    s: sp.Symbol, isobar: DecayNode\n",
+    ") -> tuple[sp.Expr, dict[sp.Symbol, complex | float]]:\n",
+    "    mass = create_mass_symbol(isobar.parent)\n",
+    "    width = sp.Symbol(Rf\"\\Gamma_{{{isobar.parent.latex}}}\", nonnegative=True)\n",
+    "    breit_wigner_expr = (mass * width) / (mass**2 - s - width * mass * sp.I)\n",
+    "    parameter_defaults: dict[sp.Symbol, complex | float] = {\n",
+    "        mass: isobar.parent.mass,\n",
+    "        width: isobar.parent.width,\n",
+    "    }\n",
+    "    return breit_wigner_expr, parameter_defaults"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_builder = DalitzPlotDecompositionBuilder(decay)\n",
+    "for chain in model_builder.decay.chains:\n",
+    "    model_builder.dynamics_choices.register_builder(chain, formulate_breit_wigner)\n",
+    "model = model_builder.formulate(reference_subsystem=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.intensity"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -373,14 +443,7 @@
    },
    "outputs": [],
    "source": [
-    "Math(aslatex(model.kinematic_variables))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Visualization"
+    "Math(aslatex(model.variables))"
    ]
   },
   {
@@ -751,7 +814,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.5"
   }
  },
  "nbformat": 4,
diff --git a/pyproject.toml b/pyproject.toml
index 1e0ef22..45228dc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -114,6 +114,7 @@ builtins-ignorelist = ["display"]
     "PLW0603",
     "S101",
     "T201",
+    "TCH",
 ]
 "docs/conf.py" = [
     "D100",

From 7e8975efb1e442db43e854289c5ec5e2b39d79de Mon Sep 17 00:00:00 2001
From: Vitor Shen <17490173+shenvitor@users.noreply.github.com>
Date: Fri, 16 Aug 2024 15:42:10 +0200
Subject: [PATCH 5/8] Finalize draft of Dalitz plot and 1d projection

---
 docs/lambda-k-pi/ampform-dpd.ipynb | 202 +++++++++++++++++------------
 pyproject.toml                     |   1 +
 2 files changed, 118 insertions(+), 85 deletions(-)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
index d6c6a08..7d318b7 100644
--- a/docs/lambda-k-pi/ampform-dpd.ipynb
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -40,7 +40,6 @@
     "import ampform\n",
     "import graphviz\n",
     "import ipywidgets as w\n",
-    "import jax\n",
     "import jax.numpy as jnp\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
@@ -49,9 +48,7 @@
     "from ampform.dynamics.builder import RelativisticBreitWignerBuilder\n",
     "from ampform.io import improve_latex_rendering\n",
     "from ampform_dpd import DalitzPlotDecompositionBuilder\n",
-    "from ampform_dpd.adapter.qrules import (\n",
-    "    to_three_body_decay,\n",
-    ")\n",
+    "from ampform_dpd.adapter.qrules import normalize_state_ids, to_three_body_decay\n",
     "from ampform_dpd.decay import DecayNode, ThreeBodyDecayChain\n",
     "from ampform_dpd.dynamics.builder import create_mass_symbol, get_mandelstam_s\n",
     "from ampform_dpd.io import aslatex\n",
@@ -59,8 +56,6 @@
     "from qrules.particle import Particle, Spin, create_particle, load_pdg\n",
     "from tensorwaves.data import (\n",
     "    SympyDataTransformer,\n",
-    "    TFPhaseSpaceGenerator,\n",
-    "    TFUniformRealNumberGenerator,\n",
     ")\n",
     "from tensorwaves.function.sympy import create_parametrized_function\n",
     "\n",
@@ -269,16 +264,14 @@
     "    max_angular_momentum=4,\n",
     "    max_spin_magnitude=4,\n",
     "    mass_conservation_factor=0,\n",
-    ")"
+    ")\n",
+    "reaction = normalize_state_ids(reaction)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "tags": [
      "hide-input"
     ]
@@ -314,7 +307,7 @@
    "source": [
     "model_builder = ampform.get_builder(reaction)\n",
     "model_builder.config.scalar_initial_state_mass = True\n",
-    "model_builder.config.stable_final_state_ids = 0, 1, 2\n",
+    "model_builder.config.stable_final_state_ids = list(reaction.final_state)\n",
     "bw_builder = RelativisticBreitWignerBuilder(\n",
     "    energy_dependent_width=False,\n",
     "    form_factor=False,\n",
@@ -446,13 +439,20 @@
     "Math(aslatex(model.variables))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Visualization"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "unfolded_expression = model.expression.doit()\n",
+    "unfolded_expression = model.full_expression.doit()\n",
     "intensity_func = create_parametrized_function(\n",
     "    expression=unfolded_expression,\n",
     "    parameters=model.parameter_defaults,\n",
@@ -460,45 +460,81 @@
     ")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "i, j = 3, 1\n",
+    "(k,) = {1, 2, 3} - {i, j}\n",
+    "σk, σk_expr = list(model.invariants.items())[k - 1]\n",
+    "Latex(aslatex({σk: σk_expr}))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
     "tags": [
-     "remove-output"
+     "hide-input"
     ]
    },
    "outputs": [],
    "source": [
-    "phsp_event = 500_000\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",
-    "phsp_momenta = phsp_generator.generate(phsp_event, rng)"
+    "resolution = 1_000\n",
+    "m = sorted(model.masses, key=str)\n",
+    "x_min = float(((m[j] + m[k]) ** 2).xreplace(model.masses))\n",
+    "x_max = float(((m[0] - m[i]) ** 2).xreplace(model.masses))\n",
+    "y_min = float(((m[i] + m[k]) ** 2).xreplace(model.masses))\n",
+    "y_max = float(((m[0] - m[j]) ** 2).xreplace(model.masses))\n",
+    "x_diff = x_max - x_min\n",
+    "y_diff = y_max - y_min\n",
+    "x_min -= 0.05 * x_diff\n",
+    "x_max += 0.05 * x_diff\n",
+    "y_min -= 0.05 * y_diff\n",
+    "y_max += 0.05 * y_diff\n",
+    "X, Y = jnp.meshgrid(\n",
+    "    jnp.linspace(x_min, x_max, num=resolution),\n",
+    "    jnp.linspace(y_min, y_max, num=resolution),\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
    "outputs": [],
    "source": [
-    "helicity_transformer = SympyDataTransformer.from_sympy(\n",
-    "    model.kinematic_variables,\n",
-    "    backend=\"jax\",\n",
-    ")\n",
-    "phsp = helicity_transformer(phsp_momenta)"
+    "definitions = dict(model.variables)\n",
+    "definitions[σk] = σk_expr\n",
+    "definitions = {\n",
+    "    symbol: expr.xreplace(definitions).xreplace(model.masses)\n",
+    "    for symbol, expr in definitions.items()\n",
+    "}\n",
+    "\n",
+    "data_transformer = SympyDataTransformer.from_sympy(definitions, backend=\"jax\")\n",
+    "dalitz_data = {\n",
+    "    f\"sigma{i}\": X,\n",
+    "    f\"sigma{j}\": Y,\n",
+    "}\n",
+    "dalitz_data.update(data_transformer(dalitz_data))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "jupyter": {
-     "source_hidden": true
-    },
     "tags": [
      "hide-output",
      "hide-input"
@@ -537,6 +573,7 @@
     "categorized_sliders_m = defaultdict(list)\n",
     "categorized_sliders_gamma = defaultdict(list)\n",
     "categorized_cphi_pair = defaultdict(list)\n",
+    "couplings_name_root = R\"\\mathcal{H}^\\mathrm{decay}\"\n",
     "\n",
     "for symbol, value in model.parameter_defaults.items():\n",
     "    if symbol.name.startswith(R\"\\Gamma_{\"):\n",
@@ -549,11 +586,11 @@
     "        )\n",
     "        sliders[symbol.name] = slider\n",
     "        if symbol.name.startswith(R\"\\Gamma_{K\"):\n",
-    "            categorized_sliders_gamma[0].append(slider)\n",
-    "        elif symbol.name.startswith(R\"\\Gamma_{\\S\"):\n",
     "            categorized_sliders_gamma[1].append(slider)\n",
-    "        elif symbol.name.startswith(R\"\\Gamma_{N\"):\n",
+    "        elif symbol.name.startswith(R\"\\Gamma_{\\S\"):\n",
     "            categorized_sliders_gamma[2].append(slider)\n",
+    "        elif symbol.name.startswith(R\"\\Gamma_{N\"):\n",
+    "            categorized_sliders_gamma[3].append(slider)\n",
     "\n",
     "    if symbol.name.startswith(\"m_{\"):\n",
     "        slider = w.FloatSlider(\n",
@@ -565,15 +602,15 @@
     "        )\n",
     "        sliders[symbol.name] = slider\n",
     "        if symbol.name.startswith(\"m_{K\"):\n",
-    "            categorized_sliders_m[0].append(slider)\n",
-    "        elif symbol.name.startswith(R\"m_{\\S\"):\n",
     "            categorized_sliders_m[1].append(slider)\n",
-    "        elif symbol.name.startswith(\"m_{N\"):\n",
+    "        elif symbol.name.startswith(R\"m_{\\S\"):\n",
     "            categorized_sliders_m[2].append(slider)\n",
+    "        elif symbol.name.startswith(\"m_{N\"):\n",
+    "            categorized_sliders_m[3].append(slider)\n",
     "\n",
-    "    if symbol.name.startswith(\"C_{\"):\n",
+    "    if symbol.name.startswith(couplings_name_root):\n",
     "        c_latex = sp.latex(symbol)\n",
-    "        phi_latex = c_latex.replace(\"C\", R\"\\phi\", 1)\n",
+    "        phi_latex = c_latex.replace(couplings_name_root, R\"\\phi\", 1)\n",
     "\n",
     "        slider_c = w.FloatSlider(\n",
     "            description=Rf\"\\({c_latex}\\)\",\n",
@@ -591,15 +628,15 @@
     "        )\n",
     "\n",
     "        sliders[symbol.name] = slider_c\n",
-    "        sliders[symbol.name.replace(\"C\", \"phi\", 1)] = slider_phi\n",
+    "        sliders[symbol.name.replace(couplings_name_root, \"phi\", 1)] = slider_phi\n",
     "\n",
     "        cphi_hbox = w.HBox([slider_c, slider_phi])\n",
     "        if \"Sigma\" in symbol.name:\n",
-    "            categorized_cphi_pair[1].append(cphi_hbox)\n",
-    "        elif R\"\\to N\" in symbol.name:\n",
     "            categorized_cphi_pair[2].append(cphi_hbox)\n",
+    "        elif \"N\" in symbol.name:\n",
+    "            categorized_cphi_pair[3].append(cphi_hbox)\n",
     "        else:\n",
-    "            categorized_cphi_pair[0].append(cphi_hbox)\n",
+    "            categorized_cphi_pair[1].append(cphi_hbox)\n",
     "\n",
     "\n",
     "assert len(categorized_sliders_gamma) == 3\n",
@@ -607,33 +644,33 @@
     "assert len(categorized_cphi_pair) == 3\n",
     "\n",
     "subtabs = {}\n",
-    "for category, resonance_list in resonances.items():\n",
-    "    subtabs[category] = []\n",
+    "for recoild_id, resonance_list in resonances.items():\n",
+    "    subtabs[recoild_id] = []\n",
     "    for particle in resonance_list:\n",
     "        m_sliders = [\n",
     "            slider\n",
-    "            for slider in categorized_sliders_m[category]\n",
+    "            for slider in categorized_sliders_m[recoild_id]\n",
     "            if particle.latex in slider.description\n",
     "        ]\n",
     "        gamma_sliders = [\n",
     "            slider\n",
-    "            for slider in categorized_sliders_gamma[category]\n",
+    "            for slider in categorized_sliders_gamma[recoild_id]\n",
     "            if particle.latex in slider.description\n",
     "        ]\n",
     "        cphi_pairs = [\n",
     "            hbox\n",
-    "            for hbox in categorized_cphi_pair[category]\n",
+    "            for hbox in categorized_cphi_pair[recoild_id]\n",
     "            if particle.latex in hbox.children[0].description\n",
     "        ]\n",
     "        pole_pair = w.HBox(m_sliders + gamma_sliders)\n",
     "        resonance_tab = w.VBox([pole_pair, *cphi_pairs])\n",
-    "        subtabs[category].append(resonance_tab)\n",
+    "        subtabs[recoild_id].append(resonance_tab)\n",
     "assert len(subtabs) == 3\n",
     "\n",
     "main_tabs = []\n",
-    "for category, slider_boxes in subtabs.items():\n",
+    "for recoild_id, slider_boxes in subtabs.items():\n",
     "    sub_tab_widget = w.Tab(children=slider_boxes)\n",
-    "    for i, particle in enumerate(resonances[category]):\n",
+    "    for i, particle in enumerate(resonances[recoild_id]):\n",
     "        sub_tab_widget.set_title(i, particle.name)\n",
     "\n",
     "    main_tabs.append(sub_tab_widget)\n",
@@ -654,10 +691,10 @@
     "def insert_phi(parameters: dict) -> dict:\n",
     "    updated_parameters = {}\n",
     "    for key, value in parameters.items():\n",
-    "        if key.startswith(\"phi_\"):\n",
+    "        if key.startswith(\"phi\"):\n",
     "            continue\n",
-    "        if key.startswith(\"C_\"):\n",
-    "            phi_key = key.replace(\"C_\", \"phi_\")\n",
+    "        if key.startswith(couplings_name_root):\n",
+    "            phi_key = key.replace(couplings_name_root, \"phi\")\n",
     "            if phi_key in parameters:\n",
     "                phi = parameters[phi_key]\n",
     "                value *= np.exp(1j * phi)  # noqa:PLW2901\n",
@@ -689,28 +726,21 @@
     "mesh = None\n",
     "\n",
     "\n",
-    "def update_histogram(**parameters):\n",
+    "def update_dalitz_plot(**parameters):\n",
     "    global mesh\n",
     "    parameters = insert_phi(parameters)\n",
     "    intensity_func.update_parameters(parameters)\n",
-    "    intensity_weights = intensity_func(phsp)\n",
-    "    bin_values, xedges, yedges = jnp.histogram2d(\n",
-    "        phsp[\"m_01\"].real ** 2,\n",
-    "        phsp[\"m_12\"].real ** 2,\n",
-    "        bins=200,\n",
-    "        weights=intensity_weights,\n",
-    "        density=True,\n",
-    "    )\n",
-    "    bin_values = jnp.where(bin_values < 1e-6, jnp.nan, bin_values)\n",
-    "    x, y = jnp.meshgrid(xedges[:-1], yedges[:-1])\n",
+    "    intensities = intensity_func(dalitz_data)  # z\n",
+    "    intensities /= jnp.nansum(intensities)  # normalization\n",
+    "\n",
     "    if mesh is None:\n",
-    "        mesh = ax_2d.pcolormesh(x, y, bin_values.T, cmap=\"jet\", vmax=0.15)\n",
+    "        mesh = ax_2d.pcolormesh(X, Y, intensities, cmap=\"jet\", vmax=3e-5)\n",
     "    else:\n",
-    "        mesh.set_array(bin_values.T)\n",
+    "        mesh.set_array(intensities)\n",
     "    fig_2d.canvas.draw_idle()\n",
     "\n",
     "\n",
-    "interactive_plot = w.interactive_output(update_histogram, sliders)\n",
+    "interactive_plot = w.interactive_output(update_dalitz_plot, sliders)\n",
     "fig_2d.tight_layout()\n",
     "fig_2d.colorbar(mesh, ax=ax_2d)\n",
     "display(UI, interactive_plot)"
@@ -734,42 +764,44 @@
     "%matplotlib widget\n",
     "%config InlineBackend.figure_formats = ['svg']\n",
     "\n",
-    "fig, axes = plt.subplots(figsize=(11, 3.5), ncols=3, sharey=True)\n",
+    "fig, axes = plt.subplots(figsize=(11, 3.5), ncols=2, sharey=True)\n",
     "fig.canvas.toolbar_visible = False\n",
     "fig.canvas.header_visible = False\n",
     "fig.canvas.footer_visible = False\n",
-    "ax1, ax2, ax3 = axes\n",
+    "ax1, ax2 = axes\n",
     "\n",
-    "for recoil_id, ax in enumerate(axes):\n",
-    "    decay_products = sorted({0, 1, 2} - {recoil_id})\n",
+    "for ax in axes:\n",
+    "    recoil_id = 3 if ax is ax1 else 1\n",
+    "    decay_products = sorted(set(reaction.final_state) - {recoil_id})\n",
     "    product_latex = \" \".join([reaction.final_state[i].latex for i in decay_products])\n",
     "    ax.set_xlabel(f\"$m({product_latex})$ [GeV]\")\n",
     "\n",
-    "LINES = 3 * [None]\n",
-    "RESONANCE_LINE = [None] * len(reaction.get_intermediate_particles())\n",
+    "LINES = defaultdict(lambda: None)\n",
+    "RESONANCE_LINE = defaultdict(lambda: None)\n",
     "\n",
     "\n",
     "def update_plot(**parameters):\n",
     "    parameters = insert_phi(parameters)\n",
     "    intensity_func.update_parameters(parameters)\n",
-    "    intensities = intensity_func(phsp)\n",
+    "    intensities = intensity_func(dalitz_data)  # z\n",
+    "    intensities /= jnp.nansum(intensities)  # normalization\n",
+    "\n",
     "    max_value = 0\n",
     "    color_id = 0\n",
-    "    for recoil_id, ax in enumerate(axes):\n",
-    "        decay_products = sorted({0, 1, 2} - {recoil_id})\n",
-    "        key = f\"m_{''.join(str(i) for i in decay_products)}\"\n",
-    "        bin_values, bin_edges = jax.numpy.histogram(\n",
-    "            phsp[key].real,\n",
-    "            bins=120,\n",
-    "            density=True,\n",
-    "            weights=intensities,\n",
-    "        )\n",
-    "        max_value = max(max_value, bin_values.max())\n",
+    "    for ax in axes:\n",
+    "        if ax is ax1:\n",
+    "            x = jnp.sqrt(X[0])\n",
+    "            y = jnp.nansum(intensities, axis=0)\n",
+    "        else:\n",
+    "            x = jnp.sqrt(Y[:, 0])\n",
+    "            y = jnp.nansum(intensities, axis=1)\n",
     "\n",
+    "        max_value = max(max_value, y.max())\n",
+    "        recoil_id = 3 if ax is ax1 else 1\n",
     "        if LINES[recoil_id] is None:\n",
-    "            LINES[recoil_id] = ax.step(bin_edges[:-1], bin_values, alpha=0.5)[0]\n",
+    "            LINES[recoil_id] = ax.plot(x, y, alpha=0.5)[0]\n",
     "        else:\n",
-    "            LINES[recoil_id].set_ydata(bin_values)\n",
+    "            LINES[recoil_id].set_ydata(y)\n",
     "\n",
     "        for resonance in resonances[recoil_id]:\n",
     "            key = f\"m_{{{resonance.latex}}}\"\n",
diff --git a/pyproject.toml b/pyproject.toml
index 45228dc..a93930e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -110,6 +110,7 @@ builtins-ignorelist = ["display"]
     "D103",
     "E303",
     "E501",
+    "PLC2401",
     "PLR2004",
     "PLW0603",
     "S101",

From fd1b0a30600e8197b0499af5ec1e4998f692206a Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Fri, 16 Aug 2024 16:10:05 +0200
Subject: [PATCH 6/8] DOC: improve Sphinx rendering

---
 docs/lambda-k-pi/ampform-dpd.ipynb | 56 ++++++++++++++++++++++++++----
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
index 7d318b7..16eba5e 100644
--- a/docs/lambda-k-pi/ampform-dpd.ipynb
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Amplitude model with `ampform-dpd` (Draft)"
+    "# Amplitude model with `ampform-dpd`"
    ]
   },
   {
@@ -19,6 +19,9 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
     "mystnb": {
      "code_prompt_show": "Import Python libraries"
     },
@@ -272,6 +275,9 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
     "tags": [
      "hide-input"
     ]
@@ -285,7 +291,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    }
+   },
    "outputs": [],
    "source": [
     "decay = to_three_body_decay(reaction.transitions)\n",
@@ -329,7 +339,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    }
+   },
    "outputs": [],
    "source": [
     "def formulate_breit_wigner(\n",
@@ -409,6 +423,9 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
     "mystnb": {
      "code_prompt_show": "Model parameters"
     },
@@ -419,7 +436,9 @@
    },
    "outputs": [],
    "source": [
-    "Math(aslatex(model.parameter_defaults))"
+    "if STATIC_PAGE:\n",
+    "    src = aslatex(model.parameter_defaults)\n",
+    "    display(Math(src))"
    ]
   },
   {
@@ -463,7 +482,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    }
+   },
    "outputs": [],
    "source": [
     "i, j = 3, 1\n",
@@ -479,6 +502,9 @@
     "jupyter": {
      "source_hidden": true
     },
+    "mystnb": {
+     "code_prompt_show": "Define meshgrid"
+    },
     "tags": [
      "hide-input"
     ]
@@ -510,6 +536,9 @@
     "jupyter": {
      "source_hidden": true
     },
+    "mystnb": {
+     "code_prompt_show": "Define data transformer"
+    },
     "tags": [
      "hide-input"
     ]
@@ -535,9 +564,16 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Collect resonance names per sub-system"
+    },
     "tags": [
      "hide-output",
-     "hide-input"
+     "hide-input",
+     "remove-output"
     ]
    },
    "outputs": [],
@@ -684,7 +720,13 @@
    "metadata": {
     "jupyter": {
      "source_hidden": true
-    }
+    },
+    "mystnb": {
+     "code_prompt_show": "Define functions for inserting phase"
+    },
+    "tags": [
+     "hide-input"
+    ]
    },
    "outputs": [],
    "source": [

From ab822a0c2ac53cf4db0a0abb716849a2bc5a89b5 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Fri, 16 Aug 2024 16:15:28 +0200
Subject: [PATCH 7/8] ENH: create plots in Sphinx without `ipympl`

---
 docs/lambda-k-pi/ampform-dpd.ipynb | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
index 16eba5e..0192a9b 100644
--- a/docs/lambda-k-pi/ampform-dpd.ipynb
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -55,7 +55,7 @@
     "from ampform_dpd.decay import DecayNode, ThreeBodyDecayChain\n",
     "from ampform_dpd.dynamics.builder import create_mass_symbol, get_mandelstam_s\n",
     "from ampform_dpd.io import aslatex\n",
-    "from IPython.display import Latex, Markdown, Math, display\n",
+    "from IPython.display import SVG, Image, Latex, Markdown, Math, display\n",
     "from qrules.particle import Particle, Spin, create_particle, load_pdg\n",
     "from tensorwaves.data import (\n",
     "    SympyDataTransformer,\n",
@@ -785,7 +785,14 @@
     "interactive_plot = w.interactive_output(update_dalitz_plot, sliders)\n",
     "fig_2d.tight_layout()\n",
     "fig_2d.colorbar(mesh, ax=ax_2d)\n",
-    "display(UI, interactive_plot)"
+    "\n",
+    "if STATIC_PAGE:\n",
+    "    filename = \"dalitz-plot-dpd.png\"\n",
+    "    fig_2d.savefig(filename)\n",
+    "    plt.close(fig_2d)\n",
+    "    display(UI, Image(filename))\n",
+    "else:\n",
+    "    display(UI, interactive_plot)"
    ]
   },
   {
@@ -867,8 +874,13 @@
     "for ax in axes:\n",
     "    ax.legend(fontsize=\"small\")\n",
     "\n",
-    "fig.tight_layout()\n",
-    "display(UI, interactive_plot)"
+    "if STATIC_PAGE:\n",
+    "    filename = \"histogram-dpd.svg\"\n",
+    "    fig.savefig(filename)\n",
+    "    plt.close(fig)\n",
+    "    display(UI, SVG(filename))\n",
+    "else:\n",
+    "    display(UI, interactive_plot)"
    ]
   }
  ],

From d7ed1ddcf026059a4c33336a3ec9a6677f343eac Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Fri, 16 Aug 2024 16:16:02 +0200
Subject: [PATCH 8/8] MAINT: replace `Latex` with `Math`

---
 docs/lambda-k-pi/ampform-dpd.ipynb | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/lambda-k-pi/ampform-dpd.ipynb b/docs/lambda-k-pi/ampform-dpd.ipynb
index 0192a9b..40b351e 100644
--- a/docs/lambda-k-pi/ampform-dpd.ipynb
+++ b/docs/lambda-k-pi/ampform-dpd.ipynb
@@ -55,7 +55,7 @@
     "from ampform_dpd.decay import DecayNode, ThreeBodyDecayChain\n",
     "from ampform_dpd.dynamics.builder import create_mass_symbol, get_mandelstam_s\n",
     "from ampform_dpd.io import aslatex\n",
-    "from IPython.display import SVG, Image, Latex, Markdown, Math, display\n",
+    "from IPython.display import SVG, Image, Markdown, Math, display\n",
     "from qrules.particle import Particle, Spin, create_particle, load_pdg\n",
     "from tensorwaves.data import (\n",
     "    SympyDataTransformer,\n",
@@ -299,7 +299,7 @@
    "outputs": [],
    "source": [
     "decay = to_three_body_decay(reaction.transitions)\n",
-    "Latex(aslatex(decay, with_jp=True))"
+    "Math(aslatex(decay, with_jp=True))"
    ]
   },
   {
@@ -492,7 +492,7 @@
     "i, j = 3, 1\n",
     "(k,) = {1, 2, 3} - {i, j}\n",
     "σk, σk_expr = list(model.invariants.items())[k - 1]\n",
-    "Latex(aslatex({σk: σk_expr}))"
+    "Math(aslatex({σk: σk_expr}))"
    ]
   },
   {