{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell",
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "# WARNING: advised to install a specific version, e.g. tensorwaves==0.1.2\n",
    "%pip install -q tensorwaves[doc,jax,viz] IPython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hideCode": true,
    "hideOutput": true,
    "hidePrompt": true,
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import os\n",
    "\n",
    "from IPython.display import display  # noqa: F401\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analytic continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import warnings\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import qrules\n",
    "from IPython.display import Math\n",
    "\n",
    "from tensorwaves.data import generate_data, generate_phsp\n",
    "from tensorwaves.data.transform import HelicityTransformer\n",
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes {mod}`qrules` finds resonances that lie outside phase space, because resonances can 'leak' into phase space. The example below is simplified and the two selected resonances are a bit unusual, but it serves to illustrate how to handle these sub-threshold resonances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = qrules.generate_transitions(\n",
    "    initial_state=\"D0\",\n",
    "    final_state=[\"K-\", \"K+\", \"K0\"],\n",
    "    allowed_intermediate_particles=[\"a(0)(980)0\", \"a(2)(1320)+\"],\n",
    "    formalism_type=\"canonical-helicity\",\n",
    ")\n",
    "dot = qrules.io.asdot(\n",
    "    result,\n",
    "    collapse_graphs=True,\n",
    "    render_node=False,\n",
    ")\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of the two resonances, $a_0(980)$ lies just below threshold ― it's mass is smaller than the the masses of the two decay products combined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "pdg = qrules.load_pdg()\n",
    "a_meson = pdg[\"a(0)(980)0\"]\n",
    "k_plus = pdg[\"K+\"]\n",
    "k_minus = pdg[\"K+\"]\n",
    "two_k_mass = round(k_minus.mass + k_plus.mass, 3)\n",
    "display(\n",
    "    Math(\n",
    "        Rf\"m_{{{a_meson.latex}}} = {a_meson.mass} \\pm\"\n",
    "        Rf\" {a_meson.width}\\;\\mathrm{{GeV}}\"\n",
    "    ),\n",
    "    Math(\n",
    "        Rf\"m_{{{k_plus.latex}}} + m_{{{k_minus.latex}}} =\"\n",
    "        Rf\" {two_k_mass}\\;\\mathrm{{GeV}}\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To correctly describe the dynamics for this resonance, we should use make use of {doc}`analytic continuation <ampform:usage/dynamics/analytic-continuation>`. As opposed to {doc}`/usage/step1`, we now use {func}`~ampform.dynamics.builder.create_analytic_breit_wigner` (a relativistic Breit-Wigner with analytic continuation) instead of {func}`~ampform.dynamics.builder.create_relativistic_breit_wigner_with_ff`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ampform.dynamics.builder import (\n",
    "    create_analytic_breit_wigner,\n",
    "    create_non_dynamic_with_ff,\n",
    "    create_relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "\n",
    "model_builder = ampform.get_builder(result)\n",
    "model_builder.set_dynamics(\n",
    "    \"J/psi(1S)\",\n",
    "    create_non_dynamic_with_ff,\n",
    ")\n",
    "model_builder.set_dynamics(\n",
    "    \"a(0)(980)0\",\n",
    "    create_analytic_breit_wigner,\n",
    ")\n",
    "model_builder.set_dynamics(\n",
    "    \"a(2)(1320)0\",\n",
    "    create_relativistic_breit_wigner_with_ff,\n",
    ")\n",
    "model = model_builder.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The effect can be seen once we generate data. Despite the fact that the resonance lies outside phase space, there is still a contribution to the intensity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy_model = SympyModel(\n",
    "    expression=model.expression,\n",
    "    parameters=model.parameter_defaults,\n",
    ")\n",
    "intensity = LambdifiedFunction(sympy_model, backend=\"jax\")\n",
    "data_converter = HelicityTransformer(model.adapter)\n",
    "phsp_sample = generate_phsp(100_000, model.adapter.reaction_info)\n",
    "data_sample = generate_data(\n",
    "    2_000, model.adapter.reaction_info, data_converter, intensity\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import cm\n",
    "\n",
    "reaction_info = model.adapter.reaction_info\n",
    "intermediate_states = sorted(\n",
    "    (\n",
    "        p\n",
    "        for p in model.particles\n",
    "        if p not in reaction_info.final_state.values()\n",
    "        and p not in reaction_info.initial_state.values()\n",
    "    ),\n",
    "    key=lambda p: p.mass,\n",
    ")\n",
    "\n",
    "evenly_spaced_interval = np.linspace(0, 1, len(intermediate_states))\n",
    "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n",
    "\n",
    "\n",
    "def indicate_masses():\n",
    "    plt.xlabel(\"$m_{12}$ [GeV]\")\n",
    "    for i, p in enumerate(intermediate_states):\n",
    "        plt.gca().axvline(\n",
    "            x=p.mass, linestyle=\"dotted\", label=p.name, color=colors[i]\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phsp_set = data_converter.transform(phsp_sample)\n",
    "data_set = data_converter.transform(data_sample)\n",
    "data_frame = pd.DataFrame(data_set)\n",
    "data_frame[\"m_12\"].hist(bins=50, alpha=0.5, density=True)\n",
    "indicate_masses()\n",
    "plt.legend();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
