{ "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": { "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 `. 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 `. TensorWaves can then {meth}`~.Model.lambdify` this expression to some computational backend. Finally, TensorWaves {ref}`'fits' ` this model to some data sample. Optionally, a data sample can be {ref}`generated ` 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 as q\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": [ "result = q.generate_transitions(\n", " initial_state=(\"J/psi(1S)\", [-1, +1]),\n", " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", " allowed_intermediate_particles=[\"f(0)(980)\", \"f(0)(1500)\", \"f(0)(1710)\"],\n", " allowed_interaction_types=[\"strong\", \"EM\"],\n", " formalism_type=\"canonical-helicity\",\n", ")\n", "dot = q.io.asdot(result, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model_builder = ampform.get_builder(result)\n", "for name in result.get_intermediate_particles().names:\n", " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", "model = model_builder.generate()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model.components[\n", " R\"A[J/\\psi(1S)_{-1} \\to f_{0}(980)_{0} \\gamma_{+1,L=2,S=1}; f_{0}(980)_{0}\"\n", " R\" \\to \\pi^{0}_{0} \\pi^{0}_{0,L=0,S=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,\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", " 10_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$ [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)\n", "data_frame[\"m_12\"].hist(bins=100, alpha=0.5, density=True)\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", " plt.hist(data, bins=bins, alpha=0.5, label=\"data\", density=True)\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": {}, "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": {}, "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" ] }, "outputs": [], "source": [ "import pandas as pd\n", "import sympy as sp\n", "\n", "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 = Rf\"\\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", "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 }