docs/jpsi2ksp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# J/ψ → K⁰ Σ⁺ p̅\n",
    "\n",
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import logging\n",
    "import os\n",
    "import warnings\n",
    "from typing import TYPE_CHECKING\n",
    "\n",
    "import graphviz\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import qrules\n",
    "import sympy as sp\n",
    "from ampform.dynamics import BlattWeisskopfSquared, EnergyDependentWidth\n",
    "from ampform.sympy import perform_cached_doit\n",
    "from IPython.display import Latex, Markdown\n",
    "from tensorwaves.data.transform import SympyDataTransformer\n",
    "from tqdm.auto import tqdm\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 State\n",
    "from ampform_dpd.dynamics import FormFactor, RelativisticBreitWigner\n",
    "from ampform_dpd.dynamics.builder import formulate_breit_wigner_with_form_factor\n",
    "from ampform_dpd.io import (\n",
    "    as_markdown_table,\n",
    "    aslatex,\n",
    "    perform_cached_lambdify,\n",
    "    simplify_latex_rendering,\n",
    ")\n",
    "\n",
    "simplify_latex_rendering()\n",
    "logging.getLogger(\"absl\").setLevel(logging.ERROR)  # mute JAX\n",
    "warnings.simplefilter(\"ignore\", category=RuntimeWarning)\n",
    "\n",
    "NO_TQDM = \"EXECUTE_NB\" in os.environ\n",
    "if NO_TQDM:\n",
    "    logging.getLogger(\"ampform_dpd.io\").setLevel(logging.ERROR)\n",
    "if TYPE_CHECKING:\n",
    "    from tensorwaves.interface import DataSample, ParametrizedFunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Decay definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We follow [this example](https://qrules.readthedocs.io/en/0.10.1/usage.html#investigate-intermediate-resonances), which was generated with QRules, and leave out the $K$-resonances and the resonances that lie far outside of phase space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "REACTION = qrules.generate_transitions(\n",
    "    initial_state=\"J/psi(1S)\",\n",
    "    final_state=[\"K0\", \"Sigma+\", \"p~\"],\n",
    "    allowed_interaction_types=\"strong\",\n",
    "    formalism=\"canonical-helicity\",\n",
    "    mass_conservation_factor=0.05,\n",
    ")\n",
    "REACTION = normalize_state_ids(REACTION)\n",
    "dot = qrules.io.asdot(REACTION, 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(REACTION.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": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(DECAY, with_jp=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Model formulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total, aligned intensity expression looks as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=False)\n",
    "for chain in model_builder.decay.chains:\n",
    "    model_builder.dynamics_choices.register_builder(\n",
    "        chain, formulate_breit_wigner_with_form_factor\n",
    "    )\n",
    "model = model_builder.formulate(reference_subsystem=2)\n",
    "model.intensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(model.variables))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each **unaligned** amplitude is defined as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input",
     "full-width",
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "Latex(aslatex(model.amplitudes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\")\n",
    "exprs = [\n",
    "    RelativisticBreitWigner(s, m0, w0, m1, m2, L, R),\n",
    "    EnergyDependentWidth(s, m0, w0, m1, m2, L, R),\n",
    "    FormFactor(s, m1, m2, L, R),\n",
    "    BlattWeisskopfSquared(z, L),\n",
    "]\n",
    "Latex(aslatex({e: e.doit(deep=False) for e in exprs}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Preparing for input data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "i, j = (3, 2)\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
    },
    "mystnb": {
     "code_prompt_show": "Define meshgrid for Dalitz plot"
    },
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "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": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Create data converter for Dalitz coordinates"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "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",
    "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": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "for key, array in dalitz_data.items():\n",
    "    assert not jnp.all(jnp.isnan(array)), f\"All values for {key} are NaN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Prepare parametrized numerical function"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "unfolded_intensity_expr = perform_cached_doit(model.intensity)\n",
    "full_intensity_expr = perform_cached_doit(\n",
    "    unfolded_intensity_expr.xreplace(model.amplitudes)\n",
    ")\n",
    "free_parameters = {\n",
    "    k: v\n",
    "    for k, v in model.parameter_defaults.items()\n",
    "    if isinstance(k, sp.Indexed)\n",
    "    if \"production\" in str(k) or \"decay\" in str(k)\n",
    "}\n",
    "fixed_parameters = {\n",
    "    k: v for k, v in model.parameter_defaults.items() if k not in free_parameters\n",
    "}\n",
    "intensity_func = perform_cached_lambdify(\n",
    "    full_intensity_expr.xreplace(fixed_parameters),\n",
    "    parameters=free_parameters,\n",
    "    backend=\"jax\",\n",
    ")\n",
    "intensities = intensity_func(dalitz_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert not jnp.all(jnp.isnan(intensities)), \"All intensities are NaN\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Dalitz plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['png']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def get_decay_products(subsystem_id: int) -> tuple[State, State]:\n",
    "    return tuple(s for s in DECAY.final_state.values() if s.index != subsystem_id)\n",
    "\n",
    "\n",
    "plt.rc(\"font\", size=18)\n",
    "I_tot = jnp.nansum(intensities)\n",
    "normalized_intensities = intensities / I_tot\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(14, 10))\n",
    "mesh = ax.pcolormesh(X, Y, normalized_intensities)\n",
    "ax.set_aspect(\"equal\")\n",
    "c_bar = plt.colorbar(mesh, ax=ax, pad=0.01)\n",
    "c_bar.ax.set_ylabel(\"Normalized intensity (a.u.)\")\n",
    "sigma_labels = {\n",
    "    i: Rf\"$\\sigma_{i} = M^2\\left({' '.join(p.latex for p in get_decay_products(i))}\\right)$\"\n",
    "    for i in (1, 2, 3)\n",
    "}\n",
    "ax.set_xlabel(sigma_labels[i])\n",
    "ax.set_ylabel(sigma_labels[j])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width",
     "scroll-input"
    ]
   },
   "outputs": [],
   "source": [
    "def compute_sub_intensity(\n",
    "    func: ParametrizedFunction, phsp: DataSample, resonance_latex: str\n",
    ") -> jnp.ndarray:\n",
    "    original_parameters = dict(func.parameters)\n",
    "    zero_parameters = {\n",
    "        k: 0\n",
    "        for k, v in func.parameters.items()\n",
    "        if R\"\\mathcal{H}\" in k\n",
    "        if resonance_latex not in k\n",
    "    }\n",
    "    func.update_parameters(zero_parameters)\n",
    "    intensities = func(phsp)\n",
    "    func.update_parameters(original_parameters)\n",
    "    return intensities\n",
    "\n",
    "\n",
    "plt.rc(\"font\", size=16)\n",
    "fig, axes = plt.subplots(figsize=(18, 6), ncols=2, sharey=True)\n",
    "fig.subplots_adjust(wspace=0.02)\n",
    "ax1, ax2 = axes\n",
    "x = jnp.sqrt(X[0])\n",
    "y = jnp.sqrt(Y[:, 0])\n",
    "ax1.fill_between(x, jnp.nansum(normalized_intensities, axis=0), alpha=0.5)\n",
    "ax2.fill_between(y, jnp.nansum(normalized_intensities, axis=1), alpha=0.5)\n",
    "for ax in axes:\n",
    "    _, y_max = ax.get_ylim()\n",
    "    ax.set_ylim(0, y_max)\n",
    "    ax.autoscale(enable=False, axis=\"x\")\n",
    "ax1.set_ylabel(\"Normalized intensity (a.u.)\")\n",
    "ax1.set_xlabel(sigma_labels[i])\n",
    "ax2.set_xlabel(sigma_labels[j])\n",
    "resonance_counter1, resonance_counter2 = 0, 0\n",
    "for chain in tqdm(model.decay.chains, disable=NO_TQDM):\n",
    "    resonance = chain.resonance\n",
    "    if set(chain.decay_products) == set(get_decay_products(subsystem_id=i)):\n",
    "        ax = ax1\n",
    "        resonance_counter1 += 1\n",
    "        color = f\"C{resonance_counter1}\"\n",
    "        projection_axis = 0\n",
    "        x_data = x\n",
    "    elif set(chain.decay_products) == set(get_decay_products(subsystem_id=j)):\n",
    "        ax = ax2\n",
    "        resonance_counter2 += 1\n",
    "        color = f\"C{resonance_counter2}\"\n",
    "        projection_axis = 1\n",
    "        x_data = y\n",
    "    else:\n",
    "        raise NotImplementedError\n",
    "    sub_intensities = compute_sub_intensity(\n",
    "        intensity_func, dalitz_data, resonance.latex\n",
    "    )\n",
    "    ax.plot(x_data, jnp.nansum(sub_intensities / I_tot, axis=projection_axis), c=color)\n",
    "    ax.axvline(resonance.mass, label=f\"${resonance.latex}$\", c=color, ls=\"dashed\")\n",
    "for ax in axes:\n",
    "    ax.legend(fontsize=12)\n",
    "plt.show()"
   ]
  }
 ],
 "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
}