{
 "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": {},
   "source": [
    "# Speed up lambdifying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import qrules\n",
    "import sympy as sp\n",
    "from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff\n",
    "from IPython.display import HTML, SVG\n",
    "\n",
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "from tensorwaves.model.sympy import optimized_lambdify, split_expression\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lambdifying a SymPy expression can take rather long when an expression is complicated. Fortunately, TensorWaves offers a way to speed up the lambdify process. The idea is to split up an an expression into sub-expressions, separate those separately, and then recombining them. Let's illustrate that idea with the following simplified example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, z = sp.symbols(\"x:z\")\n",
    "expr = x ** z + 2 * y\n",
    "expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This expression can be represented in a tree of mathematical operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(expr)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function {func}`.split_expression` can now be used to split up this expression tree into a 'top expression' plus definitions for each of the sub-expressions into which it was split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_expr, sub_expressions = split_expression(expr, max_complexity=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_expressions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original expression can easily be reconstructed with {meth}`~sympy.core.basic.Basic.subs` or {meth}`~sympy.core.basic.Basic.xreplace`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_expr.xreplace(sub_expressions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of the expression trees are now smaller than the original:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(top_expr)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "for symbol, definition in sub_expressions.items():\n",
    "    dot = sp.dotprint(definition)\n",
    "    graph = graphviz.Source(dot)\n",
    "    graph.render(filename=f\"sub_expr_{symbol.name}\", format=\"svg\")\n",
    "\n",
    "html = \"<table>\\n\"\n",
    "html += \"  <tr>\\n\"\n",
    "html += \"\".join(\n",
    "    '    <th style=\"text-align:center;'\n",
    "    f' background-color:white\">{symbol.name}</th>\\n'\n",
    "    for symbol in sub_expressions\n",
    ")\n",
    "html += \"  </tr>\\n\"\n",
    "html += \"  <tr>\\n\"\n",
    "for symbol in sub_expressions:\n",
    "    svg = SVG(f\"sub_expr_{symbol.name}.svg\").data\n",
    "    html += f'    <td style=\"background-color:white\">{svg}</td>\\n'\n",
    "html += \"  </tr>\\n\"\n",
    "html += \"</table>\"\n",
    "HTML(html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimized lambdify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generally, the lambdify time scales exponentially with the size of an expression tree. With larger expression trees, it's therefore much faster to lambdify these sub-expressions separately and to recombine them. TensorWaves offers a function that does this for you: {func}`.optimized_lambdify`. We'll use an {class}`~ampform.helicity.HelicityModel` to illustrate this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaction = qrules.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [+1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)\"],\n",
    ")\n",
    "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": {},
   "outputs": [],
   "source": [
    "expression = model.expression.doit()\n",
    "sorted_symbols = sorted(expression.free_symbols, key=lambda s: s.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "lambdified_optimized = optimized_lambdify(\n",
    "    expression,\n",
    "    sorted_symbols,\n",
    "    max_complexity=100,\n",
    "    backend=\"numpy\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "lambdified_optimized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "sp.lambdify(sorted_symbols, expression)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying complexity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the usually workflow (see {doc}`/usage`), TensorWaves uses SymPy's own {func}`~sympy.utilities.lambdify.lambdify` by default. You can change this behavior with the `max_complexity` argument of {class}`.SympyModel`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy_model = SympyModel(\n",
    "    expression=model.expression.doit(),\n",
    "    parameters=model.parameter_defaults,\n",
    "    max_complexity=100,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `max_complexity` is specified (i.e., is not {obj}`None`), {class}`.LambdifiedFunction` uses TensorWaves's {func}`.optimized_lambdify`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "intensity = LambdifiedFunction(sympy_model, backend=\"jax\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}