{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Phase space sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "mystnb": {
     "code_prompt_show": "Import Python libraries"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import sympy as sp\n",
    "from ampform.kinematics.phasespace import (\n",
    "    Kallen,\n",
    "    Kibble,\n",
    "    compute_third_mandelstam,\n",
    "    is_within_phasespace,\n",
    ")\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "from polarimetry.data import (\n",
    "    create_mass_symbol_mapping,\n",
    "    create_phase_space_filter,\n",
    "    generate_meshgrid_sample,\n",
    "    generate_phasespace_sample,\n",
    ")\n",
    "from polarimetry.io import display_doit, display_latex, mute_jax_warnings\n",
    "from polarimetry.lhcb import load_model_builder\n",
    "from polarimetry.lhcb.particle import load_particles\n",
    "from polarimetry.plot import use_mpl_latex_fonts\n",
    "\n",
    "mute_jax_warnings()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "AmpForm's {doc}`ampform:usage/kinematics` page.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "m0, mi, mj, mk = sp.symbols(\"m0 m_(i:k)\", nonnegative=True)\n",
    "σi, σj, σk = sp.symbols(\"sigma_(i:k)\", nonnegative=True)\n",
    "is_within_phasespace(σi, σj, m0, mi, mj, mk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "display_doit(Kibble(σi, σj, σk, m0, mi, mj, mk))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "display_doit(Kallen(*sp.symbols(\"x:z\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "m1, m2, m3 = sp.symbols(\"m1:4\")\n",
    "display_latex({σk: compute_third_mandelstam(σi, σj, m0, m1, m2, m3)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "model_choice = 0\n",
    "model_file = \"../../data/model-definitions.yaml\"\n",
    "particles = load_particles(\"../../data/particle-definitions.yaml\")\n",
    "amplitude_builder = load_model_builder(model_file, particles, model_choice)\n",
    "decay = amplitude_builder.decay\n",
    "display_latex(create_mass_symbol_mapping(decay))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def plot_phsp_boundary(ax, x_mandelstam: int, y_mandelstam: int):\n",
    "    ax.set_xticks([])\n",
    "    ax.set_xlabel(Rf\"$\\sigma_{x_mandelstam}$ [GeV$^2$]\")\n",
    "    ax.set_ylabel(Rf\"$\\sigma_{y_mandelstam}$ [GeV$^2$]\")\n",
    "    phsp = generate_meshgrid_sample(decay, resolution, x_mandelstam, y_mandelstam)\n",
    "    phsp_filter = create_phase_space_filter(\n",
    "        decay, x_mandelstam, y_mandelstam, outside_value=0\n",
    "    )\n",
    "    mesh = ax.contour(\n",
    "        phsp[f\"sigma{x_mandelstam}\"],\n",
    "        phsp[f\"sigma{y_mandelstam}\"],\n",
    "        phsp_filter(phsp),\n",
    "        colors=\"black\",\n",
    "    )\n",
    "    contour = mesh.collections[0]\n",
    "    contour.set_facecolor(\"lightgray\")\n",
    "\n",
    "\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "resolution = 500\n",
    "use_mpl_latex_fonts()\n",
    "plt.rc(\"font\", size=18)\n",
    "fig, ax = plt.subplots(figsize=(5, 5), tight_layout=True)\n",
    "plot_phsp_boundary(ax, 1, 2)\n",
    "fig.savefig(\"../_images/phase-space-boundary.svg\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{only} latex\n",
    "![](../_images/phase-space-boundary.svg)\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "def plot_phsp_distribution(ax, x_mandelstam: int, y_mandelstam: int):\n",
    "    ax.set_xticks([])\n",
    "    ax.set_xlabel(Rf\"$\\sigma_{x_mandelstam}$\")\n",
    "    ax.set_ylabel(Rf\"$\\sigma_{y_mandelstam}$\")\n",
    "    ax.hist2d(\n",
    "        phsp[f\"sigma{x_mandelstam}\"],\n",
    "        phsp[f\"sigma{y_mandelstam}\"],\n",
    "        bins=500,\n",
    "        norm=LogNorm(),\n",
    "    )\n",
    "\n",
    "\n",
    "%config InlineBackend.figure_formats = ['png']\n",
    "phsp = generate_phasespace_sample(decay, n_events=10_000_000, seed=0)\n",
    "fig, axes = plt.subplots(ncols=3, dpi=200, figsize=(15, 4.9), tight_layout=True)\n",
    "plot_phsp_distribution(axes[0], 1, 2)\n",
    "plot_phsp_distribution(axes[1], 2, 3)\n",
    "plot_phsp_distribution(axes[2], 3, 1)\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "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.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}