{
 "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"
    ]
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%config Completer.use_jedi = False\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import os\n",
    "\n",
    "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)\n",
    "\n",
    "# Install on Google Colab\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "install_packages = \"google.colab\" in str(get_ipython())\n",
    "if install_packages:\n",
    "    for package in [\"tensorwaves[doc]\", \"graphviz\"]:\n",
    "        subprocess.check_call(\n",
    "            [sys.executable, \"-m\", \"pip\", \"install\", package]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TensorWaves is a package for fitting general mathematical expressions to data distributions. The fundamentals behind the package are illustrated in {doc}`/usage/basics`.\n",
    "\n",
    "While general in design, the package is intended for doing {doc}`Partial Wave Analysis <pwa:index>`. First, the {mod}`ampform` package determines which transitions are allowed from some initial state to a final state. It then formulates those transitions mathematically as an {ref}`amplitude model <usage:Construct a model>`. TensorWaves can then {meth}`~.Model.lambdify` this expression to some computational backend. Finally, TensorWaves {ref}`'fits' <usage:Optimize the model>` this model to some data sample. Optionally, a data sample can be {ref}`generated <usage:Generate data sample>` from the model.\n",
    "\n",
    "This page shows a brief overview of the complete workflow. More info about each step can be found under {ref}`usage:Step-by-step workflow`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "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",
    "import sympy as sp\n",
    "from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff\n",
    "\n",
    "from tensorwaves.data import generate_data, generate_phsp\n",
    "from tensorwaves.data.transform import HelicityTransformer\n",
    "from tensorwaves.estimator import UnbinnedNLL\n",
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "from tensorwaves.optimizer.callbacks import CSVSummary\n",
    "from tensorwaves.optimizer.minuit import Minuit2\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{doc}`usage/step1`\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "\n",
    "This example fits a simple model. See {doc}`/usage/step1` for a more complicated model.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [-1, +1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\n",
    "        \"f(0)(980)\",\n",
    "        \"f(0)(1500)\",\n",
    "        \"f(0)(1710)\",\n",
    "    ],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    "    formalism=\"canonical-helicity\",\n",
    ")\n",
    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_builder = ampform.get_builder(reaction)\n",
    "for name in reaction.get_intermediate_particles().names:\n",
    "    model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n",
    "model = model_builder.formulate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "model.components[\n",
    "    R\"A_{J/\\psi(1S)_{-1} \\xrightarrow[S=1]{L=0} f_{0}(980)_{0} \\gamma_{-1};\"\n",
    "    R\" f_{0}(980)_{0} \\xrightarrow[S=0]{L=0} \\pi^{0}_{0} \\pi^{0}_{0}}\"\n",
    "].doit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate data sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{doc}`usage/step2`\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy_model = SympyModel(\n",
    "    expression=model.expression.doit(),\n",
    "    parameters=model.parameter_defaults,\n",
    ")\n",
    "intensity = LambdifiedFunction(sympy_model, backend=\"jax\")\n",
    "data_converter = HelicityTransformer(model.adapter)\n",
    "reaction_info = model.adapter.reaction_info\n",
    "initial_state_mass = reaction_info.initial_state[-1].mass\n",
    "final_state_masses = {i: p.mass for i, p in reaction_info.final_state.items()}\n",
    "phsp_sample = generate_phsp(100_000, initial_state_mass, final_state_masses)\n",
    "data_sample = generate_data(\n",
    "    10_000,\n",
    "    initial_state_mass,\n",
    "    final_state_masses,\n",
    "    data_converter,\n",
    "    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$ [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": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "phsp_set = data_converter.transform(phsp_sample)\n",
    "data_set = data_converter.transform(data_sample)\n",
    "data_frame = pd.DataFrame(data_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9, 4))\n",
    "data_frame[\"m_12\"].hist(bins=100, alpha=0.5, density=True, ax=ax)\n",
    "indicate_masses()\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimize the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "\n",
    "{doc}`usage/step3`\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def compare_model(\n",
    "    variable_name,\n",
    "    data_set,\n",
    "    phsp_set,\n",
    "    intensity_model,\n",
    "    bins=150,\n",
    "):\n",
    "    data = data_set[variable_name]\n",
    "    phsp = phsp_set[variable_name]\n",
    "    intensities = intensity_model(phsp_set)\n",
    "    fig, ax = plt.subplots(figsize=(9, 4))\n",
    "    plt.hist(\n",
    "        data,\n",
    "        bins=bins,\n",
    "        alpha=0.5,\n",
    "        label=\"data\",\n",
    "        density=True,\n",
    "    )\n",
    "    plt.hist(\n",
    "        phsp,\n",
    "        weights=intensities,\n",
    "        bins=bins,\n",
    "        histtype=\"step\",\n",
    "        color=\"red\",\n",
    "        label=\"initial fit model\",\n",
    "        density=True,\n",
    "    )\n",
    "    indicate_masses()\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = UnbinnedNLL(\n",
    "    intensity,\n",
    "    data_set,\n",
    "    phsp_set,\n",
    "    backend=\"jax\",\n",
    ")\n",
    "initial_parameters = {\n",
    "    \"m_f(0)(980)\": 0.93,\n",
    "    \"m_f(0)(1500)\": 1.45,\n",
    "    \"m_f(0)(1710)\": 1.8,\n",
    "    \"Gamma_f(0)(980)\": 0.1,\n",
    "    \"Gamma_f(0)(1710)\": 0.2,\n",
    "}\n",
    "intensity.update_parameters(initial_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "compare_model(\"m_12\", data_set, phsp_set, intensity)\n",
    "print(\"Number of free parameters:\", len(initial_parameters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "callback = CSVSummary(\"fit_traceback.csv\")\n",
    "minuit2 = Minuit2(callback)\n",
    "fit_result = minuit2.optimize(estimator, initial_parameters)\n",
    "fit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert fit_result.minimum_valid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "optimized_parameters = fit_result.parameter_values\n",
    "intensity.update_parameters(optimized_parameters)\n",
    "compare_model(\"m_12\", data_set, phsp_set, intensity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "full-width"
    ]
   },
   "outputs": [],
   "source": [
    "converters = {p: lambda s: complex(s).real for p in initial_parameters}\n",
    "fit_traceback = pd.read_csv(\"fit_traceback.csv\", converters=converters)\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    2, figsize=(7, 8), gridspec_kw={\"height_ratios\": [1, 1.8]}\n",
    ")\n",
    "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n",
    "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n",
    "ax1.set_title(\"Negative log likelihood\")\n",
    "ax2.set_title(\"Parameter values\")\n",
    "ax1.set_xlabel(\"function call\")\n",
    "ax2.set_xlabel(\"function call\")\n",
    "fig.tight_layout()\n",
    "ax1.legend().remove()\n",
    "legend_texts = ax2.legend().get_texts()\n",
    "for text in legend_texts:\n",
    "    latex = f\"${sp.latex(sp.Symbol(text.get_text()))}$\"\n",
    "    latex = latex.replace(\"\\\\\\\\\", \"\\\\\")\n",
    "    if latex[2] == \"C\":\n",
    "        latex = fR\"\\left|{latex}\\right|\"\n",
    "    text.set_text(latex)\n",
    "for line in ax2.get_lines():\n",
    "    label = line.get_label()\n",
    "    color = line.get_color()\n",
    "    ax2.axhline(\n",
    "        y=complex(sympy_model.parameters[label]).real,\n",
    "        color=color,\n",
    "        alpha=0.5,\n",
    "        linestyle=\"dotted\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-by-step workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following pages go through the usual workflow when using {mod}`tensorwaves`. The output in each of these steps is written to disk, so that each notebook can be run separately."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{toctree}\n",
    "---\n",
    "maxdepth: 2\n",
    "---\n",
    "usage/step1\n",
    "usage/step2\n",
    "usage/step3\n",
    "usage/basics\n",
    "usage/analytic-continuation\n",
    "usage/faster-lambdify\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}