2021.12.13/tensorwaves-optimized-expression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorWaves demo PANDA Seminar December 2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook accompanies [these slides](https://docs.google.com/presentation/d/e/2PACX-1vSymz5AjdhPw4Kz1pKhdFMnFGYuQvVaC8WbV_HTg770x6RDYoP-Anv9tn88DSuzvSiiQ9F4pcDGVExv/pub). They were presented during a PANDA Seminar on 13 December 2021.\n",
    "\n",
    "_This notebook is based on [`tensorwaves.ipynb`](./tensorwaves.ipynb). It demonstrates speed up the fit by simplifying the amplitude model expression with SymPy._\n",
    "\n",
    "Related notebooks for this presentation:\n",
    "- [QRules demo](./qrules.ipynb)\n",
    "- [AmpForm demo](./ampform.ipynb)\n",
    "\n",
    "For more extensive examples, see **[tensorwaves.rtfd.io](https://tensorwaves.readthedocs.io)**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "import logging\n",
    "import re\n",
    "import warnings\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import qrules\n",
    "import sympy as sp\n",
    "import tensorflow as tf\n",
    "from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff\n",
    "from jax.lib import xla_bridge\n",
    "from matplotlib import cm\n",
    "from tensorflow.python.ops.numpy_ops import np_config\n",
    "from tensorwaves.data import (\n",
    "    IntensityDistributionGenerator,\n",
    "    TFPhaseSpaceGenerator,\n",
    "    TFUniformRealNumberGenerator,\n",
    "    TFWeightedPhaseSpaceGenerator,\n",
    ")\n",
    "from tensorwaves.data.transform import SympyDataTransformer\n",
    "from tensorwaves.estimator import UnbinnedNLL\n",
    "from tensorwaves.function.sympy import create_parametrized_function\n",
    "from tensorwaves.optimizer.callbacks import CSVSummary\n",
    "from tensorwaves.optimizer.minuit import Minuit2\n",
    "\n",
    "LOGGER = logging.getLogger(\"absl\")\n",
    "LOGGER.setLevel(logging.ERROR)\n",
    "tf.get_logger().setLevel(\"WARNING\")\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "mpl.rcParams.update({\"font.size\": 14})\n",
    "np_config.enable_numpy_behavior()\n",
    "\n",
    "has_tf_gpu = bool(tf.config.list_physical_devices(\"GPU\"))\n",
    "jax_backend = xla_bridge.get_backend().platform.upper()\n",
    "print(\"TF backend: \", \"GPU\" if has_tf_gpu else \"CPU\")\n",
    "print(\"JAX backend:\", jax_backend)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some helper functions for visualizing the distributions and fit result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "def indicate_masses():\n",
    "    reaction_info = model.reaction_info\n",
    "    resonances = sorted(\n",
    "        reaction_info.get_intermediate_particles(),\n",
    "        key=lambda p: p.mass,\n",
    "    )\n",
    "    evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n",
    "    colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n",
    "    plt.xlabel(\"$m$ [GeV]\")\n",
    "    for i, p in enumerate(resonances):\n",
    "        plt.gca().axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=colors[i])\n",
    "\n",
    "\n",
    "def compare_model(\n",
    "    variable_name,\n",
    "    data_set,\n",
    "    phsp_set,\n",
    "    intensity_model,\n",
    "    bins=100,\n",
    "):\n",
    "    data = np.array(data_set[variable_name])\n",
    "    phsp = np.array(phsp_set[variable_name])\n",
    "    intensities = np.array(intensity_model(phsp_set))\n",
    "    _, ax = plt.subplots(figsize=(10, 3.5))\n",
    "    ax.hist(\n",
    "        data,\n",
    "        bins=bins,\n",
    "        alpha=0.5,\n",
    "        label=\"data\",\n",
    "        density=True,\n",
    "    )\n",
    "    ax.hist(\n",
    "        phsp,\n",
    "        weights=intensities,\n",
    "        bins=bins,\n",
    "        histtype=\"step\",\n",
    "        color=\"red\",\n",
    "        label=\"model\",\n",
    "        density=True,\n",
    "    )\n",
    "    ax.set_yticks([])\n",
    "    indicate_masses()\n",
    "    ax.legend()\n",
    "\n",
    "\n",
    "def natural_sorting(text):\n",
    "    # https://stackoverflow.com/a/5967539/13219025\n",
    "    return [\n",
    "        __attempt_number_cast(c)\n",
    "        for c in re.split(r\"[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)\", text)\n",
    "    ]\n",
    "\n",
    "\n",
    "def __attempt_number_cast(text):\n",
    "    try:\n",
    "        return float(text)\n",
    "    except ValueError:\n",
    "        return text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Formulate amplitude model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate allowed transitions for $J/\\psi \\to \\gamma f_0, f_0 \\to \\pi^0 \\pi^0$ with QRules (see [QRules demo](./qrules.ipynb)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=\"J/psi(1S)\",\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)\"],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Express the transitions as an amplitude model with the resonances parametrized as a relativistic Breit-Wigner with form factor (see [AmpForm demo](./ampform.ipynb)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = ampform.get_builder(reaction)\n",
    "resonances = reaction.get_intermediate_particles()\n",
    "for p in resonances:\n",
    "    builder.set_dynamics(p.name, create_relativistic_breit_wigner_with_ff)\n",
    "model = builder.formulate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a **deterministic** phase-space sample for this decay and an intensity-based hit-and-miss sample for this amplitude model (`intensity`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expression = model.expression.doit()\n",
    "parameter_defaults = model.parameter_defaults\n",
    "intensity = create_parametrized_function(expression, parameter_defaults, backend=\"jax\")\n",
    "helicity_transformer = SympyDataTransformer.from_sympy(\n",
    "    model.kinematic_variables, backend=\"jax\"\n",
    ")\n",
    "reaction_info = model.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",
    "rng = TFUniformRealNumberGenerator(seed=0)\n",
    "\n",
    "phsp_generator = TFPhaseSpaceGenerator(initial_state_mass, final_state_masses)\n",
    "phsp_momenta = phsp_generator.generate(1_000_000, rng)\n",
    "phsp = helicity_transformer(phsp_momenta)\n",
    "\n",
    "data_generator = IntensityDistributionGenerator(\n",
    "    function=intensity,\n",
    "    domain_generator=TFWeightedPhaseSpaceGenerator(\n",
    "        initial_state_mass, final_state_masses\n",
    "    ),\n",
    "    domain_transformer=helicity_transformer,\n",
    ")\n",
    "data_momenta = data_generator.generate(100_000, rng)\n",
    "data = helicity_transformer(data_momenta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.hist(np.array(data[\"m_12\"]), density=True, bins=100, alpha=0.5, label=\"data\")\n",
    "ax.hist(np.array(phsp[\"m_12\"]), density=True, bins=100, alpha=0.5, label=\"phsp\")\n",
    "ax.set_xlabel(R\"$m_{\\pi^0\\pi^0}$\")\n",
    "ax.set_yticks([])\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following, we give a few of the parameters in the amplitude model a little offset from the original values suggested by AmpForm. These parameters will later be optimized with regard to the generated data sample. The model with offset parameters looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_parameters = {\n",
    "    \"m_f(0)(500)\": 0.35,\n",
    "    \"m_f(0)(980)\": 0.88,\n",
    "    \"m_f(0)(1370)\": 1.22,\n",
    "    \"m_f(0)(1500)\": 1.45,\n",
    "    \"m_f(0)(1710)\": 1.83,\n",
    "    \"Gamma_f(0)(500)\": 0.3,\n",
    "    \"Gamma_f(0)(980)\": 0.1,\n",
    "    \"Gamma_f(0)(1710)\": 0.3,\n",
    "}\n",
    "original_parameters = intensity.parameters\n",
    "intensity.update_parameters(initial_parameters)\n",
    "compare_model(\"m_12\", data, phsp, intensity, bins=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having decided which parameters are to be optimized, we can apply some additional optimizations to the amplitude model expression, so that the computations during the fit are faster. The first of these optimizations is to substitute the parameters that remain fixed with their suggested parameter values. Note how this decreases the number of operations (nodes) in the expression tree:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.count_ops(expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fixed_parameters = {\n",
    "    p: v\n",
    "    for p, v in model.parameter_defaults.items()\n",
    "    if p.name not in initial_parameters\n",
    "}\n",
    "optimized_expression = expression.subs(fixed_parameters)\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, we can substitute the final state masses $m_0, m_1, m_2$ with scalar values, as the final states in this decay channel are stable. SymPy will then further simply the expression, so that the computation is less complex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_expression = optimized_expression.subs({\n",
    "    sp.Symbol(f\"m_{i}\", real=True): particle.mass\n",
    "    for i, particle in reaction.final_state.items()\n",
    "})\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there are a few irrational numbers (square roots) in the expression as well, for example in this amplitude:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_expression.args[0].args[0].args[0].args[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we substitute these values, the expression tree becomes even smaller."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for node in sp.preorder_traversal(optimized_expression):\n",
    "    if node.free_symbols:\n",
    "        continue\n",
    "    if isinstance(node, sp.Pow):\n",
    "        optimized_expression = optimized_expression.xreplace({node: node.n()})\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also note that SymPy has rendered $\\sqrt{m_{12}^2}$ in each Breit-Wigner as $|m_{12}|$. This can be substituted as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m_12 = sp.Symbol(\"m_12\", real=True)\n",
    "optimized_expression = optimized_expression.subs({sp.Abs(m_12): m_12})\n",
    "sp.count_ops(optimized_expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, note that some of the computed kinematic variable arrays are complex-valued. This can be useful in come cases, but in this decay channel, all kinematic variable values lie on the real axis. The fit can therefore be further optimized by casting the data arrays to real values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_real = {var: array.real for var, array in data.items()}\n",
    "phsp_real = {var: array.real for var, array in phsp.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Further optimizations, such as identifying [common sub-expressions](https://github.com/ComPWA/tensorwaves/pull/374), are performed when the expression is expressed as a function in the computational backend. The effect of all optimization is illustrated below.\n",
    "\n",
    "_See [`tensorwaves-caching.ipynb`](./tensorwaves-caching.ipynb) for an additional step: caching sub-trees that are independent of free parameters._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slow_function = create_parametrized_function(\n",
    "    expression, parameter_defaults, backend=\"jax\"\n",
    ")\n",
    "optimized_function = create_parametrized_function(\n",
    "    optimized_expression, parameter_defaults, backend=\"jax\"\n",
    ")\n",
    "slow_function({k: v[0] for k, v in data.items()})  # JIT-compile function\n",
    "optimized_function({k: v[0] for k, v in data.items()})  # JIT-compile function\n",
    "%timeit slow_function(data)\n",
    "%timeit optimized_function(data_real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = UnbinnedNLL(optimized_function, data_real, phsp_real, backend=\"jax\")\n",
    "optimizer = Minuit2(callback=CSVSummary(\"fit_traceback.csv\"))\n",
    "fit_result = optimizer.optimize(estimator, initial_parameters)\n",
    "fit_result.execution_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Time estimate:_<br>\n",
    "For this (deterministic) data sample and these initial parameter values, the fit requires **515 iterations** (GPU; 516 on CPU). On Google Colab, this should take **around 15 seconds** with GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "optimized_function.update_parameters(fit_result.parameter_values)\n",
    "compare_model(\"m_12\", data, phsp, optimized_function, bins=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize fit traceback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "traceback = np.genfromtxt(\"fit_traceback.csv\", delimiter=\",\", names=True)\n",
    "fig, (ax1, ax2) = plt.subplots(\n",
    "    nrows=2,\n",
    "    figsize=(7, 8),\n",
    "    gridspec_kw={\"height_ratios\": [1, 2.5]},\n",
    "    sharex=True,\n",
    ")\n",
    "ax1.set_title(\"Negative log likelihood\")\n",
    "ax2.set_title(\"Parameter values\")\n",
    "ax2.set_xlabel(\"function call\")\n",
    "fig.tight_layout()\n",
    "\n",
    "x = traceback[\"function_call\"]\n",
    "ax1.plot(x, traceback[\"estimator_value\"])\n",
    "for par in initial_parameters:\n",
    "    label = f\"${sp.latex(sp.Symbol(par))}$\"\n",
    "    key = re.sub(r\"[\\(\\)]\", \"\", par)\n",
    "    ax2.plot(x, traceback[key], label=label)\n",
    "\n",
    "legend = ax2.legend(loc=(0.77, 0.43))\n",
    "legend.get_frame().set_alpha(None)\n",
    "for par, line in zip(initial_parameters, ax2.get_lines(), strict=True):\n",
    "    label = line.get_label()\n",
    "    color = line.get_color()\n",
    "    ax2.axhline(\n",
    "        y=original_parameters[par],\n",
    "        color=color,\n",
    "        alpha=0.5,\n",
    "        linestyle=\"dotted\",\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Pyproject Local",
   "language": "python",
   "name": "pyproject_local_kernel"
  },
  "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}