docs/report/002.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "lambdification",
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Faster lambdification by splitting expressions\n",
    "TR-002\n",
    "^^^\n",
    "This notebook investigates how to speed up {func}`sympy.lambdify <sympy.utilities.lambdify.lambdify>` by splitting up the expression tree of a complicated expression into components, lambdifying those, and then combining them back again.\n",
    "+++\n",
    "✅&nbsp;[tensorwaves#281](https://github.com/ComPWA/tensorwaves/issues/281)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speed up lambdifying"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- cspell:disable -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q ampform==0.9.1 numpy==1.19.5 pandas==1.2.4 qrules==0.8.2 sympy==1.8 tensorwaves[jax]==0.2.9 git+https://github.com/zfit/phasespace@7131fbd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "from __future__ import annotations\n",
    "\n",
    "import inspect\n",
    "import logging\n",
    "import timeit\n",
    "import warnings\n",
    "from typing import Callable, Generator, Sequence\n",
    "\n",
    "import ampform\n",
    "import graphviz\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\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",
    "from tensorwaves.data import generate_phsp\n",
    "from tensorwaves.data.phasespace import TFUniformRealNumberGenerator\n",
    "from tensorwaves.data.transform import HelicityTransformer\n",
    "from tensorwaves.model import LambdifiedFunction, SympyModel\n",
    "\n",
    "LOGGER = logging.getLogger()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create dummy expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's create an amplitude model with {mod}`ampform`. We'll use this model as complicated {class}`sympy.Expr <sympy.core.expr.Expr>` in the rest of this notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "LOGGER.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "result = qrules.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [-1, +1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)(980)\"],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    "    formalism_type=\"canonical-helicity\",\n",
    ")\n",
    "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "graphviz.Source(dot).render(\"002-f0(980)-graph\", format=\"svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983331-6eb948fe-d360-40bd-a4f7-fa1aad3e296a.svg)"
   ]
  },
  {
   "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": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}],\n",
       " Gamma_f(0)(980),\n",
       " d_f(0)(980),\n",
       " m_1,\n",
       " m_12,\n",
       " m_2,\n",
       " m_f(0)(980),\n",
       " phi_1+2,\n",
       " phi_1,1+2,\n",
       " theta_1+2,\n",
       " theta_1,1+2]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free_symbols = sorted(model.expression.free_symbols, key=lambda s: s.name)\n",
    "free_symbols"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helicity model components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A {class}`~ampform.helicity.HelicityModel` has the benefit that it comes with {attr}`~ampform.helicity.HelicityModel.components` (intensities and amplitudes) that together form its {attr}`~ampform.helicity.HelicityModel.expression`. Let's separate these components into _amplitude_ and _intensity_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output",
     "full-width"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n",
       " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amplitudes = {\n",
    "    name: expr for name, expr in model.components.items() if name.startswith(\"A\")\n",
    "}\n",
    "sorted(amplitudes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensities = {\n",
    "    name: expr for name, expr in model.components.items() if name.startswith(\"I\")\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(amplitudes) + len(intensities) == len(model.components)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Component structure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each intensity consists of a subset of these amplitudes. This means that _intensities have a larger expression tree than amplitudes_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "amplitude_to_symbol = {\n",
    "    expr: sp.Symbol(f\"A{i}\") for i, expr in enumerate(amplitudes.values(), 1)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "intensity_to_symbol = {\n",
    "    expr: sp.Symbol(f\"I{i}\") for i, expr in enumerate(intensities.values(), 1)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle I_{1} + I_{2} + I_{3} + I_{4}$"
      ],
      "text/plain": [
       "I1 + I2 + I3 + I4"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "intensity_expr = model.expression.subs(intensity_to_symbol, simultaneous=True)\n",
    "intensity_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(intensity_expr)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "graphviz.Source(dot).render(\"002-collapsed-expression-tree\", format=\"svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983184-fde89791-2e75-4bd1-9c03-9a45edf24216.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left|{A_{1} + A_{2}}\\right|^{2} + \\left|{A_{3} + A_{4}}\\right|^{2} + \\left|{A_{5} + A_{6}}\\right|^{2} + \\left|{A_{7} + A_{8}}\\right|^{2}$"
      ],
      "text/plain": [
       "Abs(A1 + A2)**2 + Abs(A3 + A4)**2 + Abs(A5 + A6)**2 + Abs(A7 + A8)**2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amplitude_expr = model.expression.subs(amplitude_to_symbol, simultaneous=True)\n",
    "amplitude_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "dot = sp.dotprint(amplitude_expr)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "graphviz.Source(dot).render(\"002-simple-expression-tree\", format=\"svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983978-73d1b6a4-0f09-4a10-88d8-de9d6a055bf3.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance check"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lambdifying the whole {attr}`HelicityModel.expression <ampform.helicity.HelicityModel.expression>` is slowest. The {func}`~sympy.utilities.lambdify.lambdify` function first prints the expression as a {obj}`str` (!) with (in this case) {mod}`numpy` syntax and then uses {func}`eval` to convert that back to actual {mod}`numpy` objects:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "\n",
    "We store the time with {mod}`timeit` for section {ref}`report/002:Arbitrary expressions`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "runtime = {}\n",
    "start = timeit.default_timer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.46 s, sys: 703 µs, total: 1.46 s\n",
      "Wall time: 1.46 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "np_complete_model = sp.lambdify(free_symbols, model.expression.doit(), \"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "stop = timeit.default_timer()\n",
    "runtime[\"complete model\"] = stop - start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Printing to {obj}`str` and converting back with {func}`eval` becomes exponentially slow the larger the expression tree. This means that it's more efficient to lambdify sub-trees of the expression tree separately. Lambdifying the four _intensities_ of this model separately, the effect is not noticeable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "LOGGER.setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Lambdifying I1\n",
      "INFO:root:Lambdifying I2\n",
      "INFO:root:Lambdifying I3\n",
      "INFO:root:Lambdifying I4\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.56 s, sys: 4.94 ms, total: 1.56 s\n",
      "Wall time: 1.56 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "for expr, symbol in intensity_to_symbol.items():\n",
    "    logging.info(f\"Lambdifying {symbol.name}\")\n",
    "    start = timeit.default_timer()\n",
    "    sp.lambdify(free_symbols, expr.doit(), \"numpy\")\n",
    "    stop = timeit.default_timer()\n",
    "    runtime[symbol.name] = stop - start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...but each of the eight _amplitudes_ separately does result in a significant speed-up:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Lambdifying A1\n",
      "INFO:root:Lambdifying A2\n",
      "INFO:root:Lambdifying A3\n",
      "INFO:root:Lambdifying A4\n",
      "INFO:root:Lambdifying A5\n",
      "INFO:root:Lambdifying A6\n",
      "INFO:root:Lambdifying A7\n",
      "INFO:root:Lambdifying A8\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 547 ms, sys: 3.85 ms, total: 550 ms\n",
      "Wall time: 547 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "np_amplitudes = {}\n",
    "for expr, symbol in amplitude_to_symbol.items():\n",
    "    logging.info(f\"Lambdifying {symbol.name}\")\n",
    "    start = timeit.default_timer()\n",
    "    np_expr = sp.lambdify(free_symbols, expr.doit(), \"numpy\")\n",
    "    stop = timeit.default_timer()\n",
    "    runtime[symbol.name] = stop - start\n",
    "    np_amplitudes[symbol] = np_expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recombining components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{ref}`Recall <report/002:Component structure>` what amplitude module expressed in its amplitude components looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left|{A_{1} + A_{2}}\\right|^{2} + \\left|{A_{3} + A_{4}}\\right|^{2} + \\left|{A_{5} + A_{6}}\\right|^{2} + \\left|{A_{7} + A_{8}}\\right|^{2}$"
      ],
      "text/plain": [
       "Abs(A1 + A2)**2 + Abs(A3 + A4)**2 + Abs(A5 + A6)**2 + Abs(A7 + A8)**2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "amplitude_expr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to lambdify that top expression as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_amplitude_symbols = sorted(np_amplitudes, key=lambda s: s.name)\n",
    "np_amplitude_expr = sp.lambdify(sorted_amplitude_symbols, amplitude_expr, \"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(A1, A2, A3, A4, A5, A6, A7, A8):\n",
      "    return (abs(A1 + A2)**2 + abs(A3 + A4)**2 + abs(A5 + A6)**2 + abs(A7 + A8)**2)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "source = inspect.getsource(np_amplitude_expr)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have a lambdified expression for the complete amplitude model, as well as lambdified expressions that are to be plugged in to its arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def componentwise_lambdified(*args):\n",
    "    \"\"\"Lambdified amplitude model, recombined from its amplitude components.\n",
    "\n",
    "    .. warning:: Order of the ``args`` has to be the same as that\n",
    "        of the ``args`` of the lambdified amplitude components.\n",
    "    \"\"\"\n",
    "    amplitude_values = []\n",
    "    for amp_symbol in sorted_amplitude_symbols:\n",
    "        np_amplitude = np_amplitudes[amp_symbol]\n",
    "        values = np_amplitude(*args)\n",
    "        amplitude_values.append(values)\n",
    "    return np_amplitude_expr(*amplitude_values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Test with data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, so does all this work? Let's first generate a phase space sample with good-old {mod}`tensorwaves`. We can then use this sample as input to the component-wise lambdified function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "LOGGER.setLevel(logging.ERROR)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "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",
    "rng = TFUniformRealNumberGenerator(seed=0)\n",
    "phsp_sample = generate_phsp(10_000, model.adapter.reaction_info, random_generator=rng)\n",
    "phsp_set = data_converter.transform(phsp_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.hist(phsp_set[\"m_12\"], bins=50, alpha=0.5, density=True)\n",
    "ax.hist(\n",
    "    phsp_set[\"m_12\"],\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    weights=np.array(intensity(phsp_set)),\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "fig.savefig(\"002-histogram-m12.svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983924-9ecf9149-af1d-437b-b5f2-4a73a4d1d81b.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The arguments of the component-wise lambdified amplitude model should be covered by the entries in the phase space set and the provided parameter defaults:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinematic_variable_names = set(phsp_set)\n",
    "parameter_names = {symbol.name for symbol in model.parameter_defaults}\n",
    "free_symbol_names = {symbol.name for symbol in free_symbols}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert free_symbol_names <= kinematic_variable_names ^ parameter_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That allows us to sort the input arrays and parameter defaults so that they can be used as positional argument input to the component-wise lambdified amplitude model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_par_var_values = {\n",
    "    symbol.name: value for symbol, value in model.parameter_defaults.items()\n",
    "}\n",
    "merged_par_var_values.update(phsp_set)\n",
    "args_values = [merged_par_var_values[symbol.name] for symbol in free_symbols]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, here's the result of plugging that back into the component-wise lambdified expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00048765, 0.00033425, 0.00524706, ..., 0.00140122, 0.00714365,\n",
       "       0.00030117])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "componentwise_result = componentwise_lambdified(*args_values)\n",
    "componentwise_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it's indeed the same as that the intensity computed by {mod}`tensorwaves` (direct lambdify):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-7.307471250984975e-11"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tensorwaves_result = np.array(intensity(phsp_set))\n",
    "mean_difference = (componentwise_result - tensorwaves_result).mean()\n",
    "mean_difference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert mean_difference < 1e-9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arbitrary expressions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem with {ref}`report/002:Test with data` is that it requires a {obj}`~ampform.helicity.HelicityModel`. In {mod}`tensorwaves`, we want to work with general {class}`sympy.Expr <sympy.core.expr.Expr>`s though (see [`SympyModel`](https://tensorwaves.readthedocs.io/en/0.2.7/api/tensorwaves.model.html#tensorwaves.model.SympyModel)), where we don't have sub-{attr}`ampform.helicity.HelicityModel.components` available.\n",
    "\n",
    "Instead, we have to split up the lambdifying in a more general way that can handle arbitrary {class}`sympy.core.expr.Expr`s. For that we need:\n",
    "\n",
    "1. A general method of traversing through a SymPy expression tree. This can be done with {doc}`sympy:tutorials/intro-tutorial/manipulation`.\n",
    "2. A **fast** method to estimate the complexity of a model, so that we can decide whether a node in the expression tree is small enough to be lambdified without much runtime. The best measure for complexity is {func}`~sympy.core.function.count_ops` (\"count operations\"), see notes under {doc}`sympy:modules/simplify/simplify`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expression complexity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's tackle 2. first and use the {attr}`HelicityModel.expression <ampform.helicity.HelicityModel.expression>` and its {attr}`~ampform.helicity.HelicityModel.components` that we lambdified earlier on. Here's an overview of the number of operations versus the time it took to lambdify each component:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "keep_output",
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>operations</th>\n",
       "      <th>runtime (s)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>complete model</th>\n",
       "      <td>823</td>\n",
       "      <td>0.980456</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>I1</th>\n",
       "      <td>209</td>\n",
       "      <td>0.279897</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>I2</th>\n",
       "      <td>203</td>\n",
       "      <td>0.235227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>I3</th>\n",
       "      <td>207</td>\n",
       "      <td>0.215937</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>I4</th>\n",
       "      <td>201</td>\n",
       "      <td>0.233635</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A1</th>\n",
       "      <td>103</td>\n",
       "      <td>0.045300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A2</th>\n",
       "      <td>103</td>\n",
       "      <td>0.040710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A3</th>\n",
       "      <td>100</td>\n",
       "      <td>0.039767</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A4</th>\n",
       "      <td>100</td>\n",
       "      <td>0.035684</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A5</th>\n",
       "      <td>102</td>\n",
       "      <td>0.036551</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A6</th>\n",
       "      <td>102</td>\n",
       "      <td>0.036198</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A7</th>\n",
       "      <td>99</td>\n",
       "      <td>0.042208</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>A8</th>\n",
       "      <td>99</td>\n",
       "      <td>0.040205</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                operations  runtime (s)\n",
       "complete model         823     0.980456\n",
       "I1                     209     0.279897\n",
       "I2                     203     0.235227\n",
       "I3                     207     0.215937\n",
       "I4                     201     0.233635\n",
       "A1                     103     0.045300\n",
       "A2                     103     0.040710\n",
       "A3                     100     0.039767\n",
       "A4                     100     0.035684\n",
       "A5                     102     0.036551\n",
       "A6                     102     0.036198\n",
       "A7                      99     0.042208\n",
       "A8                      99     0.040205"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(runtime.values(), index=runtime, columns=[\"runtime (s)\"])\n",
    "operations = [sp.count_ops(model.expression)]\n",
    "operations.extend(sp.count_ops(expr) for expr in intensity_to_symbol)\n",
    "operations.extend(sp.count_ops(expr) for expr in amplitude_to_symbol)\n",
    "df.insert(0, \"operations\", operations)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From this we can already roughly see that the lambdify runtime scales roughly with the number of SymPy operations.\n",
    "\n",
    "To better visualize this, we can lambdify the expressions in {class}`~ampform.dynamics.BlattWeisskopfSquared` for each angular momentums and compute their runtime a number of times with {mod}`timeit`. Note that the {class}`~ampform.dynamics.BlattWeisskopfSquared` becomes increasingly complex the higher the angular momentum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{cases} 1 & \\text{for}\\: L = 0 \\\\\\frac{2 z}{z + 1} & \\text{for}\\: L = 1 \\\\\\frac{13 z^{2}}{9 z + \\left(z - 3\\right)^{2}} & \\text{for}\\: L = 2 \\\\\\frac{277 z^{3}}{z \\left(z - 15\\right)^{2} + \\left(2 z - 5\\right) \\left(18 z - 45\\right)} & \\text{for}\\: L = 3 \\\\\\frac{12746 z^{4}}{25 z \\left(2 z - 21\\right)^{2} + \\left(z^{2} - 45 z + 105\\right)^{2}} & \\text{for}\\: L = 4 \\\\\\frac{998881 z^{5}}{z^{5} + 15 z^{4} + 315 z^{3} + 6300 z^{2} + 99225 z + 893025} & \\text{for}\\: L = 5 \\\\\\frac{118394977 z^{6}}{z^{6} + 21 z^{5} + 630 z^{4} + 18900 z^{3} + 496125 z^{2} + 9823275 z + 108056025} & \\text{for}\\: L = 6 \\\\\\frac{19727003738 z^{7}}{z^{7} + 28 z^{6} + 1134 z^{5} + 47250 z^{4} + 1819125 z^{3} + 58939650 z^{2} + 1404728325 z + 18261468225} & \\text{for}\\: L = 7 \\\\\\frac{4392846440677 z^{8}}{z^{8} + 36 z^{7} + 1890 z^{6} + 103950 z^{5} + 5457375 z^{4} + 255405150 z^{3} + 9833098275 z^{2} + 273922023375 z + 4108830350625} & \\text{for}\\: L = 8 \\end{cases}$"
      ],
      "text/plain": [
       "Piecewise((1, Eq(L, 0)), (2*z/(z + 1), Eq(L, 1)), (13*z**2/(9*z + (z - 3)**2), Eq(L, 2)), (277*z**3/(z*(z - 15)**2 + (2*z - 5)*(18*z - 45)), Eq(L, 3)), (12746*z**4/(25*z*(2*z - 21)**2 + (z**2 - 45*z + 105)**2), Eq(L, 4)), (998881*z**5/(z**5 + 15*z**4 + 315*z**3 + 6300*z**2 + 99225*z + 893025), Eq(L, 5)), (118394977*z**6/(z**6 + 21*z**5 + 630*z**4 + 18900*z**3 + 496125*z**2 + 9823275*z + 108056025), Eq(L, 6)), (19727003738*z**7/(z**7 + 28*z**6 + 1134*z**5 + 47250*z**4 + 1819125*z**3 + 58939650*z**2 + 1404728325*z + 18261468225), Eq(L, 7)), (4392846440677*z**8/(z**8 + 36*z**7 + 1890*z**6 + 103950*z**5 + 5457375*z**4 + 255405150*z**3 + 9833098275*z**2 + 273922023375*z + 4108830350625), Eq(L, 8)))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ampform.dynamics import BlattWeisskopfSquared\n",
    "\n",
    "angular_momentum, z = sp.symbols(\"L z\")\n",
    "BlattWeisskopfSquared(angular_momentum, z).doit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "operations = []\n",
    "runtime = []\n",
    "for angular_momentum in range(9):\n",
    "    ff2 = BlattWeisskopfSquared(angular_momentum, z)\n",
    "    operations.append(sp.count_ops(ff2.doit()))\n",
    "    n_iterations = 10\n",
    "    t = timeit.timeit(\n",
    "        setup=f\"\"\"\n",
    "import sympy as sp\n",
    "from ampform.dynamics import BlattWeisskopfSquared\n",
    "z = sp.Symbol(\"z\")\n",
    "ff2 = BlattWeisskopfSquared({angular_momentum}, z)\n",
    "    \"\"\",\n",
    "        stmt='sp.lambdify(z, ff2.doit(), \"numpy\")',\n",
    "        number=n_iterations,\n",
    "    )\n",
    "    runtime.append(t / n_iterations * 1_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input",
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>operations</th>\n",
       "      <th>runtime (ms)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0.81877</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>1.24712</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>7</td>\n",
       "      <td>1.64094</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>12</td>\n",
       "      <td>2.52622</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>14</td>\n",
       "      <td>2.29422</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>16</td>\n",
       "      <td>1.88900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>19</td>\n",
       "      <td>2.24741</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>22</td>\n",
       "      <td>2.72068</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>25</td>\n",
       "      <td>3.01171</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   operations  runtime (ms)\n",
       "0           0       0.81877\n",
       "1           3       1.24712\n",
       "2           7       1.64094\n",
       "3          12       2.52622\n",
       "4          14       2.29422\n",
       "5          16       1.88900\n",
       "6          19       2.24741\n",
       "7          22       2.72068\n",
       "8          25       3.01171"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"operations\": operations,\n",
    "        \"runtime (ms)\": runtime,\n",
    "    },\n",
    ")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "plt.scatter(x=df[\"operations\"], y=df[\"runtime (ms)\"])\n",
    "ax.set_ylim(bottom=0)\n",
    "ax.set_xlabel(\"operations\")\n",
    "ax.set_ylabel(\"runtime (ms)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "fig.savefig(\"002-runtime-vs-operations.svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983940-3da1d1df-d740-42e4-8a6e-e899c5148034.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Identifying nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now imagine that we don't know anything about the {attr}`~ampform.helicity.HelicityModel.expression` that we created before other than that it is a {class}`sympy.Expr <sympy.core.expr.Expr>`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Approach 1: Generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A first attempt is to use a generator to recursively identify components in the expression that lie within a certain 'complexity' (as computed by {func}`~sympy.core.function.count_ops`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recurse_tree(\n",
    "    expression: sp.Expr, *, min_complexity: int = 0, max_complexity: int\n",
    ") -> Generator[sp.Expr, None, None]:\n",
    "    for arg in expression.args:\n",
    "        complexity = sp.count_ops(arg)\n",
    "        if complexity < max_complexity and complexity > min_complexity:\n",
    "            yield arg\n",
    "        else:\n",
    "            yield from recurse_tree(\n",
    "                arg,\n",
    "                min_complexity=min_complexity,\n",
    "                max_complexity=max_complexity,\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use this generator function to create a mapping of these sub-expressions within the expression tree to {class}`~sympy.core.symbol.Symbol`s. That mapping can then be used in {meth}`~sympy.core.basic.Basic.xreplace` to replace the sub-expressions with those symbols."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{margin}\n",
    "\n",
    "The {meth}`~sympy.core.basic.Basic.xreplace` method is **much faster** than {meth}`~sympy.core.basic.Basic.subs`, because it doesn't do any evaluation.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 314 ms, sys: 135 µs, total: 314 ms\n",
      "Wall time: 313 ms\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left|{f_{0} + f_{1}}\\right|^{2} + \\left|{f_{2} + f_{3}}\\right|^{2} + \\left|{f_{4} + f_{5}}\\right|^{2} + \\left|{f_{6} + f_{7}}\\right|^{2}$"
      ],
      "text/plain": [
       "Abs(f0 + f1)**2 + Abs(f2 + f3)**2 + Abs(f4 + f5)**2 + Abs(f6 + f7)**2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "expression = model.expression.doit()\n",
    "sub_expressions = {}\n",
    "for i, expr in enumerate(recurse_tree(expression, max_complexity=100)):\n",
    "    symbol = sp.Symbol(f\"f{i}\")\n",
    "    complexity = sp.count_ops(expr)\n",
    "    sub_expressions[expr] = symbol\n",
    "expression.xreplace(sub_expressions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Approach 2: Direct substitution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is one problem though: {meth}`~sympy.core.basic.Basic.xreplace` is not accurate for larger expressions. It would therefore be better to directly substitute the sub-expression with a symbol while we loop over the nodes in the expression tree. The following function can do that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_expression(\n",
    "    expression: sp.Expr,\n",
    "    max_complexity: int,\n",
    "    min_complexity: int = 0,\n",
    ") -> tuple[sp.Expr, dict[sp.Symbol, sp.Expr]]:\n",
    "    i = 0\n",
    "    symbol_mapping = {}\n",
    "\n",
    "    def recursive_split(sub_expression: sp.Expr) -> sp.Expr:\n",
    "        nonlocal i\n",
    "        for arg in sub_expression.args:\n",
    "            complexity = sp.count_ops(arg)\n",
    "            if complexity < max_complexity and complexity > min_complexity:\n",
    "                symbol = sp.Symbol(f\"f{i}\")\n",
    "                i += 1\n",
    "                symbol_mapping[symbol] = arg\n",
    "                sub_expression = sub_expression.xreplace({arg: symbol})\n",
    "            else:\n",
    "                new_arg = recursive_split(arg)\n",
    "                sub_expression = sub_expression.xreplace({arg: new_arg})\n",
    "        return sub_expression\n",
    "\n",
    "    top_expression = recursive_split(expression)\n",
    "    return top_expression, symbol_mapping"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And indeed, this is **much faster** than {ref}`report/002:Approach 1: Generator` (it's even possible to parallelize this for loop):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7 µs, sys: 1 µs, total: 8 µs\n",
      "Wall time: 15.5 µs\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "\n",
    "top_expression, sub_expressions = split_expression(expression, max_complexity=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left|{f_{0} + f_{1}}\\right|^{2} + \\left|{f_{2} + f_{3}}\\right|^{2} + \\left|{f_{4} + f_{5}}\\right|^{2} + \\left|{f_{6} + f_{7}}\\right|^{2}$"
      ],
      "text/plain": [
       "Abs(f0 + f1)**2 + Abs(f2 + f3)**2 + Abs(f4 + f5)**2 + Abs(f6 + f7)**2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_expression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}] \\Gamma_{f(0)(980)} m_{f(0)(980)} \\left(\\frac{\\cos{\\left(\\theta_{1+2} \\right)}}{2} + \\frac{1}{2}\\right) e^{i \\phi_{1+2}}}{- \\frac{i \\Gamma_{f(0)(980)} m_{f(0)(980)} \\sqrt{\\frac{\\left(m_{12}^{2} - \\left(m_{1} - m_{2}\\right)^{2}\\right) \\left(m_{12}^{2} - \\left(m_{1} + m_{2}\\right)^{2}\\right)}{m_{12}^{2}}} \\sqrt{m_{f(0)(980)}^{2}}}{\\sqrt{\\frac{\\left(m_{f(0)(980)}^{2} - \\left(m_{1} - m_{2}\\right)^{2}\\right) \\left(m_{f(0)(980)}^{2} - \\left(m_{1} + m_{2}\\right)^{2}\\right)}{m_{f(0)(980)}^{2}}} \\left|{m_{12}}\\right|} - m_{12}^{2} + m_{f(0)(980)}^{2}}$"
      ],
      "text/plain": [
       "C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}]*Gamma_f(0)(980)*m_f(0)(980)*(cos(theta_1+2)/2 + 1/2)*exp(I*phi_1+2)/(-I*Gamma_f(0)(980)*m_f(0)(980)*sqrt((m_12**2 - (m_1 - m_2)**2)*(m_12**2 - (m_1 + m_2)**2)/m_12**2)*sqrt(m_f(0)(980)**2)/(sqrt((m_f(0)(980)**2 - (m_1 - m_2)**2)*(m_f(0)(980)**2 - (m_1 + m_2)**2)/m_f(0)(980)**2)*Abs(m_12)) - m_12**2 + m_f(0)(980)**2)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_expressions[sp.Symbol(\"f0\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lambdify and combine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the machinery to split up arbitrary expressions by complexity, we need to lambdify the top expression as well as each of the sub-expressions and recombine them. The following function can do that and return a recombined {class}`~typing.Callable`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimized_lambdify(\n",
    "    args: Sequence[sp.Symbol],\n",
    "    expr: sp.Expr,\n",
    "    modules: str | None = None,\n",
    "    min_complexity: int = 0,\n",
    "    max_complexity: int = 100,\n",
    ") -> Callable:\n",
    "    top_expression, definitions = split_expression(\n",
    "        expression,\n",
    "        min_complexity=min_complexity,\n",
    "        max_complexity=max_complexity,\n",
    "    )\n",
    "    top_symbols = sorted(definitions, key=lambda s: s.name)\n",
    "    top_lambdified = sp.lambdify(top_symbols, top_expression, modules)\n",
    "    sub_lambdified = [\n",
    "        sp.lambdify(args, definitions[symbol], modules) for symbol in top_symbols\n",
    "    ]\n",
    "\n",
    "    def recombined_function(*args):\n",
    "        new_args = [sub_expr(*args) for sub_expr in sub_lambdified]\n",
    "        return top_lambdified(*new_args)\n",
    "\n",
    "    return recombined_function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the same input values as in {ref}`report/002:Test with data` to check that the resulting lambdified expression results in the same output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8 µs, sys: 1 µs, total: 9 µs\n",
      "Wall time: 17.4 µs\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "\n",
    "treewise_lambdified = optimized_lambdify(free_symbols, expression, \"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00048765, 0.00033425, 0.00524706, ..., 0.00140122, 0.00714365,\n",
       "       0.00030117])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "treewise_result = treewise_lambdified(*args_values)\n",
    "treewise_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And it's indeed the same as that the intensity computed by {mod}`tensorwaves` (direct lambdify):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-7.307471274905997e-11"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean_difference = (treewise_result - tensorwaves_result).mean()\n",
    "mean_difference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "assert mean_difference < 1e-9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now have a look at a slightly more complicated model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "result = qrules.generate_transitions(\n",
    "    initial_state=(\"J/psi(1S)\", [+1]),\n",
    "    final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
    "    allowed_intermediate_particles=[\"f(0)\"],\n",
    "    allowed_interaction_types=[\"strong\", \"EM\"],\n",
    "    formalism_type=\"canonical-helicity\",\n",
    ")\n",
    "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",
    "complex_model = model_builder.generate()\n",
    "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n",
    "graphviz.Source(dot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "graphviz.Source(dot).render(\"002-f0-graph\", format=\"svg\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://user-images.githubusercontent.com/29308176/164983377-f7c3c5a3-edfd-49aa-b449-08ee77cda67f.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This makes it clear that the functions defined in {ref}`report/002:Arbitrary expressions` results in a huge speed-up!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_expression = complex_model.expression.doit()\n",
    "new_free_symbols = sorted(new_expression.free_symbols, key=lambda s: s.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-skip} section\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.57 s, sys: 3.16 ms, total: 4.57 s\n",
      "Wall time: 4.57 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "np_expr = sp.lambdify(new_free_symbols, new_expression)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "keep_output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 261 ms, sys: 87 µs, total: 262 ms\n",
      "Wall time: 260 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "np_expr = optimized_lambdify(new_free_symbols, new_expression)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}